| // 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/tint/transform/manager.h" |
| |
| #include "src/tint/lang/wgsl/ast/transform/transform.h" |
| #include "src/tint/lang/wgsl/program/program_builder.h" |
| |
| #if TINT_BUILD_IR |
| #include "src/tint/lang/core/ir/from_program.h" |
| #include "src/tint/lang/core/ir/to_program.h" |
| #include "src/tint/lang/core/ir/transform/transform.h" |
| #else |
| // Declare an ir::Module class so that the transform target variant compiles. |
| namespace ir { |
| class Module; |
| } |
| #endif // TINT_BUILD_IR |
| |
| /// If set to 1 then the transform::Manager will dump the WGSL of the program |
| /// before and after each transform. Helpful for debugging bad output. |
| #define TINT_PRINT_PROGRAM_FOR_EACH_TRANSFORM 0 |
| |
| #if TINT_PRINT_PROGRAM_FOR_EACH_TRANSFORM |
| #include <iostream> |
| #include "src/tint/lang/core/ir/disassembler.h" |
| #define TINT_IF_PRINT_PROGRAM(x) x |
| #else // TINT_PRINT_PROGRAM_FOR_EACH_TRANSFORM |
| #define TINT_IF_PRINT_PROGRAM(x) |
| #endif // TINT_PRINT_PROGRAM_FOR_EACH_TRANSFORM |
| |
| namespace tint::transform { |
| |
| Manager::Manager() = default; |
| Manager::~Manager() = default; |
| |
| template <typename OUTPUT, typename INPUT> |
| OUTPUT Manager::RunTransforms(INPUT in, |
| const transform::DataMap& inputs, |
| transform::DataMap& outputs) const { |
| static_assert(std::is_same<INPUT, const Program*>() || std::is_same<INPUT, ir::Module*>()); |
| static_assert(std::is_same<OUTPUT, Program>() || std::is_same<OUTPUT, ir::Module*>()); |
| |
| // The current transform target, which could be either AST or IR. |
| std::variant<const Program*, ir::Module*> target = in; |
| // A local AST program to hold the result of AST transforms. |
| Program ast_result; |
| #if TINT_BUILD_IR |
| // A local IR module to hold the result of AST->IR conversions. |
| ir::Module ir_result; |
| #endif |
| |
| #if TINT_PRINT_PROGRAM_FOR_EACH_TRANSFORM |
| auto print_program = [&](const char* msg, const char* name) { |
| std::cout << "=========================================================" << std::endl; |
| std::cout << "== " << msg << " " << name << ":" << std::endl; |
| std::cout << "=========================================================" << std::endl; |
| if (std::holds_alternative<const Program*>(target)) { |
| auto* program = std::get<const Program*>(target); |
| auto wgsl = Program::printer(program); |
| std::cout << wgsl << std::endl; |
| if (!program->IsValid()) { |
| std::cout << "-- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- --" |
| << std::endl; |
| std::cout << program->Diagnostics().str() << std::endl; |
| } |
| } else if (std::holds_alternative<ir::Module*>(target)) { |
| #if TINT_BUILD_IR |
| ir::Disassembler dis(*std::get<ir::Module*>(target)); |
| std::cout << dis.Disassemble(); |
| #endif // TINT_BUILD_IR |
| } |
| std::cout << "=========================================================" << std::endl |
| << std::endl; |
| }; |
| #endif |
| |
| // Helper functions to get the current program state as either an AST program or IR module, |
| // performing a conversion if necessary. |
| auto get_ast = [&] { |
| #if TINT_BUILD_IR |
| if (std::holds_alternative<ir::Module*>(target)) { |
| // Convert the IR module to an AST program. |
| ast_result = ir::ToProgram(*std::get<ir::Module*>(target)); |
| target = &ast_result; |
| } |
| #endif // TINT_BUILD_IR |
| TINT_ASSERT(Transform, std::holds_alternative<const Program*>(target)); |
| return std::get<const Program*>(target); |
| }; |
| #if TINT_BUILD_IR |
| auto get_ir = [&] { |
| if (std::holds_alternative<const Program*>(target)) { |
| // Convert the AST program to an IR module. |
| auto converted = ir::FromProgram(std::get<const Program*>(target)); |
| TINT_ASSERT(Transform, converted); |
| ir_result = converted.Move(); |
| target = &ir_result; |
| } |
| TINT_ASSERT(Transform, std::holds_alternative<ir::Module*>(target)); |
| return std::get<ir::Module*>(target); |
| }; |
| #endif // TINT_BUILD_IR |
| |
| TINT_IF_PRINT_PROGRAM(print_program("Input of", "transform manager")); |
| |
| for (const auto& transform : transforms_) { |
| if (auto* ast_transform = transform->As<ast::transform::Transform>()) { |
| if (auto result = ast_transform->Apply(get_ast(), inputs, outputs)) { |
| ast_result = std::move(result.value()); |
| target = &ast_result; |
| |
| if (!ast_result.IsValid()) { |
| TINT_IF_PRINT_PROGRAM( |
| print_program("Invalid output of", transform->TypeInfo().name)); |
| break; |
| } |
| |
| TINT_IF_PRINT_PROGRAM(print_program("Output of", transform->TypeInfo().name)); |
| } else { |
| TINT_IF_PRINT_PROGRAM(std::cout << "Skipped " << transform->TypeInfo().name |
| << std::endl); |
| } |
| #if TINT_BUILD_IR |
| } else if (auto* ir_transform = transform->As<ir::transform::Transform>()) { |
| ir_transform->Run(get_ir(), inputs, outputs); |
| TINT_IF_PRINT_PROGRAM(print_program("Output of", transform->TypeInfo().name)); |
| #endif // TINT_BUILD_IR |
| } else { |
| TINT_ASSERT(Transform, false && "unhandled transform type"); |
| } |
| } |
| |
| TINT_IF_PRINT_PROGRAM(print_program("Final output of", "transform manager")); |
| |
| if constexpr (std::is_same<OUTPUT, Program>()) { |
| auto* result = get_ast(); |
| if (result == in) { |
| // AST transform pipelines are expected to return a clone of the program, so make sure |
| // the input is cloned at least once even if nothing changed. |
| ProgramBuilder b; |
| CloneContext ctx{&b, result, /* auto_clone_symbols */ true}; |
| ctx.Clone(); |
| ast_result = Program(std::move(b)); |
| } |
| return ast_result; |
| #if TINT_BUILD_IR |
| } else if constexpr (std::is_same<OUTPUT, ir::Module*>()) { |
| auto* result = get_ir(); |
| if (result == &ir_result) { |
| // IR transform pipelines are expected to mutate the module in place, so move the local |
| // temporary result to the original input. |
| *in = std::move(ir_result); |
| } |
| return in; |
| #endif // TINT_BUILD_IR |
| } |
| } |
| |
| Program Manager::Run(const Program* program, |
| const transform::DataMap& inputs, |
| transform::DataMap& outputs) const { |
| return RunTransforms<Program>(program, inputs, outputs); |
| } |
| |
| #if TINT_BUILD_IR |
| void Manager::Run(ir::Module* module, const DataMap& inputs, DataMap& outputs) const { |
| auto* output = RunTransforms<ir::Module*>(module, inputs, outputs); |
| TINT_ASSERT(Transform, output == module); |
| } |
| #endif // TINT_BUILD_IR |
| |
| } // namespace tint::transform |