[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 ↗

(On Diff #210915)

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 ↗

(On Diff #210915)

This pattern can now be replaced with getOpcodeDef()?

1276 ↗

(On Diff #210915)

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 ↗

(On Diff #210915)

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

1184 ↗

(On Diff #210915)

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 ↗

(On Diff #210915)

👍

[paquette]

llvm/lib/Target/AArch64/AArch64InstructionSelector.cpp

1184 ↗

(On Diff #210915)

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 ↗

(On Diff #210915)

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

1184 ↗

(On Diff #210915)

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/trunk/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;
 
if (AmtMI->getOpcode() != TargetOpcode::G_CONSTANT) {
       MachineIRBuilder MIB(I);
       auto Trunc = MIB.buildInstr(TargetOpcode::COPY, {SrcTy}, {})
                        .addReg(ShiftReg, 0, AArch64::sub_32);
       MRI.setRegBank(Trunc.getReg(0), RBI.getRegBank(AArch64::GPRRegBankID));
       I.getOperand(2).setReg(Trunc.getReg(0));
     }
     return;
   }
+  case TargetOpcode::G_STORE:
+    contractCrossBankCopyIntoStore(I, MRI);
+    return;
   default:
     return;
   }
 }
 
 bool AArch64InstructionSelector::earlySelectSHL(
     MachineInstr &I, MachineRegisterInfo &MRI) const {
   // We try to match the immediate variant of LSL, which is actually an alias
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 a scalar, it doesn't matter what register bank that
+  // scalar is on. All that matters is the size.
+  //
+  // So, if we see something like this (with a 32-bit scalar as an example):
+  //
+  // %x:gpr(s32) = ... something ...
+  // %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)
+    return;
+  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;
+
+  // Are the source and dst types the same size?
+  if (DefDstTy.getSizeInBits() != StoreSrcTy.getSizeInBits())
+    return;
+
+  if (RBI.getRegBank(StoreSrcReg, MRI, TRI) ==
+      RBI.getRegBank(DefDstReg, MRI, TRI))
+    return;
+
+  // We have a cross-bank copy, which is entering a store. Let's fold it.
+  I.getOperand(0).setReg(DefDstReg);
+}
+
 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) {

llvm/trunk/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 }
+...
+---
+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)