| // 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/transform/add_spirv_block_decoration.h" |
| |
| #include <unordered_map> |
| #include <unordered_set> |
| #include <utility> |
| |
| #include "src/program_builder.h" |
| #include "src/sem/variable.h" |
| #include "src/utils/map.h" |
| |
| TINT_INSTANTIATE_TYPEINFO(tint::transform::AddSpirvBlockDecoration); |
| TINT_INSTANTIATE_TYPEINFO( |
| tint::transform::AddSpirvBlockDecoration::SpirvBlockDecoration); |
| |
| namespace tint { |
| namespace transform { |
| |
| AddSpirvBlockDecoration::AddSpirvBlockDecoration() = default; |
| |
| AddSpirvBlockDecoration::~AddSpirvBlockDecoration() = default; |
| |
| void AddSpirvBlockDecoration::Run(CloneContext& ctx, const DataMap&, DataMap&) { |
| auto& sem = ctx.src->Sem(); |
| |
| // Collect the set of structs that are nested in other types. |
| std::unordered_set<const sem::Struct*> nested_structs; |
| for (auto* node : ctx.src->ASTNodes().Objects()) { |
| if (auto* arr = sem.Get<sem::Array>(node->As<ast::Array>())) { |
| if (auto* nested_str = arr->ElemType()->As<sem::Struct>()) { |
| nested_structs.insert(nested_str); |
| } |
| } else if (auto* str = sem.Get<sem::Struct>(node->As<ast::Struct>())) { |
| for (auto* member : str->Members()) { |
| if (auto* nested_str = member->Type()->As<sem::Struct>()) { |
| nested_structs.insert(nested_str); |
| } |
| } |
| } |
| } |
| |
| // A map from a struct in the source program to a block-decorated wrapper that |
| // contains it in the destination program. |
| std::unordered_map<const sem::Struct*, const ast::Struct*> wrapper_structs; |
| |
| // Process global variables that are buffers. |
| for (auto* var : ctx.src->AST().GlobalVariables()) { |
| auto* sem_var = sem.Get<sem::GlobalVariable>(var); |
| if (var->declared_storage_class != ast::StorageClass::kStorage && |
| var->declared_storage_class != ast::StorageClass::kUniform) { |
| continue; |
| } |
| |
| auto* str = sem.Get<sem::Struct>(var->type); |
| if (!str) { |
| // TODO(jrprice): We'll need to wrap these too, when WGSL supports this. |
| TINT_ICE(Transform, ctx.dst->Diagnostics()) |
| << "non-struct buffer types are not yet supported"; |
| continue; |
| } |
| |
| if (nested_structs.count(str)) { |
| const char* kInnerStructMemberName = "inner"; |
| |
| // This struct is nested somewhere else, so we need to wrap it first. |
| auto* wrapper = utils::GetOrCreate(wrapper_structs, str, [&]() { |
| auto* block = |
| ctx.dst->ASTNodes().Create<SpirvBlockDecoration>(ctx.dst->ID()); |
| auto wrapper_name = |
| ctx.src->Symbols().NameFor(str->Declaration()->name) + "_block"; |
| auto* ret = ctx.dst->create<ast::Struct>( |
| ctx.dst->Symbols().New(wrapper_name), |
| ast::StructMemberList{ctx.dst->Member(kInnerStructMemberName, |
| CreateASTTypeFor(ctx, str))}, |
| ast::DecorationList{block}); |
| ctx.InsertAfter(ctx.src->AST().GlobalDeclarations(), str->Declaration(), |
| ret); |
| return ret; |
| }); |
| ctx.Replace(var->type, ctx.dst->ty.Of(wrapper)); |
| |
| // Insert a member accessor to get the original struct from the wrapper at |
| // any usage of the original variable. |
| for (auto* user : sem_var->Users()) { |
| ctx.Replace(user->Declaration(), |
| ctx.dst->MemberAccessor(ctx.Clone(var->symbol), |
| kInnerStructMemberName)); |
| } |
| } else { |
| // Add a block decoration to this struct directly. |
| auto* block = |
| ctx.dst->ASTNodes().Create<SpirvBlockDecoration>(ctx.dst->ID()); |
| ctx.InsertFront(str->Declaration()->decorations, block); |
| } |
| } |
| |
| ctx.Clone(); |
| } |
| |
| AddSpirvBlockDecoration::SpirvBlockDecoration::SpirvBlockDecoration( |
| ProgramID pid) |
| : Base(pid) {} |
| AddSpirvBlockDecoration::SpirvBlockDecoration::~SpirvBlockDecoration() = |
| default; |
| std::string AddSpirvBlockDecoration::SpirvBlockDecoration::InternalName() |
| const { |
| return "spirv_block"; |
| } |
| |
| const AddSpirvBlockDecoration::SpirvBlockDecoration* |
| AddSpirvBlockDecoration::SpirvBlockDecoration::Clone(CloneContext* ctx) const { |
| return ctx->dst->ASTNodes() |
| .Create<AddSpirvBlockDecoration::SpirvBlockDecoration>(ctx->dst->ID()); |
| } |
| |
| } // namespace transform |
| } // namespace tint |