opt: Add SwitchDescriptorSetPass (#5375)

This is a simple pass to change DescriptorSet decoration values.

Android.mk

@@ -182,6 +182,7 @@ SPVTOOLS_OPT_SRC_FILES := \
 		source/opt/strip_debug_info_pass.cpp \
 		source/opt/strip_nonsemantic_info_pass.cpp \
 		source/opt/struct_cfg_analysis.cpp \
+		source/opt/switch_descriptorset_pass.cpp \
 		source/opt/trim_capabilities_pass.cpp \
 		source/opt/type_manager.cpp \
 		source/opt/types.cpp \

BUILD.gn

@@ -789,6 +789,8 @@ static_library("spvtools_opt") {
     "source/opt/strip_nonsemantic_info_pass.h",
     "source/opt/struct_cfg_analysis.cpp",
     "source/opt/struct_cfg_analysis.h",
+    "source/opt/switch_descriptorset_pass.cpp",
+    "source/opt/switch_descriptorset_pass.h",
     "source/opt/tree_iterator.h",
     "source/opt/trim_capabilities_pass.cpp",
     "source/opt/trim_capabilities_pass.h",

include/spirv-tools/optimizer.hpp

@@ -992,6 +992,11 @@ Optimizer::PassToken CreateFixFuncCallArgumentsPass();
 // the unknown capability interacts with one of the trimmed capabilities.
 Optimizer::PassToken CreateTrimCapabilitiesPass();
 
+// Creates a switch-descriptorset pass.
+// This pass changes any DescriptorSet decorations with the value |ds_from| to
+// use the new value |ds_to|.
+Optimizer::PassToken CreateSwitchDescriptorSetPass(uint32_t ds_from,
+                                                   uint32_t ds_to);
 }  // namespace spvtools
 
 #endif  // INCLUDE_SPIRV_TOOLS_OPTIMIZER_HPP_

source/opt/CMakeLists.txt

@@ -121,6 +121,7 @@ set(SPIRV_TOOLS_OPT_SOURCES
   strip_debug_info_pass.h
   strip_nonsemantic_info_pass.h
   struct_cfg_analysis.h
+  switch_descriptorset_pass.h
   tree_iterator.h
   trim_capabilities_pass.h
   type_manager.h
@@ -237,6 +238,7 @@ set(SPIRV_TOOLS_OPT_SOURCES
   strip_debug_info_pass.cpp
   strip_nonsemantic_info_pass.cpp
   struct_cfg_analysis.cpp
+  switch_descriptorset_pass.cpp
   trim_capabilities_pass.cpp
   type_manager.cpp
   types.cpp

source/opt/optimizer.cpp

@@ -15,6 +15,7 @@
 #include "spirv-tools/optimizer.hpp"
 
 #include <cassert>
+#include <charconv>
 #include <memory>
 #include <string>
 #include <unordered_map>
@@ -549,6 +550,39 @@ bool Optimizer::RegisterPassFromFlag(const std::string& flag) {
              pass_args.c_str());
       return false;
     }
+  } else if (pass_name == "switch-descriptorset") {
+    if (pass_args.size() == 0) {
+      Error(consumer(), nullptr, {},
+            "--switch-descriptorset requires a from:to argument.");
+      return false;
+    }
+    uint32_t from_set, to_set;
+    const char* start = pass_args.data();
+    const char* end = pass_args.data() + pass_args.size();
+
+    auto result = std::from_chars(start, end, from_set);
+    if (result.ec != std::errc()) {
+      Errorf(consumer(), nullptr, {},
+             "Invalid argument for --switch-descriptorset: %s",
+             pass_args.c_str());
+      return false;
+    }
+    start = result.ptr;
+    if (start[0] != ':') {
+      Errorf(consumer(), nullptr, {},
+             "Invalid argument for --switch-descriptorset: %s",
+             pass_args.c_str());
+      return false;
+    }
+    start++;
+    result = std::from_chars(start, end, to_set);
+    if (result.ec != std::errc() || result.ptr != end) {
+      Errorf(consumer(), nullptr, {},
+             "Invalid argument for --switch-descriptorset: %s",
+             pass_args.c_str());
+      return false;
+    }
+    RegisterPass(CreateSwitchDescriptorSetPass(from_set, to_set));
   } else {
     Errorf(consumer(), nullptr, {},
            "Unknown flag '--%s'. Use --help for a list of valid flags",
@@ -1076,6 +1110,11 @@ Optimizer::PassToken CreateTrimCapabilitiesPass() {
   return MakeUnique<Optimizer::PassToken::Impl>(
       MakeUnique<opt::TrimCapabilitiesPass>());
 }
+
+Optimizer::PassToken CreateSwitchDescriptorSetPass(uint32_t from, uint32_t to) {
+  return MakeUnique<Optimizer::PassToken::Impl>(
+      MakeUnique<opt::SwitchDescriptorSetPass>(from, to));
+}
 }  // namespace spvtools
 
 extern "C" {

source/opt/passes.h

@@ -82,6 +82,7 @@
 #include "source/opt/strength_reduction_pass.h"
 #include "source/opt/strip_debug_info_pass.h"
 #include "source/opt/strip_nonsemantic_info_pass.h"
+#include "source/opt/switch_descriptorset_pass.h"
 #include "source/opt/trim_capabilities_pass.h"
 #include "source/opt/unify_const_pass.h"
 #include "source/opt/upgrade_memory_model.h"

source/opt/switch_descriptorset_pass.cpp

@@ -0,0 +1,46 @@
+// Copyright (c) 2023 LunarG Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "source/opt/switch_descriptorset_pass.h"
+
+#include "source/opt/ir_builder.h"
+#include "source/util/string_utils.h"
+
+namespace spvtools {
+namespace opt {
+
+Pass::Status SwitchDescriptorSetPass::Process() {
+  Status status = Status::SuccessWithoutChange;
+  auto* deco_mgr = context()->get_decoration_mgr();
+
+  for (Instruction& var : context()->types_values()) {
+    if (var.opcode() != spv::Op::OpVariable) {
+      continue;
+    }
+    auto decos = deco_mgr->GetDecorationsFor(var.result_id(), false);
+    for (const auto& deco : decos) {
+      spv::Decoration d = spv::Decoration(deco->GetSingleWordInOperand(1u));
+      if (d == spv::Decoration::DescriptorSet &&
+          deco->GetSingleWordInOperand(2u) == ds_from_) {
+        deco->SetInOperand(2u, {ds_to_});
+        status = Status::SuccessWithChange;
+        break;
+      }
+    }
+  }
+  return status;
+}
+
+}  // namespace opt
+}  // namespace spvtools

source/opt/switch_descriptorset_pass.h

@@ -0,0 +1,52 @@
+// Copyright (c) 2023 LunarG Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include <cstdio>
+#include <memory>
+#include <queue>
+#include <unordered_map>
+#include <unordered_set>
+#include <vector>
+
+#include "source/opt/pass.h"
+
+namespace spvtools {
+namespace opt {
+
+// See optimizer.hpp for documentation.
+class SwitchDescriptorSetPass : public Pass {
+ public:
+  SwitchDescriptorSetPass(uint32_t ds_from, uint32_t ds_to)
+      : ds_from_(ds_from), ds_to_(ds_to) {}
+
+  const char* name() const override { return "switch-descriptorset"; }
+
+  Status Process() override;
+
+  IRContext::Analysis GetPreservedAnalyses() override {
+    // this pass preserves everything except decorations
+    uint32_t mask = ((IRContext::kAnalysisEnd << 1) - 1);
+    mask &= ~static_cast<uint32_t>(IRContext::kAnalysisDecorations);
+    return static_cast<IRContext::Analysis>(mask);
+  }
+
+ private:
+  uint32_t ds_from_;
+  uint32_t ds_to_;
+};
+
+}  // namespace opt
+}  // namespace spvtools

test/opt/CMakeLists.txt

@@ -103,6 +103,7 @@ add_spvtools_unittest(TARGET opt
        strip_debug_info_test.cpp
        strip_nonsemantic_info_test.cpp
        struct_cfg_analysis_test.cpp
+       switch_descriptorset_test.cpp
        trim_capabilities_pass_test.cpp
        type_manager_test.cpp
        types_test.cpp

test/opt/switch_descriptorset_test.cpp

@@ -0,0 +1,193 @@
+// Copyright (c) 2023 LunarG Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// Bindless Check Instrumentation Tests.
+// Tests ending with V2 use version 2 record format.
+
+#include <string>
+#include <vector>
+
+#include "test/opt/pass_fixture.h"
+#include "test/opt/pass_utils.h"
+
+namespace spvtools {
+namespace opt {
+namespace {
+
+using SwitchDescriptorSetTest = PassTest<::testing::Test>;
+
+TEST_F(SwitchDescriptorSetTest, Basic) {
+  // #version 450
+  // #extension GL_EXT_buffer_reference : enable
+  //
+  // layout(buffer_reference, buffer_reference_align = 16) buffer bufStruct;
+  //
+  // layout(set = 7, binding = 7) uniform ufoo {
+  //     bufStruct data;
+  //     uint offset;
+  // } u_info;
+  //
+  // layout(buffer_reference, std140) buffer bufStruct {
+  //     layout(offset = 0) int a[2];
+  //     layout(offset = 32) int b;
+  // };
+  //
+  // void main() {
+  //     u_info.data.b = 0xca7;
+  // }
+
+  const std::string spirv = R"(
+OpCapability Shader
+OpCapability PhysicalStorageBufferAddresses
+OpExtension "SPV_EXT_physical_storage_buffer"
+%1 = OpExtInstImport "GLSL.std.450"
+OpMemoryModel PhysicalStorageBuffer64 GLSL450
+OpEntryPoint GLCompute %main "main"
+OpExecutionMode %main LocalSize 1 1 1
+OpSource GLSL 450
+OpSourceExtension "GL_EXT_buffer_reference"
+OpName %main "main"
+OpName %ufoo "ufoo"
+OpMemberName %ufoo 0 "data"
+OpMemberName %ufoo 1 "offset"
+OpName %bufStruct "bufStruct"
+OpMemberName %bufStruct 0 "a"
+OpMemberName %bufStruct 1 "b"
+OpName %u_info "u_info"
+OpMemberDecorate %ufoo 0 Offset 0
+OpMemberDecorate %ufoo 1 Offset 8
+OpDecorate %ufoo Block
+OpDecorate %_arr_int_uint_2 ArrayStride 16
+OpMemberDecorate %bufStruct 0 Offset 0
+OpMemberDecorate %bufStruct 1 Offset 32
+OpDecorate %bufStruct Block
+OpDecorate %u_info DescriptorSet 7
+;CHECK: OpDecorate %u_info DescriptorSet 31
+OpDecorate %u_info Binding 7
+;CHECK: OpDecorate %u_info Binding 7
+%void = OpTypeVoid
+%3 = OpTypeFunction %void
+OpTypeForwardPointer %_ptr_PhysicalStorageBuffer_bufStruct PhysicalStorageBuffer
+%uint = OpTypeInt 32 0
+%ufoo = OpTypeStruct %_ptr_PhysicalStorageBuffer_bufStruct %uint
+%int = OpTypeInt 32 1
+%uint_2 = OpConstant %uint 2
+%_arr_int_uint_2 = OpTypeArray %int %uint_2
+%bufStruct = OpTypeStruct %_arr_int_uint_2 %int
+%_ptr_PhysicalStorageBuffer_bufStruct = OpTypePointer PhysicalStorageBuffer %bufStruct
+%_ptr_Uniform_ufoo = OpTypePointer Uniform %ufoo
+%u_info = OpVariable %_ptr_Uniform_ufoo Uniform
+%int_0 = OpConstant %int 0
+%_ptr_Uniform__ptr_PhysicalStorageBuffer_bufStruct = OpTypePointer Uniform %_ptr_PhysicalStorageBuffer_bufStruct
+%int_1 = OpConstant %int 1
+%int_3239 = OpConstant %int 3239
+%_ptr_PhysicalStorageBuffer_int = OpTypePointer PhysicalStorageBuffer %int
+%main = OpFunction %void None %3
+%5 = OpLabel
+%17 = OpAccessChain %_ptr_Uniform__ptr_PhysicalStorageBuffer_bufStruct %u_info %int_0
+%18 = OpLoad %_ptr_PhysicalStorageBuffer_bufStruct %17
+%22 = OpAccessChain %_ptr_PhysicalStorageBuffer_int %18 %int_1
+OpReturn
+OpFunctionEnd
+)";
+  // clang-format off
+
+  SinglePassRunAndMatch<SwitchDescriptorSetPass>(spirv, true, 7, 31);
+}
+
+
+// Make sure DescriptorSet decorations that don't match the requested number
+// are left unchanged.
+TEST_F(SwitchDescriptorSetTest, Unchanged) {
+  // #version 450
+  // #extension GL_EXT_buffer_reference : enable
+  //
+  // layout(buffer_reference, buffer_reference_align = 16) buffer bufStruct;
+  //
+  // layout(set = 11, binding = 7) uniform ufoo {
+  //     bufStruct data;
+  //     uint offset;
+  // } u_info;
+  //
+  // layout(buffer_reference, std140) buffer bufStruct {
+  //     layout(offset = 0) int a[2];
+  //     layout(offset = 32) int b;
+  // };
+  //
+  // void main() {
+  //     u_info.data.b = 0xca7;
+  // }
+
+  const std::string spirv = R"(
+OpCapability Shader
+OpCapability PhysicalStorageBufferAddresses
+OpExtension "SPV_EXT_physical_storage_buffer"
+%1 = OpExtInstImport "GLSL.std.450"
+OpMemoryModel PhysicalStorageBuffer64 GLSL450
+OpEntryPoint GLCompute %main "main"
+OpExecutionMode %main LocalSize 1 1 1
+OpSource GLSL 450
+OpSourceExtension "GL_EXT_buffer_reference"
+OpName %main "main"
+OpName %ufoo "ufoo"
+OpMemberName %ufoo 0 "data"
+OpMemberName %ufoo 1 "offset"
+OpName %bufStruct "bufStruct"
+OpMemberName %bufStruct 0 "a"
+OpMemberName %bufStruct 1 "b"
+OpName %u_info "u_info"
+OpMemberDecorate %ufoo 0 Offset 0
+OpMemberDecorate %ufoo 1 Offset 8
+OpDecorate %ufoo Block
+OpDecorate %_arr_int_uint_2 ArrayStride 16
+OpMemberDecorate %bufStruct 0 Offset 0
+OpMemberDecorate %bufStruct 1 Offset 32
+OpDecorate %bufStruct Block
+OpDecorate %u_info DescriptorSet 11
+;CHECK: OpDecorate %u_info DescriptorSet 11
+OpDecorate %u_info Binding 7
+;CHECK: OpDecorate %u_info Binding 7
+%void = OpTypeVoid
+%3 = OpTypeFunction %void
+OpTypeForwardPointer %_ptr_PhysicalStorageBuffer_bufStruct PhysicalStorageBuffer
+%uint = OpTypeInt 32 0
+%ufoo = OpTypeStruct %_ptr_PhysicalStorageBuffer_bufStruct %uint
+%int = OpTypeInt 32 1
+%uint_2 = OpConstant %uint 2
+%_arr_int_uint_2 = OpTypeArray %int %uint_2
+%bufStruct = OpTypeStruct %_arr_int_uint_2 %int
+%_ptr_PhysicalStorageBuffer_bufStruct = OpTypePointer PhysicalStorageBuffer %bufStruct
+%_ptr_Uniform_ufoo = OpTypePointer Uniform %ufoo
+%u_info = OpVariable %_ptr_Uniform_ufoo Uniform
+%int_0 = OpConstant %int 0
+%_ptr_Uniform__ptr_PhysicalStorageBuffer_bufStruct = OpTypePointer Uniform %_ptr_PhysicalStorageBuffer_bufStruct
+%int_1 = OpConstant %int 1
+%int_3239 = OpConstant %int 3239
+%_ptr_PhysicalStorageBuffer_int = OpTypePointer PhysicalStorageBuffer %int
+%main = OpFunction %void None %3
+%5 = OpLabel
+%17 = OpAccessChain %_ptr_Uniform__ptr_PhysicalStorageBuffer_bufStruct %u_info %int_0
+%18 = OpLoad %_ptr_PhysicalStorageBuffer_bufStruct %17
+%22 = OpAccessChain %_ptr_PhysicalStorageBuffer_int %18 %int_1
+OpReturn
+OpFunctionEnd
+)";
+  // clang-format off
+
+  SinglePassRunAndMatch<SwitchDescriptorSetPass>(spirv, true, 7, 31);
+}
+
+}  // namespace
+}  // namespace opt
+}  // namespace spvtools

tools/opt/opt.cpp

@@ -496,6 +496,10 @@ Options (in lexicographical order):)",
                covers reflection information defined by
                SPV_GOOGLE_hlsl_functionality1 and SPV_KHR_non_semantic_info)");
   printf(R"(
+  --switch-descriptorset=<from>:<to>
+               Switch any DescriptoSet decorations using the value <from> to
+               the new value <to>.)");
+  printf(R"(
   --target-env=<env>
                Set the target environment. Without this flag the target
                environment defaults to spv1.5. <env> must be one of