[SVE][CodeGen] Add the ExtensionType flag to MGATHER

[kmclaughlin]

Summary

Adds the ExtensionType flag, which reflects the LoadExtType of a MaskedGatherSDNode.
Also updated SelectionDAGDumper::print_details so that details of the gather
load (is signed, is scaled & extension type) are printed.

[kmclaughlin]

• Moved changes to SplitVecRes_MGATHER & SplitVecOp_MGATHER into this patch from D91092

[kmclaughlin]

Added DAG combines for the following:

• fold (and (masked_gather x)) -> (zext_masked_gather x)
• fold (sext_inreg (masked_gather x)) -> (sext_masked_gather x)

[sdesmalen]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

938 ↗

(On Diff #304868)

nit: Capitalise to ULL seems more common in LLVM.

955 ↗

(On Diff #304868)

nit: unnecessary curly braces.

5665 ↗

(On Diff #304868)

I'm surprised this doesn't break anything for other targets like X86.

In any case I think this fold (and same for the sext_masked_gather) should be guarded under TargetLoweringInfo::isVectorLoadExtDesirable, so that the target can inspect the operation and type, to make a suitable decision.

11526 ↗

(On Diff #304868)

is !LegalOperations necessary here? The combine for zext_masked_gather doesn't seem to need it.

11527 ↗

(On Diff #304868)

Is this missing a N0.hasOneUse() ?

[kmclaughlin]

• Replaced calls to TLI.isLoadExtLegal with TLI.isVectorLoadExtDesirable in the masked gather combines, which also checks the type of load operation being used
• Removed unnecessary !LegalOperations from the zext_masked_gather combine and added a check for hasOneUse(), matching the sext_masked_gather combine

[sdesmalen]

Hi @kmclaughlin, we just discussed this offline, just sharing my thoughts here as well.

The code-changes seem fine, but I think the patch series is in the wrong order. i.e. D91084 is an optimization for D91092. The better order I think would probably be to:

1. Have https://reviews.llvm.org/D91092 as a standalone patch (you'll need to update the test and remove the reference to MGT->getExtensionType()
2. Add the ExtensionType flag to MGATHER (re-adding the reference to MGT->getExtensionType() in LowerMGATHER, but the tests are probably unchanged).
3. Add the DAGCombines to merge zext/sext into the mgather itself (which then requires updating the tests)

[kmclaughlin]

Removed isConstantSplatVectorMaskForType() & the DAG combines for s/ext_masked_gather (these will be added in a follow up patch)

[kmclaughlin]

Hi @sdesmalen, I've updated the series as suggested and the patches are now in the following order:
D91092 - Lower scalable masked scatters (with references to ExtensionType in LowerMGATHER removed)
D91084 - Add the ExtensionType flag to MGATHER
D92230 - DAG combines for z/sext of a masked gather, adding ExtensionType back into LowerMGATHER

Thanks for reviewing all of these patches!

[sdesmalen]

LGTM!

llvm/include/llvm/CodeGen/SelectionDAG.h

SDValue getMaskedStore(SDValue Chain, const SDLoc &dl, SDValue Val,
                          SDValue Base, SDValue Offset, SDValue Mask, EVT MemVT,
                          MachineMemOperand *MMO, ISD::MemIndexedMode AM,
                          bool IsTruncating = false, bool IsCompressing = false);
   SDValue getIndexedMaskedStore(SDValue OrigStore, const SDLoc &dl,
                                 SDValue Base, SDValue Offset,
                                 ISD::MemIndexedMode AM);
   SDValue getMaskedGather(SDVTList VTs, EVT VT, const SDLoc &dl,
                           ArrayRef<SDValue> Ops, MachineMemOperand *MMO,
-                          ISD::MemIndexType IndexType);
+                          ISD::MemIndexType IndexType, ISD::LoadExtType ExtTy);
   SDValue getMaskedScatter(SDVTList VTs, EVT VT, const SDLoc &dl,
                            ArrayRef<SDValue> Ops, MachineMemOperand *MMO,
                            ISD::MemIndexType IndexType,
                            bool IsTruncating = false);
 
   /// Construct a node to track a Value* through the backend.
   SDValue getSrcValue(const Value *v);
 

llvm/include/llvm/CodeGen/SelectionDAGNodes.h

class LSBaseSDNodeBitfields {
     //   MaskedGatherScatterSDNode => enum ISD::MemIndexType
     uint16_t AddressingMode : 3;
   };
   enum { NumLSBaseSDNodeBits = NumMemSDNodeBits + 3 };
 
   class LoadSDNodeBitfields {
     friend class LoadSDNode;
     friend class MaskedLoadSDNode;
+    friend class MaskedGatherSDNode;
 
     uint16_t : NumLSBaseSDNodeBits;
 
     uint16_t ExtTy : 2; // enum ISD::LoadExtType
     uint16_t IsExpanding : 1;
   };
 
   class StoreSDNodeBitfields {

 
 /// This class is used to represent an MGATHER node
 ///
 class MaskedGatherSDNode : public MaskedGatherScatterSDNode {
 public:
   friend class SelectionDAG;
 
   MaskedGatherSDNode(unsigned Order, const DebugLoc &dl, SDVTList VTs,
-                     EVT MemVT, MachineMemOperand *MMO,
+                     ISD::LoadExtType ETy, EVT MemVT, MachineMemOperand *MMO,
                      ISD::MemIndexType IndexType)
       : MaskedGatherScatterSDNode(ISD::MGATHER, Order, dl, VTs, MemVT, MMO,
-                                  IndexType) {}
+                                  IndexType) {
+    LoadSDNodeBits.ExtTy = ETy;
+  }
 
   const SDValue &getPassThru() const { return getOperand(1); }
 
+  ISD::LoadExtType getExtensionType() const {
+    return ISD::LoadExtType(LoadSDNodeBits.ExtTy);
+  }
+
   static bool classof(const SDNode *N) {
     return N->getOpcode() == ISD::MGATHER;
   }
 };
 
 /// This class is used to represent an MSCATTER node
 ///
 class MaskedScatterSDNode : public MaskedGatherScatterSDNode {

llvm/lib/CodeGen/SelectionDAG/LegalizeIntegerTypes.cpp

 }
 
 SDValue DAGTypeLegalizer::PromoteIntRes_MGATHER(MaskedGatherSDNode *N) {
   EVT NVT = TLI.getTypeToTransformTo(*DAG.getContext(), N->getValueType(0));
   SDValue ExtPassThru = GetPromotedInteger(N->getPassThru());
   assert(NVT == ExtPassThru.getValueType() &&
       "Gather result type and the passThru argument type should be the same");
 
+  ISD::LoadExtType ExtType = N->getExtensionType();
+  if (ExtType == ISD::NON_EXTLOAD)
+    ExtType = ISD::EXTLOAD;
+
   SDLoc dl(N);
   SDValue Ops[] = {N->getChain(), ExtPassThru, N->getMask(), N->getBasePtr(),
                    N->getIndex(), N->getScale() };
   SDValue Res = DAG.getMaskedGather(DAG.getVTList(NVT, MVT::Other),
                                     N->getMemoryVT(), dl, Ops,
-                                    N->getMemOperand(), N->getIndexType());
+                                    N->getMemOperand(), N->getIndexType(),
+                                    ExtType);
   // Legalize the chain result - switch anything that used the old chain to
   // use the new one.
   ReplaceValueWith(SDValue(N, 1), Res.getValue(1));
   return Res;
 }
 
 /// Promote the overflow flag of an overflowing arithmetic node.
 SDValue DAGTypeLegalizer::PromoteIntRes_Overflow(SDNode *N) {

llvm/lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp

void DAGTypeLegalizer::SplitVecRes_MGATHER(MaskedGatherSDNode *MGT,
 
   SDValue Ch = MGT->getChain();
   SDValue Ptr = MGT->getBasePtr();
   SDValue Mask = MGT->getMask();
   SDValue PassThru = MGT->getPassThru();
   SDValue Index = MGT->getIndex();
   SDValue Scale = MGT->getScale();
   Align Alignment = MGT->getOriginalAlign();
+  ISD::LoadExtType ExtType = MGT->getExtensionType();
 
   // Split Mask operand
   SDValue MaskLo, MaskHi;
   if (Mask.getOpcode() == ISD::SETCC) {
     SplitVecRes_SETCC(Mask.getNode(), MaskLo, MaskHi);
   } else {
     if (getTypeAction(Mask.getValueType()) == TargetLowering::TypeSplitVector)
       GetSplitVector(Mask, MaskLo, MaskHi);
void DAGTypeLegalizer::SplitVecRes_MGATHER(MaskedGatherSDNode *MGT,
 
   MachineMemOperand *MMO = DAG.getMachineFunction().getMachineMemOperand(
       MGT->getPointerInfo(), MachineMemOperand::MOLoad,
       MemoryLocation::UnknownSize, Alignment, MGT->getAAInfo(),
       MGT->getRanges());
 
   SDValue OpsLo[] = {Ch, PassThruLo, MaskLo, Ptr, IndexLo, Scale};
   Lo = DAG.getMaskedGather(DAG.getVTList(LoVT, MVT::Other), LoVT, dl, OpsLo,
-                           MMO, MGT->getIndexType());
+                           MMO, MGT->getIndexType(), ExtType);
 
   SDValue OpsHi[] = {Ch, PassThruHi, MaskHi, Ptr, IndexHi, Scale};
   Hi = DAG.getMaskedGather(DAG.getVTList(HiVT, MVT::Other), HiVT, dl, OpsHi,
-                           MMO, MGT->getIndexType());
+                           MMO, MGT->getIndexType(), ExtType);
 
   // Build a factor node to remember that this load is independent of the
   // other one.
   Ch = DAG.getNode(ISD::TokenFactor, dl, MVT::Other, Lo.getValue(1),
                    Hi.getValue(1));
 
   // Legalize the chain result - switch anything that used the old chain to
   // use the new one.
SDValue DAGTypeLegalizer::SplitVecOp_MGATHER(MaskedGatherSDNode *MGT,
 
   SDValue Ch = MGT->getChain();
   SDValue Ptr = MGT->getBasePtr();
   SDValue Index = MGT->getIndex();
   SDValue Scale = MGT->getScale();
   SDValue Mask = MGT->getMask();
   SDValue PassThru = MGT->getPassThru();
   Align Alignment = MGT->getOriginalAlign();
+  ISD::LoadExtType ExtType = MGT->getExtensionType();
 
   SDValue MaskLo, MaskHi;
   if (getTypeAction(Mask.getValueType()) == TargetLowering::TypeSplitVector)
     // Split Mask operand
     GetSplitVector(Mask, MaskLo, MaskHi);
   else
     std::tie(MaskLo, MaskHi) = DAG.SplitVector(Mask, dl);
 
SDValue DAGTypeLegalizer::SplitVecOp_MGATHER(MaskedGatherSDNode *MGT,
 
   MachineMemOperand *MMO = DAG.getMachineFunction().getMachineMemOperand(
       MGT->getPointerInfo(), MachineMemOperand::MOLoad,
       MemoryLocation::UnknownSize, Alignment, MGT->getAAInfo(),
       MGT->getRanges());
 
   SDValue OpsLo[] = {Ch, PassThruLo, MaskLo, Ptr, IndexLo, Scale};
   SDValue Lo = DAG.getMaskedGather(DAG.getVTList(LoVT, MVT::Other), LoVT, dl,
-                                   OpsLo, MMO, MGT->getIndexType());
+                                   OpsLo, MMO, MGT->getIndexType(), ExtType);
 
   SDValue OpsHi[] = {Ch, PassThruHi, MaskHi, Ptr, IndexHi, Scale};
   SDValue Hi = DAG.getMaskedGather(DAG.getVTList(HiVT, MVT::Other), HiVT, dl,
-                                   OpsHi, MMO, MGT->getIndexType());
+                                   OpsHi, MMO, MGT->getIndexType(), ExtType);
 
   // Build a factor node to remember that this load is independent of the
   // other one.
   Ch = DAG.getNode(ISD::TokenFactor, dl, MVT::Other, Lo.getValue(1),
                    Hi.getValue(1));
 
   // Legalize the chain result - switch anything that used the old chain to
   // use the new one.
SDValue DAGTypeLegalizer::WidenVecRes_MGATHER(MaskedGatherSDNode *N) {
   EVT WideIndexVT = EVT::getVectorVT(*DAG.getContext(),
                                      Index.getValueType().getScalarType(),
                                      NumElts);
   Index = ModifyToType(Index, WideIndexVT);
   SDValue Ops[] = { N->getChain(), PassThru, Mask, N->getBasePtr(), Index,
                     Scale };
   SDValue Res = DAG.getMaskedGather(DAG.getVTList(WideVT, MVT::Other),
                                     N->getMemoryVT(), dl, Ops,
-                                    N->getMemOperand(), N->getIndexType());
+                                    N->getMemOperand(), N->getIndexType(),
+                                    N->getExtensionType());
 
   // Legalize the chain result - switch anything that used the old chain to
   // use the new one.
   ReplaceValueWith(SDValue(N, 1), Res.getValue(1));
   return Res;
 }
 
 SDValue DAGTypeLegalizer::WidenVecRes_ScalarOp(SDNode *N) {
SDValue DAGTypeLegalizer::WidenVecOp_MGATHER(SDNode *N, unsigned OpNo) {
 
   // Just widen the index. It's allowed to have extra elements.
   SDValue Index = GetWidenedVector(MG->getIndex());
 
   SDLoc dl(N);
   SDValue Ops[] = {MG->getChain(), DataOp, Mask, MG->getBasePtr(), Index,
                    Scale};
   SDValue Res = DAG.getMaskedGather(MG->getVTList(), MG->getMemoryVT(), dl, Ops,
-                                    MG->getMemOperand(), MG->getIndexType());
+                                    MG->getMemOperand(), MG->getIndexType(),
+                                    MG->getExtensionType());
   ReplaceValueWith(SDValue(N, 1), Res.getValue(1));
   ReplaceValueWith(SDValue(N, 0), Res.getValue(0));
   return SDValue();
 }
 
 SDValue DAGTypeLegalizer::WidenVecOp_MSCATTER(SDNode *N, unsigned OpNo) {
   MaskedScatterSDNode *MSC = cast<MaskedScatterSDNode>(N);
   SDValue DataOp = MSC->getValue();

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

SDValue SelectionDAG::getIndexedMaskedStore(SDValue OrigStore, const SDLoc &dl,
   return getMaskedStore(ST->getChain(), dl, ST->getValue(), Base, Offset,
                         ST->getMask(), ST->getMemoryVT(), ST->getMemOperand(),
                         AM, ST->isTruncatingStore(), ST->isCompressingStore());
 }
 
 SDValue SelectionDAG::getMaskedGather(SDVTList VTs, EVT VT, const SDLoc &dl,
                                       ArrayRef<SDValue> Ops,
                                       MachineMemOperand *MMO,
-                                      ISD::MemIndexType IndexType) {
+                                      ISD::MemIndexType IndexType,
+                                      ISD::LoadExtType ExtTy) {
   assert(Ops.size() == 6 && "Incompatible number of operands");
 
   FoldingSetNodeID ID;
   AddNodeIDNode(ID, ISD::MGATHER, VTs, Ops);
   ID.AddInteger(VT.getRawBits());
   ID.AddInteger(getSyntheticNodeSubclassData<MaskedGatherSDNode>(
-      dl.getIROrder(), VTs, VT, MMO, IndexType));
+      dl.getIROrder(), VTs, ExtTy, VT, MMO, IndexType));
   ID.AddInteger(MMO->getPointerInfo().getAddrSpace());
   void *IP = nullptr;
   if (SDNode *E = FindNodeOrInsertPos(ID, dl, IP)) {
     cast<MaskedGatherSDNode>(E)->refineAlignment(MMO);
     return SDValue(E, 0);
   }
 
   auto *N = newSDNode<MaskedGatherSDNode>(dl.getIROrder(), dl.getDebugLoc(),
-                                          VTs, VT, MMO, IndexType);
+                                          VTs, ExtTy, VT, MMO, IndexType);
   createOperands(N, Ops);
 
   assert(N->getPassThru().getValueType() == N->getValueType(0) &&
          "Incompatible type of the PassThru value in MaskedGatherSDNode");
   assert(N->getMask().getValueType().getVectorNumElements() ==
              N->getValueType(0).getVectorNumElements() &&
          "Vector width mismatch between mask and data");
   assert(N->getIndex().getValueType().getVectorNumElements() >=

llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

void SelectionDAGBuilder::visitMaskedGather(const CallInst &I) {
   if (!UniformBase) {
     Base = DAG.getConstant(0, sdl, TLI.getPointerTy(DAG.getDataLayout()));
     Index = getValue(Ptr);
     IndexType = ISD::SIGNED_SCALED;
     Scale = DAG.getTargetConstant(1, sdl, TLI.getPointerTy(DAG.getDataLayout()));
   }
   SDValue Ops[] = { Root, Src0, Mask, Base, Index, Scale };
   SDValue Gather = DAG.getMaskedGather(DAG.getVTList(VT, MVT::Other), VT, sdl,
-                                       Ops, MMO, IndexType);
+                                       Ops, MMO, IndexType, ISD::NON_EXTLOAD);
 
   PendingLoads.push_back(Gather.getValue(1));
   setValue(&I, Gather);
 }
 
 void SelectionDAGBuilder::visitAtomicCmpXchg(const AtomicCmpXchgInst &I) {
   SDLoc dl = getCurSDLoc();
   AtomicOrdering SuccessOrdering = I.getSuccessOrdering();

llvm/lib/CodeGen/SelectionDAG/SelectionDAGDumper.cpp

else if (const LoadSDNode *LD = dyn_cast<LoadSDNode>(this)) {
     const char *AM = getIndexedModeName(MSt->getAddressingMode());
     if (*AM)
       OS << ", " << AM;
 
     if (MSt->isCompressingStore())
       OS << ", compressing";
 
     OS << ">";
+  } else if (const auto *MGather = dyn_cast<MaskedGatherSDNode>(this)) {
+    OS << "<";
+    printMemOperand(OS, *MGather->getMemOperand(), G);
+
+    bool doExt = true;
+    switch (MGather->getExtensionType()) {
+    default: doExt = false; break;
+    case ISD::EXTLOAD:  OS << ", anyext"; break;
+    case ISD::SEXTLOAD: OS << ", sext"; break;
+    case ISD::ZEXTLOAD: OS << ", zext"; break;
+    }
+    if (doExt)
+      OS << " from " << MGather->getMemoryVT().getEVTString();
+
+    auto Signed = MGather->isIndexSigned() ? "signed" : "unsigned";
+    auto Scaled = MGather->isIndexScaled() ? "scaled" : "unscaled";
+    OS << ", " << Signed << " " << Scaled << " offset";
+
+    OS << ">";
   } else if (const auto *MScatter = dyn_cast<MaskedScatterSDNode>(this)) {
     OS << "<";
     printMemOperand(OS, *MScatter->getMemOperand(), G);
 
     if (MScatter->isTruncatingStore())
       OS << ", trunc to " << MScatter->getMemoryVT().getEVTString();
 
     auto Signed = MScatter->isIndexSigned() ? "signed" : "unsigned";

llvm/lib/Target/X86/X86ISelLowering.cpp

static SDValue rebuildGatherScatter(MaskedGatherScatterSDNode *GorS,
   SDLoc DL(GorS);
 
   if (auto *Gather = dyn_cast<MaskedGatherSDNode>(GorS)) {
     SDValue Ops[] = { Gather->getChain(), Gather->getPassThru(),
                       Gather->getMask(), Base, Index, Scale } ;
     return DAG.getMaskedGather(Gather->getVTList(),
                                Gather->getMemoryVT(), DL, Ops,
                                Gather->getMemOperand(),
-                               Gather->getIndexType());
+                               Gather->getIndexType(),
+                               Gather->getExtensionType());
   }
   auto *Scatter = cast<MaskedScatterSDNode>(GorS);
   SDValue Ops[] = { Scatter->getChain(), Scatter->getValue(),
                     Scatter->getMask(), Base, Index, Scale };
   return DAG.getMaskedScatter(Scatter->getVTList(),
                               Scatter->getMemoryVT(), DL,
                               Ops, Scatter->getMemOperand(),
                               Scatter->getIndexType(),