src/tint/lang/spirv/reader/lower/decompose_strided_array.cc

// 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/spirv/reader/lower/decompose_strided_array.h"

#include <utility>

#include "src/tint/lang/core/ir/builder.h"
#include "src/tint/lang/core/ir/module.h"
#include "src/tint/lang/core/ir/validator/validate.h"
#include "src/tint/lang/spirv/type/explicit_layout_array.h"

namespace tint::spirv::reader::lower {
namespace {

/// PIMPL state for the transform.
struct State {
    /// The IR module.
    core::ir::Module& ir;

    /// The IR builder.
    core::ir::Builder b{ir};

    /// The type manager.
    core::type::Manager& ty{ir.Types()};

    /// The symbol manager.
    SymbolTable& sym{ir.symbols};

    /// A map from a type to its replacement type (which may be the same as the original).
    Hashmap<const core::type::Type*, const core::type::Type*, 32> type_map{};

    /// A map from a constant to the rewritten constant.
    Hashmap<const core::constant::Value*, const core::constant::Value*, 8> constant_map{};

    /// All of the padded structure types that were created by this transform.
    Hashset<const core::type::Type*, 8> padded_structures{};

    /// Process the module.
    void Process() {
        Vector<core::ir::Access*, 8> access_worklist;
        Vector<core::ir::Construct*, 8> construct_worklist;
        for (auto* inst : ir.Instructions()) {
            // Replace all constant operands where the type will be changed, either due to being a
            // strided array, or a composite containing a strided array.
            for (uint32_t i = 0; i < inst->Operands().Length(); ++i) {
                auto* constant = As<core::ir::Constant>(inst->Operands()[i]);
                if (!constant) {
                    continue;
                }

                auto* new_constant = RewriteConstant(constant->Value());
                if (new_constant != constant->Value()) {
                    inst->SetOperand(i, b.Constant(new_constant));
                }
            }

            // Track all construct instructions that need to be updated later.
            if (auto* construct = inst->As<core::ir::Construct>()) {
                if (construct->Result(0)->Type()->Is<spirv::type::ExplicitLayoutArray>()) {
                    construct_worklist.Push(construct);
                }
            }

            // Update any instruction result that contains an array with a non-default stride.
            for (auto* result : inst->Results()) {
                result->SetType(RewriteType(result->Type()));
            }

            // Track all access instructions that may need to be updated later.
            if (auto* access = inst->As<core::ir::Access>()) {
                access_worklist.Push(access);
            }
        }

        // Update the types of any function parameters and function return types that contain arrays
        // with non-default strides.
        for (auto func : ir.functions) {
            for (auto* param : func->Params()) {
                param->SetType(RewriteType(param->Type()));
            }
            func->SetReturnType(RewriteType(func->ReturnType()));
        }

        // Wrap operands for construct instructions in padded structures.
        for (auto* construct : construct_worklist) {
            WrapConstructOperands(construct);
        }

        // Inject additional indices into access instructions where needed.
        // We do this last as it relies on us knowing whether the access chain indexes into padded
        // structures that we have added in the previous steps.
        for (auto* access : access_worklist) {
            UpdateAccessIndices(access);
        }
    }

    /// Rewrite a type to avoid using any arrays that have non-default strides.
    const core::type::Type* RewriteType(const core::type::Type* type) {
        return type_map.GetOrAdd(type, [&] {
            return tint::Switch(
                type,  //
                [&](const spirv::type::ExplicitLayoutArray* arr) {
                    auto* new_element_type = RewriteType(arr->ElemType());

                    if (!arr->IsStrideImplicit()) {
                        // The stride does not match the implicit element stride, so we need to wrap
                        // the element type in a structure that is padded to the target stride.
                        auto* element_struct =
                            ty.Struct(sym.New("tint_padded_array_element"),
                                      Vector{
                                          ty.Get<core::type::StructMember>(
                                              sym.New("tint_element"), new_element_type,
                                              /* index */ 0u,
                                              /* offset */ 0u,
                                              /* align */ new_element_type->Align(),
                                              /* size */ arr->Stride(), core::IOAttributes{}),
                                      });
                        new_element_type = element_struct;
                        padded_structures.Add(element_struct);
                    }

                    return ty.Get<core::type::Array>(new_element_type, arr->Count(), arr->Size());
                },
                [&](const core::type::Array* arr) {
                    auto* new_element_type = RewriteType(arr->ElemType());
                    return ty.Get<core::type::Array>(new_element_type, arr->Count(), arr->Size());
                },
                [&](const core::type::Struct* str) -> const core::type::Struct* {
                    // Rewrite members of the struct that contain arrays with non-default strides.
                    bool made_changes = false;
                    Vector<const core::type::StructMember*, 8> new_members;
                    new_members.Reserve(str->Members().Length());
                    for (auto* member : str->Members()) {
                        auto* new_member_type = RewriteType(member->Type());
                        if (new_member_type == member->Type()) {
                            new_members.Push(member);
                            continue;
                        }

                        new_members.Push(ty.Get<core::type::StructMember>(
                            member->Name(), new_member_type, member->Index(), member->Offset(),
                            member->Align(), member->Size(), member->Attributes()));
                        made_changes = true;
                    }
                    if (!made_changes) {
                        return str;
                    }

                    return ty.Struct(sym.New(str->Name().Name()), std::move(new_members));
                },
                [&](const core::type::Pointer* ptr) {
                    return ty.ptr(ptr->AddressSpace(), RewriteType(ptr->StoreType()),
                                  ptr->Access());
                },
                [&](Default) { return type; });
        });
    }

    /// Rewrite a constant to avoid using any arrays that have non-default strides.
    const core::constant::Value* RewriteConstant(const core::constant::Value* constant) {
        auto* new_type = RewriteType(constant->Type());
        if (new_type == constant->Type()) {
            return constant;
        }

        // Check if the constant is an array of padded structures.
        // A padded structure will only appear as the element type of an array, so we only need to
        // check the first element type.
        const core::type::Type* padded_struct_type = nullptr;
        if (padded_structures.Contains(new_type->Element(0))) {
            padded_struct_type = new_type->Element(0);
        }

        Vector<const core::constant::Value*, 16> elements;
        for (uint32_t i = 0; i < constant->NumElements(); i++) {
            auto* new_element = RewriteConstant(constant->Index(i));

            // If we are rewriting an array of padded structures, then we need to wrap the original
            // element in a constant structure.
            if (padded_struct_type) {
                new_element = ir.constant_values.Composite(padded_struct_type, Vector{new_element});
            }

            elements.Push(new_element);
        }
        return ir.constant_values.Composite(new_type, std::move(elements));
    }

    /// Wrap all operands for a construct instruction that produces an array with padded elements.
    void WrapConstructOperands(core::ir::Construct* construct) {
        auto* padded_struct_type = construct->Result(0)->Type()->Element(0);
        TINT_ASSERT(padded_struct_type && padded_structures.Contains(padded_struct_type));
        b.InsertBefore(construct, [&] {
            Vector<core::ir::Value*, 8> new_operands;
            for (auto* operand : construct->Operands()) {
                new_operands.Push(b.Construct(padded_struct_type, operand));
            }
            construct->SetOperands(new_operands);
        });
    }

    /// Inject additional `0u` indices into access instructions that index into any padded structure
    /// types that we have created.
    void UpdateAccessIndices(core::ir::Access* access) {
        Vector<core::ir::Value*, 8> new_operands{access->Object()};

        auto old_indices = access->Indices();
        auto* current_type = access->Object()->Type()->UnwrapPtr();
        for (auto* idx : old_indices) {
            new_operands.Push(idx);

            auto* const_idx = idx->As<core::ir::Constant>();
            current_type = const_idx
                               ? current_type->Element(const_idx->Value()->ValueAs<uint32_t>())
                               : current_type->Elements().type;

            if (padded_structures.Contains(current_type)) {
                new_operands.Push(b.Zero<core::u32>());
                current_type = current_type->Element(0);
            }
        }
        access->SetOperands(std::move(new_operands));
    }
};

}  // namespace

Result<SuccessType> DecomposeStridedArray(core::ir::Module& ir) {
    core::ir::AssertValid(ir, "before spirv.DecomposeStridedArray");

    State{ir}.Process();

    return Success;
}

}  // namespace tint::spirv::reader::lower