[SPIR-V] Add -fvk-bind-*-heap flag to DXC

include/dxc/Support/HLSLOptions.td

@@ -418,6 +418,12 @@ def fvk_allow_rwstructuredbuffer_arrays: Flag<["-"], "fvk-allow-rwstructuredbuff
   HelpText<"Allow arrays of RWStructuredBuffers, AppendStructuredBuffers, and ConsumeStructuredBuffers. This is in development, and the option will be removed when the feature is complete.">;
 def fspv_max_id : MultiArg<["-"], "fspv-max-id", 1>, MetaVarName<"<shift> <space>">, Group<spirv_Group>, Flags<[CoreOption, DriverOption]>,
   HelpText<"Set the maximum value for an id in the SPIR-V binary. Default is 0x3FFFFF, which is the largest value all drivers must support.">;
+def fvk_bind_resource_heap : MultiArg<["-"], "fvk-bind-resource-heap", 2>, MetaVarName<"<binding> <set>">, Group<spirv_Group>, Flags<[CoreOption, DriverOption]>,
+  HelpText<"Specify Vulkan binding number and set number for the resource heap.">;
+def fvk_bind_sampler_heap : MultiArg<["-"], "fvk-bind-sampler-heap", 2>, MetaVarName<"<binding> <set>">, Group<spirv_Group>, Flags<[CoreOption, DriverOption]>,
+  HelpText<"Specify Vulkan binding number and set number for the sampler heap.">;
+def fvk_bind_counter_heap : MultiArg<["-"], "fvk-bind-counter-heap", 2>, MetaVarName<"<binding> <set>">, Group<spirv_Group>, Flags<[CoreOption, DriverOption]>,
+  HelpText<"Specify Vulkan binding number and set number for the counter heap.">;
 // SPIRV Change Ends
 
 //////////////////////////////////////////////////////////////////////////////

include/dxc/Support/SPIRVOptions.h

@@ -20,6 +20,8 @@
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Option/ArgList.h"
 
