Add ArrayLengthFromImmediates transform
ArrayLengthFromUniform() is used to support arrayLength() built-in function.
After immediate constants supported, immediate block could hold array
length and support arrayLength() built-in.
This CL add ArrayLengthFromImmediates transform to support array length
from immediate block. And will replace ArrayLengthFromUniform in future.
Bug:366291600
Change-Id: I5e5c2df0385cd2f2e846b206b652b5e963f0aad0
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/252614
Reviewed-by: James Price <jrprice@google.com>
Commit-Queue: dan sinclair <dsinclair@chromium.org>
diff --git a/src/tint/lang/core/ir/transform/BUILD.bazel b/src/tint/lang/core/ir/transform/BUILD.bazel
index f7580cb..af40537 100644
--- a/src/tint/lang/core/ir/transform/BUILD.bazel
+++ b/src/tint/lang/core/ir/transform/BUILD.bazel
@@ -40,6 +40,7 @@
name = "transform",
srcs = [
"add_empty_entry_point.cc",
+ "array_length_from_immediate.cc",
"array_length_from_uniform.cc",
"bgra8unorm_polyfill.cc",
"binary_polyfill.cc",
@@ -73,6 +74,7 @@
],
hdrs = [
"add_empty_entry_point.h",
+ "array_length_from_immediate.h",
"array_length_from_uniform.h",
"bgra8unorm_polyfill.h",
"binary_polyfill.h",
@@ -134,6 +136,7 @@
alwayslink = True,
srcs = [
"add_empty_entry_point_test.cc",
+ "array_length_from_immediate_test.cc",
"array_length_from_uniform_test.cc",
"bgra8unorm_polyfill_test.cc",
"binary_polyfill_test.cc",
diff --git a/src/tint/lang/core/ir/transform/BUILD.cmake b/src/tint/lang/core/ir/transform/BUILD.cmake
index d143d6f..8a24efc 100644
--- a/src/tint/lang/core/ir/transform/BUILD.cmake
+++ b/src/tint/lang/core/ir/transform/BUILD.cmake
@@ -41,6 +41,8 @@
tint_add_target(tint_lang_core_ir_transform lib
lang/core/ir/transform/add_empty_entry_point.cc
lang/core/ir/transform/add_empty_entry_point.h
+ lang/core/ir/transform/array_length_from_immediate.cc
+ lang/core/ir/transform/array_length_from_immediate.h
lang/core/ir/transform/array_length_from_uniform.cc
lang/core/ir/transform/array_length_from_uniform.h
lang/core/ir/transform/bgra8unorm_polyfill.cc
@@ -135,6 +137,7 @@
################################################################################
tint_add_target(tint_lang_core_ir_transform_test test
lang/core/ir/transform/add_empty_entry_point_test.cc
+ lang/core/ir/transform/array_length_from_immediate_test.cc
lang/core/ir/transform/array_length_from_uniform_test.cc
lang/core/ir/transform/bgra8unorm_polyfill_test.cc
lang/core/ir/transform/binary_polyfill_test.cc
diff --git a/src/tint/lang/core/ir/transform/BUILD.gn b/src/tint/lang/core/ir/transform/BUILD.gn
index d969108..5e18225 100644
--- a/src/tint/lang/core/ir/transform/BUILD.gn
+++ b/src/tint/lang/core/ir/transform/BUILD.gn
@@ -47,6 +47,8 @@
sources = [
"add_empty_entry_point.cc",
"add_empty_entry_point.h",
+ "array_length_from_immediate.cc",
+ "array_length_from_immediate.h",
"array_length_from_uniform.cc",
"array_length_from_uniform.h",
"bgra8unorm_polyfill.cc",
@@ -135,6 +137,7 @@
tint_unittests_source_set("unittests") {
sources = [
"add_empty_entry_point_test.cc",
+ "array_length_from_immediate_test.cc",
"array_length_from_uniform_test.cc",
"bgra8unorm_polyfill_test.cc",
"binary_polyfill_test.cc",
diff --git a/src/tint/lang/core/ir/transform/array_length_from_immediate.cc b/src/tint/lang/core/ir/transform/array_length_from_immediate.cc
new file mode 100644
index 0000000..d5a549f
--- /dev/null
+++ b/src/tint/lang/core/ir/transform/array_length_from_immediate.cc
@@ -0,0 +1,369 @@
+// Copyright 2025 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/array_length_from_immediate.h"
+
+#include <algorithm>
+#include <utility>
+
+#include "src/tint/lang/core/ir/builder.h"
+#include "src/tint/lang/core/ir/module.h"
+#include "src/tint/lang/core/ir/transform/prepare_immediate_data.h"
+#include "src/tint/lang/core/ir/validator.h"
+
+using namespace tint::core::fluent_types; // NOLINT
+using namespace tint::core::number_suffixes; // NOLINT
+
+namespace tint::core::ir::transform {
+
+namespace {
+
+/// PIMPL state for the transform.
+struct State {
+ /// The IR module.
+ Module& ir;
+
+ /// Immediate data layout contains all immediate block info.
+ const core::ir::transform::ImmediateDataLayout& immediate_data_layout;
+
+ /// The offset in immediate block for buffer sizes array.
+ uint32_t buffer_sizes_offset = 0;
+
+ /// The total number of vec4s used to store buffer sizes provided in the immediate block.
+ uint32_t buffer_sizes_array_elements_num = 0;
+
+ /// The map from binding point to the element index which holds the size of that buffer.
+ const std::unordered_map<BindingPoint, uint32_t>& bindpoint_to_size_index;
+
+ /// The IR builder.
+ core::ir::Builder b{ir};
+
+ /// The type manager.
+ core::type::Manager& ty{ir.Types()};
+
+ /// The construct instruction that creates the array lengths structure in the entry point.
+ Construct* lengths_constructor = nullptr;
+
+ /// A map from an array function parameter to the function parameter that holds its length.
+ Hashmap<FunctionParam*, FunctionParam*, 8> array_param_to_length_param{};
+
+ /// A map from a function to the structure that holds all of the array lengths.
+ Hashmap<Function*, Value*, 8> function_to_lengths_structure{};
+
+ /// A list of structure members for the array lengths structure.
+ Vector<type::Manager::StructMemberDesc, 8> lengths_structure_members{};
+
+ /// A map from a binding point to its index in the array length structure.
+ Hashmap<BindingPoint, uint32_t, 8> bindpoint_to_length_member_index{};
+
+ /// An ordered list of binding points that map to the structure members.
+ struct BindingPointInfo {
+ BindingPoint binding_point{};
+ const type::Type* store_type = nullptr;
+ };
+ Vector<BindingPointInfo, 8> ordered_bindpoints{};
+
+ /// A map from block to its containing function.
+ Hashmap<core::ir::Block*, core::ir::Function*, 64> block_to_function{};
+
+ /// Process the module.
+ void Process() {
+ // Seed the block-to-function map with the function entry blocks.
+ // This is used to determine the owning function for any given instruction.
+ for (auto& func : ir.functions) {
+ block_to_function.Add(func->Block(), func);
+ }
+
+ // Look for and replace calls to the array length builtin.
+ for (auto* inst : ir.Instructions()) {
+ if (auto* call = inst->As<CoreBuiltinCall>()) {
+ if (call->Func() == BuiltinFn::kArrayLength) {
+ MaybeReplace(call);
+ }
+ }
+ }
+
+ // Create the lengths structure and update all of the places that need to use it.
+ // We can only do this after we have replaced all of the array length callsites, now that we
+ // know all of the structure members that we need.
+ CreateLengthsStructure();
+ }
+
+ /// Replace a call to an array length builtin, if the variable appears in the bindpoint map.
+ /// @param call the arrayLength call to replace
+ void MaybeReplace(CoreBuiltinCall* call) {
+ if (auto* length = GetComputedLength(call->Args()[0], call)) {
+ call->Result()->ReplaceAllUsesWith(length);
+ call->Destroy();
+ }
+ }
+
+ /// Get the computed length value for a runtime-sized array pointer.
+ /// @param ptr the pointer to the runtime-sized array
+ /// @param insertion_point the insertion point for new instructions
+ /// @returns the computed length, or nullptr if the original builtin should be used
+ Value* GetComputedLength(Value* ptr, Instruction* insertion_point) {
+ // Trace back from the value until we reach the originating variable.
+ while (true) {
+ if (auto* param = ptr->As<FunctionParam>()) {
+ // The length of an array pointer passed as a function parameter will be passed as
+ // an additional parameter to the function.
+ return GetArrayLengthParam(param);
+ }
+
+ if (auto* result = ptr->As<InstructionResult>()) {
+ if (auto* var = result->Instruction()->As<Var>()) {
+ // We found the originating variable, so compute its array length.
+ return ComputeArrayLength(var, insertion_point);
+ }
+ if (auto* access = result->Instruction()->As<Access>()) {
+ ptr = access->Object();
+ continue;
+ }
+ if (auto* let = result->Instruction()->As<Let>()) {
+ ptr = let->Value();
+ continue;
+ }
+ TINT_UNREACHABLE() << "unhandled source of a storage buffer pointer: "
+ << result->Instruction()->TypeInfo().name;
+ }
+ TINT_UNREACHABLE() << "unhandled source of a storage buffer pointer: "
+ << ptr->TypeInfo().name;
+ }
+ }
+
+ /// Get (or create) the array length parameter that corresponds to an array parameter.
+ /// @param array_param the array parameter
+ /// @returns the array length parameter
+ FunctionParam* GetArrayLengthParam(FunctionParam* array_param) {
+ return array_param_to_length_param.GetOrAdd(array_param, [&] {
+ // Add a new parameter to receive the array length.
+ auto* length = b.FunctionParam<u32>("tint_array_length");
+ array_param->Function()->AppendParam(length);
+
+ // Update callsites of this function to pass the array length to it.
+ array_param->Function()->ForEachUseUnsorted([&](core::ir::Usage use) {
+ if (auto* call = use.instruction->As<core::ir::UserCall>()) {
+ // Get the length of the array in the calling function and pass that.
+ auto* arg = call->Args()[array_param->Index()];
+ auto* len = GetComputedLength(arg, call);
+ if (!len) {
+ // The originating variable was not in the bindpoint map, so we need to call
+ // the original arrayLength builtin as the callee is expecting a value.
+ b.InsertBefore(call, [&] {
+ len = b.Call<u32>(BuiltinFn::kArrayLength, arg)->Result();
+ });
+ }
+ call->AppendArg(len);
+ }
+ });
+
+ return length;
+ });
+ }
+
+ /// Get (or create) the array lengths structure for a function.
+ /// @param func the function that needs the structure
+ /// @returns the array lengths structure
+ Value* GetArrayLengthsStructure(Function* func) {
+ return function_to_lengths_structure.GetOrAdd(func, [&]() -> Value* {
+ if (func->IsEntryPoint()) {
+ // Create a placeholder construct instruction for the lengths structure that will be
+ // filled in later when we know all of the structure members.
+ TINT_ASSERT(lengths_constructor == nullptr);
+ lengths_constructor = b.ConstructWithResult(ir.CreateValue<InstructionResult>());
+ lengths_constructor->InsertBefore(func->Block()->Front());
+ return lengths_constructor->Result();
+ }
+
+ // Add a new parameter to receive the array lengths structure.
+ // The type is a placeholder and will be filled in later when we create the struct type.
+ auto* lengths = b.FunctionParam("tint_array_lengths", nullptr);
+ func->AppendParam(lengths);
+
+ // Update callsites of this function to pass the structure to it.
+ func->ForEachUseUnsorted([&](core::ir::Usage use) {
+ if (auto* call = use.instruction->As<core::ir::UserCall>()) {
+ // Get the structure in the calling function and pass that.
+ auto* caller = ContainingFunction(call);
+ call->AppendArg(GetArrayLengthsStructure(caller));
+ }
+ });
+
+ return lengths;
+ });
+ }
+
+ /// Compute the array length of the runtime-sized array that is inside a storage buffer
+ /// variable. If the variable's binding point is not found in the bindpoint map, returns nullptr
+ /// to indicate that the original arrayLength builtin should be used instead.
+ ///
+ /// @param var the storage buffer variable that contains the runtime-sized array
+ /// @param insertion_point the insertion point for new instructions
+ /// @returns the length of the array, or nullptr if the original builtin should be used
+ Value* ComputeArrayLength(Var* var, Instruction* insertion_point) {
+ auto binding = var->BindingPoint();
+ TINT_ASSERT(binding);
+
+ auto idx_it = bindpoint_to_size_index.find(*binding);
+ if (idx_it == bindpoint_to_size_index.end()) {
+ // If the bindpoint_to_size_index map does not contain an entry for the storage buffer,
+ // then we preserve the arrayLength() call.
+ return nullptr;
+ }
+
+ // Get the index of the structure member that holds the length for this binding point,
+ // creating the structure member descriptor if necessary.
+ auto member_index = bindpoint_to_length_member_index.GetOrAdd(*binding, [&]() {
+ auto index = static_cast<uint32_t>(lengths_structure_members.Length());
+ auto name = "tint_array_length_" + std::to_string(binding->group) + "_" +
+ std::to_string(binding->binding);
+ lengths_structure_members.Push(type::Manager::StructMemberDesc{
+ .name = ir.symbols.Register(name),
+ .type = ty.u32(),
+ });
+ ordered_bindpoints.Push(BindingPointInfo{
+ .binding_point = *binding,
+ .store_type = var->Result()->Type()->UnwrapPtr(),
+ });
+ return index;
+ });
+
+ // Extract the length from the structure.
+ auto* length = b.Access<u32>(GetArrayLengthsStructure(ContainingFunction(insertion_point)),
+ u32(member_index));
+ length->InsertBefore(insertion_point);
+ return length->Result();
+ }
+
+ /// Create the structure to hold the array lengths and fill in the construct instruction that
+ /// sets all of the length values.
+ void CreateLengthsStructure() {
+ if (lengths_structure_members.IsEmpty()) {
+ return;
+ }
+
+ // Create the lengths structure.
+ auto* lengths_struct = ty.Struct(ir.symbols.New("tint_array_lengths_struct"),
+ std::move(lengths_structure_members));
+
+ // Update all of the function parameters that need to receive the lengths structure.
+ for (auto s : function_to_lengths_structure) {
+ if (auto* param = s.value->As<FunctionParam>()) {
+ param->SetType(lengths_struct);
+ }
+ }
+
+ // Insert code at the beginning of the entry point to initialize the array length members.
+ if (lengths_constructor == nullptr) {
+ return;
+ }
+ lengths_constructor->Result()->SetType(lengths_struct);
+ b.InsertBefore(lengths_constructor->Block()->Front(), [&] {
+ Vector<Value*, 8> constructor_values;
+ for (auto info : ordered_bindpoints) {
+ TINT_ASSERT(bindpoint_to_size_index.contains(info.binding_point));
+ TINT_ASSERT(bindpoint_to_length_member_index.Contains(info.binding_point));
+
+ // Load the total storage buffer size from the immediate block.
+ // The sizes are packed into vec4s to satisfy the 16-byte alignment requirement for
+ // array elements in immediate block, so we have to find the vector and element that
+ // correspond to the index that we want.
+ const uint32_t size_index = bindpoint_to_size_index.at(info.binding_point);
+ const uint32_t array_index = size_index / 4;
+ const uint32_t vec_index = size_index % 4;
+ auto* buffer_sizes = b.Access(
+ ty.ptr(immediate, ty.array(ty.vec4<u32>(), buffer_sizes_array_elements_num)),
+ immediate_data_layout.var,
+ u32(immediate_data_layout.IndexOf(buffer_sizes_offset)));
+ auto* vec_ptr = b.Access(ty.ptr(immediate, ty.vec4<u32>()), buffer_sizes->Result(),
+ u32(array_index));
+ auto* total_buffer_size = b.LoadVectorElement(vec_ptr, u32(vec_index))->Result();
+
+ // Calculate actual array length:
+ // total_buffer_size - array_offset
+ // array_length = --------------------------------
+ // array_stride
+ auto* array_size = total_buffer_size;
+ const type::Array* array_type = nullptr;
+ if (auto* str = info.store_type->As<core::type::Struct>()) {
+ // The variable is a struct, so subtract the byte offset of the array member.
+ auto* member = str->Members().Back();
+ array_type = member->Type()->As<core::type::Array>();
+ array_size =
+ b.Subtract<u32>(total_buffer_size, u32(member->Offset()))->Result();
+ } else {
+ array_type = info.store_type->As<core::type::Array>();
+ }
+ TINT_ASSERT(array_type);
+
+ auto* length = b.Divide<u32>(array_size, u32(array_type->Stride()))->Result();
+ constructor_values.Push(length);
+ }
+ lengths_constructor->SetOperands(std::move(constructor_values));
+ });
+ }
+
+ /// Get the function that contains an instruction.
+ /// @param inst the instruction
+ /// @returns the function
+ core::ir::Function* ContainingFunction(core::ir::Instruction* inst) {
+ return block_to_function.GetOrAdd(inst->Block(), [&] { //
+ return ContainingFunction(inst->Block()->Parent());
+ });
+ }
+
+ /// @returns true if the transformed module needs a storage buffer sizes UBO
+ bool NeedsStorageBufferSizes() {
+ return !lengths_structure_members.IsEmpty() && lengths_constructor != nullptr;
+ }
+};
+
+} // namespace
+
+Result<ArrayLengthFromImmediateResult> ArrayLengthFromImmediates(
+ Module& ir,
+ const core::ir::transform::ImmediateDataLayout& immediate_data_layout,
+ const uint32_t buffer_sizes_offset,
+ const uint32_t buffer_sizes_array_elements_num,
+ const std::unordered_map<BindingPoint, uint32_t>& bindpoint_to_size_index) {
+ auto validated = ValidateAndDumpIfNeeded(ir, "core.ArrayLengthFromImmediates",
+ kArrayLengthFromImmediateCapabilities);
+ if (validated != Success) {
+ return validated.Failure();
+ }
+
+ State state{ir, immediate_data_layout, buffer_sizes_offset, buffer_sizes_array_elements_num,
+ bindpoint_to_size_index};
+ state.Process();
+
+ ArrayLengthFromImmediateResult result;
+ result.needs_storage_buffer_sizes = state.NeedsStorageBufferSizes();
+ return result;
+}
+
+} // namespace tint::core::ir::transform
diff --git a/src/tint/lang/core/ir/transform/array_length_from_immediate.h b/src/tint/lang/core/ir/transform/array_length_from_immediate.h
new file mode 100644
index 0000000..eeacc41
--- /dev/null
+++ b/src/tint/lang/core/ir/transform/array_length_from_immediate.h
@@ -0,0 +1,89 @@
+// Copyright 2025 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_ARRAY_LENGTH_FROM_IMMEDIATE_H_
+#define SRC_TINT_LANG_CORE_IR_TRANSFORM_ARRAY_LENGTH_FROM_IMMEDIATE_H_
+
+#include <unordered_map>
+
+#include "src/tint/api/common/binding_point.h"
+#include "src/tint/lang/core/ir/validator.h"
+#include "src/tint/utils/result.h"
+
+// Forward declarations.
+namespace tint::core::ir {
+class Module;
+}
+
+namespace tint::core::ir::transform {
+struct ImmediateDataLayout;
+
+/// The capabilities that the transform can support.
+const Capabilities kArrayLengthFromImmediateCapabilities{Capability::kAllowDuplicateBindings};
+
+/// The result of running the ArrayLengthFromUniform transform.
+struct ArrayLengthFromImmediateResult {
+ /// `true` if the transformed module needs the storage buffer sizes immediate data.
+ bool needs_storage_buffer_sizes = false;
+};
+
+/// ArrayLengthFromImmediates is a transform that replaces calls to the arrayLength() builtin by
+/// calculating the array length from the total size of the storage buffer, which is received via
+/// immediate blocks.
+///
+/// The generated immediate blocks will have the form:
+/// ```
+/// @group(0) @binding(30)
+/// struct tint_immediate_data_struct {
+/// ...
+/// buffer_sizes: array<vec4<u32>, 8>; // offset is provided via config
+// };
+/// var<immediate> tint_immediate_data : tint_immediate_data_struct;
+/// ```
+/// The offset of `buffer_sizes` in the immediate block is provided by config.
+/// The transform config also defines the mapping from a storage buffer's `BindingPoint` to the
+/// element index that will be used to get the size of that buffer.
+///
+/// @param module the module to transform
+/// @param immediate_data_layout The immediate data layout information.
+/// @param bindpoint_to_size_index The map from binding point to an index which holds the size
+/// of that buffer.
+/// @param buffer_sizes_offset The offset in immediate block where buffer sizes start.
+/// @param buffer_sizes_array_elements_num the number of vec4s used to store buffer sizes that will
+/// be set into the immediate block.
+/// @returns the transform result or failure
+/// TODO(crbug.com/366291600): Replace ArrayLengthFromUniform.
+Result<ArrayLengthFromImmediateResult> ArrayLengthFromImmediates(
+ Module& module,
+ const core::ir::transform::ImmediateDataLayout& immediate_data_layout,
+ const uint32_t buffer_sizes_offset,
+ const uint32_t buffer_sizes_array_elements_num,
+ const std::unordered_map<BindingPoint, uint32_t>& bindpoint_to_size_index);
+
+} // namespace tint::core::ir::transform
+
+#endif // SRC_TINT_LANG_CORE_IR_TRANSFORM_ARRAY_LENGTH_FROM_IMMEDIATE_H_
diff --git a/src/tint/lang/core/ir/transform/array_length_from_immediate_test.cc b/src/tint/lang/core/ir/transform/array_length_from_immediate_test.cc
new file mode 100644
index 0000000..d289d9d
--- /dev/null
+++ b/src/tint/lang/core/ir/transform/array_length_from_immediate_test.cc
@@ -0,0 +1,1580 @@
+// Copyright 2025 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/array_length_from_immediate.h"
+
+#include <algorithm>
+#include <utility>
+
+#include "src/tint/lang/core/ir/transform/helper_test.h"
+#include "src/tint/lang/core/ir/transform/prepare_immediate_data.h"
+
+namespace tint::core::ir::transform {
+namespace {
+
+using namespace tint::core::fluent_types; // NOLINT
+using namespace tint::core::number_suffixes; // NOLINT
+
+using IR_ArrayLengthFromImmediatesTest = TransformTest;
+
+uint32_t GetBufferSizesNumElements(
+ const std::unordered_map<BindingPoint, uint32_t>& bindpoint_to_size_index) {
+ uint32_t max_index = 0;
+ for (auto& entry : bindpoint_to_size_index) {
+ max_index = std::max(max_index, entry.second);
+ }
+ return (max_index / 4) + 1;
+}
+
+TEST_F(IR_ArrayLengthFromImmediatesTest, NoModify_UserFunction) {
+ auto* arr = ty.array<i32>();
+ auto* arr_ptr = ty.ptr<storage>(arr);
+
+ auto* buffer = b.Var("buffer", arr_ptr);
+ buffer->SetBindingPoint(0, 0);
+ mod.root_block->Append(buffer);
+
+ auto* user_func = b.Function("arrayLength", ty.u32());
+ auto* param = b.FunctionParam("arr", arr_ptr);
+ user_func->SetParams({param});
+ b.Append(user_func->Block(), [&] { //
+ b.Return(user_func, 42_u);
+ });
+
+ auto* func = b.Function("foo", ty.void_());
+ b.Append(func->Block(), [&] {
+ b.Call(user_func, buffer);
+ b.Return(func);
+ });
+
+ auto* src = R"(
+$B1: { # root
+ %buffer:ptr<storage, array<i32>, read_write> = var undef @binding_point(0, 0)
+}
+
+%arrayLength = func(%arr:ptr<storage, array<i32>, read_write>):u32 {
+ $B2: {
+ ret 42u
+ }
+}
+%foo = func():void {
+ $B3: {
+ %5:u32 = call %arrayLength, %buffer
+ ret
+ }
+}
+)";
+ EXPECT_EQ(src, str());
+
+ auto* expect = R"(
+$B1: { # root
+ %buffer:ptr<storage, array<i32>, read_write> = var undef @binding_point(0, 0)
+}
+
+%arrayLength = func(%arr:ptr<storage, array<i32>, read_write>):u32 {
+ $B2: {
+ ret 42u
+ }
+}
+%foo = func():void {
+ $B3: {
+ %5:u32 = call %arrayLength, %buffer
+ ret
+ }
+}
+)";
+
+ std::unordered_map<BindingPoint, uint32_t> bindpoint_to_index;
+ bindpoint_to_index[{0, 0}] = 0;
+
+ core::ir::transform::PrepareImmediateDataConfig immediate_data_config;
+ auto immediate_data = PrepareImmediateData(mod, {});
+ EXPECT_EQ(immediate_data, Success);
+
+ Run(ArrayLengthFromImmediates, immediate_data.Get(), 0u, 0u, bindpoint_to_index);
+
+ EXPECT_EQ(expect, str());
+}
+
+TEST_F(IR_ArrayLengthFromImmediatesTest, DirectUse) {
+ auto* arr = ty.array<i32>();
+ auto* arr_ptr = ty.ptr<storage>(arr);
+
+ auto* buffer = b.Var("buffer", arr_ptr);
+ buffer->SetBindingPoint(0, 0);
+ mod.root_block->Append(buffer);
+
+ auto* func = b.ComputeFunction("foo");
+ b.Append(func->Block(), [&] {
+ auto* len = b.Call<u32>(BuiltinFn::kArrayLength, buffer);
+ b.Let("let", len);
+ b.Return(func);
+ });
+
+ auto* src = R"(
+$B1: { # root
+ %buffer:ptr<storage, array<i32>, read_write> = var undef @binding_point(0, 0)
+}
+
+%foo = @compute @workgroup_size(1u, 1u, 1u) func():void {
+ $B2: {
+ %3:u32 = arrayLength %buffer
+ %let:u32 = let %3
+ ret
+ }
+}
+)";
+ EXPECT_EQ(src, str());
+
+ auto* expect = R"(
+tint_immediate_data_struct = struct @align(16), @block {
+ tint_storage_buffer_sizes:array<vec4<u32>, 1> @offset(16)
+}
+
+tint_array_lengths_struct = struct @align(4) {
+ tint_array_length_0_0:u32 @offset(0)
+}
+
+$B1: { # root
+ %buffer:ptr<storage, array<i32>, read_write> = var undef @binding_point(0, 0)
+ %tint_immediate_data:ptr<immediate, tint_immediate_data_struct, read> = var undef
+}
+
+%foo = @compute @workgroup_size(1u, 1u, 1u) func():void {
+ $B2: {
+ %4:ptr<immediate, array<vec4<u32>, 1>, read> = access %tint_immediate_data, 0u
+ %5:ptr<immediate, vec4<u32>, read> = access %4, 0u
+ %6:u32 = load_vector_element %5, 0u
+ %7:u32 = div %6, 4u
+ %8:tint_array_lengths_struct = construct %7
+ %9:u32 = access %8, 0u
+ %let:u32 = let %9
+ ret
+ }
+}
+)";
+
+ std::unordered_map<BindingPoint, uint32_t> bindpoint_to_index;
+ bindpoint_to_index[{0, 0}] = 0;
+
+ core::ir::transform::PrepareImmediateDataConfig immediate_data_config;
+ constexpr uint32_t buffer_size_start_offset = 16;
+ uint32_t num_elements = GetBufferSizesNumElements(bindpoint_to_index);
+ immediate_data_config.AddInternalImmediateData(buffer_size_start_offset,
+ mod.symbols.New("tint_storage_buffer_sizes"),
+ ty.array(ty.vec4<u32>(), num_elements));
+ auto immediate_data = PrepareImmediateData(mod, immediate_data_config);
+ EXPECT_EQ(immediate_data, Success);
+
+ Run(ArrayLengthFromImmediates, immediate_data.Get(), buffer_size_start_offset, num_elements,
+ bindpoint_to_index);
+
+ EXPECT_EQ(expect, str());
+}
+
+TEST_F(IR_ArrayLengthFromImmediatesTest, DirectUse_NonZeroIndex) {
+ auto* arr = ty.array<i32>();
+ auto* arr_ptr = ty.ptr<storage>(arr);
+
+ auto* buffer = b.Var("buffer", arr_ptr);
+ buffer->SetBindingPoint(0, 0);
+ mod.root_block->Append(buffer);
+
+ auto* func = b.ComputeFunction("foo");
+ b.Append(func->Block(), [&] {
+ auto* len = b.Call<u32>(BuiltinFn::kArrayLength, buffer);
+ b.Let("let", len);
+ b.Return(func);
+ });
+
+ auto* src = R"(
+$B1: { # root
+ %buffer:ptr<storage, array<i32>, read_write> = var undef @binding_point(0, 0)
+}
+
+%foo = @compute @workgroup_size(1u, 1u, 1u) func():void {
+ $B2: {
+ %3:u32 = arrayLength %buffer
+ %let:u32 = let %3
+ ret
+ }
+}
+)";
+ EXPECT_EQ(src, str());
+
+ auto* expect = R"(
+tint_immediate_data_struct = struct @align(16), @block {
+ tint_storage_buffer_sizes:array<vec4<u32>, 2> @offset(16)
+}
+
+tint_array_lengths_struct = struct @align(4) {
+ tint_array_length_0_0:u32 @offset(0)
+}
+
+$B1: { # root
+ %buffer:ptr<storage, array<i32>, read_write> = var undef @binding_point(0, 0)
+ %tint_immediate_data:ptr<immediate, tint_immediate_data_struct, read> = var undef
+}
+
+%foo = @compute @workgroup_size(1u, 1u, 1u) func():void {
+ $B2: {
+ %4:ptr<immediate, array<vec4<u32>, 2>, read> = access %tint_immediate_data, 0u
+ %5:ptr<immediate, vec4<u32>, read> = access %4, 1u
+ %6:u32 = load_vector_element %5, 3u
+ %7:u32 = div %6, 4u
+ %8:tint_array_lengths_struct = construct %7
+ %9:u32 = access %8, 0u
+ %let:u32 = let %9
+ ret
+ }
+}
+)";
+
+ std::unordered_map<BindingPoint, uint32_t> bindpoint_to_index;
+ bindpoint_to_index[{0, 0}] = 7;
+
+ core::ir::transform::PrepareImmediateDataConfig immediate_data_config;
+ constexpr uint32_t buffer_size_start_offset = 16;
+ uint32_t num_elements = GetBufferSizesNumElements(bindpoint_to_index);
+ immediate_data_config.AddInternalImmediateData(buffer_size_start_offset,
+ mod.symbols.New("tint_storage_buffer_sizes"),
+ ty.array(ty.vec4<u32>(), num_elements));
+ auto immediate_data = PrepareImmediateData(mod, immediate_data_config);
+ EXPECT_EQ(immediate_data, Success);
+
+ Run(ArrayLengthFromImmediates, immediate_data.Get(), buffer_size_start_offset, num_elements,
+ bindpoint_to_index);
+
+ EXPECT_EQ(expect, str());
+}
+
+TEST_F(IR_ArrayLengthFromImmediatesTest, DirectUse_NotInMap) {
+ auto* arr = ty.array<i32>();
+ auto* arr_ptr = ty.ptr<storage>(arr);
+
+ auto* buffer = b.Var("buffer", arr_ptr);
+ buffer->SetBindingPoint(0, 1);
+ mod.root_block->Append(buffer);
+
+ auto* func = b.ComputeFunction("foo");
+ b.Append(func->Block(), [&] {
+ auto* len = b.Call<u32>(BuiltinFn::kArrayLength, buffer);
+ b.Let("let", len);
+ b.Return(func);
+ });
+
+ auto* src = R"(
+$B1: { # root
+ %buffer:ptr<storage, array<i32>, read_write> = var undef @binding_point(0, 1)
+}
+
+%foo = @compute @workgroup_size(1u, 1u, 1u) func():void {
+ $B2: {
+ %3:u32 = arrayLength %buffer
+ %let:u32 = let %3
+ ret
+ }
+}
+)";
+ EXPECT_EQ(src, str());
+
+ auto* expect = R"(
+$B1: { # root
+ %buffer:ptr<storage, array<i32>, read_write> = var undef @binding_point(0, 1)
+}
+
+%foo = @compute @workgroup_size(1u, 1u, 1u) func():void {
+ $B2: {
+ %3:u32 = arrayLength %buffer
+ %let:u32 = let %3
+ ret
+ }
+}
+)";
+
+ std::unordered_map<BindingPoint, uint32_t> bindpoint_to_index;
+
+ auto immediate_data = PrepareImmediateData(mod, {});
+ EXPECT_EQ(immediate_data, Success);
+
+ Run(ArrayLengthFromImmediates, immediate_data.Get(), 0u, 0u, bindpoint_to_index);
+
+ EXPECT_EQ(expect, str());
+}
+
+TEST_F(IR_ArrayLengthFromImmediatesTest, DirectUse_NoEntryPoint) {
+ auto* arr = ty.array<i32>();
+ auto* arr_ptr = ty.ptr<storage>(arr);
+
+ auto* buffer = b.Var("buffer", arr_ptr);
+ buffer->SetBindingPoint(0, 0);
+ mod.root_block->Append(buffer);
+
+ auto* bar = b.Function("bar", ty.u32());
+ b.Append(bar->Block(), [&] {
+ auto* len = b.Call<u32>(BuiltinFn::kArrayLength, buffer);
+ b.Return(bar, len);
+ });
+
+ auto* src = R"(
+$B1: { # root
+ %buffer:ptr<storage, array<i32>, read_write> = var undef @binding_point(0, 0)
+}
+
+%bar = func():u32 {
+ $B2: {
+ %3:u32 = arrayLength %buffer
+ ret %3
+ }
+}
+)";
+ EXPECT_EQ(src, str());
+
+ auto* expect = R"(
+tint_array_lengths_struct = struct @align(4) {
+ tint_array_length_0_0:u32 @offset(0)
+}
+
+$B1: { # root
+ %buffer:ptr<storage, array<i32>, read_write> = var undef @binding_point(0, 0)
+}
+
+%bar = func(%tint_array_lengths:tint_array_lengths_struct):u32 {
+ $B2: {
+ %4:u32 = access %tint_array_lengths, 0u
+ ret %4
+ }
+}
+)";
+ std::unordered_map<BindingPoint, uint32_t> bindpoint_to_index;
+ bindpoint_to_index[{0, 0}] = 0;
+
+ auto immediate_data = PrepareImmediateData(mod, {});
+ EXPECT_EQ(immediate_data, Success);
+
+ Run(ArrayLengthFromImmediates, immediate_data.Get(), 0u, 0u, bindpoint_to_index);
+
+ EXPECT_EQ(expect, str());
+}
+
+TEST_F(IR_ArrayLengthFromImmediatesTest, DirectUse_CalledByEntryPoint) {
+ auto* arr = ty.array<i32>();
+ auto* arr_ptr = ty.ptr<storage>(arr);
+
+ auto* buffer = b.Var("buffer", arr_ptr);
+ buffer->SetBindingPoint(0, 0);
+ mod.root_block->Append(buffer);
+
+ auto* bar = b.Function("bar", ty.u32());
+ b.Append(bar->Block(), [&] {
+ auto* len = b.Call<u32>(BuiltinFn::kArrayLength, buffer);
+ b.Return(bar, len);
+ });
+
+ auto* foo = b.ComputeFunction("foo");
+ b.Append(foo->Block(), [&] {
+ auto* len = b.Call<u32>(bar);
+ b.Let("let", len);
+ b.Return(foo);
+ });
+
+ auto* src = R"(
+$B1: { # root
+ %buffer:ptr<storage, array<i32>, read_write> = var undef @binding_point(0, 0)
+}
+
+%bar = func():u32 {
+ $B2: {
+ %3:u32 = arrayLength %buffer
+ ret %3
+ }
+}
+%foo = @compute @workgroup_size(1u, 1u, 1u) func():void {
+ $B3: {
+ %5:u32 = call %bar
+ %let:u32 = let %5
+ ret
+ }
+}
+)";
+ EXPECT_EQ(src, str());
+
+ auto* expect = R"(
+tint_immediate_data_struct = struct @align(16), @block {
+ tint_storage_buffer_sizes:array<vec4<u32>, 1> @offset(16)
+}
+
+tint_array_lengths_struct = struct @align(4) {
+ tint_array_length_0_0:u32 @offset(0)
+}
+
+$B1: { # root
+ %buffer:ptr<storage, array<i32>, read_write> = var undef @binding_point(0, 0)
+ %tint_immediate_data:ptr<immediate, tint_immediate_data_struct, read> = var undef
+}
+
+%bar = func(%tint_array_lengths:tint_array_lengths_struct):u32 {
+ $B2: {
+ %5:u32 = access %tint_array_lengths, 0u
+ ret %5
+ }
+}
+%foo = @compute @workgroup_size(1u, 1u, 1u) func():void {
+ $B3: {
+ %7:ptr<immediate, array<vec4<u32>, 1>, read> = access %tint_immediate_data, 0u
+ %8:ptr<immediate, vec4<u32>, read> = access %7, 0u
+ %9:u32 = load_vector_element %8, 0u
+ %10:u32 = div %9, 4u
+ %11:tint_array_lengths_struct = construct %10
+ %12:u32 = call %bar, %11
+ %let:u32 = let %12
+ ret
+ }
+}
+)";
+
+ std::unordered_map<BindingPoint, uint32_t> bindpoint_to_index;
+ bindpoint_to_index[{0, 0}] = 0;
+
+ core::ir::transform::PrepareImmediateDataConfig immediate_data_config;
+ constexpr uint32_t buffer_size_start_offset = 16;
+ uint32_t num_elements = GetBufferSizesNumElements(bindpoint_to_index);
+ immediate_data_config.AddInternalImmediateData(buffer_size_start_offset,
+ mod.symbols.New("tint_storage_buffer_sizes"),
+ ty.array(ty.vec4<u32>(), num_elements));
+ auto immediate_data = PrepareImmediateData(mod, immediate_data_config);
+ EXPECT_EQ(immediate_data, Success);
+
+ Run(ArrayLengthFromImmediates, immediate_data.Get(), buffer_size_start_offset, num_elements,
+ bindpoint_to_index);
+
+ EXPECT_EQ(expect, str());
+}
+
+TEST_F(IR_ArrayLengthFromImmediatesTest, ViaAccess_StructMember) {
+ auto* arr = ty.array<i32>();
+ auto* structure = ty.Struct(mod.symbols.New("MyStruct"), {
+ {mod.symbols.New("a"), arr},
+ });
+ auto* arr_ptr = ty.ptr<storage>(arr);
+ auto* structure_ptr = ty.ptr<storage>(structure);
+
+ auto* buffer = b.Var("buffer", structure_ptr);
+ buffer->SetBindingPoint(0, 0);
+ mod.root_block->Append(buffer);
+
+ auto* func = b.ComputeFunction("foo");
+ b.Append(func->Block(), [&] {
+ auto* len = b.Call<u32>(BuiltinFn::kArrayLength, b.Access(arr_ptr, buffer, 0_u));
+ b.Let("let", len);
+ b.Return(func);
+ });
+
+ auto* src = R"(
+MyStruct = struct @align(4) {
+ a:array<i32> @offset(0)
+}
+
+$B1: { # root
+ %buffer:ptr<storage, MyStruct, read_write> = var undef @binding_point(0, 0)
+}
+
+%foo = @compute @workgroup_size(1u, 1u, 1u) func():void {
+ $B2: {
+ %3:ptr<storage, array<i32>, read_write> = access %buffer, 0u
+ %4:u32 = arrayLength %3
+ %let:u32 = let %4
+ ret
+ }
+}
+)";
+ EXPECT_EQ(src, str());
+
+ auto* expect = R"(
+MyStruct = struct @align(4) {
+ a:array<i32> @offset(0)
+}
+
+tint_immediate_data_struct = struct @align(16), @block {
+ tint_storage_buffer_sizes:array<vec4<u32>, 1> @offset(16)
+}
+
+tint_array_lengths_struct = struct @align(4) {
+ tint_array_length_0_0:u32 @offset(0)
+}
+
+$B1: { # root
+ %buffer:ptr<storage, MyStruct, read_write> = var undef @binding_point(0, 0)
+ %tint_immediate_data:ptr<immediate, tint_immediate_data_struct, read> = var undef
+}
+
+%foo = @compute @workgroup_size(1u, 1u, 1u) func():void {
+ $B2: {
+ %4:ptr<immediate, array<vec4<u32>, 1>, read> = access %tint_immediate_data, 0u
+ %5:ptr<immediate, vec4<u32>, read> = access %4, 0u
+ %6:u32 = load_vector_element %5, 0u
+ %7:u32 = sub %6, 0u
+ %8:u32 = div %7, 4u
+ %9:tint_array_lengths_struct = construct %8
+ %10:ptr<storage, array<i32>, read_write> = access %buffer, 0u
+ %11:u32 = access %9, 0u
+ %let:u32 = let %11
+ ret
+ }
+}
+)";
+
+ std::unordered_map<BindingPoint, uint32_t> bindpoint_to_index;
+ bindpoint_to_index[{0, 0}] = 0;
+
+ core::ir::transform::PrepareImmediateDataConfig immediate_data_config;
+ constexpr uint32_t buffer_size_start_offset = 16;
+ uint32_t num_elements = GetBufferSizesNumElements(bindpoint_to_index);
+ immediate_data_config.AddInternalImmediateData(buffer_size_start_offset,
+ mod.symbols.New("tint_storage_buffer_sizes"),
+ ty.array(ty.vec4<u32>(), num_elements));
+ auto immediate_data = PrepareImmediateData(mod, immediate_data_config);
+ EXPECT_EQ(immediate_data, Success);
+
+ Run(ArrayLengthFromImmediates, immediate_data.Get(), buffer_size_start_offset, num_elements,
+ bindpoint_to_index);
+
+ EXPECT_EQ(expect, str());
+}
+
+TEST_F(IR_ArrayLengthFromImmediatesTest, ViaAccess_StructMember_NonZeroOffset) {
+ auto* arr = ty.array<i32>();
+ auto* structure = ty.Struct(mod.symbols.New("MyStruct"), {
+ {mod.symbols.New("u1"), ty.u32()},
+ {mod.symbols.New("u2"), ty.u32()},
+ {mod.symbols.New("u3"), ty.u32()},
+ {mod.symbols.New("u4"), ty.u32()},
+ {mod.symbols.New("u5"), ty.u32()},
+ {mod.symbols.New("a"), arr},
+ });
+ auto* arr_ptr = ty.ptr<storage>(arr);
+ auto* structure_ptr = ty.ptr<storage>(structure);
+
+ auto* buffer = b.Var("buffer", structure_ptr);
+ buffer->SetBindingPoint(0, 0);
+ mod.root_block->Append(buffer);
+
+ auto* func = b.ComputeFunction("foo");
+ b.Append(func->Block(), [&] {
+ auto* len = b.Call<u32>(BuiltinFn::kArrayLength, b.Access(arr_ptr, buffer, 5_u));
+ b.Let("let", len);
+ b.Return(func);
+ });
+
+ auto* src = R"(
+MyStruct = struct @align(4) {
+ u1:u32 @offset(0)
+ u2:u32 @offset(4)
+ u3:u32 @offset(8)
+ u4:u32 @offset(12)
+ u5:u32 @offset(16)
+ a:array<i32> @offset(20)
+}
+
+$B1: { # root
+ %buffer:ptr<storage, MyStruct, read_write> = var undef @binding_point(0, 0)
+}
+
+%foo = @compute @workgroup_size(1u, 1u, 1u) func():void {
+ $B2: {
+ %3:ptr<storage, array<i32>, read_write> = access %buffer, 5u
+ %4:u32 = arrayLength %3
+ %let:u32 = let %4
+ ret
+ }
+}
+)";
+ EXPECT_EQ(src, str());
+
+ auto* expect = R"(
+MyStruct = struct @align(4) {
+ u1:u32 @offset(0)
+ u2:u32 @offset(4)
+ u3:u32 @offset(8)
+ u4:u32 @offset(12)
+ u5:u32 @offset(16)
+ a:array<i32> @offset(20)
+}
+
+tint_immediate_data_struct = struct @align(16), @block {
+ tint_storage_buffer_sizes:array<vec4<u32>, 1> @offset(16)
+}
+
+tint_array_lengths_struct = struct @align(4) {
+ tint_array_length_0_0:u32 @offset(0)
+}
+
+$B1: { # root
+ %buffer:ptr<storage, MyStruct, read_write> = var undef @binding_point(0, 0)
+ %tint_immediate_data:ptr<immediate, tint_immediate_data_struct, read> = var undef
+}
+
+%foo = @compute @workgroup_size(1u, 1u, 1u) func():void {
+ $B2: {
+ %4:ptr<immediate, array<vec4<u32>, 1>, read> = access %tint_immediate_data, 0u
+ %5:ptr<immediate, vec4<u32>, read> = access %4, 0u
+ %6:u32 = load_vector_element %5, 0u
+ %7:u32 = sub %6, 20u
+ %8:u32 = div %7, 4u
+ %9:tint_array_lengths_struct = construct %8
+ %10:ptr<storage, array<i32>, read_write> = access %buffer, 5u
+ %11:u32 = access %9, 0u
+ %let:u32 = let %11
+ ret
+ }
+}
+)";
+
+ std::unordered_map<BindingPoint, uint32_t> bindpoint_to_index;
+ bindpoint_to_index[{0, 0}] = 0;
+
+ core::ir::transform::PrepareImmediateDataConfig immediate_data_config;
+ constexpr uint32_t buffer_size_start_offset = 16;
+ uint32_t num_elements = GetBufferSizesNumElements(bindpoint_to_index);
+ immediate_data_config.AddInternalImmediateData(buffer_size_start_offset,
+ mod.symbols.New("tint_storage_buffer_sizes"),
+ ty.array(ty.vec4<u32>(), num_elements));
+ auto immediate_data = PrepareImmediateData(mod, immediate_data_config);
+ EXPECT_EQ(immediate_data, Success);
+
+ Run(ArrayLengthFromImmediates, immediate_data.Get(), buffer_size_start_offset, num_elements,
+ bindpoint_to_index);
+
+ EXPECT_EQ(expect, str());
+}
+
+TEST_F(IR_ArrayLengthFromImmediatesTest, ViaLet) {
+ auto* arr = ty.array<i32>();
+ auto* arr_ptr = ty.ptr<storage>(arr);
+
+ auto* buffer = b.Var("buffer", arr_ptr);
+ buffer->SetBindingPoint(0, 0);
+ mod.root_block->Append(buffer);
+
+ auto* func = b.ComputeFunction("foo");
+ b.Append(func->Block(), [&] {
+ auto* len = b.Call<u32>(BuiltinFn::kArrayLength, b.Let("let", buffer));
+ b.Let("let", len);
+ b.Return(func);
+ });
+
+ auto* src = R"(
+$B1: { # root
+ %buffer:ptr<storage, array<i32>, read_write> = var undef @binding_point(0, 0)
+}
+
+%foo = @compute @workgroup_size(1u, 1u, 1u) func():void {
+ $B2: {
+ %let:ptr<storage, array<i32>, read_write> = let %buffer
+ %4:u32 = arrayLength %let
+ %let_1:u32 = let %4 # %let_1: 'let'
+ ret
+ }
+}
+)";
+ EXPECT_EQ(src, str());
+
+ auto* expect = R"(
+tint_immediate_data_struct = struct @align(16), @block {
+ tint_storage_buffer_sizes:array<vec4<u32>, 1> @offset(16)
+}
+
+tint_array_lengths_struct = struct @align(4) {
+ tint_array_length_0_0:u32 @offset(0)
+}
+
+$B1: { # root
+ %buffer:ptr<storage, array<i32>, read_write> = var undef @binding_point(0, 0)
+ %tint_immediate_data:ptr<immediate, tint_immediate_data_struct, read> = var undef
+}
+
+%foo = @compute @workgroup_size(1u, 1u, 1u) func():void {
+ $B2: {
+ %4:ptr<immediate, array<vec4<u32>, 1>, read> = access %tint_immediate_data, 0u
+ %5:ptr<immediate, vec4<u32>, read> = access %4, 0u
+ %6:u32 = load_vector_element %5, 0u
+ %7:u32 = div %6, 4u
+ %8:tint_array_lengths_struct = construct %7
+ %let:ptr<storage, array<i32>, read_write> = let %buffer
+ %10:u32 = access %8, 0u
+ %let_1:u32 = let %10 # %let_1: 'let'
+ ret
+ }
+}
+)";
+
+ std::unordered_map<BindingPoint, uint32_t> bindpoint_to_index;
+ bindpoint_to_index[{0, 0}] = 0;
+
+ core::ir::transform::PrepareImmediateDataConfig immediate_data_config;
+ constexpr uint32_t buffer_size_start_offset = 16;
+ uint32_t num_elements = GetBufferSizesNumElements(bindpoint_to_index);
+ immediate_data_config.AddInternalImmediateData(buffer_size_start_offset,
+ mod.symbols.New("tint_storage_buffer_sizes"),
+ ty.array(ty.vec4<u32>(), num_elements));
+ auto immediate_data = PrepareImmediateData(mod, immediate_data_config);
+ EXPECT_EQ(immediate_data, Success);
+
+ Run(ArrayLengthFromImmediates, immediate_data.Get(), buffer_size_start_offset, num_elements,
+ bindpoint_to_index);
+
+ EXPECT_EQ(expect, str());
+}
+
+TEST_F(IR_ArrayLengthFromImmediatesTest, ViaParameter) {
+ auto* arr = ty.array<i32>();
+ auto* arr_ptr = ty.ptr<storage>(arr);
+
+ auto* buffer = b.Var("buffer", arr_ptr);
+ buffer->SetBindingPoint(0, 0);
+ mod.root_block->Append(buffer);
+
+ auto* bar = b.Function("bar", ty.u32());
+ auto* param = b.FunctionParam("param", arr_ptr);
+ bar->SetParams({param});
+ b.Append(bar->Block(), [&] {
+ auto* len = b.Call<u32>(BuiltinFn::kArrayLength, param);
+ b.Return(bar, len);
+ });
+
+ auto* foo = b.ComputeFunction("foo");
+ b.Append(foo->Block(), [&] {
+ auto* len = b.Call<u32>(bar, buffer);
+ b.Let("let", len);
+ b.Return(foo);
+ });
+
+ auto* src = R"(
+$B1: { # root
+ %buffer:ptr<storage, array<i32>, read_write> = var undef @binding_point(0, 0)
+}
+
+%bar = func(%param:ptr<storage, array<i32>, read_write>):u32 {
+ $B2: {
+ %4:u32 = arrayLength %param
+ ret %4
+ }
+}
+%foo = @compute @workgroup_size(1u, 1u, 1u) func():void {
+ $B3: {
+ %6:u32 = call %bar, %buffer
+ %let:u32 = let %6
+ ret
+ }
+}
+)";
+ EXPECT_EQ(src, str());
+
+ auto* expect = R"(
+tint_immediate_data_struct = struct @align(16), @block {
+ tint_storage_buffer_sizes:array<vec4<u32>, 1> @offset(16)
+}
+
+tint_array_lengths_struct = struct @align(4) {
+ tint_array_length_0_0:u32 @offset(0)
+}
+
+$B1: { # root
+ %buffer:ptr<storage, array<i32>, read_write> = var undef @binding_point(0, 0)
+ %tint_immediate_data:ptr<immediate, tint_immediate_data_struct, read> = var undef
+}
+
+%bar = func(%param:ptr<storage, array<i32>, read_write>, %tint_array_length:u32):u32 {
+ $B2: {
+ ret %tint_array_length
+ }
+}
+%foo = @compute @workgroup_size(1u, 1u, 1u) func():void {
+ $B3: {
+ %7:ptr<immediate, array<vec4<u32>, 1>, read> = access %tint_immediate_data, 0u
+ %8:ptr<immediate, vec4<u32>, read> = access %7, 0u
+ %9:u32 = load_vector_element %8, 0u
+ %10:u32 = div %9, 4u
+ %11:tint_array_lengths_struct = construct %10
+ %12:u32 = access %11, 0u
+ %13:u32 = call %bar, %buffer, %12
+ %let:u32 = let %13
+ ret
+ }
+}
+)";
+
+ std::unordered_map<BindingPoint, uint32_t> bindpoint_to_index;
+ bindpoint_to_index[{0, 0}] = 0;
+
+ core::ir::transform::PrepareImmediateDataConfig immediate_data_config;
+ constexpr uint32_t buffer_size_start_offset = 16;
+ uint32_t num_elements = GetBufferSizesNumElements(bindpoint_to_index);
+ immediate_data_config.AddInternalImmediateData(buffer_size_start_offset,
+ mod.symbols.New("tint_storage_buffer_sizes"),
+ ty.array(ty.vec4<u32>(), num_elements));
+ auto immediate_data = PrepareImmediateData(mod, immediate_data_config);
+ EXPECT_EQ(immediate_data, Success);
+
+ Run(ArrayLengthFromImmediates, immediate_data.Get(), buffer_size_start_offset, num_elements,
+ bindpoint_to_index);
+
+ EXPECT_EQ(expect, str());
+}
+
+TEST_F(IR_ArrayLengthFromImmediatesTest, ViaParameterChain) {
+ auto* arr = ty.array<i32>();
+ auto* arr_ptr = ty.ptr<storage>(arr);
+
+ auto* buffer = b.Var("buffer", arr_ptr);
+ buffer->SetBindingPoint(0, 0);
+ mod.root_block->Append(buffer);
+
+ auto* zoo = b.Function("foo", ty.u32());
+ auto* param_zoo = b.FunctionParam("param_zoo", arr_ptr);
+ zoo->SetParams({param_zoo});
+ b.Append(zoo->Block(), [&] {
+ auto* len = b.Call<u32>(BuiltinFn::kArrayLength, param_zoo);
+ b.Return(zoo, len);
+ });
+
+ auto* bar = b.Function("foo", ty.u32());
+ auto* param_bar = b.FunctionParam("param_bar", arr_ptr);
+ bar->SetParams({param_bar});
+ b.Append(bar->Block(), [&] {
+ auto* len = b.Call<u32>(zoo, param_bar);
+ b.Return(bar, len);
+ });
+
+ auto* foo = b.ComputeFunction("foo");
+ b.Append(foo->Block(), [&] {
+ auto* len = b.Call<u32>(bar, buffer);
+ b.Let("let", len);
+ b.Return(foo);
+ });
+
+ auto* src = R"(
+$B1: { # root
+ %buffer:ptr<storage, array<i32>, read_write> = var undef @binding_point(0, 0)
+}
+
+%foo = func(%param_zoo:ptr<storage, array<i32>, read_write>):u32 {
+ $B2: {
+ %4:u32 = arrayLength %param_zoo
+ ret %4
+ }
+}
+%foo_1 = func(%param_bar:ptr<storage, array<i32>, read_write>):u32 { # %foo_1: 'foo'
+ $B3: {
+ %7:u32 = call %foo, %param_bar
+ ret %7
+ }
+}
+%foo_2 = @compute @workgroup_size(1u, 1u, 1u) func():void { # %foo_2: 'foo'
+ $B4: {
+ %9:u32 = call %foo_1, %buffer
+ %let:u32 = let %9
+ ret
+ }
+}
+)";
+ EXPECT_EQ(src, str());
+
+ auto* expect = R"(
+tint_immediate_data_struct = struct @align(16), @block {
+ tint_storage_buffer_sizes:array<vec4<u32>, 1> @offset(16)
+}
+
+tint_array_lengths_struct = struct @align(4) {
+ tint_array_length_0_0:u32 @offset(0)
+}
+
+$B1: { # root
+ %buffer:ptr<storage, array<i32>, read_write> = var undef @binding_point(0, 0)
+ %tint_immediate_data:ptr<immediate, tint_immediate_data_struct, read> = var undef
+}
+
+%foo = func(%param_zoo:ptr<storage, array<i32>, read_write>, %tint_array_length:u32):u32 {
+ $B2: {
+ ret %tint_array_length
+ }
+}
+%foo_1 = func(%param_bar:ptr<storage, array<i32>, read_write>, %tint_array_length_1:u32):u32 { # %foo_1: 'foo', %tint_array_length_1: 'tint_array_length'
+ $B3: {
+ %9:u32 = call %foo, %param_bar, %tint_array_length_1
+ ret %9
+ }
+}
+%foo_2 = @compute @workgroup_size(1u, 1u, 1u) func():void { # %foo_2: 'foo'
+ $B4: {
+ %11:ptr<immediate, array<vec4<u32>, 1>, read> = access %tint_immediate_data, 0u
+ %12:ptr<immediate, vec4<u32>, read> = access %11, 0u
+ %13:u32 = load_vector_element %12, 0u
+ %14:u32 = div %13, 4u
+ %15:tint_array_lengths_struct = construct %14
+ %16:u32 = access %15, 0u
+ %17:u32 = call %foo_1, %buffer, %16
+ %let:u32 = let %17
+ ret
+ }
+}
+)";
+
+ std::unordered_map<BindingPoint, uint32_t> bindpoint_to_index;
+ bindpoint_to_index[{0, 0}] = 0;
+
+ core::ir::transform::PrepareImmediateDataConfig immediate_data_config;
+ constexpr uint32_t buffer_size_start_offset = 16;
+ uint32_t num_elements = GetBufferSizesNumElements(bindpoint_to_index);
+ immediate_data_config.AddInternalImmediateData(buffer_size_start_offset,
+ mod.symbols.New("tint_storage_buffer_sizes"),
+ ty.array(ty.vec4<u32>(), num_elements));
+ auto immediate_data = PrepareImmediateData(mod, immediate_data_config);
+ EXPECT_EQ(immediate_data, Success);
+
+ Run(ArrayLengthFromImmediates, immediate_data.Get(), buffer_size_start_offset, num_elements,
+ bindpoint_to_index);
+
+ EXPECT_EQ(expect, str());
+}
+
+// Test that when we arrayLength is called on a parameter but the originating variable at the
+// callsite is not in the bindpoint map, we reintroduce an arrayLength call instead of passing
+// undef to the callee.
+TEST_F(IR_ArrayLengthFromImmediatesTest, ViaParameter_NotInMap) {
+ auto* arr = ty.array<i32>();
+ auto* arr_ptr = ty.ptr<storage>(arr);
+
+ auto* buffer = b.Var("buffer", arr_ptr);
+ buffer->SetBindingPoint(0, 0);
+ mod.root_block->Append(buffer);
+
+ auto* bar = b.Function("bar", ty.u32());
+ auto* param = b.FunctionParam("param", arr_ptr);
+ bar->SetParams({param});
+ b.Append(bar->Block(), [&] {
+ auto* len = b.Call<u32>(BuiltinFn::kArrayLength, param);
+ b.Return(bar, len);
+ });
+
+ auto* foo = b.ComputeFunction("foo");
+ b.Append(foo->Block(), [&] {
+ auto* len = b.Call<u32>(bar, buffer);
+ b.Let("let", len);
+ b.Return(foo);
+ });
+
+ auto* src = R"(
+$B1: { # root
+ %buffer:ptr<storage, array<i32>, read_write> = var undef @binding_point(0, 0)
+}
+
+%bar = func(%param:ptr<storage, array<i32>, read_write>):u32 {
+ $B2: {
+ %4:u32 = arrayLength %param
+ ret %4
+ }
+}
+%foo = @compute @workgroup_size(1u, 1u, 1u) func():void {
+ $B3: {
+ %6:u32 = call %bar, %buffer
+ %let:u32 = let %6
+ ret
+ }
+}
+)";
+ EXPECT_EQ(src, str());
+
+ auto* expect = R"(
+$B1: { # root
+ %buffer:ptr<storage, array<i32>, read_write> = var undef @binding_point(0, 0)
+}
+
+%bar = func(%param:ptr<storage, array<i32>, read_write>, %tint_array_length:u32):u32 {
+ $B2: {
+ ret %tint_array_length
+ }
+}
+%foo = @compute @workgroup_size(1u, 1u, 1u) func():void {
+ $B3: {
+ %6:u32 = arrayLength %buffer
+ %7:u32 = call %bar, %buffer, %6
+ %let:u32 = let %7
+ ret
+ }
+}
+)";
+
+ std::unordered_map<BindingPoint, uint32_t> bindpoint_to_index;
+ auto immediate_data = PrepareImmediateData(mod, {});
+ EXPECT_EQ(immediate_data, Success);
+
+ Run(ArrayLengthFromImmediates, immediate_data.Get(), 0u, 0u, bindpoint_to_index);
+
+ EXPECT_EQ(expect, str());
+}
+
+// Test that we reuse the length parameter for multiple arrayLength calls on the same parameter.
+TEST_F(IR_ArrayLengthFromImmediatesTest, ViaParameter_MultipleCallsSameParameter) {
+ auto* arr = ty.array<i32>();
+ auto* arr_ptr = ty.ptr<storage>(arr);
+
+ auto* buffer = b.Var("buffer", arr_ptr);
+ buffer->SetBindingPoint(0, 0);
+ mod.root_block->Append(buffer);
+
+ auto* bar = b.Function("bar", ty.u32());
+ auto* param = b.FunctionParam("param", arr_ptr);
+ bar->SetParams({param});
+ b.Append(bar->Block(), [&] {
+ auto* len_a = b.Call<u32>(BuiltinFn::kArrayLength, param);
+ auto* len_b = b.Call<u32>(BuiltinFn::kArrayLength, param);
+ auto* len_c = b.Call<u32>(BuiltinFn::kArrayLength, param);
+ b.Return(bar, b.Add<u32>(len_a, b.Add<u32>(len_b, len_c)));
+ });
+
+ auto* foo = b.ComputeFunction("foo");
+ b.Append(foo->Block(), [&] {
+ auto* len = b.Call<u32>(bar, buffer);
+ b.Let("let", len);
+ b.Return(foo);
+ });
+
+ auto* src = R"(
+$B1: { # root
+ %buffer:ptr<storage, array<i32>, read_write> = var undef @binding_point(0, 0)
+}
+
+%bar = func(%param:ptr<storage, array<i32>, read_write>):u32 {
+ $B2: {
+ %4:u32 = arrayLength %param
+ %5:u32 = arrayLength %param
+ %6:u32 = arrayLength %param
+ %7:u32 = add %5, %6
+ %8:u32 = add %4, %7
+ ret %8
+ }
+}
+%foo = @compute @workgroup_size(1u, 1u, 1u) func():void {
+ $B3: {
+ %10:u32 = call %bar, %buffer
+ %let:u32 = let %10
+ ret
+ }
+}
+)";
+ EXPECT_EQ(src, str());
+
+ auto* expect = R"(
+tint_immediate_data_struct = struct @align(16), @block {
+ tint_storage_buffer_sizes:array<vec4<u32>, 1> @offset(16)
+}
+
+tint_array_lengths_struct = struct @align(4) {
+ tint_array_length_0_0:u32 @offset(0)
+}
+
+$B1: { # root
+ %buffer:ptr<storage, array<i32>, read_write> = var undef @binding_point(0, 0)
+ %tint_immediate_data:ptr<immediate, tint_immediate_data_struct, read> = var undef
+}
+
+%bar = func(%param:ptr<storage, array<i32>, read_write>, %tint_array_length:u32):u32 {
+ $B2: {
+ %6:u32 = add %tint_array_length, %tint_array_length
+ %7:u32 = add %tint_array_length, %6
+ ret %7
+ }
+}
+%foo = @compute @workgroup_size(1u, 1u, 1u) func():void {
+ $B3: {
+ %9:ptr<immediate, array<vec4<u32>, 1>, read> = access %tint_immediate_data, 0u
+ %10:ptr<immediate, vec4<u32>, read> = access %9, 0u
+ %11:u32 = load_vector_element %10, 0u
+ %12:u32 = div %11, 4u
+ %13:tint_array_lengths_struct = construct %12
+ %14:u32 = access %13, 0u
+ %15:u32 = call %bar, %buffer, %14
+ %let:u32 = let %15
+ ret
+ }
+}
+)";
+
+ std::unordered_map<BindingPoint, uint32_t> bindpoint_to_index;
+ bindpoint_to_index[{0, 0}] = 0;
+
+ core::ir::transform::PrepareImmediateDataConfig immediate_data_config;
+ constexpr uint32_t buffer_size_start_offset = 16;
+ uint32_t num_elements = GetBufferSizesNumElements(bindpoint_to_index);
+ immediate_data_config.AddInternalImmediateData(buffer_size_start_offset,
+ mod.symbols.New("tint_storage_buffer_sizes"),
+ ty.array(ty.vec4<u32>(), num_elements));
+ auto immediate_data = PrepareImmediateData(mod, immediate_data_config);
+ EXPECT_EQ(immediate_data, Success);
+
+ Run(ArrayLengthFromImmediates, immediate_data.Get(), buffer_size_start_offset, num_elements,
+ bindpoint_to_index);
+
+ EXPECT_EQ(expect, str());
+}
+
+TEST_F(IR_ArrayLengthFromImmediatesTest, ViaParameter_MultipleCallsDifferentParameters) {
+ auto* arr = ty.array<i32>();
+ auto* arr_ptr = ty.ptr<storage>(arr);
+
+ auto* buffer = b.Var("buffer", arr_ptr);
+ buffer->SetBindingPoint(0, 0);
+ mod.root_block->Append(buffer);
+
+ auto* bar = b.Function("bar", ty.u32());
+ auto* param_a = b.FunctionParam("param_a", arr_ptr);
+ auto* param_b = b.FunctionParam("param_b", arr_ptr);
+ auto* param_c = b.FunctionParam("param_c", arr_ptr);
+ bar->SetParams({param_a, param_b, param_c});
+ b.Append(bar->Block(), [&] {
+ auto* len_a = b.Call<u32>(BuiltinFn::kArrayLength, param_a);
+ auto* len_b = b.Call<u32>(BuiltinFn::kArrayLength, param_b);
+ auto* len_c = b.Call<u32>(BuiltinFn::kArrayLength, param_c);
+ b.Return(bar, b.Add<u32>(len_a, b.Add<u32>(len_b, len_c)));
+ });
+
+ auto* foo = b.ComputeFunction("foo");
+ b.Append(foo->Block(), [&] {
+ auto* len = b.Call<u32>(bar, buffer, buffer, buffer);
+ b.Let("let", len);
+ b.Return(foo);
+ });
+
+ auto* src = R"(
+$B1: { # root
+ %buffer:ptr<storage, array<i32>, read_write> = var undef @binding_point(0, 0)
+}
+
+%bar = func(%param_a:ptr<storage, array<i32>, read_write>, %param_b:ptr<storage, array<i32>, read_write>, %param_c:ptr<storage, array<i32>, read_write>):u32 {
+ $B2: {
+ %6:u32 = arrayLength %param_a
+ %7:u32 = arrayLength %param_b
+ %8:u32 = arrayLength %param_c
+ %9:u32 = add %7, %8
+ %10:u32 = add %6, %9
+ ret %10
+ }
+}
+%foo = @compute @workgroup_size(1u, 1u, 1u) func():void {
+ $B3: {
+ %12:u32 = call %bar, %buffer, %buffer, %buffer
+ %let:u32 = let %12
+ ret
+ }
+}
+)";
+ EXPECT_EQ(src, str());
+
+ auto* expect = R"(
+tint_immediate_data_struct = struct @align(16), @block {
+ tint_storage_buffer_sizes:array<vec4<u32>, 1> @offset(16)
+}
+
+tint_array_lengths_struct = struct @align(4) {
+ tint_array_length_0_0:u32 @offset(0)
+}
+
+$B1: { # root
+ %buffer:ptr<storage, array<i32>, read_write> = var undef @binding_point(0, 0)
+ %tint_immediate_data:ptr<immediate, tint_immediate_data_struct, read> = var undef
+}
+
+%bar = func(%param_a:ptr<storage, array<i32>, read_write>, %param_b:ptr<storage, array<i32>, read_write>, %param_c:ptr<storage, array<i32>, read_write>, %tint_array_length:u32, %tint_array_length_1:u32, %tint_array_length_2:u32):u32 { # %tint_array_length_1: 'tint_array_length', %tint_array_length_2: 'tint_array_length'
+ $B2: {
+ %10:u32 = add %tint_array_length_1, %tint_array_length_2
+ %11:u32 = add %tint_array_length, %10
+ ret %11
+ }
+}
+%foo = @compute @workgroup_size(1u, 1u, 1u) func():void {
+ $B3: {
+ %13:ptr<immediate, array<vec4<u32>, 1>, read> = access %tint_immediate_data, 0u
+ %14:ptr<immediate, vec4<u32>, read> = access %13, 0u
+ %15:u32 = load_vector_element %14, 0u
+ %16:u32 = div %15, 4u
+ %17:tint_array_lengths_struct = construct %16
+ %18:u32 = access %17, 0u
+ %19:u32 = access %17, 0u
+ %20:u32 = access %17, 0u
+ %21:u32 = call %bar, %buffer, %buffer, %buffer, %18, %19, %20
+ %let:u32 = let %21
+ ret
+ }
+}
+)";
+
+ std::unordered_map<BindingPoint, uint32_t> bindpoint_to_index;
+ bindpoint_to_index[{0, 0}] = 0;
+
+ core::ir::transform::PrepareImmediateDataConfig immediate_data_config;
+ constexpr uint32_t buffer_size_start_offset = 16;
+ uint32_t num_elements = GetBufferSizesNumElements(bindpoint_to_index);
+ immediate_data_config.AddInternalImmediateData(buffer_size_start_offset,
+ mod.symbols.New("tint_storage_buffer_sizes"),
+ ty.array(ty.vec4<u32>(), num_elements));
+ auto immediate_data = PrepareImmediateData(mod, immediate_data_config);
+ EXPECT_EQ(immediate_data, Success);
+ Run(ArrayLengthFromImmediates, immediate_data.Get(), buffer_size_start_offset, num_elements,
+ bindpoint_to_index);
+
+ EXPECT_EQ(expect, str());
+}
+
+TEST_F(IR_ArrayLengthFromImmediatesTest, ViaComplexChain) {
+ auto* arr = ty.array<i32>();
+ auto* structure = ty.Struct(mod.symbols.New("MyStruct"), {
+ {mod.symbols.New("u1"), ty.u32()},
+ {mod.symbols.New("a"), arr},
+ });
+ auto* arr_ptr = ty.ptr<storage>(arr);
+ auto* structure_ptr = ty.ptr<storage>(structure);
+
+ auto* buffer = b.Var("buffer", structure_ptr);
+ buffer->SetBindingPoint(0, 0);
+ mod.root_block->Append(buffer);
+
+ auto* bar = b.Function("bar", ty.u32());
+ auto* param = b.FunctionParam("param", arr_ptr);
+ bar->SetParams({param});
+ b.Append(bar->Block(), [&] {
+ auto* let1 = b.Let("let1", param);
+ auto* let2 = b.Let("let2", let1);
+ auto* len = b.Call<u32>(BuiltinFn::kArrayLength, let2);
+ b.Return(bar, len);
+ });
+
+ auto* foo = b.ComputeFunction("foo");
+ b.Append(foo->Block(), [&] {
+ auto* access = b.Access(arr_ptr, buffer, 1_u);
+ auto* let = b.Let("let", access);
+ auto* len = b.Call<u32>(bar, let);
+ b.Let("let", len);
+ b.Return(foo);
+ });
+
+ auto* src = R"(
+MyStruct = struct @align(4) {
+ u1:u32 @offset(0)
+ a:array<i32> @offset(4)
+}
+
+$B1: { # root
+ %buffer:ptr<storage, MyStruct, read_write> = var undef @binding_point(0, 0)
+}
+
+%bar = func(%param:ptr<storage, array<i32>, read_write>):u32 {
+ $B2: {
+ %let1:ptr<storage, array<i32>, read_write> = let %param
+ %let2:ptr<storage, array<i32>, read_write> = let %let1
+ %6:u32 = arrayLength %let2
+ ret %6
+ }
+}
+%foo = @compute @workgroup_size(1u, 1u, 1u) func():void {
+ $B3: {
+ %8:ptr<storage, array<i32>, read_write> = access %buffer, 1u
+ %let:ptr<storage, array<i32>, read_write> = let %8
+ %10:u32 = call %bar, %let
+ %let_1:u32 = let %10 # %let_1: 'let'
+ ret
+ }
+}
+)";
+ EXPECT_EQ(src, str());
+
+ auto* expect = R"(
+MyStruct = struct @align(4) {
+ u1:u32 @offset(0)
+ a:array<i32> @offset(4)
+}
+
+tint_immediate_data_struct = struct @align(16), @block {
+ tint_storage_buffer_sizes:array<vec4<u32>, 1> @offset(16)
+}
+
+tint_array_lengths_struct = struct @align(4) {
+ tint_array_length_0_0:u32 @offset(0)
+}
+
+$B1: { # root
+ %buffer:ptr<storage, MyStruct, read_write> = var undef @binding_point(0, 0)
+ %tint_immediate_data:ptr<immediate, tint_immediate_data_struct, read> = var undef
+}
+
+%bar = func(%param:ptr<storage, array<i32>, read_write>, %tint_array_length:u32):u32 {
+ $B2: {
+ %let1:ptr<storage, array<i32>, read_write> = let %param
+ %let2:ptr<storage, array<i32>, read_write> = let %let1
+ ret %tint_array_length
+ }
+}
+%foo = @compute @workgroup_size(1u, 1u, 1u) func():void {
+ $B3: {
+ %9:ptr<immediate, array<vec4<u32>, 1>, read> = access %tint_immediate_data, 0u
+ %10:ptr<immediate, vec4<u32>, read> = access %9, 0u
+ %11:u32 = load_vector_element %10, 0u
+ %12:u32 = sub %11, 4u
+ %13:u32 = div %12, 4u
+ %14:tint_array_lengths_struct = construct %13
+ %15:ptr<storage, array<i32>, read_write> = access %buffer, 1u
+ %let:ptr<storage, array<i32>, read_write> = let %15
+ %17:u32 = access %14, 0u
+ %18:u32 = call %bar, %let, %17
+ %let_1:u32 = let %18 # %let_1: 'let'
+ ret
+ }
+}
+)";
+
+ std::unordered_map<BindingPoint, uint32_t> bindpoint_to_index;
+ bindpoint_to_index[{0, 0}] = 0;
+
+ core::ir::transform::PrepareImmediateDataConfig immediate_data_config;
+ constexpr uint32_t buffer_size_start_offset = 16;
+ uint32_t num_elements = GetBufferSizesNumElements(bindpoint_to_index);
+ immediate_data_config.AddInternalImmediateData(buffer_size_start_offset,
+ mod.symbols.New("tint_storage_buffer_sizes"),
+ ty.array(ty.vec4<u32>(), num_elements));
+ auto immediate_data = PrepareImmediateData(mod, immediate_data_config);
+ EXPECT_EQ(immediate_data, Success);
+ Run(ArrayLengthFromImmediates, immediate_data.Get(), buffer_size_start_offset, num_elements,
+ bindpoint_to_index);
+
+ EXPECT_EQ(expect, str());
+}
+
+TEST_F(IR_ArrayLengthFromImmediatesTest, ElementStrideLargerThanSize) {
+ auto* arr = ty.array<vec3<i32>>();
+ auto* arr_ptr = ty.ptr<storage>(arr);
+
+ auto* buffer = b.Var("buffer", arr_ptr);
+ buffer->SetBindingPoint(0, 0);
+ mod.root_block->Append(buffer);
+
+ auto* func = b.ComputeFunction("foo");
+ b.Append(func->Block(), [&] {
+ auto* len = b.Call<u32>(BuiltinFn::kArrayLength, buffer);
+ b.Let("let", len);
+ b.Return(func);
+ });
+
+ auto* src = R"(
+$B1: { # root
+ %buffer:ptr<storage, array<vec3<i32>>, read_write> = var undef @binding_point(0, 0)
+}
+
+%foo = @compute @workgroup_size(1u, 1u, 1u) func():void {
+ $B2: {
+ %3:u32 = arrayLength %buffer
+ %let:u32 = let %3
+ ret
+ }
+}
+)";
+ EXPECT_EQ(src, str());
+
+ auto* expect = R"(
+tint_immediate_data_struct = struct @align(16), @block {
+ tint_storage_buffer_sizes:array<vec4<u32>, 1> @offset(16)
+}
+
+tint_array_lengths_struct = struct @align(4) {
+ tint_array_length_0_0:u32 @offset(0)
+}
+
+$B1: { # root
+ %buffer:ptr<storage, array<vec3<i32>>, read_write> = var undef @binding_point(0, 0)
+ %tint_immediate_data:ptr<immediate, tint_immediate_data_struct, read> = var undef
+}
+
+%foo = @compute @workgroup_size(1u, 1u, 1u) func():void {
+ $B2: {
+ %4:ptr<immediate, array<vec4<u32>, 1>, read> = access %tint_immediate_data, 0u
+ %5:ptr<immediate, vec4<u32>, read> = access %4, 0u
+ %6:u32 = load_vector_element %5, 0u
+ %7:u32 = div %6, 16u
+ %8:tint_array_lengths_struct = construct %7
+ %9:u32 = access %8, 0u
+ %let:u32 = let %9
+ ret
+ }
+}
+)";
+
+ std::unordered_map<BindingPoint, uint32_t> bindpoint_to_index;
+ bindpoint_to_index[{0, 0}] = 0;
+
+ core::ir::transform::PrepareImmediateDataConfig immediate_data_config;
+ constexpr uint32_t buffer_size_start_offset = 16;
+ uint32_t num_elements = GetBufferSizesNumElements(bindpoint_to_index);
+ immediate_data_config.AddInternalImmediateData(buffer_size_start_offset,
+ mod.symbols.New("tint_storage_buffer_sizes"),
+ ty.array(ty.vec4<u32>(), num_elements));
+ auto immediate_data = PrepareImmediateData(mod, immediate_data_config);
+ EXPECT_EQ(immediate_data, Success);
+ Run(ArrayLengthFromImmediates, immediate_data.Get(), buffer_size_start_offset, num_elements,
+ bindpoint_to_index);
+
+ EXPECT_EQ(expect, str());
+}
+
+TEST_F(IR_ArrayLengthFromImmediatesTest, MultipleVars) {
+ auto* arr = ty.array<i32>();
+ auto* arr_ptr = ty.ptr<storage>(arr);
+
+ auto* buffer_a = b.Var("buffer_a", arr_ptr);
+ auto* buffer_b = b.Var("buffer_b", arr_ptr);
+ auto* buffer_c = b.Var("buffer_c", arr_ptr);
+ auto* buffer_d = b.Var("buffer_d", arr_ptr);
+ auto* buffer_e = b.Var("buffer_e", arr_ptr);
+ buffer_a->SetBindingPoint(0, 0);
+ buffer_b->SetBindingPoint(0, 1);
+ buffer_c->SetBindingPoint(1, 0);
+ buffer_d->SetBindingPoint(1, 1);
+ buffer_e->SetBindingPoint(2, 3);
+ mod.root_block->Append(buffer_a);
+ mod.root_block->Append(buffer_b);
+ mod.root_block->Append(buffer_c);
+ mod.root_block->Append(buffer_d);
+ mod.root_block->Append(buffer_e);
+
+ auto* func = b.ComputeFunction("foo");
+ b.Append(func->Block(), [&] {
+ b.Call<u32>(BuiltinFn::kArrayLength, buffer_a);
+ b.Call<u32>(BuiltinFn::kArrayLength, buffer_b);
+ b.Call<u32>(BuiltinFn::kArrayLength, buffer_c);
+ b.Call<u32>(BuiltinFn::kArrayLength, buffer_d);
+ b.Call<u32>(BuiltinFn::kArrayLength, buffer_e);
+ b.Return(func);
+ });
+
+ auto* src = R"(
+$B1: { # root
+ %buffer_a:ptr<storage, array<i32>, read_write> = var undef @binding_point(0, 0)
+ %buffer_b:ptr<storage, array<i32>, read_write> = var undef @binding_point(0, 1)
+ %buffer_c:ptr<storage, array<i32>, read_write> = var undef @binding_point(1, 0)
+ %buffer_d:ptr<storage, array<i32>, read_write> = var undef @binding_point(1, 1)
+ %buffer_e:ptr<storage, array<i32>, read_write> = var undef @binding_point(2, 3)
+}
+
+%foo = @compute @workgroup_size(1u, 1u, 1u) func():void {
+ $B2: {
+ %7:u32 = arrayLength %buffer_a
+ %8:u32 = arrayLength %buffer_b
+ %9:u32 = arrayLength %buffer_c
+ %10:u32 = arrayLength %buffer_d
+ %11:u32 = arrayLength %buffer_e
+ ret
+ }
+}
+)";
+ EXPECT_EQ(src, str());
+
+ auto* expect = R"(
+tint_immediate_data_struct = struct @align(16), @block {
+ tint_storage_buffer_sizes:array<vec4<u32>, 2> @offset(16)
+}
+
+tint_array_lengths_struct = struct @align(4) {
+ tint_array_length_0_0:u32 @offset(0)
+ tint_array_length_0_1:u32 @offset(4)
+ tint_array_length_1_0:u32 @offset(8)
+ tint_array_length_1_1:u32 @offset(12)
+ tint_array_length_2_3:u32 @offset(16)
+}
+
+$B1: { # root
+ %buffer_a:ptr<storage, array<i32>, read_write> = var undef @binding_point(0, 0)
+ %buffer_b:ptr<storage, array<i32>, read_write> = var undef @binding_point(0, 1)
+ %buffer_c:ptr<storage, array<i32>, read_write> = var undef @binding_point(1, 0)
+ %buffer_d:ptr<storage, array<i32>, read_write> = var undef @binding_point(1, 1)
+ %buffer_e:ptr<storage, array<i32>, read_write> = var undef @binding_point(2, 3)
+ %tint_immediate_data:ptr<immediate, tint_immediate_data_struct, read> = var undef
+}
+
+%foo = @compute @workgroup_size(1u, 1u, 1u) func():void {
+ $B2: {
+ %8:ptr<immediate, array<vec4<u32>, 2>, read> = access %tint_immediate_data, 0u
+ %9:ptr<immediate, vec4<u32>, read> = access %8, 0u
+ %10:u32 = load_vector_element %9, 0u
+ %11:u32 = div %10, 4u
+ %12:ptr<immediate, array<vec4<u32>, 2>, read> = access %tint_immediate_data, 0u
+ %13:ptr<immediate, vec4<u32>, read> = access %12, 1u
+ %14:u32 = load_vector_element %13, 1u
+ %15:u32 = div %14, 4u
+ %16:ptr<immediate, array<vec4<u32>, 2>, read> = access %tint_immediate_data, 0u
+ %17:ptr<immediate, vec4<u32>, read> = access %16, 0u
+ %18:u32 = load_vector_element %17, 3u
+ %19:u32 = div %18, 4u
+ %20:ptr<immediate, array<vec4<u32>, 2>, read> = access %tint_immediate_data, 0u
+ %21:ptr<immediate, vec4<u32>, read> = access %20, 0u
+ %22:u32 = load_vector_element %21, 2u
+ %23:u32 = div %22, 4u
+ %24:ptr<immediate, array<vec4<u32>, 2>, read> = access %tint_immediate_data, 0u
+ %25:ptr<immediate, vec4<u32>, read> = access %24, 1u
+ %26:u32 = load_vector_element %25, 0u
+ %27:u32 = div %26, 4u
+ %28:tint_array_lengths_struct = construct %11, %15, %19, %23, %27
+ %29:u32 = access %28, 0u
+ %30:u32 = access %28, 1u
+ %31:u32 = access %28, 2u
+ %32:u32 = access %28, 3u
+ %33:u32 = access %28, 4u
+ ret
+ }
+}
+)";
+
+ std::unordered_map<BindingPoint, uint32_t> bindpoint_to_index;
+ bindpoint_to_index[{0, 0}] = 0;
+ bindpoint_to_index[{0, 1}] = 5;
+ bindpoint_to_index[{1, 0}] = 3;
+ bindpoint_to_index[{1, 1}] = 2;
+ bindpoint_to_index[{2, 3}] = 4;
+
+ core::ir::transform::PrepareImmediateDataConfig immediate_data_config;
+ constexpr uint32_t buffer_size_start_offset = 16;
+ uint32_t num_elements = GetBufferSizesNumElements(bindpoint_to_index);
+ immediate_data_config.AddInternalImmediateData(buffer_size_start_offset,
+ mod.symbols.New("tint_storage_buffer_sizes"),
+ ty.array(ty.vec4<u32>(), num_elements));
+ auto immediate_data = PrepareImmediateData(mod, immediate_data_config);
+ EXPECT_EQ(immediate_data, Success);
+ Run(ArrayLengthFromImmediates, immediate_data.Get(), buffer_size_start_offset, num_elements,
+ bindpoint_to_index);
+
+ EXPECT_EQ(expect, str());
+}
+
+} // namespace
+} // namespace tint::core::ir::transform
diff --git a/src/tint/lang/core/ir/transform/prepare_immediate_data.h b/src/tint/lang/core/ir/transform/prepare_immediate_data.h
index 6d26694..29c2253 100644
--- a/src/tint/lang/core/ir/transform/prepare_immediate_data.h
+++ b/src/tint/lang/core/ir/transform/prepare_immediate_data.h
@@ -58,7 +58,7 @@
core::ir::Var* var = nullptr;
/// A map from member offset to member index.
- Hashmap<uint32_t, uint32_t, 4> offset_to_index;
+ Hashmap<uint32_t, uint32_t, 6> offset_to_index;
/// @returns the member index of the constant at @p offset
uint32_t IndexOf(uint32_t offset) const {
