[SPIR-V] Add -fvk-bind-*-heap flag to DXC (#6919)

Those flags allow the user to hardcode the binding ID for the heaps.
This allow the user to either force heap aliasing, or simply have more
control over the heaps.

Fixes #6913

---------

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

include/dxc/Support/HLSLOptions.td

@@ -416,6 +416,12 @@ def fspv_preserve_interface : Flag<["-"], "fspv-preserve-interface">, Group<spir
   HelpText<"Preserves all interface variables in the entry point, even when those variables are unused">;
 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 {
 
@@ -97,6 +99,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) {
@@ -1109,6 +1152,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());
   }
@@ -1245,7 +1300,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();
+}