+#include <optional>
+
 namespace clang {
 namespace spirv {
 
@@ -98,6 +100,15 @@ struct SpirvCodeGenOptions {
   std::string entrypointName;
   std::string floatDenormalMode; // OPT_denorm
 
+  // User-defined bindings/set numbers for resource/sampler/counter heaps.
+  struct BindingInfo {
+    size_t binding;
+    size_t set;
+  };
+  std::optional<BindingInfo> resourceHeapBinding;
+  std::optional<BindingInfo> samplerHeapBinding;
+  std::optional<BindingInfo> counterHeapBinding;
+
   bool signaturePacking; ///< Whether signature packing is enabled or not
 
   bool printAll; // Dump SPIR-V module before each pass and after the last one.

lib/DxcSupport/HLSLOptions.cpp

@@ -25,10 +25,13 @@
 #include "llvm/Support/Path.h"
 #include "llvm/Support/raw_ostream.h"
 
+#include <optional>
+
 using namespace llvm::opt;
 using namespace dxc;
 using namespace hlsl;
 using namespace hlsl::options;
+using namespace clang::spirv;
 
 #define PREFIX(NAME, VALUE) static const char *const NAME[] = VALUE;
 #include "dxc/Support/HLSLOptions.inc"
@@ -320,6 +323,46 @@ static bool handleVkShiftArgs(const InputArgList &args, OptSpecifier id,
   return true;
 }
 
+// Parses the given flag |id| in |args|. If present and valid, sets |info| to
+// the correct value. Returns true if parsing succeeded. Returns false if
+// parsing failed, and outputs in |errors| a message using |name| as pretty name
+// for the flag.
+static bool
+handleFixedBinding(const InputArgList &args, OptSpecifier id,
+                   std::optional<SpirvCodeGenOptions::BindingInfo> *info,
+                   llvm::StringRef name, llvm::raw_ostream &errors) {
+  const auto values = args.getAllArgValues(id);
+  if (values.size() == 0) {
+    *info = std::nullopt;
+    return true;
+  }
+
+  if (!args.hasArg(OPT_spirv)) {
+    errors << name << " requires -spirv";
+    return false;
+  }
+
+  assert(values.size() == 2);
+
+  size_t output[2] = {0, 0};
+  for (unsigned i = 0; i < 2; ++i) {
+    int number = 0;
+    if (llvm::StringRef(values[i]).getAsInteger(10, number)) {
+      errors << "invalid " << name << " argument: '" << values[i] << "'";
+      return false;
+    }
+    if (number < 0) {
+      errors << "expected positive integer for " << name
+             << ", got: " << values[i];
+      return false;
+    }
+    output[i] = number;
+  }
+
+  *info = {output[0], output[1]};
+  return true;
+}
+
 // Check if any options that are unsupported with SPIR-V are used.
 static bool hasUnsupportedSpirvOption(const InputArgList &args,
                                       llvm::raw_ostream &errors) {
@@ -1111,6 +1154,18 @@ int ReadDxcOpts(const OptTable *optionTable, unsigned flagsToInclude,
     return 1;
   }
 
+  if (!handleFixedBinding(Args, OPT_fvk_bind_resource_heap,
+                          &opts.SpirvOptions.resourceHeapBinding,
+                          "-fvk-bind-resource-heap", errors) ||
+      !handleFixedBinding(Args, OPT_fvk_bind_sampler_heap,
+                          &opts.SpirvOptions.samplerHeapBinding,
+                          "-fvk-bind-sampler-heap", errors) ||
+      !handleFixedBinding(Args, OPT_fvk_bind_counter_heap,
+                          &opts.SpirvOptions.counterHeapBinding,
+                          "-fvk-bind-counter-heap", errors)) {
+    return 1;
+  }
+
   for (const Arg *A : Args.filtered(OPT_fspv_extension_EQ)) {
     opts.SpirvOptions.allowedExtensions.push_back(A->getValue());
   }
@@ -1249,7 +1304,10 @@ int ReadDxcOpts(const OptTable *optionTable, unsigned flagsToInclude,
       !Args.getLastArgValue(OPT_fvk_b_shift).empty() ||
       !Args.getLastArgValue(OPT_fvk_t_shift).empty() ||
       !Args.getLastArgValue(OPT_fvk_s_shift).empty() ||
-      !Args.getLastArgValue(OPT_fvk_u_shift).empty()) {
+      !Args.getLastArgValue(OPT_fvk_u_shift).empty() ||
+      !Args.getLastArgValue(OPT_fvk_bind_resource_heap).empty() ||
+      !Args.getLastArgValue(OPT_fvk_bind_sampler_heap).empty() ||
+      !Args.getLastArgValue(OPT_fvk_bind_counter_heap).empty()) {
     errors << "SPIR-V CodeGen not available. "
               "Please recompile with -DENABLE_SPIRV_CODEGEN=ON.";
     return 1;

tools/clang/lib/SPIRV/DeclResultIdMapper.cpp

@@ -15,6 +15,7 @@
 
 #include "dxc/DXIL/DxilConstants.h"
 #include "dxc/DXIL/DxilTypeSystem.h"
+#include "dxc/Support/SPIRVOptions.h"
 #include "clang/AST/Expr.h"
 #include "clang/AST/HlslTypes.h"
 #include "clang/SPIRV/AstTypeProbe.h"
@@ -2506,6 +2507,17 @@ bool DeclResultIdMapper::decorateResourceBindings() {
 
   BindingSet bindingSet;
 
+  // If some bindings are reserved for heaps, mark those are used.
+  if (spirvOptions.resourceHeapBinding)
+    bindingSet.useBinding(spirvOptions.resourceHeapBinding->binding,
+                          spirvOptions.resourceHeapBinding->set);
+  if (spirvOptions.samplerHeapBinding)
+    bindingSet.useBinding(spirvOptions.samplerHeapBinding->binding,
+                          spirvOptions.samplerHeapBinding->set);
+  if (spirvOptions.counterHeapBinding)
+    bindingSet.useBinding(spirvOptions.counterHeapBinding->binding,
+                          spirvOptions.counterHeapBinding->set);
+
   // Decorates the given varId of the given category with set number
   // setNo, binding number bindingNo. Ignores overlaps.
   const auto tryToDecorate = [this, &bindingSet](const ResourceVar &var,
@@ -2698,6 +2710,15 @@ bool DeclResultIdMapper::decorateResourceBindings() {
   return true;
 }
 
+SpirvCodeGenOptions::BindingInfo DeclResultIdMapper::getBindingInfo(
+    BindingSet &bindingSet,
+    const std::optional<SpirvCodeGenOptions::BindingInfo> &userProvidedInfo) {
+  if (userProvidedInfo.has_value()) {
+    return *userProvidedInfo;
+  }
+  return {bindingSet.useNextBinding(0), /* set= */ 0};
+}
+
 void DeclResultIdMapper::decorateResourceHeapsBindings(BindingSet &bindingSet) {
   bool hasResource = false;
   bool hasSamplers = false;
@@ -2725,12 +2746,21 @@ void DeclResultIdMapper::decorateResourceHeapsBindings(BindingSet &bindingSet) {
   // Allocate bindings only for used resources. The order of this allocation is
   // important:
   //  - First resource heaps, then sampler heaps, and finally counter heaps.
-  const uint32_t resourceBinding =
-      hasResource ? bindingSet.useNextBinding(0) : 0;
-  const uint32_t samplersBinding =
-      hasSamplers ? bindingSet.useNextBinding(0) : 0;
-  const uint32_t countersBinding =
-      hasCounters ? bindingSet.useNextBinding(0) : 0;
+  SpirvCodeGenOptions::BindingInfo resourceBinding = {/* binding= */ 0,
+                                                      /* set= */ 0};
+  SpirvCodeGenOptions::BindingInfo samplersBinding = {/* binding= */ 0,
+                                                      /* set= */ 0};
+  SpirvCodeGenOptions::BindingInfo countersBinding = {/* binding= */ 0,
+                                                      /* set= */ 0};
+  if (hasResource)
+    resourceBinding =
+        getBindingInfo(bindingSet, spirvOptions.resourceHeapBinding);
+  if (hasSamplers)
+    samplersBinding =
+        getBindingInfo(bindingSet, spirvOptions.samplerHeapBinding);
+  if (hasCounters)
+    countersBinding =
+        getBindingInfo(bindingSet, spirvOptions.counterHeapBinding);
 
   for (const auto &var : resourceVars) {
     if (!var.getDeclaration())
@@ -2739,13 +2769,14 @@ void DeclResultIdMapper::decorateResourceHeapsBindings(BindingSet &bindingSet) {
     if (!decl)
       continue;
 
-    if (isResourceDescriptorHeap(decl->getType()))
-      spvBuilder.decorateDSetBinding(var.getSpirvInstr(), /* set= */ 0,
-                                     var.isCounter() ? countersBinding
-                                                     : resourceBinding);
-    else if (isSamplerDescriptorHeap(decl->getType()))
-      spvBuilder.decorateDSetBinding(var.getSpirvInstr(), /* set= */ 0,
-                                     samplersBinding);
+    const bool isResourceHeap = isResourceDescriptorHeap(decl->getType());
+    const bool isSamplerHeap = isSamplerDescriptorHeap(decl->getType());
+    if (!isSamplerHeap && !isResourceHeap)
+      continue;
+    const SpirvCodeGenOptions::BindingInfo &info =
+        isSamplerHeap ? samplersBinding
+                      : (var.isCounter() ? countersBinding : resourceBinding);
+    spvBuilder.decorateDSetBinding(var.getSpirvInstr(), info.set, info.binding);
   }
 }
 

tools/clang/lib/SPIRV/DeclResultIdMapper.h

@@ -643,7 +643,13 @@ class DeclResultIdMapper {
       llvm::DenseSet<StageVariableLocationInfo, StageVariableLocationInfo>
           *stageVariableLocationInfo);
 
-  /// \bried Decorates used Resource/Sampler descriptor heaps with the correct
+  /// \brief Get a valid BindingInfo. If no user provided binding info is given,
+  /// allocates a new binding and returns it.
+  static SpirvCodeGenOptions::BindingInfo getBindingInfo(
+      BindingSet &bindingSet,
+      const std::optional<SpirvCodeGenOptions::BindingInfo> &userProvidedInfo);
+
+  /// \brief Decorates used Resource/Sampler descriptor heaps with the correct
   /// binding/set decorations.
   void decorateResourceHeapsBindings(BindingSet &bindingSet);
 

tools/clang/test/CodeGenSPIRV/sm6_6.descriptorheap.aliasing.hlsl

@@ -0,0 +1,13 @@
+// RUN: %dxc -T ps_6_6 -spirv %s -fvk-bind-sampler-heap 4 7 -fvk-bind-resource-heap 4 7 | FileCheck %s
+
+// CHECK-DAG: OpDecorate %ResourceDescriptorHeap DescriptorSet 7
+// CHECK-DAG: OpDecorate %ResourceDescriptorHeap Binding 4
+// CHECK-DAG: OpDecorate %SamplerDescriptorHeap DescriptorSet 7
+// CHECK-DAG: OpDecorate %SamplerDescriptorHeap Binding 4
+
+float4 main() : SV_Target {
+  SamplerState Sampler = SamplerDescriptorHeap[2];
+  Texture2D Texture = ResourceDescriptorHeap[3];
+  return Texture.Sample(Sampler, float2(0, 0));
+}
+

tools/clang/test/CodeGenSPIRV/sm6_6.descriptorheap.flag.invalid.hlsl

@@ -0,0 +1,17 @@
+// RUN: not %dxc -T cs_6_6 -spirv %s 2>&1 -fvk-bind-sampler-heap  a 2 | FileCheck %s --check-prefix=CHECK-SAMPLER-A
+// RUN: not %dxc -T cs_6_6 -spirv %s 2>&1 -fvk-bind-sampler-heap  1 b | FileCheck %s --check-prefix=CHECK-SAMPLER-B
+// RUN: not %dxc -T cs_6_6 -spirv %s 2>&1 -fvk-bind-resource-heap a 2 | FileCheck %s --check-prefix=CHECK-RESOURCE-A
+// RUN: not %dxc -T cs_6_6 -spirv %s 2>&1 -fvk-bind-resource-heap 1 b | FileCheck %s --check-prefix=CHECK-RESOURCE-B
+// RUN: not %dxc -T cs_6_6 -spirv %s 2>&1 -fvk-bind-counter-heap  a 2 | FileCheck %s --check-prefix=CHECK-COUNTER-A
+// RUN: not %dxc -T cs_6_6 -spirv %s 2>&1 -fvk-bind-counter-heap  1 b | FileCheck %s --check-prefix=CHECK-COUNTER-B
+
+//  CHECK-SAMPLER-A: invalid -fvk-bind-sampler-heap argument: 'a'
+// CHECK-RESOURCE-A: invalid -fvk-bind-resource-heap argument: 'a'
+//  CHECK-COUNTER-A: invalid -fvk-bind-counter-heap argument: 'a'
+//  CHECK-SAMPLER-B: invalid -fvk-bind-sampler-heap argument: 'b'
+// CHECK-RESOURCE-B: invalid -fvk-bind-resource-heap argument: 'b'
+//  CHECK-COUNTER-B: invalid -fvk-bind-counter-heap argument: 'b'
+[numthreads(1, 1, 1)]
+void main() {
+  SamplerState Sampler = SamplerDescriptorHeap[0];
+}

tools/clang/test/CodeGenSPIRV/sm6_6.descriptorheap.flag.missing.hlsl

@@ -0,0 +1,14 @@
+// RUN: not %dxc -T cs_6_6 -spirv %s 2>&1 -fvk-bind-sampler-heap 1 | FileCheck %s --check-prefix=CHECK-SAMPLER
+// RUN: not %dxc -T cs_6_6 -spirv %s 2>&1 -fvk-bind-sampler-heap   | FileCheck %s --check-prefix=CHECK-SAMPLER
+// RUN: not %dxc -T cs_6_6 -spirv %s 2>&1 -fvk-bind-resource-heap 1 | FileCheck %s --check-prefix=CHECK-RESOURCE
+// RUN: not %dxc -T cs_6_6 -spirv %s 2>&1 -fvk-bind-resource-heap   | FileCheck %s --check-prefix=CHECK-RESOURCE
+// RUN: not %dxc -T cs_6_6 -spirv %s 2>&1 -fvk-bind-counter-heap 1 | FileCheck %s --check-prefix=CHECK-COUNTER
+// RUN: not %dxc -T cs_6_6 -spirv %s 2>&1 -fvk-bind-counter-heap   | FileCheck %s --check-prefix=CHECK-COUNTER
+
+// CHECK-SAMPLER: Argument to '-fvk-bind-sampler-heap' is missing.
+// CHECK-RESOURCE: Argument to '-fvk-bind-resource-heap' is missing.
+// CHECK-COUNTER: Argument to '-fvk-bind-counter-heap' is missing.
+[numthreads(1, 1, 1)]
+void main() {
+  SamplerState Sampler = SamplerDescriptorHeap[0];
+}

tools/clang/test/CodeGenSPIRV/sm6_6.descriptorheap.flag.negative.hlsl

@@ -0,0 +1,14 @@
+// RUN: not %dxc -T cs_6_6 -spirv %s 2>&1 -fvk-bind-sampler-heap  -1  2 | FileCheck %s --check-prefix=CHECK-SAMPLER
+// RUN: not %dxc -T cs_6_6 -spirv %s 2>&1 -fvk-bind-sampler-heap   1 -1 | FileCheck %s --check-prefix=CHECK-SAMPLER
+// RUN: not %dxc -T cs_6_6 -spirv %s 2>&1 -fvk-bind-resource-heap -1  2 | FileCheck %s --check-prefix=CHECK-RESOURCE
+// RUN: not %dxc -T cs_6_6 -spirv %s 2>&1 -fvk-bind-resource-heap  1 -1 | FileCheck %s --check-prefix=CHECK-RESOURCE
+// RUN: not %dxc -T cs_6_6 -spirv %s 2>&1 -fvk-bind-counter-heap  -1  2 | FileCheck %s --check-prefix=CHECK-COUNTER
+// RUN: not %dxc -T cs_6_6 -spirv %s 2>&1 -fvk-bind-counter-heap   1 -1 | FileCheck %s --check-prefix=CHECK-COUNTER
+
+//  CHECK-SAMPLER: expected positive integer for -fvk-bind-sampler-heap, got: -1
+// CHECK-RESOURCE: expected positive integer for -fvk-bind-resource-heap, got: -1
+//  CHECK-COUNTER: expected positive integer for -fvk-bind-counter-heap, got: -1
+[numthreads(1, 1, 1)]
+void main() {
+  SamplerState Sampler = SamplerDescriptorHeap[0];
+}

tools/clang/test/CodeGenSPIRV/sm6_6.descriptorheap.implicit-bindings.cbuffer.flag.hlsl

@@ -0,0 +1,16 @@
+// RUN: %dxc -T ps_6_6 -E main %s -spirv -fvk-bind-resource-heap 6 7 | FileCheck %s
+
+// CHECK-DAG: OpDecorate %a DescriptorSet 0
+// CHECK-DAG: OpDecorate %a Binding 0
+cbuffer a {
+  float4 value;
+}
+
+float4 main() : SV_Target {
+// CHECK-DAG: OpDecorate %ResourceDescriptorHeap DescriptorSet 7
+// CHECK-DAG: OpDecorate %ResourceDescriptorHeap Binding 6
+  RWStructuredBuffer<float4> buffer = ResourceDescriptorHeap[1];
+  buffer[0] = value;
+  return value;
+}
+

tools/clang/test/CodeGenSPIRV/sm6_6.descriptorheap.implicit-bindings.cbuffer.hlsl

@@ -1,4 +1,4 @@
-// RUN: %dxc -T ps_6_6 -E main %s -spirv  2>&1 | FileCheck %s
+// RUN: %dxc -T ps_6_6 -E main %s -spirv | FileCheck %s
 
 // CHECK-DAG: OpDecorate %a DescriptorSet 0
 // CHECK-DAG: OpDecorate %a Binding 0

tools/clang/test/CodeGenSPIRV/sm6_6.descriptorheap.implicit-bindings.counter.flag.hlsl

@@ -0,0 +1,21 @@
+// RUN: %dxc -T cs_6_6 -E main -fcgl %s -spirv -fvk-bind-counter-heap 4 5 | FileCheck %s
+
+// CHECK-DAG: OpDecorate %a DescriptorSet 4
+// CHECK-DAG: OpDecorate %a Binding 3
+// CHECK-DAG: OpDecorate %counter_var_a DescriptorSet 4
+// CHECK-DAG: OpDecorate %counter_var_a Binding 0
+[[vk::binding(3, 4)]]
+RWStructuredBuffer<uint> a;
+
+[numthreads(1, 1, 1)]
+void main() {
+// CHECK-DAG: OpDecorate %ResourceDescriptorHeap DescriptorSet 0
+// CHECK-DAG: OpDecorate %ResourceDescriptorHeap Binding 0
+// CHECK-DAG: OpDecorate %counter_var_ResourceDescriptorHeap DescriptorSet 5
+// CHECK-DAG: OpDecorate %counter_var_ResourceDescriptorHeap Binding 4
+  RWStructuredBuffer<uint> b = ResourceDescriptorHeap[1];
+
+  a.IncrementCounter();
+  b.IncrementCounter();
+}
+

tools/clang/test/CodeGenSPIRV/sm6_6.descriptorheap.implicit-bindings.counter.hlsl

@@ -1,4 +1,4 @@
-// RUN: %dxc -T cs_6_6 -E main -fcgl %s -spirv 2>&1 | FileCheck %s
+// RUN: %dxc -T cs_6_6 -E main -fcgl %s -spirv | FileCheck %s
 
 // CHECK-DAG: OpDecorate %a DescriptorSet 4
 // CHECK-DAG: OpDecorate %a Binding 3

tools/clang/test/CodeGenSPIRV/sm6_6.descriptorheap.implicit-bindings.flag.hlsl

@@ -0,0 +1,13 @@
+// RUN: %dxc -T ps_6_6 -E main -fcgl %s -spirv -fvk-bind-sampler-heap 3 4 | FileCheck %s
+
+// CHECK-DAG: OpDecorate %Texture DescriptorSet 0
+// CHECK-DAG: OpDecorate %Texture Binding 0
+Texture2D Texture;
+
+float4 main() : SV_Target {
+// CHECK-DAG: OpDecorate %SamplerDescriptorHeap DescriptorSet 4
+// CHECK-DAG: OpDecorate %SamplerDescriptorHeap Binding 3
+  SamplerState Sampler = SamplerDescriptorHeap[0];
+  return Texture.Sample(Sampler, float2(0, 0));
+}
+

tools/clang/test/CodeGenSPIRV/sm6_6.descriptorheap.implicit-bindings.hlsl

@@ -1,4 +1,4 @@
-// RUN: %dxc -T ps_6_6 -E main -fcgl %s -spirv  2>&1 | FileCheck %s
+// RUN: %dxc -T ps_6_6 -E main -fcgl %s -spirv | FileCheck %s
 
 // CHECK-DAG: OpDecorate %Texture DescriptorSet 0
 // CHECK-DAG: OpDecorate %Texture Binding 0

tools/clang/test/CodeGenSPIRV/sm6_6.descriptorheap.implicit-bindings.reserved.hlsl

@@ -0,0 +1,18 @@
+// RUN: %dxc -T cs_6_6 -E main -fcgl %s -spirv | FileCheck %s --check-prefix=CHECK-DEFAULT
+// RUN: %dxc -T cs_6_6 -E main -fcgl %s -spirv -fvk-bind-resource-heap 0 0 -fvk-bind-counter-heap 1 0 -fvk-bind-sampler-heap 2 0 | FileCheck %s --check-prefix=CHECK-RESERVED
+
+// CHECK-DEFAULT-DAG: OpDecorate %a DescriptorSet 0
+// CHECK-DEFAULT-DAG: OpDecorate %a Binding 0
+// CHECK-DEFAULT-DAG: OpDecorate %counter_var_a DescriptorSet 0
+// CHECK-DEFAULT-DAG: OpDecorate %counter_var_a Binding 1
+
+// CHECK-RESERVED-DAG: OpDecorate %a DescriptorSet 0
+// CHECK-RESERVED-DAG: OpDecorate %a Binding 3
+// CHECK-RESERVED-DAG: OpDecorate %counter_var_a DescriptorSet 0
+// CHECK-RESERVED-DAG: OpDecorate %counter_var_a Binding 4
+RWStructuredBuffer<uint> a;
+
+[numthreads(1, 1, 1)]
+void main() {
+  a.IncrementCounter();
+}