src/tint/transform/decompose_strided_matrix.cc - dawn - Git at Google

 // Copyright 2021 The Tint Authors.
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
 // You may obtain a copy of the License at
 //
 //     http://www.apache.org/licenses/LICENSE-2.0
 //
 // Unless required by applicable law or agreed to in writing, software
 // distributed under the License is distributed on an "AS IS" BASIS,
 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 // See the License for the specific language governing permissions and
 // limitations under the License.

 #include "src/tint/transform/decompose_strided_matrix.h"

 #include <unordered_map>
 #include <utility>
 #include <vector>

 #include "src/tint/program_builder.h"
 #include "src/tint/sem/expression.h"
 #include "src/tint/sem/member_accessor_expression.h"
 #include "src/tint/transform/simplify_pointers.h"
 #include "src/tint/utils/hash.h"
 #include "src/tint/utils/map.h"

 TINT_INSTANTIATE_TYPEINFO(tint::transform::DecomposeStridedMatrix);

 namespace tint::transform {
 namespace {

 /// MatrixInfo describes a matrix member with a custom stride
 struct MatrixInfo {
     /// The stride in bytes between columns of the matrix
     uint32_t stride = 0;
     /// The type of the matrix
     const sem::Matrix* matrix = nullptr;

     /// @returns a new ast::Array that holds an vector column for each row of the
     /// matrix.
     const ast::Array* array(ProgramBuilder* b) const {
         return b->ty.array(b->ty.vec<f32>(matrix->rows()), u32(matrix->columns()), stride);
     }

     /// Equality operator
     bool operator==(const MatrixInfo& info) const {
         return stride == info.stride && matrix == info.matrix;
     }
     /// Hash function
     struct Hasher {
         size_t operator()(const MatrixInfo& t) const { return utils::Hash(t.stride, t.matrix); }
     };
 };

 /// Return type of the callback function of GatherCustomStrideMatrixMembers
 enum GatherResult { kContinue, kStop };

 /// GatherCustomStrideMatrixMembers scans `program` for all matrix members of
 /// storage and uniform structs, which are of a matrix type, and have a custom
 /// matrix stride attribute. For each matrix member found, `callback` is called.
 /// `callback` is a function with the signature:
 ///      GatherResult(const sem::StructMember* member,
 ///                   sem::Matrix* matrix,
 ///                   uint32_t stride)
 /// If `callback` return GatherResult::kStop, then the scanning will immediately
 /// terminate, and GatherCustomStrideMatrixMembers() will return, otherwise
 /// scanning will continue.
 template <typename F>
 void GatherCustomStrideMatrixMembers(const Program* program, F&& callback) {
     for (auto* node : program->ASTNodes().Objects()) {
         if (auto* str = node->As<ast::Struct>()) {
             auto* str_ty = program->Sem().Get(str);
             if (!str_ty->UsedAs(ast::StorageClass::kUniform) &&
                 !str_ty->UsedAs(ast::StorageClass::kStorage)) {
                 continue;
             }
             for (auto* member : str_ty->Members()) {
                 auto* matrix = member->Type()->As<sem::Matrix>();
                 if (!matrix) {
                     continue;
                 }
                 auto* attr =
                     ast::GetAttribute<ast::StrideAttribute>(member->Declaration()->attributes);
                 if (!attr) {
                     continue;
                 }
                 uint32_t stride = attr->stride;
                 if (matrix->ColumnStride() == stride) {
                     continue;
                 }
                 if (callback(member, matrix, stride) == GatherResult::kStop) {
                     return;
                 }
             }
         }
     }
 }

 }  // namespace

 DecomposeStridedMatrix::DecomposeStridedMatrix() = default;

 DecomposeStridedMatrix::~DecomposeStridedMatrix() = default;

 bool DecomposeStridedMatrix::ShouldRun(const Program* program, const DataMap&) const {
     bool should_run = false;
     GatherCustomStrideMatrixMembers(program,
                                     [&](const sem::StructMember*, const sem::Matrix*, uint32_t) {
                                         should_run = true;
                                         return GatherResult::kStop;
                                     });
     return should_run;
 }

 void DecomposeStridedMatrix::Run(CloneContext& ctx, const DataMap&, DataMap&) const {
     // Scan the program for all storage and uniform structure matrix members with
     // a custom stride attribute. Replace these matrices with an equivalent array,
     // and populate the `decomposed` map with the members that have been replaced.
     std::unordered_map<const ast::StructMember*, MatrixInfo> decomposed;
     GatherCustomStrideMatrixMembers(
         ctx.src, [&](const sem::StructMember* member, const sem::Matrix* matrix, uint32_t stride) {
             // We've got ourselves a struct member of a matrix type with a custom
             // stride. Replace this with an array of column vectors.
             MatrixInfo info{stride, matrix};
             auto* replacement =
                 ctx.dst->Member(member->Offset(), ctx.Clone(member->Name()), info.array(ctx.dst));
             ctx.Replace(member->Declaration(), replacement);
             decomposed.emplace(member->Declaration(), info);
             return GatherResult::kContinue;
         });

     // For all expressions where a single matrix column vector was indexed, we can
     // preserve these without calling conversion functions.
     // Example:
     //   ssbo.mat[2] -> ssbo.mat[2]
     ctx.ReplaceAll(
         [&](const ast::IndexAccessorExpression* expr) -> const ast::IndexAccessorExpression* {
             if (auto* access = ctx.src->Sem().Get<sem::StructMemberAccess>(expr->object)) {
                 auto it = decomposed.find(access->Member()->Declaration());
                 if (it != decomposed.end()) {
                     auto* obj = ctx.CloneWithoutTransform(expr->object);
                     auto* idx = ctx.Clone(expr->index);
                     return ctx.dst->IndexAccessor(obj, idx);
                 }
             }
             return nullptr;
         });

     // For all struct member accesses to the matrix on the LHS of an assignment,
     // we need to convert the matrix to the array before assigning to the
     // structure.
     // Example:
     //   ssbo.mat = mat_to_arr(m)
     std::unordered_map<MatrixInfo, Symbol, MatrixInfo::Hasher> mat_to_arr;
     ctx.ReplaceAll([&](const ast::AssignmentStatement* stmt) -> const ast::Statement* {
         if (auto* access = ctx.src->Sem().Get<sem::StructMemberAccess>(stmt->lhs)) {
             auto it = decomposed.find(access->Member()->Declaration());
             if (it == decomposed.end()) {
                 return nullptr;
             }
             MatrixInfo info = it->second;
             auto fn = utils::GetOrCreate(mat_to_arr, info, [&] {
                 auto name =
                     ctx.dst->Symbols().New("mat" + std::to_string(info.matrix->columns()) + "x" +
                                            std::to_string(info.matrix->rows()) + "_stride_" +
                                            std::to_string(info.stride) + "_to_arr");

                 auto matrix = [&] { return CreateASTTypeFor(ctx, info.matrix); };
                 auto array = [&] { return info.array(ctx.dst); };

                 auto mat = ctx.dst->Sym("m");
                 utils::Vector<const ast::Expression*, 4> columns;
                 for (uint32_t i = 0; i < static_cast<uint32_t>(info.matrix->columns()); i++) {
                     columns.Push(ctx.dst->IndexAccessor(mat, u32(i)));
                 }
                 ctx.dst->Func(name,
                               utils::Vector{
                                   ctx.dst->Param(mat, matrix()),
                               },
                               array(),
                               utils::Vector{
                                   ctx.dst->Return(ctx.dst->Construct(array(), columns)),
                               });
                 return name;
             });
             auto* lhs = ctx.CloneWithoutTransform(stmt->lhs);
             auto* rhs = ctx.dst->Call(fn, ctx.Clone(stmt->rhs));
             return ctx.dst->Assign(lhs, rhs);
         }
         return nullptr;
     });

     // For all other struct member accesses, we need to convert the array to the
     // matrix type. Example:
     //   m = arr_to_mat(ssbo.mat)
     std::unordered_map<MatrixInfo, Symbol, MatrixInfo::Hasher> arr_to_mat;
     ctx.ReplaceAll([&](const ast::MemberAccessorExpression* expr) -> const ast::Expression* {
         if (auto* access = ctx.src->Sem().Get<sem::StructMemberAccess>(expr)) {
             auto it = decomposed.find(access->Member()->Declaration());
             if (it == decomposed.end()) {
                 return nullptr;
             }
             MatrixInfo info = it->second;
             auto fn = utils::GetOrCreate(arr_to_mat, info, [&] {
                 auto name = ctx.dst->Symbols().New(
                     "arr_to_mat" + std::to_string(info.matrix->columns()) + "x" +
                     std::to_string(info.matrix->rows()) + "_stride_" + std::to_string(info.stride));

                 auto matrix = [&] { return CreateASTTypeFor(ctx, info.matrix); };
                 auto array = [&] { return info.array(ctx.dst); };

                 auto arr = ctx.dst->Sym("arr");
                 utils::Vector<const ast::Expression*, 4> columns;
                 for (uint32_t i = 0; i < static_cast<uint32_t>(info.matrix->columns()); i++) {
                     columns.Push(ctx.dst->IndexAccessor(arr, u32(i)));
                 }
                 ctx.dst->Func(name,
                               utils::Vector{
                                   ctx.dst->Param(arr, array()),
                               },
                               matrix(),
                               utils::Vector{
                                   ctx.dst->Return(ctx.dst->Construct(matrix(), columns)),
                               });
                 return name;
             });
             return ctx.dst->Call(fn, ctx.CloneWithoutTransform(expr));
         }
         return nullptr;
     });

     ctx.Clone();
 }

 }  // namespace tint::transform
	// Copyright 2021 The Tint Authors.
	//
	// Licensed under the Apache License, Version 2.0 (the "License");
	// you may not use this file except in compliance with the License.
	// You may obtain a copy of the License at
	//
	// http://www.apache.org/licenses/LICENSE-2.0
	//
	// Unless required by applicable law or agreed to in writing, software
	// distributed under the License is distributed on an "AS IS" BASIS,
	// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	// See the License for the specific language governing permissions and
	// limitations under the License.

	#include "src/tint/transform/decompose_strided_matrix.h"

	#include <unordered_map>
	#include <utility>
	#include <vector>

	#include "src/tint/program_builder.h"
	#include "src/tint/sem/expression.h"
	#include "src/tint/sem/member_accessor_expression.h"
	#include "src/tint/transform/simplify_pointers.h"
	#include "src/tint/utils/hash.h"
	#include "src/tint/utils/map.h"

	TINT_INSTANTIATE_TYPEINFO(tint::transform::DecomposeStridedMatrix);

	namespace tint::transform {
	namespace {

	/// MatrixInfo describes a matrix member with a custom stride
	struct MatrixInfo {
	/// The stride in bytes between columns of the matrix
	uint32_t stride = 0;
	/// The type of the matrix
	const sem::Matrix* matrix = nullptr;

	/// @returns a new ast::Array that holds an vector column for each row of the
	/// matrix.
	const ast::Array* array(ProgramBuilder* b) const {
	return b->ty.array(b->ty.vec<f32>(matrix->rows()), u32(matrix->columns()), stride);
	}

	/// Equality operator
	bool operator==(const MatrixInfo& info) const {
	return stride == info.stride && matrix == info.matrix;
	}
	/// Hash function
	struct Hasher {
	size_t operator()(const MatrixInfo& t) const { return utils::Hash(t.stride, t.matrix); }
	};
	};

	/// Return type of the callback function of GatherCustomStrideMatrixMembers
	enum GatherResult { kContinue, kStop };

	/// GatherCustomStrideMatrixMembers scans `program` for all matrix members of
	/// storage and uniform structs, which are of a matrix type, and have a custom
	/// matrix stride attribute. For each matrix member found, `callback` is called.
	/// `callback` is a function with the signature:
	/// GatherResult(const sem::StructMember* member,
	/// sem::Matrix* matrix,
	/// uint32_t stride)
	/// If `callback` return GatherResult::kStop, then the scanning will immediately
	/// terminate, and GatherCustomStrideMatrixMembers() will return, otherwise
	/// scanning will continue.
	template <typename F>
	void GatherCustomStrideMatrixMembers(const Program* program, F&& callback) {
	for (auto* node : program->ASTNodes().Objects()) {
	if (auto* str = node->As<ast::Struct>()) {
	auto* str_ty = program->Sem().Get(str);
	if (!str_ty->UsedAs(ast::StorageClass::kUniform) &&
	!str_ty->UsedAs(ast::StorageClass::kStorage)) {
	continue;
	}
	for (auto* member : str_ty->Members()) {
	auto* matrix = member->Type()->As<sem::Matrix>();
	if (!matrix) {
	continue;
	}
	auto* attr =
	ast::GetAttribute<ast::StrideAttribute>(member->Declaration()->attributes);
	if (!attr) {
	continue;
	}
	uint32_t stride = attr->stride;
	if (matrix->ColumnStride() == stride) {
	continue;
	}
	if (callback(member, matrix, stride) == GatherResult::kStop) {
	return;
	}
	}
	}
	}
	}

	} // namespace

	DecomposeStridedMatrix::DecomposeStridedMatrix() = default;

	DecomposeStridedMatrix::~DecomposeStridedMatrix() = default;

	bool DecomposeStridedMatrix::ShouldRun(const Program* program, const DataMap&) const {
	bool should_run = false;
	GatherCustomStrideMatrixMembers(program,
	[&](const sem::StructMember, const sem::Matrix, uint32_t) {
	should_run = true;
	return GatherResult::kStop;
	});
	return should_run;
	}

	void DecomposeStridedMatrix::Run(CloneContext& ctx, const DataMap&, DataMap&) const {
	// Scan the program for all storage and uniform structure matrix members with
	// a custom stride attribute. Replace these matrices with an equivalent array,
	// and populate the `decomposed` map with the members that have been replaced.
	std::unordered_map<const ast::StructMember*, MatrixInfo> decomposed;
	GatherCustomStrideMatrixMembers(
	ctx.src, [&](const sem::StructMember* member, const sem::Matrix* matrix, uint32_t stride) {
	// We've got ourselves a struct member of a matrix type with a custom
	// stride. Replace this with an array of column vectors.
	MatrixInfo info{stride, matrix};
	auto* replacement =
	ctx.dst->Member(member->Offset(), ctx.Clone(member->Name()), info.array(ctx.dst));
	ctx.Replace(member->Declaration(), replacement);
	decomposed.emplace(member->Declaration(), info);
	return GatherResult::kContinue;
	});

	// For all expressions where a single matrix column vector was indexed, we can
	// preserve these without calling conversion functions.
	// Example:
	// ssbo.mat[2] -> ssbo.mat[2]
	ctx.ReplaceAll(
	[&](const ast::IndexAccessorExpression* expr) -> const ast::IndexAccessorExpression* {
	if (auto* access = ctx.src->Sem().Get<sem::StructMemberAccess>(expr->object)) {
	auto it = decomposed.find(access->Member()->Declaration());
	if (it != decomposed.end()) {
	auto* obj = ctx.CloneWithoutTransform(expr->object);
	auto* idx = ctx.Clone(expr->index);
	return ctx.dst->IndexAccessor(obj, idx);
	}
	}
	return nullptr;
	});

	// For all struct member accesses to the matrix on the LHS of an assignment,
	// we need to convert the matrix to the array before assigning to the
	// structure.
	// Example:
	// ssbo.mat = mat_to_arr(m)
	std::unordered_map<MatrixInfo, Symbol, MatrixInfo::Hasher> mat_to_arr;
	ctx.ReplaceAll([&](const ast::AssignmentStatement* stmt) -> const ast::Statement* {
	if (auto* access = ctx.src->Sem().Get<sem::StructMemberAccess>(stmt->lhs)) {
	auto it = decomposed.find(access->Member()->Declaration());
	if (it == decomposed.end()) {
	return nullptr;
	}
	MatrixInfo info = it->second;
	auto fn = utils::GetOrCreate(mat_to_arr, info, [&] {
	auto name =
	ctx.dst->Symbols().New("mat" + std::to_string(info.matrix->columns()) + "x" +
	std::to_string(info.matrix->rows()) + "_stride_" +
	std::to_string(info.stride) + "_to_arr");

	auto matrix = [&] { return CreateASTTypeFor(ctx, info.matrix); };
	auto array = [&] { return info.array(ctx.dst); };

	auto mat = ctx.dst->Sym("m");
	utils::Vector<const ast::Expression*, 4> columns;
	for (uint32_t i = 0; i < static_cast<uint32_t>(info.matrix->columns()); i++) {
	columns.Push(ctx.dst->IndexAccessor(mat, u32(i)));
	}
	ctx.dst->Func(name,
	utils::Vector{
	ctx.dst->Param(mat, matrix()),
	},
	array(),
	utils::Vector{
	ctx.dst->Return(ctx.dst->Construct(array(), columns)),
	});
	return name;
	});
	auto* lhs = ctx.CloneWithoutTransform(stmt->lhs);
	auto* rhs = ctx.dst->Call(fn, ctx.Clone(stmt->rhs));
	return ctx.dst->Assign(lhs, rhs);
	}
	return nullptr;
	});

	// For all other struct member accesses, we need to convert the array to the
	// matrix type. Example:
	// m = arr_to_mat(ssbo.mat)
	std::unordered_map<MatrixInfo, Symbol, MatrixInfo::Hasher> arr_to_mat;
	ctx.ReplaceAll([&](const ast::MemberAccessorExpression* expr) -> const ast::Expression* {
	if (auto* access = ctx.src->Sem().Get<sem::StructMemberAccess>(expr)) {
	auto it = decomposed.find(access->Member()->Declaration());
	if (it == decomposed.end()) {
	return nullptr;
	}
	MatrixInfo info = it->second;
	auto fn = utils::GetOrCreate(arr_to_mat, info, [&] {
	auto name = ctx.dst->Symbols().New(
	"arr_to_mat" + std::to_string(info.matrix->columns()) + "x" +
	std::to_string(info.matrix->rows()) + "_stride_" + std::to_string(info.stride));

	auto matrix = [&] { return CreateASTTypeFor(ctx, info.matrix); };
	auto array = [&] { return info.array(ctx.dst); };

	auto arr = ctx.dst->Sym("arr");
	utils::Vector<const ast::Expression*, 4> columns;
	for (uint32_t i = 0; i < static_cast<uint32_t>(info.matrix->columns()); i++) {
	columns.Push(ctx.dst->IndexAccessor(arr, u32(i)));
	}
	ctx.dst->Func(name,
	utils::Vector{
	ctx.dst->Param(arr, array()),
	},
	matrix(),
	utils::Vector{
	ctx.dst->Return(ctx.dst->Construct(matrix(), columns)),
	});
	return name;
	});
	return ctx.dst->Call(fn, ctx.CloneWithoutTransform(expr));
	}
	return nullptr;
	});

	ctx.Clone();
	}

	} // namespace tint::transform