[spirv-reader] Add DecomposeStridedMatrix transform Rewrite struct members that are strided matrices to use strided arrays instead. These will then be rewritten by the DecomposeStridedArray transform. Bug: 417682206 Change-Id: I160336d64ea221ba20abdb7f8989302804b89fbc Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/250534 Reviewed-by: dan sinclair <dsinclair@chromium.org> Commit-Queue: James Price <jrprice@google.com>
diff --git a/src/tint/lang/core/ir/store.h b/src/tint/lang/core/ir/store.h index 13ff0e8..24e7ee7 100644 --- a/src/tint/lang/core/ir/store.h +++ b/src/tint/lang/core/ir/store.h
@@ -80,6 +80,9 @@ /// @returns the value being stored const Value* From() const { return Operand(kFromOperandOffset); } + /// @param from the value being stored + void SetFrom(Value* from) { SetOperand(kFromOperandOffset, from); } + /// @returns the friendly name for the instruction std::string FriendlyName() const override { return "store"; }
diff --git a/src/tint/lang/core/ir/transform/dead_code_elimination.h b/src/tint/lang/core/ir/transform/dead_code_elimination.h index c37de06..f2e540a 100644 --- a/src/tint/lang/core/ir/transform/dead_code_elimination.h +++ b/src/tint/lang/core/ir/transform/dead_code_elimination.h
@@ -45,6 +45,7 @@ core::ir::Capability::kAllowPhonyInstructions, core::ir::Capability::kAllowUnannotatedModuleIOVariables, core::ir::Capability::kAllowNonCoreTypes, + core::ir::Capability::kAllowStructMatrixDecorations, }; /// DeadCodeElimination is a transform that removes dead code from the given IR module.
diff --git a/src/tint/lang/spirv/reader/lower/BUILD.bazel b/src/tint/lang/spirv/reader/lower/BUILD.bazel index 7275f09..581de23 100644 --- a/src/tint/lang/spirv/reader/lower/BUILD.bazel +++ b/src/tint/lang/spirv/reader/lower/BUILD.bazel
@@ -42,6 +42,7 @@ "atomics.cc", "builtins.cc", "decompose_strided_array.cc", + "decompose_strided_matrix.cc", "lower.cc", "shader_io.cc", "texture.cc", @@ -51,6 +52,7 @@ "atomics.h", "builtins.h", "decompose_strided_array.h", + "decompose_strided_matrix.h", "lower.h", "shader_io.h", "texture.h", @@ -90,6 +92,7 @@ "atomics_test.cc", "builtins_test.cc", "decompose_strided_array_test.cc", + "decompose_strided_matrix_test.cc", "shader_io_test.cc", "texture_test.cc", "vector_element_pointer_test.cc",
diff --git a/src/tint/lang/spirv/reader/lower/BUILD.cmake b/src/tint/lang/spirv/reader/lower/BUILD.cmake index a2980d7..760f1ab 100644 --- a/src/tint/lang/spirv/reader/lower/BUILD.cmake +++ b/src/tint/lang/spirv/reader/lower/BUILD.cmake
@@ -45,6 +45,8 @@ lang/spirv/reader/lower/builtins.h lang/spirv/reader/lower/decompose_strided_array.cc lang/spirv/reader/lower/decompose_strided_array.h + lang/spirv/reader/lower/decompose_strided_matrix.cc + lang/spirv/reader/lower/decompose_strided_matrix.h lang/spirv/reader/lower/lower.cc lang/spirv/reader/lower/lower.h lang/spirv/reader/lower/shader_io.cc @@ -91,6 +93,7 @@ lang/spirv/reader/lower/atomics_test.cc lang/spirv/reader/lower/builtins_test.cc lang/spirv/reader/lower/decompose_strided_array_test.cc + lang/spirv/reader/lower/decompose_strided_matrix_test.cc lang/spirv/reader/lower/shader_io_test.cc lang/spirv/reader/lower/texture_test.cc lang/spirv/reader/lower/vector_element_pointer_test.cc
diff --git a/src/tint/lang/spirv/reader/lower/BUILD.gn b/src/tint/lang/spirv/reader/lower/BUILD.gn index 9d45b0b..29c9276 100644 --- a/src/tint/lang/spirv/reader/lower/BUILD.gn +++ b/src/tint/lang/spirv/reader/lower/BUILD.gn
@@ -51,6 +51,8 @@ "builtins.h", "decompose_strided_array.cc", "decompose_strided_array.h", + "decompose_strided_matrix.cc", + "decompose_strided_matrix.h", "lower.cc", "lower.h", "shader_io.cc", @@ -91,6 +93,7 @@ "atomics_test.cc", "builtins_test.cc", "decompose_strided_array_test.cc", + "decompose_strided_matrix_test.cc", "shader_io_test.cc", "texture_test.cc", "vector_element_pointer_test.cc",
diff --git a/src/tint/lang/spirv/reader/lower/builtins.cc b/src/tint/lang/spirv/reader/lower/builtins.cc index bfd3fb3..be0220c 100644 --- a/src/tint/lang/spirv/reader/lower/builtins.cc +++ b/src/tint/lang/spirv/reader/lower/builtins.cc
@@ -1174,6 +1174,7 @@ core::ir::Capabilities{ core::ir::Capability::kAllowOverrides, core::ir::Capability::kAllowNonCoreTypes, + core::ir::Capability::kAllowStructMatrixDecorations, }); if (result != Success) { return result.Failure();
diff --git a/src/tint/lang/spirv/reader/lower/decompose_strided_matrix.cc b/src/tint/lang/spirv/reader/lower/decompose_strided_matrix.cc new file mode 100644 index 0000000..e6b64e7 --- /dev/null +++ b/src/tint/lang/spirv/reader/lower/decompose_strided_matrix.cc
@@ -0,0 +1,381 @@ +// 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_matrix.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.h" +#include "src/tint/lang/core/type/matrix.h" +#include "src/tint/lang/spirv/type/explicit_layout_array.h" + +namespace tint::spirv::reader::lower { +namespace { + +using namespace tint::core::fluent_types; // NOLINT + +/// 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). + struct TypeAndStride { + const core::type::Type* type; + uint32_t stride; + + bool operator==(const TypeAndStride& other) const { + return type == other.type && stride == other.stride; + } + + tint::HashCode HashCode() const { return Hash(type, stride); } + }; + Hashmap<TypeAndStride, const core::type::Type*, 32> type_map{}; + + /// A map from rewritten structs to original structs. + Hashmap<const core::type::Struct*, const core::type::Struct*, 4> struct_to_original{}; + + /// Process the module. + void Process() { + Vector<core::ir::Access*, 32> access_worklist; + Vector<core::ir::Construct*, 32> construct_worklist; + for (auto* inst : ir.Instructions()) { + // Replace all constant operands where the type will be changed due to it containing a + // structure that uses a matrix stride attribute. + for (uint32_t i = 0; i < inst->Operands().Length(); ++i) { + if (auto* constant = As<core::ir::Constant>(inst->Operands()[i])) { + auto* new_constant = RewriteConstant(constant->Value()); + if (new_constant != constant->Value()) { + inst->SetOperand(i, b.Constant(new_constant)); + } + } + } + + // Update any instruction result that contains a matrix stride attribute. + for (auto* result : inst->Results()) { + result->SetType(RewriteType(result->Type())); + } + + // Track instructions that may need to be updated later. + if (auto* access = inst->As<core::ir::Access>()) { + access_worklist.Push(access); + } + if (auto* construct = inst->As<core::ir::Construct>()) { + construct_worklist.Push(construct); + } + } + + // Update the types of any function parameters and function return types that contain + // matrices with non-default strides. + for (auto func : ir.functions) { + for (auto* param : func->Params()) { + param->SetType(RewriteType(param->Type())); + } + func->SetReturnType(RewriteType(func->ReturnType())); + } + + // Update any access instructions that produce strided matrices. + for (auto* access : access_worklist) { + UpdateAccessInstruction(access, /* source_is_strided */ false); + } + + // Convert strided matrix operands for construct instructions. + for (auto* construct : construct_worklist) { + ConvertConstructOperands(construct); + } + } + + /// Rewrite a type to replace structure members that have matrix strides. + const core::type::Type* RewriteType(const core::type::Type* type, uint32_t stride = 0) { + return type_map.GetOrAdd(TypeAndStride{type, stride}, [&] { + return tint::Switch( + type, // + [&](const core::type::Matrix* mat) -> const core::type::Type* { + if (stride == 0 || stride == mat->ColumnStride()) { + return mat; + } + // Replace the matrix with a strided array of column vectors. + TINT_ASSERT(stride % mat->ColumnStride() == 0); + return ty.Get<spirv::type::ExplicitLayoutArray>( + mat->ColumnType(), ty.Get<core::type::ConstantArrayCount>(mat->Columns()), + stride, stride * mat->Columns(), stride); + }, + [&](const core::type::Array* arr) { return RewriteArray(arr, stride); }, + [&](const core::type::Struct* str) { return RewriteStruct(str); }, + [&](const core::type::Pointer* ptr) { + return ty.ptr(ptr->AddressSpace(), RewriteType(ptr->StoreType()), + ptr->Access()); + }, + [&](Default) { return type; }); + }); + } + + /// Rewrite an array type if necessary. + const core::type::Array* RewriteArray(const core::type::Array* arr, uint32_t matrix_stride) { + auto* new_element_type = RewriteType(arr->ElemType(), matrix_stride); + if (new_element_type == arr->ElemType()) { + return arr; + } + + // The element type is the only thing that will change. That does not affect the stride of + // the array itself, which may either be the natural stride or an larger stride in the case + // of an explicitly laid out array. + if (arr->Is<spirv::type::ExplicitLayoutArray>()) { + return ty.Get<spirv::type::ExplicitLayoutArray>( + new_element_type, arr->Count(), arr->Align(), arr->Size(), arr->Stride()); + } + return ty.Get<core::type::Array>(new_element_type, arr->Count(), arr->Align(), arr->Size(), + arr->Stride(), arr->Stride()); + } + + /// Rewrite a structure type to replace structure members that have matrix stride attributes. + const core::type::Struct* RewriteStruct(const core::type::Struct* old_struct) { + bool made_changes = false; + + Vector<const core::type::StructMember*, 8> new_members; + new_members.Reserve(old_struct->Members().Length()); + for (auto* member : old_struct->Members()) { + auto* new_member_type = RewriteType(member->Type(), member->MatrixStride()); + if (member->HasMatrixStride() || new_member_type != member->Type()) { + // Recreate the struct member without the stride attribute, and using the new type. + 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; + } else { + new_members.Push(member); + } + } + if (!made_changes) { + return old_struct; + } + + // Create the new struct and record the mapping to the old struct. + auto* new_struct = ty.Struct(sym.New(old_struct->Name().Name()), std::move(new_members)); + struct_to_original.Add(new_struct, old_struct); + return new_struct; + } + + /// Rewrite a constant to replace strided matrix constants with the equivalent strided array + /// of column vector constants. + const core::constant::Value* RewriteConstant(const core::constant::Value* constant, + uint32_t stride = 0) { + auto* new_type = RewriteType(constant->Type(), stride); + if (new_type == constant->Type()) { + return constant; + } + + Vector<const core::constant::Value*, 16> elements; + for (uint32_t i = 0; i < constant->NumElements(); i++) { + auto* value = constant->Index(i); + + // If this is a struct member, we need to check if the type has changed. + if (auto* new_struct_type = new_type->As<core::type::Struct>()) { + auto* new_member_type = new_struct_type->Members()[i]->Type(); + if (new_member_type != value->Type()) { + // Create a new constant using the strided array type. + // If the type changed, it must have had a MatrixStride decoration and will have + // been rewritten as an array type (or it already was an array). + auto* array = new_member_type->As<core::type::Array>(); + TINT_ASSERT(array); + + auto* old_struct_type = constant->Type()->As<core::type::Struct>(); + auto member_stride = old_struct_type->Members()[i]->MatrixStride(); + + Vector<const core::constant::Value*, 4> new_elements; + for (uint32_t j = 0; j < array->ConstantCount().value(); j++) { + new_elements.Push(RewriteConstant(value->Index(j), member_stride)); + } + value = ir.constant_values.Composite(array, std::move(new_elements)); + } + } + + elements.Push(RewriteConstant(value, stride)); + } + return ir.constant_values.Composite(new_type, std::move(elements)); + } + + /// Convert strided matrix operands to strided arrays for a construct instruction. + void ConvertConstructOperands(core::ir::Construct* construct) { + auto* struct_type = construct->Result()->Type()->As<core::type::Struct>(); + if (!struct_type) { + return; + } + + b.InsertBefore(construct, [&] { + Vector<core::ir::Value*, 8> new_operands; + for (uint32_t i = 0; i < construct->Operands().Length(); i++) { + auto* operand = construct->Operands()[i]; + auto* member_type = struct_type->Members()[i]->Type(); + if (member_type != operand->Type()) { + new_operands.Push(Convert(member_type, operand)); + } else { + new_operands.Push(operand); + } + } + construct->SetOperands(new_operands); + }); + } + + /// Update the result type of an access instruction if needed, and the uses of that result. + void UpdateAccessInstruction(core::ir::Access* access, bool source_is_strided) { + // Determine the result type based on the potentially modified object type. + bool indexed_through_strided_member = source_is_strided; + auto* current_type = access->Object()->Type()->UnwrapPtr(); + for (auto* idx : access->Indices()) { + if (auto* struct_type = current_type->As<core::type::Struct>()) { + auto const_idx = idx->As<core::ir::Constant>()->Value()->ValueAs<uint32_t>(); + current_type = current_type->Element(const_idx); + + // Check if we are indexing into a member that has a non-natural matrix stride. + auto* original_struct = struct_to_original.GetOr(struct_type, nullptr); + if (!original_struct) { + // The structure type has not changed so cannot have any matrix strides. + continue; + } + auto* member = original_struct->Members()[const_idx]; + if (member->HasMatrixStride() && current_type != member->Type()) { + indexed_through_strided_member = true; + } + } else { + current_type = current_type->Elements().type; + } + } + if (!indexed_through_strided_member || current_type->Is<core::type::Vector>()) { + return; + } + + if (auto* ptr = access->Result()->Type()->As<core::type::Pointer>()) { + ReplaceMatrixPointerWithArrayPointer(ptr, current_type, access); + } else { + // We were extracting a strided matrix from a structure, so we need to convert the + // strided array back to that matrix type. + b.InsertAfter(access, [&] { + auto* extracted_array = b.InstructionResult(current_type); + access->Result()->ReplaceAllUsesWith( + Convert(access->Result()->Type(), extracted_array)); + access->SetResult(extracted_array); + }); + } + } + + /// Change the type of a pointer instruction result that contains a strided matrix, and then + /// update any instructions that use that result. + void ReplaceMatrixPointerWithArrayPointer(const core::type::Pointer* old_ptr, + const core::type::Type* new_store_type, + core::ir::Instruction* instruction) { + auto* old_store_type = old_ptr->StoreType(); + auto* new_ptr = ty.ptr(old_ptr->AddressSpace(), new_store_type, old_ptr->Access()); + + Vector<core::ir::Instruction*, 8> worklist{instruction}; + while (!worklist.IsEmpty()) { + auto* inst = worklist.Pop(); + inst->Result()->SetType(new_ptr); + inst->Result()->ForEachUseUnsorted([&](const core::ir::Usage& use) { + tint::Switch( + use.instruction, // + [&](core::ir::Access* access) { + UpdateAccessInstruction(access, /* source_is_strided */ true); + }, + [&](core::ir::Let* let) { worklist.Push(let); }, + [&](core::ir::Load* load) { + // Convert the value to the original type. + b.InsertAfter(load, [&] { + auto* new_load_result = b.InstructionResult(new_store_type); + auto* converted = Convert(old_store_type, new_load_result); + load->Result()->ReplaceAllUsesWith(converted); + load->SetResult(new_load_result); + }); + }, + [&](core::ir::Store* store) { + // Convert the value to the new type. + b.InsertBefore(store, [&] { // + store->SetFrom(Convert(new_store_type, store->From())); + }); + }, + TINT_ICE_ON_NO_MATCH); + }); + } + } + + /// Convert a value between an [array of] strided matrix and an [array of] strided array. + core::ir::Value* Convert(const core::type::Type* dst, core::ir::Value* src) { + auto dst_elements = dst->Elements(); + auto src_elements = src->Type()->Elements(); + TINT_ASSERT(dst_elements.count == src_elements.count); + Vector<core::ir::Value*, 8> elements; + elements.Reserve(dst_elements.count); + for (uint32_t i = 0; i < dst_elements.count; i++) { + // Extract the element from the source value. + core::ir::Value* el = nullptr; + if (auto* constant = src->As<core::ir::Constant>()) { + el = b.Constant(constant->Value()->Index(i)); + } else { + el = b.Access(src_elements.type, src, u32(i))->Result(); + } + + // Recurse to convert strided matrices nested in arrays if needed. + if (src_elements.type != dst_elements.type) { + el = Convert(dst_elements.type, el); + } + + elements.Push(el); + } + return b.Construct(dst, std::move(elements))->Result(); + } +}; + +} // namespace + +Result<SuccessType> DecomposeStridedMatrix(core::ir::Module& ir) { + auto result = ValidateAndDumpIfNeeded(ir, "spirv.DecomposeStridedMatrix", + core::ir::Capabilities{ + core::ir::Capability::kAllowStructMatrixDecorations, + core::ir::Capability::kAllowNonCoreTypes, + core::ir::Capability::kAllowOverrides, + }); + if (result != Success) { + return result.Failure(); + } + + State{ir}.Process(); + + return Success; +} + +} // namespace tint::spirv::reader::lower
diff --git a/src/tint/lang/spirv/reader/lower/decompose_strided_matrix.h b/src/tint/lang/spirv/reader/lower/decompose_strided_matrix.h new file mode 100644 index 0000000..ae4a724 --- /dev/null +++ b/src/tint/lang/spirv/reader/lower/decompose_strided_matrix.h
@@ -0,0 +1,48 @@ +// 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_SPIRV_READER_LOWER_DECOMPOSE_STRIDED_MATRIX_H_ +#define SRC_TINT_LANG_SPIRV_READER_LOWER_DECOMPOSE_STRIDED_MATRIX_H_ + +#include "src/tint/utils/result.h" + +// Forward declarations. +namespace tint::core::ir { +class Module; +} + +namespace tint::spirv::reader::lower { + +/// DecomposeStridedMatrix is a transform that replaces structure members that have matrix stride +/// attributes with arrays of vectors that have non-default strides. +/// @param module the module to transform +/// @returns success or failure +Result<SuccessType> DecomposeStridedMatrix(core::ir::Module& module); + +} // namespace tint::spirv::reader::lower + +#endif // SRC_TINT_LANG_SPIRV_READER_LOWER_DECOMPOSE_STRIDED_MATRIX_H_
diff --git a/src/tint/lang/spirv/reader/lower/decompose_strided_matrix_test.cc b/src/tint/lang/spirv/reader/lower/decompose_strided_matrix_test.cc new file mode 100644 index 0000000..f39af21 --- /dev/null +++ b/src/tint/lang/spirv/reader/lower/decompose_strided_matrix_test.cc
@@ -0,0 +1,2270 @@ +// 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_matrix.h" + +#include "src/tint/lang/core/ir/transform/helper_test.h" +#include "src/tint/lang/spirv/type/explicit_layout_array.h" + +namespace tint::spirv::reader::lower { +namespace { + +using namespace tint::core::fluent_types; // NOLINT +using namespace tint::core::number_suffixes; // NOLINT + +class SpirvReader_DecomposeStridedMatrixTest : public core::ir::transform::TransformTest { + protected: + void SetUp() override { capabilities.Add(core::ir::Capability::kAllowNonCoreTypes); } + + /// Create a struct that has a matrix member sandwiched between two u32 members, optionally + /// nested inside one or more arrays. + const core::type::Struct* Struct(const core::type::Matrix* matrix_type, + uint32_t matrix_stride, + std::initializer_list<uint32_t> array_counts = {}) { + uint32_t member_size = matrix_stride * matrix_type->Columns(); + const core::type::Type* member_type = matrix_type; + for (uint32_t count : array_counts) { + member_type = ty.array(member_type, count); + member_size *= count; + } + auto* matrix_member = + ty.Get<core::type::StructMember>(mod.symbols.New("b"), member_type, 1u, matrix_stride, + matrix_stride, member_size, core::IOAttributes{}); + matrix_member->SetMatrixStride(matrix_stride); + return ty.Struct( + mod.symbols.New("S"), + Vector{ + ty.Get<core::type::StructMember>(mod.symbols.New("a"), ty.u32(), 0u, 0u, 4u, 4u, + core::IOAttributes{}), + matrix_member, + ty.Get<core::type::StructMember>(mod.symbols.New("c"), ty.u32(), 2u, + matrix_member->Offset() + matrix_member->Size(), + 4u, 4u, core::IOAttributes{}), + }); + } +}; + +TEST_F(SpirvReader_DecomposeStridedMatrixTest, NaturalStride_CreateConstant) { + auto* matrix_type = ty.mat4x4<f32>(); + auto* struct_type = Struct(matrix_type, 16); + + auto* f = b.ComputeFunction("foo"); + b.Append(f->Block(), [&] { + b.Let("value", b.Composite(struct_type, 42_u, b.Zero(matrix_type), 42_u)); + b.Return(f); + }); + + auto* before = R"( +S = struct @align(16) { + a:u32 @offset(0) + b:mat4x4<f32> @offset(16), @matrix_stride(16) + c:u32 @offset(80) +} + +%foo = @compute @workgroup_size(1u, 1u, 1u) func():void { + $B1: { + %value:S = let S(42u, mat4x4<f32>(vec4<f32>(0.0f)), 42u) + ret + } +} +)"; + auto* after = R"( +S = struct @align(16) { + a:u32 @offset(0) + b:mat4x4<f32> @offset(16), @matrix_stride(16) + c:u32 @offset(80) +} + +S_1 = struct @align(16) { + a:u32 @offset(0) + b:mat4x4<f32> @offset(16) + c:u32 @offset(80) +} + +%foo = @compute @workgroup_size(1u, 1u, 1u) func():void { + $B1: { + %value:S_1 = let S_1(42u, mat4x4<f32>(vec4<f32>(0.0f)), 42u) + ret + } +} +)"; + + ASSERT_EQ(before, str()); + Run(DecomposeStridedMatrix); + ASSERT_EQ(after, str()); +} + +TEST_F(SpirvReader_DecomposeStridedMatrixTest, NaturalStride_LoadMatrix) { + auto* matrix_type = ty.mat4x4<f32>(); + auto* struct_type = Struct(matrix_type, 16); + + auto* var = b.Var("var", ty.ptr<private_>(struct_type)); + mod.root_block->Append(var); + + auto* f = b.ComputeFunction("foo"); + b.Append(f->Block(), [&] { + b.Let("value", b.Load(b.Access<ptr<private_, mat4x4<f32>>>(var, 1_u))); + b.Return(f); + }); + + auto* before = R"( +S = struct @align(16) { + a:u32 @offset(0) + b:mat4x4<f32> @offset(16), @matrix_stride(16) + c:u32 @offset(80) +} + +$B1: { # root + %var:ptr<private, S, read_write> = var undef +} + +%foo = @compute @workgroup_size(1u, 1u, 1u) func():void { + $B2: { + %3:ptr<private, mat4x4<f32>, read_write> = access %var, 1u + %4:mat4x4<f32> = load %3 + %value:mat4x4<f32> = let %4 + ret + } +} +)"; + auto* after = R"( +S = struct @align(16) { + a:u32 @offset(0) + b:mat4x4<f32> @offset(16), @matrix_stride(16) + c:u32 @offset(80) +} + +S_1 = struct @align(16) { + a:u32 @offset(0) + b:mat4x4<f32> @offset(16) + c:u32 @offset(80) +} + +$B1: { # root + %var:ptr<private, S_1, read_write> = var undef +} + +%foo = @compute @workgroup_size(1u, 1u, 1u) func():void { + $B2: { + %3:ptr<private, mat4x4<f32>, read_write> = access %var, 1u + %4:mat4x4<f32> = load %3 + %value:mat4x4<f32> = let %4 + ret + } +} +)"; + + ASSERT_EQ(before, str()); + Run(DecomposeStridedMatrix); + ASSERT_EQ(after, str()); +} + +TEST_F(SpirvReader_DecomposeStridedMatrixTest, NaturalStride_ExtractMatrix) { + auto* matrix_type = ty.mat4x4<f32>(); + auto* struct_type = Struct(matrix_type, 16); + + auto* f = b.ComputeFunction("foo"); + b.Append(f->Block(), [&] { + auto* s = b.Let("s", b.Zero(struct_type)); + b.Let("value", b.Access<mat4x4<f32>>(s, 1_u)); + b.Return(f); + }); + + auto* before = R"( +S = struct @align(16) { + a:u32 @offset(0) + b:mat4x4<f32> @offset(16), @matrix_stride(16) + c:u32 @offset(80) +} + +%foo = @compute @workgroup_size(1u, 1u, 1u) func():void { + $B1: { + %s:S = let S(0u, mat4x4<f32>(vec4<f32>(0.0f)), 0u) + %3:mat4x4<f32> = access %s, 1u + %value:mat4x4<f32> = let %3 + ret + } +} +)"; + auto* after = R"( +S = struct @align(16) { + a:u32 @offset(0) + b:mat4x4<f32> @offset(16), @matrix_stride(16) + c:u32 @offset(80) +} + +S_1 = struct @align(16) { + a:u32 @offset(0) + b:mat4x4<f32> @offset(16) + c:u32 @offset(80) +} + +%foo = @compute @workgroup_size(1u, 1u, 1u) func():void { + $B1: { + %s:S_1 = let S_1(0u, mat4x4<f32>(vec4<f32>(0.0f)), 0u) + %3:mat4x4<f32> = access %s, 1u + %value:mat4x4<f32> = let %3 + ret + } +} +)"; + + ASSERT_EQ(before, str()); + Run(DecomposeStridedMatrix); + ASSERT_EQ(after, str()); +} + +TEST_F(SpirvReader_DecomposeStridedMatrixTest, NaturalStride_Construct) { + auto* matrix_type = ty.mat4x4<f32>(); + auto* struct_type = Struct(matrix_type, 16); + + auto* f = b.ComputeFunction("foo"); + b.Append(f->Block(), [&] { + auto* m = b.Let("m", b.Zero(matrix_type)); + b.Let("value", b.Construct(struct_type, 42_u, m, 42_u)); + b.Return(f); + }); + + auto* before = R"( +S = struct @align(16) { + a:u32 @offset(0) + b:mat4x4<f32> @offset(16), @matrix_stride(16) + c:u32 @offset(80) +} + +%foo = @compute @workgroup_size(1u, 1u, 1u) func():void { + $B1: { + %m:mat4x4<f32> = let mat4x4<f32>(vec4<f32>(0.0f)) + %3:S = construct 42u, %m, 42u + %value:S = let %3 + ret + } +} +)"; + auto* after = R"( +S = struct @align(16) { + a:u32 @offset(0) + b:mat4x4<f32> @offset(16), @matrix_stride(16) + c:u32 @offset(80) +} + +S_1 = struct @align(16) { + a:u32 @offset(0) + b:mat4x4<f32> @offset(16) + c:u32 @offset(80) +} + +%foo = @compute @workgroup_size(1u, 1u, 1u) func():void { + $B1: { + %m:mat4x4<f32> = let mat4x4<f32>(vec4<f32>(0.0f)) + %3:S_1 = construct 42u, %m, 42u + %value:S_1 = let %3 + ret + } +} +)"; + + ASSERT_EQ(before, str()); + Run(DecomposeStridedMatrix); + ASSERT_EQ(after, str()); +} + +TEST_F(SpirvReader_DecomposeStridedMatrixTest, LoadMatrixElement) { + auto* matrix_type = ty.mat4x4<f32>(); + auto* struct_type = Struct(matrix_type, 64); + + auto* var = b.Var("var", ty.ptr<private_>(struct_type)); + mod.root_block->Append(var); + + auto* f = b.ComputeFunction("foo"); + b.Append(f->Block(), [&] { + auto* access = b.Access<ptr<private_, vec4<f32>>>(var, 1_u, 3_u); + b.Let("value", b.LoadVectorElement(access, 2_u)); + b.Return(f); + }); + + auto* before = R"( +S = struct @align(64) { + a:u32 @offset(0) + b:mat4x4<f32> @offset(64) @size(256), @matrix_stride(64) + c:u32 @offset(320) +} + +$B1: { # root + %var:ptr<private, S, read_write> = var undef +} + +%foo = @compute @workgroup_size(1u, 1u, 1u) func():void { + $B2: { + %3:ptr<private, vec4<f32>, read_write> = access %var, 1u, 3u + %4:f32 = load_vector_element %3, 2u + %value:f32 = let %4 + ret + } +} +)"; + auto* after = R"( +S = struct @align(64) { + a:u32 @offset(0) + b:mat4x4<f32> @offset(64) @size(256), @matrix_stride(64) + c:u32 @offset(320) +} + +S_1 = struct @align(64) { + a:u32 @offset(0) + b:spirv.explicit_layout_array<vec4<f32>, 4, stride=64> @offset(64) + c:u32 @offset(320) +} + +$B1: { # root + %var:ptr<private, S_1, read_write> = var undef +} + +%foo = @compute @workgroup_size(1u, 1u, 1u) func():void { + $B2: { + %3:ptr<private, vec4<f32>, read_write> = access %var, 1u, 3u + %4:f32 = load_vector_element %3, 2u + %value:f32 = let %4 + ret + } +} +)"; + + ASSERT_EQ(before, str()); + Run(DecomposeStridedMatrix); + ASSERT_EQ(after, str()); +} + +TEST_F(SpirvReader_DecomposeStridedMatrixTest, LoadMatrixColumn) { + auto* matrix_type = ty.mat4x4<f32>(); + auto* struct_type = Struct(matrix_type, 64); + + auto* var = b.Var("var", ty.ptr<private_>(struct_type)); + mod.root_block->Append(var); + + auto* f = b.ComputeFunction("foo"); + b.Append(f->Block(), [&] { + b.Let("value", b.Load(b.Access<ptr<private_, vec4<f32>>>(var, 1_u, 2_u))); + b.Return(f); + }); + + auto* before = R"( +S = struct @align(64) { + a:u32 @offset(0) + b:mat4x4<f32> @offset(64) @size(256), @matrix_stride(64) + c:u32 @offset(320) +} + +$B1: { # root + %var:ptr<private, S, read_write> = var undef +} + +%foo = @compute @workgroup_size(1u, 1u, 1u) func():void { + $B2: { + %3:ptr<private, vec4<f32>, read_write> = access %var, 1u, 2u + %4:vec4<f32> = load %3 + %value:vec4<f32> = let %4 + ret + } +} +)"; + auto* after = R"( +S = struct @align(64) { + a:u32 @offset(0) + b:mat4x4<f32> @offset(64) @size(256), @matrix_stride(64) + c:u32 @offset(320) +} + +S_1 = struct @align(64) { + a:u32 @offset(0) + b:spirv.explicit_layout_array<vec4<f32>, 4, stride=64> @offset(64) + c:u32 @offset(320) +} + +$B1: { # root + %var:ptr<private, S_1, read_write> = var undef +} + +%foo = @compute @workgroup_size(1u, 1u, 1u) func():void { + $B2: { + %3:ptr<private, vec4<f32>, read_write> = access %var, 1u, 2u + %4:vec4<f32> = load %3 + %value:vec4<f32> = let %4 + ret + } +} +)"; + + ASSERT_EQ(before, str()); + Run(DecomposeStridedMatrix); + ASSERT_EQ(after, str()); +} + +TEST_F(SpirvReader_DecomposeStridedMatrixTest, LoadMatrix) { + auto* matrix_type = ty.mat4x4<f32>(); + auto* struct_type = Struct(matrix_type, 64); + + auto* var = b.Var("var", ty.ptr<private_>(struct_type)); + mod.root_block->Append(var); + + auto* f = b.ComputeFunction("foo"); + b.Append(f->Block(), [&] { + b.Let("value", b.Load(b.Access<ptr<private_, mat4x4<f32>>>(var, 1_u))); + b.Return(f); + }); + + auto* before = R"( +S = struct @align(64) { + a:u32 @offset(0) + b:mat4x4<f32> @offset(64) @size(256), @matrix_stride(64) + c:u32 @offset(320) +} + +$B1: { # root + %var:ptr<private, S, read_write> = var undef +} + +%foo = @compute @workgroup_size(1u, 1u, 1u) func():void { + $B2: { + %3:ptr<private, mat4x4<f32>, read_write> = access %var, 1u + %4:mat4x4<f32> = load %3 + %value:mat4x4<f32> = let %4 + ret + } +} +)"; + auto* after = R"( +S = struct @align(64) { + a:u32 @offset(0) + b:mat4x4<f32> @offset(64) @size(256), @matrix_stride(64) + c:u32 @offset(320) +} + +S_1 = struct @align(64) { + a:u32 @offset(0) + b:spirv.explicit_layout_array<vec4<f32>, 4, stride=64> @offset(64) + c:u32 @offset(320) +} + +$B1: { # root + %var:ptr<private, S_1, read_write> = var undef +} + +%foo = @compute @workgroup_size(1u, 1u, 1u) func():void { + $B2: { + %3:ptr<private, spirv.explicit_layout_array<vec4<f32>, 4, stride=64>, read_write> = access %var, 1u + %4:spirv.explicit_layout_array<vec4<f32>, 4, stride=64> = load %3 + %5:vec4<f32> = access %4, 0u + %6:vec4<f32> = access %4, 1u + %7:vec4<f32> = access %4, 2u + %8:vec4<f32> = access %4, 3u + %9:mat4x4<f32> = construct %5, %6, %7, %8 + %value:mat4x4<f32> = let %9 + ret + } +} +)"; + + ASSERT_EQ(before, str()); + Run(DecomposeStridedMatrix); + ASSERT_EQ(after, str()); +} + +TEST_F(SpirvReader_DecomposeStridedMatrixTest, LoadMatrix_ViaLet) { + auto* matrix_type = ty.mat4x4<f32>(); + auto* struct_type = Struct(matrix_type, 64); + + auto* var = b.Var("var", ty.ptr<private_>(struct_type)); + mod.root_block->Append(var); + + auto* f = b.ComputeFunction("foo"); + b.Append(f->Block(), [&] { + auto* let = b.Let("ptr", b.Access<ptr<private_, mat4x4<f32>>>(var, 1_u)); + b.Let("value", b.Load(let)); + b.Return(f); + }); + + auto* before = R"( +S = struct @align(64) { + a:u32 @offset(0) + b:mat4x4<f32> @offset(64) @size(256), @matrix_stride(64) + c:u32 @offset(320) +} + +$B1: { # root + %var:ptr<private, S, read_write> = var undef +} + +%foo = @compute @workgroup_size(1u, 1u, 1u) func():void { + $B2: { + %3:ptr<private, mat4x4<f32>, read_write> = access %var, 1u + %ptr:ptr<private, mat4x4<f32>, read_write> = let %3 + %5:mat4x4<f32> = load %ptr + %value:mat4x4<f32> = let %5 + ret + } +} +)"; + auto* after = R"( +S = struct @align(64) { + a:u32 @offset(0) + b:mat4x4<f32> @offset(64) @size(256), @matrix_stride(64) + c:u32 @offset(320) +} + +S_1 = struct @align(64) { + a:u32 @offset(0) + b:spirv.explicit_layout_array<vec4<f32>, 4, stride=64> @offset(64) + c:u32 @offset(320) +} + +$B1: { # root + %var:ptr<private, S_1, read_write> = var undef +} + +%foo = @compute @workgroup_size(1u, 1u, 1u) func():void { + $B2: { + %3:ptr<private, spirv.explicit_layout_array<vec4<f32>, 4, stride=64>, read_write> = access %var, 1u + %ptr:ptr<private, spirv.explicit_layout_array<vec4<f32>, 4, stride=64>, read_write> = let %3 + %5:spirv.explicit_layout_array<vec4<f32>, 4, stride=64> = load %ptr + %6:vec4<f32> = access %5, 0u + %7:vec4<f32> = access %5, 1u + %8:vec4<f32> = access %5, 2u + %9:vec4<f32> = access %5, 3u + %10:mat4x4<f32> = construct %6, %7, %8, %9 + %value:mat4x4<f32> = let %10 + ret + } +} +)"; + + ASSERT_EQ(before, str()); + Run(DecomposeStridedMatrix); + ASSERT_EQ(after, str()); +} + +TEST_F(SpirvReader_DecomposeStridedMatrixTest, LoadStruct) { + auto* matrix_type = ty.mat4x4<f32>(); + auto* struct_type = Struct(matrix_type, 64); + + auto* var = b.Var("var", ty.ptr<private_>(struct_type)); + mod.root_block->Append(var); + + auto* f = b.ComputeFunction("foo"); + b.Append(f->Block(), [&] { + b.Let("value", b.Load(var)); + b.Return(f); + }); + + auto* before = R"( +S = struct @align(64) { + a:u32 @offset(0) + b:mat4x4<f32> @offset(64) @size(256), @matrix_stride(64) + c:u32 @offset(320) +} + +$B1: { # root + %var:ptr<private, S, read_write> = var undef +} + +%foo = @compute @workgroup_size(1u, 1u, 1u) func():void { + $B2: { + %3:S = load %var + %value:S = let %3 + ret + } +} +)"; + auto* after = R"( +S = struct @align(64) { + a:u32 @offset(0) + b:mat4x4<f32> @offset(64) @size(256), @matrix_stride(64) + c:u32 @offset(320) +} + +S_1 = struct @align(64) { + a:u32 @offset(0) + b:spirv.explicit_layout_array<vec4<f32>, 4, stride=64> @offset(64) + c:u32 @offset(320) +} + +$B1: { # root + %var:ptr<private, S_1, read_write> = var undef +} + +%foo = @compute @workgroup_size(1u, 1u, 1u) func():void { + $B2: { + %3:S_1 = load %var + %value:S_1 = let %3 + ret + } +} +)"; + + ASSERT_EQ(before, str()); + Run(DecomposeStridedMatrix); + ASSERT_EQ(after, str()); +} + +TEST_F(SpirvReader_DecomposeStridedMatrixTest, LoadStruct_ViaLet) { + auto* matrix_type = ty.mat4x4<f32>(); + auto* struct_type = Struct(matrix_type, 64); + + auto* var = b.Var("var", ty.ptr<private_>(struct_type)); + mod.root_block->Append(var); + + auto* f = b.ComputeFunction("foo"); + b.Append(f->Block(), [&] { + auto* let = b.Let("ptr", var); + b.Let("value", b.Load(let)); + b.Return(f); + }); + + auto* before = R"( +S = struct @align(64) { + a:u32 @offset(0) + b:mat4x4<f32> @offset(64) @size(256), @matrix_stride(64) + c:u32 @offset(320) +} + +$B1: { # root + %var:ptr<private, S, read_write> = var undef +} + +%foo = @compute @workgroup_size(1u, 1u, 1u) func():void { + $B2: { + %ptr:ptr<private, S, read_write> = let %var + %4:S = load %ptr + %value:S = let %4 + ret + } +} +)"; + auto* after = R"( +S = struct @align(64) { + a:u32 @offset(0) + b:mat4x4<f32> @offset(64) @size(256), @matrix_stride(64) + c:u32 @offset(320) +} + +S_1 = struct @align(64) { + a:u32 @offset(0) + b:spirv.explicit_layout_array<vec4<f32>, 4, stride=64> @offset(64) + c:u32 @offset(320) +} + +$B1: { # root + %var:ptr<private, S_1, read_write> = var undef +} + +%foo = @compute @workgroup_size(1u, 1u, 1u) func():void { + $B2: { + %ptr:ptr<private, S_1, read_write> = let %var + %4:S_1 = load %ptr + %value:S_1 = let %4 + ret + } +} +)"; + + ASSERT_EQ(before, str()); + Run(DecomposeStridedMatrix); + ASSERT_EQ(after, str()); +} + +TEST_F(SpirvReader_DecomposeStridedMatrixTest, LoadStruct_ExtractMatrix) { + auto* matrix_type = ty.mat4x4<f32>(); + auto* struct_type = Struct(matrix_type, 64); + + auto* var = b.Var("var", ty.ptr<private_>(struct_type)); + mod.root_block->Append(var); + + auto* f = b.ComputeFunction("foo"); + b.Append(f->Block(), [&] { + auto* struct_value = b.Let("struct_value", b.Load(var)); + b.Let("matrix_value", b.Access(matrix_type, struct_value, 1_u)); + b.Return(f); + }); + + auto* before = R"( +S = struct @align(64) { + a:u32 @offset(0) + b:mat4x4<f32> @offset(64) @size(256), @matrix_stride(64) + c:u32 @offset(320) +} + +$B1: { # root + %var:ptr<private, S, read_write> = var undef +} + +%foo = @compute @workgroup_size(1u, 1u, 1u) func():void { + $B2: { + %3:S = load %var + %struct_value:S = let %3 + %5:mat4x4<f32> = access %struct_value, 1u + %matrix_value:mat4x4<f32> = let %5 + ret + } +} +)"; + auto* after = R"( +S = struct @align(64) { + a:u32 @offset(0) + b:mat4x4<f32> @offset(64) @size(256), @matrix_stride(64) + c:u32 @offset(320) +} + +S_1 = struct @align(64) { + a:u32 @offset(0) + b:spirv.explicit_layout_array<vec4<f32>, 4, stride=64> @offset(64) + c:u32 @offset(320) +} + +$B1: { # root + %var:ptr<private, S_1, read_write> = var undef +} + +%foo = @compute @workgroup_size(1u, 1u, 1u) func():void { + $B2: { + %3:S_1 = load %var + %struct_value:S_1 = let %3 + %5:spirv.explicit_layout_array<vec4<f32>, 4, stride=64> = access %struct_value, 1u + %6:vec4<f32> = access %5, 0u + %7:vec4<f32> = access %5, 1u + %8:vec4<f32> = access %5, 2u + %9:vec4<f32> = access %5, 3u + %10:mat4x4<f32> = construct %6, %7, %8, %9 + %matrix_value:mat4x4<f32> = let %10 + ret + } +} +)"; + + ASSERT_EQ(before, str()); + Run(DecomposeStridedMatrix); + ASSERT_EQ(after, str()); +} + +TEST_F(SpirvReader_DecomposeStridedMatrixTest, StoreMatrixElement) { + auto* matrix_type = ty.mat4x4<f32>(); + auto* struct_type = Struct(matrix_type, 64); + + auto* var = b.Var("var", ty.ptr<private_>(struct_type)); + mod.root_block->Append(var); + + auto* f = b.ComputeFunction("foo"); + b.Append(f->Block(), [&] { + auto* access = b.Access<ptr<private_, vec4<f32>>>(var, 1_u, 3_u); + b.StoreVectorElement(access, 2_u, 42_f); + b.Return(f); + }); + + auto* before = R"( +S = struct @align(64) { + a:u32 @offset(0) + b:mat4x4<f32> @offset(64) @size(256), @matrix_stride(64) + c:u32 @offset(320) +} + +$B1: { # root + %var:ptr<private, S, read_write> = var undef +} + +%foo = @compute @workgroup_size(1u, 1u, 1u) func():void { + $B2: { + %3:ptr<private, vec4<f32>, read_write> = access %var, 1u, 3u + store_vector_element %3, 2u, 42.0f + ret + } +} +)"; + auto* after = R"( +S = struct @align(64) { + a:u32 @offset(0) + b:mat4x4<f32> @offset(64) @size(256), @matrix_stride(64) + c:u32 @offset(320) +} + +S_1 = struct @align(64) { + a:u32 @offset(0) + b:spirv.explicit_layout_array<vec4<f32>, 4, stride=64> @offset(64) + c:u32 @offset(320) +} + +$B1: { # root + %var:ptr<private, S_1, read_write> = var undef +} + +%foo = @compute @workgroup_size(1u, 1u, 1u) func():void { + $B2: { + %3:ptr<private, vec4<f32>, read_write> = access %var, 1u, 3u + store_vector_element %3, 2u, 42.0f + ret + } +} +)"; + + ASSERT_EQ(before, str()); + Run(DecomposeStridedMatrix); + ASSERT_EQ(after, str()); +} + +TEST_F(SpirvReader_DecomposeStridedMatrixTest, StoreMatrixColumn) { + auto* matrix_type = ty.mat4x4<f32>(); + auto* struct_type = Struct(matrix_type, 64); + + auto* var = b.Var("var", ty.ptr<private_>(struct_type)); + mod.root_block->Append(var); + + auto* f = b.ComputeFunction("foo"); + b.Append(f->Block(), [&] { + b.Store(b.Access<ptr<private_, vec4<f32>>>(var, 1_u, 2_u), b.Zero<vec4<f32>>()); + b.Return(f); + }); + + auto* before = R"( +S = struct @align(64) { + a:u32 @offset(0) + b:mat4x4<f32> @offset(64) @size(256), @matrix_stride(64) + c:u32 @offset(320) +} + +$B1: { # root + %var:ptr<private, S, read_write> = var undef +} + +%foo = @compute @workgroup_size(1u, 1u, 1u) func():void { + $B2: { + %3:ptr<private, vec4<f32>, read_write> = access %var, 1u, 2u + store %3, vec4<f32>(0.0f) + ret + } +} +)"; + auto* after = R"( +S = struct @align(64) { + a:u32 @offset(0) + b:mat4x4<f32> @offset(64) @size(256), @matrix_stride(64) + c:u32 @offset(320) +} + +S_1 = struct @align(64) { + a:u32 @offset(0) + b:spirv.explicit_layout_array<vec4<f32>, 4, stride=64> @offset(64) + c:u32 @offset(320) +} + +$B1: { # root + %var:ptr<private, S_1, read_write> = var undef +} + +%foo = @compute @workgroup_size(1u, 1u, 1u) func():void { + $B2: { + %3:ptr<private, vec4<f32>, read_write> = access %var, 1u, 2u + store %3, vec4<f32>(0.0f) + ret + } +} +)"; + + ASSERT_EQ(before, str()); + Run(DecomposeStridedMatrix); + ASSERT_EQ(after, str()); +} + +TEST_F(SpirvReader_DecomposeStridedMatrixTest, StoreMatrix) { + auto* matrix_type = ty.mat4x4<f32>(); + auto* struct_type = Struct(matrix_type, 64); + + auto* var = b.Var("var", ty.ptr<private_>(struct_type)); + mod.root_block->Append(var); + + auto* f = b.ComputeFunction("foo"); + b.Append(f->Block(), [&] { + b.Store(b.Access<ptr<private_, mat4x4<f32>>>(var, 1_u), b.Zero<mat4x4<f32>>()); + b.Return(f); + }); + + auto* before = R"( +S = struct @align(64) { + a:u32 @offset(0) + b:mat4x4<f32> @offset(64) @size(256), @matrix_stride(64) + c:u32 @offset(320) +} + +$B1: { # root + %var:ptr<private, S, read_write> = var undef +} + +%foo = @compute @workgroup_size(1u, 1u, 1u) func():void { + $B2: { + %3:ptr<private, mat4x4<f32>, read_write> = access %var, 1u + store %3, mat4x4<f32>(vec4<f32>(0.0f)) + ret + } +} +)"; + auto* after = R"( +S = struct @align(64) { + a:u32 @offset(0) + b:mat4x4<f32> @offset(64) @size(256), @matrix_stride(64) + c:u32 @offset(320) +} + +S_1 = struct @align(64) { + a:u32 @offset(0) + b:spirv.explicit_layout_array<vec4<f32>, 4, stride=64> @offset(64) + c:u32 @offset(320) +} + +$B1: { # root + %var:ptr<private, S_1, read_write> = var undef +} + +%foo = @compute @workgroup_size(1u, 1u, 1u) func():void { + $B2: { + %3:ptr<private, spirv.explicit_layout_array<vec4<f32>, 4, stride=64>, read_write> = access %var, 1u + %4:spirv.explicit_layout_array<vec4<f32>, 4, stride=64> = construct vec4<f32>(0.0f), vec4<f32>(0.0f), vec4<f32>(0.0f), vec4<f32>(0.0f) + store %3, %4 + ret + } +} +)"; + + ASSERT_EQ(before, str()); + Run(DecomposeStridedMatrix); + ASSERT_EQ(after, str()); +} + +TEST_F(SpirvReader_DecomposeStridedMatrixTest, StoreStruct) { + auto* matrix_type = ty.mat4x4<f32>(); + auto* struct_type = Struct(matrix_type, 64); + + auto* var = b.Var("var", ty.ptr<private_>(struct_type)); + mod.root_block->Append(var); + + auto* f = b.ComputeFunction("foo"); + b.Append(f->Block(), [&] { + b.Store(var, b.Zero(struct_type)); + b.Return(f); + }); + + auto* before = R"( +S = struct @align(64) { + a:u32 @offset(0) + b:mat4x4<f32> @offset(64) @size(256), @matrix_stride(64) + c:u32 @offset(320) +} + +$B1: { # root + %var:ptr<private, S, read_write> = var undef +} + +%foo = @compute @workgroup_size(1u, 1u, 1u) func():void { + $B2: { + store %var, S(0u, mat4x4<f32>(vec4<f32>(0.0f)), 0u) + ret + } +} +)"; + auto* after = R"( +S = struct @align(64) { + a:u32 @offset(0) + b:mat4x4<f32> @offset(64) @size(256), @matrix_stride(64) + c:u32 @offset(320) +} + +S_1 = struct @align(64) { + a:u32 @offset(0) + b:spirv.explicit_layout_array<vec4<f32>, 4, stride=64> @offset(64) + c:u32 @offset(320) +} + +$B1: { # root + %var:ptr<private, S_1, read_write> = var undef +} + +%foo = @compute @workgroup_size(1u, 1u, 1u) func():void { + $B2: { + store %var, S_1(0u, spirv.explicit_layout_array<vec4<f32>, 4, stride=64>(vec4<f32>(0.0f)), 0u) + ret + } +} +)"; + + ASSERT_EQ(before, str()); + Run(DecomposeStridedMatrix); + ASSERT_EQ(after, str()); +} + +TEST_F(SpirvReader_DecomposeStridedMatrixTest, LoadMatrixFromFuncParam) { + auto* matrix_type = ty.mat4x4<f32>(); + auto* struct_type = Struct(matrix_type, 64); + + auto* foo = b.Function("foo", ty.void_()); + auto* param = b.FunctionParam("param", ty.ptr(function, struct_type)); + foo->SetParams({param}); + b.Append(foo->Block(), [&] { + b.Let("value", b.Load(b.Access<ptr<function, mat4x4<f32>>>(param, 1_u))); + b.Return(foo); + }); + + auto* bar = b.Function("bar", ty.void_()); + b.Append(bar->Block(), [&] { + auto* var = b.Var("var", ty.ptr(function, struct_type)); + b.Call(foo, var); + b.Return(bar); + }); + + auto* before = R"( +S = struct @align(64) { + a:u32 @offset(0) + b:mat4x4<f32> @offset(64) @size(256), @matrix_stride(64) + c:u32 @offset(320) +} + +%foo = func(%param:ptr<function, S, read_write>):void { + $B1: { + %3:ptr<function, mat4x4<f32>, read_write> = access %param, 1u + %4:mat4x4<f32> = load %3 + %value:mat4x4<f32> = let %4 + ret + } +} +%bar = func():void { + $B2: { + %var:ptr<function, S, read_write> = var undef + %8:void = call %foo, %var + ret + } +} +)"; + auto* after = R"( +S = struct @align(64) { + a:u32 @offset(0) + b:mat4x4<f32> @offset(64) @size(256), @matrix_stride(64) + c:u32 @offset(320) +} + +S_1 = struct @align(64) { + a:u32 @offset(0) + b:spirv.explicit_layout_array<vec4<f32>, 4, stride=64> @offset(64) + c:u32 @offset(320) +} + +%foo = func(%param:ptr<function, S_1, read_write>):void { + $B1: { + %3:ptr<function, spirv.explicit_layout_array<vec4<f32>, 4, stride=64>, read_write> = access %param, 1u + %4:spirv.explicit_layout_array<vec4<f32>, 4, stride=64> = load %3 + %5:vec4<f32> = access %4, 0u + %6:vec4<f32> = access %4, 1u + %7:vec4<f32> = access %4, 2u + %8:vec4<f32> = access %4, 3u + %9:mat4x4<f32> = construct %5, %6, %7, %8 + %value:mat4x4<f32> = let %9 + ret + } +} +%bar = func():void { + $B2: { + %var:ptr<function, S_1, read_write> = var undef + %13:void = call %foo, %var + ret + } +} +)"; + + ASSERT_EQ(before, str()); + Run(DecomposeStridedMatrix); + ASSERT_EQ(after, str()); +} + +TEST_F(SpirvReader_DecomposeStridedMatrixTest, ReturnStructFromFunction) { + auto* matrix_type = ty.mat4x4<f32>(); + auto* struct_type = Struct(matrix_type, 64); + + auto* f = b.Function("foo", struct_type); + b.Append(f->Block(), [&] { + auto* var = b.Var("var", ty.ptr(function, struct_type)); + b.Return(f, b.Load(var)); + }); + + auto* before = R"( +S = struct @align(64) { + a:u32 @offset(0) + b:mat4x4<f32> @offset(64) @size(256), @matrix_stride(64) + c:u32 @offset(320) +} + +%foo = func():S { + $B1: { + %var:ptr<function, S, read_write> = var undef + %3:S = load %var + ret %3 + } +} +)"; + auto* after = R"( +S = struct @align(64) { + a:u32 @offset(0) + b:mat4x4<f32> @offset(64) @size(256), @matrix_stride(64) + c:u32 @offset(320) +} + +S_1 = struct @align(64) { + a:u32 @offset(0) + b:spirv.explicit_layout_array<vec4<f32>, 4, stride=64> @offset(64) + c:u32 @offset(320) +} + +%foo = func():S_1 { + $B1: { + %var:ptr<function, S_1, read_write> = var undef + %3:S_1 = load %var + ret %3 + } +} +)"; + + ASSERT_EQ(before, str()); + Run(DecomposeStridedMatrix); + ASSERT_EQ(after, str()); +} + +TEST_F(SpirvReader_DecomposeStridedMatrixTest, LoadMatrix_StructNestedInArray) { + auto* matrix_type = ty.mat4x4<f32>(); + auto* struct_type = Struct(matrix_type, 64); + auto* array_type = ty.array(struct_type, 4); + + auto* var = b.Var("var", ty.ptr<private_>(array_type)); + mod.root_block->Append(var); + + auto* f = b.ComputeFunction("foo"); + b.Append(f->Block(), [&] { + b.Let("value", b.Load(b.Access<ptr<private_, mat4x4<f32>>>(var, 2_u, 1_u))); + b.Return(f); + }); + + auto* before = R"( +S = struct @align(64) { + a:u32 @offset(0) + b:mat4x4<f32> @offset(64) @size(256), @matrix_stride(64) + c:u32 @offset(320) +} + +$B1: { # root + %var:ptr<private, array<S, 4>, read_write> = var undef +} + +%foo = @compute @workgroup_size(1u, 1u, 1u) func():void { + $B2: { + %3:ptr<private, mat4x4<f32>, read_write> = access %var, 2u, 1u + %4:mat4x4<f32> = load %3 + %value:mat4x4<f32> = let %4 + ret + } +} +)"; + auto* after = R"( +S = struct @align(64) { + a:u32 @offset(0) + b:mat4x4<f32> @offset(64) @size(256), @matrix_stride(64) + c:u32 @offset(320) +} + +S_1 = struct @align(64) { + a:u32 @offset(0) + b:spirv.explicit_layout_array<vec4<f32>, 4, stride=64> @offset(64) + c:u32 @offset(320) +} + +$B1: { # root + %var:ptr<private, array<S_1, 4>, read_write> = var undef +} + +%foo = @compute @workgroup_size(1u, 1u, 1u) func():void { + $B2: { + %3:ptr<private, spirv.explicit_layout_array<vec4<f32>, 4, stride=64>, read_write> = access %var, 2u, 1u + %4:spirv.explicit_layout_array<vec4<f32>, 4, stride=64> = load %3 + %5:vec4<f32> = access %4, 0u + %6:vec4<f32> = access %4, 1u + %7:vec4<f32> = access %4, 2u + %8:vec4<f32> = access %4, 3u + %9:mat4x4<f32> = construct %5, %6, %7, %8 + %value:mat4x4<f32> = let %9 + ret + } +} +)"; + + ASSERT_EQ(before, str()); + Run(DecomposeStridedMatrix); + ASSERT_EQ(after, str()); +} + +TEST_F(SpirvReader_DecomposeStridedMatrixTest, LoadMatrix_StructNestedInStridedArray) { + auto* matrix_type = ty.mat4x4<f32>(); + auto* struct_type = Struct(matrix_type, 64); + auto* strided_array = ty.Get<spirv::type::ExplicitLayoutArray>( + struct_type, ty.Get<core::type::ConstantArrayCount>(4u), 16u, 1024u, 256u); + + auto* var = b.Var("var", ty.ptr<private_>(strided_array)); + mod.root_block->Append(var); + + auto* f = b.ComputeFunction("foo"); + b.Append(f->Block(), [&] { + b.Let("value", b.Load(b.Access<ptr<private_, mat4x4<f32>>>(var, 2_u, 1_u))); + b.Return(f); + }); + + auto* before = R"( +S = struct @align(64) { + a:u32 @offset(0) + b:mat4x4<f32> @offset(64) @size(256), @matrix_stride(64) + c:u32 @offset(320) +} + +$B1: { # root + %var:ptr<private, spirv.explicit_layout_array<S, 4, stride=256>, read_write> = var undef +} + +%foo = @compute @workgroup_size(1u, 1u, 1u) func():void { + $B2: { + %3:ptr<private, mat4x4<f32>, read_write> = access %var, 2u, 1u + %4:mat4x4<f32> = load %3 + %value:mat4x4<f32> = let %4 + ret + } +} +)"; + auto* after = R"( +S = struct @align(64) { + a:u32 @offset(0) + b:mat4x4<f32> @offset(64) @size(256), @matrix_stride(64) + c:u32 @offset(320) +} + +S_1 = struct @align(64) { + a:u32 @offset(0) + b:spirv.explicit_layout_array<vec4<f32>, 4, stride=64> @offset(64) + c:u32 @offset(320) +} + +$B1: { # root + %var:ptr<private, spirv.explicit_layout_array<S_1, 4, stride=256>, read_write> = var undef +} + +%foo = @compute @workgroup_size(1u, 1u, 1u) func():void { + $B2: { + %3:ptr<private, spirv.explicit_layout_array<vec4<f32>, 4, stride=64>, read_write> = access %var, 2u, 1u + %4:spirv.explicit_layout_array<vec4<f32>, 4, stride=64> = load %3 + %5:vec4<f32> = access %4, 0u + %6:vec4<f32> = access %4, 1u + %7:vec4<f32> = access %4, 2u + %8:vec4<f32> = access %4, 3u + %9:mat4x4<f32> = construct %5, %6, %7, %8 + %value:mat4x4<f32> = let %9 + ret + } +} +)"; + + ASSERT_EQ(before, str()); + Run(DecomposeStridedMatrix); + ASSERT_EQ(after, str()); +} + +TEST_F(SpirvReader_DecomposeStridedMatrixTest, LoadMatrix_StructNestedInStruct) { + auto* matrix_type = ty.mat4x4<f32>(); + auto* inner_struct_type = Struct(matrix_type, 64); + auto* outer_struct_type = + ty.Struct(mod.symbols.New("Outer"), { + {mod.symbols.New("a"), ty.u32()}, + {mod.symbols.New("b"), inner_struct_type}, + }); + + auto* var = b.Var("var", ty.ptr<private_>(outer_struct_type)); + mod.root_block->Append(var); + + auto* f = b.ComputeFunction("foo"); + b.Append(f->Block(), [&] { + b.Let("value", b.Load(b.Access<ptr<private_, mat4x4<f32>>>(var, 1_u, 1_u))); + b.Return(f); + }); + + auto* before = R"( +S = struct @align(64) { + a:u32 @offset(0) + b:mat4x4<f32> @offset(64) @size(256), @matrix_stride(64) + c:u32 @offset(320) +} + +Outer = struct @align(64) { + a_1:u32 @offset(0) + b_1:S @offset(64) +} + +$B1: { # root + %var:ptr<private, Outer, read_write> = var undef +} + +%foo = @compute @workgroup_size(1u, 1u, 1u) func():void { + $B2: { + %3:ptr<private, mat4x4<f32>, read_write> = access %var, 1u, 1u + %4:mat4x4<f32> = load %3 + %value:mat4x4<f32> = let %4 + ret + } +} +)"; + auto* after = R"( +S = struct @align(64) { + a:u32 @offset(0) + b:mat4x4<f32> @offset(64) @size(256), @matrix_stride(64) + c:u32 @offset(320) +} + +Outer = struct @align(64) { + a_1:u32 @offset(0) + b_1:S @offset(64) +} + +S_1 = struct @align(64) { + a:u32 @offset(0) + b:spirv.explicit_layout_array<vec4<f32>, 4, stride=64> @offset(64) + c:u32 @offset(320) +} + +Outer_1 = struct @align(64) { + a_1:u32 @offset(0) + b_1:S_1 @offset(64) +} + +$B1: { # root + %var:ptr<private, Outer_1, read_write> = var undef +} + +%foo = @compute @workgroup_size(1u, 1u, 1u) func():void { + $B2: { + %3:ptr<private, spirv.explicit_layout_array<vec4<f32>, 4, stride=64>, read_write> = access %var, 1u, 1u + %4:spirv.explicit_layout_array<vec4<f32>, 4, stride=64> = load %3 + %5:vec4<f32> = access %4, 0u + %6:vec4<f32> = access %4, 1u + %7:vec4<f32> = access %4, 2u + %8:vec4<f32> = access %4, 3u + %9:mat4x4<f32> = construct %5, %6, %7, %8 + %value:mat4x4<f32> = let %9 + ret + } +} +)"; + + ASSERT_EQ(before, str()); + Run(DecomposeStridedMatrix); + ASSERT_EQ(after, str()); +} + +TEST_F(SpirvReader_DecomposeStridedMatrixTest, ConstructAndAccess) { + auto* matrix_type = ty.mat4x4<f32>(); + auto* struct_type = Struct(matrix_type, 64); + + auto* f = b.Function("foo", matrix_type); + auto* p0 = b.FunctionParam("p0", ty.u32()); + auto* p1 = b.FunctionParam("p1", matrix_type); + auto* p2 = b.FunctionParam("p2", ty.u32()); + f->SetParams({p0, p1, p2}); + + b.Append(f->Block(), [&] { + auto* c = b.Construct(struct_type, p0, p1, p2); + auto* a = b.Access(matrix_type, c, 1_u); + b.Return(f, a); + }); + + auto* before = R"( +S = struct @align(64) { + a:u32 @offset(0) + b:mat4x4<f32> @offset(64) @size(256), @matrix_stride(64) + c:u32 @offset(320) +} + +%foo = func(%p0:u32, %p1:mat4x4<f32>, %p2:u32):mat4x4<f32> { + $B1: { + %5:S = construct %p0, %p1, %p2 + %6:mat4x4<f32> = access %5, 1u + ret %6 + } +} +)"; + auto* after = R"( +S = struct @align(64) { + a:u32 @offset(0) + b:mat4x4<f32> @offset(64) @size(256), @matrix_stride(64) + c:u32 @offset(320) +} + +S_1 = struct @align(64) { + a:u32 @offset(0) + b:spirv.explicit_layout_array<vec4<f32>, 4, stride=64> @offset(64) + c:u32 @offset(320) +} + +%foo = func(%p0:u32, %p1:mat4x4<f32>, %p2:u32):mat4x4<f32> { + $B1: { + %5:vec4<f32> = access %p1, 0u + %6:vec4<f32> = access %p1, 1u + %7:vec4<f32> = access %p1, 2u + %8:vec4<f32> = access %p1, 3u + %9:spirv.explicit_layout_array<vec4<f32>, 4, stride=64> = construct %5, %6, %7, %8 + %10:S_1 = construct %p0, %9, %p2 + %11:spirv.explicit_layout_array<vec4<f32>, 4, stride=64> = access %10, 1u + %12:vec4<f32> = access %11, 0u + %13:vec4<f32> = access %11, 1u + %14:vec4<f32> = access %11, 2u + %15:vec4<f32> = access %11, 3u + %16:mat4x4<f32> = construct %12, %13, %14, %15 + ret %16 + } +} +)"; + + ASSERT_EQ(before, str()); + Run(DecomposeStridedMatrix); + ASSERT_EQ(after, str()); +} + +TEST_F(SpirvReader_DecomposeStridedMatrixTest, ExistingStridedArray) { + auto* strided_array = ty.Get<spirv::type::ExplicitLayoutArray>( + ty.vec4<f32>(), ty.Get<core::type::ConstantArrayCount>(4u), 16u, 64u, 16u); + auto* struct_ty = + ty.Struct(mod.symbols.New("MyStruct"), { + {mod.symbols.New("a"), ty.u32()}, + {mod.symbols.New("b"), strided_array}, + }); + + auto* var = b.Var("var", ty.ptr<private_>(struct_ty)); + mod.root_block->Append(var); + + auto* f = b.ComputeFunction("foo"); + b.Append(f->Block(), [&] { + auto* array_ptr = b.Access(ty.ptr(private_, strided_array), var, 1_u); + b.Let("value", b.Load(array_ptr)); + b.Return(f); + }); + + auto* before = R"( +MyStruct = struct @align(16) { + a:u32 @offset(0) + b:spirv.explicit_layout_array<vec4<f32>, 4, stride=16> @offset(16) +} + +$B1: { # root + %var:ptr<private, MyStruct, read_write> = var undef +} + +%foo = @compute @workgroup_size(1u, 1u, 1u) func():void { + $B2: { + %3:ptr<private, spirv.explicit_layout_array<vec4<f32>, 4, stride=16>, read_write> = access %var, 1u + %4:spirv.explicit_layout_array<vec4<f32>, 4, stride=16> = load %3 + %value:spirv.explicit_layout_array<vec4<f32>, 4, stride=16> = let %4 + ret + } +} +)"; + auto* after = before; + + ASSERT_EQ(before, str()); + Run(DecomposeStridedMatrix); + ASSERT_EQ(after, str()); +} + +TEST_F(SpirvReader_DecomposeStridedMatrixTest, ArrayOfStridedMatrix_LoadMatrix) { + auto* matrix_type = ty.mat4x4<f32>(); + auto* struct_type = Struct(matrix_type, 64, {2, 3}); + + auto* var = b.Var("var", ty.ptr<private_>(struct_type)); + mod.root_block->Append(var); + + auto* f = b.ComputeFunction("foo"); + b.Append(f->Block(), [&] { + b.Let("value", b.Load(b.Access<ptr<private_, mat4x4<f32>>>(var, 1_u, 2_u, 1_u))); + b.Return(f); + }); + + auto* before = R"( +S = struct @align(64) { + a:u32 @offset(0) + b:array<array<mat4x4<f32>, 2>, 3> @offset(64) @size(1536), @matrix_stride(64) + c:u32 @offset(1600) +} + +$B1: { # root + %var:ptr<private, S, read_write> = var undef +} + +%foo = @compute @workgroup_size(1u, 1u, 1u) func():void { + $B2: { + %3:ptr<private, mat4x4<f32>, read_write> = access %var, 1u, 2u, 1u + %4:mat4x4<f32> = load %3 + %value:mat4x4<f32> = let %4 + ret + } +} +)"; + auto* after = R"( +S = struct @align(64) { + a:u32 @offset(0) + b:array<array<mat4x4<f32>, 2>, 3> @offset(64) @size(1536), @matrix_stride(64) + c:u32 @offset(1600) +} + +S_1 = struct @align(64) { + a:u32 @offset(0) + b:array<array<spirv.explicit_layout_array<vec4<f32>, 4, stride=64>, 2>, 3> @offset(64) @size(1536) + c:u32 @offset(1600) +} + +$B1: { # root + %var:ptr<private, S_1, read_write> = var undef +} + +%foo = @compute @workgroup_size(1u, 1u, 1u) func():void { + $B2: { + %3:ptr<private, spirv.explicit_layout_array<vec4<f32>, 4, stride=64>, read_write> = access %var, 1u, 2u, 1u + %4:spirv.explicit_layout_array<vec4<f32>, 4, stride=64> = load %3 + %5:vec4<f32> = access %4, 0u + %6:vec4<f32> = access %4, 1u + %7:vec4<f32> = access %4, 2u + %8:vec4<f32> = access %4, 3u + %9:mat4x4<f32> = construct %5, %6, %7, %8 + %value:mat4x4<f32> = let %9 + ret + } +} +)"; + + ASSERT_EQ(before, str()); + Run(DecomposeStridedMatrix); + ASSERT_EQ(after, str()); +} + +TEST_F(SpirvReader_DecomposeStridedMatrixTest, ArrayOfStridedMatrix_LoadArray) { + auto* matrix_type = ty.mat4x4<f32>(); + auto* struct_type = Struct(matrix_type, 64, {2, 3}); + + auto* var = b.Var("var", ty.ptr<private_>(struct_type)); + mod.root_block->Append(var); + + auto* f = b.ComputeFunction("foo"); + b.Append(f->Block(), [&] { + b.Let("value", b.Load(b.Access<ptr<private_, array<array<mat4x4<f32>, 2>, 3>>>(var, 1_u))); + b.Return(f); + }); + + auto* before = R"( +S = struct @align(64) { + a:u32 @offset(0) + b:array<array<mat4x4<f32>, 2>, 3> @offset(64) @size(1536), @matrix_stride(64) + c:u32 @offset(1600) +} + +$B1: { # root + %var:ptr<private, S, read_write> = var undef +} + +%foo = @compute @workgroup_size(1u, 1u, 1u) func():void { + $B2: { + %3:ptr<private, array<array<mat4x4<f32>, 2>, 3>, read_write> = access %var, 1u + %4:array<array<mat4x4<f32>, 2>, 3> = load %3 + %value:array<array<mat4x4<f32>, 2>, 3> = let %4 + ret + } +} +)"; + auto* after = R"( +S = struct @align(64) { + a:u32 @offset(0) + b:array<array<mat4x4<f32>, 2>, 3> @offset(64) @size(1536), @matrix_stride(64) + c:u32 @offset(1600) +} + +S_1 = struct @align(64) { + a:u32 @offset(0) + b:array<array<spirv.explicit_layout_array<vec4<f32>, 4, stride=64>, 2>, 3> @offset(64) @size(1536) + c:u32 @offset(1600) +} + +$B1: { # root + %var:ptr<private, S_1, read_write> = var undef +} + +%foo = @compute @workgroup_size(1u, 1u, 1u) func():void { + $B2: { + %3:ptr<private, array<array<spirv.explicit_layout_array<vec4<f32>, 4, stride=64>, 2>, 3>, read_write> = access %var, 1u + %4:array<array<spirv.explicit_layout_array<vec4<f32>, 4, stride=64>, 2>, 3> = load %3 + %5:array<spirv.explicit_layout_array<vec4<f32>, 4, stride=64>, 2> = access %4, 0u + %6:spirv.explicit_layout_array<vec4<f32>, 4, stride=64> = access %5, 0u + %7:vec4<f32> = access %6, 0u + %8:vec4<f32> = access %6, 1u + %9:vec4<f32> = access %6, 2u + %10:vec4<f32> = access %6, 3u + %11:mat4x4<f32> = construct %7, %8, %9, %10 + %12:spirv.explicit_layout_array<vec4<f32>, 4, stride=64> = access %5, 1u + %13:vec4<f32> = access %12, 0u + %14:vec4<f32> = access %12, 1u + %15:vec4<f32> = access %12, 2u + %16:vec4<f32> = access %12, 3u + %17:mat4x4<f32> = construct %13, %14, %15, %16 + %18:array<mat4x4<f32>, 2> = construct %11, %17 + %19:array<spirv.explicit_layout_array<vec4<f32>, 4, stride=64>, 2> = access %4, 1u + %20:spirv.explicit_layout_array<vec4<f32>, 4, stride=64> = access %19, 0u + %21:vec4<f32> = access %20, 0u + %22:vec4<f32> = access %20, 1u + %23:vec4<f32> = access %20, 2u + %24:vec4<f32> = access %20, 3u + %25:mat4x4<f32> = construct %21, %22, %23, %24 + %26:spirv.explicit_layout_array<vec4<f32>, 4, stride=64> = access %19, 1u + %27:vec4<f32> = access %26, 0u + %28:vec4<f32> = access %26, 1u + %29:vec4<f32> = access %26, 2u + %30:vec4<f32> = access %26, 3u + %31:mat4x4<f32> = construct %27, %28, %29, %30 + %32:array<mat4x4<f32>, 2> = construct %25, %31 + %33:array<spirv.explicit_layout_array<vec4<f32>, 4, stride=64>, 2> = access %4, 2u + %34:spirv.explicit_layout_array<vec4<f32>, 4, stride=64> = access %33, 0u + %35:vec4<f32> = access %34, 0u + %36:vec4<f32> = access %34, 1u + %37:vec4<f32> = access %34, 2u + %38:vec4<f32> = access %34, 3u + %39:mat4x4<f32> = construct %35, %36, %37, %38 + %40:spirv.explicit_layout_array<vec4<f32>, 4, stride=64> = access %33, 1u + %41:vec4<f32> = access %40, 0u + %42:vec4<f32> = access %40, 1u + %43:vec4<f32> = access %40, 2u + %44:vec4<f32> = access %40, 3u + %45:mat4x4<f32> = construct %41, %42, %43, %44 + %46:array<mat4x4<f32>, 2> = construct %39, %45 + %47:array<array<mat4x4<f32>, 2>, 3> = construct %18, %32, %46 + %value:array<array<mat4x4<f32>, 2>, 3> = let %47 + ret + } +} +)"; + + ASSERT_EQ(before, str()); + Run(DecomposeStridedMatrix); + ASSERT_EQ(after, str()); +} + +TEST_F(SpirvReader_DecomposeStridedMatrixTest, ArrayOfStridedMatrix_StoreMatrix) { + auto* matrix_type = ty.mat4x4<f32>(); + auto* struct_type = Struct(matrix_type, 64, {2, 3}); + + auto* var = b.Var("var", ty.ptr<private_>(struct_type)); + mod.root_block->Append(var); + + auto* f = b.ComputeFunction("foo"); + b.Append(f->Block(), [&] { + b.Store(b.Access<ptr<private_, mat4x4<f32>>>(var, 1_u, 2_u, 1_u), b.Zero<mat4x4<f32>>()); + b.Return(f); + }); + + auto* before = R"( +S = struct @align(64) { + a:u32 @offset(0) + b:array<array<mat4x4<f32>, 2>, 3> @offset(64) @size(1536), @matrix_stride(64) + c:u32 @offset(1600) +} + +$B1: { # root + %var:ptr<private, S, read_write> = var undef +} + +%foo = @compute @workgroup_size(1u, 1u, 1u) func():void { + $B2: { + %3:ptr<private, mat4x4<f32>, read_write> = access %var, 1u, 2u, 1u + store %3, mat4x4<f32>(vec4<f32>(0.0f)) + ret + } +} +)"; + auto* after = R"( +S = struct @align(64) { + a:u32 @offset(0) + b:array<array<mat4x4<f32>, 2>, 3> @offset(64) @size(1536), @matrix_stride(64) + c:u32 @offset(1600) +} + +S_1 = struct @align(64) { + a:u32 @offset(0) + b:array<array<spirv.explicit_layout_array<vec4<f32>, 4, stride=64>, 2>, 3> @offset(64) @size(1536) + c:u32 @offset(1600) +} + +$B1: { # root + %var:ptr<private, S_1, read_write> = var undef +} + +%foo = @compute @workgroup_size(1u, 1u, 1u) func():void { + $B2: { + %3:ptr<private, spirv.explicit_layout_array<vec4<f32>, 4, stride=64>, read_write> = access %var, 1u, 2u, 1u + %4:spirv.explicit_layout_array<vec4<f32>, 4, stride=64> = construct vec4<f32>(0.0f), vec4<f32>(0.0f), vec4<f32>(0.0f), vec4<f32>(0.0f) + store %3, %4 + ret + } +} +)"; + + ASSERT_EQ(before, str()); + Run(DecomposeStridedMatrix); + ASSERT_EQ(after, str()); +} + +TEST_F(SpirvReader_DecomposeStridedMatrixTest, ArrayOfStridedMatrix_StoreArray) { + auto* matrix_type = ty.mat4x4<f32>(); + auto* struct_type = Struct(matrix_type, 64, {2, 3}); + + auto* var = b.Var("var", ty.ptr<private_>(struct_type)); + mod.root_block->Append(var); + + auto* f = b.ComputeFunction("foo"); + b.Append(f->Block(), [&] { + b.Store(b.Access<ptr<private_, array<array<mat4x4<f32>, 2>, 3>>>(var, 1_u), + b.Zero<array<array<mat4x4<f32>, 2>, 3>>()); + b.Return(f); + }); + + auto* before = R"( +S = struct @align(64) { + a:u32 @offset(0) + b:array<array<mat4x4<f32>, 2>, 3> @offset(64) @size(1536), @matrix_stride(64) + c:u32 @offset(1600) +} + +$B1: { # root + %var:ptr<private, S, read_write> = var undef +} + +%foo = @compute @workgroup_size(1u, 1u, 1u) func():void { + $B2: { + %3:ptr<private, array<array<mat4x4<f32>, 2>, 3>, read_write> = access %var, 1u + store %3, array<array<mat4x4<f32>, 2>, 3>(array<mat4x4<f32>, 2>(mat4x4<f32>(vec4<f32>(0.0f)))) + ret + } +} +)"; + auto* after = R"( +S = struct @align(64) { + a:u32 @offset(0) + b:array<array<mat4x4<f32>, 2>, 3> @offset(64) @size(1536), @matrix_stride(64) + c:u32 @offset(1600) +} + +S_1 = struct @align(64) { + a:u32 @offset(0) + b:array<array<spirv.explicit_layout_array<vec4<f32>, 4, stride=64>, 2>, 3> @offset(64) @size(1536) + c:u32 @offset(1600) +} + +$B1: { # root + %var:ptr<private, S_1, read_write> = var undef +} + +%foo = @compute @workgroup_size(1u, 1u, 1u) func():void { + $B2: { + %3:ptr<private, array<array<spirv.explicit_layout_array<vec4<f32>, 4, stride=64>, 2>, 3>, read_write> = access %var, 1u + %4:spirv.explicit_layout_array<vec4<f32>, 4, stride=64> = construct vec4<f32>(0.0f), vec4<f32>(0.0f), vec4<f32>(0.0f), vec4<f32>(0.0f) + %5:spirv.explicit_layout_array<vec4<f32>, 4, stride=64> = construct vec4<f32>(0.0f), vec4<f32>(0.0f), vec4<f32>(0.0f), vec4<f32>(0.0f) + %6:array<spirv.explicit_layout_array<vec4<f32>, 4, stride=64>, 2> = construct %4, %5 + %7:spirv.explicit_layout_array<vec4<f32>, 4, stride=64> = construct vec4<f32>(0.0f), vec4<f32>(0.0f), vec4<f32>(0.0f), vec4<f32>(0.0f) + %8:spirv.explicit_layout_array<vec4<f32>, 4, stride=64> = construct vec4<f32>(0.0f), vec4<f32>(0.0f), vec4<f32>(0.0f), vec4<f32>(0.0f) + %9:array<spirv.explicit_layout_array<vec4<f32>, 4, stride=64>, 2> = construct %7, %8 + %10:spirv.explicit_layout_array<vec4<f32>, 4, stride=64> = construct vec4<f32>(0.0f), vec4<f32>(0.0f), vec4<f32>(0.0f), vec4<f32>(0.0f) + %11:spirv.explicit_layout_array<vec4<f32>, 4, stride=64> = construct vec4<f32>(0.0f), vec4<f32>(0.0f), vec4<f32>(0.0f), vec4<f32>(0.0f) + %12:array<spirv.explicit_layout_array<vec4<f32>, 4, stride=64>, 2> = construct %10, %11 + %13:array<array<spirv.explicit_layout_array<vec4<f32>, 4, stride=64>, 2>, 3> = construct %6, %9, %12 + store %3, %13 + ret + } +} +)"; + + ASSERT_EQ(before, str()); + Run(DecomposeStridedMatrix); + ASSERT_EQ(after, str()); +} + +TEST_F(SpirvReader_DecomposeStridedMatrixTest, ArrayOfStridedMatrix_StoreStruct) { + auto* matrix_type = ty.mat4x4<f32>(); + auto* struct_type = Struct(matrix_type, 64, {2, 3}); + + auto* var = b.Var("var", ty.ptr<private_>(struct_type)); + mod.root_block->Append(var); + + auto* f = b.ComputeFunction("foo"); + b.Append(f->Block(), [&] { + b.Store(var, b.Zero(struct_type)); + b.Return(f); + }); + + auto* before = R"( +S = struct @align(64) { + a:u32 @offset(0) + b:array<array<mat4x4<f32>, 2>, 3> @offset(64) @size(1536), @matrix_stride(64) + c:u32 @offset(1600) +} + +$B1: { # root + %var:ptr<private, S, read_write> = var undef +} + +%foo = @compute @workgroup_size(1u, 1u, 1u) func():void { + $B2: { + store %var, S(0u, array<array<mat4x4<f32>, 2>, 3>(array<mat4x4<f32>, 2>(mat4x4<f32>(vec4<f32>(0.0f)))), 0u) + ret + } +} +)"; + auto* after = R"( +S = struct @align(64) { + a:u32 @offset(0) + b:array<array<mat4x4<f32>, 2>, 3> @offset(64) @size(1536), @matrix_stride(64) + c:u32 @offset(1600) +} + +S_1 = struct @align(64) { + a:u32 @offset(0) + b:array<array<spirv.explicit_layout_array<vec4<f32>, 4, stride=64>, 2>, 3> @offset(64) @size(1536) + c:u32 @offset(1600) +} + +$B1: { # root + %var:ptr<private, S_1, read_write> = var undef +} + +%foo = @compute @workgroup_size(1u, 1u, 1u) func():void { + $B2: { + store %var, S_1(0u, array<array<spirv.explicit_layout_array<vec4<f32>, 4, stride=64>, 2>, 3>(array<spirv.explicit_layout_array<vec4<f32>, 4, stride=64>, 2>(spirv.explicit_layout_array<vec4<f32>, 4, stride=64>(vec4<f32>(0.0f)))), 0u) + ret + } +} +)"; + + ASSERT_EQ(before, str()); + Run(DecomposeStridedMatrix); + ASSERT_EQ(after, str()); +} + +TEST_F(SpirvReader_DecomposeStridedMatrixTest, ArrayOfStridedMatrix_ConstructAndAccess) { + auto* matrix_type = ty.mat4x4<f32>(); + auto* array_type = ty.array(ty.array(matrix_type, 2), 3); + auto* struct_type = Struct(matrix_type, 64, {2, 3}); + + auto* f = b.Function("foo", array_type); + auto* p0 = b.FunctionParam("p0", ty.u32()); + auto* p1 = b.FunctionParam("p1", array_type); + auto* p2 = b.FunctionParam("p2", ty.u32()); + f->SetParams({p0, p1, p2}); + + b.Append(f->Block(), [&] { + auto* c = b.Construct(struct_type, p0, p1, p2); + auto* a = b.Access(array_type, c, 1_u); + b.Return(f, a); + }); + + auto* before = R"( +S = struct @align(64) { + a:u32 @offset(0) + b:array<array<mat4x4<f32>, 2>, 3> @offset(64) @size(1536), @matrix_stride(64) + c:u32 @offset(1600) +} + +%foo = func(%p0:u32, %p1:array<array<mat4x4<f32>, 2>, 3>, %p2:u32):array<array<mat4x4<f32>, 2>, 3> { + $B1: { + %5:S = construct %p0, %p1, %p2 + %6:array<array<mat4x4<f32>, 2>, 3> = access %5, 1u + ret %6 + } +} +)"; + auto* after = R"( +S = struct @align(64) { + a:u32 @offset(0) + b:array<array<mat4x4<f32>, 2>, 3> @offset(64) @size(1536), @matrix_stride(64) + c:u32 @offset(1600) +} + +S_1 = struct @align(64) { + a:u32 @offset(0) + b:array<array<spirv.explicit_layout_array<vec4<f32>, 4, stride=64>, 2>, 3> @offset(64) @size(1536) + c:u32 @offset(1600) +} + +%foo = func(%p0:u32, %p1:array<array<mat4x4<f32>, 2>, 3>, %p2:u32):array<array<mat4x4<f32>, 2>, 3> { + $B1: { + %5:array<mat4x4<f32>, 2> = access %p1, 0u + %6:mat4x4<f32> = access %5, 0u + %7:vec4<f32> = access %6, 0u + %8:vec4<f32> = access %6, 1u + %9:vec4<f32> = access %6, 2u + %10:vec4<f32> = access %6, 3u + %11:spirv.explicit_layout_array<vec4<f32>, 4, stride=64> = construct %7, %8, %9, %10 + %12:mat4x4<f32> = access %5, 1u + %13:vec4<f32> = access %12, 0u + %14:vec4<f32> = access %12, 1u + %15:vec4<f32> = access %12, 2u + %16:vec4<f32> = access %12, 3u + %17:spirv.explicit_layout_array<vec4<f32>, 4, stride=64> = construct %13, %14, %15, %16 + %18:array<spirv.explicit_layout_array<vec4<f32>, 4, stride=64>, 2> = construct %11, %17 + %19:array<mat4x4<f32>, 2> = access %p1, 1u + %20:mat4x4<f32> = access %19, 0u + %21:vec4<f32> = access %20, 0u + %22:vec4<f32> = access %20, 1u + %23:vec4<f32> = access %20, 2u + %24:vec4<f32> = access %20, 3u + %25:spirv.explicit_layout_array<vec4<f32>, 4, stride=64> = construct %21, %22, %23, %24 + %26:mat4x4<f32> = access %19, 1u + %27:vec4<f32> = access %26, 0u + %28:vec4<f32> = access %26, 1u + %29:vec4<f32> = access %26, 2u + %30:vec4<f32> = access %26, 3u + %31:spirv.explicit_layout_array<vec4<f32>, 4, stride=64> = construct %27, %28, %29, %30 + %32:array<spirv.explicit_layout_array<vec4<f32>, 4, stride=64>, 2> = construct %25, %31 + %33:array<mat4x4<f32>, 2> = access %p1, 2u + %34:mat4x4<f32> = access %33, 0u + %35:vec4<f32> = access %34, 0u + %36:vec4<f32> = access %34, 1u + %37:vec4<f32> = access %34, 2u + %38:vec4<f32> = access %34, 3u + %39:spirv.explicit_layout_array<vec4<f32>, 4, stride=64> = construct %35, %36, %37, %38 + %40:mat4x4<f32> = access %33, 1u + %41:vec4<f32> = access %40, 0u + %42:vec4<f32> = access %40, 1u + %43:vec4<f32> = access %40, 2u + %44:vec4<f32> = access %40, 3u + %45:spirv.explicit_layout_array<vec4<f32>, 4, stride=64> = construct %41, %42, %43, %44 + %46:array<spirv.explicit_layout_array<vec4<f32>, 4, stride=64>, 2> = construct %39, %45 + %47:array<array<spirv.explicit_layout_array<vec4<f32>, 4, stride=64>, 2>, 3> = construct %18, %32, %46 + %48:S_1 = construct %p0, %47, %p2 + %49:array<array<spirv.explicit_layout_array<vec4<f32>, 4, stride=64>, 2>, 3> = access %48, 1u + %50:array<spirv.explicit_layout_array<vec4<f32>, 4, stride=64>, 2> = access %49, 0u + %51:spirv.explicit_layout_array<vec4<f32>, 4, stride=64> = access %50, 0u + %52:vec4<f32> = access %51, 0u + %53:vec4<f32> = access %51, 1u + %54:vec4<f32> = access %51, 2u + %55:vec4<f32> = access %51, 3u + %56:mat4x4<f32> = construct %52, %53, %54, %55 + %57:spirv.explicit_layout_array<vec4<f32>, 4, stride=64> = access %50, 1u + %58:vec4<f32> = access %57, 0u + %59:vec4<f32> = access %57, 1u + %60:vec4<f32> = access %57, 2u + %61:vec4<f32> = access %57, 3u + %62:mat4x4<f32> = construct %58, %59, %60, %61 + %63:array<mat4x4<f32>, 2> = construct %56, %62 + %64:array<spirv.explicit_layout_array<vec4<f32>, 4, stride=64>, 2> = access %49, 1u + %65:spirv.explicit_layout_array<vec4<f32>, 4, stride=64> = access %64, 0u + %66:vec4<f32> = access %65, 0u + %67:vec4<f32> = access %65, 1u + %68:vec4<f32> = access %65, 2u + %69:vec4<f32> = access %65, 3u + %70:mat4x4<f32> = construct %66, %67, %68, %69 + %71:spirv.explicit_layout_array<vec4<f32>, 4, stride=64> = access %64, 1u + %72:vec4<f32> = access %71, 0u + %73:vec4<f32> = access %71, 1u + %74:vec4<f32> = access %71, 2u + %75:vec4<f32> = access %71, 3u + %76:mat4x4<f32> = construct %72, %73, %74, %75 + %77:array<mat4x4<f32>, 2> = construct %70, %76 + %78:array<spirv.explicit_layout_array<vec4<f32>, 4, stride=64>, 2> = access %49, 2u + %79:spirv.explicit_layout_array<vec4<f32>, 4, stride=64> = access %78, 0u + %80:vec4<f32> = access %79, 0u + %81:vec4<f32> = access %79, 1u + %82:vec4<f32> = access %79, 2u + %83:vec4<f32> = access %79, 3u + %84:mat4x4<f32> = construct %80, %81, %82, %83 + %85:spirv.explicit_layout_array<vec4<f32>, 4, stride=64> = access %78, 1u + %86:vec4<f32> = access %85, 0u + %87:vec4<f32> = access %85, 1u + %88:vec4<f32> = access %85, 2u + %89:vec4<f32> = access %85, 3u + %90:mat4x4<f32> = construct %86, %87, %88, %89 + %91:array<mat4x4<f32>, 2> = construct %84, %90 + %92:array<array<mat4x4<f32>, 2>, 3> = construct %63, %77, %91 + ret %92 + } +} +)"; + + ASSERT_EQ(before, str()); + Run(DecomposeStridedMatrix); + ASSERT_EQ(after, str()); +} + +TEST_F(SpirvReader_DecomposeStridedMatrixTest, ArrayOfStridedMatrix_LoadMatrix_AccessChain) { + auto* matrix_type = ty.mat4x4<f32>(); + auto* struct_type = Struct(matrix_type, 64, {2, 3}); + + auto* var = b.Var("var", ty.ptr<private_>(struct_type)); + mod.root_block->Append(var); + + auto* f = b.ComputeFunction("foo"); + b.Append(f->Block(), [&] { + auto* outer_array_ptr = + b.Access(ty.ptr(private_, ty.array(ty.array(matrix_type, 2), 3)), var, 1_u); + auto* inner_array_ptr = + b.Access(ty.ptr(private_, ty.array(matrix_type, 2)), outer_array_ptr, 2_u); + auto* matrix_ptr = b.Access(ty.ptr(private_, matrix_type), inner_array_ptr, 1_u); + b.Let("value", b.Load(matrix_ptr)); + b.Return(f); + }); + + auto* before = R"( +S = struct @align(64) { + a:u32 @offset(0) + b:array<array<mat4x4<f32>, 2>, 3> @offset(64) @size(1536), @matrix_stride(64) + c:u32 @offset(1600) +} + +$B1: { # root + %var:ptr<private, S, read_write> = var undef +} + +%foo = @compute @workgroup_size(1u, 1u, 1u) func():void { + $B2: { + %3:ptr<private, array<array<mat4x4<f32>, 2>, 3>, read_write> = access %var, 1u + %4:ptr<private, array<mat4x4<f32>, 2>, read_write> = access %3, 2u + %5:ptr<private, mat4x4<f32>, read_write> = access %4, 1u + %6:mat4x4<f32> = load %5 + %value:mat4x4<f32> = let %6 + ret + } +} +)"; + auto* after = R"( +S = struct @align(64) { + a:u32 @offset(0) + b:array<array<mat4x4<f32>, 2>, 3> @offset(64) @size(1536), @matrix_stride(64) + c:u32 @offset(1600) +} + +S_1 = struct @align(64) { + a:u32 @offset(0) + b:array<array<spirv.explicit_layout_array<vec4<f32>, 4, stride=64>, 2>, 3> @offset(64) @size(1536) + c:u32 @offset(1600) +} + +$B1: { # root + %var:ptr<private, S_1, read_write> = var undef +} + +%foo = @compute @workgroup_size(1u, 1u, 1u) func():void { + $B2: { + %3:ptr<private, array<array<spirv.explicit_layout_array<vec4<f32>, 4, stride=64>, 2>, 3>, read_write> = access %var, 1u + %4:ptr<private, array<spirv.explicit_layout_array<vec4<f32>, 4, stride=64>, 2>, read_write> = access %3, 2u + %5:ptr<private, spirv.explicit_layout_array<vec4<f32>, 4, stride=64>, read_write> = access %4, 1u + %6:spirv.explicit_layout_array<vec4<f32>, 4, stride=64> = load %5 + %7:vec4<f32> = access %6, 0u + %8:vec4<f32> = access %6, 1u + %9:vec4<f32> = access %6, 2u + %10:vec4<f32> = access %6, 3u + %11:mat4x4<f32> = construct %7, %8, %9, %10 + %value:mat4x4<f32> = let %11 + ret + } +} +)"; + + ASSERT_EQ(before, str()); + Run(DecomposeStridedMatrix); + ASSERT_EQ(after, str()); +} + +TEST_F(SpirvReader_DecomposeStridedMatrixTest, AccessNonMatrixPointer) { + auto* matrix_type = ty.mat4x4<f32>(); + auto* inner_struct_type = Struct(matrix_type, 64); + auto* outer_struct_type = + ty.Struct(mod.symbols.New("Outer"), { + {mod.symbols.New("s"), inner_struct_type}, + }); + + auto* var = b.Var("var", ty.ptr<private_>(outer_struct_type)); + mod.root_block->Append(var); + + auto* f = b.ComputeFunction("foo"); + b.Append(f->Block(), [&] { + auto* access = b.Access<ptr<private_, u32>>(var, 0_u, 2_u); + b.Let("value", b.Load(access)); + b.Return(f); + }); + + auto* before = R"( +S = struct @align(64) { + a:u32 @offset(0) + b:mat4x4<f32> @offset(64) @size(256), @matrix_stride(64) + c:u32 @offset(320) +} + +Outer = struct @align(64) { + s:S @offset(0) +} + +$B1: { # root + %var:ptr<private, Outer, read_write> = var undef +} + +%foo = @compute @workgroup_size(1u, 1u, 1u) func():void { + $B2: { + %3:ptr<private, u32, read_write> = access %var, 0u, 2u + %4:u32 = load %3 + %value:u32 = let %4 + ret + } +} +)"; + auto* after = R"( +S = struct @align(64) { + a:u32 @offset(0) + b:mat4x4<f32> @offset(64) @size(256), @matrix_stride(64) + c:u32 @offset(320) +} + +Outer = struct @align(64) { + s:S @offset(0) +} + +S_1 = struct @align(64) { + a:u32 @offset(0) + b:spirv.explicit_layout_array<vec4<f32>, 4, stride=64> @offset(64) + c:u32 @offset(320) +} + +Outer_1 = struct @align(64) { + s:S_1 @offset(0) +} + +$B1: { # root + %var:ptr<private, Outer_1, read_write> = var undef +} + +%foo = @compute @workgroup_size(1u, 1u, 1u) func():void { + $B2: { + %3:ptr<private, u32, read_write> = access %var, 0u, 2u + %4:u32 = load %3 + %value:u32 = let %4 + ret + } +} +)"; + + ASSERT_EQ(before, str()); + Run(DecomposeStridedMatrix); + ASSERT_EQ(after, str()); +} + +TEST_F(SpirvReader_DecomposeStridedMatrixTest, AccessNonMatrixValue) { + auto* matrix_type = ty.mat4x4<f32>(); + auto* inner_struct_type = Struct(matrix_type, 64); + auto* outer_struct_type = + ty.Struct(mod.symbols.New("Outer"), { + {mod.symbols.New("s"), inner_struct_type}, + }); + + auto* var = b.Var("var", ty.ptr<private_>(outer_struct_type)); + mod.root_block->Append(var); + + auto* f = b.ComputeFunction("foo"); + b.Append(f->Block(), [&] { + auto* struct_value = b.Load(var); + b.Let("value", b.Access<u32>(struct_value, 0_u, 2_u)); + b.Return(f); + }); + + auto* before = R"( +S = struct @align(64) { + a:u32 @offset(0) + b:mat4x4<f32> @offset(64) @size(256), @matrix_stride(64) + c:u32 @offset(320) +} + +Outer = struct @align(64) { + s:S @offset(0) +} + +$B1: { # root + %var:ptr<private, Outer, read_write> = var undef +} + +%foo = @compute @workgroup_size(1u, 1u, 1u) func():void { + $B2: { + %3:Outer = load %var + %4:u32 = access %3, 0u, 2u + %value:u32 = let %4 + ret + } +} +)"; + auto* after = R"( +S = struct @align(64) { + a:u32 @offset(0) + b:mat4x4<f32> @offset(64) @size(256), @matrix_stride(64) + c:u32 @offset(320) +} + +Outer = struct @align(64) { + s:S @offset(0) +} + +S_1 = struct @align(64) { + a:u32 @offset(0) + b:spirv.explicit_layout_array<vec4<f32>, 4, stride=64> @offset(64) + c:u32 @offset(320) +} + +Outer_1 = struct @align(64) { + s:S_1 @offset(0) +} + +$B1: { # root + %var:ptr<private, Outer_1, read_write> = var undef +} + +%foo = @compute @workgroup_size(1u, 1u, 1u) func():void { + $B2: { + %3:Outer_1 = load %var + %4:u32 = access %3, 0u, 2u + %value:u32 = let %4 + ret + } +} +)"; + + ASSERT_EQ(before, str()); + Run(DecomposeStridedMatrix); + ASSERT_EQ(after, str()); +} + +} // namespace +} // namespace tint::spirv::reader::lower
diff --git a/src/tint/lang/spirv/reader/lower/lower.cc b/src/tint/lang/spirv/reader/lower/lower.cc index 5efc21b..7b01506 100644 --- a/src/tint/lang/spirv/reader/lower/lower.cc +++ b/src/tint/lang/spirv/reader/lower/lower.cc
@@ -33,6 +33,7 @@ #include "src/tint/lang/spirv/reader/lower/atomics.h" #include "src/tint/lang/spirv/reader/lower/builtins.h" #include "src/tint/lang/spirv/reader/lower/decompose_strided_array.h" +#include "src/tint/lang/spirv/reader/lower/decompose_strided_matrix.h" #include "src/tint/lang/spirv/reader/lower/shader_io.h" #include "src/tint/lang/spirv/reader/lower/texture.h" #include "src/tint/lang/spirv/reader/lower/vector_element_pointer.h" @@ -52,6 +53,9 @@ RUN_TRANSFORM(lower::VectorElementPointer, mod); RUN_TRANSFORM(lower::ShaderIO, mod); RUN_TRANSFORM(lower::Builtins, mod); + // DecomposeStridedMatrix must come before DecomposeStridedArray, as it introduces strided + // arrays that need to be replaced. + RUN_TRANSFORM(lower::DecomposeStridedMatrix, mod); RUN_TRANSFORM(lower::DecomposeStridedArray, mod); RUN_TRANSFORM(lower::Atomics, mod); RUN_TRANSFORM(lower::Texture, mod);
diff --git a/src/tint/lang/spirv/reader/lower/shader_io.cc b/src/tint/lang/spirv/reader/lower/shader_io.cc index 8750942..fdce4f5 100644 --- a/src/tint/lang/spirv/reader/lower/shader_io.cc +++ b/src/tint/lang/spirv/reader/lower/shader_io.cc
@@ -695,6 +695,7 @@ core::ir::Capability::kAllowOverrides, core::ir::Capability::kAllowPhonyInstructions, core::ir::Capability::kAllowNonCoreTypes, + core::ir::Capability::kAllowStructMatrixDecorations, }); if (result != Success) { return result.Failure();
diff --git a/src/tint/lang/spirv/reader/lower/vector_element_pointer.cc b/src/tint/lang/spirv/reader/lower/vector_element_pointer.cc index ae0c258..ed278a3 100644 --- a/src/tint/lang/spirv/reader/lower/vector_element_pointer.cc +++ b/src/tint/lang/spirv/reader/lower/vector_element_pointer.cc
@@ -157,6 +157,7 @@ core::ir::Capability::kAllowVectorElementPointer, core::ir::Capability::kAllowPhonyInstructions, core::ir::Capability::kAllowNonCoreTypes, + core::ir::Capability::kAllowStructMatrixDecorations, }); if (result != Success) { return result.Failure();