[AArch64][GlobalISel] Avoid copies to target register bank for subregister copies

[tambre]

Summary

Previously for any copy from a register bigger than the destination:

1. Copied to a same-sized register in the destination register bank.
2. Subregister copy of that to the destination.

This fails for copies from 128-bit FPRs to GPRs because the GPR register bank can't accomodate 128-bit values.

Instead of special-casing such copies to perform the truncation beforehand in the source register bank, generalize this:
a) Perform a subregister copy straight from source register whenever possible. This results in shorter MIR and fixes the above problem.
b) Perform a full copy to target bank and then do a subregister copy only if source bank can't support target's size. E.g. GPR to 8-bit FPR copy.

[tambre]

This patch fixes bug 45029.

[tambre]

Re-generate patch on Linux. Seems like line endings might've been messed up.

[tambre]

Jessica: Adding you as a reviewer, as you seem to have worked with this code.

[paquette]

I think this looks good; just a couple nits on the testcase.

llvm/lib/Target/AArch64/AArch64InstructionSelector.cpp

776

On calls to getMinClassForRegBank, can we add a comment like below?

getMinClassForRegBank(SrcRegBank, DstSize, /* GetAllRegSet = */ true)

llvm/test/CodeGen/AArch64/GlobalISel/subreg-copy.mir

5

This isn't necessary, because the runline already specifies the triple.

9–19

The actual IR here isn't necessary; it can probably just be

define void @test_128_fpr_truncation() { ret void }

[paquette]

llvm/lib/Target/AArch64/AArch64InstructionSelector.cpp

459

This should be an llvm_unreachable, I think. We never want to return anything here.

[paquette]

llvm/lib/Target/AArch64/AArch64InstructionSelector.cpp

662–663

While we're here, I think it would be good to add a couple asserts:

assert(SrcReg.isValid() && "Expected a valid source register?");
assert(To && "Destination register class cannot be null");
assert(SubReg && "Expected a valid subregister");

[paquette]

llvm/test/CodeGen/AArch64/GlobalISel/subreg-copy.mir

38

Do you have to use G_GLOBAL_VALUE here? Would this work instead?

bb.0:
liveins: $x0
  %1:gpr(p0) = COPY $x0

Then you could replace the G_LOAD with

%0:fpr(s128) = G_LOAD %1:gpr(p0) :: (load 16)

If you can do that, then you can delete the entire IR section in the testcase.

(You might have to add tracksRegLiveness: true to make this work)

44–50

Is anything after the COPY necessary to test the behaviour here? Can we just return the result of the COPY? E.g.

$x1 = COPY %8:gpr(s64)
RET_ReallyLR implicit $x1

[tambre]

Address Jessica's review comments.

[tambre]

Thanks for the review Jessica! I addressed your comments, except for being unable to reduce the testcase much further.
Note that this change is a fix for bug 45029.

llvm/test/CodeGen/AArch64/GlobalISel/subreg-copy.mir

9–19

It seemed to me that the IR was usually included for reference. Seems not! :)

44–50

Changing anything after the COPY results in the testcase no longer crashing prior to this change. I tried reducing in various ways, but unfortunately no dice.

[paquette]

LGTM!

[tambre]

Jessica, please land the patch for me. As this is my first contribution, I probably lack the permissions to do that myself.

llvm/lib/Target/AArch64/AArch64InstructionSelector.cpp

default:
     LLVM_DEBUG(
         dbgs() << "Couldn't find appropriate subregister for register class.");
     return false;
   }
 
   return true;
 }
 
