[SPIRV] Lower spirv.Layout type

Add this type in accordance with the [design doc](llvm/wg-hlsl#171).

Fixes #138276

Author: s-perron

llvm/lib/Target/SPIRV/SPIRVBuiltins.cpp

@@ -3163,6 +3163,12 @@ static SPIRVType *getVulkanBufferType(const TargetExtType *ExtensionType,
   return GR->getOrCreateVulkanBufferType(MIRBuilder, T, SC, IsWritable);
 }
 
+static SPIRVType *getLayoutType(const TargetExtType *ExtensionType,
+                                MachineIRBuilder &MIRBuilder,
+                                SPIRVGlobalRegistry *GR) {
+  return GR->getOrCreateLayoutType(MIRBuilder, ExtensionType);
+}
+
 namespace SPIRV {
 TargetExtType *parseBuiltinTypeNameToTargetExtType(std::string TypeName,
                                                    LLVMContext &Context) {
@@ -3240,6 +3246,8 @@ SPIRVType *lowerBuiltinType(const Type *OpaqueType,
     TargetType = getInlineSpirvType(BuiltinType, MIRBuilder, GR);
   } else if (Name == "spirv.VulkanBuffer") {
     TargetType = getVulkanBufferType(BuiltinType, MIRBuilder, GR);
+  } else if (Name == "spirv.Layout") {
+    TargetType = getLayoutType(BuiltinType, MIRBuilder, GR);
   } else {
     // Lookup the demangled builtin type in the TableGen records.
     const SPIRV::BuiltinType *TypeRecord = SPIRV::lookupBuiltinType(Name);

llvm/lib/Target/SPIRV/SPIRVEmitIntrinsics.cpp

@@ -673,10 +673,16 @@ Type *SPIRVEmitIntrinsics::deduceElementTypeHelper(
       } else if (HandleType->getTargetExtName() == "spirv.VulkanBuffer") {
         // This call is supposed to index into an array
         Ty = HandleType->getTypeParameter(0);
-        assert(Ty->isArrayTy() &&
-               "spv_resource_getpointer indexes into an array, so the type of "
-               "the buffer should be an array.");
-        Ty = Ty->getArrayElementType();
+        if (Ty->isArrayTy())
+          Ty = Ty->getArrayElementType();
+        else {
+          TargetExtType *BufferTy = cast<TargetExtType>(Ty);
+          assert(BufferTy->getTargetExtName() == "spirv.Layout");
+          Ty = BufferTy->getTypeParameter(0);
+          assert(Ty && Ty->isStructTy());
+          uint32_t Index = cast<ConstantInt>(II->getOperand(1))->getZExtValue();
+          Ty = cast<StructType>(Ty)->getElementType(Index);
+        }
       } else {
         llvm_unreachable("Unknown handle type for spv_resource_getpointer.");
       }

llvm/lib/Target/SPIRV/SPIRVGlobalRegistry.cpp

@@ -957,7 +957,7 @@ SPIRVType *SPIRVGlobalRegistry::getOpTypeOpaque(const StructType *Ty,
 SPIRVType *SPIRVGlobalRegistry::getOpTypeStruct(
     const StructType *Ty, MachineIRBuilder &MIRBuilder,
     SPIRV::AccessQualifier::AccessQualifier AccQual,
-    bool ExplicitLayoutRequired, bool EmitIR) {
+    StructOffsetDecorator Decorator, bool EmitIR) {
   const SPIRVSubtarget &ST =
       cast<SPIRVSubtarget>(MIRBuilder.getMF().getSubtarget());
   SmallVector<Register, 4> FieldTypes;
@@ -974,7 +974,7 @@ SPIRVType *SPIRVGlobalRegistry::getOpTypeStruct(
 
   for (const auto &Elem : Ty->elements()) {
     SPIRVType *ElemTy = findSPIRVType(toTypedPointer(Elem), MIRBuilder, AccQual,
-                                      ExplicitLayoutRequired, EmitIR);
+                                      Decorator != nullptr, EmitIR);
     assert(ElemTy && ElemTy->getOpcode() != SPIRV::OpTypeVoid &&
            "Invalid struct element type");
     FieldTypes.push_back(getSPIRVTypeID(ElemTy));
@@ -1001,9 +1001,8 @@ SPIRVType *SPIRVGlobalRegistry::getOpTypeStruct(
         return MIBStruct;
       });
 
-  if (ExplicitLayoutRequired)
-    addStructOffsetDecorations(SPVType->defs().begin()->getReg(),
-                               const_cast<StructType *>(Ty), MIRBuilder);
+  if (Decorator)
+    Decorator(SPVType->defs().begin()->getReg());
 
   return SPVType;
 }
@@ -1142,8 +1141,15 @@ SPIRVType *SPIRVGlobalRegistry::createSPIRVType(
   if (auto SType = dyn_cast<StructType>(Ty)) {
     if (SType->isOpaque())
       return getOpTypeOpaque(SType, MIRBuilder);
-    return getOpTypeStruct(SType, MIRBuilder, AccQual, ExplicitLayoutRequired,
-                           EmitIR);
+
+    StructOffsetDecorator Decorator = nullptr;
+    if (ExplicitLayoutRequired) {
+      Decorator = [&MIRBuilder, SType, this](Register Reg) {
+        addStructOffsetDecorations(Reg, const_cast<StructType *>(SType),
+                                   MIRBuilder);
+      };
+    }
+    return getOpTypeStruct(SType, MIRBuilder, AccQual, Decorator, EmitIR);
   }
   if (auto FType = dyn_cast<FunctionType>(Ty)) {
     SPIRVType *RetTy = findSPIRVType(FType->getReturnType(), MIRBuilder,
@@ -1460,6 +1466,32 @@ SPIRVType *SPIRVGlobalRegistry::getOrCreateVulkanBufferType(
   return R;
 }
 
+SPIRVType *SPIRVGlobalRegistry::getOrCreateLayoutType(
+    MachineIRBuilder &MIRBuilder, const TargetExtType *T, bool EmitIr) {
+  auto Key = SPIRV::handle(T);
+  if (const MachineInstr *MI = findMI(Key, &MIRBuilder.getMF()))
+    return MI;
+
+  StructType *ST = cast<StructType>(T->getTypeParameter(0));
+  ArrayRef<uint32_t> Offsets = T->int_params().slice(1);
+  assert(ST->getNumElements() == Offsets.size());
+
+  StructOffsetDecorator Decorator = [&MIRBuilder, &Offsets](Register Reg) {
+    for (uint32_t I = 0; I < Offsets.size(); ++I) {
+      buildOpMemberDecorate(Reg, MIRBuilder, SPIRV::Decoration::Offset, I,
+                            {Offsets[I]});
+    }
+  };
+
+  // We need a new OpTypeStruct instruction because decorations will be
+  // different from a struct with an explicit layout created from a different
+  // entry point.
+  SPIRVType *SPIRVStructType = getOpTypeStruct(
+      ST, MIRBuilder, SPIRV::AccessQualifier::None, Decorator, EmitIr);
+  add(Key, SPIRVStructType);
+  return SPIRVStructType;
+}
+
 SPIRVType *SPIRVGlobalRegistry::getOrCreateOpTypeImage(
     MachineIRBuilder &MIRBuilder, SPIRVType *SampledType, SPIRV::Dim::Dim Dim,
     uint32_t Depth, uint32_t Arrayed, uint32_t Multisampled, uint32_t Sampled,

llvm/lib/Target/SPIRV/SPIRVGlobalRegistry.h

@@ -26,6 +26,7 @@
 namespace llvm {
 class SPIRVSubtarget;
 using SPIRVType = const MachineInstr;
+using StructOffsetDecorator = std::function<void(Register)>;
 
 class SPIRVGlobalRegistry : public SPIRVIRMapping {
   // Registers holding values which have types associated with them.
@@ -451,7 +452,7 @@ class SPIRVGlobalRegistry : public SPIRVIRMapping {
 
   SPIRVType *getOpTypeStruct(const StructType *Ty, MachineIRBuilder &MIRBuilder,
                              SPIRV::AccessQualifier::AccessQualifier AccQual,
-                             bool ExplicitLayoutRequired, bool EmitIR);
+                             StructOffsetDecorator Decorator, bool EmitIR);
 
   SPIRVType *getOpTypePointer(SPIRV::StorageClass::StorageClass SC,
                               SPIRVType *ElemType, MachineIRBuilder &MIRBuilder,
@@ -601,6 +602,9 @@ class SPIRVGlobalRegistry : public SPIRVIRMapping {
                                          SPIRV::StorageClass::StorageClass SC,
                                          bool IsWritable, bool EmitIr = false);
 
+  SPIRVType *getOrCreateLayoutType(MachineIRBuilder &MIRBuilder,
+                                   const TargetExtType *T, bool EmitIr = false);
+
   SPIRVType *
   getOrCreateOpTypeImage(MachineIRBuilder &MIRBuilder, SPIRVType *SampledType,
                          SPIRV::Dim::Dim Dim, uint32_t Depth, uint32_t Arrayed,

llvm/lib/Target/SPIRV/SPIRVLegalizerInfo.cpp

@@ -86,14 +86,15 @@ SPIRVLegalizerInfo::SPIRVLegalizerInfo(const SPIRVSubtarget &ST) {
   const LLT p8 = LLT::pointer(8, PSize); // Output
   const LLT p10 = LLT::pointer(10, PSize); // Private
   const LLT p11 = LLT::pointer(11, PSize); // StorageBuffer
+  const LLT p12 = LLT::pointer(12, PSize); // Uniform
 
   // TODO: remove copy-pasting here by using concatenation in some way.
   auto allPtrsScalarsAndVectors = {
-      p0,    p1,    p2,     p3,     p4,    p5,    p6,    p7,    p8,
-      p10,   p11,   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,   p11,   p12,   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,
@@ -120,10 +121,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,  p7,  p8, p10, p11};
+  auto allFloatAndIntScalarsAndPtrs = {s8, s16, s32, s64, p0, p1,  p2,  p3,
+                                       p4, p5,  p6,  p7,  p8, p10, p11, p12};
 
-  auto allPtrs = {p0, p1, p2, p3, p4, p5, p6, p7, p8, p10, p11};
+  auto allPtrs = {p0, p1, p2, p3, p4, p5, p6, p7, p8, p10, p11, p12};
 
   bool IsExtendedInts =
       ST.canUseExtension(

llvm/lib/Target/SPIRV/SPIRVUtils.cpp

@@ -262,6 +262,8 @@ addressSpaceToStorageClass(unsigned AddrSpace, const SPIRVSubtarget &STI) {
     return SPIRV::StorageClass::Private;
   case 11:
     return SPIRV::StorageClass::StorageBuffer;
+  case 12:
+    return SPIRV::StorageClass::Uniform;
   default:
     report_fatal_error("Unknown address space");
   }

llvm/lib/Target/SPIRV/SPIRVUtils.h

@@ -210,6 +210,8 @@ storageClassToAddressSpace(SPIRV::StorageClass::StorageClass SC) {
     return 10;
   case SPIRV::StorageClass::StorageBuffer:
     return 11;
+  case SPIRV::StorageClass::Uniform:
+    return 12;
   default:
     report_fatal_error("Unable to get address space id");
   }

llvm/test/CodeGen/SPIRV/hlsl-resources/spirv.layout.type.ll

@@ -0,0 +1,55 @@
+; RUN: llc -O0 -verify-machineinstrs -mtriple=spirv1.6-unknown-vulkan1.3-library %s -o - | FileCheck %s
+; RUN: %if spirv-tools %{ llc -O0 -mtriple=spirv1.6-unknown-vulkan1.3-library %s -o - -filetype=obj | spirv-val %}
+
+target datalayout = "e-i64:64-v16:16-v24:32-v32:32-v48:64-v96:128-v192:256-v256:256-v512:512-v1024:1024-n8:16:32:64-G10"
+
+; CHECK-DAG: OpName [[standard_layout:%[0-9]+]] "standard_layout"
+; CHECK-DAG: OpMemberDecorate [[standard_layout]] 0 Offset 0
+; CHECK-DAG: OpMemberDecorate [[standard_layout]] 1 Offset 4
+
+; CHECK-DAG: OpName [[standard_layout_with_different_offset:%[0-9]+]] "standard_layout"
+; CHECK-DAG: OpMemberDecorate [[standard_layout_with_different_offset]] 0 Offset 0
+; CHECK-DAG: OpMemberDecorate [[standard_layout_with_different_offset]] 1 Offset 8
+%standard_layout = type { i32, i32 }
+
+; CHECK-DAG: OpName [[backwards_layout:%[0-9]+]] "backwards_layout"
+; CHECK-DAG: OpMemberDecorate [[backwards_layout]] 0 Offset 4
+; CHECK-DAG: OpMemberDecorate [[backwards_layout]] 1 Offset 0
+%backwards_layout = type { i32, i32 }
+
+; CHECK-DAG: OpName [[large_gap:%[0-9]+]] "large_gap"
+; CHECK-DAG: OpMemberDecorate [[large_gap]] 0 Offset 0
+; CHECK-DAG: OpMemberDecorate [[large_gap]] 1 Offset 64
+; CHECK-DAG: OpMemberDecorate [[large_gap]] 2 Offset 1020
+; CHECK-DAG: OpMemberDecorate [[large_gap]] 3 Offset 4
+%large_gap = type { i32, i32, i32, i32 }
+
+; CHECK-DAG: OpName [[mixed_layout:%[0-9]+]] "mixed_layout"
+; CHECK-DAG: OpMemberDecorate [[mixed_layout]] 0 Offset 0
+; CHECK-DAG: OpMemberDecorate [[mixed_layout]] 1 Offset 8
+; CHECK-DAG: OpMemberDecorate [[mixed_layout]] 2 Offset 4
+; CHECK-DAG: OpMemberDecorate [[mixed_layout]] 3 Offset 12
+%mixed_layout = type { i32, i32, i32, i32 }
+
+define void @main() local_unnamed_addr #1 {
+entry:
+  %standard_handle = tail call target("spirv.VulkanBuffer", target("spirv.Layout", %standard_layout, 8, 0, 4), 2, 0) @llvm.spv.resource.handlefrombinding.tspirv.VulkanBuffer_tspirv.Layout_s___cblayout_Bs_8_0_4t_2_0t(i32 0, i32 1, i32 1, i32 0, i1 false)
+  %standard_handle_with_different_offset = tail call target("spirv.VulkanBuffer", target("spirv.Layout", %standard_layout, 12, 0, 8), 2, 0) @llvm.spv.resource.handlefrombinding.tspirv.VulkanBuffer_tspirv.Layout_s___cblayout_Bs_8_0_4t_2_0t(i32 0, i32 1, i32 1, i32 0, i1 false)
+  %backwards_handle = tail call target("spirv.VulkanBuffer", target("spirv.Layout", %backwards_layout, 8, 4, 0), 2, 0) @llvm.spv.resource.handlefrombinding.tspirv.VulkanBuffer_tspirv.Layout_s___cblayout_Bs_8_0_4t_2_0t(i32 0, i32 1, i32 1, i32 0, i1 false)
+  %large_gap_handle = tail call target("spirv.VulkanBuffer", target("spirv.Layout", %large_gap, 1024, 0, 64, 1020, 4), 2, 0) @llvm.spv.resource.handlefrombinding.tspirv.VulkanBuffer_tspirv.Layout_s___cblayout_Bs_8_0_4t_2_0t(i32 0, i32 1, i32 1, i32 0, i1 false)
+  %mixed_handle = tail call target("spirv.VulkanBuffer", target("spirv.Layout", %mixed_layout, 16, 0, 8, 4, 12), 2, 0) @llvm.spv.resource.handlefrombinding.tspirv.VulkanBuffer_tspirv.Layout_s___cblayout_Bs_8_0_4t_2_0t(i32 0, i32 1, i32 1, i32 0, i1 false)
+  ret void
+}
+
+attributes #1 = { mustprogress nofree noinline norecurse nosync nounwind willreturn memory(readwrite, argmem: write, inaccessiblemem: none) "approx-func-fp-math"="false" "frame-pointer"="all" "hlsl.numthreads"="1,1,1" "hlsl.shader"="compute" "no-infs-fp-math"="false" "no-nans-fp-math"="false" "no-signed-zeros-fp-math"="false" "no-trapping-math"="true" "stack-protector-buffer-size"="8" }
+
+!llvm.module.flags = !{!0, !1}
+!llvm.ident = !{!2}
+
+!0 = !{i32 1, !"wchar_size", i32 4}
+!1 = !{i32 7, !"frame-pointer", i32 2}
+!2 = !{!"clang version 21.0.0git ([email protected]:s-perron/llvm-project.git b02f2e80567af09576692554bc7ce048326dfd06)"}
+!3 = !{!4, !4, i64 0}
+!4 = !{!"int", !5, i64 0}
+!5 = !{!"omnipotent char", !6, i64 0}
+!6 = !{!"Simple C++ TBAA"}