[AArch64][GlobalISel] Select unscaled loads/stores in manual selector

[paquette]

Summary

For smaller types (e.g. 16-bit), we don't import patterns for unscaled loads and stores.

Teach the selector to do this.

We didn't really have any testcases for the unscaled load/store stuff anyway, so add testcases showing what we can select.

(Nothing for STURBi or LDURBi, since we need more legalizer support for s8s to get those.)

This is a 0.2% code size improvement on 7zip at -O3. (Geomean -0.0%)

[paquette]

Remove assert because I forgot that there are a million different generic load instruction variants

[aemerson]

llvm/lib/Target/AArch64/GISel/AArch64InstructionSelector.cpp

2358

selectUnscaledLoadStoreOp can take the bool IsStore for the first arg instead of a full opcode.

2367

When would this happen?

llvm/lib/Target/AArch64/GISel/AArch64InstructionSelector.cpp

case AArch64::FPRRegBankID:
     case 64:
       return isStore ? AArch64::STRDui : AArch64::LDRDui;
     }
     break;
   }
   return GenericOpc;
 }
 
+/// Select the AArch64 opcode for a G_LOAD or G_STORE in the unscaled addressing
+/// mode.
+/// \p GenericOpc is one of G_LOAD and G_STORE.
+/// \p RegBankID is the register bank
+/// \p OpSize is the memory access size.
+/// \returns None if the combination is unsupported.
+static Optional<unsigned> selectUnscaledLoadStoreOp(unsigned GenericOpc,
+                                          unsigned RegBankID, unsigned OpSize) {
+  const bool isStore = GenericOpc == TargetOpcode::G_STORE;
+  switch (RegBankID) {
+  default:
+    return None;
+  case AArch64::GPRRegBankID:
+    switch (OpSize) {
+    default:
+      return None;
+    case 8:
+      return isStore ? AArch64::STURBBi : AArch64::LDURBBi;
+    case 16:
+      return isStore ? AArch64::STURHHi : AArch64::LDURHHi;
+    case 32:
+      return isStore ? AArch64::STURWi : AArch64::LDURWi;
+    case 64:
+      return isStore ? AArch64::STURXi : AArch64::LDURXi;
+    }
+  case AArch64::FPRRegBankID:
+    switch (OpSize) {
+    default:
+      return None;
+    case 8:
+      return isStore ? AArch64::STURBi : AArch64::LDURBi;
+    case 16:
+      return isStore ? AArch64::STURHi : AArch64::LDURHi;
+    case 32:
+      return isStore ? AArch64::STURSi : AArch64::LDURSi;
+    case 64:
+      return isStore ? AArch64::STURDi : AArch64::LDURDi;
+    }
+  }
+}
+
 #ifndef NDEBUG
 /// Helper function that verifies that we have a valid copy at the end of
 /// selectCopy. Verifies that the source and dest have the expected sizes and
 /// then returns true.
 static bool isValidCopy(const MachineInstr &I, const RegisterBank &DstBank,
                         const MachineRegisterInfo &MRI,
                         const TargetRegisterInfo &TRI,
                         const RegisterBankInfo &RBI) {

 #endif
 
     const Register ValReg = I.getOperand(0).getReg();
     const RegisterBank &RB = *RBI.getRegBank(ValReg, MRI, TRI);
 
     // Helper lambda for partially selecting I. Either returns the original
     // instruction with an updated opcode, or a new instruction.
     auto SelectLoadStoreAddressingMode = [&]() -> MachineInstr * {
-      bool IsStore = I.getOpcode() == TargetOpcode::G_STORE;
-      const unsigned NewOpc =
-          selectLoadStoreUIOp(I.getOpcode(), RB.getID(), MemSizeInBits);
-      if (NewOpc == I.getOpcode())
-        return nullptr;
+      unsigned GenericOpc = I.getOpcode();
+      bool IsStore = GenericOpc == TargetOpcode::G_STORE;
+      MachineOperand &PtrOp = I.getOperand(1);
+      unsigned RegBankID = RB.getID();
+
+      // By default, we'll use the UI opcode.
+      unsigned NewOpc =
+          selectLoadStoreUIOp(GenericOpc, RegBankID, MemSizeInBits);
+
       // Check if we can fold anything into the addressing mode.
-      auto AddrModeFns =
-          selectAddrModeIndexed(I.getOperand(1), MemSizeInBytes);
+      ComplexRendererFns AddrModeFns;
+      if (auto UnscaledOpc =
+              selectUnscaledLoadStoreOp(GenericOpc, RegBankID, MemSizeInBits)) {
+        AddrModeFns = selectAddrModeUnscaled(PtrOp, MemSizeInBytes);
+        if (AddrModeFns)
+          NewOpc = *UnscaledOpc;
+      }
+
+      if (!AddrModeFns)
+        AddrModeFns = selectAddrModeIndexed(PtrOp, MemSizeInBytes);
+
+      if (NewOpc == GenericOpc)
+        return nullptr;
+
       if (!AddrModeFns) {
-        // Can't fold anything. Use the original instruction.
+        // Can't fold anything. Use the original instruction, using the UI
+        // opcode.
         I.setDesc(TII.get(NewOpc));
         I.addOperand(MachineOperand::CreateImm(0));
         return &I;
       }
 
       // Folded something. Create a new instruction and return it.
       auto NewInst = MIB.buildInstr(NewOpc, {}, {}, I.getFlags());
       IsStore ? NewInst.addUse(ValReg) : NewInst.addDef(ValReg);

llvm/test/CodeGen/AArch64/GlobalISel/select-load.mir

 # NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
 # RUN: llc -mtriple=aarch64-- -run-pass=instruction-select -verify-machineinstrs %s -o - | FileCheck %s
 
 --- |
   target datalayout = "e-m:o-i64:64-i128:128-n32:64-S128"
+  @x = external hidden local_unnamed_addr global i32*, align 8
 
   define void @load_s64_gpr(i64* %addr) { ret void }
   define void @load_s32_gpr(i32* %addr) { ret void }
   define void @load_s16_gpr_anyext(i16* %addr) { ret void }
   define void @load_s16_gpr(i16* %addr) { ret void }
   define void @load_s8_gpr_anyext(i8* %addr) { ret void }
   define void @load_s8_gpr(i8* %addr) { ret void }
 
--- |
   define void @load_v2s32(i64 *%addr) { ret void }
   define void @load_v2s64(i64 *%addr) { ret void }
 
   define void @load_4xi16(<4 x i16>* %ptr) { ret void }
   define void @load_4xi32(<4 x i32>* %ptr) { ret void }
   define void @load_8xi16(<8 x i16>* %ptr) { ret void }
   define void @load_16xi8(<16 x i8>* %ptr) { ret void }
 
+  ; Check selection for unscaled addressing modes.
+  define void @ldurwi() { ret void }
+  define void @ldursi() { ret void }
+  define void @ldurxi() { ret void }
+  define void @ldurdi() { ret void }
+  define void @ldurhhi() { ret void }
+  define void @ldurhi() { ret void }
+  define void @ldurbbi() { ret void }
+
 ...
 
 ---
 name:            load_s64_gpr
 legalized:       true
 regBankSelected: true
 
 registers:
bb.1 (%ir-block.0):
     ; CHECK: $q0 = COPY [[LDRQui]]
     ; CHECK: RET_ReallyLR implicit $q0
     %0:gpr(p0) = COPY $x0
     %1:fpr(<16 x s8>) = G_LOAD %0(p0) :: (load 16 from %ir.ptr)
     $q0 = COPY %1(<16 x s8>)
     RET_ReallyLR implicit $q0
 
 ...
+---
+name:            ldurwi
+legalized:       true
+regBankSelected: true
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $x0, $w1
+    ; CHECK-LABEL: name: ldurwi
+    ; CHECK: liveins: $x0, $w1
+    ; CHECK: %ptr:gpr64sp = COPY $x0
+    ; CHECK: %val:gpr32 = LDURWi %ptr, -256 :: (load 4 from @x)
+    ; CHECK: $w1 = COPY %val
+    ; CHECK: RET_ReallyLR implicit $w1
+    %ptr:gpr(p0) = COPY $x0
+    %cst:gpr(s64) = G_CONSTANT i64 -256
+    %ptr_add:gpr(p0) = G_PTR_ADD %ptr(p0), %cst
+    %val:gpr(s32) = G_LOAD %ptr_add(p0) :: (load 4 from @x)
+    $w1 = COPY %val(s32)
+    RET_ReallyLR implicit $w1
+
+...
+---
+name:            ldursi
+legalized:       true
+regBankSelected: true
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $x0, $s1
+    ; CHECK-LABEL: name: ldursi
+    ; CHECK: liveins: $x0, $s1
+    ; CHECK: %ptr:gpr64sp = COPY $x0
+    ; CHECK: %val:fpr32 = LDURSi %ptr, -256 :: (load 4 from @x)
+    ; CHECK: $s1 = COPY %val
+    ; CHECK: RET_ReallyLR implicit $s1
+    %ptr:gpr(p0) = COPY $x0
+    %cst:gpr(s64) = G_CONSTANT i64 -256
+    %ptr_add:gpr(p0) = G_PTR_ADD %ptr(p0), %cst
+    %val:fpr(s32) = G_LOAD %ptr_add(p0) :: (load 4 from @x)
+    $s1 = COPY %val(s32)
+    RET_ReallyLR implicit $s1
+
+...
+---
+name:            ldurxi
+legalized:       true
+regBankSelected: true
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $x0, $x1
+    ; CHECK-LABEL: name: ldurxi
+    ; CHECK: liveins: $x0, $x1
+    ; CHECK: %ptr:gpr64sp = COPY $x0
+    ; CHECK: %val:gpr64 = LDURXi %ptr, -256 :: (load 8 from @x)
+    ; CHECK: $x1 = COPY %val
+    ; CHECK: RET_ReallyLR implicit $x1
+    %ptr:gpr(p0) = COPY $x0
+    %cst:gpr(s64) = G_CONSTANT i64 -256
+    %ptr_add:gpr(p0) = G_PTR_ADD %ptr(p0), %cst
+    %val:gpr(s64) = G_LOAD %ptr_add(p0) :: (load 8 from @x)
+    $x1 = COPY %val(s64)
+    RET_ReallyLR implicit $x1
+
+...
+---
+name:            ldurdi
+legalized:       true
+regBankSelected: true
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $x0, $d1
+    ; CHECK-LABEL: name: ldurdi
+    ; CHECK: liveins: $x0, $d1
+    ; CHECK: %ptr:gpr64sp = COPY $x0
+    ; CHECK: %val:fpr64 = LDURDi %ptr, -256 :: (load 8 from @x)
+    ; CHECK: $d1 = COPY %val
+    ; CHECK: RET_ReallyLR implicit $d1
+    %ptr:gpr(p0) = COPY $x0
+    %cst:gpr(s64) = G_CONSTANT i64 -256
+    %ptr_add:gpr(p0) = G_PTR_ADD %ptr(p0), %cst
+    %val:fpr(s64) = G_LOAD %ptr_add(p0) :: (load 8 from @x)
+    $d1 = COPY %val(s64)
+    RET_ReallyLR implicit $d1
+
+...
+---
+name:            ldurhhi
+legalized:       true
+regBankSelected: true
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $x0, $w1
+    ; CHECK-LABEL: name: ldurhhi
+    ; CHECK: liveins: $x0, $w1
+    ; CHECK: %ptr:gpr64sp = COPY $x0
+    ; CHECK: %val:gpr32 = LDURHHi %ptr, -256 :: (load 2 from @x)
+    ; CHECK: %val_wide:gpr32all = COPY %val
+    ; CHECK: $w1 = COPY %val_wide
+    ; CHECK: RET_ReallyLR implicit $w1
+    %ptr:gpr(p0) = COPY $x0
+    %cst:gpr(s64) = G_CONSTANT i64 -256
+    %ptr_add:gpr(p0) = G_PTR_ADD %ptr(p0), %cst
+    %val:gpr(s16) = G_LOAD %ptr_add(p0) :: (load 2 from @x)
+    %val_wide:gpr(s32) = G_ZEXT %val(s16)
+    $w1 = COPY %val_wide(s32)
+    RET_ReallyLR implicit $w1
+
+...
+---
+name:            ldurhi
+legalized:       true
+regBankSelected: true
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $x0, $h1
+    ; CHECK-LABEL: name: ldurhi
+    ; CHECK: liveins: $x0, $h1
+    ; CHECK: %ptr:gpr64sp = COPY $x0
+    ; CHECK: %val:fpr16 = LDURHi %ptr, -256 :: (load 2 from @x)
+    ; CHECK: $h1 = COPY %val
+    ; CHECK: RET_ReallyLR implicit $h1
+    %ptr:gpr(p0) = COPY $x0
+    %cst:gpr(s64) = G_CONSTANT i64 -256
+    %ptr_add:gpr(p0) = G_PTR_ADD %ptr(p0), %cst
+    %val:fpr(s16) = G_LOAD %ptr_add(p0) :: (load 2 from @x)
+    $h1 = COPY %val(s16)
+    RET_ReallyLR implicit $h1
+
+...
+---
+name:            ldurbbi
+legalized:       true
+regBankSelected: true
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $x0, $w1
+    ; CHECK-LABEL: name: ldurbbi
+    ; CHECK: liveins: $x0, $w1
+    ; CHECK: %ptr:gpr64sp = COPY $x0
+    ; CHECK: %val:gpr32 = LDURBBi %ptr, -256 :: (load 1 from @x)
+    ; CHECK: %val_wide:gpr32all = COPY %val
+    ; CHECK: $w1 = COPY %val_wide
+    ; CHECK: RET_ReallyLR implicit $w1
+    %ptr:gpr(p0) = COPY $x0
+    %cst:gpr(s64) = G_CONSTANT i64 -256
+    %ptr_add:gpr(p0) = G_PTR_ADD %ptr(p0), %cst
+    %val:gpr(s8) = G_LOAD %ptr_add(p0) :: (load 1 from @x)
+    %val_wide:gpr(s32) = G_ZEXT %val(s8)
+    $w1 = COPY %val_wide(s32)
+    RET_ReallyLR implicit $w1

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

--- |
   define void @store_4xi16(<4 x i16> %v, <4 x i16>* %ptr) { ret void }
   define void @store_4xi32(<4 x i32> %v, <4 x i32>* %ptr) { ret void }
   define void @store_8xi16(<8 x i16> %v, <8 x i16>* %ptr) { ret void }
   define void @store_16xi8(<16 x i8> %v, <16 x i8>* %ptr) { ret void }
 
   @x = external hidden local_unnamed_addr global i32*, align 8
   define void @store_adrp_add_low() { ret void }
 
+  ; Check selection for unscaled addressing modes.
+  define void @sturwi() { ret void }
+  define void @stursi() { ret void }
+  define void @sturxi() { ret void }
+  define void @sturdi() { ret void }
+  define void @sturhhi() { ret void }
+  define void @sturhi() { ret void }
+  define void @sturbbi() { ret void }
+
 ...
 
 ---
 name:            store_s64_gpr
 legalized:       true
 regBankSelected: true
 
 registers:
bb.0:
     ; CHECK: liveins: $x0
     ; CHECK: %copy:gpr64 = COPY $x0
     ; CHECK: %adrp:gpr64common = ADRP target-flags(aarch64-page) @x
     ; CHECK: STRXui %copy, %adrp, target-flags(aarch64-pageoff, aarch64-nc) @x :: (store 8 into @x)
     %copy:gpr(p0) = COPY $x0
     %adrp:gpr64(p0) = ADRP target-flags(aarch64-page) @x
     %add_low:gpr(p0) = G_ADD_LOW %adrp(p0), target-flags(aarch64-pageoff, aarch64-nc) @x
     G_STORE %copy(p0), %add_low(p0) :: (store 8 into @x)
+
+...
+---
+name:            sturwi
+legalized:       true
+regBankSelected: true
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $x0, $w1
+    ; CHECK-LABEL: name: sturwi
+    ; CHECK: liveins: $x0, $w1
+    ; CHECK: %ptr:gpr64sp = COPY $x0
+    ; CHECK: %val:gpr32 = COPY $w1
+    ; CHECK: STURWi %val, %ptr, -256 :: (store 4 into @x)
+    %ptr:gpr(p0) = COPY $x0
+    %val:gpr(s32) = COPY $w1
+    %cst:gpr(s64) = G_CONSTANT i64 -256
+    %ptr_add:gpr(p0) = G_PTR_ADD %ptr(p0), %cst
+    G_STORE %val(s32), %ptr_add(p0) :: (store 4 into @x)
+
+...
+---
+name:            stursi
+legalized:       true
+regBankSelected: true
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $x0, $s1
+    ; CHECK-LABEL: name: stursi
+    ; CHECK: liveins: $x0, $s1
+    ; CHECK: %ptr:gpr64sp = COPY $x0
+    ; CHECK: %val:fpr32 = COPY $s1
+    ; CHECK: STURSi %val, %ptr, -256 :: (store 4 into @x)
+    %ptr:gpr(p0) = COPY $x0
+    %val:fpr(s32) = COPY $s1
+    %cst:gpr(s64) = G_CONSTANT i64 -256
+    %ptr_add:gpr(p0) = G_PTR_ADD %ptr(p0), %cst
+    G_STORE %val(s32), %ptr_add(p0) :: (store 4 into @x)
+
+...
+---
+name:            sturxi
+legalized:       true
+regBankSelected: true
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $x0, $x1
+    ; CHECK-LABEL: name: sturxi
+    ; CHECK: liveins: $x0, $x1
+    ; CHECK: %ptr:gpr64sp = COPY $x0
+    ; CHECK: %val:gpr64 = COPY $x1
+    ; CHECK: STURXi %val, %ptr, -256 :: (store 8 into @x)
+    %ptr:gpr(p0) = COPY $x0
+    %val:gpr(s64) = COPY $x1
+    %cst:gpr(s64) = G_CONSTANT i64 -256
+    %ptr_add:gpr(p0) = G_PTR_ADD %ptr(p0), %cst
+    G_STORE %val(s64), %ptr_add(p0) :: (store 8 into @x)
+
+...
+---
+name:            sturdi
+legalized:       true
+regBankSelected: true
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $x0, $d1
+    ; CHECK-LABEL: name: sturdi
+    ; CHECK: liveins: $x0, $d1
+    ; CHECK: %ptr:gpr64sp = COPY $x0
+    ; CHECK: %val:fpr64 = COPY $d1
+    ; CHECK: STURDi %val, %ptr, -256 :: (store 8 into @x)
+    %ptr:gpr(p0) = COPY $x0
+    %val:fpr(s64) = COPY $d1
+    %cst:gpr(s64) = G_CONSTANT i64 -256
+    %ptr_add:gpr(p0) = G_PTR_ADD %ptr(p0), %cst
+    G_STORE %val(s64), %ptr_add(p0) :: (store 8 into @x)
+
+...
+---
+name:            sturhhi
+legalized:       true
+regBankSelected: true
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $x0, $x1
+    ; CHECK-LABEL: name: sturhhi
+    ; CHECK: liveins: $x0, $x1
+    ; CHECK: %ptr:gpr64sp = COPY $x0
+    ; CHECK: %val:gpr64 = COPY $x1
+    ; CHECK: %val_trunc:gpr32 = COPY %val.sub_32
+    ; CHECK: STURHHi %val_trunc, %ptr, -256 :: (store 2 into @x)
+    %ptr:gpr(p0) = COPY $x0
+    %val:gpr(s64) = COPY $x1
+    %val_trunc:gpr(s16) = G_TRUNC %val(s64)
+    %cst:gpr(s64) = G_CONSTANT i64 -256
+    %ptr_add:gpr(p0) = G_PTR_ADD %ptr(p0), %cst
+    G_STORE %val_trunc(s16), %ptr_add(p0) :: (store 2 into @x)
+
+...
+---
+name:            sturhi
+legalized:       true
+regBankSelected: true
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $x0, $h1
+    ; CHECK-LABEL: name: sturhi
+    ; CHECK: liveins: $x0, $h1
+    ; CHECK: %ptr:gpr64sp = COPY $x0
+    ; CHECK: %val:fpr16 = COPY $h1
+    ; CHECK: STURHi %val, %ptr, -256 :: (store 2 into @x)
+    %ptr:gpr(p0) = COPY $x0
+    %val:fpr(s16) = COPY $h1
+    %cst:gpr(s64) = G_CONSTANT i64 -256
+    %ptr_add:gpr(p0) = G_PTR_ADD %ptr(p0), %cst
+    G_STORE %val(s16), %ptr_add(p0) :: (store 2 into @x)
+
+...
+---
+name:            sturbbi
+legalized:       true
+regBankSelected: true
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $x0, $x1
+    ; CHECK-LABEL: name: sturbbi
+    ; CHECK: liveins: $x0, $x1
+    ; CHECK: %ptr:gpr64sp = COPY $x0
+    ; CHECK: %val:gpr64 = COPY $x1
+    ; CHECK: %val_trunc:gpr32 = COPY %val.sub_32
+    ; CHECK: STURBBi %val_trunc, %ptr, -256 :: (store 1 into @x)
+    %ptr:gpr(p0) = COPY $x0
+    %val:gpr(s64) = COPY $x1
+    %val_trunc:gpr(s8) = G_TRUNC %val(s64)
+    %cst:gpr(s64) = G_CONSTANT i64 -256
+    %ptr_add:gpr(p0) = G_PTR_ADD %ptr(p0), %cst
+    G_STORE %val_trunc(s8), %ptr_add(p0) :: (store 1 into @x)