[AArch64][GlobalISel] Legalize vector G_PTR_ADD and enable selection.

[aemerson]

Summary

For selection, we need to add support for vectors in the PTR_ADD conversion function.

[paquette]

llvm/test/CodeGen/AArch64/GlobalISel/select-ptr-add.mir

2

IOS prefix isn't used?

3–10

Do you actually need the IR here?

[aemerson]

Update test.

[paquette]

LGTM

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

bool AArch64InstructionSelector::convertPtrAddToAdd(
     MachineInstr &I, MachineRegisterInfo &MRI) {
   assert(I.getOpcode() == TargetOpcode::G_PTR_ADD && "Expected G_PTR_ADD");
   Register DstReg = I.getOperand(0).getReg();
   Register AddOp1Reg = I.getOperand(1).getReg();
   const LLT PtrTy = MRI.getType(DstReg);
   if (PtrTy.getAddressSpace() != 0)
     return false;
 
-  // Only do this for scalars for now.
-  if (PtrTy.isVector())
-    return false;
-
   MachineIRBuilder MIB(I);
-  const LLT s64 = LLT::scalar(64);
-  auto PtrToInt = MIB.buildPtrToInt(s64, AddOp1Reg);
+  const LLT CastPtrTy = PtrTy.isVector() ? LLT::vector(2, 64) : LLT::scalar(64);
+  auto PtrToInt = MIB.buildPtrToInt(CastPtrTy, AddOp1Reg);
   // Set regbanks on the registers.
+  if (PtrTy.isVector())
+    MRI.setRegBank(PtrToInt.getReg(0), RBI.getRegBank(AArch64::FPRRegBankID));
+  else
     MRI.setRegBank(PtrToInt.getReg(0), RBI.getRegBank(AArch64::GPRRegBankID));
 
   // Now turn the %dst(p0) = G_PTR_ADD %base, off into:
-  // %dst(s64) = G_ADD %intbase, off
+  // %dst(intty) = G_ADD %intbase, off
   I.setDesc(TII.get(TargetOpcode::G_ADD));
-  MRI.setType(DstReg, s64);
+  MRI.setType(DstReg, CastPtrTy);
   I.getOperand(1).setReg(PtrToInt.getReg(0));
   if (!select(*PtrToInt)) {
     LLVM_DEBUG(dbgs() << "Failed to select G_PTRTOINT in convertPtrAddToAdd");
     return false;
   }
   return true;
 }
 
if (DstRB.getID() == AArch64::GPRRegBankID) {
         MachineIRBuilder MIB(I);
         MachineInstr *Extract = emitExtractVectorElt(
             DstReg, DstRB, LLT::scalar(DstTy.getSizeInBits()), SrcReg, 0, MIB);
         if (!Extract)
           return false;
         I.eraseFromParent();
         return true;
       }
