[CodeGen] Add getTgtMemIntrinsic overload for multiple memory operands (NFC)

llvm/include/llvm/CodeGen/GlobalISel/IRTranslator.h

@@ -299,7 +299,7 @@ class IRTranslator : public MachineFunctionPass {
 
   bool translateIntrinsic(
       const CallBase &CB, Intrinsic::ID ID, MachineIRBuilder &MIRBuilder,
-      const TargetLowering::IntrinsicInfo *TgtMemIntrinsicInfo = nullptr);
+      ArrayRef<TargetLowering::IntrinsicInfo> TgtMemIntrinsicInfos = {});
 
   /// When an invoke or a cleanupret unwinds to the next EH pad, there are
   /// many places it could ultimately go. In the IR, we have a single unwind

llvm/include/llvm/CodeGen/SelectionDAG.h

@@ -438,10 +438,18 @@ class SelectionDAG {
 
   template <typename SDNodeTy>
   static uint16_t getSyntheticNodeSubclassData(unsigned Opc, unsigned Order,
-                                                SDVTList VTs, EVT MemoryVT,
-                                                MachineMemOperand *MMO) {
+                                               SDVTList VTs, EVT MemoryVT,
+                                               MachineMemOperand *MMO) {
     return SDNodeTy(Opc, Order, DebugLoc(), VTs, MemoryVT, MMO)
-         .getRawSubclassData();
+        .getRawSubclassData();
+  }
+
+  template <typename SDNodeTy>
+  static uint16_t getSyntheticNodeSubclassData(
+      unsigned Opc, unsigned Order, SDVTList VTs, EVT MemoryVT,
+      PointerUnion<MachineMemOperand *, MachineMemOperand **> MemRefs) {
+    return SDNodeTy(Opc, Order, DebugLoc(), VTs, MemoryVT, MemRefs)
+        .getRawSubclassData();
   }
 
   void createOperands(SDNode *Node, ArrayRef<SDValue> Vals);
@@ -1481,6 +1489,12 @@ class SelectionDAG {
                                        SDVTList VTList, ArrayRef<SDValue> Ops,
                                        EVT MemVT, MachineMemOperand *MMO);
 
+  /// getMemIntrinsicNode - Creates a MemIntrinsicNode with multiple MMOs.
+  LLVM_ABI SDValue getMemIntrinsicNode(unsigned Opcode, const SDLoc &dl,
+                                       SDVTList VTList, ArrayRef<SDValue> Ops,
+                                       EVT MemVT,
+                                       ArrayRef<MachineMemOperand *> MMOs);
+
   /// Creates a LifetimeSDNode that starts (`IsStart==true`) or ends
   /// (`IsStart==false`) the lifetime of the `FrameIndex`.
   LLVM_ABI SDValue getLifetimeNode(bool IsStart, const SDLoc &dl, SDValue Chain,

llvm/include/llvm/CodeGen/SelectionDAGNodes.h

@@ -1411,19 +1411,26 @@ class MemSDNode : public SDNode {
   EVT MemoryVT;
 
 protected:
-  /// Memory reference information.
-  MachineMemOperand *MMO;
+  /// Memory reference information. Must always have at least one MMO.
+  /// - MachineMemOperand*: exactly 1 MMO (common case)
+  /// - MachineMemOperand**: pointer to array, size at offset -1
+  PointerUnion<MachineMemOperand *, MachineMemOperand **> MemRefs;
 
 public:
-  LLVM_ABI MemSDNode(unsigned Opc, unsigned Order, const DebugLoc &dl,
-                     SDVTList VTs, EVT memvt, MachineMemOperand *MMO);
+  /// Constructor that supports single or multiple MMOs. For single MMO, pass
+  /// the MMO pointer directly. For multiple MMOs, pre-allocate storage with
+  /// count at offset -1 and pass pointer to array.
+  LLVM_ABI
+  MemSDNode(unsigned Opc, unsigned Order, const DebugLoc &dl, SDVTList VTs,
+            EVT memvt,
+            PointerUnion<MachineMemOperand *, MachineMemOperand **> memrefs);
 
-  bool readMem() const { return MMO->isLoad(); }
-  bool writeMem() const { return MMO->isStore(); }
+  bool readMem() const { return getMemOperand()->isLoad(); }
+  bool writeMem() const { return getMemOperand()->isStore(); }
 
   /// Returns alignment and volatility of the memory access
-  Align getBaseAlign() const { return MMO->getBaseAlign(); }
-  Align getAlign() const { return MMO->getAlign(); }
+  Align getBaseAlign() const { return getMemOperand()->getBaseAlign(); }
+  Align getAlign() const { return getMemOperand()->getAlign(); }
 
   /// Return the SubclassData value, without HasDebugValue. This contains an
   /// encoding of the volatile flag, as well as bits used by subclasses. This
@@ -1450,49 +1457,78 @@ class MemSDNode : public SDNode {
   bool isInvariant() const { return MemSDNodeBits.IsInvariant; }
 
   // Returns the offset from the location of the access.
-  int64_t getSrcValueOffset() const { return MMO->getOffset(); }
+  int64_t getSrcValueOffset() const { return getMemOperand()->getOffset(); }
 
   /// Returns the AA info that describes the dereference.
-  AAMDNodes getAAInfo() const { return MMO->getAAInfo(); }
+  AAMDNodes getAAInfo() const { return getMemOperand()->getAAInfo(); }
 
   /// Returns the Ranges that describes the dereference.
-  const MDNode *getRanges() const { return MMO->getRanges(); }
+  const MDNode *getRanges() const { return getMemOperand()->getRanges(); }
 
   /// Returns the synchronization scope ID for this memory operation.
-  SyncScope::ID getSyncScopeID() const { return MMO->getSyncScopeID(); }
+  SyncScope::ID getSyncScopeID() const {
+    return getMemOperand()->getSyncScopeID();
+  }
 
   /// Return the atomic ordering requirements for this memory operation. For
   /// cmpxchg atomic operations, return the atomic ordering requirements when
   /// store occurs.
   AtomicOrdering getSuccessOrdering() const {
-    return MMO->getSuccessOrdering();
+    return getMemOperand()->getSuccessOrdering();
   }
 
   /// Return a single atomic ordering that is at least as strong as both the
   /// success and failure orderings for an atomic operation.  (For operations
   /// other than cmpxchg, this is equivalent to getSuccessOrdering().)
-  AtomicOrdering getMergedOrdering() const { return MMO->getMergedOrdering(); }
+  AtomicOrdering getMergedOrdering() const {
+    return getMemOperand()->getMergedOrdering();
+  }
 
   /// Return true if the memory operation ordering is Unordered or higher.
-  bool isAtomic() const { return MMO->isAtomic(); }
+  bool isAtomic() const { return getMemOperand()->isAtomic(); }
 
   /// Returns true if the memory operation doesn't imply any ordering
   /// constraints on surrounding memory operations beyond the normal memory
   /// aliasing rules.
-  bool isUnordered() const { return MMO->isUnordered(); }
+  bool isUnordered() const { return getMemOperand()->isUnordered(); }
 
   /// Returns true if the memory operation is neither atomic or volatile.
   bool isSimple() const { return !isAtomic() && !isVolatile(); }
 
   /// Return the type of the in-memory value.
   EVT getMemoryVT() const { return MemoryVT; }
 
-  /// Return a MachineMemOperand object describing the memory
+  /// Return the unique MachineMemOperand object describing the memory
   /// reference performed by operation.
-  MachineMemOperand *getMemOperand() const { return MMO; }
+  /// Asserts if multiple MMOs are present - use memoperands() instead.
+  MachineMemOperand *getMemOperand() const {
+    assert(!isa<MachineMemOperand **>(MemRefs) &&
+           "Use memoperands() for nodes with multiple memory operands");
+    return cast<MachineMemOperand *>(MemRefs);
+  }
+
+  /// Return the number of memory operands.
+  size_t getNumMemOperands() const {
+    if (isa<MachineMemOperand *>(MemRefs))
+      return 1;
+    MachineMemOperand **Array = cast<MachineMemOperand **>(MemRefs);
+    return reinterpret_cast<size_t *>(Array)[-1];
+  }
+
+  /// Return true if this node has exactly one memory operand.
+  bool hasUniqueMemOperand() const { return isa<MachineMemOperand *>(MemRefs); }
+
+  /// Return the memory operands for this node.
+  ArrayRef<MachineMemOperand *> memoperands() const {
+    if (isa<MachineMemOperand *>(MemRefs))
+      return ArrayRef(MemRefs.getAddrOfPtr1(), 1);
+    MachineMemOperand **Array = cast<MachineMemOperand **>(MemRefs);
+    size_t Count = reinterpret_cast<size_t *>(Array)[-1];
+    return ArrayRef(Array, Count);
+  }
 
   const MachinePointerInfo &getPointerInfo() const {
-    return MMO->getPointerInfo();
+    return getMemOperand()->getPointerInfo();
   }
 
   /// Return the address space for the associated pointer
@@ -1501,19 +1537,35 @@ class MemSDNode : public SDNode {
   }
 
   /// Update this MemSDNode's MachineMemOperand information
-  /// to reflect the alignment of NewMMO, if it has a greater alignment.
+  /// to reflect the alignment of NewMMOs, if they have greater alignment.
   /// This must only be used when the new alignment applies to all users of
-  /// this MachineMemOperand.
-  void refineAlignment(const MachineMemOperand *NewMMO) {
-    MMO->refineAlignment(NewMMO);
+  /// these MachineMemOperands. The NewMMOs array must parallel memoperands().
+  void refineAlignment(ArrayRef<MachineMemOperand *> NewMMOs) {
+    ArrayRef<MachineMemOperand *> MMOs = memoperands();
+    assert(NewMMOs.size() == MMOs.size() && "MMO count mismatch");
+    for (auto [MMO, NewMMO] : zip(MMOs, NewMMOs))
+      MMO->refineAlignment(NewMMO);
+  }
+
+  void refineAlignment(MachineMemOperand *NewMMO) {
+    refineAlignment(ArrayRef(NewMMO));
   }
 
-  void refineRanges(const MachineMemOperand *NewMMO) {
-    // If this node has range metadata that is different than NewMMO, clear the
-    // range metadata.
+  /// Refine range metadata for all MMOs. The NewMMOs array must parallel
+  /// memoperands(). For each pair, if ranges differ, the stored range is
+  /// cleared.
+  void refineRanges(ArrayRef<MachineMemOperand *> NewMMOs) {
+    ArrayRef<MachineMemOperand *> MMOs = memoperands();
+    assert(NewMMOs.size() == MMOs.size() && "MMO count mismatch");
     // FIXME: Union the ranges instead?
-    if (getRanges() && getRanges() != NewMMO->getRanges())
-      MMO->clearRanges();
+    for (auto [MMO, NewMMO] : zip(MMOs, NewMMOs)) {
+      if (MMO->getRanges() && MMO->getRanges() != NewMMO->getRanges())
+        MMO->clearRanges();
+    }
+  }
+
+  void refineRanges(MachineMemOperand *NewMMO) {
+    refineRanges(ArrayRef(NewMMO));
   }
 
   const SDValue &getChain() const { return getOperand(0); }
@@ -1626,7 +1678,7 @@ class AtomicSDNode : public MemSDNode {
   /// when store does not occur.
   AtomicOrdering getFailureOrdering() const {
     assert(isCompareAndSwap() && "Must be cmpxchg operation");
-    return MMO->getFailureOrdering();
+    return getMemOperand()->getFailureOrdering();
   }
 
   // Methods to support isa and dyn_cast
@@ -1666,9 +1718,11 @@ class AtomicSDNode : public MemSDNode {
 /// opcode (see `SelectionDAGTargetInfo::isTargetMemoryOpcode`).
 class MemIntrinsicSDNode : public MemSDNode {
 public:
-  MemIntrinsicSDNode(unsigned Opc, unsigned Order, const DebugLoc &dl,
-                     SDVTList VTs, EVT MemoryVT, MachineMemOperand *MMO)
-      : MemSDNode(Opc, Order, dl, VTs, MemoryVT, MMO) {
+  MemIntrinsicSDNode(
+      unsigned Opc, unsigned Order, const DebugLoc &dl, SDVTList VTs,
+      EVT MemoryVT,
+      PointerUnion<MachineMemOperand *, MachineMemOperand **> MemRefs)
+      : MemSDNode(Opc, Order, dl, VTs, MemoryVT, MemRefs) {
     SDNodeBits.IsMemIntrinsic = true;
   }
 

llvm/include/llvm/CodeGen/TargetLowering.h

@@ -1244,15 +1244,32 @@ class LLVM_ABI TargetLoweringBase {
   };
 
   /// Given an intrinsic, checks if on the target the intrinsic will need to map
-  /// to a MemIntrinsicNode (touches memory). If this is the case, it returns
-  /// true and store the intrinsic information into the IntrinsicInfo that was
-  /// passed to the function.
+  /// to a MemIntrinsicNode (touches memory). If this is the case, it stores
+  /// the intrinsic information into the IntrinsicInfo vector passed to the
+  /// function. The vector may contain multiple entries for intrinsics that
+  /// access multiple memory locations.
+  virtual void getTgtMemIntrinsic(SmallVectorImpl<IntrinsicInfo> &Infos,
+                                  const CallBase &I, MachineFunction &MF,
+                                  unsigned Intrinsic) const {
+    // The default implementation forwards to the legacy single-info overload
+    // for compatibility.
+    IntrinsicInfo Info;
+    if (getTgtMemIntrinsic(Info, I, MF, Intrinsic))
+      Infos.push_back(Info);
+  }
+
+protected:
+  /// This is a legacy single-info overload. New code should override the
+  /// SmallVectorImpl overload instead to support multiple memory operands.
+  ///
+  /// TODO: Remove this once the refactoring is complete.
   virtual bool getTgtMemIntrinsic(IntrinsicInfo &, const CallBase &,
                                   MachineFunction &,
                                   unsigned /*Intrinsic*/) const {
     return false;
   }
 
+public:
   /// Returns true if the target can instruction select the specified FP
   /// immediate natively. If false, the legalizer will materialize the FP
   /// immediate as a load from a constant pool.

llvm/lib/CodeGen/GlobalISel/IRTranslator.cpp

@@ -2819,20 +2819,16 @@ bool IRTranslator::translateCall(const User &U, MachineIRBuilder &MIRBuilder) {
   if (translateKnownIntrinsic(CI, ID, MIRBuilder))
     return true;
 
-  TargetLowering::IntrinsicInfo Info;
-  bool IsTgtMemIntrinsic = TLI->getTgtMemIntrinsic(Info, CI, *MF, ID);
+  SmallVector<TargetLowering::IntrinsicInfo> Infos;
+  TLI->getTgtMemIntrinsic(Infos, CI, *MF, ID);
 
-  return translateIntrinsic(CI, ID, MIRBuilder,
-                            IsTgtMemIntrinsic ? &Info : nullptr);
+  return translateIntrinsic(CI, ID, MIRBuilder, Infos);
 }
 
 /// Translate a call or callbr to an intrinsic.
-/// Depending on whether TLI->getTgtMemIntrinsic() is true, TgtMemIntrinsicInfo
-/// is a pointer to the correspondingly populated IntrinsicInfo object.
-/// Otherwise, this pointer is null.
 bool IRTranslator::translateIntrinsic(
     const CallBase &CB, Intrinsic::ID ID, MachineIRBuilder &MIRBuilder,
-    const TargetLowering::IntrinsicInfo *TgtMemIntrinsicInfo) {
+    ArrayRef<TargetLowering::IntrinsicInfo> TgtMemIntrinsicInfos) {
   ArrayRef<Register> ResultRegs;
   if (!CB.getType()->isVoidTy())
     ResultRegs = getOrCreateVRegs(CB);
@@ -2874,30 +2870,25 @@ bool IRTranslator::translateIntrinsic(
     }
   }
 
-  // Add a MachineMemOperand if it is a target mem intrinsic.
-  if (TgtMemIntrinsicInfo) {
-    const Function *F = CB.getCalledFunction();
+  // Add MachineMemOperands for each memory access described by the target.
+  for (const auto &Info : TgtMemIntrinsicInfos) {
+    Align Alignment = Info.align.value_or(
+        DL->getABITypeAlign(Info.memVT.getTypeForEVT(CB.getContext())));
+    LLT MemTy = Info.memVT.isSimple()
+                    ? getLLTForMVT(Info.memVT.getSimpleVT())
+                    : LLT::scalar(Info.memVT.getStoreSizeInBits());
 
-    Align Alignment = TgtMemIntrinsicInfo->align.value_or(DL->getABITypeAlign(
-        TgtMemIntrinsicInfo->memVT.getTypeForEVT(F->getContext())));
-    LLT MemTy =
-        TgtMemIntrinsicInfo->memVT.isSimple()
-            ? getLLTForMVT(TgtMemIntrinsicInfo->memVT.getSimpleVT())
-            : LLT::scalar(TgtMemIntrinsicInfo->memVT.getStoreSizeInBits());
-
-    // TODO: We currently just fallback to address space 0 if getTgtMemIntrinsic
-    //       didn't yield anything useful.
+    // TODO: We currently just fallback to address space 0 if
+    // getTgtMemIntrinsic didn't yield anything useful.
     MachinePointerInfo MPI;
-    if (TgtMemIntrinsicInfo->ptrVal) {
-      MPI = MachinePointerInfo(TgtMemIntrinsicInfo->ptrVal,
-                               TgtMemIntrinsicInfo->offset);
-    } else if (TgtMemIntrinsicInfo->fallbackAddressSpace) {
-      MPI = MachinePointerInfo(*TgtMemIntrinsicInfo->fallbackAddressSpace);
+    if (Info.ptrVal) {
+      MPI = MachinePointerInfo(Info.ptrVal, Info.offset);
+    } else if (Info.fallbackAddressSpace) {
+      MPI = MachinePointerInfo(*Info.fallbackAddressSpace);
     }
     MIB.addMemOperand(MF->getMachineMemOperand(
-        MPI, TgtMemIntrinsicInfo->flags, MemTy, Alignment, CB.getAAMetadata(),
-        /*Ranges=*/nullptr, TgtMemIntrinsicInfo->ssid,
-        TgtMemIntrinsicInfo->order, TgtMemIntrinsicInfo->failureOrder));
+        MPI, Info.flags, MemTy, Alignment, CB.getAAMetadata(),
+        /*Ranges=*/nullptr, Info.ssid, Info.order, Info.failureOrder));
   }
 
   if (CB.isConvergent()) {

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

@@ -1208,7 +1208,7 @@ bool DAGCombiner::reassociationCanBreakAddressingModePattern(unsigned Opc,
 
     for (SDNode *Node : N->users()) {
       auto *LoadStore = dyn_cast<MemSDNode>(Node);
-      if (!LoadStore)
+      if (!LoadStore || !LoadStore->hasUniqueMemOperand())
         return false;
 
       // Is x[offset2] a legal addressing mode? If so then