blob: bd4af0ed6afa16a7ddd543e5613b17b0bc1326f7 [file] [log] [blame]
// 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