Implement WebGPU->Vulkan initializer conversion for 'Function' variables (#2513)

WebGPU requires certain variables to be initialized, whereas there are
known issues with using initializers in Vulkan. This PR is the first
of three implementing a pass to decompose initialized variables into
a variable declaration followed by a store. This has been broken up
into multiple PRs, because there 3 distinct cases that need to be
handled, which require separate implementations.

This first PR implements the basic infrastructure that is needed, and
handling of Function storage class variables. Private and Output will
be handled in future PRs.

This is part of resolving #2388

Android.mk

@@ -92,6 +92,7 @@ SPVTOOLS_OPT_SRC_FILES := \
 		source/opt/dead_branch_elim_pass.cpp \
 		source/opt/dead_insert_elim_pass.cpp \
 		source/opt/dead_variable_elimination.cpp \
+		source/opt/decompose_initialized_variables_pass.cpp \
 		source/opt/decoration_manager.cpp \
 		source/opt/def_use_manager.cpp \
 		source/opt/dominator_analysis.cpp \

BUILD.gn

@@ -484,6 +484,8 @@ static_library("spvtools_opt") {
     "source/opt/dead_insert_elim_pass.h",
     "source/opt/dead_variable_elimination.cpp",
     "source/opt/dead_variable_elimination.h",
+    "source/opt/decompose_initialized_variables_pass.cpp",
+    "source/opt/decompose_initialized_variables_pass.h",
     "source/opt/decoration_manager.cpp",
     "source/opt/decoration_manager.h",
     "source/opt/def_use_manager.cpp",

include/spirv-tools/optimizer.hpp

@@ -768,6 +768,10 @@ Optimizer::PassToken CreateFixStorageClassPass();
 // converts those literals to 0.
 Optimizer::PassToken CreateLegalizeVectorShufflePass();
 
+// Create a pass to decompose initialized variables into a seperate variable
+// declaration and an initial store.
+Optimizer::PassToken CreateDecomposeInitializedVariablesPass();
+
 }  // namespace spvtools
 
 #endif  // INCLUDE_SPIRV_TOOLS_OPTIMIZER_HPP_

source/opt/CMakeLists.txt

@@ -31,6 +31,7 @@ set(SPIRV_TOOLS_OPT_SOURCES
   dead_branch_elim_pass.h
   dead_insert_elim_pass.h
   dead_variable_elimination.h
+  decompose_initialized_variables_pass.h
   decoration_manager.h
   def_use_manager.h
   dominator_analysis.h
@@ -130,6 +131,7 @@ set(SPIRV_TOOLS_OPT_SOURCES
   dead_branch_elim_pass.cpp
   dead_insert_elim_pass.cpp
   dead_variable_elimination.cpp
+  decompose_initialized_variables_pass.cpp
   decoration_manager.cpp
   def_use_manager.cpp
   dominator_analysis.cpp

source/opt/decompose_initialized_variables_pass.cpp

@@ -0,0 +1,77 @@
+// Copyright (c) 2019 Google LLC.
+//
+// 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/decompose_initialized_variables_pass.h"
+
+#include "source/opt/ir_context.h"
+
+namespace spvtools {
+namespace opt {
+
+using inst_iterator = InstructionList::iterator;
+
+namespace {
+
+bool HasInitializer(Instruction* inst) {
+  if (inst->opcode() != SpvOpVariable) return false;
+  if (inst->NumOperands() < 4) return false;
+
+  return true;
+}
+
+}  // namespace
+
+Pass::Status DecomposeInitializedVariablesPass::Process() {
+  auto* module = context()->module();
+  bool changed = false;
+
+  // TODO(zoddicus): Handle 'Output' variables
+  // TODO(zoddicus): Handle 'Private' variables
+
+  // Handle 'Function' variables
+  for (auto func = module->begin(); func != module->end(); ++func) {
+    auto block = func->entry().get();
+    std::vector<Instruction*> new_stores;
+
+    auto last_var = block->begin();
+    for (auto iter = block->begin();
+         iter != block->end() && iter->opcode() == SpvOpVariable; ++iter) {
+      last_var = iter;
+      Instruction* inst = &(*iter);
+      if (!HasInitializer(inst)) continue;
+
+      changed = true;
+      auto var_id = inst->result_id();
+      auto val_id = inst->GetOperand(3).words[0];
+      Instruction* store_inst = new Instruction(
+          context(), SpvOpStore, 0, 0,
+          {{SPV_OPERAND_TYPE_ID, {var_id}}, {SPV_OPERAND_TYPE_ID, {val_id}}});
+      new_stores.push_back(store_inst);
+      iter->RemoveOperand(3);
+      get_def_use_mgr()->UpdateDefUse(&*iter);
+    }
+
+    for (auto store = new_stores.begin(); store != new_stores.end(); ++store) {
+      context()->AnalyzeDefUse(*store);
+      context()->set_instr_block(*store, block);
+      (*store)->InsertAfter(&*last_var);
+      last_var = *store;
+    }
+  }
+
+  return changed ? Status::SuccessWithChange : Status::SuccessWithoutChange;
+}
+
+}  // namespace opt
+}  // namespace spvtools

source/opt/decompose_initialized_variables_pass.h

@@ -0,0 +1,57 @@
+// Copyright (c) 2019 Google LLC.
+//
+// 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.
+
+#ifndef SOURCE_OPT_DECOMPOSE_INITALIZED_VAIRABLES_PASS_H_
+#define SOURCE_OPT_DECOMPOSE_INITALIZED_VAIRABLES_PASS_H_
+
+#include "source/opt/ir_context.h"
+#include "source/opt/module.h"
+#include "source/opt/pass.h"
+
+namespace spvtools {
+namespace opt {
+
+// Converts variable declartions with initializers into seperate declaration and
+// assignment statements. This is done due to known issues with some Vulkan
+// implementations' handling of initialized variables.
+//
+// Only decomposes variables with storage classes that are valid in Vulkan
+// execution environments; Output, Private, and Function.
+// Currently only Function is implemented.
+class DecomposeInitializedVariablesPass : public Pass {
+ public:
+  const char* name() const override {
+    return "decompose-initialized-variables";
+  }
+  Status Process() override;
+
+  IRContext::Analysis GetPreservedAnalyses() override {
+    return IRContext::kAnalysisInstrToBlockMapping |
+           IRContext::kAnalysisDecorations | IRContext::kAnalysisCombinators |
+           IRContext::kAnalysisCFG | IRContext::kAnalysisDominatorAnalysis |
+           IRContext::kAnalysisLoopAnalysis | IRContext::kAnalysisNameMap |
+           IRContext::kAnalysisScalarEvolution |
+           IRContext::kAnalysisRegisterPressure |
+           IRContext::kAnalysisValueNumberTable |
+           IRContext::kAnalysisStructuredCFG |
+           IRContext::kAnalysisBuiltinVarId |
+           IRContext::kAnalysisIdToFuncMapping | IRContext::kAnalysisTypes |
+           IRContext::kAnalysisDefUse | IRContext::kAnalysisConstants;
+  }
+};
+
+}  // namespace opt
+}  // namespace spvtools
+
+#endif  // SOURCE_OPT_DECOMPOSE_INITALIZED_VAIRABLES_PASS_H_

source/opt/generate_webgpu_initializers_pass.cpp

@@ -1,4 +1,4 @@
-// Copyright (c) 2019 Google Inc.
+// Copyright (c) 2019 Google LLC.
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.

source/opt/generate_webgpu_initializers_pass.h

@@ -1,4 +1,4 @@
-// Copyright (c) 2019 Google Inc.
+// Copyright (c) 2019 Google LLC.
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.

source/opt/legalize_vector_shuffle_pass.cpp

@@ -1,4 +1,4 @@
-// Copyright (c) 2019 Google Inc.
+// Copyright (c) 2019 Google LLC.
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.

source/opt/legalize_vector_shuffle_pass.h

@@ -1,4 +1,4 @@
-// Copyright (c) 2019 Google Inc.
+// Copyright (c) 2019 Google LLC.
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.

source/opt/optimizer.cpp

@@ -231,7 +231,9 @@ Optimizer& Optimizer::RegisterVulkanToWebGPUPasses() {
       .RegisterPass(CreateDeadBranchElimPass());
 }
 
-Optimizer& Optimizer::RegisterWebGPUToVulkanPasses() { return *this; }
+Optimizer& Optimizer::RegisterWebGPUToVulkanPasses() {
+  return RegisterPass(CreateDecomposeInitializedVariablesPass());
+}
 
 bool Optimizer::RegisterPassesFromFlags(const std::vector<std::string>& flags) {
   for (const auto& flag : flags) {
@@ -866,4 +868,9 @@ Optimizer::PassToken CreateLegalizeVectorShufflePass() {
       MakeUnique<opt::LegalizeVectorShufflePass>());
 }
 
+Optimizer::PassToken CreateDecomposeInitializedVariablesPass() {
+  return MakeUnique<Optimizer::PassToken::Impl>(
+      MakeUnique<opt::DecomposeInitializedVariablesPass>());
+}
+
 }  // namespace spvtools

source/opt/passes.h

@@ -29,6 +29,7 @@
 #include "source/opt/dead_branch_elim_pass.h"
 #include "source/opt/dead_insert_elim_pass.h"
 #include "source/opt/dead_variable_elimination.h"
+#include "source/opt/decompose_initialized_variables_pass.h"
 #include "source/opt/eliminate_dead_constant_pass.h"
 #include "source/opt/eliminate_dead_functions_pass.h"
 #include "source/opt/eliminate_dead_members_pass.h"

source/opt/strip_atomic_counter_memory_pass.cpp

@@ -1,4 +1,4 @@
-// Copyright (c) 2019 Google Inc.
+// Copyright (c) 2019 Google LLC.
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.

source/opt/strip_atomic_counter_memory_pass.h

@@ -1,4 +1,4 @@
-// Copyright (c) 2019 Google Inc.
+// Copyright (c) 2019 Google LLC.
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.

source/reduce/remove_selection_reduction_opportunity_finder.cpp

@@ -1,4 +1,4 @@
-// Copyright (c) 2019 Google Inc.
+// Copyright (c) 2019 Google LLC.
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.

test/opt/CMakeLists.txt

@@ -30,6 +30,7 @@ add_spvtools_unittest(TARGET opt
        dead_branch_elim_test.cpp
        dead_insert_elim_test.cpp
        dead_variable_elim_test.cpp
+       decompose_initialized_variables_test.cpp
        decoration_manager_test.cpp
        def_use_test.cpp
        eliminate_dead_const_test.cpp

test/opt/decompose_initialized_variables_test.cpp

@@ -0,0 +1,85 @@
+// Copyright (c) 2019 Google LLC.
+//
+// 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 <vector>
+
+#include "test/opt/pass_fixture.h"
+#include "test/opt/pass_utils.h"
+
+namespace spvtools {
+namespace opt {
+namespace {
+
+using DecomposeInitializedVariablesTest = PassTest<::testing::Test>;
+
+void operator+=(std::vector<const char*>& lhs,
+                const std::vector<const char*> rhs) {
+  for (auto elem : rhs) lhs.push_back(elem);
+}
+
+std::vector<const char*> header = {
+    "OpCapability Shader",
+    "OpCapability VulkanMemoryModelKHR",
+    "OpExtension \"SPV_KHR_vulkan_memory_model\"",
+    "OpMemoryModel Logical VulkanKHR",
+    "OpEntryPoint Vertex %1 \"shader\"",
+    "%uint = OpTypeInt 32 0",
+    "%uint_1 = OpConstant %uint 1",
+    "%4 = OpConstantNull %uint",
+    "%void = OpTypeVoid",
+    "%6 = OpTypeFunction %void"};
+
+std::string GetFunctionTest(std::vector<const char*> body) {
+  auto result = header;
+  result += {"%_ptr_Function_uint = OpTypePointer Function %uint",
+             "%1 = OpFunction %void None %6", "%8 = OpLabel"};
+  result += body;
+  result += {"OpReturn", "OpFunctionEnd"};
+  return JoinAllInsts(result);
+}
+
+TEST_F(DecomposeInitializedVariablesTest, FunctionChanged) {
+  std::string input =
+      GetFunctionTest({"%9 = OpVariable %_ptr_Function_uint Function %uint_1"});
+  std::string expected = GetFunctionTest(
+      {"%9 = OpVariable %_ptr_Function_uint Function", "OpStore %9 %uint_1"});
+
+  SinglePassRunAndCheck<DecomposeInitializedVariablesPass>(
+      input, expected, /* skip_nop = */ false);
+}
+
+TEST_F(DecomposeInitializedVariablesTest, FunctionUnchanged) {
+  std::string input =
+      GetFunctionTest({"%9 = OpVariable %_ptr_Function_uint Function"});
+
+  SinglePassRunAndCheck<DecomposeInitializedVariablesPass>(
+      input, input, /* skip_nop = */ false);
+}
+
+TEST_F(DecomposeInitializedVariablesTest, FunctionMultipleVariables) {
+  std::string input =
+      GetFunctionTest({"%9 = OpVariable %_ptr_Function_uint Function %uint_1",
+                       "%10 = OpVariable %_ptr_Function_uint Function %4"});
+  std::string expected =
+      GetFunctionTest({"%9 = OpVariable %_ptr_Function_uint Function",
+                       "%10 = OpVariable %_ptr_Function_uint Function",
+                       "OpStore %9 %uint_1", "OpStore %10 %4"});
+
+  SinglePassRunAndCheck<DecomposeInitializedVariablesPass>(
+      input, expected, /* skip_nop = */ false);
+}
+
+}  // namespace
+}  // namespace opt
+}  // namespace spvtools

test/opt/generate_webgpu_initializers_test.cpp

@@ -1,4 +1,4 @@
-// Copyright (c) 2019 Google Inc.
+// Copyright (c) 2019 Google LLC.
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.

test/opt/legalize_vector_shuffle_test.cpp

@@ -1,4 +1,4 @@
-// Copyright (c) 2019 Google Inc.
+// Copyright (c) 2019 Google LLC.
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.

test/opt/optimizer_test.cpp

@@ -222,7 +222,7 @@ TEST(Optimizer, CanRegisterPassesFromFlags) {
   EXPECT_EQ(msg_level, SPV_MSG_ERROR);
 }
 
-TEST(Optimizer, VulkanToWebGPUModeSetsCorrectPasses) {
+TEST(Optimizer, VulkanToWebGPUSetsCorrectPasses) {
   Optimizer opt(SPV_ENV_WEBGPU_0);
   opt.RegisterVulkanToWebGPUPasses();
   std::vector<const char*> pass_names = opt.GetPassNames();
@@ -260,9 +260,11 @@ using VulkanToWebGPUPassTest =
     PassTest<::testing::TestWithParam<VulkanToWebGPUPassCase>>;
 
 TEST_P(VulkanToWebGPUPassTest, Ran) {
-  SpirvTools tools(SPV_ENV_WEBGPU_0);
   std::vector<uint32_t> binary;
-  tools.Assemble(GetParam().input, &binary);
+  {
+    SpirvTools tools(SPV_ENV_VULKAN_1_1);
+    tools.Assemble(GetParam().input, &binary);
+  }
 
   Optimizer opt(SPV_ENV_WEBGPU_0);
   opt.RegisterVulkanToWebGPUPasses();
@@ -272,7 +274,10 @@ TEST_P(VulkanToWebGPUPassTest, Ran) {
   ASSERT_TRUE(opt.Run(binary.data(), binary.size(), &optimized,
                       validator_options, true));
   std::string disassembly;
-  tools.Disassemble(optimized.data(), optimized.size(), &disassembly);
+  {
+    SpirvTools tools(SPV_ENV_WEBGPU_0);
+    tools.Disassemble(optimized.data(), optimized.size(), &disassembly);
+  }
 
   EXPECT_EQ(GetParam().expected, disassembly)
       << "Was expecting pass '" << GetParam().pass << "' to have been run.\n";
@@ -521,6 +526,101 @@ INSTANTIATE_TEST_SUITE_P(
          // pass
          "legalize-vector-shuffle"}}));
 
+TEST(Optimizer, WebGPUToVulkanSetsCorrectPasses) {
+  Optimizer opt(SPV_ENV_VULKAN_1_1);
+  opt.RegisterWebGPUToVulkanPasses();
+  std::vector<const char*> pass_names = opt.GetPassNames();
+
+  std::vector<std::string> registered_passes;
+  for (auto name = pass_names.begin(); name != pass_names.end(); ++name)
+    registered_passes.push_back(*name);
+
+  std::vector<std::string> expected_passes = {
+      "decompose-initialized-variables"};
+  std::sort(registered_passes.begin(), registered_passes.end());
+  std::sort(expected_passes.begin(), expected_passes.end());
+
+  ASSERT_EQ(registered_passes.size(), expected_passes.size());
+  for (size_t i = 0; i < registered_passes.size(); i++)
+    EXPECT_EQ(registered_passes[i], expected_passes[i]);
+}
+
+struct WebGPUToVulkanPassCase {
+  // Input SPIR-V
+  std::string input;
+  // Expected result SPIR-V
+  std::string expected;
+  // Specific pass under test, used for logging messages.
+  std::string pass;
+};
+
+using WebGPUToVulkanPassTest =
+    PassTest<::testing::TestWithParam<WebGPUToVulkanPassCase>>;
+
+TEST_P(WebGPUToVulkanPassTest, Ran) {
+  std::vector<uint32_t> binary;
+  {
+    SpirvTools tools(SPV_ENV_WEBGPU_0);
+    tools.Assemble(GetParam().input, &binary);
+  }
+
+  Optimizer opt(SPV_ENV_VULKAN_1_1);
+  opt.RegisterWebGPUToVulkanPasses();
+
+  std::vector<uint32_t> optimized;
+  class ValidatorOptions validator_options;
+  ASSERT_TRUE(opt.Run(binary.data(), binary.size(), &optimized,
+                      validator_options, true));
+  std::string disassembly;
+  {
+    SpirvTools tools(SPV_ENV_VULKAN_1_1);
+    tools.Disassemble(optimized.data(), optimized.size(), &disassembly);
+  }
+
+  EXPECT_EQ(GetParam().expected, disassembly)
+      << "Was expecting pass '" << GetParam().pass << "' to have been run.\n";
+}
+
+INSTANTIATE_TEST_SUITE_P(
+    Optimizer, WebGPUToVulkanPassTest,
+    ::testing::ValuesIn(std::vector<WebGPUToVulkanPassCase>{
+        // Decompose Initialized Variables
+        {// input
+         "OpCapability Shader\n"
+         "OpCapability VulkanMemoryModelKHR\n"
+         "OpExtension \"SPV_KHR_vulkan_memory_model\"\n"
+         "OpMemoryModel Logical VulkanKHR\n"
+         "OpEntryPoint Vertex %1 \"shader\"\n"
+         "%uint = OpTypeInt 32 0\n"
+         "%_ptr_Function_uint = OpTypePointer Function %uint\n"
+         "%4 = OpConstantNull %uint\n"
+         "%void = OpTypeVoid\n"
+         "%6 = OpTypeFunction %void\n"
+         "%1 = OpFunction %void None %6\n"
+         "%7 = OpLabel\n"
+         "%8 = OpVariable %_ptr_Function_uint Function %4\n"
+         "OpReturn\n"
+         "OpFunctionEnd\n",
+         // expected
+         "OpCapability Shader\n"
+         "OpCapability VulkanMemoryModelKHR\n"
+         "OpExtension \"SPV_KHR_vulkan_memory_model\"\n"
+         "OpMemoryModel Logical VulkanKHR\n"
+         "OpEntryPoint Vertex %1 \"shader\"\n"
+         "%uint = OpTypeInt 32 0\n"
+         "%_ptr_Function_uint = OpTypePointer Function %uint\n"
+         "%4 = OpConstantNull %uint\n"
+         "%void = OpTypeVoid\n"
+         "%6 = OpTypeFunction %void\n"
+         "%1 = OpFunction %void None %6\n"
+         "%7 = OpLabel\n"
+         "%8 = OpVariable %_ptr_Function_uint Function\n"
+         "OpStore %8 %4\n"
+         "OpReturn\n"
+         "OpFunctionEnd\n",
+         // pass
+         "decompose-initialized-variables"}}));
+
 }  // namespace
 }  // namespace opt
 }  // namespace spvtools

test/opt/strip_atomic_counter_memory_test.cpp

@@ -1,4 +1,4 @@
-// Copyright (c) 2019 Google Inc.
+// Copyright (c) 2019 Google LLC.
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.