+
+      // We might have a vector G_PTRTOINT, in which case just emit a COPY.
+      if (Opcode == TargetOpcode::G_PTRTOINT) {
+        assert(DstTy.isVector() && "Expected an FPR ptrtoint to be a vector");
+        I.setDesc(TII.get(TargetOpcode::COPY));
+        return true;
+      }
     }
 
     return false;
   }
 
   case TargetOpcode::G_ANYEXT: {
     const Register DstReg = I.getOperand(0).getReg();
     const Register SrcReg = I.getOperand(1).getReg();

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

getActionDefinitionsBuilder(G_SHL)
     .clampScalar(0, s32, s64)
     .widenScalarToNextPow2(0)
     .clampNumElements(0, v2s32, v4s32)
     .clampNumElements(0, v2s64, v2s64)
     .moreElementsToNextPow2(0)
     .minScalarSameAs(1, 0);
 
   getActionDefinitionsBuilder(G_PTR_ADD)
-      .legalFor({{p0, s64}})
+      .legalFor({{p0, s64}, {v2p0, v2s64}})
       .clampScalar(1, s64, s64);
 
   getActionDefinitionsBuilder(G_PTRMASK).legalFor({{p0, s64}});
 
   getActionDefinitionsBuilder({G_SDIV, G_UDIV})
       .legalFor({s32, s64})
       .libcallFor({s128})
       .clampScalar(0, s32, s64)

llvm/test/CodeGen/AArch64/GlobalISel/legalize-ptr-add.mir

 # NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
 # RUN: llc -march=aarch64 -run-pass=legalizer %s -o - | FileCheck %s
 ---
-name:            test_gep_small
+name:            test_ptr_add_small
 body:             |
   bb.0.entry:
-    ; CHECK-LABEL: name: test_gep_small
+    ; CHECK-LABEL: name: test_ptr_add_small
     ; CHECK: [[COPY:%[0-9]+]]:_(p0) = COPY $x0
     ; CHECK: [[COPY1:%[0-9]+]]:_(s64) = COPY $x1
     ; CHECK: [[COPY2:%[0-9]+]]:_(s64) = COPY [[COPY1]](s64)
     ; CHECK: [[C:%[0-9]+]]:_(s64) = G_CONSTANT i64 56
     ; CHECK: [[SHL:%[0-9]+]]:_(s64) = G_SHL [[COPY2]], [[C]](s64)
     ; CHECK: [[ASHR:%[0-9]+]]:_(s64) = G_ASHR [[SHL]], [[C]](s64)
-    ; CHECK: [[GEP:%[0-9]+]]:_(p0) = G_PTR_ADD [[COPY]], [[ASHR]](s64)
-    ; CHECK: $x0 = COPY [[GEP]](p0)
+    ; CHECK: [[PTR_ADD:%[0-9]+]]:_(p0) = G_PTR_ADD [[COPY]], [[ASHR]](s64)
+    ; CHECK: $x0 = COPY [[PTR_ADD]](p0)
     %0:_(p0) = COPY $x0
     %1:_(s64) = COPY $x1
     %2:_(s8) = G_TRUNC %1(s64)
     %3:_(p0) = G_PTR_ADD %0, %2(s8)
     $x0 = COPY %3(p0)
 
 ...
+---
+name:            test_ptr_add_vec_p0
+body:             |
+  bb.0.entry:
+    ; CHECK-LABEL: name: test_ptr_add_vec_p0
+    ; CHECK: [[COPY:%[0-9]+]]:_(<2 x p0>) = COPY $q0
+    ; CHECK: [[COPY1:%[0-9]+]]:_(<2 x s64>) = COPY $q1
+    ; CHECK: [[PTR_ADD:%[0-9]+]]:_(<2 x p0>) = G_PTR_ADD [[COPY]], [[COPY1]](<2 x s64>)
+    ; CHECK: $q0 = COPY [[PTR_ADD]](<2 x p0>)
+    %0:_(<2 x p0>) = COPY $q0
+    %1:_(<2 x s64>) = COPY $q1
+    %3:_(<2 x p0>) = G_PTR_ADD %0, %1(<2 x s64>)
+    $q0 = COPY %3(<2 x p0>)
+
+...

llvm/test/CodeGen/AArch64/GlobalISel/select-ptr-add.mir

+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
+# RUN: llc -O0 -mtriple=aarch64-apple-ios -run-pass=instruction-select -verify-machineinstrs %s -o - | FileCheck %s -check-prefix=CHECK -check-prefix=IOS
+--- |
+  target datalayout = "e-m:o-i64:64-i128:128-n32:64-S128"
+
+  define i8* @ptr_add(i8* %in) { ret i8* undef }
+  define i8* @ptr_add_no_constant(i8* %in) { ret i8* undef }
+  define i8* @ptr_add_bad_imm(i8* %in) { ret i8* undef }
+  define i8* @ptr_add_vec(i8* %in) { ret i8* undef }
+
+...
+
+---
+name:            ptr_add
+legalized:       true
+regBankSelected: true
+body:             |
+  bb.0:
+      liveins: $x0
+    ; CHECK-LABEL: name: ptr_add
+    ; CHECK: [[COPY:%[0-9]+]]:gpr64 = COPY $x0
+    ; CHECK: [[COPY1:%[0-9]+]]:gpr64common = COPY [[COPY]]
+    ; CHECK: [[ADDXri:%[0-9]+]]:gpr64sp = ADDXri [[COPY1]], 42, 0
+    ; CHECK: $x0 = COPY [[ADDXri]]
+    %0:gpr(p0) = COPY $x0
+    %1:gpr(s64) = G_CONSTANT i64 42
+    %2:gpr(p0) = G_PTR_ADD %0, %1(s64)
+    $x0 = COPY %2(p0)
+...
+
+---
+name:            ptr_add_no_constant
+legalized:       true
+regBankSelected: true
+body:             |
+  bb.0:
+      liveins: $x0, $x1
+    ; CHECK-LABEL: name: ptr_add_no_constant
+    ; CHECK: [[COPY:%[0-9]+]]:gpr64 = COPY $x0
+    ; CHECK: [[COPY1:%[0-9]+]]:gpr64 = COPY $x1
+    ; CHECK: [[ADDXrr:%[0-9]+]]:gpr64 = ADDXrr [[COPY]], [[COPY1]]
+    ; CHECK: $x0 = COPY [[ADDXrr]]
+    %0:gpr(p0) = COPY $x0
+    %1:gpr(s64) = COPY $x1
+    %2:gpr(p0) = G_PTR_ADD %0, %1(s64)
+    $x0 = COPY %2(p0)
+...
+
+---
+name:            ptr_add_bad_imm
+legalized:       true
+regBankSelected: true
+body:             |
+  bb.0:
+      liveins: $x0, $x1
+    ; CHECK-LABEL: name: ptr_add_bad_imm
+    ; CHECK: [[COPY:%[0-9]+]]:gpr64 = COPY $x0
+    ; CHECK: [[MOVi32imm:%[0-9]+]]:gpr32 = MOVi32imm 10000
+    ; CHECK: [[SUBREG_TO_REG:%[0-9]+]]:gpr64 = SUBREG_TO_REG 0, [[MOVi32imm]], %subreg.sub_32
+    ; CHECK: [[ADDXrr:%[0-9]+]]:gpr64 = ADDXrr [[COPY]], [[SUBREG_TO_REG]]
+    ; CHECK: $x0 = COPY [[ADDXrr]]
+    %0:gpr(p0) = COPY $x0
+    %1:gpr(s64) = G_CONSTANT i64 10000
+    %2:gpr(p0) = G_PTR_ADD %0, %1(s64)
+    $x0 = COPY %2(p0)
+...
+
+---
+name:            ptr_add_vec
+legalized:       true
+regBankSelected: true
+body:             |
+  bb.0:
+      liveins: $q0, $q1
+    ; CHECK-LABEL: name: ptr_add_vec
+    ; CHECK: [[COPY:%[0-9]+]]:fpr128 = COPY $q0
+    ; CHECK: [[COPY1:%[0-9]+]]:fpr128 = COPY $q1
+    ; CHECK: [[ADDv2i64_:%[0-9]+]]:fpr128 = ADDv2i64 [[COPY]], [[COPY1]]
+    ; CHECK: $q0 = COPY [[ADDv2i64_]]
+    %0:fpr(<2 x p0>) = COPY $q0
+    %1:fpr(<2 x s64>) = COPY $q1
+    %3:fpr(<2 x p0>) = G_PTR_ADD %0, %1(<2 x s64>)
+    $q0 = COPY %3(<2 x p0>)
+...

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

 # RUN: llc -O0 -mtriple=aarch64-apple-ios -run-pass=instruction-select -verify-machineinstrs %s -o - | FileCheck %s -check-prefix=CHECK -check-prefix=IOS
 # RUN: llc -O0 -mtriple=aarch64-linux-gnu -run-pass=instruction-select -verify-machineinstrs %s -o - | FileCheck %s -check-prefix=CHECK -check-prefix=LINUX-DEFAULT
 # RUN: llc -O0 -mtriple=aarch64-linux-gnu -relocation-model=pic -run-pass=instruction-select -verify-machineinstrs %s -o - | FileCheck %s -check-prefix=CHECK -check-prefix=LINUX-PIC
 
 --- |
   target datalayout = "e-m:o-i64:64-i128:128-n32:64-S128"
 
   define void @frame_index() {
     %ptr0 = alloca i64
     ret void
   }
 
-  define i8* @ptr_add(i8* %in) { ret i8* undef }
-  define i8* @ptr_add_no_constant(i8* %in) { ret i8* undef }
-  define i8* @ptr_add_bad_imm(i8* %in) { ret i8* undef }
-
   define i8* @ptr_mask(i8* %in) { ret i8* undef }
 
   @var_local = global i8 0
   define i8* @global_local() { ret i8* undef }
 
   @var_got = external global i8
   define i8* @global_got() { ret i8* undef }
 

 # CHECK: %0:gpr64sp = ADDXri %stack.0.ptr0, 0, 0
 body:             |
   bb.0:
     %0(p0) = G_FRAME_INDEX %stack.0.ptr0
     $x0 = COPY %0(p0)
 ...
 
 ---
-# CHECK-LABEL: name: ptr_add
-name:            ptr_add
-legalized:       true
-regBankSelected: true
-registers:
-  - { id: 0, class: gpr }
-  - { id: 1, class: gpr }
-  - { id: 2, class: gpr }
-
-# CHECK:  body:
-# CHECK: %{{[0-9]+}}:gpr64sp = ADDXri %{{[0-9]+}}, 42, 0
-body:             |
-  bb.0:
-      liveins: $x0
-    %0(p0) = COPY $x0
-    %1(s64) = G_CONSTANT i64 42
-    %2(p0) = G_PTR_ADD %0, %1(s64)
-    $x0 = COPY %2(p0)
-...
-
----
-# CHECK-LABEL: name: ptr_add_no_constant
-name:            ptr_add_no_constant
-legalized:       true
-regBankSelected: true
-registers:
-  - { id: 0, class: gpr }
-  - { id: 1, class: gpr }
-  - { id: 2, class: gpr }
-
-# CHECK:  body:
-# CHECK: %{{[0-9]+}}:gpr64 = ADDXrr %{{[0-9]+}}, %{{[0-9]+}}
-body:             |
-  bb.0:
-      liveins: $x0, $x1
-    %0(p0) = COPY $x0
-    %1(s64) = COPY $x1
-    %2(p0) = G_PTR_ADD %0, %1(s64)
-    $x0 = COPY %2(p0)
-...
-
----
-# CHECK-LABEL: name: ptr_add_bad_imm
-name:            ptr_add_bad_imm
-legalized:       true
-regBankSelected: true
-registers:
-  - { id: 0, class: gpr }
-  - { id: 1, class: gpr }
-  - { id: 2, class: gpr }
-
-# CHECK:  body:
-# CHECK: %0:gpr64 = COPY $x0
-# CHECK: %4:gpr32 = MOVi32imm 10000
-# CHECK: %1:gpr64 = SUBREG_TO_REG 0, %4, %subreg.sub_32
-# CHECK: %{{[0-9]+}}:gpr64 = ADDXrr %0, %1
-body:             |
-  bb.0:
-      liveins: $x0, $x1
-    %0(p0) = COPY $x0
-    %1(s64) = G_CONSTANT i64 10000
-    %2(p0) = G_PTR_ADD %0, %1(s64)
-    $x0 = COPY %2(p0)
-...
 
 ---
 # CHECK-LABEL: name: ptr_mask
 name:            ptr_mask
 legalized:       true
 regBankSelected: true
 
 # CHECK:  body: