| // Copyright 2021 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/wgsl/ast/transform/array_length_from_uniform.h" |
| |
| #include <cstdint> |
| #include <memory> |
| #include <string> |
| #include <string_view> |
| #include <utility> |
| |
| #include "src/tint/lang/core/fluent_types.h" |
| #include "src/tint/lang/core/unary_op.h" |
| #include "src/tint/lang/wgsl/ast/expression.h" |
| #include "src/tint/lang/wgsl/ast/transform/simplify_pointers.h" |
| #include "src/tint/lang/wgsl/ast/unary_op_expression.h" |
| #include "src/tint/lang/wgsl/ast/variable.h" |
| #include "src/tint/lang/wgsl/builtin_fn.h" |
| #include "src/tint/lang/wgsl/program/clone_context.h" |
| #include "src/tint/lang/wgsl/program/program_builder.h" |
| #include "src/tint/lang/wgsl/resolver/resolve.h" |
| #include "src/tint/lang/wgsl/sem/array.h" |
| #include "src/tint/lang/wgsl/sem/builtin_fn.h" |
| #include "src/tint/lang/wgsl/sem/call.h" |
| #include "src/tint/lang/wgsl/sem/expression.h" |
| #include "src/tint/lang/wgsl/sem/function.h" |
| #include "src/tint/lang/wgsl/sem/member_accessor_expression.h" |
| #include "src/tint/lang/wgsl/sem/statement.h" |
| #include "src/tint/lang/wgsl/sem/variable.h" |
| #include "src/tint/utils/containers/unique_vector.h" |
| #include "src/tint/utils/diagnostic/diagnostic.h" |
| #include "src/tint/utils/ice/ice.h" |
| #include "src/tint/utils/rtti/switch.h" |
| #include "src/tint/utils/text/text_style.h" |
| |
| TINT_INSTANTIATE_TYPEINFO(tint::ast::transform::ArrayLengthFromUniform); |
| TINT_INSTANTIATE_TYPEINFO(tint::ast::transform::ArrayLengthFromUniform::Config); |
| TINT_INSTANTIATE_TYPEINFO(tint::ast::transform::ArrayLengthFromUniform::Result); |
| |
| using namespace tint::core::fluent_types; // NOLINT |
| |
| namespace tint::ast::transform { |
| namespace { |
| |
| bool ShouldRun(const Program& program) { |
| for (auto* fn : program.AST().Functions()) { |
| if (auto* sem_fn = program.Sem().Get(fn)) { |
| for (auto* builtin : sem_fn->DirectlyCalledBuiltins()) { |
| if (builtin->Fn() == wgsl::BuiltinFn::kArrayLength) { |
| return true; |
| } |
| } |
| } |
| } |
| return false; |
| } |
| |
| } // namespace |
| |
| ArrayLengthFromUniform::ArrayLengthFromUniform() = default; |
| ArrayLengthFromUniform::~ArrayLengthFromUniform() = default; |
| |
| /// PIMPL state for the transform |
| struct ArrayLengthFromUniform::State { |
| /// Constructor |
| /// @param program the source program |
| /// @param in the input transform data |
| /// @param out the output transform data |
| State(const Program& program, const DataMap& in, DataMap& out) |
| : src(program), outputs(out), cfg(in.Get<Config>()) {} |
| |
| /// Runs the transform |
| /// @returns the new program or SkipTransform if the transform is not required |
| ApplyResult Run() { |
| if (cfg == nullptr) { |
| b.Diagnostics().AddError(Source{}) << "missing transform data for " |
| << tint::TypeInfo::Of<ArrayLengthFromUniform>().name; |
| return resolver::Resolve(b); |
| } |
| |
| if (cfg->bindpoint_to_size_index.empty() || !ShouldRun(src)) { |
| return SkipTransform; |
| } |
| |
| // Create the name of the array lengths uniform variable. |
| array_lengths_var = b.Symbols().New("tint_array_lengths"); |
| |
| // Replace all the arrayLength() calls. |
| for (auto* fn : src.AST().Functions()) { |
| if (auto* sem_fn = sem.Get(fn)) { |
| for (auto* call : sem_fn->DirectCalls()) { |
| if (auto* target = call->Target()->As<sem::BuiltinFn>()) { |
| if (target->Fn() == wgsl::BuiltinFn::kArrayLength) { |
| ReplaceArrayLengthCall(call); |
| } |
| } |
| } |
| } |
| } |
| |
| // Add the necessary array-length arguments to all the newly created array-length |
| // parameters. |
| while (!len_params_needing_args.IsEmpty()) { |
| AddArrayLengthArguments(len_params_needing_args.Pop()); |
| } |
| |
| // Add the tint_array_lengths module-scope uniform variable. |
| AddArrayLengthsUniformVar(); |
| |
| outputs.Add<Result>(used_size_indices); |
| |
| ctx.Clone(); |
| return resolver::Resolve(b); |
| } |
| |
| private: |
| // Replaces the arrayLength() builtin call with an array-length expression passed via a uniform |
| // buffer. |
| void ReplaceArrayLengthCall(const sem::Call* call) { |
| if (auto* replacement = ArrayLengthOf(call->Arguments()[0])) { |
| ctx.Replace(call->Declaration(), replacement); |
| } |
| } |
| |
| /// @returns an AST expression that is equal to the arrayLength() of the runtime-sized array |
| /// accessed by the pointer expression @p expr, or nullptr on error or if the array is not in |
| /// the Config::bindpoint_to_size_index map. |
| const ast::Expression* ArrayLengthOf(const sem::Expression* expr) { |
| const ast::Expression* len = nullptr; |
| while (expr) { |
| expr = Switch( |
| expr, // |
| [&](const sem::VariableUser* user) { |
| len = ArrayLengthOf(user->Variable()); |
| return nullptr; |
| }, |
| [&](const sem::MemberAccessorExpression* access) { |
| return access->Object(); // Follow the object |
| }, |
| [&](const sem::Expression* e) { |
| return Switch( |
| e->Declaration(), // |
| [&](const ast::UnaryOpExpression* unary) -> const sem::Expression* { |
| switch (unary->op) { |
| case core::UnaryOp::kAddressOf: |
| case core::UnaryOp::kIndirection: |
| return sem.Get(unary->expr); // Follow the object |
| default: |
| TINT_ICE() << "unexpected unary op: " << unary->op; |
| return nullptr; |
| } |
| }, |
| TINT_ICE_ON_NO_MATCH); |
| }, |
| TINT_ICE_ON_NO_MATCH); |
| } |
| return len; |
| } |
| |
| /// @returns an AST expression that is equal to the arrayLength() of the runtime-sized array |
| /// held by the module-scope variable or parameter @p var, or nullptr on error or if the array |
| /// is not in the Config::bindpoint_to_size_index map. |
| const ast::Expression* ArrayLengthOf(const sem::Variable* var) { |
| return Switch( |
| var, // |
| [&](const sem::GlobalVariable* global) { return ArrayLengthOf(global); }, |
| [&](const sem::Parameter* param) { return ArrayLengthOf(param); }, |
| TINT_ICE_ON_NO_MATCH); |
| } |
| |
| /// @returns an AST expression that is equal to the arrayLength() of the runtime-sized array |
| /// held by the module scope variable @p global, or nullptr on error or if the array is not in |
| /// the Config::bindpoint_to_size_index map. |
| const ast::Expression* ArrayLengthOf(const sem::GlobalVariable* global) { |
| auto binding = global->Attributes().binding_point; |
| TINT_ASSERT_OR_RETURN_VALUE(binding, nullptr); |
| |
| auto idx_it = cfg->bindpoint_to_size_index.find(*binding); |
| if (idx_it == cfg->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; |
| } |
| |
| uint32_t size_index = idx_it->second; |
| used_size_indices.insert(size_index); |
| |
| // Load the total storage buffer size from the UBO. |
| uint32_t array_index = size_index / 4; |
| auto* vec_expr = b.IndexAccessor( |
| b.MemberAccessor(array_lengths_var, kArrayLengthsMemberName), u32(array_index)); |
| uint32_t vec_index = size_index % 4; |
| auto* total_storage_buffer_size = b.IndexAccessor(vec_expr, u32(vec_index)); |
| |
| // Calculate actual array length |
| // total_storage_buffer_size - array_offset |
| // array_length = ---------------------------------------- |
| // array_stride |
| const Expression* total_size = total_storage_buffer_size; |
| if (TINT_UNLIKELY(global->Type()->Is<core::type::Pointer>())) { |
| TINT_ICE() << "storage buffer variable should not be a pointer. " |
| "These should have been removed by the SimplifyPointers transform"; |
| return nullptr; |
| } |
| auto* storage_buffer_type = global->Type()->UnwrapRef(); |
| const core::type::Array* array_type = nullptr; |
| if (auto* str = storage_buffer_type->As<core::type::Struct>()) { |
| // The variable is a struct, so subtract the byte offset of the |
| // array member. |
| auto* array_member_sem = str->Members().Back(); |
| array_type = array_member_sem->Type()->As<core::type::Array>(); |
| total_size = b.Sub(total_storage_buffer_size, u32(array_member_sem->Offset())); |
| } else if (auto* arr = storage_buffer_type->As<core::type::Array>()) { |
| array_type = arr; |
| } else { |
| TINT_ICE() << "expected form of arrayLength argument to be &array_var or " |
| "&struct_var.array_member"; |
| return nullptr; |
| } |
| return b.Div(total_size, u32(array_type->Stride())); |
| } |
| |
| /// @returns an AST expression that is equal to the arrayLength() of the runtime-sized array |
| /// held by the object pointed to by the pointer parameter @p param. |
| const ast::Expression* ArrayLengthOf(const sem::Parameter* param) { |
| // Pointer originates from a parameter. |
| // Add a new array length parameter to the function, and use that. |
| auto len_name = param_lengths.GetOrAdd(param, [&] { |
| auto* fn = param->Owner()->As<sem::Function>(); |
| auto name = b.Symbols().New(param->Declaration()->name->symbol.Name() + "_length"); |
| auto* len_param = b.Param(name, b.ty.u32()); |
| ctx.InsertAfter(fn->Declaration()->params, param->Declaration(), len_param); |
| len_params_needing_args.Add(param); |
| return name; |
| }); |
| return b.Expr(len_name); |
| } |
| |
| /// Constructs the uniform buffer variable that will hold the array lengths. |
| void AddArrayLengthsUniformVar() { |
| // Calculate the highest index in the array lengths array |
| uint32_t highest_index = 0; |
| for (auto idx : used_size_indices) { |
| if (idx > highest_index) { |
| highest_index = idx; |
| } |
| } |
| |
| // Emit an array<vec4<u32>, N>, where N is 1/4 number of elements. |
| // We do this because UBOs require an element stride that is 16-byte aligned. |
| auto* buffer_size_struct = |
| b.Structure(b.Symbols().New("TintArrayLengths"), |
| tint::Vector{ |
| b.Member(kArrayLengthsMemberName, |
| b.ty.array(b.ty.vec4<u32>(), u32((highest_index / 4) + 1))), |
| }); |
| b.GlobalVar(array_lengths_var, b.ty.Of(buffer_size_struct), core::AddressSpace::kUniform, |
| b.Group(AInt(cfg->ubo_binding.group)), |
| b.Binding(AInt(cfg->ubo_binding.binding))); |
| } |
| |
| /// Adds an additional array-length argument to all the calls to the function that owns the |
| /// pointer parameter @p param. This may add new entries to #len_params_needing_args. |
| void AddArrayLengthArguments(const sem::Parameter* param) { |
| auto* fn = param->Owner()->As<sem::Function>(); |
| for (auto* call : fn->CallSites()) { |
| auto* arg = call->Arguments()[param->Index()]; |
| if (auto* len = ArrayLengthOf(arg); len) { |
| ctx.InsertAfter(call->Declaration()->args, arg->Declaration(), len); |
| } else { |
| // Callee expects an array length, but there's no binding for it. |
| // Call arrayLength() at the call-site. |
| len = b.Call(wgsl::BuiltinFn::kArrayLength, ctx.Clone(arg->Declaration())); |
| ctx.InsertAfter(call->Declaration()->args, arg->Declaration(), len); |
| } |
| } |
| } |
| |
| /// Name of the array-lengths struct member that holds all the array lengths. |
| static constexpr std::string_view kArrayLengthsMemberName = "array_lengths"; |
| |
| /// The source program |
| const Program& src; |
| /// The transform outputs |
| DataMap& outputs; |
| /// The transform config |
| const Config* const cfg; |
| /// The target program builder |
| ProgramBuilder b; |
| /// The clone context |
| program::CloneContext ctx = {&b, &src, /* auto_clone_symbols */ true}; |
| /// Alias to src.Sem() |
| const sem::Info& sem = src.Sem(); |
| /// Name of the uniform buffer variable that holds the array lengths |
| Symbol array_lengths_var; |
| /// A map of pointer-parameter to the name of the new array-length parameter. |
| Hashmap<const sem::Parameter*, Symbol, 8> param_lengths; |
| /// Indices into the uniform buffer array indices that are statically used. |
| std::unordered_set<uint32_t> used_size_indices; |
| /// A vector of array-length parameters which need corresponding array-length arguments for all |
| /// callsites. |
| UniqueVector<const sem::Parameter*, 8> len_params_needing_args; |
| }; |
| |
| Transform::ApplyResult ArrayLengthFromUniform::Apply(const Program& src, |
| const DataMap& inputs, |
| DataMap& outputs) const { |
| return State{src, inputs, outputs}.Run(); |
| } |
| |
| ArrayLengthFromUniform::Config::Config() = default; |
| ArrayLengthFromUniform::Config::Config(BindingPoint ubo_bp) : ubo_binding(ubo_bp) {} |
| ArrayLengthFromUniform::Config::Config(const Config&) = default; |
| ArrayLengthFromUniform::Config& ArrayLengthFromUniform::Config::operator=(const Config&) = default; |
| ArrayLengthFromUniform::Config::~Config() = default; |
| |
| ArrayLengthFromUniform::Result::Result(std::unordered_set<uint32_t> used_size_indices_in) |
| : used_size_indices(std::move(used_size_indices_in)) {} |
| ArrayLengthFromUniform::Result::Result(const Result&) = default; |
| ArrayLengthFromUniform::Result::~Result() = default; |
| |
| } // namespace tint::ast::transform |