Update VLD.CONV inst-select to lookthrough bitcast

abhinay-anubola · abhinay-anubola · commit ddb9b1a7ad30 · 2025-02-27T14:07:35.000+05:30
diff --git a/llvm/lib/Target/AIE/AIEBaseInstructionSelector.cpp b/llvm/lib/Target/AIE/AIEBaseInstructionSelector.cpp
@@ -677,7 +677,8 @@ bool AIEBaseInstructionSelector::canCombineCONVLoad(MachineInstr &MemOp,
 bool AIEBaseInstructionSelector::selectG_AIE_LOAD_CONV(
     MachineInstr &CONVI, MachineRegisterInfo &MRI) {
   Register LoadResult = (std::next(CONVI.uses().begin()))->getReg();
-  MachineInstr *LoadOp = getDefIgnoringCopiesAndBitcasts(LoadResult, MRI);
+  MachineInstr *LoadOp =
+      getDefIgnoringCopiesAndBitcasts(LoadResult, false, MRI);
   assert(LoadOp && "Expected SSA.");
 
   // Do not try to combine if one of the load's defs is used by another
diff --git a/llvm/lib/Target/AIE/AIECombinerHelper.cpp b/llvm/lib/Target/AIE/AIECombinerHelper.cpp
@@ -170,8 +170,8 @@ bool isNonCoalesceableUseOf(const MachineInstr &MemI,
       MRI.hasOneNonDBGUse(InBetweenMI.getOperand(0).getReg())) {
     const MachineInstr *CopyOrignMI =
         MRI.getVRegDef(InBetweenMI.getOperand(1).getReg());
-    const MachineInstr *CopyDestMI =
-        &*MRI.use_instr_nodbg_begin(InBetweenMI.getOperand(0).getReg());
+    const MachineInstr *CopyDestMI = getUserIgnoringCopiesAndBitcasts(
+        InBetweenMI.getOperand(0).getReg(), MRI);
     if (CopyOrignMI == &MemI && CopyDestMI == &Dest)
       return false;
   }
@@ -248,6 +248,29 @@ llvm::getDefIgnoringCopiesAndBitcasts(Register Reg, bool AllowMultiUse,
   return DefInstr;
 }
 
+/// Find the use instruction for \p Reg, folding away any trivial copies and
+/// bitcasts. May return nullptr if \p Reg is not a generic virtual register.
+MachineInstr *
+llvm::getUserIgnoringCopiesAndBitcasts(Register Reg,
+                                       const MachineRegisterInfo &MRI) {
+  MachineInstr *User = &*MRI.use_instr_nodbg_begin(Reg);
+
+  auto IsSingleUseCopyOrBitcast = [&](const MachineInstr *MI) {
+    return (MI->isCopy() || (MI->getOpcode() == TargetOpcode::G_BITCAST)) &&
+           MRI.hasOneNonDBGUse(MI->getOperand(0).getReg());
+  };
+
+  auto UseVirtReg = [&](const MachineInstr *MI) {
+    return MI->getOperand(1).getReg().isVirtual();
+  };
+
+  // Stop if we reach an use of a physical register.
+  while (User && IsSingleUseCopyOrBitcast(User) && UseVirtReg(User))
+    User = &*MRI.use_instr_nodbg_begin(User->getOperand(0).getReg());
+
+  return User;
+}
+
 MachineInstr *findLastRegUseInBB(Register Reg, MachineInstr &IgnoreUser,
                                  MachineRegisterInfo &MRI,
                                  CombinerHelper &Helper,
diff --git a/llvm/lib/Target/AIE/AIECombinerHelper.h b/llvm/lib/Target/AIE/AIECombinerHelper.h
@@ -115,6 +115,10 @@ bool canAdvanceOp(MachineInstr &MemI, MachineInstr &Dest,
 /// Reg, has multiple uses.
 MachineInstr *getDefIgnoringCopiesAndBitcasts(Register Reg, bool AllowMultiUse,
                                               const MachineRegisterInfo &MRI);
+/// Find the use instruction for \p Reg, folding away any trivial copies and
+/// bitcasts. May return nullptr if \p Reg is not a generic virtual register.
+MachineInstr *getUserIgnoringCopiesAndBitcasts(Register Reg,
+                                               const MachineRegisterInfo &MRI);
 
 class InstrNode {
   MachineInstr *BaseNode;
diff --git a/llvm/lib/Target/AIE/aie2p/AIE2PInstructionSelector.cpp b/llvm/lib/Target/AIE/aie2p/AIE2PInstructionSelector.cpp
@@ -1304,7 +1304,8 @@ bool AIE2PInstructionSelector::selectG_AIE_LOAD_UPS(MachineInstr &UPSI,
 
   // First use is the G_INTRINSIC_W_SIDE_EFFECTS ID
   Register LoadResult = (std::next(UPSI.uses().begin()))->getReg();
-  MachineInstr *LoadOp = getDefIgnoringCopiesAndBitcasts(LoadResult, MRI);
+  MachineInstr *LoadOp =
+      getDefIgnoringCopiesAndBitcasts(LoadResult, false, MRI);
 
   assert(LoadOp && "Expected SSA.");
 
@@ -3371,7 +3372,8 @@ bool AIE2PInstructionSelector::selectG_AIE_STORE_PACK(
     MachineInstr &StoreI, MachineRegisterInfo &MRI) {
 
   Register PackResult = (StoreI.uses().begin())->getReg();
-  MachineInstr *PackOp = getDefIgnoringCopiesAndBitcasts(PackResult, MRI);
+  MachineInstr *PackOp =
+      getDefIgnoringCopiesAndBitcasts(PackResult, false, MRI);
 
   if (!canCombinePACK(StoreI, *PackOp, MRI))
     return false;
@@ -3451,7 +3453,7 @@ bool AIE2PInstructionSelector::selectG_AIE_STORE_SRS(MachineInstr &StoreI,
                                                      MachineRegisterInfo &MRI) {
 
   Register SrsResult = (StoreI.uses().begin())->getReg();
-  MachineInstr *SrsOp = getDefIgnoringCopiesAndBitcasts(SrsResult, MRI);
+  MachineInstr *SrsOp = getDefIgnoringCopiesAndBitcasts(SrsResult, false, MRI);
 
   assert(SrsOp && "Expected SSA.");
 
@@ -3735,7 +3737,8 @@ bool AIE2PInstructionSelector::selectG_AIE_STORE_CONV(
   // differ.
 
   Register ConvResult = (StoreI.uses().begin())->getReg();
-  MachineInstr *ConvOp = getDefIgnoringCopiesAndBitcasts(ConvResult, MRI);
+  MachineInstr *ConvOp =
+      getDefIgnoringCopiesAndBitcasts(ConvResult, false, MRI);
 
   assert(ConvOp && "Expected SSA.");
 
@@ -3981,7 +3984,8 @@ bool AIE2PInstructionSelector::canCombineUNPACKLoad(MachineInstr &MemOp,
 bool AIE2PInstructionSelector::selectG_AIE_LOAD_UNPACK(
     MachineInstr &UNPACKI, MachineRegisterInfo &MRI) {
   Register LoadResult = (std::next(UNPACKI.uses().begin()))->getReg();
-  MachineInstr *LoadOp = getDefIgnoringCopiesAndBitcasts(LoadResult, MRI);
+  MachineInstr *LoadOp =
+      getDefIgnoringCopiesAndBitcasts(LoadResult, false, MRI);
   // Should we build the instruction at load's position?
   bool ShouldAdvanceOp = false;
 
@@ -4936,7 +4940,8 @@ unsigned int getStoreFifoOpcode(MachineInstr &I) {
 bool AIE2PInstructionSelector::selectVST_FIFO_CONV(MachineInstr &StoreI,
                                                    MachineRegisterInfo &MRI) {
   Register ConvResult = StoreI.getOperand(5).getReg();
-  MachineInstr *ConvOp = getDefIgnoringCopiesAndBitcasts(ConvResult, MRI);
+  MachineInstr *ConvOp =
+      getDefIgnoringCopiesAndBitcasts(ConvResult, false, MRI);
   assert(ConvOp && "Expected SSA.");
 
   if (!canCombineCONV(StoreI, *ConvOp) ||
diff --git a/llvm/test/CodeGen/AIE/aie2/GlobalISel/inst-select-vlda_conv.mir b/llvm/test/CodeGen/AIE/aie2/GlobalISel/inst-select-vlda_conv.mir
@@ -380,3 +380,25 @@ body: |
     %3:accregbank(<8 x s64>) = G_INTRINSIC intrinsic(@llvm.aie2.v16bf16.to.v16accfloat), %2:vregbank(<16 x s16>)
     PseudoRET implicit $lr, implicit %3, implicit %2
 ...
+
+---
+name:            VLDA_CONV_COPY_bitcast
+legalized:       true
+regBankSelected: true
+tracksRegLiveness: true
+body: |
+  bb.0:
+    liveins: $p0, $r0
+    ; CHECK-LABEL: name: VLDA_CONV_COPY_bitcast
+    ; CHECK: liveins: $p0, $r0
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:ep = COPY $p0
+    ; CHECK-NEXT: [[VLDA_CONV_FP32_BF16_ag_idx_imm:%[0-9]+]]:acc512 = VLDA_CONV_FP32_BF16_ag_idx_imm [[COPY]], 0 :: (load (<8 x s32>))
+    ; CHECK-NEXT: PseudoRET implicit $lr, implicit [[VLDA_CONV_FP32_BF16_ag_idx_imm]]
+    %0:ptrregbank(p0) = COPY $p0
+    %1:vregbank(<8 x s32>) = G_LOAD %0:ptrregbank(p0) :: (load (<8 x s32>))
+    %2:vregbank(<8 x s32>) = COPY %1
+    %3:vregbank(<16 x s16>) = G_BITCAST %2(<8 x s32>)
+    %4:accregbank(<8 x s64>) = G_INTRINSIC intrinsic(@llvm.aie2.v16bf16.to.v16accfloat), %3:vregbank(<16 x s16>)
+    PseudoRET implicit $lr, implicit %4
+...
