Make TargetLowering::getValueType() virtual to fix <N x ptr(7)> crash

Even though ptr addrspace(7) is rewritten away into more primitive
constructs before reaching SelectionDAG or GlobalISel, we stil
sometimes need to query properties like how many registers it will
require.

We already had an existing workaroung that would map ptr
addrspace(7) (and addrspace(9)) to MVT::{5,6}i32, their ultimate
in-register representations, in overloads of
TargetLowering::getPointerType(DL, AddressSpace).

However, whenever TargetLowering::getValue() tried to construct a
vector VT out of those vector MVTs, the vector constructor would
assert because you can't have a vector of vectors.

This commit solves the crash by manually overriding getValueTy() and
getMemValueType() in the AMDGPU TargetLowering.

This is something of a big hammer, and I'm open to suggestions for a
more precise change.

Author: krzysz00

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)