[spirv] Support shifting all sets with -fvk-*-shift N all (#1224)

docs/SPIR-V.rst

@@ -2683,8 +2683,9 @@ codegen for Vulkan:
   sets its Vulkan descriptor set to ``M`` and binding number to ``X + N``. If
   you need to shift the inferred binding numbers for more than one space,
   provide more than one such option. If more than one such option is provided
-  for the same space, the last one takes effect. See `HLSL register and Vulkan
-  binding`_ for explanation and examples.
+  for the same space, the last one takes effect. If you need to shift the
+  inferred binding numbers for all sets, use ``all`` as ``M``.
+  See `HLSL register and Vulkan binding`_ for explanation and examples.
 - ``-fvk-t-shift N M``, similar to ``-fvk-b-shift``, but for t-type registers.
 - ``-fvk-s-shift N M``, similar to ``-fvk-b-shift``, but for s-type registers.
 - ``-fvk-u-shift N M``, similar to ``-fvk-b-shift``, but for u-type registers.

include/dxc/Support/HLSLOptions.h

@@ -167,10 +167,10 @@ public:
   bool VkUseDxLayout;                      // OPT_fvk_use_dx_layout
   bool SpvEnableReflect;                   // OPT_fspv_reflect
   llvm::StringRef VkStageIoOrder;          // OPT_fvk_stage_io_order
-  llvm::SmallVector<uint32_t, 4> VkBShift; // OPT_fvk_b_shift
-  llvm::SmallVector<uint32_t, 4> VkTShift; // OPT_fvk_t_shift
-  llvm::SmallVector<uint32_t, 4> VkSShift; // OPT_fvk_s_shift
-  llvm::SmallVector<uint32_t, 4> VkUShift; // OPT_fvk_u_shift
+  llvm::SmallVector<int32_t, 4> VkBShift;  // OPT_fvk_b_shift
+  llvm::SmallVector<int32_t, 4> VkTShift;  // OPT_fvk_t_shift
+  llvm::SmallVector<int32_t, 4> VkSShift;  // OPT_fvk_s_shift
+  llvm::SmallVector<int32_t, 4> VkUShift;  // OPT_fvk_u_shift
   llvm::SmallVector<llvm::StringRef, 4> SpvExtensions; // OPT_fspv_extension
   llvm::StringRef SpvTargetEnv;                        // OPT_fspv_target_env
 #endif

lib/DxcSupport/HLSLOptions.cpp

@@ -490,26 +490,43 @@ int ReadDxcOpts(const OptTable *optionTable, unsigned flagsToInclude,
   opts.VkIgnoreUnusedResources = Args.hasFlag(OPT_fvk_ignore_unused_resources, OPT_INVALID, false);
 
   // Collects the arguments for -fvk-{b|s|t|u}-shift.
-  const auto handleVkShiftArgs = [genSpirv, &Args, &errors](
-      OptSpecifier id, const char* name, llvm::SmallVectorImpl<uint32_t>* shifts) {
-    const auto values = Args.getAllArgValues(id);
-
-    if (!genSpirv && !values.empty()) {
-      errors << "-fvk-" << name << "-shift requires -spirv";
-      return false;
-    }
-
-    shifts->clear();
-    for (const auto& val : values) {
-      uint32_t number = 0;
-      if (llvm::StringRef(val).getAsInteger(10, number)) {
-        errors << "invalid -fvk-" << name << "-shift argument: " << val;
-        return false;
-      }
-      shifts->push_back(number);
-    }
-    return true;
-  };
+  const auto handleVkShiftArgs =
+      [genSpirv, &Args, &errors](OptSpecifier id, const char *name,
+                                 llvm::SmallVectorImpl<int32_t> *shifts) {
+        const auto values = Args.getAllArgValues(id);
+
+        if (!genSpirv && !values.empty()) {
+          errors << "-fvk-" << name << "-shift requires -spirv";
+          return false;
+        }
+
+        shifts->clear();
+        bool setForAll = false;
+
+        for (const auto &val : values) {
+          int32_t number = 0;
+          if (val == "all") {
+            number = -1;
+            setForAll = true;
+          } else {
+            if (llvm::StringRef(val).getAsInteger(10, number)) {
+              errors << "invalid -fvk-" << name << "-shift argument: " << val;
+              return false;
+            }
+            if (number < 0) {
+              errors << "negative -fvk-" << name << "-shift argument: " << val;
+              return false;
+            }
+          }
+          shifts->push_back(number);
+        }
+        if (setForAll && shifts->size() > 2) {
+          errors << "setting all sets via -fvk-" << name
+                 << "-shift argument should be used alone";
+          return false;
+        }
+        return true;
+      };
 
   if (!handleVkShiftArgs(OPT_fvk_b_shift, "b", &opts.VkBShift) ||
       !handleVkShiftArgs(OPT_fvk_t_shift, "t", &opts.VkTShift) ||

tools/clang/include/clang/SPIRV/EmitSPIRVOptions.h

@@ -40,10 +40,10 @@ struct EmitSPIRVOptions {
   bool enableReflect;
   bool enableDebugInfo;
   llvm::StringRef stageIoOrder;
-  llvm::SmallVector<uint32_t, 4> bShift;
-  llvm::SmallVector<uint32_t, 4> tShift;
-  llvm::SmallVector<uint32_t, 4> sShift;
-  llvm::SmallVector<uint32_t, 4> uShift;
+  llvm::SmallVector<int32_t, 4> bShift;
+  llvm::SmallVector<int32_t, 4> tShift;
+  llvm::SmallVector<int32_t, 4> sShift;
+  llvm::SmallVector<int32_t, 4> uShift;
   llvm::SmallVector<llvm::StringRef, 4> allowedExtensions;
   llvm::StringRef targetEnv;
   spirv::LayoutRule cBufferLayoutRule;

tools/clang/lib/SPIRV/DeclResultIdMapper.cpp

@@ -1015,15 +1015,19 @@ namespace {
 /// sets.
 class BindingShiftMapper {
 public:
-  explicit BindingShiftMapper(const llvm::SmallVectorImpl<uint32_t> &shifts)
+  explicit BindingShiftMapper(const llvm::SmallVectorImpl<int32_t> &shifts)
       : masterShift(0) {
     assert(shifts.size() % 2 == 0);
-    for (uint32_t i = 0; i < shifts.size(); i += 2)
-      perSetShift[shifts[i + 1]] = shifts[i];
+    if (shifts.size() == 2 && shifts[1] == -1) {
+      masterShift = shifts[0];
+    } else {
+      for (uint32_t i = 0; i < shifts.size(); i += 2)
+        perSetShift[shifts[i + 1]] = shifts[i];
+    }
   }
 
   /// Returns the shift amount for the given set.
-  uint32_t getShiftForSet(uint32_t set) const {
+  int32_t getShiftForSet(int32_t set) const {
     const auto found = perSetShift.find(set);
     if (found != perSetShift.end())
       return found->second;
@@ -1032,7 +1036,7 @@ public:
 
 private:
   uint32_t masterShift; /// Shift amount applies to all sets.
-  llvm::DenseMap<uint32_t, uint32_t> perSetShift;
+  llvm::DenseMap<int32_t, int32_t> perSetShift;
 };
 } // namespace
 

tools/clang/test/CodeGenSPIRV/vk.binding.cl.all-sets.hlsl

@@ -0,0 +1,53 @@
+// Run: %dxc -T ps_6_0 -E main -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.
+
+// CHECK: OpDecorate %cbuffer3 DescriptorSet 0
+// CHECK: OpDecorate %cbuffer3 Binding 42
+[[vk::binding(42)]]
+ConstantBuffer<S> cbuffer3 : register(b10, space2);
+
+// CHECK: OpDecorate %cbuffer1 DescriptorSet 0
+// 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 0
+// CHECK: OpDecorate %texture2 Binding 2001
+Texture2D<float4> texture2: register(t1);
+
+// CHECK: OpDecorate %sampler1 DescriptorSet 0
+// 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 0
+// CHECK: OpDecorate %rwbuffer2 Binding 4003
+RWBuffer<float4> rwbuffer2 : register(u3);
+
+// Lacking binding assignment is unaffacted.
+
+// CHECK: OpDecorate %cbuffer4 DescriptorSet 0
+// CHECK: OpDecorate %cbuffer4 Binding 0
+ConstantBuffer<S> cbuffer4;
+// CHECK: OpDecorate %cbuffer5 DescriptorSet 0
+// CHECK: OpDecorate %cbuffer5 Binding 1
+ConstantBuffer<S> cbuffer5;
+
+float4 main() : SV_Target {
+    return cbuffer1.f;
+}

tools/clang/unittests/SPIRV/CodeGenSPIRVTest.cpp

@@ -1321,10 +1321,15 @@ TEST_F(FileTest, VulkanRegisterBinding) {
   runFileTest("vk.binding.register.hlsl");
 }
 TEST_F(FileTest, VulkanRegisterBindingShift) {
-  // Resource binding from :register() and with shift specified via
+  // Resource binding from :register() with shift specified via
   // command line option
   runFileTest("vk.binding.cl.hlsl");
 }
+TEST_F(FileTest, VulkanRegisterBindingShiftAllSets) {
+  // Resource binding from :register() with shift specified for all sets via
+  // command line option
+  runFileTest("vk.binding.cl.all-sets.hlsl");
+}
 TEST_F(FileTest, VulkanStructuredBufferCounter) {
   // [[vk::counter_binding()]] for RWStructuredBuffer, AppendStructuredBuffer,
   // and ConsumeStructuredBuffer