[GISel] Enforce G_PTR_ADD RHS type matching index size for addr space

llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp

@@ -4004,7 +4004,14 @@ Register LegalizerHelper::getVectorElementPointer(Register VecPtr, LLT VecTy,
 
   Index = clampVectorIndex(MIRBuilder, Index, VecTy);
 
-  LLT IdxTy = MRI.getType(Index);
+  // Convert index to the correct size for the address space.
+  const DataLayout &DL = MIRBuilder.getDataLayout();
+  unsigned AS = MRI.getType(VecPtr).getAddressSpace();
+  unsigned IndexSizeInBits = DL.getIndexSize(AS) * 8;
+  LLT IdxTy = MRI.getType(Index).changeElementSize(IndexSizeInBits);
+  if (IdxTy != MRI.getType(Index))
+    Index = MIRBuilder.buildSExtOrTrunc(IdxTy, Index).getReg(0);
+
   auto Mul = MIRBuilder.buildMul(IdxTy, Index,
                                  MIRBuilder.buildConstant(IdxTy, EltSize));
 

llvm/lib/CodeGen/MachineVerifier.cpp

@@ -1301,6 +1301,16 @@ void MachineVerifier::verifyPreISelGenericInstruction(const MachineInstr *MI) {
     if (OffsetTy.isPointerOrPointerVector())
       report("gep offset operand must not be a pointer", MI);
 
+    if (PtrTy.isPointerOrPointerVector()) {
+      const DataLayout &DL = MF->getDataLayout();
+      unsigned AS = PtrTy.getAddressSpace();
+      unsigned IndexSizeInBits = DL.getIndexSize(AS) * 8;
+      if (OffsetTy.getScalarSizeInBits() != IndexSizeInBits) {
+        report("gep offset operand must match index size for address space",
+               MI);
+      }
+    }
+
     // TODO: Is the offset allowed to be a scalar with a vector?
     break;
   }

llvm/test/CodeGen/AArch64/GlobalISel/combine-ptradd-int2ptr.mir

@@ -11,7 +11,7 @@ body:             |
     ; CHECK: [[C:%[0-9]+]]:_(p64) = G_CONSTANT i64 44
     ; CHECK: [[PTRTOINT:%[0-9]+]]:_(s64) = G_PTRTOINT [[C]](p64)
     ; CHECK: $x0 = COPY [[PTRTOINT]](s64)
-    %1:_(s32) = G_CONSTANT i32 42
+    %1:_(s64) = G_CONSTANT i64 42
     %2:_(s32) = G_CONSTANT i32 2
     %3:_(p64) = G_INTTOPTR %2
     %4:_(p64) = G_PTR_ADD %3, %1
@@ -26,7 +26,7 @@ body:             |
     ; CHECK-LABEL: name: agc.test_combine_ptradd_constants_ptrres
     ; CHECK: [[C:%[0-9]+]]:_(p64) = G_CONSTANT i64 44
     ; CHECK: $x0 = COPY [[C]](p64)
-    %1:_(s32) = G_CONSTANT i32 42
+    %1:_(s64) = G_CONSTANT i64 42
     %2:_(s32) = G_CONSTANT i32 2
     %3:_(p64) = G_INTTOPTR %2
     %4:_(p64) = G_PTR_ADD %3, %1
@@ -39,12 +39,12 @@ body:             |
     liveins: $x0, $x1
     ; Ensure non-constant G_PTR_ADDs are not folded.
     ; CHECK-LABEL: name: agc.test_not_combine_variable_ptradd
-    ; CHECK: [[C:%[0-9]+]]:_(s32) = G_CONSTANT i32 42
+    ; CHECK: [[C:%[0-9]+]]:_(s64) = G_CONSTANT i64 42
     ; CHECK: [[COPY:%[0-9]+]]:_(p64) = COPY $x1
-    ; CHECK: [[PTR_ADD:%[0-9]+]]:_(p64) = G_PTR_ADD [[COPY]], [[C]](s32)
+    ; CHECK: [[PTR_ADD:%[0-9]+]]:_(p64) = G_PTR_ADD [[COPY]], [[C]](s64)
     ; CHECK: [[PTRTOINT:%[0-9]+]]:_(s64) = G_PTRTOINT [[PTR_ADD]](p64)
     ; CHECK: $x0 = COPY [[PTRTOINT]](s64)
-    %1:_(s32) = G_CONSTANT i32 42
+    %1:_(s64) = G_CONSTANT i64 42
     %2:_(p64) = COPY $x1
     %3:_(p64) = G_PTR_ADD %2, %1
     %4:_(s64) = G_PTRTOINT %3

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

@@ -1,23 +1,6 @@
 # NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
 # RUN: llc -mtriple=aarch64 -run-pass=legalizer %s -o - | FileCheck %s
 ---
-name:            test_ptr_add_small
-body:             |
-  bb.0.entry:
-    ; CHECK-LABEL: name: test_ptr_add_small
-    ; CHECK: [[COPY:%[0-9]+]]:_(p0) = COPY $x0
-    ; CHECK: [[COPY1:%[0-9]+]]:_(s64) = COPY $x1
-    ; CHECK: [[SEXT_INREG:%[0-9]+]]:_(s64) = G_SEXT_INREG [[COPY1]], 8
-    ; CHECK: [[PTR_ADD:%[0-9]+]]:_(p0) = G_PTR_ADD [[COPY]], [[SEXT_INREG]](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: