[AMDGPU] Override getRegUsageForType() to fix <N x ptr(7)> crash

llvm/include/llvm/CodeGen/TargetLowering.h

@@ -1675,8 +1675,8 @@ class TargetLoweringBase {
   /// operations except for the pointer size.  If AllowUnknown is true, this
   /// will return MVT::Other for types with no EVT counterpart (e.g. structs),
   /// otherwise it will assert.
-  EVT getValueType(const DataLayout &DL, Type *Ty,
-                   bool AllowUnknown = false) const {
+  virtual EVT getValueType(const DataLayout &DL, Type *Ty,
+                           bool AllowUnknown = false) const {
     // Lower scalar pointers to native pointer types.
     if (auto *PTy = dyn_cast<PointerType>(Ty))
       return getPointerTy(DL, PTy->getAddressSpace());
@@ -1695,8 +1695,8 @@ class TargetLoweringBase {
     return EVT::getEVT(Ty, AllowUnknown);
   }
 
-  EVT getMemValueType(const DataLayout &DL, Type *Ty,
-                      bool AllowUnknown = false) const {
+  virtual EVT getMemValueType(const DataLayout &DL, Type *Ty,
+                              bool AllowUnknown = false) const {
     // Lower scalar pointers to native pointer types.
     if (auto *PTy = dyn_cast<PointerType>(Ty))
       return getPointerMemTy(DL, PTy->getAddressSpace());
@@ -1714,7 +1714,6 @@ class TargetLoweringBase {
     return getValueType(DL, Ty, AllowUnknown);
   }
 
-
   /// Return the MVT corresponding to this LLVM type. See getValueType.
   MVT getSimpleValueType(const DataLayout &DL, Type *Ty,
                          bool AllowUnknown = false) const {

llvm/lib/Target/AMDGPU/SIISelLowering.cpp

@@ -1203,6 +1203,41 @@ MVT SITargetLowering::getPointerMemTy(const DataLayout &DL, unsigned AS) const {
   return AMDGPUTargetLowering::getPointerMemTy(DL, AS);
 }
 
+/// Passes like the loop vectorizer will, for example, try to query the size in
+/// registers of buffer fat pointer. They don't exist by the time we reach
+/// codegen, but these queries can still come in. Unfortunately, something like
+/// <2 x ptr addrspace(7)> will get lowered to <2 x v5i32> by the workarounds
+/// above, which causes a crash. Handle this case here.
+EVT SITargetLowering::getValueType(const DataLayout &DL, Type *Ty,
+                                   bool AllowUnknown) const {
+  if (auto *VT = dyn_cast<VectorType>(Ty)) {
+    if (auto *PT = dyn_cast<PointerType>(VT->getElementType())) {
+      MVT MET = getPointerTy(DL, PT->getAddressSpace());
+      if (MET.isVector() && MET.getVectorElementType() == MVT::i32) {
+        return EVT::getVectorVT(
+            Ty->getContext(), EVT(MET.getVectorElementType()),
+            VT->getElementCount() * MET.getVectorNumElements());
+      }
+    }
+  }
+  return AMDGPUTargetLowering::getValueType(DL, Ty, AllowUnknown);
+}
+
+EVT SITargetLowering::getMemValueType(const DataLayout &DL, Type *Ty,
+                                      bool AllowUnknown) const {
+  if (auto *VT = dyn_cast<VectorType>(Ty)) {
+    if (auto *PT = dyn_cast<PointerType>(VT->getElementType())) {
+      MVT ScalarTy = getPointerMemTy(DL, PT->getAddressSpace());
+      if (ScalarTy.isVector() && ScalarTy.getVectorElementType() == MVT::i32) {
+        return EVT::getVectorVT(
+            Ty->getContext(), EVT(ScalarTy.getVectorElementType()),
+            VT->getElementCount() * ScalarTy.getVectorNumElements());
+      }
+    }
+  }
+  return AMDGPUTargetLowering::getValueType(DL, Ty, AllowUnknown);
+}
+
 bool SITargetLowering::getTgtMemIntrinsic(IntrinsicInfo &Info,
                                           const CallInst &CI,
                                           MachineFunction &MF,

llvm/lib/Target/AMDGPU/SIISelLowering.h

@@ -306,6 +306,10 @@ class SITargetLowering final : public AMDGPUTargetLowering {
   // so, to work around the lack of i160, map it to v5i32.
   MVT getPointerTy(const DataLayout &DL, unsigned AS) const override;
   MVT getPointerMemTy(const DataLayout &DL, unsigned AS) const override;
+  EVT getValueType(const DataLayout &DL, Type *Ty,
+                   bool AllowUnknown = false) const override;
+  EVT getMemValueType(const DataLayout &DL, Type *Ty,
+                      bool AllowUnknown = false) const override;
 
   bool getTgtMemIntrinsic(IntrinsicInfo &, const CallInst &,
                           MachineFunction &MF,

llvm/test/Transforms/LoopVectorize/AMDGPU/buffer-fat-pointer.ll

@@ -0,0 +1,39 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; RUN: opt -mtriple=amdgcn-amd-amdhsa -mcpu=gfx942 -passes=loop-vectorize -S < %s | FileCheck %s
+
+; Reduced from a crash, variables added to make things more realistic.
+; This is a roundabout test for TargetLowering::getValueType() returning
+; a reasonable value for <N x p7> instead of asserting.
+define amdgpu_kernel void @_dynamic_pack_simple_dispatch_0_pack_i32(ptr addrspace(1) %.ptr, i64 %0) {
+; CHECK-LABEL: define amdgpu_kernel void @_dynamic_pack_simple_dispatch_0_pack_i32(
+; CHECK-SAME: ptr addrspace(1) [[DOTPTR:%.*]], i64 [[TMP0:%.*]]) #[[ATTR0:[0-9]+]] {
+; CHECK-NEXT:  [[_LR_PH5:.*:]]
+; CHECK-NEXT:    [[DOTRSRC:%.*]] = call ptr addrspace(8) @llvm.amdgcn.make.buffer.rsrc.p1(ptr addrspace(1) [[DOTPTR]], i16 0, i32 -2147483648, i32 159744)
+; CHECK-NEXT:    [[TMP1:%.*]] = addrspacecast ptr addrspace(8) [[DOTRSRC]] to ptr addrspace(7)
+; CHECK-NEXT:    br label %[[BB2:.*]]
+; CHECK:       [[BB2]]:
+; CHECK-NEXT:    [[TMP3:%.*]] = phi i64 [ 0, [[DOTLR_PH5:%.*]] ], [ [[TMP5:%.*]], %[[BB2]] ]
+; CHECK-NEXT:    [[TMP4:%.*]] = getelementptr i32, ptr addrspace(7) [[TMP1]], i32 0
+; CHECK-NEXT:    [[TMP5]] = add i64 [[TMP3]], 1
+; CHECK-NEXT:    [[EXITCOND_NOT:%.*]] = icmp eq i64 [[TMP3]], [[TMP0]]
+; CHECK-NEXT:    br i1 [[EXITCOND_NOT]], [[DOT_CRIT_EDGE_LOOPEXIT:label %.*]], label %[[BB2]]
+; CHECK:       [[__CRIT_EDGE_LOOPEXIT:.*:]]
+; CHECK-NEXT:    ret void
+;
+.lr.ph5:
+  %.rsrc = call ptr addrspace(8) @llvm.amdgcn.make.buffer.rsrc.p1(ptr addrspace(1) %.ptr, i16 0, i32 2147483648, i32 159744)
+  %1 = addrspacecast ptr addrspace(8) %.rsrc to ptr addrspace(7)
+  br label %2
+
+2:                                                ; preds = %2, %.lr.ph5
+  %3 = phi i64 [ 0, %.lr.ph5 ], [ %5, %2 ]
+  %4 = getelementptr i32, ptr addrspace(7) %1, i32 0
+  %5 = add i64 %3, 1
+  %exitcond.not = icmp eq i64 %3, %0
+  br i1 %exitcond.not, label %._crit_edge.loopexit, label %2
+
+._crit_edge.loopexit:                             ; preds = %2
+  ret void
+}
+
+declare ptr addrspace(8) @llvm.amdgcn.make.buffer.rsrc.p1(ptr addrspace(1) readnone, i16, i32, i32)