[GlobalISel][AArch64] Fix contract cross-bank copies with SIMD instructions

[scw]

Summary

contractCrossBankCopyIntoStore() finds the instruction defines the
source register and uses its output to replace the register. There are,
however, instructions that have multiple outputs, e.g. G_UNMERGE_VALUES.
Current implementation hardcodes to operand 0 and has no way of knowing
which output should be used.

This change adds another function to directly return the register that
is the source of the register and use that for folding.

This fixes https://bugs.llvm.org/show_bug.cgi?id=44783

[paquette]

Thanks for fixing this!

Can you attach a MIR testcase based off the one from the bugzilla bug?

llvm/lib/CodeGen/GlobalISel/Utils.cpp

303–312

Instead of a std::pair, can we do something similar to getConstantVRegValWithLookThrough, which returns an Optional struct?

Then when we bail out, we can just return None, which is a little nicer IMO.

[scw]

Updated test.

[scw]

llvm/lib/CodeGen/GlobalISel/Utils.cpp

303–312

Done.

[MaskRay]

llvm/lib/CodeGen/GlobalISel/Utils.cpp

303

namespace {
struct ...
}

per https://llvm.org/docs/CodingStandards.html#anonymous-namespaces

324

{ -> {

[scw]

Update style.

[paquette]

LGTM

[scw]

Submitted as ce9633633c0

llvm/include/llvm/CodeGen/GlobalISel/Utils.h

MachineInstr *getOpcodeDef(unsigned Opcode, Register Reg,
                            const MachineRegisterInfo &MRI);
 
 /// Find the def instruction for \p Reg, folding away any trivial copies. Note
 /// it may still return a COPY, if it changes the type. May return nullptr if \p
 /// Reg is not a generic virtual register.
 MachineInstr *getDefIgnoringCopies(Register Reg,
                                    const MachineRegisterInfo &MRI);
 
+/// Find the source register for \p Reg, folding away any trivial copies. It
+/// will be an output register of the instruction that getDefIgnoringCopies
+/// returns. May return an invalid register if \p Reg is not a generic virtual
+/// register.
+Register getSrcRegIgnoringCopies(Register Reg,
+                                 const MachineRegisterInfo &MRI);
+
 /// Returns an APFloat from Val converted to the appropriate size.
 APFloat getAPFloatFromSize(double Val, unsigned Size);
 
 /// Modify analysis usage so it preserves passes required for the SelectionDAG
 /// fallback.
 void getSelectionDAGFallbackAnalysisUsage(AnalysisUsage &AU);
 
 Optional<APInt> ConstantFoldBinOp(unsigned Opcode, const unsigned Op1,

llvm/lib/CodeGen/GlobalISel/Utils.cpp

 const llvm::ConstantFP* llvm::getConstantFPVRegVal(Register VReg,
                                        const MachineRegisterInfo &MRI) {
   MachineInstr *MI = MRI.getVRegDef(VReg);
   if (TargetOpcode::G_FCONSTANT != MI->getOpcode())
     return nullptr;
   return MI->getOperand(1).getFPImm();
 }
 
-llvm::MachineInstr *llvm::getDefIgnoringCopies(Register Reg,
-                                               const MachineRegisterInfo &MRI) {
+struct DefinitionAndSourceRegister {
+  llvm::MachineInstr *MI;
+  Register Reg;
+};
+
+static llvm::Optional<DefinitionAndSourceRegister> getDefSrcRegIgnoringCopies(
+    Register Reg,
+  Register DefSrcReg = Reg;
   auto *DefMI = MRI.getVRegDef(Reg);
   auto DstTy = MRI.getType(DefMI->getOperand(0).getReg());
   if (!DstTy.isValid())
-    return nullptr;
+    return None;
   while (DefMI->getOpcode() == TargetOpcode::COPY) {
     Register SrcReg = DefMI->getOperand(1).getReg();
     auto SrcTy = MRI.getType(SrcReg);
     if (!SrcTy.isValid() || SrcTy != DstTy)
       break;
     DefMI = MRI.getVRegDef(SrcReg);
-  }
-  return DefMI;
+    DefSrcReg = SrcReg;
+  }
+  return DefinitionAndSourceRegister{ DefMI, DefSrcReg };
+
+llvm::MachineInstr *llvm::getDefIgnoringCopies(Register Reg,
+                                               const MachineRegisterInfo &MRI) {
+  Optional<DefinitionAndSourceRegister> DefSrcReg =
+      getDefSrcRegIgnoringCopies(Reg, MRI);
+  return DefSrcReg ? DefSrcReg->MI : nullptr;
+}
+
+Register llvm::getSrcRegIgnoringCopies(Register Reg,
+                                       const MachineRegisterInfo &MRI) {
+  Optional<DefinitionAndSourceRegister> DefSrcReg =
+      getDefSrcRegIgnoringCopies(Reg, MRI);
+  return DefSrcReg ? DefSrcReg->Reg : Register();
 }
 
 llvm::MachineInstr *llvm::getOpcodeDef(unsigned Opcode, Register Reg,
                                        const MachineRegisterInfo &MRI) {
   MachineInstr *DefMI = getDefIgnoringCopies(Reg, MRI);
   return DefMI && DefMI->getOpcode() == Opcode ? DefMI : nullptr;
 }
 

llvm/lib/Target/AArch64/AArch64InstructionSelector.cpp

bool AArch64InstructionSelector::contractCrossBankCopyIntoStore(
   // %y:fpr(s32) = COPY %x:gpr(s32)
   // G_STORE %y:fpr(s32)
   //
   // We can fix this up into something like this:
   //
   // G_STORE %x:gpr(s32)
   //
   // And then continue the selection process normally.
-  MachineInstr *Def = getDefIgnoringCopies(I.getOperand(0).getReg(), MRI);
-  if (!Def)
+  Register DefDstReg = getSrcRegIgnoringCopies(I.getOperand(0).getReg(), MRI);
+  if (!DefDstReg.isValid())
     return false;
-  Register DefDstReg = Def->getOperand(0).getReg();
   LLT DefDstTy = MRI.getType(DefDstReg);
   Register StoreSrcReg = I.getOperand(0).getReg();
   LLT StoreSrcTy = MRI.getType(StoreSrcReg);
 
   // If we get something strange like a physical register, then we shouldn't
   // go any further.
   if (!DefDstTy.isValid())
     return false;

llvm/test/CodeGen/AArch64/GlobalISel/contract-store.mir

 # NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
 # RUN: llc -mtriple=aarch64-unknown-unknown -run-pass=instruction-select -verify-machineinstrs %s -o - | FileCheck %s
 
 --- |
   define void @contract_s64_gpr(i64* %addr) { ret void }
   define void @contract_s32_gpr(i32* %addr) { ret void }
   define void @contract_s64_fpr(i64* %addr) { ret void }
   define void @contract_s32_fpr(i32* %addr) { ret void }
   define void @contract_s16_fpr(i16* %addr) { ret void }
+  define void @contract_g_unmerge_values_first(i16* %addr) { ret void }
+  define void @contract_g_unmerge_values_second(i16* %addr) { ret void }
 ...
 ---
 name:            contract_s64_gpr
 legalized:       true
 regBankSelected: true
 body:             |
   bb.0:
     liveins: $x0, $x1
bb.0:
     ; CHECK-LABEL: name: contract_s16_fpr
     ; CHECK: [[COPY:%[0-9]+]]:gpr64sp = COPY $x0
     ; CHECK: [[COPY1:%[0-9]+]]:fpr16 = COPY $h1
     ; CHECK: STRHui [[COPY1]], [[COPY]], 0 :: (store 2 into %ir.addr)
     %0:gpr(p0) = COPY $x0
     %1:fpr(s16) = COPY $h1
     %2:gpr(s16) = COPY %1
     G_STORE  %2:gpr(s16), %0 :: (store 2 into %ir.addr)
+...
+---
+name:            contract_g_unmerge_values_first
+legalized:       true
+regBankSelected: true
+body:             |
+  bb.0:
+    liveins: $x0, $x1
+    ; CHECK-LABEL: name: contract_g_unmerge_values_first
+    ; CHECK: [[COPY:%[0-9]+]]:gpr64sp = COPY $x0
+    ; CHECK: [[LOAD:%[0-9]+]]:fpr128 = LDRQui [[COPY]], 0
+    ; CHECK: [[COPY1:%[0-9]+]]:fpr64 = COPY [[LOAD]].dsub
+    ; CHECK: STRDui [[COPY1]], [[COPY]], 0 :: (store 8 into %ir.addr)
+    %0:gpr(p0) = COPY $x0
+    %1:fpr(<2 x s64>) = G_LOAD %0:gpr(p0) :: (dereferenceable load 16 from %ir.addr)
+    %2:fpr(s64), %3:fpr(s64) = G_UNMERGE_VALUES %1:fpr(<2 x s64>)
+    %4:gpr(s64) = COPY %2
+    %5:gpr(s64) = COPY %3
+    G_STORE  %4:gpr(s64), %0 :: (store 8 into %ir.addr)
+...
+---
+name:            contract_g_unmerge_values_second
+legalized:       true
+regBankSelected: true
+body:             |
+  bb.0:
+    liveins: $x0, $x1
+    ; CHECK-LABEL: name: contract_g_unmerge_values_second
+    ; CHECK: [[COPY:%[0-9]+]]:gpr64sp = COPY $x0
+    ; CHECK: [[LOAD:%[0-9]+]]:fpr128 = LDRQui [[COPY]], 0
+    ; CHECK: [[COPY1:%[0-9]+]]:fpr64 = CPYi64 [[LOAD]], 1
+    ; CHECK: STRDui [[COPY1]], [[COPY]], 0 :: (store 8 into %ir.addr)
+    %0:gpr(p0) = COPY $x0
+    %1:fpr(<2 x s64>) = G_LOAD %0:gpr(p0) :: (dereferenceable load 16 from %ir.addr)
+    %2:fpr(s64), %3:fpr(s64) = G_UNMERGE_VALUES %1:fpr(<2 x s64>)
+    %4:gpr(s64) = COPY %2
+    %5:gpr(s64) = COPY %3
+    G_STORE  %5:gpr(s64), %0 :: (store 8 into %ir.addr)