[AIE] Add combine_lag_ptr_add to AIE2PS and AIE2P

llvm/lib/Target/AIE/AIEBaseInstrInfo.h

@@ -63,6 +63,14 @@ struct OperandRegInfo {
       : Reg(R), RC(RegClass) {}
 };
 
+/// Operand indices that describe a pointer-post-increment instruction's key
+/// operands. All indices are absolute (defs precede uses in MachineInstr).
+struct PtrPostIncOpInfo {
+  int PostPtrDefIdx;  ///< Operand index of the post-increment pointer def.
+  int PrePtrUseIdx;   ///< Operand index of the pre-increment pointer use.
+  int ModifierUseIdx; ///< Operand index of the s20 step/modifier use.
+};
+
 // Structure representing a schedule class variant.
 struct SchedVariantInfo {
   unsigned SchedClass;
@@ -344,6 +352,12 @@ struct AIEBaseInstrInfo : public TargetInstrInfo {
   virtual unsigned getGenericPostIncStoreOpcode() const {
     llvm_unreachable("Target didn't implement getGenericPostIncStoreOpcode!");
   }
+  /// Map a POSTINC opcode to its key operand indices, or std::nullopt if the
+  /// opcode is not a pointer-post-increment instruction for this target.
+  virtual std::optional<PtrPostIncOpInfo>
+  getPtrPostIncOpInfo(unsigned Opcode) const {
+    return std::nullopt;
+  }
   /// Return the opcode to be used for subvector extraction.
   virtual unsigned getGenericExtractSubvectorOpcode() const {
     llvm_unreachable(

llvm/lib/Target/AIE/AIECombine.td

@@ -371,6 +371,13 @@ def combine_postincloadstore_ptradd_with_trunc : GICombineRule<
   [{ return matchPostIncLoadStorePtrAddWithTrunc(*${root}, MRI, Helper, (const AIEBaseInstrInfo &)B.getTII(), Observer, ${matchinfo}); }]),
   (apply [{ Helper.applyBuildFnNoErase(*${root}, ${matchinfo}); }])>;
 
+def combine_lag_ptr_add_matchdata : GIDefMatchData<"LagPtrAddMatchInfo">;
+def combine_lag_ptr_add : GICombineRule<
+  (defs root:$root, combine_lag_ptr_add_matchdata:$matchinfo),
+  (match (wip_match_opcode G_PTR_ADD):$root,
+         [{ return matchLagPtrAdd(*${root}, MRI, (const AIEBaseInstrInfo &)B.getTII(), Helper, ${matchinfo}); }]),
+  (apply [{ applyLagPtrAdd(*${root}, MRI, B, Observer, ${matchinfo}); }])>;
+
 def combine_trunc_load : GICombineRule<
   (defs root:$root, build_fn_matchinfo:$matchinfo),
   (match (wip_match_opcode G_LOAD): $root,
@@ -589,6 +596,7 @@ def aie2_postlegalizer_custom_combines : GICombineGroup<[
 
 // Post-legalizer custom combines for AIE2P and more recent targets (AIE2PS, etc.).
 def aie2p_plus_postlegalizer_custom_shared_combines : GICombineGroup<[
+  combine_lag_ptr_add,
   combine_extract_vector_assert_combine,
   combine_extract_broadcast_to_scalar,
   combine_vshift_chain_to_copy,

llvm/lib/Target/AIE/AIECombinerHelper.cpp

@@ -1544,6 +1544,127 @@ bool llvm::matchChainedPtrAddWithNonConstOffsets(MachineInstr &MI,
   return true;
 }
 
+// Prologue for combine_lag_ptr_add:
+//   * bail out when the root's def feeds a real (non-side-effecting) load/store
+//   * require a compile-time constant offset on operand 2
+// On success, returns {PreReg (pre-increment register, operand 1),
+//                      C (constant offset)}; std::nullopt otherwise.
+static std::optional<std::pair<Register, int64_t>>
+getLagRootPreReg(const MachineInstr &MI, MachineRegisterInfo &MRI) {
+  Register DefReg = MI.getOperand(0).getReg();
+  for (const MachineInstr &UseMI : MRI.use_nodbg_instructions(DefReg)) {
+    if (UseMI.mayLoadOrStore() && !UseMI.hasUnmodeledSideEffects())
+      return std::nullopt;
+  }
+  auto MaybeCst =
+      getIConstantVRegValWithLookThrough(MI.getOperand(2).getReg(), MRI);
+  if (!MaybeCst)
+    return std::nullopt;
+  return std::make_pair(MI.getOperand(1).getReg(),
+                        MaybeCst->Value.getSExtValue());
+}
+
+// POSTINC pointer-modifying AIE intrinsics take a "pre" (pre-increment)
+// base pointer and a constant %step (S), and produce a "post"
+// (post-increment) pointer holding pre + S. So with the dataflow:
+//
+//   %post_ptr = G_AIE_..._POSTINC_... %pre_ptr, ..., %step   ; %post = %pre + S
+//   ...                                                       ; (dominates the
+//   G_PTR_ADD)
+//   %add      = G_PTR_ADD %pre_ptr, C
+//
+// the G_PTR_ADD is "lagging": it recomputes %pre + C even though %post
+// = %pre + S is already available. Substituting:
+//
+//   %pre + C = (%post - S) + C = %post + (C - S)
+//
+// lets the G_PTR_ADD reuse %post and free %pre. When C == S the
+// remaining offset is zero, so %post IS the result and we emit a COPY
+// (later folded by register coalescing, which fuses the live ranges
+// of %pre and %post):
+//
+//   %post_ptr = G_AIE_..._POSTINC_... %pre_ptr, ..., %step
+//   ...
+//   %add      = COPY %post_ptr
+//
+// Otherwise emit the adjusted G_PTR_ADD (requires (C - S) to fit in
+// the original offset type, else the rebuilt G_CONSTANT would wrap):
+//
+//   %post_ptr = G_AIE_..._POSTINC_... %pre_ptr, ..., %step
+//   ...
+//   %add      = G_PTR_ADD %post_ptr, (C - S)
+//
+// Bails out when the G_PTR_ADD's def is consumed by a real (non-side-
+// effecting) load/store, since those are rewritten by the memory-op-
+// rooted siblings.
+bool llvm::matchLagPtrAdd(MachineInstr &MI, MachineRegisterInfo &MRI,
+                          const AIEBaseInstrInfo &TII, CombinerHelper &Helper,
+                          LagPtrAddMatchInfo &MatchInfo) {
+  assert(MI.getOpcode() == TargetOpcode::G_PTR_ADD && "Expected G_PTR_ADD");
+
+  auto Root = getLagRootPreReg(MI, MRI);
+  if (!Root)
+    return false;
+  auto [PreReg, C] = *Root;
+
+  for (MachineInstr &UseMI : MRI.use_nodbg_instructions(PreReg)) {
+    // UseMI (the POSTINC) must come before MI (the G_PTR_ADD root) in program
+    // order.
+    if (!Helper.dominates(UseMI, MI))
+      continue;
+    // Find a POSTINC instruction with PreReg as base pointer.
+    auto MaybeInfo = TII.getPtrPostIncOpInfo(UseMI.getOpcode());
+    if (!MaybeInfo)
+      continue;
+    if (!UseMI.getOperand(MaybeInfo->PrePtrUseIdx).isReg() ||
+        UseMI.getOperand(MaybeInfo->PrePtrUseIdx).getReg() != PreReg)
+      continue;
+    // Step must be a compile-time constant.
+    auto MaybeStep = getIConstantVRegValWithLookThrough(
+        UseMI.getOperand(MaybeInfo->ModifierUseIdx).getReg(), MRI);
+    if (!MaybeStep)
+      continue;
+    // Replacement offset NewC = C - step must fit in the original offset
+    // type; otherwise the G_CONSTANT built in apply silently wraps.
+    const int64_t S = MaybeStep->Value.getSExtValue();
+    LLT OffTy = MRI.getType(MI.getOperand(2).getReg());
+    if (!isIntN(OffTy.getSizeInBits(), C - S))
+      continue;
+    MatchInfo = {UseMI.getOperand(MaybeInfo->PostPtrDefIdx).getReg(), S};
+    return true;
+  }
+  return false;
+}
+
+void llvm::applyLagPtrAdd(MachineInstr &MI, MachineRegisterInfo &MRI,
+                          MachineIRBuilder &B, GISelChangeObserver &Observer,
+                          LagPtrAddMatchInfo &MatchInfo) {
+  Register DstReg = MI.getOperand(0).getReg();
+  auto MaybeCst =
+      getIConstantVRegValWithLookThrough(MI.getOperand(2).getReg(), MRI);
+  assert(MaybeCst && "matchLagPtrAdd should have verified constant");
+  int64_t NewC = MaybeCst->Value.getSExtValue() - MatchInfo.Step;
+
+  B.setInsertPt(*MI.getParent(), MI.getIterator());
+  B.setDebugLoc(MI.getDebugLoc());
+
+  if (NewC == 0) {
+    // C == step: %post_ptr IS the result; insert a COPY that copy-propagation
+    // will later eliminate.
+    B.buildCopy(DstReg, MatchInfo.PostReg);
+  } else {
+    // C != step: adjust from %post_ptr by the remaining offset. Preserve the
+    // type of the original offset operand.
+    LLT OffTy = MRI.getType(MI.getOperand(2).getReg());
+    auto NewCst = B.buildConstant(OffTy, NewC);
+    B.buildInstr(TargetOpcode::G_PTR_ADD, {DstReg},
+                 {MatchInfo.PostReg, NewCst});
+  }
+
+  Observer.erasingInstr(MI);
+  MI.eraseFromParent();
+}
+
 /// Match a pattern of G_AIE_POSTINC_LOAD/STORE followed by G_PTR_ADD where both
 /// offsets come from G_TRUNC of s32 values. Combines them by updating the
 /// POSTINC to use the combined offset.

llvm/lib/Target/AIE/AIECombinerHelper.h

@@ -495,6 +495,26 @@ bool matchPostIncLoadStorePtrAddWithTrunc(MachineInstr &MI,
                                           GISelChangeObserver &Observer,
                                           BuildFnTy &MatchInfo);
 
+/// Match data for combine_lag_ptr_add: the post-increment pointer register
+/// produced by a dominating POSTINC instruction and the step constant S it
+/// adds to the pre-increment pointer.
+struct LagPtrAddMatchInfo {
+  Register PostReg; ///< %post_ptr produced by the dominating POSTINC.
+  int64_t Step;     ///< Step constant S added by the POSTINC (post = pre + S).
+};
+
+/// Eliminate lag-register copies by rewriting G_PTR_ADD uses of the
+/// pre-increment pointer register in terms of the POSTINC's post-pointer def.
+/// Rewrites  %dst = G_PTR_ADD %pre, C
+/// to        %dst = COPY %post           (when C == S)
+/// or        %dst = G_PTR_ADD %post, (C-S)  (otherwise)
+bool matchLagPtrAdd(MachineInstr &MI, MachineRegisterInfo &MRI,
+                    const AIEBaseInstrInfo &TII, CombinerHelper &Helper,
+                    LagPtrAddMatchInfo &MatchInfo);
+void applyLagPtrAdd(MachineInstr &MI, MachineRegisterInfo &MRI,
+                    MachineIRBuilder &B, GISelChangeObserver &Observer,
+                    LagPtrAddMatchInfo &MatchInfo);
+
 } // namespace llvm
 
 #endif

llvm/lib/Target/AIE/aie2p/AIE2PInstrInfo.cpp

@@ -1895,6 +1895,43 @@ unsigned AIE2PInstrInfo::getGenericPostIncStoreOpcode() const {
   return AIE2P::G_AIE_POSTINC_STORE;
 }
 
+std::optional<PtrPostIncOpInfo>
+AIE2PInstrInfo::getPtrPostIncOpInfo(unsigned Opcode) const {
+  switch (Opcode) {
+  // %val(0), %ptrDst(1) = G_AIE_POSTINC_LOAD/ZEXTLOAD/SEXTLOAD %ptrSrc(2),
+  // %modifier(3)
+  case AIE2P::G_AIE_POSTINC_LOAD:
+  case AIE2P::G_AIE_POSTINC_ZEXTLOAD:
+  case AIE2P::G_AIE_POSTINC_SEXTLOAD:
+    return PtrPostIncOpInfo{1, 2, 3};
+  // %ptrDst(0) = G_AIE_POSTINC_STORE %val(1), %ptrSrc(2), %modifier(3)
+  case AIE2P::G_AIE_POSTINC_STORE:
+    return PtrPostIncOpInfo{0, 2, 3};
+  // %val(0), %ptrDst(1), %cntDst(2) = G_AIE_POSTINC_2D_LOAD* %ptrSrc(3),
+  // %modifier(4), ...
+  case AIE2P::G_AIE_POSTINC_2D_LOAD:
+  case AIE2P::G_AIE_POSTINC_2D_ZEXTLOAD:
+  case AIE2P::G_AIE_POSTINC_2D_SEXTLOAD:
+    return PtrPostIncOpInfo{1, 3, 4};
+  // %ptrDst(0), %cntDst(1) = G_AIE_POSTINC_2D_STORE %val(2), %ptrSrc(3),
+  // %modifier(4), ...
+  case AIE2P::G_AIE_POSTINC_2D_STORE:
+    return PtrPostIncOpInfo{0, 3, 4};
+  // %val(0), %ptrDst(1), %cnt1Dst(2), %cnt2Dst(3) = G_AIE_POSTINC_3D_LOAD*
+  // %ptrSrc(4), %modifier(5), ...
+  case AIE2P::G_AIE_POSTINC_3D_LOAD:
+  case AIE2P::G_AIE_POSTINC_3D_ZEXTLOAD:
+  case AIE2P::G_AIE_POSTINC_3D_SEXTLOAD:
+    return PtrPostIncOpInfo{1, 4, 5};
+  // %ptrDst(0), %cnt1Dst(1), %cnt2Dst(2) = G_AIE_POSTINC_3D_STORE %val(3),
+  // %ptrSrc(4), %modifier(5), ...
+  case AIE2P::G_AIE_POSTINC_3D_STORE:
+    return PtrPostIncOpInfo{0, 4, 5};
+  default:
+    return std::nullopt;
+  }
+}
+
 unsigned AIE2PInstrInfo::getGenericVSelOpcode() const {
   return AIE2P::G_AIE_VSEL;
 }

llvm/lib/Target/AIE/aie2p/AIE2PInstrInfo.h

@@ -61,6 +61,8 @@ class AIE2PInstrInfo : public AIE2PGenInstrInfo {
   unsigned getGenericBroadcastVectorOpcode() const override;
   unsigned getGenericPostIncLoadOpcode() const override;
   unsigned getGenericPostIncStoreOpcode() const override;
+  std::optional<PtrPostIncOpInfo>
+  getPtrPostIncOpInfo(unsigned Opcode) const override;
   unsigned getGenericVSelOpcode() const override;
   unsigned getGenericVShiftOpcode() const override;
   unsigned getGenericShuffleVectorOpcode() const override;

llvm/lib/Target/AIE/aie2ps/AIE2PSInstrInfo.cpp

@@ -1518,6 +1518,43 @@ unsigned AIE2PSInstrInfo::getGenericPostIncStoreOpcode() const {
   return AIE2PS::G_AIE_POSTINC_STORE;
 }
 
+std::optional<PtrPostIncOpInfo>
+AIE2PSInstrInfo::getPtrPostIncOpInfo(unsigned Opcode) const {
+  switch (Opcode) {
+  // %val(0), %ptrDst(1) = G_AIE_POSTINC_LOAD/ZEXTLOAD/SEXTLOAD %ptrSrc(2),
+  // %modifier(3)
+  case AIE2PS::G_AIE_POSTINC_LOAD:
+  case AIE2PS::G_AIE_POSTINC_ZEXTLOAD:
+  case AIE2PS::G_AIE_POSTINC_SEXTLOAD:
+    return PtrPostIncOpInfo{1, 2, 3};
+  // %ptrDst(0) = G_AIE_POSTINC_STORE %val(1), %ptrSrc(2), %modifier(3)
+  case AIE2PS::G_AIE_POSTINC_STORE:
+    return PtrPostIncOpInfo{0, 2, 3};
+  // %val(0), %ptrDst(1), %cntDst(2) = G_AIE_POSTINC_2D_LOAD* %ptrSrc(3),
+  // %modifier(4), ...
+  case AIE2PS::G_AIE_POSTINC_2D_LOAD:
+  case AIE2PS::G_AIE_POSTINC_2D_ZEXTLOAD:
+  case AIE2PS::G_AIE_POSTINC_2D_SEXTLOAD:
+    return PtrPostIncOpInfo{1, 3, 4};
+  // %ptrDst(0), %cntDst(1) = G_AIE_POSTINC_2D_STORE %val(2), %ptrSrc(3),
+  // %modifier(4), ...
+  case AIE2PS::G_AIE_POSTINC_2D_STORE:
+    return PtrPostIncOpInfo{0, 3, 4};
+  // %val(0), %ptrDst(1), %cnt1Dst(2), %cnt2Dst(3) = G_AIE_POSTINC_3D_LOAD*
+  // %ptrSrc(4), %modifier(5), ...
+  case AIE2PS::G_AIE_POSTINC_3D_LOAD:
+  case AIE2PS::G_AIE_POSTINC_3D_ZEXTLOAD:
+  case AIE2PS::G_AIE_POSTINC_3D_SEXTLOAD:
+    return PtrPostIncOpInfo{1, 4, 5};
+  // %ptrDst(0), %cnt1Dst(1), %cnt2Dst(2) = G_AIE_POSTINC_3D_STORE %val(3),
+  // %ptrSrc(4), %modifier(5), ...
+  case AIE2PS::G_AIE_POSTINC_3D_STORE:
+    return PtrPostIncOpInfo{0, 4, 5};
+  default:
+    return std::nullopt;
+  }
+}
+
 unsigned AIE2PSInstrInfo::getGenericShuffleVectorOpcode() const {
   return AIE2PS::G_AIE_SHUFFLE_VECTOR;
 }

llvm/lib/Target/AIE/aie2ps/AIE2PSInstrInfo.h

@@ -145,6 +145,8 @@ class AIE2PSInstrInfo : public AIE2PSGenInstrInfo {
   unsigned getGenericShuffleVectorOpcode() const override;
   unsigned getGenericPostIncLoadOpcode() const override;
   unsigned getGenericPostIncStoreOpcode() const override;
+  std::optional<PtrPostIncOpInfo>
+  getPtrPostIncOpInfo(unsigned Opcode) const override;
   unsigned getGenericExtractSubvectorOpcode() const override;
   unsigned getGenericIntegerComparisonOpcode() const override;