[spirv] Added support for [[vk::shader_record_ext]] (#3426)

I added support for the Shader Record buffer for the Khronos extension.
The code is adapted from [this](https://github.com/microsoft/DirectXShaderCompiler/pull/2179) previous merge.

It should fix this issue #3177

docs/SPIR-V.rst

@@ -247,6 +247,10 @@ Please note as per the requirements of VK_NV_ray_tracing, "there must be no
 more than one shader_record_nv block statically used per shader entry point
 otherwise results are undefined."
 
+The official Khronos ray tracing extension also comes with a SPIR-V storage class
+that has the same functionality. The ``[[vk::shader_record_ext]]`` annotation can 
+be used when targeting the SPV_KHR_ray_tracing extension.
+
 Builtin variables
 ~~~~~~~~~~~~~~~~~
 

tools/clang/include/clang/Basic/Attr.td

@@ -1037,6 +1037,15 @@ def VKShaderRecordNV : InheritableAttr {
   let LangOpts = [SPIRV];
   let Documentation = [Undocumented];
 }
+
+def VKShaderRecordEXT : InheritableAttr {
+  let Spellings = [CXX11<"vk", "shader_record_ext">];
+  let Subjects = SubjectList<[StructGlobalVar, HLSLBuffer, ConstantBuffer],
+                             ErrorDiag, "ExpectedCTBuffer">;
+  let Args = [];
+  let LangOpts = [SPIRV];
+  let Documentation = [Undocumented];
+}
 // SPIRV Change Ends
 
 def C11NoReturn : InheritableAttr {

tools/clang/lib/SPIRV/DeclResultIdMapper.cpp

@@ -912,6 +912,8 @@ SpirvVariable *DeclResultIdMapper::createStructOrStructArrayVarOfExplicitLayout(
   const bool forPC = usageKind == ContextUsageKind::PushConstant;
   const bool forShaderRecordNV =
       usageKind == ContextUsageKind::ShaderRecordBufferNV;
+  const bool forShaderRecordEXT =
+      usageKind == ContextUsageKind::ShaderRecordBufferEXT;
 
   const auto &declGroup = collectDeclsInDeclContext(decl);
 
@@ -960,7 +962,9 @@ SpirvVariable *DeclResultIdMapper::createStructOrStructArrayVarOfExplicitLayout(
   const auto sc = forPC ? spv::StorageClass::PushConstant
                         : forShaderRecordNV
                               ? spv::StorageClass::ShaderRecordBufferNV
-                              : spv::StorageClass::Uniform;
+                              : forShaderRecordEXT
+                                    ? spv::StorageClass::ShaderRecordBufferKHR
+                                    : spv::StorageClass::Uniform;
 
   // Create the variable for the whole struct / struct array.
   // The fields may be 'precise', but the structure itself is not.
@@ -1117,16 +1121,24 @@ SpirvVariable *DeclResultIdMapper::createPushConstant(const VarDecl *decl) {
 }
 
 SpirvVariable *
-DeclResultIdMapper::createShaderRecordBufferNV(const VarDecl *decl) {
+DeclResultIdMapper::createShaderRecordBuffer(const VarDecl *decl, 
+                                                ContextUsageKind kind) {
   const auto *recordType =
       hlsl::GetHLSLResourceResultType(decl->getType())->getAs<RecordType>();
   assert(recordType);
 
+  assert(kind == ContextUsageKind::ShaderRecordBufferEXT ||
+         kind == ContextUsageKind::ShaderRecordBufferNV);
+
+  const auto typeName = kind == ContextUsageKind::ShaderRecordBufferEXT
+                            ? "type.ShaderRecordBufferEXT."
+                            : "type.ShaderRecordBufferNV.";
+
   const std::string structName =
-      "type.ShaderRecordBufferNV." + recordType->getDecl()->getName().str();
+      typeName + recordType->getDecl()->getName().str();
   SpirvVariable *var = createStructOrStructArrayVarOfExplicitLayout(
       recordType->getDecl(), /*arraySize*/ 0,
-      ContextUsageKind::ShaderRecordBufferNV, structName, decl->getName());
+      kind, structName, decl->getName());
 
   // Register the VarDecl
   astDecls[decl] = DeclSpirvInfo(var);
@@ -1138,13 +1150,20 @@ DeclResultIdMapper::createShaderRecordBufferNV(const VarDecl *decl) {
 }
 
 SpirvVariable *
-DeclResultIdMapper::createShaderRecordBufferNV(const HLSLBufferDecl *decl) {
+DeclResultIdMapper::createShaderRecordBuffer(const HLSLBufferDecl *decl,
+                                               ContextUsageKind kind) {
+  assert(kind == ContextUsageKind::ShaderRecordBufferEXT ||
+         kind == ContextUsageKind::ShaderRecordBufferNV);
+
+  const auto typeName = kind == ContextUsageKind::ShaderRecordBufferEXT
+                            ? "type.ShaderRecordBufferEXT."
+                            : "type.ShaderRecordBufferNV.";
 
   const std::string structName =
-      "type.ShaderRecordBufferNV." + decl->getName().str();
+      typeName + decl->getName().str();
   // The front-end does not allow arrays of cbuffer/tbuffer.
   SpirvVariable *bufferVar = createStructOrStructArrayVarOfExplicitLayout(
-      decl, /*arraySize*/ 0, ContextUsageKind::ShaderRecordBufferNV, structName,
+      decl, /*arraySize*/ 0, kind, structName,
       decl->getName());
 
   // We still register all VarDecls seperately here. All the VarDecls are

tools/clang/lib/SPIRV/DeclResultIdMapper.h

@@ -416,10 +416,20 @@ public:
                                               const QualType retType,
                                               uint32_t numOutputControlPoints);
 
+  /// \brief An enum class for representing what the DeclContext is used for
+  enum class ContextUsageKind {
+    CBuffer,
+    TBuffer,
+    PushConstant,
+    Globals,
+    ShaderRecordBufferNV,
+    ShaderRecordBufferEXT
+  };
+
   /// Raytracing specific functions
-  /// \brief Creates a ShaderRecordBufferNV block from the given decl.
-  SpirvVariable *createShaderRecordBufferNV(const VarDecl *decl);
-  SpirvVariable *createShaderRecordBufferNV(const HLSLBufferDecl *decl);
+  /// \brief Creates a ShaderRecordBufferEXT or ShaderRecordBufferNV block from the given decl.
+  SpirvVariable *createShaderRecordBuffer(const VarDecl *decl, ContextUsageKind kind);
+  SpirvVariable *createShaderRecordBuffer(const HLSLBufferDecl *decl, ContextUsageKind kind);
 
 private:
   /// The struct containing SPIR-V information of a AST Decl.
@@ -600,15 +610,6 @@ private:
   /// construction.
   bool finalizeStageIOLocations(bool forInput);
 
-  /// \brief An enum class for representing what the DeclContext is used for
-  enum class ContextUsageKind {
-    CBuffer,
-    TBuffer,
-    PushConstant,
-    Globals,
-    ShaderRecordBufferNV,
-  };
-
   /// Creates a variable of struct type with explicit layout decorations.
   /// The sub-Decls in the given DeclContext will be treated as the struct
   /// fields. The struct type will be named as typeName, and the variable

tools/clang/lib/SPIRV/SpirvEmitter.cpp

@@ -1290,6 +1290,32 @@ bool SpirvEmitter::validateVKAttributes(const NamedDecl *decl) {
     }
   }
 
+  // vk::shader_record_ext is supported only on cbuffer/ConstantBuffer
+  if (const auto *srbAttr = decl->getAttr<VKShaderRecordEXTAttr>()) {
+    const auto loc = srbAttr->getLocation();
+    const HLSLBufferDecl *bufDecl = nullptr;
+    bool isValidType = false;
+    if ((bufDecl = dyn_cast<HLSLBufferDecl>(decl)))
+      isValidType = bufDecl->isCBuffer();
+    else if ((bufDecl = dyn_cast<HLSLBufferDecl>(decl->getDeclContext())))
+      isValidType = bufDecl->isCBuffer();
+    else if (isa<VarDecl>(decl))
+      isValidType = isConstantBuffer(dyn_cast<VarDecl>(decl)->getType());
+
+    if (!isValidType) {
+      emitError(
+          "vk::shader_record_ext can be applied only to cbuffer/ConstantBuffer",
+          loc);
+      success = false;
+    }
+    if (decl->hasAttr<VKBindingAttr>()) {
+      emitError("vk::shader_record_ext attribute cannot be used together with "
+                "vk::binding attribute",
+                loc);
+      success = false;
+    }
+  }
+
   return success;
 }
 
@@ -1319,7 +1345,12 @@ void SpirvEmitter::doHLSLBufferDecl(const HLSLBufferDecl *bufferDecl) {
   if (!validateVKAttributes(bufferDecl))
     return;
   if (bufferDecl->hasAttr<VKShaderRecordNVAttr>()) {
-    (void)declIdMapper.createShaderRecordBufferNV(bufferDecl);
+    (void)declIdMapper.createShaderRecordBuffer(
+        bufferDecl, DeclResultIdMapper::ContextUsageKind::ShaderRecordBufferNV);
+  } else if (bufferDecl->hasAttr<VKShaderRecordEXTAttr>()) {
+    (void)declIdMapper.createShaderRecordBuffer(
+        bufferDecl,
+        DeclResultIdMapper::ContextUsageKind::ShaderRecordBufferEXT);
   } else {
     (void)declIdMapper.createCTBuffer(bufferDecl);
   }
@@ -1401,7 +1432,14 @@ void SpirvEmitter::doVarDecl(const VarDecl *decl) {
   }
 
   if (decl->hasAttr<VKShaderRecordNVAttr>()) {
-    (void)declIdMapper.createShaderRecordBufferNV(decl);
+    (void)declIdMapper.createShaderRecordBuffer(
+        decl, DeclResultIdMapper::ContextUsageKind::ShaderRecordBufferNV);
+    return;
+  }
+
+  if (decl->hasAttr<VKShaderRecordEXTAttr>()) {
+    (void)declIdMapper.createShaderRecordBuffer(
+        decl, DeclResultIdMapper::ContextUsageKind::ShaderRecordBufferEXT);
     return;
   }
 

tools/clang/lib/Sema/SemaHLSL.cpp

@@ -11698,6 +11698,9 @@ void hlsl::HandleDeclAttributeForHLSL(Sema &S, Decl *D, const AttributeList &A,
   case AttributeList::AT_VKShaderRecordNV:
     declAttr = ::new (S.Context) VKShaderRecordNVAttr(A.getRange(), S.Context, A.getAttributeSpellingListIndex());
     break;
+  case AttributeList::AT_VKShaderRecordEXT:
+    declAttr = ::new (S.Context) VKShaderRecordEXTAttr(A.getRange(), S.Context, A.getAttributeSpellingListIndex());
+    break;
   default:
     Handled = false;
     return;
@@ -13088,6 +13091,7 @@ bool hlsl::IsHLSLAttr(clang::attr::Kind AttrKind) {
   case clang::attr::VKOffset:
   case clang::attr::VKPushConstant:
   case clang::attr::VKShaderRecordNV:
+  case clang::attr::VKShaderRecordEXT:
     return true;
   default:
     // Only HLSL/VK Attributes return true. Only used for printPretty(), which doesn't support them.

tools/clang/test/CodeGenSPIRV/vk.attribute.error.hlsl

@@ -42,19 +42,33 @@ float foo([[vk::push_constant]] int param) // error
 [[vk::push_constant(5)]]
 T pcs;
 
-struct SRB {
+struct SRB_NV {
     [[vk::shader_record_nv]] // error
     float4 f;
 };
 
 [[vk::shader_record_nv]] // error
-float foosrb([[vk::shader_record_nv]] int param) // error
+float foosrbnv([[vk::shader_record_nv]] int param) // error
 {
     return param;
 }
 
 [[vk::shader_record_nv(5)]]
-SRB recordBuf;
+SRB_NV recordBufNV;
+
+struct SRB_EXT {
+    [[vk::shader_record_ext]] // error
+    float4 f;
+};
+
+[[vk::shader_record_ext]] // error
+float foosrbext([[vk::shader_record_ext]] int param)
+{
+    return param;
+}
+
+[[vk::shader_record_ext(5)]]
+SRB_EXT recordBufEXT;
 
 // CHECK:   :4:7: error: 'binding' attribute only applies to global variables, cbuffers, and tbuffers
 // CHECK:   :6:7: error: 'counter_binding' attribute only applies to RWStructuredBuffers, AppendStructuredBuffers, and ConsumeStructuredBuffers
@@ -70,6 +84,10 @@ SRB recordBuf;
 // CHECK:  :36:3: error: 'push_constant' attribute only applies to global variables of struct type
 // CHECK:  :42:3: error: 'push_constant' attribute takes no arguments
 // CHECK:  :46:7: error: 'shader_record_nv' attribute only applies to cbuffer or ConstantBuffer
-// CHECK: :51:16: error: 'shader_record_nv' attribute only applies to cbuffer or ConstantBuffer
+// CHECK: :51:18: error: 'shader_record_nv' attribute only applies to cbuffer or ConstantBuffer
 // CHECK:  :50:3: error: 'shader_record_nv' attribute only applies to cbuffer or ConstantBuffer
-// CHECK:  :56:3: error: 'shader_record_nv' attribute takes no arguments
+// CHECK:  :56:3: error: 'shader_record_nv' attribute takes no arguments
+// CHECK:  :60:7: error: 'shader_record_ext' attribute only applies to cbuffer or ConstantBuffer
+// CHECK: :65:19: error: 'shader_record_ext' attribute only applies to cbuffer or ConstantBuffer
+// CHECK:  :64:3: error: 'shader_record_ext' attribute only applies to cbuffer or ConstantBuffer
+// CHECK:  :70:3: error: 'shader_record_ext' attribute takes no arguments

tools/clang/test/CodeGenSPIRV/vk.attribute.shader-record-ext.invalid.hlsl

@@ -0,0 +1,14 @@
+// Run: %dxc -T vs_6_0 -E main
+
+struct S {
+    float f;
+};
+
+[[vk::shader_record_ext, vk::binding(6)]]
+ConstantBuffer<S> recordBuf;
+
+float main() : A {
+    return 1.0;
+}
+
+// CHECK: :7:3: error: vk::shader_record_ext attribute cannot be used together with vk::binding attribute

tools/clang/test/CodeGenSPIRV/vk.layout.shader-record-ext.std430.hlsl

@@ -0,0 +1,78 @@
+// Run: %dxc -T lib_6_3 -fspv-target-env=vulkan1.2
+
+// CHECK: OpDecorate %_arr_v2float_uint_3 ArrayStride 8
+// CHECK: OpDecorate %_arr_mat3v2float_uint_2 ArrayStride 32
+// CHECK: OpDecorate %_arr_v2int_uint_3 ArrayStride 8
+// CHECK: OpDecorate %_arr__arr_v2int_uint_3_uint_2 ArrayStride 24
+// CHECK: OpDecorate %_arr_mat3v2float_uint_2_0 ArrayStride 24
+// CHECK-NOT: OpDecorate %cbuf DescriptorSet
+// CHECK-NOT: OpDecorate %cbuf Binding
+// CHECK-NOT: OpDecorate %block DescriptorSet
+// CHECK-NOT: OpDecorate %block Binding
+
+// CHECK: OpMemberDecorate %T 0 Offset 0
+// CHECK: OpMemberDecorate %T 1 Offset 32
+// CHECK: OpMemberDecorate %T 1 MatrixStride 16
+// CHECK: OpMemberDecorate %T 1 RowMajor
+// CHECK: OpMemberDecorate %T 2 Offset 96
+// CHECK: OpMemberDecorate %T 3 Offset 144
+// CHECK: OpMemberDecorate %T 3 MatrixStride 8
+// CHECK: OpMemberDecorate %T 3 ColMajor
+struct T {
+                 float2   f1[3];
+    column_major float3x2 f2[2];
+    row_major    int3x2   f4[2];
+    row_major    float3x2 f3[2];
+};
+
+// CHECK: OpMemberDecorate %type_ShaderRecordBufferEXT_S 0 Offset 0
+// CHECK: OpMemberDecorate %type_ShaderRecordBufferEXT_S 1 Offset 4
+// CHECK: OpMemberDecorate %type_ShaderRecordBufferEXT_S 2 Offset 16
+// CHECK: OpMemberDecorate %type_ShaderRecordBufferEXT_S 3 Offset 208
+// CHECK: OpMemberDecorate %type_ShaderRecordBufferEXT_S 4 Offset 240
+// CHECK: OpMemberDecorate %type_ShaderRecordBufferEXT_S 4 MatrixStride 16
+// CHECK: OpMemberDecorate %type_ShaderRecordBufferEXT_S 4 ColMajor
+
+struct S {
+              float    f1;
+              float3   f2;
+              T        f4;
+    row_major int2x3   f5;
+    row_major float2x3 f3;
+};
+
+[[vk::shader_record_ext]]
+ConstantBuffer<S> cbuf;
+
+// CHECK: OpDecorate %type_ShaderRecordBufferEXT_S Block
+// CHECK: OpMemberDecorate %type_ShaderRecordBufferEXT_block 0 Offset 0
+// CHECK: OpMemberDecorate %type_ShaderRecordBufferEXT_block 1 Offset 4
+// CHECK: OpMemberDecorate %type_ShaderRecordBufferEXT_block 2 Offset 16
+// CHECK: OpMemberDecorate %type_ShaderRecordBufferEXT_block 3 Offset 208
+// CHECK: OpMemberDecorate %type_ShaderRecordBufferEXT_block 4 Offset 240
+// CHECK: OpMemberDecorate %type_ShaderRecordBufferEXT_block 4 MatrixStride 16
+// CHECK: OpMemberDecorate %type_ShaderRecordBufferEXT_block 4 ColMajor
+
+
+[[vk::shader_record_ext]]
+cbuffer block {
+              float    f1;
+              float3   f2;
+              T        f4;
+    row_major int2x3   f5;
+    row_major float2x3 f3;
+}
+
+// CHECK: OpDecorate %type_ShaderRecordBufferEXT_block Block
+struct Payload { float p; };
+struct Attr    { float a; };
+
+[shader("closesthit")]
+void chs1(inout Payload P, in Attr A) {
+    P.p = cbuf.f1;
+}
+
+[shader("closesthit")]
+void chs2(inout Payload P, in Attr A) {
+    P.p = f1;
+}

tools/clang/test/CodeGenSPIRV/vk.shader-record-ext.hlsl

@@ -0,0 +1,45 @@
+// Run: %dxc -T lib_6_3 -fspv-target-env=vulkan1.2
+
+struct T {
+    float2 val[3];
+};
+
+// CHECK: OpName %type_ShaderRecordBufferEXT_S "type.ShaderRecordBufferEXT.S"
+// CHECK: OpMemberName %type_ShaderRecordBufferEXT_S 0 "f1"
+// CHECK: OpMemberName %type_ShaderRecordBufferEXT_S 1 "f2"
+// CHECK: OpMemberName %type_ShaderRecordBufferEXT_S 2 "f3"
+// CHECK: OpMemberName %type_ShaderRecordBufferEXT_S 3 "f4"
+// CHECK-NOT: OpDecorate %srb DescriptorSet
+// CHECK-NOT: OpDecorate %srb Binding
+
+// CHECK: %type_ShaderRecordBufferEXT_S = OpTypeStruct %float %v3float %mat2v3float %T
+struct S {
+    float    f1;
+    float3   f2;
+    float2x3 f3;
+    T        f4;
+};
+// CHECK: %_ptr_ShaderRecordBufferNV_type_ShaderRecordBufferEXT_S = OpTypePointer ShaderRecordBufferNV %type_ShaderRecordBufferEXT_S
+
+// CHECK: %srb = OpVariable %_ptr_ShaderRecordBufferNV_type_ShaderRecordBufferEXT_S ShaderRecordBufferNV
+[[vk::shader_record_ext]]
+ConstantBuffer<S> srb;
+
+struct Payload { float p; };
+struct Attribute { float a; };
+
+[shader("miss")]
+void main(inout Payload P) 
+{
+   P.p = 
+// CHECK:     {{%\d+}} = OpAccessChain %_ptr_ShaderRecordBufferNV_float %srb %int_0
+        srb.f1 +
+// CHECK: [[ptr:%\d+]] = OpAccessChain %_ptr_ShaderRecordBufferNV_v3float %srb %int_1
+// CHECK:     {{%\d+}} = OpAccessChain %_ptr_ShaderRecordBufferNV_float [[ptr]] %int_2
+        srb.f2.z +
+// CHECK:     {{%\d+}} = OpAccessChain %_ptr_ShaderRecordBufferNV_float %srb %int_2 %uint_1 %uint_2
+        srb.f3[1][2] +
+// CHECK: [[ptr:%\d+]] = OpAccessChain %_ptr_ShaderRecordBufferNV_v2float %srb %int_3 %int_0 %int_2
+// CHECK:     {{%\d+}} = OpAccessChain %_ptr_ShaderRecordBufferNV_float [[ptr]] %int_1
+        srb.f4.val[2].y;
+}

tools/clang/test/CodeGenSPIRV/vk.shader-record-ext.offset.hlsl

@@ -0,0 +1,23 @@
+// Run: %dxc -T lib_6_3 -fspv-target-env=vulkan1.2
+
+// CHECK: OpMemberDecorate %type_ShaderRecordBufferEXT_S 0 Offset 0
+// CHECK: OpMemberDecorate %type_ShaderRecordBufferEXT_S 1 Offset 8
+// CHECK: OpMemberDecorate %type_ShaderRecordBufferEXT_S 2 Offset 32
+
+struct S {
+    float a;
+    [[vk::offset(8)]]
+    float2 b;
+    [[vk::offset(32)]]
+    float4 f;
+};
+
+[[vk::shader_record_ext]]
+ConstantBuffer<S> srb;
+
+struct Payload { float p; };
+struct Attr    { float a; };
+[shader("closesthit")]
+void main(inout Payload P, in Attr a) {
+    P.p = srb.a;
+}

tools/clang/unittests/SPIRV/CodeGenSpirvTest.cpp

@@ -1695,6 +1695,9 @@ TEST_F(FileTest, VulkanAttributePushConstantInvalidUsages) {
 TEST_F(FileTest, VulkanAttributeShaderRecordNVInvalidUsages) {
   runFileTest("vk.attribute.shader-record-nv.invalid.hlsl", Expect::Failure);
 }
+TEST_F(FileTest, VulkanAttributeShaderRecordEXTInvalidUsages) {
+  runFileTest("vk.attribute.shader-record-ext.invalid.hlsl", Expect::Failure);
+}
 
 TEST_F(FileTest, VulkanCLOptionInvertYVS) {
   runFileTest("vk.cloption.invert-y.vs.hlsl");
@@ -2335,6 +2338,18 @@ TEST_F(FileTest, VulkanShaderRecordBufferNVOffset) {
   // Checks the behavior of [[vk::offset]] with [[vk::shader_record_nv]]
   runFileTest("vk.shader-record-nv.offset.hlsl");
 }
+// Tests for [[vk::shader_record_ext]]
+TEST_F(FileTest, VulkanShaderRecordBufferEXT) {
+  runFileTest("vk.shader-record-ext.hlsl");
+}
+TEST_F(FileTest, VulkanLayoutShaderRecordBufferEXTStd430) {
+  setGlLayout();
+  runFileTest("vk.layout.shader-record-ext.std430.hlsl");
+}
+TEST_F(FileTest, VulkanShaderRecordBufferEXTOffset) {
+  // Checks the behavior of [[vk::offset]] with [[vk::shader_record_ext]]
+  runFileTest("vk.shader-record-ext.offset.hlsl");
+}
 TEST_F(FileTest, VulkanShadingRate) { runFileTest("vk.shading-rate.hlsl"); }
 TEST_F(FileTest, VulkanShadingRateError) {
   runFileTest("vk.shading-rate.vs-error.hlsl", Expect::Failure);