src/tint/transform/calculate_array_length.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/calculate_array_length.h"

 #include <unordered_map>
 #include <utility>

 #include "src/tint/ast/call_statement.h"
 #include "src/tint/ast/disable_validation_attribute.h"
 #include "src/tint/program_builder.h"
 #include "src/tint/sem/block_statement.h"
 #include "src/tint/sem/call.h"
 #include "src/tint/sem/function.h"
 #include "src/tint/sem/statement.h"
 #include "src/tint/sem/struct.h"
 #include "src/tint/sem/variable.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::CalculateArrayLength);
 TINT_INSTANTIATE_TYPEINFO(
     tint::transform::CalculateArrayLength::BufferSizeIntrinsic);

 namespace tint {
 namespace transform {

 namespace {

 /// ArrayUsage describes a runtime array usage.
 /// It is used as a key by the array_length_by_usage map.
 struct ArrayUsage {
   ast::BlockStatement const* const block;
   sem::Variable const* const buffer;
   bool operator==(const ArrayUsage& rhs) const {
     return block == rhs.block && buffer == rhs.buffer;
   }
   struct Hasher {
     inline std::size_t operator()(const ArrayUsage& u) const {
       return utils::Hash(u.block, u.buffer);
     }
   };
 };

 }  // namespace

 CalculateArrayLength::BufferSizeIntrinsic::BufferSizeIntrinsic(ProgramID pid)
     : Base(pid) {}
 CalculateArrayLength::BufferSizeIntrinsic::~BufferSizeIntrinsic() = default;
 std::string CalculateArrayLength::BufferSizeIntrinsic::InternalName() const {
   return "intrinsic_buffer_size";
 }

 const CalculateArrayLength::BufferSizeIntrinsic*
 CalculateArrayLength::BufferSizeIntrinsic::Clone(CloneContext* ctx) const {
   return ctx->dst->ASTNodes().Create<CalculateArrayLength::BufferSizeIntrinsic>(
       ctx->dst->ID());
 }

 CalculateArrayLength::CalculateArrayLength() = default;
 CalculateArrayLength::~CalculateArrayLength() = default;

 bool CalculateArrayLength::ShouldRun(const Program* program,
                                      const DataMap&) const {
   for (auto* fn : program->AST().Functions()) {
     if (auto* sem_fn = program->Sem().Get(fn)) {
       for (auto* builtin : sem_fn->DirectlyCalledBuiltins()) {
         if (builtin->Type() == sem::BuiltinType::kArrayLength) {
           return true;
         }
       }
     }
   }
   return false;
 }

 void CalculateArrayLength::Run(CloneContext& ctx,
                                const DataMap&,
                                DataMap&) const {
   auto& sem = ctx.src->Sem();

   // get_buffer_size_intrinsic() emits the function decorated with
   // BufferSizeIntrinsic that is transformed by the HLSL writer into a call to
   // [RW]ByteAddressBuffer.GetDimensions().
   std::unordered_map<const sem::Type*, Symbol> buffer_size_intrinsics;
   auto get_buffer_size_intrinsic = [&](const sem::Type* buffer_type) {
     return utils::GetOrCreate(buffer_size_intrinsics, buffer_type, [&] {
       auto name = ctx.dst->Sym();
       auto* type = CreateASTTypeFor(ctx, buffer_type);
       auto* disable_validation = ctx.dst->Disable(
           ast::DisabledValidation::kIgnoreConstructibleFunctionParameter);
       ctx.dst->AST().AddFunction(ctx.dst->create<ast::Function>(
           name,
           ast::VariableList{
               // Note: The buffer parameter requires the kStorage StorageClass
               // in order for HLSL to emit this as a ByteAddressBuffer.
               ctx.dst->create<ast::Variable>(
                   ctx.dst->Sym("buffer"), ast::StorageClass::kStorage,
                   ast::Access::kUndefined, type, true, false, nullptr,
                   ast::AttributeList{disable_validation}),
               ctx.dst->Param("result",
                              ctx.dst->ty.pointer(ctx.dst->ty.u32(),
                                                  ast::StorageClass::kFunction)),
           },
           ctx.dst->ty.void_(), nullptr,
           ast::AttributeList{
               ctx.dst->ASTNodes().Create<BufferSizeIntrinsic>(ctx.dst->ID()),
           },
           ast::AttributeList{}));

       return name;
     });
   };

   std::unordered_map<ArrayUsage, Symbol, ArrayUsage::Hasher>
       array_length_by_usage;

   // Find all the arrayLength() calls...
   for (auto* node : ctx.src->ASTNodes().Objects()) {
     if (auto* call_expr = node->As<ast::CallExpression>()) {
       auto* call = sem.Get(call_expr);
       if (auto* builtin = call->Target()->As<sem::Builtin>()) {
         if (builtin->Type() == sem::BuiltinType::kArrayLength) {
           // We're dealing with an arrayLength() call

           // A runtime-sized array can only appear as the store type of a
           // variable, or the last element of a structure (which cannot itself
           // be nested). Given that we require SimplifyPointers, we can assume
           // that the arrayLength() call has one of two forms:
           //   arrayLength(&struct_var.array_member)
           //   arrayLength(&array_var)
           auto* arg = call_expr->args[0];
           auto* address_of = arg->As<ast::UnaryOpExpression>();
           if (!address_of || address_of->op != ast::UnaryOp::kAddressOf) {
             TINT_ICE(Transform, ctx.dst->Diagnostics())
                 << "arrayLength() expected address-of, got "
                 << arg->TypeInfo().name;
           }
           auto* storage_buffer_expr = address_of->expr;
           if (auto* accessor =
                   storage_buffer_expr->As<ast::MemberAccessorExpression>()) {
             storage_buffer_expr = accessor->structure;
           }
           auto* storage_buffer_sem =
               sem.Get<sem::VariableUser>(storage_buffer_expr);
           if (!storage_buffer_sem) {
             TINT_ICE(Transform, ctx.dst->Diagnostics())
                 << "expected form of arrayLength argument to be &array_var or "
                    "&struct_var.array_member";
             break;
           }
           auto* storage_buffer_var = storage_buffer_sem->Variable();
           auto* storage_buffer_type = storage_buffer_sem->Type()->UnwrapRef();

           // Generate BufferSizeIntrinsic for this storage type if we haven't
           // already
           auto buffer_size = get_buffer_size_intrinsic(storage_buffer_type);

           // Find the current statement block
           auto* block = call->Stmt()->Block()->Declaration();

           auto array_length = utils::GetOrCreate(
               array_length_by_usage, {block, storage_buffer_var}, [&] {
                 // First time this array length is used for this block.
                 // Let's calculate it.

                 // Construct the variable that'll hold the result of
                 // RWByteAddressBuffer.GetDimensions()
                 auto* buffer_size_result = ctx.dst->Decl(
                     ctx.dst->Var(ctx.dst->Sym(), ctx.dst->ty.u32(),
                                  ast::StorageClass::kNone, ctx.dst->Expr(0u)));

                 // Call storage_buffer.GetDimensions(&buffer_size_result)
                 auto* call_get_dims = ctx.dst->CallStmt(ctx.dst->Call(
                     // BufferSizeIntrinsic(X, ARGS...) is
                     // translated to:
                     //  X.GetDimensions(ARGS..) by the writer
                     buffer_size, ctx.Clone(storage_buffer_expr),
                     ctx.dst->AddressOf(
                         ctx.dst->Expr(buffer_size_result->variable->symbol))));

                 // Calculate actual array length
                 //                total_storage_buffer_size - array_offset
                 // array_length = ----------------------------------------
                 //                             array_stride
                 auto name = ctx.dst->Sym();
                 const ast::Expression* total_size =
                     ctx.dst->Expr(buffer_size_result->variable);
                 const sem::Array* array_type = nullptr;
                 if (auto* str = storage_buffer_type->As<sem::Struct>()) {
                   // The variable is a struct, so subtract the byte offset of
                   // the array member.
                   auto* array_member_sem = str->Members().back();
                   array_type = array_member_sem->Type()->As<sem::Array>();
                   total_size =
                       ctx.dst->Sub(total_size, array_member_sem->Offset());
                 } else if (auto* arr = storage_buffer_type->As<sem::Array>()) {
                   array_type = arr;
                 } else {
                   TINT_ICE(Transform, ctx.dst->Diagnostics())
                       << "expected form of arrayLength argument to be "
                          "&array_var or &struct_var.array_member";
                   return name;
                 }
                 uint32_t array_stride = array_type->Size();
                 auto* array_length_var = ctx.dst->Decl(
                     ctx.dst->Const(name, ctx.dst->ty.u32(),
                                    ctx.dst->Div(total_size, array_stride)));

                 // Insert the array length calculations at the top of the block
                 ctx.InsertBefore(block->statements, block->statements[0],
                                  buffer_size_result);
                 ctx.InsertBefore(block->statements, block->statements[0],
                                  call_get_dims);
                 ctx.InsertBefore(block->statements, block->statements[0],
                                  array_length_var);
                 return name;
               });

           // Replace the call to arrayLength() with the array length variable
           ctx.Replace(call_expr, ctx.dst->Expr(array_length));
         }
       }
     }
   }

   ctx.Clone();
 }

 }  // namespace transform
 }  // namespace tint
	// 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/calculate_array_length.h"

	#include <unordered_map>
	#include <utility>

	#include "src/tint/ast/call_statement.h"
	#include "src/tint/ast/disable_validation_attribute.h"
	#include "src/tint/program_builder.h"
	#include "src/tint/sem/block_statement.h"
	#include "src/tint/sem/call.h"
	#include "src/tint/sem/function.h"
	#include "src/tint/sem/statement.h"
	#include "src/tint/sem/struct.h"
	#include "src/tint/sem/variable.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::CalculateArrayLength);
	TINT_INSTANTIATE_TYPEINFO(
	tint::transform::CalculateArrayLength::BufferSizeIntrinsic);

	namespace tint {
	namespace transform {

	namespace {

	/// ArrayUsage describes a runtime array usage.
	/// It is used as a key by the array_length_by_usage map.
	struct ArrayUsage {
	ast::BlockStatement const* const block;
	sem::Variable const* const buffer;
	bool operator==(const ArrayUsage& rhs) const {
	return block == rhs.block && buffer == rhs.buffer;
	}
	struct Hasher {
	inline std::size_t operator()(const ArrayUsage& u) const {
	return utils::Hash(u.block, u.buffer);
	}
	};
	};

	} // namespace

	CalculateArrayLength::BufferSizeIntrinsic::BufferSizeIntrinsic(ProgramID pid)
	: Base(pid) {}
	CalculateArrayLength::BufferSizeIntrinsic::~BufferSizeIntrinsic() = default;
	std::string CalculateArrayLength::BufferSizeIntrinsic::InternalName() const {
	return "intrinsic_buffer_size";
	}

	const CalculateArrayLength::BufferSizeIntrinsic*
	CalculateArrayLength::BufferSizeIntrinsic::Clone(CloneContext* ctx) const {
	return ctx->dst->ASTNodes().Create<CalculateArrayLength::BufferSizeIntrinsic>(
	ctx->dst->ID());
	}

	CalculateArrayLength::CalculateArrayLength() = default;
	CalculateArrayLength::~CalculateArrayLength() = default;

	bool CalculateArrayLength::ShouldRun(const Program* program,
	const DataMap&) const {
	for (auto* fn : program->AST().Functions()) {
	if (auto* sem_fn = program->Sem().Get(fn)) {
	for (auto* builtin : sem_fn->DirectlyCalledBuiltins()) {
	if (builtin->Type() == sem::BuiltinType::kArrayLength) {
	return true;
	}
	}
	}
	}
	return false;
	}

	void CalculateArrayLength::Run(CloneContext& ctx,
	const DataMap&,
	DataMap&) const {
	auto& sem = ctx.src->Sem();

	// get_buffer_size_intrinsic() emits the function decorated with
	// BufferSizeIntrinsic that is transformed by the HLSL writer into a call to
	// [RW]ByteAddressBuffer.GetDimensions().
	std::unordered_map<const sem::Type*, Symbol> buffer_size_intrinsics;
	auto get_buffer_size_intrinsic = [&](const sem::Type* buffer_type) {
	return utils::GetOrCreate(buffer_size_intrinsics, buffer_type, [&] {
	auto name = ctx.dst->Sym();
	auto* type = CreateASTTypeFor(ctx, buffer_type);
	auto* disable_validation = ctx.dst->Disable(
	ast::DisabledValidation::kIgnoreConstructibleFunctionParameter);
	ctx.dst->AST().AddFunction(ctx.dst->create<ast::Function>(
	name,
	ast::VariableList{
	// Note: The buffer parameter requires the kStorage StorageClass
	// in order for HLSL to emit this as a ByteAddressBuffer.
	ctx.dst->create<ast::Variable>(
	ctx.dst->Sym("buffer"), ast::StorageClass::kStorage,
	ast::Access::kUndefined, type, true, false, nullptr,
	ast::AttributeList{disable_validation}),
	ctx.dst->Param("result",
	ctx.dst->ty.pointer(ctx.dst->ty.u32(),
	ast::StorageClass::kFunction)),
	},
	ctx.dst->ty.void_(), nullptr,
	ast::AttributeList{
	ctx.dst->ASTNodes().Create<BufferSizeIntrinsic>(ctx.dst->ID()),
	},
	ast::AttributeList{}));

	return name;
	});
	};

	std::unordered_map<ArrayUsage, Symbol, ArrayUsage::Hasher>
	array_length_by_usage;

	// Find all the arrayLength() calls...
	for (auto* node : ctx.src->ASTNodes().Objects()) {
	if (auto* call_expr = node->As<ast::CallExpression>()) {
	auto* call = sem.Get(call_expr);
	if (auto* builtin = call->Target()->As<sem::Builtin>()) {
	if (builtin->Type() == sem::BuiltinType::kArrayLength) {
	// We're dealing with an arrayLength() call

	// A runtime-sized array can only appear as the store type of a
	// variable, or the last element of a structure (which cannot itself
	// be nested). Given that we require SimplifyPointers, we can assume
	// that the arrayLength() call has one of two forms:
	// arrayLength(&struct_var.array_member)
	// arrayLength(&array_var)
	auto* arg = call_expr->args[0];
	auto* address_of = arg->As<ast::UnaryOpExpression>();
	if (!address_of \|\| address_of->op != ast::UnaryOp::kAddressOf) {
	TINT_ICE(Transform, ctx.dst->Diagnostics())
	<< "arrayLength() expected address-of, got "
	<< arg->TypeInfo().name;
	}
	auto* storage_buffer_expr = address_of->expr;
	if (auto* accessor =
	storage_buffer_expr->As<ast::MemberAccessorExpression>()) {
	storage_buffer_expr = accessor->structure;
	}
	auto* storage_buffer_sem =
	sem.Get<sem::VariableUser>(storage_buffer_expr);
	if (!storage_buffer_sem) {
	TINT_ICE(Transform, ctx.dst->Diagnostics())
	<< "expected form of arrayLength argument to be &array_var or "
	"&struct_var.array_member";
	break;
	}
	auto* storage_buffer_var = storage_buffer_sem->Variable();
	auto* storage_buffer_type = storage_buffer_sem->Type()->UnwrapRef();

	// Generate BufferSizeIntrinsic for this storage type if we haven't
	// already
	auto buffer_size = get_buffer_size_intrinsic(storage_buffer_type);

	// Find the current statement block
	auto* block = call->Stmt()->Block()->Declaration();

	auto array_length = utils::GetOrCreate(
	array_length_by_usage, {block, storage_buffer_var}, [&] {
	// First time this array length is used for this block.
	// Let's calculate it.

	// Construct the variable that'll hold the result of
	// RWByteAddressBuffer.GetDimensions()
	auto* buffer_size_result = ctx.dst->Decl(
	ctx.dst->Var(ctx.dst->Sym(), ctx.dst->ty.u32(),
	ast::StorageClass::kNone, ctx.dst->Expr(0u)));

	// Call storage_buffer.GetDimensions(&buffer_size_result)
	auto* call_get_dims = ctx.dst->CallStmt(ctx.dst->Call(
	// BufferSizeIntrinsic(X, ARGS...) is
	// translated to:
	// X.GetDimensions(ARGS..) by the writer
	buffer_size, ctx.Clone(storage_buffer_expr),
	ctx.dst->AddressOf(
	ctx.dst->Expr(buffer_size_result->variable->symbol))));

	// Calculate actual array length
	// total_storage_buffer_size - array_offset
	// array_length = ----------------------------------------
	// array_stride
	auto name = ctx.dst->Sym();
	const ast::Expression* total_size =
	ctx.dst->Expr(buffer_size_result->variable);
	const sem::Array* array_type = nullptr;
	if (auto* str = storage_buffer_type->As<sem::Struct>()) {
	// The variable is a struct, so subtract the byte offset of
	// the array member.
	auto* array_member_sem = str->Members().back();
	array_type = array_member_sem->Type()->As<sem::Array>();
	total_size =
	ctx.dst->Sub(total_size, array_member_sem->Offset());
	} else if (auto* arr = storage_buffer_type->As<sem::Array>()) {
	array_type = arr;
	} else {
	TINT_ICE(Transform, ctx.dst->Diagnostics())
	<< "expected form of arrayLength argument to be "
	"&array_var or &struct_var.array_member";
	return name;
	}
	uint32_t array_stride = array_type->Size();
	auto* array_length_var = ctx.dst->Decl(
	ctx.dst->Const(name, ctx.dst->ty.u32(),
	ctx.dst->Div(total_size, array_stride)));

	// Insert the array length calculations at the top of the block
	ctx.InsertBefore(block->statements, block->statements[0],
	buffer_size_result);
	ctx.InsertBefore(block->statements, block->statements[0],
	call_get_dims);
	ctx.InsertBefore(block->statements, block->statements[0],
	array_length_var);
	return name;
	});

	// Replace the call to arrayLength() with the array length variable
	ctx.Replace(call_expr, ctx.dst->Expr(array_length));
	}
	}
	}
	}

	ctx.Clone();
	}

	} // namespace transform
	} // namespace tint