[SPIR-V] Fixup storage class for global private (#118318)

Re-land of #116636
Adds a new address spaces: hlsl_private. Variables with such address
space will be emitted with a Private storage class.
This is useful for variables global to a SPIR-V module, since up to now,
they were still emitted with a Function storage class, which is wrong.

---------

Signed-off-by: Nathan Gauër <[email protected]>

Author: Keenuts

llvm/lib/Target/SPIRV/SPIRVInstructionSelector.cpp

@@ -1460,6 +1460,16 @@ bool SPIRVInstructionSelector::selectAddrSpaceCast(Register ResVReg,
         .addUse(SrcPtr)
         .constrainAllUses(TII, TRI, RBI);
 
+  if ((SrcSC == SPIRV::StorageClass::Function &&
+       DstSC == SPIRV::StorageClass::Private) ||
+      (DstSC == SPIRV::StorageClass::Function &&
+       SrcSC == SPIRV::StorageClass::Private)) {
+    return BuildMI(BB, I, DL, TII.get(TargetOpcode::COPY))
+        .addDef(ResVReg)
+        .addUse(SrcPtr)
+        .constrainAllUses(TII, TRI, RBI);
+  }
+
   // Casting from an eligible pointer to Generic.
   if (DstSC == SPIRV::StorageClass::Generic && isGenericCastablePtr(SrcSC))
     return selectUnOp(ResVReg, ResType, I, SPIRV::OpPtrCastToGeneric);
@@ -3461,11 +3471,7 @@ bool SPIRVInstructionSelector::selectGlobalValue(
   if (HasInit && !Init)
     return true;
 
-  unsigned AddrSpace = GV->getAddressSpace();
-  SPIRV::StorageClass::StorageClass Storage =
-      addressSpaceToStorageClass(AddrSpace, STI);
-  bool HasLnkTy = GV->getLinkage() != GlobalValue::InternalLinkage &&
-                  Storage != SPIRV::StorageClass::Function;
+  bool HasLnkTy = GV->getLinkage() != GlobalValue::InternalLinkage;
   SPIRV::LinkageType::LinkageType LnkType =
       (GV->isDeclaration() || GV->hasAvailableExternallyLinkage())
           ? SPIRV::LinkageType::Import
@@ -3474,12 +3480,14 @@ bool SPIRVInstructionSelector::selectGlobalValue(
                  ? SPIRV::LinkageType::LinkOnceODR
                  : SPIRV::LinkageType::Export);
 
-  SPIRVType *ResType = GR.getOrCreateSPIRVPointerType(
-      PointerBaseType, I, TII,
-      addressSpaceToStorageClass(GV->getAddressSpace(), STI));
-  Register Reg = GR.buildGlobalVariable(ResVReg, ResType, GlobalIdent, GV,
-                                        Storage, Init, GlobalVar->isConstant(),
-                                        HasLnkTy, LnkType, MIRBuilder, true);
+  const unsigned AddrSpace = GV->getAddressSpace();
+  SPIRV::StorageClass::StorageClass StorageClass =
+      addressSpaceToStorageClass(AddrSpace, STI);
+  SPIRVType *ResType =
+      GR.getOrCreateSPIRVPointerType(PointerBaseType, I, TII, StorageClass);
+  Register Reg = GR.buildGlobalVariable(
+      ResVReg, ResType, GlobalIdent, GV, StorageClass, Init,
+      GlobalVar->isConstant(), HasLnkTy, LnkType, MIRBuilder, true);
   return Reg.isValid();
 }
 

llvm/lib/Target/SPIRV/SPIRVLegalizerInfo.cpp

@@ -112,13 +112,16 @@ SPIRVLegalizerInfo::SPIRVLegalizerInfo(const SPIRVSubtarget &ST) {
   const LLT p5 =
       LLT::pointer(5, PSize); // Input, SPV_INTEL_usm_storage_classes (Device)
   const LLT p6 = LLT::pointer(6, PSize); // SPV_INTEL_usm_storage_classes (Host)
+  const LLT p7 = LLT::pointer(7, PSize); // Input
+  const LLT p8 = LLT::pointer(8, PSize); // Output
+  const LLT p10 = LLT::pointer(10, PSize); // Private
 
   // TODO: remove copy-pasting here by using concatenation in some way.
   auto allPtrsScalarsAndVectors = {
-      p0,    p1,    p2,    p3,    p4,     p5,     p6,    s1,   s8,   s16,
-      s32,   s64,   v2s1,  v2s8,  v2s16,  v2s32,  v2s64, v3s1, v3s8, v3s16,
-      v3s32, v3s64, v4s1,  v4s8,  v4s16,  v4s32,  v4s64, v8s1, v8s8, v8s16,
-      v8s32, v8s64, v16s1, v16s8, v16s16, v16s32, v16s64};
+      p0,   p1,   p2,    p3,    p4,    p5,    p6,    p7,     p8,     p10,
+      s1,   s8,   s16,   s32,   s64,   v2s1,  v2s8,  v2s16,  v2s32,  v2s64,
+      v3s1, v3s8, v3s16, v3s32, v3s64, v4s1,  v4s8,  v4s16,  v4s32,  v4s64,
+      v8s1, v8s8, v8s16, v8s32, v8s64, v16s1, v16s8, v16s16, v16s32, v16s64};
 
   auto allVectors = {v2s1,  v2s8,   v2s16,  v2s32, v2s64, v3s1,  v3s8,
                      v3s16, v3s32,  v3s64,  v4s1,  v4s8,  v4s16, v4s32,
@@ -145,10 +148,10 @@ SPIRVLegalizerInfo::SPIRVLegalizerInfo(const SPIRVSubtarget &ST) {
       s16,   s32,   s64,   v2s16, v2s32, v2s64, v3s16,  v3s32,  v3s64,
       v4s16, v4s32, v4s64, v8s16, v8s32, v8s64, v16s16, v16s32, v16s64};
 
-  auto allFloatAndIntScalarsAndPtrs = {s8, s16, s32, s64, p0, p1,
-                                       p2, p3,  p4,  p5,  p6};
+  auto allFloatAndIntScalarsAndPtrs = {s8, s16, s32, s64, p0, p1, p2,
+                                       p3, p4,  p5,  p6,  p7, p8, p10};
 
-  auto allPtrs = {p0, p1, p2, p3, p4, p5, p6};
+  auto allPtrs = {p0, p1, p2, p3, p4, p5, p6, p7, p8, p10};
 
   bool IsExtendedInts =
       ST.canUseExtension(

llvm/lib/Target/SPIRV/SPIRVUtils.cpp

@@ -207,8 +207,12 @@ addressSpaceToStorageClass(unsigned AddrSpace, const SPIRVSubtarget &STI) {
                : SPIRV::StorageClass::CrossWorkgroup;
   case 7:
     return SPIRV::StorageClass::Input;
+  case 8:
+    return SPIRV::StorageClass::Output;
   case 9:
     return SPIRV::StorageClass::CodeSectionINTEL;
+  case 10:
+    return SPIRV::StorageClass::Private;
   default:
     report_fatal_error("Unknown address space");
   }

llvm/lib/Target/SPIRV/SPIRVUtils.h

@@ -170,8 +170,12 @@ storageClassToAddressSpace(SPIRV::StorageClass::StorageClass SC) {
     return 6;
   case SPIRV::StorageClass::Input:
     return 7;
+  case SPIRV::StorageClass::Output:
+    return 8;
   case SPIRV::StorageClass::CodeSectionINTEL:
     return 9;
+  case SPIRV::StorageClass::Private:
+    return 10;
   default:
     report_fatal_error("Unable to get address space id");
   }

llvm/test/CodeGen/SPIRV/pointers/global-addrspacecast.ll

@@ -0,0 +1,17 @@
+; RUN: llc -verify-machineinstrs -O0 -mtriple=spirv64-unknown-unknown %s -o - | FileCheck %s
+; RUN: %if spirv-tools %{ llc -O0 -mtriple=spirv64-unknown-unknown %s -o - -filetype=obj | spirv-val %}
+
+@PrivInternal = internal addrspace(10) global i32 456
+; CHECK-DAG:  %[[#type:]] = OpTypeInt 32 0
+; CHECK-DAG: %[[#ptrty:]] = OpTypePointer Private %[[#type]]
+; CHECK-DAG: %[[#value:]] = OpConstant %[[#type]] 456
+; CHECK-DAG:   %[[#var:]] = OpVariable %[[#ptrty]] Private %[[#value]]
+
+define spir_kernel void @Foo() {
+  %p = addrspacecast ptr addrspace(10) @PrivInternal to ptr
+  %v = load i32, ptr %p, align 4
+  ret void
+; CHECK:      OpLabel
+; CHECK-NEXT: OpLoad %[[#type]] %[[#var]] Aligned 4
+; CHECK-Next: OpReturn
+}

llvm/test/CodeGen/SPIRV/pointers/variables-storage-class-vk.ll

@@ -0,0 +1,15 @@
+; RUN: llc -verify-machineinstrs -O0 -mtriple=spirv-unknown-vulkan1.3-compute %s -o - | FileCheck %s
+; RUN: %if spirv-tools %{ llc -O0 -mtriple=spirv-unknown-vulkan1.3-compute %s -o - -filetype=obj | spirv-val %}
+
+; CHECK-DAG: %[[#U32:]] = OpTypeInt 32 0
+
+; CHECK-DAG: %[[#VAL:]] = OpConstant %[[#U32]] 456
+; CHECK-DAG: %[[#VTYPE:]] = OpTypePointer Private %[[#U32]]
+; CHECK-DAG: %[[#VAR:]] = OpVariable %[[#VTYPE]] Private %[[#VAL]]
+; CHECK-NOT: OpDecorate %[[#VAR]] LinkageAttributes
+@PrivInternal = internal addrspace(10) global i32 456
+
+define void @main() {
+  %l = load i32, ptr addrspace(10) @PrivInternal
+  ret void
+}

llvm/test/CodeGen/SPIRV/pointers/variables-storage-class.ll

@@ -1,17 +1,29 @@
 ; RUN: llc -verify-machineinstrs -O0 -mtriple=spirv64-unknown-unknown %s -o - | FileCheck %s
 ; RUN: %if spirv-tools %{ llc -O0 -mtriple=spirv64-unknown-unknown %s -o - -filetype=obj | spirv-val %}
 
+; CHECK-DAG: %[[#U8:]] = OpTypeInt 8 0
+; CHECK-DAG: %[[#U32:]] = OpTypeInt 32 0
+
+; CHECK-DAG: %[[#TYPE:]] = OpTypePointer CrossWorkgroup %[[#U8]]
+; CHECK-DAG: %[[#VAL:]] = OpConstantNull %[[#TYPE]]
+; CHECK-DAG: %[[#VTYPE:]] = OpTypePointer CrossWorkgroup %[[#TYPE]]
+; CHECK-DAG: %[[#PTR:]] = OpVariable %[[#VTYPE]] CrossWorkgroup %[[#VAL]]
 @Ptr = addrspace(1) global ptr addrspace(1) null
-@Init = private addrspace(2) constant i32 123
 
-; CHECK-DAG: %[[#PTR:]] = OpVariable %[[#]] UniformConstant %[[#]]
-; CHECK-DAG: %[[#INIT:]] = OpVariable %[[#]] CrossWorkgroup %[[#]]
+; CHECK-DAG: %[[#VAL:]] = OpConstant %[[#U32]] 123
+; CHECK-DAG: %[[#VTYPE:]] = OpTypePointer UniformConstant %[[#U32]]
+; CHECK-DAG: %[[#INIT:]] = OpVariable %[[#VTYPE]] UniformConstant %[[#VAL]]
+@Init = private addrspace(2) constant i32 123
 
-; CHECK: %[[#]] = OpLoad %[[#]] %[[#INIT]] Aligned 8
-; CHECK: OpCopyMemorySized %[[#]] %[[#PTR]] %[[#]] Aligned 4
+; CHECK-DAG: %[[#VAL:]] = OpConstant %[[#U32]] 456
+; CHECK-DAG: %[[#VTYPE:]] = OpTypePointer Private %[[#U32]]
+; CHECK-DAG: %[[#]] = OpVariable %[[#VTYPE]] Private %[[#VAL]]
+@PrivInternal = internal addrspace(10) global i32 456
 
 define spir_kernel void @Foo() {
+  ; CHECK: %[[#]] = OpLoad %[[#]] %[[#PTR]] Aligned 8
   %l = load ptr addrspace(1), ptr addrspace(1) @Ptr, align 8
+  ; CHECK: OpCopyMemorySized %[[#]] %[[#INIT]] %[[#]] Aligned 4
   call void @llvm.memcpy.p1.p2.i64(ptr addrspace(1) align 4 %l, ptr addrspace(2) align 1 @Init, i64 4, i1 false)
   ret void
 }