[GlobalISel][AArch64] Contract trivial same-size cross-bank copies into G_STOREs

[paquette]

Summary

Sometimes, you can end up with cross-bank copies between same-sized GPRs and FPRs, which feed into G_STOREs. When these copies feed only into stores, they aren't necessary; we can just store using the original register bank.

This provides some minor code size savings for some floating point SPEC benchmarks. (Around 0.2% for 453.povray and 450.soplex)

This issue doesn't seem to show up due to regbankselect or anything similar. So, this patch introduces an early select function, contractCrossBankCopyIntoStore which performs the contraction when possible. The selector then continues normally and selects the correct store opcode, eliminating needless copies along the way.

[aemerson]

llvm/lib/Target/AArch64/AArch64InstructionSelector.cpp

1170

You mention 32 bits here but the example uses 64 bit types. Also, it doesn't matter whether we use 32 or 64 bit stores right?

1184

This pattern can now be replaced with getOpcodeDef()?

1276

You're returning false here because we want selection to continue, even though you changed the instruction. This is what the preISelLower() hook I added is for.

[paquette]

llvm/lib/Target/AArch64/AArch64InstructionSelector.cpp

1170

Right, it doesn't matter. I'll update the comment for consistency though.

1184

It's not quite equivalent, since it looks through copies. I think that it would still work though, if I rejigger the rest of the function a bit.

1276

👍

[paquette]

llvm/lib/Target/AArch64/AArch64InstructionSelector.cpp

1184

Oh, wait, it looks for a specific opcode. I don't think that will work too well with copies.

[aemerson]

Pre-emptive LGTM with the minor changes done.

llvm/lib/Target/AArch64/AArch64InstructionSelector.cpp

1170

Maybe replace the "storing 32 bits" with something like "storing a scalar".

1184

You're right, it specifically skips copies so it's no good here.

[paquette]

Address review comments.

I switched the copy check over to walking through copies to find the def.

This is equivalent, but does prevent us from checking if all users are stores. This doesn't seem to cause any code size issues though.

[paquette]

remove unused IR function from test

[scw]

https://bugs.llvm.org/show_bug.cgi?id=44783 points to contractCrossBankCopyIntoStore() not handling G_UNMERGE_VALUES correctly.

llvm/lib/Target/AArch64/AArch64InstructionSelector.cpp

private:
   void preISelLower(MachineInstr &I) const;
 
   // An early selection function that runs before the selectImpl() call.
   bool earlySelect(MachineInstr &I) const;
 
   bool earlySelectSHL(MachineInstr &I, MachineRegisterInfo &MRI) const;
   bool earlySelectLoad(MachineInstr &I, MachineRegisterInfo &MRI) const;
 
+  /// Eliminate same-sized cross-bank copies into stores before selectImpl().
+  void contractCrossBankCopyIntoStore(MachineInstr &I,
+                                      MachineRegisterInfo &MRI) const;
+
   bool selectVaStartAAPCS(MachineInstr &I, MachineFunction &MF,
                           MachineRegisterInfo &MRI) const;
   bool selectVaStartDarwin(MachineInstr &I, MachineFunction &MF,
                            MachineRegisterInfo &MRI) const;
 
   bool selectCompareBranch(MachineInstr &I, MachineFunction &MF,
                            MachineRegisterInfo &MRI) const;
 
for (auto &RenderFn : *Imm1Fn)
     RenderFn(NewI);
   for (auto &RenderFn : *Imm2Fn)
     RenderFn(NewI);
 
   I.eraseFromParent();
   return constrainSelectedInstRegOperands(*NewI, TII, TRI, RBI);
 }
 
