[vk] Move substitute overrides to backend raise.

Move the call to SubstituteOverrides out of the ShaderModuleVK and into
the SPIR-V Raise routine.

Change-Id: Ib05ac14aa48b757081c07b146d488e9842aeb846
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/268036
Reviewed-by: James Price <jrprice@google.com>
Commit-Queue: dan sinclair <dsinclair@chromium.org>
diff --git a/src/dawn/native/vulkan/ShaderModuleVk.cpp b/src/dawn/native/vulkan/ShaderModuleVk.cpp
index e2843e6..81bf134 100644
--- a/src/dawn/native/vulkan/ShaderModuleVk.cpp
+++ b/src/dawn/native/vulkan/ShaderModuleVk.cpp
@@ -103,13 +103,10 @@
 
 #if TINT_BUILD_SPV_WRITER
 
-using SubstituteOverrideConfig = std::unordered_map<tint::OverrideId, double>;
-
 #define SPIRV_COMPILATION_REQUEST_MEMBERS(X)                                         \
     X(SingleShaderStage, stage)                                                      \
     X(ShaderModuleBase::ShaderModuleHash, shaderModuleHash)                          \
     X(UnsafeUnserializedValue<ShaderModuleBase::ScopedUseTintProgram>, inputProgram) \
-    X(SubstituteOverrideConfig, substituteOverrideConfig)                            \
     X(LimitsForCompilationRequest, limits)                                           \
     X(UnsafeUnserializedValue<LimitsForCompilationRequest>, adapterSupportedLimits)  \
     X(uint32_t, maxSubgroupSize)                                                     \
@@ -202,7 +199,6 @@
     req.inputProgram = UnsafeUnserializedValue(UseTintProgram());
     req.entryPointName = programmableStage.entryPoint;
     req.platform = UnsafeUnserializedValue(GetDevice()->GetPlatform());
-    req.substituteOverrideConfig = BuildSubstituteOverridesTransformConfig(programmableStage);
     req.usesSubgroupMatrix = programmableStage.metadata->usesSubgroupMatrix;
 
     req.tintOptions.remapped_entry_point_name = GetDevice()->GetIsolatedEntryPointName();
@@ -213,6 +209,9 @@
     req.tintOptions.disable_robustness = !GetDevice()->IsRobustnessEnabled();
     req.tintOptions.emit_vertex_point_size = emitPointSize;
 
+    req.tintOptions.substitute_overrides_config = {
+        .map = BuildSubstituteOverridesTransformConfig(programmableStage),
+    };
     req.tintOptions.disable_workgroup_init =
         GetDevice()->IsToggleEnabled(Toggle::DisableWorkgroupInit);
     // The only possible alternative for the vulkan demote to helper extension is
@@ -308,20 +307,6 @@
                                 singleEntryPointResult.Failure().reason);
             }
 
-            {
-                SCOPED_DAWN_HISTOGRAM_TIMER_MICROS(r.platform.UnsafeGetValue(),
-                                                   "ShaderModuleSubstituteOverrides");
-                // this needs to run after SingleEntryPoint transform which removes unused
-                // overrides for the current entry point.
-                tint::SubstituteOverridesConfig cfg;
-                cfg.map = std::move(r.substituteOverrideConfig);
-                auto substituteOverridesResult =
-                    tint::core::ir::transform::SubstituteOverrides(ir.Get(), cfg);
-                DAWN_INVALID_IF(substituteOverridesResult != tint::Success,
-                                "Pipeline override substitution (IR) failed:\n%s",
-                                substituteOverridesResult.Failure().reason);
-            }
-
             tint::Result<tint::spirv::writer::Output> tintResult;
             {
                 SCOPED_DAWN_HISTOGRAM_TIMER_MICROS(r.platform.UnsafeGetValue(),
diff --git a/src/tint/cmd/tint/main.cc b/src/tint/cmd/tint/main.cc
index cd2e6c3..d8fbd18 100644
--- a/src/tint/cmd/tint/main.cc
+++ b/src/tint/cmd/tint/main.cc
@@ -857,6 +857,15 @@
 
     auto entry_point = inspector.GetEntryPoint(options.ep_name);
 
+    // Run SubstituteOverrides to replace override instructions with constants.
+    // This needs to run after SingleEntryPoint which removes unused overrides.
+    auto substitute_override_cfg = CreateOverrideMap(options, inspector);
+    if (substitute_override_cfg != tint::Success) {
+        std::cerr << "Failed to create override map: " << substitute_override_cfg.Failure() << "\n";
+        return false;
+    }
+    gen_options.substitute_overrides_config = substitute_override_cfg.Get();
+
     // Immediate data Offset must be 4-byte aligned.
     uint32_t offset = tint::RoundUp(4u, entry_point.immediate_data_size);
 
@@ -878,7 +887,7 @@
     }
 
     // Check that the module and options are supported by the backend.
-    auto check = tint::spirv::writer::CanGenerate(ir, gen_options);
+    auto check = tint::spirv::writer::CanGenerate(ir, gen_options, options.ep_name);
     if (check != tint::Success) {
         std::cerr << check.Failure() << "\n";
         return false;
@@ -1358,6 +1367,14 @@
         }
     }
 
