[spirv] Respect the -auto-binding-space option in SPIR-V backend (#2211)

* [spirv] Respect the -auto-binding-space option in SPIR-V backend.

* [spirv] Use default space if vk::binding doesn't provide dset.

* Address code review comments.

docs/SPIR-V.rst

@@ -1452,11 +1452,13 @@ Explicit binding number assignment
 
 ``[[vk::binding(X[, Y])]]`` can be attached to global variables to specify the
 descriptor set as ``Y`` and binding number as ``X``. The descriptor set number
-is optional; if missing, it will be zero. RW/append/consume structured buffers
-have associated counters, which will occupy their own Vulkan descriptors.
-``[vk::counter_binding(Z)]`` can be attached to a RW/append/consume structured
-buffers to specify the binding number for the associated counter to ``Z``. Note
-that the set number of the counter is always the same as the main buffer.
+is optional; if missing, it will be zero (If ``-auto-binding-space N`` command
+line option is used, then descriptor set #N will be used instead of descriptor
+set #0). RW/append/consume structured buffers have associated counters, which
+will occupy their own Vulkan descriptors. ``[vk::counter_binding(Z)]`` can be
+attached to a RW/append/consume structured buffers to specify the binding number
+for the associated counter to ``Z``. Note that the set number of the counter is
+always the same as the main buffer.
 
 Implicit binding number assignment
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
@@ -1476,7 +1478,8 @@ respectively.
 
 If there is no register specification, the corresponding resource will be
 assigned to the next available binding number, starting from 0, in descriptor
-set #0.
+set #0 (If ``-auto-binding-space N`` command line option is used, then
+descriptor set #N will be used instead of descriptor set #0).
 
 Summary
 ~~~~~~~

tools/clang/lib/SPIRV/DeclResultIdMapper.cpp

@@ -28,6 +28,15 @@ namespace spirv {
 
 namespace {
 
+uint32_t getVkBindingAttrSet(const VKBindingAttr *attr, uint32_t defaultSet) {
+  // If the [[vk::binding(x)]] attribute is provided without the descriptor set,
+  // we should use the default descriptor set.
+  if (attr->getSet() == INT_MIN) {
+    return defaultSet;
+  }
+  return attr->getSet();
+}
+
 /// Returns the :packoffset() annotation on the given decl. Returns nullptr if
 /// the decl does not have one.
 hlsl::ConstantPacking *getPackOffset(const clang::NamedDecl *decl) {
@@ -1383,11 +1392,11 @@ public:
 
   /// Returns true and set the correct set and binding number if we can find a
   /// descriptor setting for the given register. False otherwise.
-  bool getSetBinding(const hlsl::RegisterAssignment *regAttr, int *setNo,
-                     int *bindNo) const {
+  bool getSetBinding(const hlsl::RegisterAssignment *regAttr,
+                     uint32_t defaultSpace, int *setNo, int *bindNo) const {
     std::ostringstream iss;
-    iss << regAttr->RegisterSpace.getValueOr(0) << regAttr->RegisterType
-        << regAttr->RegisterNumber;
+    iss << regAttr->RegisterSpace.getValueOr(defaultSpace)
+        << regAttr->RegisterType << regAttr->RegisterNumber;
 
     auto found = mapping.find(iss.str());
     if (found != mapping.end()) {
@@ -1424,6 +1433,13 @@ bool DeclResultIdMapper::decorateResourceBindings() {
   // - m2
   // - m3, m4, mX * c2
 
+  // The "-auto-binding-space" command line option can be used to specify a
+  // certain space as default. UINT_MAX means the user has not provided this
+  // option. If not provided, the SPIR-V backend uses space "0" as default.
+  auto defaultSpaceOpt =
+      theEmitter.getCompilerInstance().getCodeGenOpts().HLSLDefaultSpace;
+  uint32_t defaultSpace = (defaultSpaceOpt == UINT_MAX) ? 0 : defaultSpaceOpt;
+
   const bool bindGlobals = !spirvOptions.bindGlobals.empty();
   int32_t globalsBindNo = -1, globalsSetNo = -1;
   if (bindGlobals) {
@@ -1459,11 +1475,12 @@ bool DeclResultIdMapper::decorateResourceBindings() {
       if (const auto *regAttr = var.getRegister()) {
         if (var.isCounter()) {
           emitError("-fvk-bind-register for RW/Append/Consume StructuredBuffer "
-                    "umimplemented",
+                    "unimplemented",
                     var.getSourceLocation());
         } else {
           int setNo = 0, bindNo = 0;
-          if (!bindingMapper.getSetBinding(regAttr, &setNo, &bindNo)) {
+          if (!bindingMapper.getSetBinding(regAttr, defaultSpace, &setNo,
+                                           &bindNo)) {
             emitError("missing -fvk-bind-register for resource",
                       var.getSourceLocation());
             return false;
@@ -1498,18 +1515,19 @@ bool DeclResultIdMapper::decorateResourceBindings() {
     if (var.isCounter()) {
       if (const auto *vkCBinding = var.getCounterBinding()) {
         // Process mX * c1
-        uint32_t set = 0;
+        uint32_t set = defaultSpace;
         if (const auto *vkBinding = var.getBinding())
-          set = vkBinding->getSet();
+          set = getVkBindingAttrSet(vkBinding, defaultSpace);
         else if (const auto *reg = var.getRegister())
-          set = reg->RegisterSpace.getValueOr(0);
+          set = reg->RegisterSpace.getValueOr(defaultSpace);
 
         tryToDecorate(var.getSpirvInstr(), set, vkCBinding->getBinding());
       }
     } else {
       if (const auto *vkBinding = var.getBinding()) {
         // Process m1
-        tryToDecorate(var.getSpirvInstr(), vkBinding->getSet(),
+        tryToDecorate(var.getSpirvInstr(),
+                      getVkBindingAttrSet(vkBinding, defaultSpace),
                       vkBinding->getBinding());
       }
     }
@@ -1528,7 +1546,7 @@ bool DeclResultIdMapper::decorateResourceBindings() {
         if (reg->isSpaceOnly())
           continue;
 
-        const uint32_t set = reg->RegisterSpace.getValueOr(0);
+        const uint32_t set = reg->RegisterSpace.getValueOr(defaultSpace);
         uint32_t binding = reg->RegisterNumber;
         switch (reg->RegisterType) {
         case 'b':
@@ -1557,11 +1575,11 @@ bool DeclResultIdMapper::decorateResourceBindings() {
     if (var.isCounter()) {
       if (!var.getCounterBinding()) {
         // Process mX * c2
-        uint32_t set = 0;
+        uint32_t set = defaultSpace;
         if (const auto *vkBinding = var.getBinding())
-          set = vkBinding->getSet();
+          set = getVkBindingAttrSet(vkBinding, defaultSpace);
         else if (const auto *reg = var.getRegister())
-          set = reg->RegisterSpace.getValueOr(0);
+          set = reg->RegisterSpace.getValueOr(defaultSpace);
 
         spvBuilder.decorateDSetBinding(var.getSpirvInstr(), set,
                                        bindingSet.useNextBinding(set));
@@ -1569,7 +1587,7 @@ bool DeclResultIdMapper::decorateResourceBindings() {
     } else if (!var.getBinding()) {
       const auto *reg = var.getRegister();
       if (reg && reg->isSpaceOnly()) {
-        const uint32_t set = reg->RegisterSpace.getValueOr(0);
+        const uint32_t set = reg->RegisterSpace.getValueOr(defaultSpace);
         spvBuilder.decorateDSetBinding(var.getSpirvInstr(), set,
                                        bindingSet.useNextBinding(set));
       } else if (!reg) {
@@ -1584,8 +1602,9 @@ bool DeclResultIdMapper::decorateResourceBindings() {
         }
         // The normal case
         else {
-          spvBuilder.decorateDSetBinding(var.getSpirvInstr(), 0,
-                                         bindingSet.useNextBinding(0));
+          spvBuilder.decorateDSetBinding(
+              var.getSpirvInstr(), defaultSpace,
+              bindingSet.useNextBinding(defaultSpace));
         }
       }
     }

tools/clang/lib/SPIRV/SpirvEmitter.h

@@ -51,6 +51,7 @@ public:
   ASTContext &getASTContext() { return astContext; }
   SpirvBuilder &getSpirvBuilder() { return spvBuilder; }
   DiagnosticsEngine &getDiagnosticsEngine() { return diags; }
+  CompilerInstance &getCompilerInstance() { return theCompilerInstance; }
 
   void doDecl(const Decl *decl);
   void doStmt(const Stmt *stmt, llvm::ArrayRef<const Attr *> attrs = {});

tools/clang/lib/Sema/SemaHLSL.cpp

@@ -11029,9 +11029,10 @@ void hlsl::HandleDeclAttributeForHLSL(Sema &S, Decl *D, const AttributeList &A,
       ValidateAttributeIntArg(S, A), A.getAttributeSpellingListIndex());
     break;
   case AttributeList::AT_VKBinding:
-    declAttr = ::new (S.Context) VKBindingAttr(A.getRange(), S.Context,
-      ValidateAttributeIntArg(S, A), ValidateAttributeIntArg(S, A, 1),
-      A.getAttributeSpellingListIndex());
+    declAttr = ::new (S.Context) VKBindingAttr(
+        A.getRange(), S.Context, ValidateAttributeIntArg(S, A),
+        A.getNumArgs() < 2 ? INT_MIN : ValidateAttributeIntArg(S, A, 1),
+        A.getAttributeSpellingListIndex());
     break;
   case AttributeList::AT_VKCounterBinding:
     declAttr = ::new (S.Context) VKCounterBindingAttr(A.getRange(), S.Context,

tools/clang/test/CodeGenSPIRV/vk.binding.cl.register.counter.hlsl

@@ -7,4 +7,4 @@ float4 main() : SV_Target {
   return MyBuffer[0].val;
 }
 
-// CHECK: :4:24: error: -fvk-bind-register for RW/Append/Consume StructuredBuffer umimplemented
+// CHECK: :4:24: error: -fvk-bind-register for RW/Append/Consume StructuredBuffer unimplemented

tools/clang/test/CodeGenSPIRV/vk.binding.default-space.explicit.hlsl

@@ -0,0 +1,35 @@
+// Run: %dxc -T ps_6_0 -E main -auto-binding-space 77
+
+// Since both the descriptor set and binding are mentioned via "[[vk::binding]]",
+// the "-auto-binding-space" has no effect:
+//
+// CHECK:      OpDecorate %sampler1 DescriptorSet 6
+// CHECK-NEXT: OpDecorate %sampler1 Binding 5
+[[vk::binding(5,6)]]
+SamplerState sampler1 : register(s1, space2);
+
+// Since the space is NOT provided in "[[vk::binding]]",
+// we use the default space specified by "-auto-binding-space"
+//
+// CHECK:      OpDecorate %sampler2 DescriptorSet 77
+// CHECK-NEXT: OpDecorate %sampler2 Binding 8
+[[vk::binding(8)]]
+SamplerState sampler2 : register(s4);
+
+// Since both the descriptor set and binding are mentioned via ":register",
+// the "-auto-binding-space" has no effect:
+//
+// CHECK:      OpDecorate %sampler3 DescriptorSet 9
+// CHECK-NEXT: OpDecorate %sampler3 Binding 1
+SamplerState sampler3 : register(s1, space9);
+
+// Since the space is NOT provided in ":register",
+// we use the default space specified by "-auto-binding-space"
+//
+// CHECK:      OpDecorate %sampler4 DescriptorSet 77
+// CHECK-NEXT: OpDecorate %sampler4 Binding 3
+SamplerState sampler4 : register(s3);
+
+float4 main() : SV_Target {
+  return 1.0;
+}

tools/clang/test/CodeGenSPIRV/vk.binding.default-space.implicit.hlsl

@@ -0,0 +1,64 @@
+// Run: %dxc -T ps_6_0 -E main -auto-binding-space 77
+
+// CHECK:      OpDecorate %sampler1 DescriptorSet 77
+// CHECK-NEXT: OpDecorate %sampler1 Binding 0
+SamplerState sampler1;
+
+// CHECK:      OpDecorate %texture1 DescriptorSet 77
+// CHECK-NEXT: OpDecorate %texture1 Binding 1
+Texture2D<float4> texture1;
+
+// CHECK:      OpDecorate %texture2 DescriptorSet 77
+// CHECK-NEXT: OpDecorate %texture2 Binding 2
+Texture3D<float4> texture2;
+
+// CHECK:      OpDecorate %sampler2 DescriptorSet 77
+// CHECK-NEXT: OpDecorate %sampler2 Binding 3
+SamplerState sampler2;
+
+// CHECK:      OpDecorate %myCbuffer DescriptorSet 77
+// CHECK-NEXT: OpDecorate %myCbuffer Binding 4
+cbuffer myCbuffer {
+    float4 stuff;
+}
+
+// CHECK:      OpDecorate %myBuffer DescriptorSet 77
+// CHECK-NEXT: OpDecorate %myBuffer Binding 5
+Buffer<int> myBuffer;
+
+// CHECK:      OpDecorate %myRWBuffer DescriptorSet 77
+// CHECK-NEXT: OpDecorate %myRWBuffer Binding 6
+RWBuffer<float4> myRWBuffer;
+
+struct S {
+    float4 f;
+};
+
+// CHECK:      OpDecorate %myCbuffer2 DescriptorSet 77
+// CHECK-NEXT: OpDecorate %myCbuffer2 Binding 7
+ConstantBuffer<S> myCbuffer2;
+
+// CHECK:      OpDecorate %sbuffer1 DescriptorSet 77
+// CHECK-NEXT: OpDecorate %sbuffer1 Binding 8
+  StructuredBuffer<S> sbuffer1;
+// CHECK:      OpDecorate %sbuffer2 DescriptorSet 77
+// CHECK-NEXT: OpDecorate %sbuffer2 Binding 9
+// CHECK-NEXT: OpDecorate %counter_var_sbuffer2 DescriptorSet 77
+// CHECK-NEXT: OpDecorate %counter_var_sbuffer2 Binding 10
+RWStructuredBuffer<S> sbuffer2;
+
+// CHECK:      OpDecorate %abuffer DescriptorSet 77
+// CHECK-NEXT: OpDecorate %abuffer Binding 11
+// CHECK-NEXT: OpDecorate %counter_var_abuffer DescriptorSet 77
+// CHECK-NEXT: OpDecorate %counter_var_abuffer Binding 12
+AppendStructuredBuffer<S> abuffer;
+
+// CHECK:      OpDecorate %csbuffer DescriptorSet 77
+// CHECK-NEXT: OpDecorate %csbuffer Binding 13
+// CHECK-NEXT: OpDecorate %counter_var_csbuffer DescriptorSet 77
+// CHECK-NEXT: OpDecorate %counter_var_csbuffer Binding 14
+ConsumeStructuredBuffer<S> csbuffer;
+
+float4 main() : SV_Target {
+    return 1.0;
+}

tools/clang/test/CodeGenSPIRV/vk.binding.default-space.with-shift-all.hlsl

@@ -0,0 +1,54 @@
+// Run: %dxc -T ps_6_0 -E main -auto-binding-space 77 -fvk-b-shift 1000 all -fvk-t-shift 2000 all -fvk-s-shift 3000 all -fvk-u-shift 4000 all
+
+struct S {
+    float4 f;
+};
+
+// Explicit binding assignment is unaffected by the shift, but the default
+// space is used when the descriptor set is not provided in vk::binding.
+
+// CHECK: OpDecorate %cbuffer3 DescriptorSet 77
+// CHECK: OpDecorate %cbuffer3 Binding 42
+[[vk::binding(42)]]
+ConstantBuffer<S> cbuffer3 : register(b10, space2);
+
+// CHECK: OpDecorate %cbuffer1 DescriptorSet 77
+// CHECK: OpDecorate %cbuffer1 Binding 1000
+ConstantBuffer<S> cbuffer1 : register(b0);
+// CHECK: OpDecorate %cbuffer2 DescriptorSet 2
+// CHECK: OpDecorate %cbuffer2 Binding 1000
+ConstantBuffer<S> cbuffer2 : register(b0, space2);
+
+// CHECK: OpDecorate %texture1 DescriptorSet 1
+// CHECK: OpDecorate %texture1 Binding 2001
+Texture2D<float4> texture1: register(t1, space1);
+// CHECK: OpDecorate %texture2 DescriptorSet 77
+// CHECK: OpDecorate %texture2 Binding 2001
+Texture2D<float4> texture2: register(t1);
+
+// CHECK: OpDecorate %sampler1 DescriptorSet 77
+// CHECK: OpDecorate %sampler1 Binding 3000
+// CHECK: OpDecorate %sampler2 DescriptorSet 2
+// CHECK: OpDecorate %sampler2 Binding 3000
+SamplerState sampler1: register(s0);
+SamplerState sampler2: register(s0, space2);
+
+// CHECK: OpDecorate %rwbuffer1 DescriptorSet 3
+// CHECK: OpDecorate %rwbuffer1 Binding 4003
+RWBuffer<float4> rwbuffer1 : register(u3, space3);
+// CHECK: OpDecorate %rwbuffer2 DescriptorSet 77
+// CHECK: OpDecorate %rwbuffer2 Binding 4003
+RWBuffer<float4> rwbuffer2 : register(u3);
+
+// Lacking binding assignment is unaffacted.
+
+// CHECK: OpDecorate %cbuffer4 DescriptorSet 77
+// CHECK: OpDecorate %cbuffer4 Binding 0
+ConstantBuffer<S> cbuffer4;
+// CHECK: OpDecorate %cbuffer5 DescriptorSet 77
+// CHECK: OpDecorate %cbuffer5 Binding 1
+ConstantBuffer<S> cbuffer5;
+
+float4 main() : SV_Target {
+    return cbuffer1.f;
+}

tools/clang/test/CodeGenSPIRV/vk.binding.default-space.with-shift.hlsl

@@ -0,0 +1,57 @@
+// Run: %dxc -T ps_6_0 -E main -auto-binding-space 77 -fvk-b-shift 100 77 -fvk-b-shift 200 2 -fvk-t-shift 400 77 -fvk-t-shift 300 1 -fvk-s-shift 500 77 -fvk-s-shift 600 2 -fvk-u-shift 700 77 -fvk-u-shift 800 3
+
+// Tests that we can set shift for more than one sets of the same register type
+// Tests that we can override shift for the same set
+
+struct S {
+    float4 f;
+};
+
+// Explicit binding assignment is unaffected by the shift,
+// but the default space is used if no space is provided by vk::binding.
+
+// CHECK: OpDecorate %cbuffer3 DescriptorSet 77
+// CHECK: OpDecorate %cbuffer3 Binding 42
+[[vk::binding(42)]]
+ConstantBuffer<S> cbuffer3 : register(b10, space2);
+
+// CHECK: OpDecorate %cbuffer1 DescriptorSet 77
+// CHECK: OpDecorate %cbuffer1 Binding 100
+ConstantBuffer<S> cbuffer1 : register(b0);
+// CHECK: OpDecorate %cbuffer2 DescriptorSet 2
+// CHECK: OpDecorate %cbuffer2 Binding 200
+ConstantBuffer<S> cbuffer2 : register(b0, space2);
+
+// CHECK: OpDecorate %texture1 DescriptorSet 1
+// CHECK: OpDecorate %texture1 Binding 301
+Texture2D<float4> texture1: register(t1, space1);
+// CHECK: OpDecorate %texture2 DescriptorSet 77
+// CHECK: OpDecorate %texture2 Binding 401
+Texture2D<float4> texture2: register(t1);
+
+// CHECK: OpDecorate %sampler1 DescriptorSet 77
+// CHECK: OpDecorate %sampler1 Binding 500
+SamplerState sampler1: register(s0);
+// CHECK: OpDecorate %sampler2 DescriptorSet 2
+// CHECK: OpDecorate %sampler2 Binding 600
+SamplerState sampler2: register(s0, space2);
+
+// CHECK: OpDecorate %rwbuffer1 DescriptorSet 3
+// CHECK: OpDecorate %rwbuffer1 Binding 803
+RWBuffer<float4> rwbuffer1 : register(u3, space3);
+// CHECK: OpDecorate %rwbuffer2 DescriptorSet 77
+// CHECK: OpDecorate %rwbuffer2 Binding 703
+RWBuffer<float4> rwbuffer2 : register(u3);
+
+// Lacking binding assignment is unaffacted by the shift.
+
+// CHECK: OpDecorate %cbuffer4 DescriptorSet 77
+// CHECK: OpDecorate %cbuffer4 Binding 0
+ConstantBuffer<S> cbuffer4;
+// CHECK: OpDecorate %cbuffer5 DescriptorSet 77
+// CHECK: OpDecorate %cbuffer5 Binding 1
+ConstantBuffer<S> cbuffer5;
+
+float4 main() : SV_Target {
+    return cbuffer1.f;
+}

tools/clang/unittests/SPIRV/CodeGenSpirvTest.cpp

@@ -1577,6 +1577,23 @@ TEST_F(FileTest, VulkanRegisterBinding1to1MappingMissingCLOption) {
 TEST_F(FileTest, VulkanRegisterBinding1to1MappingAssociatedCounter) {
   runFileTest("vk.binding.cl.register.counter.hlsl", Expect::Failure);
 }
+
+// For testing the "-auto-binding-space" command line option which specifies the
+// "default space" for resources.
+TEST_F(FileTest, VulkanRegisterBindingDefaultSpaceImplicit) {
+  runFileTest("vk.binding.default-space.implicit.hlsl");
+}
+TEST_F(FileTest, VulkanRegisterBindingDefaultSpaceExplicit) {
+  runFileTest("vk.binding.default-space.explicit.hlsl");
+}
+// Testing combinations of "-auto-binding-space" and other options.
+TEST_F(FileTest, VulkanRegisterBindingDefaultSpaceWithShift) {
+  runFileTest("vk.binding.default-space.with-shift.hlsl");
+}
+TEST_F(FileTest, VulkanRegisterBindingDefaultSpaceWithShiftAll) {
+  runFileTest("vk.binding.default-space.with-shift-all.hlsl");
+}
+
 TEST_F(FileTest, VulkanStructuredBufferCounter) {
   // [[vk::counter_binding()]] for RWStructuredBuffer, AppendStructuredBuffer,
   // and ConsumeStructuredBuffer