[spirv] Add support for [[vk::shader_record_nv]] (#2179)

* Add support for [[vk::shader_record_nv]]

* Fixes after first review of ehsannas.

docs/SPIR-V.rst

@@ -234,6 +234,19 @@ To use Vulkan specialization constants, annotate global constants with the
   [[vk::constant_id(2)]] const int   specConstInt   = 42;
   [[vk::constant_id(3)]] const float specConstFloat = 1.5;
 
+Shader Record Buffer
+~~~~~~~~~~~~~~~~~~~~
+
+SPV_NV_ray_tracing exposes user managed buffer in shader binding table by
+using storage class ShaderRecordBufferNV. ConstantBuffer or cbuffer blocks
+can now be mapped to this storage class under HLSL by using
+``[[vk::shader_record_nv]]`` annotation. It is applicable only on ConstantBuffer
+and cbuffer declarations.
+
+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."
+
 Builtin variables
 ~~~~~~~~~~~~~~~~~
 

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

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

tools/clang/include/clang/Basic/DiagnosticSemaKinds.td

@@ -2332,6 +2332,7 @@ def warn_attribute_wrong_decl_type : Warning<
   "global variables, cbuffers, and tbuffers|"
   "RWStructuredBuffers, AppendStructuredBuffers, and ConsumeStructuredBuffers|"
   "SubpassInput, SubpassInputMS|"
+  "cbuffer or ConstantBuffer|"
   // SPIRV Change Ends
   // HLSL Change Starts - add 3 more enum values
   "variables and parameters|functions, parameters, and fields|"

tools/clang/include/clang/Sema/AttributeList.h

@@ -862,6 +862,7 @@ enum AttributeDeclKind {
   ,ExpectedGlobalVarOrCTBuffer
   ,ExpectedCounterStructuredBuffer
   ,ExpectedSubpassInput
+  ,ExpectedCTBuffer
   // SPIRV Change Ends
   // HLSL Change Begins - add attribute decl combinations
   ,ExpectedVariableOrParam,

tools/clang/lib/SPIRV/DeclResultIdMapper.cpp

@@ -697,6 +697,8 @@ SpirvVariable *DeclResultIdMapper::createStructOrStructArrayVarOfExplicitLayout(
   const bool forTBuffer = usageKind == ContextUsageKind::TBuffer;
   const bool forGlobals = usageKind == ContextUsageKind::Globals;
   const bool forPC = usageKind == ContextUsageKind::PushConstant;
+  const bool forShaderRecordNV =
+      usageKind == ContextUsageKind::ShaderRecordBufferNV;
 
   const llvm::SmallVector<const Decl *, 4> &declGroup =
       collectDeclsInDeclContext(decl);
@@ -743,8 +745,10 @@ SpirvVariable *DeclResultIdMapper::createStructOrStructArrayVarOfExplicitLayout(
   // Register the <type-id> for this decl
   ctBufferPCTypes[decl] = resultType;
 
-  const auto sc =
-      forPC ? spv::StorageClass::PushConstant : spv::StorageClass::Uniform;
+  const auto sc = forPC ? spv::StorageClass::PushConstant
+                        : forShaderRecordNV
+                              ? spv::StorageClass::ShaderRecordBufferNV
+                              : spv::StorageClass::Uniform;
 
   // Create the variable for the whole struct / struct array.
   // The fields may be 'precise', but the structure itself is not.
@@ -852,6 +856,51 @@ SpirvVariable *DeclResultIdMapper::createPushConstant(const VarDecl *decl) {
   return var;
 }
 
+SpirvVariable *
+DeclResultIdMapper::createShaderRecordBufferNV(const VarDecl *decl) {
+  const auto *recordType = decl->getType()->getAs<RecordType>();
+  assert(recordType);
+
+  const std::string structName =
+      "type.ShaderRecordBufferNV." + recordType->getDecl()->getName().str();
+  SpirvVariable *var = createStructOrStructArrayVarOfExplicitLayout(
+      recordType->getDecl(), /*arraySize*/ 0,
+      ContextUsageKind::ShaderRecordBufferNV, structName, decl->getName());
+
+  // Register the VarDecl
+  astDecls[decl] = DeclSpirvInfo(var);
+
+  // Do not push this variable into resourceVars since it does not need
+  // descriptor set.
+
+  return var;
+}
+
+SpirvVariable *
+DeclResultIdMapper::createShaderRecordBufferNV(const HLSLBufferDecl *decl) {
+
+  const std::string structName =
+      "type.ShaderRecordBufferNV." + decl->getName().str();
+  // The front-end does not allow arrays of cbuffer/tbuffer.
+  SpirvVariable *bufferVar = createStructOrStructArrayVarOfExplicitLayout(
+      decl, /*arraySize*/ 0, ContextUsageKind::ShaderRecordBufferNV, structName,
+      decl->getName());
+
+  // We still register all VarDecls seperately here. All the VarDecls are
+  // mapped to the <result-id> of the buffer object, which means when
+  // querying the <result-id> for a certain VarDecl, we need to do an extra
+  // OpAccessChain.
+  int index = 0;
+  for (const auto *subDecl : decl->decls()) {
+    if (shouldSkipInStructLayout(subDecl))
+      continue;
+
+    const auto *varDecl = cast<VarDecl>(subDecl);
+    astDecls[varDecl] = DeclSpirvInfo(bufferVar, index++);
+  }
+  return bufferVar;
+}
+
 void DeclResultIdMapper::createGlobalsCBuffer(const VarDecl *var) {
   if (astDecls.count(var) != 0)
     return;

tools/clang/lib/SPIRV/DeclResultIdMapper.h

@@ -370,6 +370,10 @@ public:
   /// stages.
   void createRayTracingNVImplicitVar(const VarDecl *varDecl);
 
+  /// \brief Creates a ShaderRecordBufferNV block from the given decl.
+  SpirvVariable *createShaderRecordBufferNV(const VarDecl *decl);
+  SpirvVariable *createShaderRecordBufferNV(const HLSLBufferDecl *decl);
+
 private:
   /// The struct containing SPIR-V information of a AST Decl.
   struct DeclSpirvInfo {
@@ -531,6 +535,7 @@ private:
     TBuffer,
     PushConstant,
     Globals,
+    ShaderRecordBufferNV,
   };
 
   /// Creates a variable of struct type with explicit layout decorations.

tools/clang/lib/SPIRV/SpirvEmitter.cpp

@@ -1151,6 +1151,30 @@ bool SpirvEmitter::validateVKAttributes(const NamedDecl *decl) {
     }
   }
 
+  // vk::shader_record_nv is supported only on cbuffer/ConstantBuffer
+  if (const auto *srbAttr = decl->getAttr<VKShaderRecordNVAttr>()) {
+    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();
+
+    if (!isValidType) {
+      emitError(
+          "vk::shader_record_nv can be applied only to cbuffer/ConstantBuffer",
+          loc);
+      success = false;
+    }
+    if (decl->hasAttr<VKBindingAttr>()) {
+      emitError("vk::shader_record_nv attribute cannot be used together with "
+                "vk::binding attribute",
+                loc);
+      success = false;
+    }
+  }
+
   return success;
 }
 
@@ -1179,7 +1203,11 @@ void SpirvEmitter::doHLSLBufferDecl(const HLSLBufferDecl *bufferDecl) {
   }
   if (!validateVKAttributes(bufferDecl))
     return;
-  (void)declIdMapper.createCTBuffer(bufferDecl);
+  if (bufferDecl->hasAttr<VKShaderRecordNVAttr>()) {
+    (void)declIdMapper.createShaderRecordBufferNV(bufferDecl);
+  } else {
+    (void)declIdMapper.createCTBuffer(bufferDecl);
+  }
 }
 
 void SpirvEmitter::doImplicitDecl(const Decl *decl) {
@@ -1253,6 +1281,11 @@ void SpirvEmitter::doVarDecl(const VarDecl *decl) {
     return;
   }
 
+  if (decl->hasAttr<VKShaderRecordNVAttr>()) {
+    (void)declIdMapper.createShaderRecordBufferNV(decl);
+    return;
+  }
+
   if (isa<HLSLBufferDecl>(decl->getDeclContext())) {
     // This is a VarDecl of a ConstantBuffer/TextureBuffer type.
     (void)declIdMapper.createCTBuffer(decl);

tools/clang/lib/Sema/SemaHLSL.cpp

@@ -11057,6 +11057,9 @@ void hlsl::HandleDeclAttributeForHLSL(Sema &S, Decl *D, const AttributeList &A,
   case AttributeList::AT_VKPostDepthCoverage:
     declAttr = ::new (S.Context) VKPostDepthCoverageAttr(A.getRange(), S.Context, A.getAttributeSpellingListIndex());
     break;
+  case AttributeList::AT_VKShaderRecordNV:
+    declAttr = ::new (S.Context) VKShaderRecordNVAttr(A.getRange(), S.Context, A.getAttributeSpellingListIndex());
+    break;
   default:
     Handled = false;
     return;
@@ -12372,6 +12375,7 @@ bool hlsl::IsHLSLAttr(clang::attr::Kind AttrKind) {
   case clang::attr::VKLocation:
   case clang::attr::VKOffset:
   case clang::attr::VKPushConstant:
+  case clang::attr::VKShaderRecordNV:
     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,6 +42,20 @@ float foo([[vk::push_constant]] int param) // error
 [[vk::push_constant(5)]]
 T pcs;
 
+struct SRB {
+    [[vk::shader_record_nv]] // error
+    float4 f;
+};
+
+[[vk::shader_record_nv]] // error
+float foosrb([[vk::shader_record_nv]] int param) // error
+{
+    return param;
+}
+
+[[vk::shader_record_nv(5)]]
+SRB recordBuf;
+
 // 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
 // CHECK:  :10:3: error: 'counter_binding' attribute only applies to RWStructuredBuffers, AppendStructuredBuffers, and ConsumeStructuredBuffers
@@ -55,3 +69,7 @@ T pcs;
 // CHECK: :37:13: error: 'push_constant' attribute only applies to global variables of struct type
 // 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:  :50:3: error: 'shader_record_nv' attribute only applies to cbuffer or ConstantBuffer
+// CHECK:  :56:3: error: 'shader_record_nv' attribute takes no arguments

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

@@ -7,8 +7,12 @@ struct S {
 [[vk::push_constant, vk::binding(5)]]
 S pcs;
 
+[[vk::shader_record_nv, vk::binding(6)]]
+ConstantBuffer<S> recordBuf;
+
 float main() : A {
     return 1.0;
 }
 
 // CHECK: :7:3: error: vk::push_constant attribute cannot be used together with vk::binding attribute
+// CHECK: :10:3: error: vk::shader_record_nv attribute cannot be used together with vk::binding attribute

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

@@ -0,0 +1,78 @@
+// Run: %dxc -T lib_6_3 
+
+// 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_ShaderRecordBufferNV_S 0 Offset 0
+// CHECK: OpMemberDecorate %type_ShaderRecordBufferNV_S 1 Offset 4
+// CHECK: OpMemberDecorate %type_ShaderRecordBufferNV_S 2 Offset 16
+// CHECK: OpMemberDecorate %type_ShaderRecordBufferNV_S 3 Offset 208
+// CHECK: OpMemberDecorate %type_ShaderRecordBufferNV_S 4 Offset 240
+// CHECK: OpMemberDecorate %type_ShaderRecordBufferNV_S 4 MatrixStride 16
+// CHECK: OpMemberDecorate %type_ShaderRecordBufferNV_S 4 ColMajor
+
+struct S {
+              float    f1;
+              float3   f2;
+              T        f4;
+    row_major int2x3   f5;
+    row_major float2x3 f3;
+};
+
+[[vk::shader_record_nv]]
+ConstantBuffer<S> cbuf;
+
+// CHECK: OpDecorate %type_ShaderRecordBufferNV_S Block
+// CHECK: OpMemberDecorate %type_ShaderRecordBufferNV_block 0 Offset 0
+// CHECK: OpMemberDecorate %type_ShaderRecordBufferNV_block 1 Offset 4
+// CHECK: OpMemberDecorate %type_ShaderRecordBufferNV_block 2 Offset 16
+// CHECK: OpMemberDecorate %type_ShaderRecordBufferNV_block 3 Offset 208
+// CHECK: OpMemberDecorate %type_ShaderRecordBufferNV_block 4 Offset 240
+// CHECK: OpMemberDecorate %type_ShaderRecordBufferNV_block 4 MatrixStride 16
+// CHECK: OpMemberDecorate %type_ShaderRecordBufferNV_block 4 ColMajor
+
+
+[[vk::shader_record_nv]]
+cbuffer block {
+              float    f1;
+              float3   f2;
+              T        f4;
+    row_major int2x3   f5;
+    row_major float2x3 f3;
+}
+
+// CHECK: OpDecorate %type_ShaderRecordBufferNV_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-nv.hlsl

@@ -0,0 +1,45 @@
+// Run: %dxc -T lib_6_3
+
+struct T {
+    float2 val[3];
+};
+
+// CHECK: OpName %type_ShaderRecordBufferNV_S "type.ShaderRecordBufferNV.S"
+// CHECK: OpMemberName %type_ShaderRecordBufferNV_S 0 "f1"
+// CHECK: OpMemberName %type_ShaderRecordBufferNV_S 1 "f2"
+// CHECK: OpMemberName %type_ShaderRecordBufferNV_S 2 "f3"
+// CHECK: OpMemberName %type_ShaderRecordBufferNV_S 3 "f4"
+// CHECK-NOT: OpDecorate %srb DescriptorSet
+// CHECK-NOT: OpDecorate %srb Binding
+
+// CHECK: %type_ShaderRecordBufferNV_S = OpTypeStruct %float %v3float %mat2v3float %T
+struct S {
+    float    f1;
+    float3   f2;
+    float2x3 f3;
+    T        f4;
+};
+// CHECK: %_ptr_ShaderRecordBufferNV_type_ShaderRecordBufferNV_S = OpTypePointer ShaderRecordBufferNV %type_ShaderRecordBufferNV_S
+
+// CHECK: %srb = OpVariable %_ptr_ShaderRecordBufferNV_type_ShaderRecordBufferNV_S ShaderRecordBufferNV
+[[vk::shader_record_nv]]
+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-nv.offset.hlsl

@@ -0,0 +1,23 @@
+// Run: %dxc -T lib_6_3
+
+// CHECK: OpMemberDecorate %type_ShaderRecordBufferNV_S 0 Offset 0
+// CHECK: OpMemberDecorate %type_ShaderRecordBufferNV_S 1 Offset 8
+// CHECK: OpMemberDecorate %type_ShaderRecordBufferNV_S 2 Offset 32
+
+struct S {
+    float a;
+    [[vk::offset(8)]]
+    float2 b;
+    [[vk::offset(32)]]
+    float4 f;
+};
+
+[[vk::shader_record_nv]]
+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

@@ -1913,4 +1913,17 @@ TEST_F(FileTest, DecorationNoContractionStageVars) {
 // For pragmas
 TEST_F(FileTest, PragmaPackMatrix) { runFileTest("pragma.pack_matrix.hlsl"); }
 
+// Tests for [[vk::shader_record_nv]]
+TEST_F(FileTest, VulkanShaderRecordBufferNV) {
+  runFileTest("vk.shader-record-nv.hlsl");
+}
+TEST_F(FileTest, VulkanLayoutShaderRecordBufferNVStd430) {
+  setGlLayout();
+  runFileTest("vk.layout.shader-record-nv.std430.hlsl");
+}
+TEST_F(FileTest, VulkanShaderRecordBufferNVOffset) {
+  // Checks the behavior of [[vk::offset]] with [[vk::shader_record_nv]]
+  runFileTest("vk.shader-record-nv.offset.hlsl");
+}
+
 } // namespace