[ir] Add ReferencedModuleVars helper
This is a stateful helper that returns (and caches) the module-scope
variables that are transitively referenced by a function. Some work is
done on creation to determine which variables are directly referenced
by which blocks. The remaining work is done when the references for a
function are actually requested.
Update the ZeroInitializeWorkgroupMemory transform to use this helper.
Change-Id: I522d6521a61a94bf255ca9d9ff3d7d5998a2b69e
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/169222
Reviewed-by: dan sinclair <dsinclair@chromium.org>
Reviewed-by: Ben Clayton <bclayton@google.com>
Commit-Queue: James Price <jrprice@google.com>
diff --git a/src/tint/cmd/test/BUILD.bazel b/src/tint/cmd/test/BUILD.bazel
index c99263c..b6d308b 100644
--- a/src/tint/cmd/test/BUILD.bazel
+++ b/src/tint/cmd/test/BUILD.bazel
@@ -48,6 +48,7 @@
"//src/tint/cmd/common:test",
"//src/tint/lang/core/constant:test",
"//src/tint/lang/core/intrinsic:test",
+ "//src/tint/lang/core/ir/transform/common:test",
"//src/tint/lang/core/ir/transform:test",
"//src/tint/lang/core/ir:test",
"//src/tint/lang/core/type:test",
diff --git a/src/tint/cmd/test/BUILD.cmake b/src/tint/cmd/test/BUILD.cmake
index 0cb839e..94c5d90 100644
--- a/src/tint/cmd/test/BUILD.cmake
+++ b/src/tint/cmd/test/BUILD.cmake
@@ -49,6 +49,7 @@
tint_cmd_common_test
tint_lang_core_constant_test
tint_lang_core_intrinsic_test
+ tint_lang_core_ir_transform_common_test
tint_lang_core_ir_transform_test
tint_lang_core_ir_test
tint_lang_core_type_test
diff --git a/src/tint/cmd/test/BUILD.gn b/src/tint/cmd/test/BUILD.gn
index 54897e3..b8647d4 100644
--- a/src/tint/cmd/test/BUILD.gn
+++ b/src/tint/cmd/test/BUILD.gn
@@ -57,6 +57,7 @@
"${tint_src_dir}/lang/core/intrinsic:unittests",
"${tint_src_dir}/lang/core/ir:unittests",
"${tint_src_dir}/lang/core/ir/transform:unittests",
+ "${tint_src_dir}/lang/core/ir/transform/common:unittests",
"${tint_src_dir}/lang/core/type:unittests",
"${tint_src_dir}/lang/hlsl/writer/common:unittests",
"${tint_src_dir}/lang/msl/ir:unittests",
diff --git a/src/tint/lang/core/ir/transform/BUILD.bazel b/src/tint/lang/core/ir/transform/BUILD.bazel
index 467e451..3ea71c4 100644
--- a/src/tint/lang/core/ir/transform/BUILD.bazel
+++ b/src/tint/lang/core/ir/transform/BUILD.bazel
@@ -85,6 +85,7 @@
"//src/tint/lang/core/constant",
"//src/tint/lang/core/intrinsic",
"//src/tint/lang/core/ir",
+ "//src/tint/lang/core/ir/transform/common",
"//src/tint/lang/core/type",
"//src/tint/utils/containers",
"//src/tint/utils/diagnostic",
diff --git a/src/tint/lang/core/ir/transform/BUILD.cmake b/src/tint/lang/core/ir/transform/BUILD.cmake
index 1ae8efb..f3a8b32 100644
--- a/src/tint/lang/core/ir/transform/BUILD.cmake
+++ b/src/tint/lang/core/ir/transform/BUILD.cmake
@@ -34,6 +34,8 @@
# Do not modify this file directly
################################################################################
+include(lang/core/ir/transform/common/BUILD.cmake)
+
################################################################################
# Target: tint_lang_core_ir_transform
# Kind: lib
@@ -84,6 +86,7 @@
tint_lang_core_constant
tint_lang_core_intrinsic
tint_lang_core_ir
+ tint_lang_core_ir_transform_common
tint_lang_core_type
tint_utils_containers
tint_utils_diagnostic
diff --git a/src/tint/lang/core/ir/transform/BUILD.gn b/src/tint/lang/core/ir/transform/BUILD.gn
index 95de9fc..db194b7 100644
--- a/src/tint/lang/core/ir/transform/BUILD.gn
+++ b/src/tint/lang/core/ir/transform/BUILD.gn
@@ -88,6 +88,7 @@
"${tint_src_dir}/lang/core/constant",
"${tint_src_dir}/lang/core/intrinsic",
"${tint_src_dir}/lang/core/ir",
+ "${tint_src_dir}/lang/core/ir/transform/common",
"${tint_src_dir}/lang/core/type",
"${tint_src_dir}/utils/containers",
"${tint_src_dir}/utils/diagnostic",
diff --git a/src/tint/lang/core/ir/transform/common/BUILD.bazel b/src/tint/lang/core/ir/transform/common/BUILD.bazel
new file mode 100644
index 0000000..f05642f
--- /dev/null
+++ b/src/tint/lang/core/ir/transform/common/BUILD.bazel
@@ -0,0 +1,103 @@
+# Copyright 2024 The Dawn & Tint Authors
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+#
+# 1. Redistributions of source code must retain the above copyright notice, this
+# list of conditions and the following disclaimer.
+#
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+# this list of conditions and the following disclaimer in the documentation
+# and/or other materials provided with the distribution.
+#
+# 3. Neither the name of the copyright holder nor the names of its
+# contributors may be used to endorse or promote products derived from
+# this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+################################################################################
+# File generated by 'tools/src/cmd/gen' using the template:
+# tools/src/cmd/gen/build/BUILD.bazel.tmpl
+#
+# To regenerate run: './tools/run gen'
+#
+# Do not modify this file directly
+################################################################################
+
+load("//src/tint:flags.bzl", "COPTS")
+load("@bazel_skylib//lib:selects.bzl", "selects")
+cc_library(
+ name = "common",
+ srcs = [
+ "referenced_module_vars.cc",
+ ],
+ hdrs = [
+ "referenced_module_vars.h",
+ ],
+ deps = [
+ "//src/tint/api/common",
+ "//src/tint/lang/core",
+ "//src/tint/lang/core/constant",
+ "//src/tint/lang/core/ir",
+ "//src/tint/lang/core/type",
+ "//src/tint/utils/containers",
+ "//src/tint/utils/diagnostic",
+ "//src/tint/utils/ice",
+ "//src/tint/utils/id",
+ "//src/tint/utils/macros",
+ "//src/tint/utils/math",
+ "//src/tint/utils/memory",
+ "//src/tint/utils/reflection",
+ "//src/tint/utils/result",
+ "//src/tint/utils/rtti",
+ "//src/tint/utils/symbol",
+ "//src/tint/utils/text",
+ "//src/tint/utils/traits",
+ ],
+ copts = COPTS,
+ visibility = ["//visibility:public"],
+)
+cc_library(
+ name = "test",
+ alwayslink = True,
+ srcs = [
+ "referenced_module_vars_test.cc",
+ ],
+ deps = [
+ "//src/tint/api/common",
+ "//src/tint/lang/core",
+ "//src/tint/lang/core/constant",
+ "//src/tint/lang/core/intrinsic",
+ "//src/tint/lang/core/ir",
+ "//src/tint/lang/core/ir/transform/common",
+ "//src/tint/lang/core/ir:test",
+ "//src/tint/lang/core/type",
+ "//src/tint/utils/containers",
+ "//src/tint/utils/diagnostic",
+ "//src/tint/utils/ice",
+ "//src/tint/utils/id",
+ "//src/tint/utils/macros",
+ "//src/tint/utils/math",
+ "//src/tint/utils/memory",
+ "//src/tint/utils/reflection",
+ "//src/tint/utils/result",
+ "//src/tint/utils/rtti",
+ "//src/tint/utils/symbol",
+ "//src/tint/utils/text",
+ "//src/tint/utils/traits",
+ "@gtest",
+ ],
+ copts = COPTS,
+ visibility = ["//visibility:public"],
+)
+
diff --git a/src/tint/lang/core/ir/transform/common/BUILD.cmake b/src/tint/lang/core/ir/transform/common/BUILD.cmake
new file mode 100644
index 0000000..7004d86
--- /dev/null
+++ b/src/tint/lang/core/ir/transform/common/BUILD.cmake
@@ -0,0 +1,101 @@
+# Copyright 2024 The Dawn & Tint Authors
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+#
+# 1. Redistributions of source code must retain the above copyright notice, this
+# list of conditions and the following disclaimer.
+#
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+# this list of conditions and the following disclaimer in the documentation
+# and/or other materials provided with the distribution.
+#
+# 3. Neither the name of the copyright holder nor the names of its
+# contributors may be used to endorse or promote products derived from
+# this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+################################################################################
+# File generated by 'tools/src/cmd/gen' using the template:
+# tools/src/cmd/gen/build/BUILD.cmake.tmpl
+#
+# To regenerate run: './tools/run gen'
+#
+# Do not modify this file directly
+################################################################################
+
+################################################################################
+# Target: tint_lang_core_ir_transform_common
+# Kind: lib
+################################################################################
+tint_add_target(tint_lang_core_ir_transform_common lib
+ lang/core/ir/transform/common/referenced_module_vars.cc
+ lang/core/ir/transform/common/referenced_module_vars.h
+)
+
+tint_target_add_dependencies(tint_lang_core_ir_transform_common lib
+ tint_api_common
+ tint_lang_core
+ tint_lang_core_constant
+ tint_lang_core_ir
+ tint_lang_core_type
+ tint_utils_containers
+ tint_utils_diagnostic
+ tint_utils_ice
+ tint_utils_id
+ tint_utils_macros
+ tint_utils_math
+ tint_utils_memory
+ tint_utils_reflection
+ tint_utils_result
+ tint_utils_rtti
+ tint_utils_symbol
+ tint_utils_text
+ tint_utils_traits
+)
+
+################################################################################
+# Target: tint_lang_core_ir_transform_common_test
+# Kind: test
+################################################################################
+tint_add_target(tint_lang_core_ir_transform_common_test test
+ lang/core/ir/transform/common/referenced_module_vars_test.cc
+)
+
+tint_target_add_dependencies(tint_lang_core_ir_transform_common_test test
+ tint_api_common
+ tint_lang_core
+ tint_lang_core_constant
+ tint_lang_core_intrinsic
+ tint_lang_core_ir
+ tint_lang_core_ir_transform_common
+ tint_lang_core_ir_test
+ tint_lang_core_type
+ tint_utils_containers
+ tint_utils_diagnostic
+ tint_utils_ice
+ tint_utils_id
+ tint_utils_macros
+ tint_utils_math
+ tint_utils_memory
+ tint_utils_reflection
+ tint_utils_result
+ tint_utils_rtti
+ tint_utils_symbol
+ tint_utils_text
+ tint_utils_traits
+)
+
+tint_target_add_external_dependencies(tint_lang_core_ir_transform_common_test test
+ "gtest"
+)
diff --git a/src/tint/lang/core/ir/transform/common/BUILD.gn b/src/tint/lang/core/ir/transform/common/BUILD.gn
new file mode 100644
index 0000000..d050caf
--- /dev/null
+++ b/src/tint/lang/core/ir/transform/common/BUILD.gn
@@ -0,0 +1,99 @@
+# Copyright 2024 The Dawn & Tint Authors
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+#
+# 1. Redistributions of source code must retain the above copyright notice, this
+# list of conditions and the following disclaimer.
+#
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+# this list of conditions and the following disclaimer in the documentation
+# and/or other materials provided with the distribution.
+#
+# 3. Neither the name of the copyright holder nor the names of its
+# contributors may be used to endorse or promote products derived from
+# this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+################################################################################
+# File generated by 'tools/src/cmd/gen' using the template:
+# tools/src/cmd/gen/build/BUILD.gn.tmpl
+#
+# To regenerate run: './tools/run gen'
+#
+# Do not modify this file directly
+################################################################################
+
+import("../../../../../../../scripts/tint_overrides_with_defaults.gni")
+
+import("${tint_src_dir}/tint.gni")
+
+if (tint_build_unittests || tint_build_benchmarks) {
+ import("//testing/test.gni")
+}
+
+libtint_source_set("common") {
+ sources = [
+ "referenced_module_vars.cc",
+ "referenced_module_vars.h",
+ ]
+ deps = [
+ "${tint_src_dir}/api/common",
+ "${tint_src_dir}/lang/core",
+ "${tint_src_dir}/lang/core/constant",
+ "${tint_src_dir}/lang/core/ir",
+ "${tint_src_dir}/lang/core/type",
+ "${tint_src_dir}/utils/containers",
+ "${tint_src_dir}/utils/diagnostic",
+ "${tint_src_dir}/utils/ice",
+ "${tint_src_dir}/utils/id",
+ "${tint_src_dir}/utils/macros",
+ "${tint_src_dir}/utils/math",
+ "${tint_src_dir}/utils/memory",
+ "${tint_src_dir}/utils/reflection",
+ "${tint_src_dir}/utils/result",
+ "${tint_src_dir}/utils/rtti",
+ "${tint_src_dir}/utils/symbol",
+ "${tint_src_dir}/utils/text",
+ "${tint_src_dir}/utils/traits",
+ ]
+}
+if (tint_build_unittests) {
+ tint_unittests_source_set("unittests") {
+ sources = [ "referenced_module_vars_test.cc" ]
+ deps = [
+ "${tint_src_dir}:gmock_and_gtest",
+ "${tint_src_dir}/api/common",
+ "${tint_src_dir}/lang/core",
+ "${tint_src_dir}/lang/core/constant",
+ "${tint_src_dir}/lang/core/intrinsic",
+ "${tint_src_dir}/lang/core/ir",
+ "${tint_src_dir}/lang/core/ir:unittests",
+ "${tint_src_dir}/lang/core/ir/transform/common",
+ "${tint_src_dir}/lang/core/type",
+ "${tint_src_dir}/utils/containers",
+ "${tint_src_dir}/utils/diagnostic",
+ "${tint_src_dir}/utils/ice",
+ "${tint_src_dir}/utils/id",
+ "${tint_src_dir}/utils/macros",
+ "${tint_src_dir}/utils/math",
+ "${tint_src_dir}/utils/memory",
+ "${tint_src_dir}/utils/reflection",
+ "${tint_src_dir}/utils/result",
+ "${tint_src_dir}/utils/rtti",
+ "${tint_src_dir}/utils/symbol",
+ "${tint_src_dir}/utils/text",
+ "${tint_src_dir}/utils/traits",
+ ]
+ }
+}
diff --git a/src/tint/lang/core/ir/transform/common/referenced_module_vars.cc b/src/tint/lang/core/ir/transform/common/referenced_module_vars.cc
new file mode 100644
index 0000000..b97b9dd
--- /dev/null
+++ b/src/tint/lang/core/ir/transform/common/referenced_module_vars.cc
@@ -0,0 +1,75 @@
+// Copyright 2024 The Dawn & Tint Authors
+//
+// Redistribution and use in source and binary forms, with or without
+// modification, are permitted provided that the following conditions are met:
+//
+// 1. Redistributions of source code must retain the above copyright notice, this
+// list of conditions and the following disclaimer.
+//
+// 2. Redistributions in binary form must reproduce the above copyright notice,
+// this list of conditions and the following disclaimer in the documentation
+// and/or other materials provided with the distribution.
+//
+// 3. Neither the name of the copyright holder nor the names of its
+// contributors may be used to endorse or promote products derived from
+// this software without specific prior written permission.
+//
+// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+// FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+// DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+// CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+// OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+#include "src/tint/lang/core/ir/transform/common/referenced_module_vars.h"
+
+#include "src/tint/lang/core/ir/control_instruction.h"
+#include "src/tint/lang/core/ir/module.h"
+#include "src/tint/lang/core/ir/user_call.h"
+#include "src/tint/lang/core/ir/var.h"
+#include "src/tint/utils/rtti/switch.h"
+
+namespace tint::core::ir {
+
+const ReferencedModuleVars::VarSet& ReferencedModuleVars::TransitiveReferences(Function* func) {
+ return transitive_references_.GetOrAdd(func, [&] {
+ VarSet vars;
+ GetTransitiveReferences(func->Block(), vars);
+ return vars;
+ });
+}
+
+/// Get the set of variables transitively referenced by @p block.
+/// @param block the block
+/// @param vars the set of transitively referenced variables to populate
+void ReferencedModuleVars::GetTransitiveReferences(Block* block, VarSet& vars) {
+ // Add directly referenced vars.
+ if (auto itr = block_to_direct_vars_.Get(block)) {
+ for (auto& var : *itr) {
+ vars.Add(var);
+ }
+ }
+
+ // Loop over instructions in the block to find indirectly referenced vars.
+ for (auto* inst : *block) {
+ tint::Switch(
+ inst,
+ [&](UserCall* call) {
+ // Get variables referenced by a function called from this block.
+ const auto& callee_vars = TransitiveReferences(call->Target());
+ for (auto* var : callee_vars) {
+ vars.Add(var);
+ }
+ },
+ [&](ControlInstruction* ctrl) {
+ // Recurse into control instructions and gather their referenced vars.
+ ctrl->ForeachBlock([&](Block* blk) { GetTransitiveReferences(blk, vars); });
+ });
+ }
+}
+
+} // namespace tint::core::ir
diff --git a/src/tint/lang/core/ir/transform/common/referenced_module_vars.h b/src/tint/lang/core/ir/transform/common/referenced_module_vars.h
new file mode 100644
index 0000000..a48b0a7
--- /dev/null
+++ b/src/tint/lang/core/ir/transform/common/referenced_module_vars.h
@@ -0,0 +1,105 @@
+// Copyright 2024 The Dawn & Tint Authors
+//
+// Redistribution and use in source and binary forms, with or without
+// modification, are permitted provided that the following conditions are met:
+//
+// 1. Redistributions of source code must retain the above copyright notice, this
+// list of conditions and the following disclaimer.
+//
+// 2. Redistributions in binary form must reproduce the above copyright notice,
+// this list of conditions and the following disclaimer in the documentation
+// and/or other materials provided with the distribution.
+//
+// 3. Neither the name of the copyright holder nor the names of its
+// contributors may be used to endorse or promote products derived from
+// this software without specific prior written permission.
+//
+// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+// FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+// DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+// CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+// OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+#ifndef SRC_TINT_LANG_CORE_IR_TRANSFORM_COMMON_REFERENCED_MODULE_VARS_H_
+#define SRC_TINT_LANG_CORE_IR_TRANSFORM_COMMON_REFERENCED_MODULE_VARS_H_
+
+#include <functional>
+
+#include "src/tint/lang/core/ir/module.h"
+#include "src/tint/lang/core/ir/var.h"
+#include "src/tint/utils/containers/hashmap.h"
+#include "src/tint/utils/containers/unique_vector.h"
+
+// Forward declarations.
+namespace tint::core::ir {
+class Block;
+class Function;
+} // namespace tint::core::ir
+
+namespace tint::core::ir {
+
+/// ReferencedModuleVars is a helper to determine the set of module-scope variables that are
+/// transitively referenced by functions in a module.
+/// References are determined lazily and cached for future requests.
+///
+/// Note: changes to the module can invalidate the cached data. This is intended to be created by
+/// a transform that need this information, and discarded when that transform completes. Tracking
+/// this information inside the IR module would add overhead any time an instruction is added or
+/// removed from the module. Since only a few transforms need this information, we expect it to be
+/// more efficient to generate it as and when needed instead.
+class ReferencedModuleVars {
+ public:
+ /// The signature of a predicate used to filter variables.
+ /// A predicate function should return `true` when the variable should be added to the set.
+ using Predicate = std::function<bool(const Var*)>;
+
+ /// A set of a variables referenced by a function (in declaration order).
+ using VarSet = UniqueVector<Var*, 16>;
+
+ /// Constructor.
+ /// @param ir the module
+ /// @param pred an optional predicate function for filtering variables
+ /// Note: @p pred is not stored by the class, so can be a lambda that captures by reference.
+ explicit ReferencedModuleVars(Module& ir, Predicate&& pred = {}) : ir_(ir) {
+ // Loop over module-scope variables, recording the blocks that they are referenced from.
+ for (auto inst : *ir_.root_block) {
+ if (auto* var = inst->As<Var>()) {
+ if (!pred || pred(var)) {
+ var->Result(0)->ForEachUse([&](const Usage& use) {
+ block_to_direct_vars_.GetOrAddZero(use.instruction->Block()).Add(var);
+ });
+ }
+ }
+ }
+ }
+
+ /// Get the set of transitively referenced module-scope variables for a function, filtered by
+ /// the predicate function if provided.
+ /// @param func the function
+ /// @returns the set of (possibly filtered) transitively reference module-scope variables
+ const VarSet& TransitiveReferences(Function* func);
+
+ private:
+ /// The module.
+ Module& ir_;
+
+ /// A map from blocks to their directly referenced variables.
+ Hashmap<Block*, VarSet, 64> block_to_direct_vars_{};
+
+ /// A map from functions to their transitively referenced variables.
+ Hashmap<Function*, VarSet, 8> transitive_references_;
+
+ /// Get the set of transitively referenced module-scope variables for a block.
+ /// @param block the block
+ /// @param vars the set of transitively reference module-scope variables to populate
+ void GetTransitiveReferences(Block* block, VarSet& vars);
+};
+
+} // namespace tint::core::ir
+
+#endif // SRC_TINT_LANG_CORE_IR_TRANSFORM_COMMON_REFERENCED_MODULE_VARS_H_
diff --git a/src/tint/lang/core/ir/transform/common/referenced_module_vars_test.cc b/src/tint/lang/core/ir/transform/common/referenced_module_vars_test.cc
new file mode 100644
index 0000000..3a5a032
--- /dev/null
+++ b/src/tint/lang/core/ir/transform/common/referenced_module_vars_test.cc
@@ -0,0 +1,422 @@
+// Copyright 2024 The Dawn & Tint Authors
+//
+// Redistribution and use in source and binary forms, with or without
+// modification, are permitted provided that the following conditions are met:
+//
+// 1. Redistributions of source code must retain the above copyright notice, this
+// list of conditions and the following disclaimer.
+//
+// 2. Redistributions in binary form must reproduce the above copyright notice,
+// this list of conditions and the following disclaimer in the documentation
+// and/or other materials provided with the distribution.
+//
+// 3. Neither the name of the copyright holder nor the names of its
+// contributors may be used to endorse or promote products derived from
+// this software without specific prior written permission.
+//
+// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+// FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+// DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+// CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+// OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+#include "src/tint/lang/core/ir/transform/common/referenced_module_vars.h"
+
+#include <string>
+
+#include "gmock/gmock.h"
+#include "src/tint/lang/core/ir/disassembly.h"
+#include "src/tint/lang/core/ir/ir_helper_test.h"
+
+namespace tint::core::ir {
+namespace {
+
+using ::testing::ElementsAre;
+
+using namespace tint::core::fluent_types; // NOLINT
+using namespace tint::core::number_suffixes; // NOLINT
+
+class IR_ReferencedModuleVarsTest : public IRTestHelper {
+ protected:
+ /// @returns the module as a disassembled string
+ std::string Disassemble() const { return "\n" + ir::Disassemble(mod).Plain(); }
+};
+
+TEST_F(IR_ReferencedModuleVarsTest, EmptyRootBlock) {
+ auto* foo = b.Function("foo", ty.void_());
+ b.Append(foo->Block(), [&] { //
+ b.Return(foo);
+ });
+
+ auto* src = R"(
+%foo = func():void {
+ $B1: {
+ ret
+ }
+}
+)";
+ EXPECT_EQ(src, Disassemble());
+
+ ReferencedModuleVars vars(mod);
+ auto& foo_vars = vars.TransitiveReferences(foo);
+ EXPECT_TRUE(foo_vars.IsEmpty());
+}
+
+TEST_F(IR_ReferencedModuleVarsTest, DirectUse) {
+ // Referenced.
+ auto* var_a = mod.root_block->Append(b.Var<workgroup, u32>("a"));
+ auto* var_b = mod.root_block->Append(b.Var<workgroup, u32>("b"));
+ // Not referenced.
+ mod.root_block->Append(b.Var<workgroup, u32>("c"));
+
+ auto* foo = b.Function("foo", ty.void_());
+ b.Append(foo->Block(), [&] { //
+ b.Load(var_a);
+ b.Load(var_b);
+ b.Return(foo);
+ });
+
+ auto* src = R"(
+$B1: { # root
+ %a:ptr<workgroup, u32, read_write> = var
+ %b:ptr<workgroup, u32, read_write> = var
+ %c:ptr<workgroup, u32, read_write> = var
+}
+
+%foo = func():void {
+ $B2: {
+ %5:u32 = load %a
+ %6:u32 = load %b
+ ret
+ }
+}
+)";
+ EXPECT_EQ(src, Disassemble());
+
+ ReferencedModuleVars vars(mod);
+ EXPECT_THAT(vars.TransitiveReferences(foo), ElementsAre(var_a, var_b));
+}
+
+TEST_F(IR_ReferencedModuleVarsTest, DirectUse_DeclarationOrder) {
+ auto* var_a = mod.root_block->Append(b.Var<workgroup, u32>("a"));
+ auto* var_b = mod.root_block->Append(b.Var<workgroup, u32>("b"));
+ auto* var_c = mod.root_block->Append(b.Var<workgroup, u32>("c"));
+ auto* var_d = mod.root_block->Append(b.Var<workgroup, u32>("d"));
+ auto* var_e = mod.root_block->Append(b.Var<workgroup, u32>("e"));
+
+ auto* foo = b.Function("foo", ty.void_());
+ b.Append(foo->Block(), [&] { //
+ b.Load(var_b);
+ b.Load(var_e);
+ b.Load(var_d);
+ b.Load(var_c);
+ b.Load(var_a);
+ b.Return(foo);
+ });
+
+ auto* src = R"(
+$B1: { # root
+ %a:ptr<workgroup, u32, read_write> = var
+ %b:ptr<workgroup, u32, read_write> = var
+ %c:ptr<workgroup, u32, read_write> = var
+ %d:ptr<workgroup, u32, read_write> = var
+ %e:ptr<workgroup, u32, read_write> = var
+}
+
+%foo = func():void {
+ $B2: {
+ %7:u32 = load %b
+ %8:u32 = load %e
+ %9:u32 = load %d
+ %10:u32 = load %c
+ %11:u32 = load %a
+ ret
+ }
+}
+)";
+ EXPECT_EQ(src, Disassemble());
+
+ ReferencedModuleVars vars(mod);
+ EXPECT_THAT(vars.TransitiveReferences(foo), ElementsAre(var_a, var_b, var_c, var_d, var_e));
+}
+
+TEST_F(IR_ReferencedModuleVarsTest, DirectUse_MultipleFunctions) {
+ auto* var_a = mod.root_block->Append(b.Var<workgroup, u32>("a"));
+ auto* var_b = mod.root_block->Append(b.Var<workgroup, u32>("b"));
+ auto* var_c = mod.root_block->Append(b.Var<workgroup, u32>("c"));
+
+ auto* foo = b.Function("foo", ty.void_());
+ b.Append(foo->Block(), [&] { //
+ b.Load(var_a);
+ b.Load(var_b);
+ b.Return(foo);
+ });
+
+ auto* bar = b.Function("bar", ty.void_());
+ b.Append(bar->Block(), [&] { //
+ b.Load(var_a);
+ b.Load(var_c);
+ b.Return(bar);
+ });
+
+ auto* zoo = b.Function("zoo", ty.void_());
+ b.Append(zoo->Block(), [&] { //
+ b.Return(zoo);
+ });
+
+ auto* src = R"(
+$B1: { # root
+ %a:ptr<workgroup, u32, read_write> = var
+ %b:ptr<workgroup, u32, read_write> = var
+ %c:ptr<workgroup, u32, read_write> = var
+}
+
+%foo = func():void {
+ $B2: {
+ %5:u32 = load %a
+ %6:u32 = load %b
+ ret
+ }
+}
+%bar = func():void {
+ $B3: {
+ %8:u32 = load %a
+ %9:u32 = load %c
+ ret
+ }
+}
+%zoo = func():void {
+ $B4: {
+ ret
+ }
+}
+)";
+ EXPECT_EQ(src, Disassemble());
+
+ ReferencedModuleVars vars(mod);
+ EXPECT_THAT(vars.TransitiveReferences(foo), ElementsAre(var_a, var_b));
+ EXPECT_THAT(vars.TransitiveReferences(bar), ElementsAre(var_a, var_c));
+ EXPECT_TRUE(vars.TransitiveReferences(zoo).IsEmpty());
+}
+
+TEST_F(IR_ReferencedModuleVarsTest, DirectUse_NestedInControlFlow) {
+ auto* var_a = mod.root_block->Append(b.Var<workgroup, u32>("a"));
+ auto* var_b = mod.root_block->Append(b.Var<workgroup, u32>("b"));
+ auto* var_c = mod.root_block->Append(b.Var<workgroup, u32>("c"));
+ auto* var_d = mod.root_block->Append(b.Var<workgroup, u32>("c"));
+
+ auto* foo = b.Function("foo", ty.void_());
+ b.Append(foo->Block(), [&] { //
+ auto* ifelse = b.If(true);
+ b.Append(ifelse->True(), [&] {
+ b.Load(var_a);
+ b.ExitIf(ifelse);
+ });
+ b.Append(ifelse->False(), [&] {
+ auto* loop = b.Loop();
+ b.Append(loop->Initializer(), [&] {
+ b.Load(var_b);
+ b.NextIteration(loop);
+ });
+ b.Append(loop->Body(), [&] {
+ b.Load(var_c);
+ b.Continue(loop);
+ });
+ b.Append(loop->Continuing(), [&] {
+ b.Load(var_d);
+ b.NextIteration(loop);
+ });
+ b.ExitIf(ifelse);
+ });
+ b.Return(foo);
+ });
+
+ auto* src = R"(
+$B1: { # root
+ %a:ptr<workgroup, u32, read_write> = var
+ %b:ptr<workgroup, u32, read_write> = var
+ %c:ptr<workgroup, u32, read_write> = var
+ %c_1:ptr<workgroup, u32, read_write> = var # %c_1: 'c'
+}
+
+%foo = func():void {
+ $B2: {
+ if true [t: $B3, f: $B4] { # if_1
+ $B3: { # true
+ %6:u32 = load %a
+ exit_if # if_1
+ }
+ $B4: { # false
+ loop [i: $B5, b: $B6, c: $B7] { # loop_1
+ $B5: { # initializer
+ %7:u32 = load %b
+ next_iteration # -> $B6
+ }
+ $B6: { # body
+ %8:u32 = load %c
+ continue # -> $B7
+ }
+ $B7: { # continuing
+ %9:u32 = load %c_1
+ next_iteration # -> $B6
+ }
+ }
+ exit_if # if_1
+ }
+ }
+ ret
+ }
+}
+)";
+ EXPECT_EQ(src, Disassemble());
+
+ ReferencedModuleVars vars(mod);
+ EXPECT_THAT(vars.TransitiveReferences(foo), ElementsAre(var_a, var_b, var_c, var_d));
+}
+
+TEST_F(IR_ReferencedModuleVarsTest, IndirectUse) {
+ // Directly used by foo.
+ auto* var_a = mod.root_block->Append(b.Var<workgroup, u32>("a"));
+ // Directly used by bar, called by zoo and foo.
+ auto* var_b = mod.root_block->Append(b.Var<workgroup, u32>("b"));
+ // Not used.
+ mod.root_block->Append(b.Var<workgroup, u32>("c"));
+
+ auto* bar = b.Function("bar", ty.void_());
+ b.Append(bar->Block(), [&] { //
+ b.Load(var_b);
+ b.Return(bar);
+ });
+
+ auto* zoo = b.Function("zoo", ty.void_());
+ b.Append(zoo->Block(), [&] { //
+ b.Call(bar);
+ b.Return(zoo);
+ });
+
+ auto* foo = b.Function("foo", ty.void_());
+ b.Append(foo->Block(), [&] { //
+ b.Load(var_a);
+ b.Call(zoo);
+ b.Return(foo);
+ });
+
+ auto* src = R"(
+$B1: { # root
+ %a:ptr<workgroup, u32, read_write> = var
+ %b:ptr<workgroup, u32, read_write> = var
+ %c:ptr<workgroup, u32, read_write> = var
+}
+
+%bar = func():void {
+ $B2: {
+ %5:u32 = load %b
+ ret
+ }
+}
+%zoo = func():void {
+ $B3: {
+ %7:void = call %bar
+ ret
+ }
+}
+%foo = func():void {
+ $B4: {
+ %9:u32 = load %a
+ %10:void = call %zoo
+ ret
+ }
+}
+)";
+ EXPECT_EQ(src, Disassemble());
+
+ ReferencedModuleVars vars(mod);
+ EXPECT_THAT(vars.TransitiveReferences(bar), ElementsAre(var_b));
+ EXPECT_THAT(vars.TransitiveReferences(zoo), ElementsAre(var_b));
+ EXPECT_THAT(vars.TransitiveReferences(foo), ElementsAre(var_a, var_b));
+}
+
+TEST_F(IR_ReferencedModuleVarsTest, NoFunctionVars) {
+ auto* var_a = mod.root_block->Append(b.Var<workgroup, u32>("a"));
+
+ auto* foo = b.Function("foo", ty.void_());
+ b.Append(foo->Block(), [&] { //
+ auto* var_b = b.Var<function, u32>("b");
+ b.Load(var_a);
+ b.Load(var_b);
+ b.Return(foo);
+ });
+
+ auto* src = R"(
+$B1: { # root
+ %a:ptr<workgroup, u32, read_write> = var
+}
+
+%foo = func():void {
+ $B2: {
+ %b:ptr<function, u32, read_write> = var
+ %4:u32 = load %a
+ %5:u32 = load %b
+ ret
+ }
+}
+)";
+ EXPECT_EQ(src, Disassemble());
+
+ ReferencedModuleVars vars(mod);
+ EXPECT_THAT(vars.TransitiveReferences(foo), ElementsAre(var_a));
+}
+
+TEST_F(IR_ReferencedModuleVarsTest, Predicate) {
+ auto* var_a = mod.root_block->Append(b.Var<workgroup, u32>("a"));
+ auto* var_b = mod.root_block->Append(b.Var<private_, u32>("b"));
+ auto* var_c = mod.root_block->Append(b.Var<workgroup, u32>("c"));
+ auto* var_d = mod.root_block->Append(b.Var<private_, u32>("d"));
+ auto* var_e = mod.root_block->Append(b.Var<workgroup, u32>("e"));
+
+ auto* foo = b.Function("foo", ty.void_());
+ b.Append(foo->Block(), [&] { //
+ b.Load(var_a);
+ b.Load(var_b);
+ b.Load(var_c);
+ b.Load(var_d);
+ b.Load(var_e);
+ b.Return(foo);
+ });
+
+ auto* src = R"(
+$B1: { # root
+ %a:ptr<workgroup, u32, read_write> = var
+ %b:ptr<private, u32, read_write> = var
+ %c:ptr<workgroup, u32, read_write> = var
+ %d:ptr<private, u32, read_write> = var
+ %e:ptr<workgroup, u32, read_write> = var
+}
+
+%foo = func():void {
+ $B2: {
+ %7:u32 = load %a
+ %8:u32 = load %b
+ %9:u32 = load %c
+ %10:u32 = load %d
+ %11:u32 = load %e
+ ret
+ }
+}
+)";
+ EXPECT_EQ(src, Disassemble());
+
+ ReferencedModuleVars vars(mod, [](const Var* var) {
+ auto* view = var->Result(0)->Type()->As<type::MemoryView>();
+ return view->AddressSpace() == AddressSpace::kPrivate;
+ });
+ EXPECT_THAT(vars.TransitiveReferences(foo), ElementsAre(var_b, var_d));
+}
+
+} // namespace
+} // namespace tint::core::ir
diff --git a/src/tint/lang/core/ir/transform/zero_init_workgroup_memory.cc b/src/tint/lang/core/ir/transform/zero_init_workgroup_memory.cc
index 26cbf76..8969c9f 100644
--- a/src/tint/lang/core/ir/transform/zero_init_workgroup_memory.cc
+++ b/src/tint/lang/core/ir/transform/zero_init_workgroup_memory.cc
@@ -32,6 +32,7 @@
#include "src/tint/lang/core/ir/builder.h"
#include "src/tint/lang/core/ir/module.h"
+#include "src/tint/lang/core/ir/transform/common/referenced_module_vars.h"
#include "src/tint/lang/core/ir/validator.h"
#include "src/tint/utils/containers/reverse.h"
@@ -61,17 +62,12 @@
/// The type manager.
core::type::Manager& ty{ir.Types()};
- /// VarSet is a hash set of workgroup variables.
- using VarSet = Hashset<Var*, 8>;
-
- /// A map from variable to an ID used for sorting.
- Hashmap<Var*, uint32_t, 8> var_to_id{};
-
- /// A map from blocks to their directly referenced workgroup variables.
- Hashmap<Block*, VarSet, 64> block_to_direct_vars{};
-
- /// A map from functions to their transitively referenced workgroup variables.
- Hashmap<Function*, VarSet, 8> function_to_transitive_vars{};
+ /// The mapping from functions to their transitively referenced workgroup variables.
+ ReferencedModuleVars referenced_module_vars_{
+ ir, [](const Var* var) {
+ auto* view = var->Result(0)->Type()->As<type::MemoryView>();
+ return view && view->AddressSpace() == AddressSpace::kWorkgroup;
+ }};
/// ArrayIndex represents a required array index for an access instruction.
struct ArrayIndex {
@@ -104,22 +100,6 @@
if (ir.root_block->IsEmpty()) {
return;
}
-
- // Loop over module-scope variables, looking for workgroup variables.
- uint32_t next_id = 0;
- for (auto inst : *ir.root_block) {
- if (auto* var = inst->As<Var>()) {
- auto* ptr = var->Result(0)->Type()->As<core::type::Pointer>();
- if (ptr && ptr->AddressSpace() == core::AddressSpace::kWorkgroup) {
- // Record the usage of the variable for each block that references it.
- var->Result(0)->ForEachUse([&](const Usage& use) {
- block_to_direct_vars.GetOrAddZero(use.instruction->Block()).Add(var);
- });
- var_to_id.Add(var, next_id++);
- }
- }
- }
-
// Process each entry point function.
for (auto& func : ir.functions) {
if (func->Stage() == Function::PipelineStage::kCompute) {
@@ -132,20 +112,14 @@
/// @param func the entry point function
void ProcessEntryPoint(Function* func) {
// Get list of transitively referenced workgroup variables.
- auto vars = GetReferencedVars(func);
+ const auto& vars = referenced_module_vars_.TransitiveReferences(func);
if (vars.IsEmpty()) {
return;
}
- // Sort the variables to get deterministic output in tests.
- auto sorted_vars = vars.Vector();
- sorted_vars.Sort([&](Var* first, Var* second) {
- return *var_to_id.Get(first) < *var_to_id.Get(second);
- });
-
// Build list of store descriptors for all workgroup variables.
StoreMap stores;
- for (auto* var : sorted_vars) {
+ for (auto* var : vars) {
PrepareStores(var, var->Result(0)->Type()->UnwrapPtr(), 1, {}, stores);
}
@@ -188,46 +162,6 @@
});
}
- /// Get the set of workgroup variables transitively referenced by @p func.
- /// @param func the function
- /// @returns the set of transitively referenced workgroup variables
- VarSet GetReferencedVars(Function* func) {
- return function_to_transitive_vars.GetOrAdd(func, [&] {
- VarSet vars;
- GetReferencedVars(func->Block(), vars);
- return vars;
- });
- }
-
- /// Get the set of workgroup variables transitively referenced by @p block.
- /// @param block the block
- /// @param vars the set of transitively referenced workgroup variables to populate
- void GetReferencedVars(Block* block, VarSet& vars) {
- // Add directly referenced vars.
- if (auto itr = block_to_direct_vars.Get(block)) {
- for (auto& var : *itr) {
- vars.Add(var);
- }
- }
-
- // Loop over instructions in the block.
- for (auto* inst : *block) {
- tint::Switch(
- inst,
- [&](UserCall* call) {
- // Get variables referenced by a function called from this block.
- auto callee_vars = GetReferencedVars(call->Target());
- for (auto& var : callee_vars) {
- vars.Add(var);
- }
- },
- [&](ControlInstruction* ctrl) {
- // Recurse into control instructions and gather their referenced vars.
- ctrl->ForeachBlock([&](Block* blk) { GetReferencedVars(blk, vars); });
- });
- }
- }
-
/// Recursively generate store descriptors for a workgroup variable.
/// Determines the combined array iteration count of each inner element.
/// @param var the workgroup variable
diff --git a/src/tint/lang/core/ir/transform/zero_init_workgroup_memory_test.cc b/src/tint/lang/core/ir/transform/zero_init_workgroup_memory_test.cc
index bf1a82b..1843665 100644
--- a/src/tint/lang/core/ir/transform/zero_init_workgroup_memory_test.cc
+++ b/src/tint/lang/core/ir/transform/zero_init_workgroup_memory_test.cc
@@ -102,6 +102,48 @@
EXPECT_EQ(expect, str());
}
+TEST_F(IR_ZeroInitWorkgroupMemoryTest, NonWorkgroupVar) {
+ auto* var = b.Var("pvar", ty.ptr(private_, ty.bool_()));
+ mod.root_block->Append(var);
+
+ auto* func = MakeEntryPoint("main", 1, 1, 1);
+ b.Append(func->Block(), [&] { //
+ b.Load(var);
+ b.Return(func);
+ });
+
+ auto* src = R"(
+$B1: { # root
+ %pvar:ptr<private, bool, read_write> = var
+}
+
+%main = @compute @workgroup_size(1, 1, 1) func():void {
+ $B2: {
+ %3:bool = load %pvar
+ ret
+ }
+}
+)";
+ EXPECT_EQ(src, str());
+
+ auto* expect = R"(
+$B1: { # root
+ %pvar:ptr<private, bool, read_write> = var
+}
+
+%main = @compute @workgroup_size(1, 1, 1) func():void {
+ $B2: {
+ %3:bool = load %pvar
+ ret
+ }
+}
+)";
+
+ Run(ZeroInitWorkgroupMemory);
+
+ EXPECT_EQ(expect, str());
+}
+
TEST_F(IR_ZeroInitWorkgroupMemoryTest, ScalarBool) {
auto* var = MakeVar("wgvar", ty.bool_());