+/// Returns the minimum size the given register bank can hold.
+static unsigned getMinSizeForRegBank(const RegisterBank &RB) {
+  switch (RB.getID()) {
+  case AArch64::GPRRegBankID:
+    return 32;
+  case AArch64::FPRRegBankID:
+    return 8;
+  default:
+    llvm_unreachable("Tried to get minimum size for unknown register bank.");
+  }
+}
+
 /// Check whether \p I is a currently unsupported binary operation:
 /// - it has an unsized type
 /// - an operand is not a vreg
 /// - all operands are not in the same bank
 /// These are checks that should someday live in the verifier, but right now,
 /// these are mostly limitations of the aarch64 selector.
 static bool unsupportedBinOp(const MachineInstr &I,
                              const AArch64RegisterBankInfo &RBI,
static bool isValidCopy(const MachineInstr &I, const RegisterBank &DstBank,
   // Check the size of the destination.
   assert((DstSize <= 64 || DstBank.getID() == AArch64::FPRRegBankID) &&
          "GPRs cannot get more than 64-bit width values");
 
   return true;
 }
 #endif
 
-/// Helper function for selectCopy. Inserts a subregister copy from
-/// \p *From to \p *To, linking it up to \p I.
-///
-/// e.g, given I = "Dst = COPY SrcReg", we'll transform that into
+/// Helper function for selectCopy. Inserts a subregister copy from \p SrcReg
+/// to \p *To.
 ///
-/// CopyReg (From class) = COPY SrcReg
-/// SubRegCopy (To class) = COPY CopyReg:SubReg
-/// Dst = COPY SubRegCopy
-static bool selectSubregisterCopy(MachineInstr &I, MachineRegisterInfo &MRI,
-                                  const RegisterBankInfo &RBI, Register SrcReg,
-                                  const TargetRegisterClass *From,
-                                  const TargetRegisterClass *To,
-                                  unsigned SubReg) {
+/// E.g "To = COPY SrcReg:SubReg"
+static bool copySubReg(MachineInstr &I, MachineRegisterInfo &MRI,
+                       const TargetRegisterClass *To, unsigned SubReg) {
+  assert(SrcReg.isValid() && "Expected a valid source register?");
+  assert(To && "Destination register class cannot be null");
+  assert(SubReg && "Expected a valid subregister");
+
   MachineIRBuilder MIB(I);
-  auto Copy = MIB.buildCopy({From}, {SrcReg});
-  auto SubRegCopy = MIB.buildInstr(TargetOpcode::COPY, {To}, {})
-                        .addReg(Copy.getReg(0), 0, SubReg);
+  auto SubRegCopy =
+      MIB.buildInstr(TargetOpcode::COPY, {To}, {}).addReg(SrcReg, 0, SubReg);
   MachineOperand &RegOp = I.getOperand(1);
   RegOp.setReg(SubRegCopy.getReg(0));
 
   // It's possible that the destination register won't be constrained. Make
   // sure that happens.
   if (!Register::isPhysicalRegister(I.getOperand(0).getReg()))
     RBI.constrainGenericRegister(I.getOperand(0).getReg(), *To, MRI);
 
if (I.isCopy()) {
     if (!SrcRC) {
       LLVM_DEBUG(dbgs() << "Couldn't determine source register class\n");
       return false;
     }
 
     unsigned SrcSize = TRI.getRegSizeInBits(*SrcRC);
     unsigned DstSize = TRI.getRegSizeInBits(*DstRC);
 
-    // If we're doing a cross-bank copy on different-sized registers, we need
-    // to do a bit more work.
+    // If the source register is bigger than the destination we need to perform
+    // a subregister copy.
     if (SrcSize > DstSize) {
-      // We're doing a cross-bank copy into a smaller register. We need a
-      // subregister copy. First, get a register class that's on the same bank
-      // as the destination, but the same size as the source.
-      const TargetRegisterClass *SubregRC =
-          getMinClassForRegBank(DstRegBank, SrcSize, true);
-      assert(SubregRC && "Didn't get a register class for subreg?");
-
-      // Get the appropriate subregister for the destination.
       unsigned SubReg = 0;
-      if (!getSubRegForClass(DstRC, TRI, SubReg)) {
-        LLVM_DEBUG(dbgs() << "Couldn't determine subregister for copy.\n");
-        return false;
+
+      // If the source bank doesn't support a subregister copy small enough,
+      // then we first need to copy to the destination bank.
+      if (getMinSizeForRegBank(SrcRegBank) > DstSize) {
+        const TargetRegisterClass *SubregRC = getMinClassForRegBank(
+            DstRegBank, SrcSize, /* GetAllRegSet = */ true);
+        getSubRegForClass(DstRC, TRI, SubReg);
+
+        MachineIRBuilder MIB(I);
+        auto Copy = MIB.buildCopy({SubregRC}, {SrcReg});
+        copySubReg(I, MRI, RBI, Copy.getReg(0), DstRC, SubReg);
+      } else {
+        const TargetRegisterClass *SubregRC = getMinClassForRegBank(
+            SrcRegBank, DstSize, /* GetAllRegSet = */ true);
+        getSubRegForClass(SubregRC, TRI, SubReg);
+        copySubReg(I, MRI, RBI, SrcReg, DstRC, SubReg);
       }
 
-      // Now, insert a subregister copy using the new register class.
-      selectSubregisterCopy(I, MRI, RBI, SrcReg, SubregRC, DstRC, SubReg);
       return CheckCopy();
     }
 
     // Is this a cross-bank copy?
     if (DstRegBank.getID() != SrcRegBank.getID()) {
       if (DstRegBank.getID() == AArch64::GPRRegBankID && DstSize == 32 &&
           SrcSize == 16) {
         // Special case for FPR16 to GPR32.

llvm/test/CodeGen/AArch64/GlobalISel/load-wro-addressing-modes.mir

bb.0:
     liveins: $x0
 
     ; The mask used for the AND here is legal for producing a roW load.
 
     ; CHECK-LABEL: name: and_uxtw
     ; CHECK: liveins: $x0
     ; CHECK: %base:gpr64sp = COPY $x0
     ; CHECK: %imp:gpr64 = IMPLICIT_DEF
-    ; CHECK: [[COPY:%[0-9]+]]:gpr64all = COPY %imp
-    ; CHECK: [[COPY1:%[0-9]+]]:gpr32all = COPY [[COPY]].sub_32
-    ; CHECK: [[COPY2:%[0-9]+]]:gpr32 = COPY [[COPY1]]
-    ; CHECK: %load:gpr64 = LDRXroW %base, [[COPY2]], 0, 1 :: (load 8)
+    ; CHECK: [[COPY:%[0-9]+]]:gpr32all = COPY %imp.sub_32
+    ; CHECK: [[COPY1:%[0-9]+]]:gpr32 = COPY [[COPY]]
+    ; CHECK: %load:gpr64 = LDRXroW %base, [[COPY1]], 0, 1 :: (load 8)
     ; CHECK: $x1 = COPY %load
     ; CHECK: RET_ReallyLR implicit $x1
     %base:gpr(p0) = COPY $x0
     %imp:gpr(s64) = G_IMPLICIT_DEF
     %mask:gpr(s64) = G_CONSTANT i64 4294967295
     %and:gpr(s64) = G_AND %imp, %mask
     %c:gpr(s64) = G_CONSTANT i64 8
     %mul:gpr(s64) = G_MUL %c, %and
     %ptr:gpr(p0) = G_PTR_ADD %base, %mul(s64)
     %load:gpr(s64) = G_LOAD %ptr(p0) :: (load 8)
     $x1 = COPY %load(s64)
     RET_ReallyLR implicit $x1
 ...

llvm/test/CodeGen/AArch64/GlobalISel/opt-fold-and-tbz-tbnz.mir

 legalized:       true
 regBankSelected: true
 body:             |
   ; CHECK-LABEL: name: dont_fold_and
   ; CHECK: bb.0:
   ; CHECK:   successors: %bb.0(0x40000000), %bb.1(0x40000000)
   ; CHECK:   %copy:gpr64 = COPY $x0
   ; CHECK:   %fold_me:gpr64sp = ANDXri %copy, 4098
-  ; CHECK:   [[COPY:%[0-9]+]]:gpr64all = COPY %fold_me
-  ; CHECK:   [[COPY1:%[0-9]+]]:gpr32all = COPY [[COPY]].sub_32
-  ; CHECK:   [[COPY2:%[0-9]+]]:gpr32 = COPY [[COPY1]]
-  ; CHECK:   TBNZW [[COPY2]], 3, %bb.1
+  ; CHECK:   [[COPY:%[0-9]+]]:gpr32all = COPY %fold_me.sub_32
+  ; CHECK:   [[COPY1:%[0-9]+]]:gpr32 = COPY [[COPY]]
+  ; CHECK:   TBNZW [[COPY1]], 3, %bb.1
   ; CHECK:   B %bb.0
   ; CHECK: bb.1:
   ; CHECK:   RET_ReallyLR
   bb.0:
     successors: %bb.0, %bb.1
     liveins: $x0
     %copy:gpr(s64) = COPY $x0
     %bit:gpr(s64) = G_CONSTANT i64 8

llvm/test/CodeGen/AArch64/GlobalISel/opt-fold-ext-tbz-tbnz.mir

 body:             |
   ; CHECK-LABEL: name: dont_fold_more_than_one_use
   ; CHECK: bb.0:
   ; CHECK:   successors: %bb.0(0x40000000), %bb.1(0x40000000)
   ; CHECK:   liveins: $x0
   ; CHECK:   %copy:gpr32 = COPY $w0
   ; CHECK:   [[SUBREG_TO_REG:%[0-9]+]]:gpr64 = SUBREG_TO_REG 0, %copy, %subreg.sub_32
   ; CHECK:   %zext:gpr64 = UBFMXri [[SUBREG_TO_REG]], 0, 31
-  ; CHECK:   [[COPY:%[0-9]+]]:gpr64all = COPY %zext
-  ; CHECK:   [[COPY1:%[0-9]+]]:gpr32all = COPY [[COPY]].sub_32
-  ; CHECK:   [[COPY2:%[0-9]+]]:gpr32 = COPY [[COPY1]]
-  ; CHECK:   TBNZW [[COPY2]], 3, %bb.1
+  ; CHECK:   [[COPY:%[0-9]+]]:gpr32all = COPY %zext.sub_32
+  ; CHECK:   [[COPY1:%[0-9]+]]:gpr32 = COPY [[COPY]]
+  ; CHECK:   TBNZW [[COPY1]], 3, %bb.1
   ; CHECK:   B %bb.0
   ; CHECK: bb.1:
   ; CHECK:   $x0 = COPY %zext
   ; CHECK:   RET_ReallyLR implicit $x0
   bb.0:
     successors: %bb.0, %bb.1
     liveins: $x0
     %copy:gpr(s32) = COPY $w0

llvm/test/CodeGen/AArch64/GlobalISel/opt-fold-shift-tbz-tbnz.mir

 legalized:       true
 regBankSelected: true
 body:             |
   ; CHECK-LABEL: name: dont_fold_shl_1
   ; CHECK: bb.0:
   ; CHECK:   successors: %bb.0(0x40000000), %bb.1(0x40000000)
   ; CHECK:   %copy:gpr64 = COPY $x0
   ; CHECK:   %fold_me:gpr64 = UBFMXri %copy, 59, 58
-  ; CHECK:   [[COPY:%[0-9]+]]:gpr64all = COPY %fold_me
-  ; CHECK:   [[COPY1:%[0-9]+]]:gpr32all = COPY [[COPY]].sub_32
-  ; CHECK:   [[COPY2:%[0-9]+]]:gpr32 = COPY [[COPY1]]
-  ; CHECK:   TBNZW [[COPY2]], 3, %bb.1
+  ; CHECK:   [[COPY:%[0-9]+]]:gpr32all = COPY %fold_me.sub_32
+  ; CHECK:   [[COPY1:%[0-9]+]]:gpr32 = COPY [[COPY]]
+  ; CHECK:   TBNZW [[COPY1]], 3, %bb.1
   ; CHECK:   B %bb.0
   ; CHECK: bb.1:
   ; CHECK:   RET_ReallyLR
   bb.0:
     successors: %bb.0, %bb.1
     liveins: $x0
     %copy:gpr(s64) = COPY $x0
     %bit:gpr(s64) = G_CONSTANT i64 8

 regBankSelected: true
 body:             |
   ; CHECK-LABEL: name: dont_fold_shl_2
   ; CHECK: bb.0:
   ; CHECK:   successors: %bb.0(0x40000000), %bb.1(0x40000000)
   ; CHECK:   %copy:gpr64 = COPY $x0
   ; CHECK:   %fold_cst:gpr64 = MOVi64imm -5
   ; CHECK:   %fold_me:gpr64 = LSLVXr %copy, %fold_cst
-  ; CHECK:   [[COPY:%[0-9]+]]:gpr64all = COPY %fold_me
-  ; CHECK:   [[COPY1:%[0-9]+]]:gpr32all = COPY [[COPY]].sub_32
-  ; CHECK:   [[COPY2:%[0-9]+]]:gpr32 = COPY [[COPY1]]
-  ; CHECK:   TBNZW [[COPY2]], 3, %bb.1
+  ; CHECK:   [[COPY:%[0-9]+]]:gpr32all = COPY %fold_me.sub_32
+  ; CHECK:   [[COPY1:%[0-9]+]]:gpr32 = COPY [[COPY]]
+  ; CHECK:   TBNZW [[COPY1]], 3, %bb.1
   ; CHECK:   B %bb.0
   ; CHECK: bb.1:
   ; CHECK:   RET_ReallyLR
   bb.0:
     successors: %bb.0, %bb.1
     liveins: $x0
     %copy:gpr(s64) = COPY $x0
     %bit:gpr(s64) = G_CONSTANT i64 8

 legalized:       true
 regBankSelected: true
 body:             |
   ; CHECK-LABEL: name: dont_fold_shl_3
   ; CHECK: bb.0:
   ; CHECK:   successors: %bb.0(0x40000000), %bb.1(0x40000000)
   ; CHECK:   %copy:gpr64 = COPY $x0
   ; CHECK:   %shl:gpr64 = UBFMXri %copy, 62, 61
-  ; CHECK:   [[COPY:%[0-9]+]]:gpr64all = COPY %shl
-  ; CHECK:   [[COPY1:%[0-9]+]]:gpr32all = COPY [[COPY]].sub_32
-  ; CHECK:   [[COPY2:%[0-9]+]]:gpr32 = COPY [[COPY1]]
-  ; CHECK:   TBNZW [[COPY2]], 3, %bb.1
+  ; CHECK:   [[COPY:%[0-9]+]]:gpr32all = COPY %shl.sub_32
+  ; CHECK:   [[COPY1:%[0-9]+]]:gpr32 = COPY [[COPY]]
+  ; CHECK:   TBNZW [[COPY1]], 3, %bb.1
   ; CHECK:   B %bb.0
   ; CHECK: bb.1:
   ; CHECK:   %second_use:gpr64sp = ORRXri %shl, 8000
   ; CHECK:   $x0 = COPY %second_use
   ; CHECK:   RET_ReallyLR implicit $x0
   bb.0:
     successors: %bb.0, %bb.1
     liveins: $x0

llvm/test/CodeGen/AArch64/GlobalISel/subreg-copy.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
+---
+name:              test_128_fpr_truncation
+alignment:         4
+legalized:         true
+regBankSelected:   true
+tracksRegLiveness: true
+body: |
+  ; CHECK-LABEL: name: test_128_fpr_truncation
+  ; CHECK: bb.0:
+  ; CHECK:   successors: %bb.1(0x80000000)
+  ; CHECK:   liveins: $x0
+  ; CHECK:   [[COPY:%[0-9]+]]:gpr64sp = COPY $x0
+  ; CHECK:   [[LDRQui:%[0-9]+]]:fpr128 = LDRQui [[COPY]], 0 :: (load 16)
+  ; CHECK:   [[COPY1:%[0-9]+]]:gpr32all = COPY [[LDRQui]].ssub
+  ; CHECK:   [[COPY2:%[0-9]+]]:gpr32 = COPY [[COPY1]]
+  ; CHECK:   TBNZW [[COPY2]], 0, %bb.1
+  ; CHECK: bb.1:
+  ; CHECK:   RET_ReallyLR
+  bb.0:
+    liveins: $x0
+    %1:gpr(p0) = COPY $x0
+    %3:gpr(s64) = G_CONSTANT i64 1
+    %5:gpr(s64) = G_CONSTANT i64 0
+    %0:fpr(s128) = G_LOAD %1:gpr(p0) :: (load 16)
+    %2:fpr(s64) = G_TRUNC %0:fpr(s128)
+    %8:gpr(s64) = COPY %2:fpr(s64)
+    %4:gpr(s64) = G_AND %8:gpr, %3:gpr
+    %7:gpr(s32) = G_ICMP intpred(ne), %4:gpr(s64), %5:gpr
+    %6:gpr(s1) = G_TRUNC %7:gpr(s32)
+    G_BRCOND %6:gpr(s1), %bb.1
+
+  bb.1:
+    RET_ReallyLR
+...