+    switch (options.format) {
+        case Format::kSpirv:
+        case Format::kSpvAsm:
+            return GenerateSpirv(options, inspector, ir.Get());
+        default:
+            break;
+    }
+
     // Run SubstituteOverrides to replace override instructions with constants.
     // This needs to run after SingleEntryPoint which removes unused overrides.
     auto substitute_override_cfg = CreateOverrideMap(options, inspector);
@@ -1373,9 +1390,6 @@
     }
 
     switch (options.format) {
-        case Format::kSpirv:
-        case Format::kSpvAsm:
-            return GenerateSpirv(options, inspector, ir.Get());
         case Format::kMsl:
             return GenerateMsl(options, inspector, ir.Get());
         case Format::kHlsl:
diff --git a/src/tint/lang/spirv/writer/common/helper_test.h b/src/tint/lang/spirv/writer/common/helper_test.h
index 4b419a3..3f50e15 100644
--- a/src/tint/lang/spirv/writer/common/helper_test.h
+++ b/src/tint/lang/spirv/writer/common/helper_test.h
@@ -219,7 +219,7 @@
             case kI32:
                 return b.Composite(MakeVectorType(type), 42_i, -10_i);
             case kU32:
-                return b.Composite(MakeVectorType(type), 42_u, 10_u);
+                return b.Composite(MakeVectorType(type), 31_u, 10_u);
             case kF32:
                 return b.Composite(MakeVectorType(type), 42_f, -0.5_f);
             case kF16:
diff --git a/src/tint/lang/spirv/writer/common/options.h b/src/tint/lang/spirv/writer/common/options.h
index c60c0f3..df4b60d 100644
--- a/src/tint/lang/spirv/writer/common/options.h
+++ b/src/tint/lang/spirv/writer/common/options.h
@@ -35,6 +35,7 @@
 #include "src/tint/api/common/binding_point.h"
 #include "src/tint/api/common/bindings.h"
 #include "src/tint/api/common/resource_binding_config.h"
+#include "src/tint/api/common/substitute_overrides_config.h"
 #include "src/tint/utils/reflection.h"
 
 namespace tint::spirv::writer {
@@ -151,6 +152,9 @@
     /// Resource binding information
     std::optional<ResourceBindingConfig> resource_binding = std::nullopt;
 
+    // Configuration for substitute overrides
+    SubstituteOverridesConfig substitute_overrides_config = {};
+
     /// Reflect the fields of this class so that it can be used by tint::ForeachField()
     TINT_REFLECT(Options,
                  remapped_entry_point_name,
@@ -179,7 +183,8 @@
                  decompose_uniform_buffers,
                  depth_range_offsets,
                  spirv_version,
-                 resource_binding);
+                 resource_binding,
+                 substitute_overrides_config);
 };
 
 }  // namespace tint::spirv::writer
diff --git a/src/tint/lang/spirv/writer/raise/raise.cc b/src/tint/lang/spirv/writer/raise/raise.cc
index 649424c..72cdfff 100644
--- a/src/tint/lang/spirv/writer/raise/raise.cc
+++ b/src/tint/lang/spirv/writer/raise/raise.cc
@@ -49,6 +49,7 @@
 #include "src/tint/lang/core/ir/transform/robustness.h"
 #include "src/tint/lang/core/ir/transform/signed_integer_polyfill.h"
 #include "src/tint/lang/core/ir/transform/std140.h"
+#include "src/tint/lang/core/ir/transform/substitute_overrides.h"
 #include "src/tint/lang/core/ir/transform/vectorize_scalar_matrix_constructors.h"
 #include "src/tint/lang/core/ir/transform/zero_init_workgroup_memory.h"
 #include "src/tint/lang/core/type/f32.h"
@@ -77,6 +78,9 @@
         }                                \
     } while (false)
 
+    RUN_TRANSFORM(core::ir::transform::SubstituteOverrides, module,
+                  options.substitute_overrides_config);
+
     tint::transform::multiplanar::BindingsMap multiplanar_map{};
     RemapperData remapper_data{};
     PopulateRemapperAndMultiplanarOptions(options, remapper_data, multiplanar_map);
diff --git a/src/tint/lang/spirv/writer/writer.cc b/src/tint/lang/spirv/writer/writer.cc
index 72e896d..1553531 100644
--- a/src/tint/lang/spirv/writer/writer.cc
+++ b/src/tint/lang/spirv/writer/writer.cc
@@ -31,6 +31,7 @@
 #include <utility>
 #include <vector>
 
+#include "src/tint/lang/core/ir/referenced_module_vars.h"
 #include "src/tint/lang/core/ir/validator.h"
 #include "src/tint/lang/core/ir/var.h"
 #include "src/tint/lang/core/type/binding_array.h"
