[GlobalISel][AArch64] Fold G_GEP into LDR/STR ui addressing mode.

We're not to the point of supporting the load/store patterns yet
(because they extensively use PatFrags).

But in the meantime, we can implement some of the simplest addressing
modes.

llvm-svn: 298863

3 files changed, 487 insertions, 1 deletions

diff --git a/llvm/lib/Target/AArch64/AArch64InstructionSelector.cpp b/llvm/lib/Target/AArch64/AArch64InstructionSelector.cpp
index 146eea07a56..be0f6eb44d1 100644
--- a/llvm/lib/Target/AArch64/AArch64InstructionSelector.cpp
+++ b/llvm/lib/Target/AArch64/AArch64InstructionSelector.cpp
@@ -774,7 +774,25 @@ bool AArch64InstructionSelector::select(MachineInstr &I) const {
 
     I.setDesc(TII.get(NewOpc));
 
-    I.addOperand(MachineOperand::CreateImm(0));
+    uint64_t Offset = 0;
+    auto *PtrMI = MRI.getVRegDef(PtrReg);
+
+    // Try to fold a GEP into our unsigned immediate addressing mode.
+    if (PtrMI->getOpcode() == TargetOpcode::G_GEP) {
+      if (auto COff = getConstantVRegVal(PtrMI->getOperand(2).getReg(), MRI)) {
+        int64_t Imm = *COff;
+        const unsigned Size = MemTy.getSizeInBits() / 8;
+        const unsigned Scale = Log2_32(Size);
+        if ((Imm & (Size - 1)) == 0 && Imm >= 0 && Imm < (0x1000 << Scale)) {
+          unsigned Ptr2Reg = PtrMI->getOperand(1).getReg();
+          I.getOperand(1).setReg(Ptr2Reg);
+          PtrMI = MRI.getVRegDef(Ptr2Reg);
+          Offset = Imm / Size;
+        }
+      }
+    }
+
+    I.addOperand(MachineOperand::CreateImm(Offset));
 
     // If we're storing a 0, use WZR/XZR.
     if (auto CVal = getConstantVRegVal(ValReg, MRI)) {
diff --git a/llvm/test/CodeGen/AArch64/GlobalISel/select-load.mir b/llvm/test/CodeGen/AArch64/GlobalISel/select-load.mir
index 0821fc76446..09b40201f9a 100644
--- a/llvm/test/CodeGen/AArch64/GlobalISel/select-load.mir
+++ b/llvm/test/CodeGen/AArch64/GlobalISel/select-load.mir
@@ -7,10 +7,22 @@
   define void @load_s32_gpr(i32* %addr) { ret void }
   define void @load_s16_gpr(i16* %addr) { ret void }
   define void @load_s8_gpr(i8* %addr) { ret void }
+
+  define void @load_gep_128_s64_gpr(i64* %addr) { ret void }
+  define void @load_gep_512_s32_gpr(i32* %addr) { ret void }
+  define void @load_gep_64_s16_gpr(i16* %addr) { ret void }
+  define void @load_gep_1_s8_gpr(i8* %addr) { ret void }
+
   define void @load_s64_fpr(i64* %addr) { ret void }
   define void @load_s32_fpr(i32* %addr) { ret void }
   define void @load_s16_fpr(i16* %addr) { ret void }
   define void @load_s8_fpr(i8* %addr) { ret void }
+
+  define void @load_gep_8_s64_fpr(i64* %addr) { ret void }
+  define void @load_gep_16_s32_fpr(i32* %addr) { ret void }
+  define void @load_gep_64_s16_fpr(i16* %addr) { ret void }
+  define void @load_gep_32_s8_fpr(i8* %addr) { ret void }
+
 ...
 
 ---
@@ -114,6 +126,134 @@ body:             |
 ...
 
 ---
+# CHECK-LABEL: name: load_gep_128_s64_gpr
+name:            load_gep_128_s64_gpr
+legalized:       true
+regBankSelected: true
+
+# CHECK:      registers:
+# CHECK-NEXT:  - { id: 0, class: gpr64sp }
+# CHECK-NEXT:  - { id: 1, class: gpr }
+# CHECK-NEXT:  - { id: 2, class: gpr }
+# CHECK-NEXT:  - { id: 3, class: gpr64 }
+registers:
+  - { id: 0, class: gpr }
+  - { id: 1, class: gpr }
+  - { id: 2, class: gpr }
+  - { id: 3, class: gpr }
+
+# CHECK:  body:
+# CHECK: %0 = COPY %x0
+# CHECK: %3 = LDRXui %0, 16 :: (load 8 from %ir.addr)
+# CHECK: %x0 = COPY %3
+body:             |
+  bb.0:
+    liveins: %x0
+
+    %0(p0) = COPY %x0
+    %1(s64) = G_CONSTANT i64 128
+    %2(p0) = G_GEP %0, %1
+    %3(s64) = G_LOAD %2 :: (load 8 from %ir.addr)
+    %x0 = COPY %3
+...
+
+---
+# CHECK-LABEL: name: load_gep_512_s32_gpr
+name:            load_gep_512_s32_gpr
+legalized:       true
+regBankSelected: true
+
+# CHECK:      registers:
+# CHECK-NEXT:  - { id: 0, class: gpr64sp }
+# CHECK-NEXT:  - { id: 1, class: gpr }
+# CHECK-NEXT:  - { id: 2, class: gpr }
+# CHECK-NEXT:  - { id: 3, class: gpr32 }
+registers:
+  - { id: 0, class: gpr }
+  - { id: 1, class: gpr }
+  - { id: 2, class: gpr }
+  - { id: 3, class: gpr }
+
+# CHECK:  body:
+# CHECK: %0 = COPY %x0
+# CHECK: %3 = LDRWui %0, 128 :: (load 4 from %ir.addr)
+# CHECK: %w0 = COPY %3
+body:             |
+  bb.0:
+    liveins: %x0
+
+    %0(p0) = COPY %x0
+    %1(s64) = G_CONSTANT i64 512
+    %2(p0) = G_GEP %0, %1
+    %3(s32) = G_LOAD %2 :: (load 4 from %ir.addr)
+    %w0 = COPY %3
+...
+
+---
+# CHECK-LABEL: name: load_gep_64_s16_gpr
+name:            load_gep_64_s16_gpr
+legalized:       true
+regBankSelected: true
+
+# CHECK:      registers:
+# CHECK-NEXT:  - { id: 0, class: gpr64sp }
+# CHECK-NEXT:  - { id: 1, class: gpr }
+# CHECK-NEXT:  - { id: 2, class: gpr }
+# CHECK-NEXT:  - { id: 3, class: gpr32 }
+registers:
+  - { id: 0, class: gpr }
+  - { id: 1, class: gpr }
+  - { id: 2, class: gpr }
+  - { id: 3, class: gpr }
+
+# CHECK:  body:
+# CHECK: %0 = COPY %x0
+# CHECK: %3 = LDRHHui %0, 32 :: (load 2 from %ir.addr)
+# CHECK: %w0 = COPY %3
+body:             |
+  bb.0:
+    liveins: %x0
+
+    %0(p0) = COPY %x0
+    %1(s64) = G_CONSTANT i64 64
+    %2(p0) = G_GEP %0, %1
+    %3(s16) = G_LOAD %2 :: (load 2 from %ir.addr)
+    %w0 = COPY %3
+...
+
+---
+# CHECK-LABEL: name: load_gep_1_s8_gpr
+name:            load_gep_1_s8_gpr
+legalized:       true
+regBankSelected: true
+
+# CHECK:      registers:
+# CHECK-NEXT:  - { id: 0, class: gpr64sp }
+# CHECK-NEXT:  - { id: 1, class: gpr }
+# CHECK-NEXT:  - { id: 2, class: gpr }
+# CHECK-NEXT:  - { id: 3, class: gpr32 }
+registers:
+  - { id: 0, class: gpr }
+  - { id: 1, class: gpr }
+  - { id: 2, class: gpr }
+  - { id: 3, class: gpr }
+
+# CHECK:  body:
+# CHECK: %0 = COPY %x0
+# CHECK: %3 = LDRBBui %0, 1 :: (load 1 from %ir.addr)
+# CHECK: %w0 = COPY %3
+body:             |
+  bb.0:
+    liveins: %x0
+
+    %0(p0) = COPY %x0
+    %1(s64) = G_CONSTANT i64 1
+    %2(p0) = G_GEP %0, %1
+    %3(s8) = G_LOAD %2 :: (load 1 from %ir.addr)
+    %w0 = COPY %3
+...
+
+---
 # CHECK-LABEL: name: load_s64_fpr
 name:            load_s64_fpr
 legalized:       true
@@ -212,3 +352,131 @@ body:             |
     %1(s8) = G_LOAD  %0 :: (load 1 from %ir.addr)
     %b0 = COPY %1(s8)
 ...
+
+---
+# CHECK-LABEL: name: load_gep_8_s64_fpr
+name:            load_gep_8_s64_fpr
+legalized:       true
+regBankSelected: true
+
+# CHECK:      registers:
+# CHECK-NEXT:  - { id: 0, class: gpr64sp }
+# CHECK-NEXT:  - { id: 1, class: gpr }
+# CHECK-NEXT:  - { id: 2, class: gpr }
+# CHECK-NEXT:  - { id: 3, class: fpr64 }
+registers:
+  - { id: 0, class: gpr }
+  - { id: 1, class: gpr }
+  - { id: 2, class: gpr }
+  - { id: 3, class: fpr }
+
+# CHECK:  body:
+# CHECK: %0 = COPY %x0
+# CHECK: %3 = LDRDui %0, 1 :: (load 8 from %ir.addr)
+# CHECK: %d0 = COPY %3
+body:             |
+  bb.0:
+    liveins: %x0
+
+    %0(p0) = COPY %x0
+    %1(s64) = G_CONSTANT i64 8
+    %2(p0) = G_GEP %0, %1
+    %3(s64) = G_LOAD %2 :: (load 8 from %ir.addr)
+    %d0 = COPY %3
+...
+
+---
+# CHECK-LABEL: name: load_gep_16_s32_fpr
+name:            load_gep_16_s32_fpr
+legalized:       true
+regBankSelected: true
+
+# CHECK:      registers:
+# CHECK-NEXT:  - { id: 0, class: gpr64sp }
+# CHECK-NEXT:  - { id: 1, class: gpr }
+# CHECK-NEXT:  - { id: 2, class: gpr }
+# CHECK-NEXT:  - { id: 3, class: fpr32 }
+registers:
+  - { id: 0, class: gpr }
+  - { id: 1, class: gpr }
+  - { id: 2, class: gpr }
+  - { id: 3, class: fpr }
+
+# CHECK:  body:
+# CHECK: %0 = COPY %x0
+# CHECK: %3 = LDRSui %0, 4 :: (load 4 from %ir.addr)
+# CHECK: %s0 = COPY %3
+body:             |
+  bb.0:
+    liveins: %x0
+
+    %0(p0) = COPY %x0
+    %1(s64) = G_CONSTANT i64 16
+    %2(p0) = G_GEP %0, %1
+    %3(s32) = G_LOAD %2 :: (load 4 from %ir.addr)
+    %s0 = COPY %3
+...
+
+---
+# CHECK-LABEL: name: load_gep_64_s16_fpr
+name:            load_gep_64_s16_fpr
+legalized:       true
+regBankSelected: true
+
+# CHECK:      registers:
+# CHECK-NEXT:  - { id: 0, class: gpr64sp }
+# CHECK-NEXT:  - { id: 1, class: gpr }
+# CHECK-NEXT:  - { id: 2, class: gpr }
+# CHECK-NEXT:  - { id: 3, class: fpr16 }
+registers:
+  - { id: 0, class: gpr }
+  - { id: 1, class: gpr }
+  - { id: 2, class: gpr }
+  - { id: 3, class: fpr }
+
+# CHECK:  body:
+# CHECK: %0 = COPY %x0
+# CHECK: %3 = LDRHui %0, 32 :: (load 2 from %ir.addr)
+# CHECK: %h0 = COPY %3
+body:             |
+  bb.0:
+    liveins: %x0
+
+    %0(p0) = COPY %x0
+    %1(s64) = G_CONSTANT i64 64
+    %2(p0) = G_GEP %0, %1
+    %3(s16) = G_LOAD %2 :: (load 2 from %ir.addr)
+    %h0 = COPY %3
+...
+
+---
+# CHECK-LABEL: name: load_gep_32_s8_fpr
+name:            load_gep_32_s8_fpr
+legalized:       true
+regBankSelected: true
+
+# CHECK:      registers:
+# CHECK-NEXT:  - { id: 0, class: gpr64sp }
+# CHECK-NEXT:  - { id: 1, class: gpr }
+# CHECK-NEXT:  - { id: 2, class: gpr }
+# CHECK-NEXT:  - { id: 3, class: fpr8 }
+registers:
+  - { id: 0, class: gpr }
+  - { id: 1, class: gpr }
+  - { id: 2, class: gpr }
+  - { id: 3, class: fpr }
+
+# CHECK:  body:
+# CHECK: %0 = COPY %x0
+# CHECK: %3 = LDRBui %0, 32 :: (load 1 from %ir.addr)
+# CHECK: %b0 = COPY %3
+body:             |
+  bb.0:
+    liveins: %x0
+
+    %0(p0) = COPY %x0
+    %1(s64) = G_CONSTANT i64 32
+    %2(p0) = G_GEP %0, %1
+    %3(s8) = G_LOAD %2 :: (load 1 from %ir.addr)
+    %b0 = COPY %3
+...
diff --git a/llvm/test/CodeGen/AArch64/GlobalISel/select-store.mir b/llvm/test/CodeGen/AArch64/GlobalISel/select-store.mir
index 95690f98690..d40a1231df2 100644
--- a/llvm/test/CodeGen/AArch64/GlobalISel/select-store.mir
+++ b/llvm/test/CodeGen/AArch64/GlobalISel/select-store.mir
@@ -11,8 +11,16 @@
   define void @store_zero_s64_gpr(i64* %addr) { ret void }
   define void @store_zero_s32_gpr(i32* %addr) { ret void }
 
+  define void @store_gep_128_s64_gpr(i64* %addr) { ret void }
+  define void @store_gep_512_s32_gpr(i32* %addr) { ret void }
+  define void @store_gep_64_s16_gpr(i16* %addr) { ret void }
+  define void @store_gep_1_s8_gpr(i8* %addr) { ret void }
+
   define void @store_s64_fpr(i64* %addr) { ret void }
   define void @store_s32_fpr(i32* %addr) { ret void }
+
+  define void @store_gep_8_s64_fpr(i64* %addr) { ret void }
+  define void @store_gep_8_s32_fpr(i32* %addr) { ret void }
 ...
 
 ---
@@ -176,6 +184,134 @@ body:             |
 ...
 
 ---
+# CHECK-LABEL: name: store_gep_128_s64_gpr
+name:            store_gep_128_s64_gpr
+legalized:       true
+regBankSelected: true
+
+# CHECK:      registers:
+# CHECK-NEXT:  - { id: 0, class: gpr64sp }
+# CHECK-NEXT:  - { id: 1, class: gpr64 }
+# CHECK-NEXT:  - { id: 2, class: gpr }
+# CHECK-NEXT:  - { id: 3, class: gpr }
+registers:
+  - { id: 0, class: gpr }
+  - { id: 1, class: gpr }
+  - { id: 2, class: gpr }
+  - { id: 3, class: gpr }
+
+# CHECK:  body:
+# CHECK: %0 = COPY %x0
+# CHECK: %1 = COPY %x1
+# CHECK: STRXui %1, %0, 16 :: (store 8 into %ir.addr)
+body:             |
+  bb.0:
+    liveins: %x0, %x1
+
+    %0(p0) = COPY %x0
+    %1(s64) = COPY %x1
+    %2(s64) = G_CONSTANT i64 128
+    %3(p0) = G_GEP %0, %2
+    G_STORE %1, %3 :: (store 8 into %ir.addr)
+...
+
+---
+# CHECK-LABEL: name: store_gep_512_s32_gpr
+name:            store_gep_512_s32_gpr
+legalized:       true
+regBankSelected: true
+
+# CHECK:      registers:
+# CHECK-NEXT:  - { id: 0, class: gpr64sp }
+# CHECK-NEXT:  - { id: 1, class: gpr32 }
+# CHECK-NEXT:  - { id: 2, class: gpr }
+# CHECK-NEXT:  - { id: 3, class: gpr }
+registers:
+  - { id: 0, class: gpr }
+  - { id: 1, class: gpr }
+  - { id: 2, class: gpr }
+  - { id: 3, class: gpr }
+
+# CHECK:  body:
+# CHECK: %0 = COPY %x0
+# CHECK: %1 = COPY %w1
+# CHECK: STRWui %1, %0, 128 :: (store 4 into %ir.addr)
+body:             |
+  bb.0:
+    liveins: %x0, %w1
+
+    %0(p0) = COPY %x0
+    %1(s32) = COPY %w1
+    %2(s64) = G_CONSTANT i64 512
+    %3(p0) = G_GEP %0, %2
+    G_STORE %1, %3 :: (store 4 into %ir.addr)
+...
+
+---
+# CHECK-LABEL: name: store_gep_64_s16_gpr
+name:            store_gep_64_s16_gpr
+legalized:       true
+regBankSelected: true
+
+# CHECK:      registers:
+# CHECK-NEXT:  - { id: 0, class: gpr64sp }
+# CHECK-NEXT:  - { id: 1, class: gpr32 }
+# CHECK-NEXT:  - { id: 2, class: gpr }
+# CHECK-NEXT:  - { id: 3, class: gpr }
+registers:
+  - { id: 0, class: gpr }
+  - { id: 1, class: gpr }
+  - { id: 2, class: gpr }
+  - { id: 3, class: gpr }
+
+# CHECK:  body:
+# CHECK: %0 = COPY %x0
+# CHECK: %1 = COPY %w1
+# CHECK: STRHHui %1, %0, 32 :: (store 2 into %ir.addr)
+body:             |
+  bb.0:
+    liveins: %x0, %w1
+
+    %0(p0) = COPY %x0
+    %1(s16) = COPY %w1
+    %2(s64) = G_CONSTANT i64 64
+    %3(p0) = G_GEP %0, %2
+    G_STORE %1, %3 :: (store 2 into %ir.addr)
+...
+
+---
+# CHECK-LABEL: name: store_gep_1_s8_gpr
+name:            store_gep_1_s8_gpr
+legalized:       true
+regBankSelected: true
+
+# CHECK:      registers:
+# CHECK-NEXT:  - { id: 0, class: gpr64sp }
+# CHECK-NEXT:  - { id: 1, class: gpr32 }
+# CHECK-NEXT:  - { id: 2, class: gpr }
+# CHECK-NEXT:  - { id: 3, class: gpr }
+registers:
+  - { id: 0, class: gpr }
+  - { id: 1, class: gpr }
+  - { id: 2, class: gpr }
+  - { id: 3, class: gpr }
+
+# CHECK:  body:
+# CHECK: %0 = COPY %x0
+# CHECK: %1 = COPY %w1
+# CHECK: STRBBui %1, %0, 1 :: (store 1 into %ir.addr)
+body:             |
+  bb.0:
+    liveins: %x0, %w1
+
+    %0(p0) = COPY %x0
+    %1(s8) = COPY %w1
+    %2(s64) = G_CONSTANT i64 1
+    %3(p0) = G_GEP %0, %2
+    G_STORE %1, %3 :: (store 1 into %ir.addr)
+...
+
+---
 # CHECK-LABEL: name: store_s64_fpr
 name:            store_s64_fpr
 legalized:       true
@@ -228,3 +364,67 @@ body:             |
     G_STORE  %1, %0 :: (store 4 into %ir.addr)
 
 ...
+
+---
+# CHECK-LABEL: name: store_gep_8_s64_fpr
+name:            store_gep_8_s64_fpr
+legalized:       true
+regBankSelected: true
+
+# CHECK:      registers:
+# CHECK-NEXT:  - { id: 0, class: gpr64sp }
+# CHECK-NEXT:  - { id: 1, class: fpr64 }
+# CHECK-NEXT:  - { id: 2, class: gpr }
+# CHECK-NEXT:  - { id: 3, class: gpr }
+registers:
+  - { id: 0, class: gpr }
+  - { id: 1, class: fpr }
+  - { id: 2, class: gpr }
+  - { id: 3, class: gpr }
+
+# CHECK:  body:
+# CHECK: %0 = COPY %x0
+# CHECK: %1 = COPY %d1
+# CHECK: STRDui %1, %0, 1 :: (store 8 into %ir.addr)
+body:             |
+  bb.0:
+    liveins: %x0, %d1
+
+    %0(p0) = COPY %x0
+    %1(s64) = COPY %d1
+    %2(s64) = G_CONSTANT i64 8
+    %3(p0) = G_GEP %0, %2
+    G_STORE %1, %3 :: (store 8 into %ir.addr)
+...
+
+---
+# CHECK-LABEL: name: store_gep_8_s32_fpr
+name:            store_gep_8_s32_fpr
+legalized:       true
+regBankSelected: true
+
+# CHECK:      registers:
+# CHECK-NEXT:  - { id: 0, class: gpr64sp }
+# CHECK-NEXT:  - { id: 1, class: fpr32 }
+# CHECK-NEXT:  - { id: 2, class: gpr }
+# CHECK-NEXT:  - { id: 3, class: gpr }
+registers:
+  - { id: 0, class: gpr }
+  - { id: 1, class: fpr }
+  - { id: 2, class: gpr }
+  - { id: 3, class: gpr }
+
+# CHECK:  body:
+# CHECK: %0 = COPY %x0
+# CHECK: %1 = COPY %s1
+# CHECK: STRSui %1, %0, 2 :: (store 4 into %ir.addr)
+body:             |
+  bb.0:
+    liveins: %x0, %s1
+
+    %0(p0) = COPY %x0
+    %1(s32) = COPY %s1
+    %2(s64) = G_CONSTANT i64 8
+    %3(p0) = G_GEP %0, %2
+    G_STORE %1, %3 :: (store 4 into %ir.addr)
+...