Rework the patch, v3

Author: krzysz00

llvm/lib/Target/AMDGPU/AMDGPUTargetTransformInfo.cpp

@@ -313,6 +313,24 @@ bool GCNTTIImpl::hasBranchDivergence(const Function *F) const {
   return !F || !ST->isSingleLaneExecution(*F);
 }
 
+unsigned GCNTTIImpl::getRegUsageForType(Type *Ty) {
+  if (auto *VT = dyn_cast<FixedVectorType>(Ty)) {
+    if (auto *PT = dyn_cast<PointerType>(VT->getElementType())) {
+      switch (PT->getAddressSpace()) {
+      // Assume that the resource parts of the vector being asked about are the
+      // same.
+      case AMDGPUAS::BUFFER_FAT_POINTER:
+        return 4 + VT->getNumElements();
+      case AMDGPUAS::BUFFER_STRIDED_POINTER:
+        return 4 + 2 * VT->getNumElements();
+      default:
+        break;
+      }
+    }
+  }
+  return BaseT::getRegUsageForType(Ty);
+}
+
 unsigned GCNTTIImpl::getNumberOfRegisters(unsigned RCID) const {
   // NB: RCID is not an RCID. In fact it is 0 or 1 for scalar or vector
   // registers. See getRegisterClassForType for the implementation.

llvm/lib/Target/AMDGPU/AMDGPUTargetTransformInfo.h

@@ -113,6 +113,14 @@ class GCNTTIImpl final : public BasicTTIImplBase<GCNTTIImpl> {
   void getPeelingPreferences(Loop *L, ScalarEvolution &SE,
                              TTI::PeelingPreferences &PP);
 
+  // Vectorization will query for the number of registers needed for
+  // <N x ptr addrspace(7/9)> and the default implementation will cause crashes,
+  // so override it here. This also lets us account for the fact that, in the
+  // context of loop vectorization (which is what uses this API), the number of
+  // registers needed for fat pointers is lower because they'll share a resource
+  // part.
+  unsigned getRegUsageForType(Type *Ty);
+
   TTI::PopcntSupportKind getPopcntSupport(unsigned TyWidth) {
     assert(isPowerOf2_32(TyWidth) && "Ty width must be power of 2");
     return TTI::PSK_FastHardware;