src/transform/msl.cc - tint - Git at Google

 // Copyright 2020 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/transform/msl.h"

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

 #include "src/ast/disable_validation_decoration.h"
 #include "src/program_builder.h"
 #include "src/sem/call.h"
 #include "src/sem/function.h"
 #include "src/sem/statement.h"
 #include "src/sem/variable.h"
 #include "src/transform/canonicalize_entry_point_io.h"
 #include "src/transform/external_texture_transform.h"
 #include "src/transform/manager.h"
 #include "src/transform/promote_initializers_to_const_var.h"

 namespace tint {
 namespace transform {

 Msl::Msl() = default;
 Msl::~Msl() = default;

 Output Msl::Run(const Program* in, const DataMap&) {
   Manager manager;
   DataMap data;
   manager.Add<CanonicalizeEntryPointIO>();
   manager.Add<ExternalTextureTransform>();
   manager.Add<PromoteInitializersToConstVar>();
   data.Add<CanonicalizeEntryPointIO::Config>(
       CanonicalizeEntryPointIO::BuiltinStyle::kParameter);
   auto out = manager.Run(in, data);
   if (!out.program.IsValid()) {
     return out;
   }

   ProgramBuilder builder;
   CloneContext ctx(&builder, &out.program);
   // TODO(jrprice): Consider making this a standalone transform, with target
   // storage class(es) as transform options.
   HandlePrivateAndWorkgroupVariables(ctx);
   ctx.Clone();
   return Output{Program(std::move(builder))};
 }

 void Msl::HandlePrivateAndWorkgroupVariables(CloneContext& ctx) const {
   // MSL does not allow private and workgroup variables at module-scope, so we
   // push these declarations into the entry point function and then pass them as
   // pointer parameters to any function that references them.
   //
   // Since WGSL does not allow function-scope variables to have these storage
   // classes, we annotate the new variable declarations with an attribute that
   // bypasses that validation rule.
   //
   // Before:
   // ```
   // var<private> v : f32 = 2.0;
   //
   // fn foo() {
   //   v = v + 1.0;
   // }
   //
   // [[stage(compute)]]
   // fn main() {
   //   foo();
   // }
   // ```
   //
   // After:
   // ```
   // fn foo(v : ptr<private, f32>) {
   //   *v = *v + 1.0;
   // }
   //
   // [[stage(compute)]]
   // fn main() {
   //   var<private> v : f32 = 2.0;
   //   let v_ptr : ptr<private, f32> = &f32;
   //   foo(v_ptr);
   // }
   // ```

   // Predetermine the list of function calls that need to be replaced.
   using CallList = std::vector<const ast::CallExpression*>;
   std::unordered_map<const ast::Function*, CallList> calls_to_replace;

   std::vector<ast::Function*> functions_to_process;

   // Build a list of functions that transitively reference any private or
   // workgroup variables.
   for (auto* func_ast : ctx.src->AST().Functions()) {
     auto* func_sem = ctx.src->Sem().Get(func_ast);

     bool needs_processing = false;
     for (auto* var : func_sem->ReferencedModuleVariables()) {
       if (var->StorageClass() == ast::StorageClass::kPrivate ||
           var->StorageClass() == ast::StorageClass::kWorkgroup) {
         needs_processing = true;
         break;
       }
     }

     if (needs_processing) {
       functions_to_process.push_back(func_ast);

       // Find all of the calls to this function that will need to be replaced.
       for (auto* call : func_sem->CallSites()) {
         auto* call_sem = ctx.src->Sem().Get(call);
         calls_to_replace[call_sem->Stmt()->Function()].push_back(call);
       }
     }
   }

   for (auto* func_ast : functions_to_process) {
     auto* func_sem = ctx.src->Sem().Get(func_ast);

     // Map module-scope variables onto their function-scope replacement.
     std::unordered_map<const sem::Variable*, Symbol> var_to_symbol;

     for (auto* var : func_sem->ReferencedModuleVariables()) {
       if (var->StorageClass() != ast::StorageClass::kPrivate &&
           var->StorageClass() != ast::StorageClass::kWorkgroup) {
         continue;
       }

       // This is the symbol for the pointer that replaces the module-scope var.
       auto new_var_symbol = ctx.dst->Sym();

       auto* store_type = CreateASTTypeFor(&ctx, var->Type()->UnwrapRef());

       if (func_ast->IsEntryPoint()) {
         // For an entry point, redeclare the variable at function-scope.
         // Disable storage class validation on this variable.
         auto* disable_validation =
             ctx.dst->ASTNodes().Create<ast::DisableValidationDecoration>(
                 ctx.dst->ID(),
                 ast::DisabledValidation::kFunctionVarStorageClass);
         auto* constructor = ctx.Clone(var->Declaration()->constructor());
         auto* local_var =
             ctx.dst->Var(ctx.dst->Sym(), store_type, var->StorageClass(),
                          constructor, ast::DecorationList{disable_validation});
         ctx.InsertBefore(func_ast->body()->statements(),
                          *func_ast->body()->begin(), ctx.dst->Decl(local_var));

         // Now take the address of the variable.
         auto* ptr = ctx.dst->Const(new_var_symbol, nullptr,
                                    ctx.dst->AddressOf(local_var));
         ctx.InsertBefore(func_ast->body()->statements(),
                          *func_ast->body()->begin(), ctx.dst->Decl(ptr));
       } else {
         // For a regular function, redeclare the variable as a pointer function
         // parameter.
         auto* ptr_type = ctx.dst->ty.pointer(store_type, var->StorageClass());
         ctx.InsertBack(func_ast->params(),
                        ctx.dst->Param(new_var_symbol, ptr_type));
       }

       // Replace all uses of the module-scope variable with the pointer
       // replacement (dereferenced).
       for (auto* user : var->Users()) {
         if (user->Stmt()->Function() == func_ast) {
           ctx.Replace(user->Declaration(), ctx.dst->Deref(new_var_symbol));
         }
       }

       var_to_symbol[var] = new_var_symbol;
     }

     // Pass the pointers through to any functions that need them.
     for (auto* call : calls_to_replace[func_ast]) {
       auto* target = ctx.src->AST().Functions().Find(call->func()->symbol());
       auto* target_sem = ctx.src->Sem().Get(target);

       // Add new arguments for any referenced private and workgroup variables.
       for (auto* target_var : target_sem->ReferencedModuleVariables()) {
         if (target_var->StorageClass() == ast::StorageClass::kPrivate ||
             target_var->StorageClass() == ast::StorageClass::kWorkgroup) {
           ctx.InsertBack(call->params(),
                          ctx.dst->Expr(var_to_symbol[target_var]));
         }
       }
     }
   }

   // Now remove all module-scope private and workgroup variables.
   for (auto* var : ctx.src->AST().GlobalVariables()) {
     if (var->declared_storage_class() == ast::StorageClass::kPrivate ||
         var->declared_storage_class() == ast::StorageClass::kWorkgroup) {
       ctx.Remove(ctx.src->AST().GlobalDeclarations(), var);
     }
   }
 }

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

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

	#include "src/ast/disable_validation_decoration.h"
	#include "src/program_builder.h"
	#include "src/sem/call.h"
	#include "src/sem/function.h"
	#include "src/sem/statement.h"
	#include "src/sem/variable.h"
	#include "src/transform/canonicalize_entry_point_io.h"
	#include "src/transform/external_texture_transform.h"
	#include "src/transform/manager.h"
	#include "src/transform/promote_initializers_to_const_var.h"

	namespace tint {
	namespace transform {

	Msl::Msl() = default;
	Msl::~Msl() = default;

	Output Msl::Run(const Program* in, const DataMap&) {
	Manager manager;
	DataMap data;
	manager.Add<CanonicalizeEntryPointIO>();
	manager.Add<ExternalTextureTransform>();
	manager.Add<PromoteInitializersToConstVar>();
	data.Add<CanonicalizeEntryPointIO::Config>(
	CanonicalizeEntryPointIO::BuiltinStyle::kParameter);
	auto out = manager.Run(in, data);
	if (!out.program.IsValid()) {
	return out;
	}

	ProgramBuilder builder;
	CloneContext ctx(&builder, &out.program);
	// TODO(jrprice): Consider making this a standalone transform, with target
	// storage class(es) as transform options.
	HandlePrivateAndWorkgroupVariables(ctx);
	ctx.Clone();
	return Output{Program(std::move(builder))};
	}

	void Msl::HandlePrivateAndWorkgroupVariables(CloneContext& ctx) const {
	// MSL does not allow private and workgroup variables at module-scope, so we
	// push these declarations into the entry point function and then pass them as
	// pointer parameters to any function that references them.
	//
	// Since WGSL does not allow function-scope variables to have these storage
	// classes, we annotate the new variable declarations with an attribute that
	// bypasses that validation rule.
	//
	// Before:
	// ```
	// var<private> v : f32 = 2.0;
	//
	// fn foo() {
	// v = v + 1.0;
	// }
	//
	// [[stage(compute)]]
	// fn main() {
	// foo();
	// }
	// ```
	//
	// After:
	// ```
	// fn foo(v : ptr<private, f32>) {
	// v = v + 1.0;
	// }
	//
	// [[stage(compute)]]
	// fn main() {
	// var<private> v : f32 = 2.0;
	// let v_ptr : ptr<private, f32> = &f32;
	// foo(v_ptr);
	// }
	// ```

	// Predetermine the list of function calls that need to be replaced.
	using CallList = std::vector<const ast::CallExpression*>;
	std::unordered_map<const ast::Function*, CallList> calls_to_replace;

	std::vector<ast::Function*> functions_to_process;

	// Build a list of functions that transitively reference any private or
	// workgroup variables.
	for (auto* func_ast : ctx.src->AST().Functions()) {
	auto* func_sem = ctx.src->Sem().Get(func_ast);

	bool needs_processing = false;
	for (auto* var : func_sem->ReferencedModuleVariables()) {
	if (var->StorageClass() == ast::StorageClass::kPrivate \|\|
	var->StorageClass() == ast::StorageClass::kWorkgroup) {
	needs_processing = true;
	break;
	}
	}

	if (needs_processing) {
	functions_to_process.push_back(func_ast);

	// Find all of the calls to this function that will need to be replaced.
	for (auto* call : func_sem->CallSites()) {
	auto* call_sem = ctx.src->Sem().Get(call);
	calls_to_replace[call_sem->Stmt()->Function()].push_back(call);
	}
	}
	}

	for (auto* func_ast : functions_to_process) {
	auto* func_sem = ctx.src->Sem().Get(func_ast);

	// Map module-scope variables onto their function-scope replacement.
	std::unordered_map<const sem::Variable*, Symbol> var_to_symbol;

	for (auto* var : func_sem->ReferencedModuleVariables()) {
	if (var->StorageClass() != ast::StorageClass::kPrivate &&
	var->StorageClass() != ast::StorageClass::kWorkgroup) {
	continue;
	}

	// This is the symbol for the pointer that replaces the module-scope var.
	auto new_var_symbol = ctx.dst->Sym();

	auto* store_type = CreateASTTypeFor(&ctx, var->Type()->UnwrapRef());

	if (func_ast->IsEntryPoint()) {
	// For an entry point, redeclare the variable at function-scope.
	// Disable storage class validation on this variable.
	auto* disable_validation =
	ctx.dst->ASTNodes().Create<ast::DisableValidationDecoration>(
	ctx.dst->ID(),
	ast::DisabledValidation::kFunctionVarStorageClass);
	auto* constructor = ctx.Clone(var->Declaration()->constructor());
	auto* local_var =
	ctx.dst->Var(ctx.dst->Sym(), store_type, var->StorageClass(),
	constructor, ast::DecorationList{disable_validation});
	ctx.InsertBefore(func_ast->body()->statements(),
	*func_ast->body()->begin(), ctx.dst->Decl(local_var));

	// Now take the address of the variable.
	auto* ptr = ctx.dst->Const(new_var_symbol, nullptr,
	ctx.dst->AddressOf(local_var));
	ctx.InsertBefore(func_ast->body()->statements(),
	*func_ast->body()->begin(), ctx.dst->Decl(ptr));
	} else {
	// For a regular function, redeclare the variable as a pointer function
	// parameter.
	auto* ptr_type = ctx.dst->ty.pointer(store_type, var->StorageClass());
	ctx.InsertBack(func_ast->params(),
	ctx.dst->Param(new_var_symbol, ptr_type));
	}

	// Replace all uses of the module-scope variable with the pointer
	// replacement (dereferenced).
	for (auto* user : var->Users()) {
	if (user->Stmt()->Function() == func_ast) {
	ctx.Replace(user->Declaration(), ctx.dst->Deref(new_var_symbol));
	}
	}

	var_to_symbol[var] = new_var_symbol;
	}

	// Pass the pointers through to any functions that need them.
	for (auto* call : calls_to_replace[func_ast]) {
	auto* target = ctx.src->AST().Functions().Find(call->func()->symbol());
	auto* target_sem = ctx.src->Sem().Get(target);

	// Add new arguments for any referenced private and workgroup variables.
	for (auto* target_var : target_sem->ReferencedModuleVariables()) {
	if (target_var->StorageClass() == ast::StorageClass::kPrivate \|\|
	target_var->StorageClass() == ast::StorageClass::kWorkgroup) {
	ctx.InsertBack(call->params(),
	ctx.dst->Expr(var_to_symbol[target_var]));
	}
	}
	}
	}

	// Now remove all module-scope private and workgroup variables.
	for (auto* var : ctx.src->AST().GlobalVariables()) {
	if (var->declared_storage_class() == ast::StorageClass::kPrivate \|\|
	var->declared_storage_class() == ast::StorageClass::kWorkgroup) {
	ctx.Remove(ctx.src->AST().GlobalDeclarations(), var);
	}
	}
	}

	} // namespace transform
	} // namespace tint