+void AArch64InstructionSelector::contractCrossBankCopyIntoStore(
+    MachineInstr &I, MachineRegisterInfo &MRI) const {
+  assert(I.getOpcode() == TargetOpcode::G_STORE && "Expected G_STORE");
+  // If we're storing 32 bits, it doesn't matter if it's 32 bits on a FPR or a
+  // GPR. So, if we run into something like this:
+  //
+  // %x:gpr(s64) = ... something ...
+  // %y:fpr(s64) = COPY %x:gpr(s64)
+  // G_STORE %y:fpr(s64)
+  //
+  // We can fix this up into something like this:
+  //
+  // G_STORE %x:gpr(s64)
+  //
+  // And then continue the selection process normally.
+  MachineInstr *Copy = MRI.getVRegDef(I.getOperand(0).getReg());
+  if (!Copy || Copy->getOpcode() != TargetOpcode::COPY)
+    return;
+  Register CopyDstReg = Copy->getOperand(0).getReg();
+  LLT CopyDstTy = MRI.getType(CopyDstReg);
+  Register CopySrcReg = Copy->getOperand(1).getReg();
+  LLT CopySrcTy = MRI.getType(CopySrcReg);
+
+  // If we get something strange like a physical register, then we shouldn't
+  // go any further.
+  if (!CopySrcTy.isValid())
+    return;
+
+  // Are the source and dst types the same size?
+  if (CopyDstTy.getSizeInBits() != CopySrcTy.getSizeInBits())
+    return;
+
+  if (RBI.getRegBank(CopySrcReg, MRI, TRI) ==
+      RBI.getRegBank(CopyDstReg, MRI, TRI))
+    return;
+
+  // Is this copy used in anything other than a store? If it is, then we don't
+  // want to fold anything, since we'll need the copy anyway.
+  if (any_of(MRI.use_instructions(CopyDstReg),
+             [](const MachineInstr &MI) { return !MI.mayStore(); }))
+    return;
+
+  // We have a cross-bank copy, which is entering a store. Let's fold it.
+  I.getOperand(0).setReg(CopySrcReg);
+}
+
 bool AArch64InstructionSelector::earlySelectLoad(
     MachineInstr &I, MachineRegisterInfo &MRI) const {
   // Try to fold in shifts, etc into the addressing mode of a load.
   assert(I.getOpcode() == TargetOpcode::G_LOAD && "unexpected op");
 
   // Don't handle atomic loads/stores yet.
   auto &MemOp = **I.memoperands_begin();
   if (MemOp.getOrdering() != AtomicOrdering::NotAtomic) {
bool AArch64InstructionSelector::earlySelect(MachineInstr &I) const {
 
   MachineBasicBlock &MBB = *I.getParent();
   MachineFunction &MF = *MBB.getParent();
   MachineRegisterInfo &MRI = MF.getRegInfo();
 
   switch (I.getOpcode()) {
   case TargetOpcode::G_SHL:
     return earlySelectSHL(I, MRI);
+  case TargetOpcode::G_STORE:
+    contractCrossBankCopyIntoStore(I, MRI);
+    return false;
   case TargetOpcode::G_LOAD:
     return earlySelectLoad(I, MRI);
   default:
     return false;
   }
 }
 
 bool AArch64InstructionSelector::select(MachineInstr &I,

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 @dont_contract(i64* %addr) { ret void }
+...
+---
+name:            contract_s64_gpr
+legalized:       true
+regBankSelected: true
+body:             |
+  bb.0:
+    liveins: $x0, $x1
+    ; CHECK-LABEL: name: contract_s64_gpr
+    ; CHECK: [[COPY:%[0-9]+]]:gpr64sp = COPY $x0
+    ; CHECK: [[COPY1:%[0-9]+]]:gpr64 = COPY $x1
+    ; CHECK: STRXui [[COPY1]], [[COPY]], 0 :: (store 8 into %ir.addr)
+    %0:gpr(p0) = COPY $x0
+    %1:gpr(s64) = COPY $x1
+    %2:fpr(s64) = COPY %1
+    G_STORE  %2:fpr(s64), %0 :: (store 8 into %ir.addr)
+...
+---
+name:            contract_s32_gpr
+legalized:       true
+regBankSelected: true
+body:             |
+  bb.0:
+    liveins: $x0, $w1
+    ; CHECK-LABEL: name: contract_s32_gpr
+    ; CHECK: [[COPY:%[0-9]+]]:gpr64sp = COPY $x0
+    ; CHECK: [[COPY1:%[0-9]+]]:gpr32 = COPY $w1
+    ; CHECK: STRWui [[COPY1]], [[COPY]], 0 :: (store 4 into %ir.addr)
+    %0:gpr(p0) = COPY $x0
+    %1:gpr(s32) = COPY $w1
+    %2:fpr(s32) = COPY %1
+    G_STORE  %2:fpr(s32), %0 :: (store 4 into %ir.addr)
+...
+---
+name:            contract_s64_fpr
+legalized:       true
+regBankSelected: true
+body:             |
+  bb.0:
+    liveins: $x0, $d1
+    ; CHECK-LABEL: name: contract_s64_fpr
+    ; CHECK: [[COPY:%[0-9]+]]:gpr64sp = COPY $x0
+    ; CHECK: [[COPY1:%[0-9]+]]:fpr64 = COPY $d1
+    ; CHECK: STRDui [[COPY1]], [[COPY]], 0 :: (store 8 into %ir.addr)
+    %0:gpr(p0) = COPY $x0
+    %1:fpr(s64) = COPY $d1
+    %2:gpr(s64) = COPY %1
+    G_STORE %2:gpr(s64), %0 :: (store 8 into %ir.addr)
+...
+---
+name:            contract_s32_fpr
+legalized:       true
+regBankSelected: true
+body:             |
+  bb.0:
+    liveins: $x0, $s1
+    ; CHECK-LABEL: name: contract_s32_fpr
+    ; CHECK: [[COPY:%[0-9]+]]:gpr64sp = COPY $x0
+    ; CHECK: [[COPY1:%[0-9]+]]:fpr32 = COPY $s1
+    ; CHECK: STRSui [[COPY1]], [[COPY]], 0 :: (store 4 into %ir.addr)
+    %0:gpr(p0) = COPY $x0
+    %1:fpr(s32) = COPY $s1
+    %2:gpr(s32) = COPY %1
+    G_STORE %2:gpr(s32), %0 :: (store 4 into %ir.addr)
+...
+---
+name:            contract_s16_fpr
+legalized:       true
+regBankSelected: true
+body:             |
+  bb.0:
+    liveins: $x0, $h1
+    ; 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:            dont_contract
+legalized:       true
+regBankSelected: true
+body:             |
+  bb.0:
+    ; We need to copy %2 into $d2 in this case, so there's no reason to fold
+    ; away the copy of %1 into %2.
+    liveins: $x0, $x1, $d2
+    ; CHECK-LABEL: name: dont_contract
+    ; CHECK: [[COPY:%[0-9]+]]:gpr64sp = COPY $x0
+    ; CHECK: [[COPY1:%[0-9]+]]:gpr64all = COPY $x1
+    ; CHECK: [[COPY2:%[0-9]+]]:fpr64 = COPY [[COPY1]]
+    ; CHECK: STRDui [[COPY2]], [[COPY]], 0 :: (store 8 into %ir.addr)
+    ; CHECK: $d2 = COPY [[COPY2]]
+    ; CHECK: RET_ReallyLR implicit $d2
+    %0:gpr(p0) = COPY $x0
+    %1:gpr(s64) = COPY $x1
+    %2:fpr(s64) = COPY %1
+    G_STORE  %2:fpr(s64), %0 :: (store 8 into %ir.addr)
+    $d2 = COPY %2:fpr(s64)
+    RET_ReallyLR implicit $d2
+...