@@ -44,7 +45,9 @@
 
 namespace tint::spirv::writer {
 
-Result<SuccessType> CanGenerate(const core::ir::Module& ir, const Options& options) {
+Result<SuccessType> CanGenerate(const core::ir::Module& ir,
+                                const Options& options,
+                                const std::string& ep_name) {
     // The enum is accessible in the API so ensure we have a valid value.
     switch (options.spirv_version) {
         case SpvVersion::kSpv13:
@@ -84,10 +87,27 @@
         }
     }
 
+    core::ir::Function* ep_func = nullptr;
+    for (auto* f : ir.functions) {
+        if (!f->IsEntryPoint()) {
+            continue;
+        }
+        if (ir.NameOf(f).NameView() == ep_name) {
+            ep_func = f;
+            break;
+        }
+    }
+    // No entrypoint, so no bindings needed
+    if (!ep_func) {
+        return Failure("entry point not found");
+    }
+
+    core::ir::ReferencedModuleVars<const core::ir::Module> referenced_module_vars{ir};
+    auto& refs = referenced_module_vars.TransitiveReferences(ep_func);
+
     // Check for unsupported module-scope variable address spaces and ensure at most one user
     // immediate.
-    for (auto* inst : *ir.root_block) {
-        auto* var = inst->As<core::ir::Var>();
+    for (auto* var : refs) {
         auto* ptr = var->Result()->Type()->As<core::type::Pointer>();
         if (ptr->AddressSpace() == core::AddressSpace::kPixelLocal) {
             return Failure("pixel_local address space is not supported by the SPIR-V backend");
diff --git a/src/tint/lang/spirv/writer/writer.h b/src/tint/lang/spirv/writer/writer.h
index 33c7873..2e0f89e 100644
--- a/src/tint/lang/spirv/writer/writer.h
+++ b/src/tint/lang/spirv/writer/writer.h
@@ -28,6 +28,8 @@
 #ifndef SRC_TINT_LANG_SPIRV_WRITER_WRITER_H_
 #define SRC_TINT_LANG_SPIRV_WRITER_WRITER_H_
 
+#include <string>
+
 #include "src/tint/lang/core/ir/module.h"
 #include "src/tint/lang/spirv/writer/common/options.h"
 #include "src/tint/lang/spirv/writer/common/output.h"
@@ -38,8 +40,11 @@
 /// Check if the module @p ir is supported by the SPIR-V backend with @p options.
 /// @param ir the module
 /// @param options the writer options
+/// @param ep_name the entry point name
 /// @returns Success or a failure message indicating why SPIR-V generation would fail
-Result<SuccessType> CanGenerate(const core::ir::Module& ir, const Options& options);
+Result<SuccessType> CanGenerate(const core::ir::Module& ir,
+                                const Options& options,
+                                const std::string& ep_name);
 
 /// Generate SPIR-V for a program, according to a set of configuration options.
 /// The result will contain the SPIR-V or failure.
diff --git a/src/tint/lang/spirv/writer/writer_fuzz.cc b/src/tint/lang/spirv/writer/writer_fuzz.cc
index d30c178..cc76836 100644
--- a/src/tint/lang/spirv/writer/writer_fuzz.cc
+++ b/src/tint/lang/spirv/writer/writer_fuzz.cc
@@ -58,7 +58,7 @@
         return Success;
     }
 
-    auto check = CanGenerate(module, options);
+    auto check = CanGenerate(module, options, ep_name);
     if (check != Success) {
         return Failure{check.Failure().reason};
     }
diff --git a/src/tint/lang/spirv/writer/writer_test.cc b/src/tint/lang/spirv/writer/writer_test.cc
index c470700..ca468e1 100644
--- a/src/tint/lang/spirv/writer/writer_test.cc
+++ b/src/tint/lang/spirv/writer/writer_test.cc
@@ -53,11 +53,19 @@
 }
 
 TEST_F(SpirvWriterTest, CanGenerate_SubgroupMatrixRequiresVulkanMemoryModel) {
-    mod.root_block->Append(b.Var(ty.ptr<private_>(ty.subgroup_matrix_result(ty.f32(), 8, 8))));
+    core::ir::Var* v = nullptr;
+    b.Append(mod.root_block,
+             [&] { v = b.Var(ty.ptr<private_>(ty.subgroup_matrix_result(ty.f32(), 8, 8))); });
+
+    auto* ep = b.ComputeFunction("main");
+    b.Append(ep->Block(), [&] {
+        b.Let("x", v);
+        b.Return(ep);
+    });
 
     Options options;
     options.use_vulkan_memory_model = false;
-    auto result = CanGenerate(mod, options);
+    auto result = CanGenerate(mod, options, "main");
     ASSERT_NE(result, Success);
     EXPECT_THAT(result.Failure().reason,
                 testing::HasSubstr("using subgroup matrices requires the Vulkan Memory Model"));