[AIE2P] Support wide vector postinc 2D/3D, and offset load/store

llvm/lib/Target/AIE/AIE2InstrInfo.cpp

@@ -223,6 +223,13 @@ unsigned AIE2InstrInfo::getOffsetMemOpcode(unsigned BaseMemOpcode) const {
   llvm_unreachable("not a generic load/store");
 }
 
+bool AIE2InstrInfo::isGenericOffsetMemOpcode(unsigned Opcode) const {
+  return ((Opcode == AIE2::G_AIE_OFFSET_STORE) ||
+          (Opcode == AIE2::G_AIE_OFFSET_LOAD) ||
+          (Opcode == AIE2::G_AIE_OFFSET_SEXTLOAD) ||
+          (Opcode == AIE2::G_AIE_OFFSET_ZEXTLOAD));
+}
+
 std::optional<unsigned> AIE2InstrInfo::getCombinedPostIncOpcode(
     MachineInstr &BaseMemI, MachineInstr &PostIncI, TypeSize Size) const {
   switch (PostIncI.getOpcode()) {

llvm/lib/Target/AIE/AIE2InstrInfo.h

@@ -83,6 +83,7 @@ class AIE2InstrInfo : public AIE2GenInstrInfo {
   bool isBooleanNoOp(unsigned Opc) const override;
   bool isBooleanNot(unsigned Opc) const override;
   bool isConstStep(const MachineInstr &MI, int64_t &Step) const override;
+  bool isGenericOffsetMemOpcode(unsigned Opcode) const override;
 
   bool verifyGenericInstruction(const MachineInstr &MI,
                                 StringRef &ErrInfo) const override;

llvm/lib/Target/AIE/AIE2InstructionSelector.cpp

@@ -52,7 +52,6 @@ class AIE2InstructionSelector : public AIEBaseInstructionSelector {
                                MachineRegisterInfo &MRI) override;
   Register createSparseRegSequence(Register Vec, Register Mask,
                                    MachineRegisterInfo &MRI);
-  void insertPtrAddForOffset(MachineRegisterInfo &MRI, MachineInstr &MemI);
   void setCtrlRegister(MachineInstr &I, MachineRegisterInfo &MRI,
                        Register CRReg, Register ValueReg);
 
@@ -994,37 +993,6 @@ bool AIE2InstructionSelector::selectG_AIE_LOAD_UNPACK(
   return constrainSelectedInstRegOperands(*NewInstr.getInstr(), TII, TRI, RBI);
 }
 
-void AIE2InstructionSelector::insertPtrAddForOffset(MachineRegisterInfo &MRI,
-                                                    MachineInstr &MemI) {
-  // The offset is not an immediate or the immediate does not fit the immediate
-  // range. Instruction select PTR_ADD for the splitting of instruction. E.g.:
-  // $x0 = G_AIE_OFFSET_LOAD %ptr, %offset has to be selected to
-  // %new_ptr = PTR_ADD %ptr, %offset
-  // $wh0 = VLDA_dmw_lda_w_ag_idx_imm %new_ptr, #32
-  // $wl0 = VLDA_dmw_lda_w_ag_idx_imm %new_ptr, #0
-
-  // This function only gets called for G_AIE_OFFSET_LOAD AND G_AIE_OFFSET_STORE
-  // Both instruction have the pointer and the offset in the same operands
-  assert((MemI.getOpcode() == AIE2::G_AIE_OFFSET_LOAD ||
-          MemI.getOpcode() == AIE2::G_AIE_OFFSET_STORE) &&
-         "Unexpected instruction in instrPtrAddForOffset");
-  const unsigned PointerRegIndex = 1;
-  const unsigned OffsetRegIndex = 2;
-
-  Register NewPtrReg =
-      MRI.cloneVirtualRegister(MemI.getOperand(PointerRegIndex).getReg());
-  MachineInstrBuilder NewPtr =
-      MIB.buildInstr(TargetOpcode::G_PTR_ADD)
-          .addDef(NewPtrReg)
-          .addReg(MemI.getOperand(PointerRegIndex).getReg())
-          .addReg(MemI.getOperand(OffsetRegIndex).getReg());
-
-  if (!selectImpl(*NewPtr.getInstr(), *CoverageInfo))
-    llvm_unreachable("Unexpected failure selecting G_PTR_ADD");
-
-  MemI.getOperand(PointerRegIndex).setReg(NewPtrReg);
-}
-
 std::optional<LoadStoreOpcodes>
 AIE2InstructionSelector::getCombinedOpcodeSRSUPS(const MachineInstr &MemOp,
                                                  const MachineInstr &CombOp,

llvm/lib/Target/AIE/AIEBaseInstrInfo.h

@@ -232,6 +232,8 @@ struct AIEBaseInstrInfo : public TargetInstrInfo {
     return false;
   }
 
+  virtual bool isGenericOffsetMemOpcode(unsigned Opcode) const { return false; }
+
   // Used for Load/Store combiners
   virtual unsigned getOffsetMemOpcode(unsigned BaseMemOpcode) const {
     llvm_unreachable("Target didn't implement getOffsetMemOpcode");

llvm/lib/Target/AIE/AIEBaseInstructionSelector.cpp

@@ -755,3 +755,33 @@ void AIEBaseInstructionSelector::makeDeadMI(MachineInstr &MI,
     Def->setReg(NewReg);
   }
 }
+
+void AIEBaseInstructionSelector::insertPtrAddForOffset(MachineRegisterInfo &MRI,
+                                                       MachineInstr &MemI) {
+  // The offset is not an immediate or the immediate does not fit the immediate
+  // range. Instruction select PTR_ADD for the splitting of instruction. E.g.:
+  // $x0 = G_AIE_OFFSET_LOAD %ptr, %offset has to be selected to
+  // %new_ptr = PTR_ADD %ptr, %offset
+  // $wh0 = VLDA_dmw_lda_w_ag_idx_imm %new_ptr, #32
+  // $wl0 = VLDA_dmw_lda_w_ag_idx_imm %new_ptr, #0
+
+  // This function only gets called for G_AIE_OFFSET_LOAD AND G_AIE_OFFSET_STORE
+  // Both instruction have the pointer and the offset in the same operands
+  assert(TII.isGenericOffsetMemOpcode(MemI.getOpcode()) &&
+         "Unexpected instruction in instrPtrAddForOffset");
+  const unsigned PointerRegIndex = 1;
+  const unsigned OffsetRegIndex = 2;
+
+  Register NewPtrReg =
+      MRI.cloneVirtualRegister(MemI.getOperand(PointerRegIndex).getReg());
+  MachineInstrBuilder NewPtr =
+      MIB.buildInstr(TargetOpcode::G_PTR_ADD)
+          .addDef(NewPtrReg)
+          .addReg(MemI.getOperand(PointerRegIndex).getReg())
+          .addReg(MemI.getOperand(OffsetRegIndex).getReg());
+
+  if (!selectImpl(*NewPtr.getInstr(), *CoverageInfo))
+    llvm_unreachable("Unexpected failure selecting G_PTR_ADD");
+
+  MemI.getOperand(PointerRegIndex).setReg(NewPtrReg);
+}

llvm/lib/Target/AIE/AIEBaseInstructionSelector.h

@@ -157,6 +157,8 @@ class AIEBaseInstructionSelector : public InstructionSelector {
   bool selectVSUB_MIN_MAX(MachineInstr &I, MachineRegisterInfo &MRI,
                           MachineIRBuilder &MIB);
 
+  void insertPtrAddForOffset(MachineRegisterInfo &MRI, MachineInstr &MemI);
+
 protected:
   void makeDeadMI(MachineInstr &MI, MachineRegisterInfo &MRI);
   virtual std::optional<AddressingModeInfo>

llvm/lib/Target/AIE/AIECombine.td

@@ -203,5 +203,7 @@ def AIE2PostLegalizerCustomCombiner
 
 def AIE2PPostLegalizerCustomCombiner
     : GICombiner<"AIE2PPostLegalizerCustomCombinerImpl", [ combine_load_store_increment,
-                                                          combine_add_vector_elt_undef ]> {
+                                                           combine_offset_load_store_ptradd,
+                                                           combine_offset_load_store_share_ptradd,
+                                                           combine_add_vector_elt_undef ]> {
 }

llvm/lib/Target/AIE/AIECombinerHelper.cpp

@@ -58,7 +58,11 @@ MachineInstr *findPreIncMatch(MachineInstr &MemI, MachineRegisterInfo &MRI,
                               const AIEBaseInstrInfo &TII) {
   // This is currently done with patterns in instruction selection.
   // No need to do it here.
-  if (MRI.getType(MemI.getOperand(0).getReg()).getSizeInBits() >= 1024)
+  MachineFunction &MF = *MemI.getMF();
+  const Triple &TT = MF.getTarget().getTargetTriple();
+  const unsigned VecSize =
+      MRI.getType(MemI.getOperand(0).getReg()).getSizeInBits();
+  if ((TT.isAIE2() && VecSize >= 1024) || (TT.isAIE2P() && VecSize > 2048))
     return nullptr;
   if (!EnableOffsetCombine)
     return nullptr;
@@ -320,9 +324,12 @@ MachineInstr *findPostIncMatch(MachineInstr &MemI, MachineRegisterInfo &MRI,
                                const AIEBaseInstrInfo &TII) {
   if (!EnablePostIncCombine)
     return nullptr;
-  if (MRI.getType(MemI.getOperand(0).getReg()).getSizeInBits() >= 1024)
+  MachineFunction &MF = *MemI.getMF();
+  const Triple &TT = MF.getTarget().getTargetTriple();
+  const unsigned VecSize =
+      MRI.getType(MemI.getOperand(0).getReg()).getSizeInBits();
+  if ((TT.isAIE2() && VecSize >= 1024) || (TT.isAIE2P() && VecSize > 2048))
     return nullptr;
-
   Register Addr = MemI.getOperand(1).getReg();
   for (auto &PtrInc : MRI.use_nodbg_instructions(Addr)) {
     if (MemI.getParent() != PtrInc.getParent())

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

@@ -253,6 +253,13 @@ unsigned AIE2PInstrInfo::getOffsetMemOpcode(unsigned BaseMemOpcode) const {
   llvm_unreachable("not a generic load/store");
 }
 
+bool AIE2PInstrInfo::isGenericOffsetMemOpcode(unsigned Opcode) const {
+  return ((Opcode == AIE2P::G_AIE_OFFSET_STORE) ||
+          (Opcode == AIE2P::G_AIE_OFFSET_LOAD) ||
+          (Opcode == AIE2P::G_AIE_OFFSET_SEXTLOAD) ||
+          (Opcode == AIE2P::G_AIE_OFFSET_ZEXTLOAD));
+}
+
 std::optional<unsigned> AIE2PInstrInfo::getCombinedPostIncOpcode(
     MachineInstr &BaseMemI, MachineInstr &PostIncI, TypeSize Size) const {
   switch (PostIncI.getOpcode()) {

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

@@ -83,6 +83,7 @@ class AIE2PInstrInfo : public AIE2PGenInstrInfo {
   bool isBooleanNoOp(unsigned Opc) const override;
   bool isBooleanNot(unsigned Opc) const override;
   bool isConstStep(const MachineInstr &MI, int64_t &Step) const override;
+  bool isGenericOffsetMemOpcode(unsigned Opcode) const override;
 
   bool verifyGenericInstruction(const MachineInstr &MI,
                                 StringRef &ErrInfo) const override;