413 lines · cpp
1//===- mlir-transform-opt.cpp -----------------------------------*- C++ -*-===//2//3// This file is licensed under the Apache License v2.0 with LLVM Exceptions.4// See https://llvm.org/LICENSE.txt for license information.5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception6//7//===----------------------------------------------------------------------===//8 9#include "mlir/Dialect/Transform/IR/TransformDialect.h"10#include "mlir/Dialect/Transform/IR/Utils.h"11#include "mlir/Dialect/Transform/Transforms/TransformInterpreterUtils.h"12#include "mlir/IR/AsmState.h"13#include "mlir/IR/BuiltinOps.h"14#include "mlir/IR/Diagnostics.h"15#include "mlir/IR/DialectRegistry.h"16#include "mlir/IR/MLIRContext.h"17#include "mlir/InitAllDialects.h"18#include "mlir/InitAllExtensions.h"19#include "mlir/InitAllPasses.h"20#include "mlir/Parser/Parser.h"21#include "mlir/Support/FileUtilities.h"22#include "mlir/Tools/mlir-opt/MlirOptMain.h"23#include "llvm/Support/CommandLine.h"24#include "llvm/Support/InitLLVM.h"25#include "llvm/Support/ManagedStatic.h"26#include "llvm/Support/SourceMgr.h"27#include "llvm/Support/ToolOutputFile.h"28#include <cstdlib>29 30namespace {31 32using namespace llvm;33 34/// Structure containing command line options for the tool, these will get35/// initialized when an instance is created.36struct MlirTransformOptCLOptions {37 cl::opt<bool> allowUnregisteredDialects{38 "allow-unregistered-dialect",39 cl::desc("Allow operations coming from an unregistered dialect"),40 cl::init(false)};41 42 cl::opt<mlir::SourceMgrDiagnosticVerifierHandler::Level> verifyDiagnostics{43 "verify-diagnostics", llvm::cl::ValueOptional,44 cl::desc("Check that emitted diagnostics match expected-* lines on the "45 "corresponding line"),46 cl::values(47 clEnumValN(48 mlir::SourceMgrDiagnosticVerifierHandler::Level::All, "all",49 "Check all diagnostics (expected, unexpected, near-misses)"),50 // Implicit value: when passed with no arguments, e.g.51 // `--verify-diagnostics` or `--verify-diagnostics=`.52 clEnumValN(53 mlir::SourceMgrDiagnosticVerifierHandler::Level::All, "",54 "Check all diagnostics (expected, unexpected, near-misses)"),55 clEnumValN(56 mlir::SourceMgrDiagnosticVerifierHandler::Level::OnlyExpected,57 "only-expected", "Check only expected diagnostics"))};58 59 cl::opt<std::string> payloadFilename{cl::Positional, cl::desc("<input file>"),60 cl::init("-")};61 62 cl::opt<std::string> outputFilename{"o", cl::desc("Output filename"),63 cl::value_desc("filename"),64 cl::init("-")};65 66 cl::opt<std::string> transformMainFilename{67 "transform",68 cl::desc("File containing entry point of the transform script, if "69 "different from the input file"),70 cl::value_desc("filename"), cl::init("")};71 72 cl::list<std::string> transformLibraryFilenames{73 "transform-library", cl::desc("File(s) containing definitions of "74 "additional transform script symbols")};75 76 cl::opt<std::string> transformEntryPoint{77 "transform-entry-point",78 cl::desc("Name of the entry point transform symbol"),79 cl::init(mlir::transform::TransformDialect::kTransformEntryPointSymbolName80 .str())};81 82 cl::opt<bool> disableExpensiveChecks{83 "disable-expensive-checks",84 cl::desc("Disables potentially expensive checks in the transform "85 "interpreter, providing more speed at the expense of "86 "potential memory problems and silent corruptions"),87 cl::init(false)};88 89 cl::opt<bool> dumpLibraryModule{90 "dump-library-module",91 cl::desc("Prints the combined library module before the output"),92 cl::init(false)};93};94} // namespace95 96/// "Managed" static instance of the command-line options structure. This makes97/// them locally-scoped and explicitly initialized/deinitialized. While this is98/// not strictly necessary in the tool source file that is not being used as a99/// library (where the options would pollute the global list of options), it is100/// good practice to follow this.101static llvm::ManagedStatic<MlirTransformOptCLOptions> clOptions;102 103/// Explicitly registers command-line options.104static void registerCLOptions() { *clOptions; }105 106namespace {107/// A wrapper class for source managers diagnostic. This provides both unique108/// ownership and virtual function-like overload for a pair of109/// inheritance-related classes that do not use virtual functions.110class DiagnosticHandlerWrapper {111public:112 /// Kind of the diagnostic handler to use.113 enum class Kind { EmitDiagnostics, VerifyDiagnostics };114 115 /// Constructs the diagnostic handler of the specified kind of the given116 /// source manager and context.117 DiagnosticHandlerWrapper(118 Kind kind, llvm::SourceMgr &mgr, mlir::MLIRContext *context,119 std::optional<mlir::SourceMgrDiagnosticVerifierHandler::Level> level =120 {}) {121 if (kind == Kind::EmitDiagnostics) {122 handler = new mlir::SourceMgrDiagnosticHandler(mgr, context);123 } else {124 assert(level.has_value() && "expected level");125 handler =126 new mlir::SourceMgrDiagnosticVerifierHandler(mgr, context, *level);127 }128 }129 130 /// This object is non-copyable but movable.131 DiagnosticHandlerWrapper(const DiagnosticHandlerWrapper &) = delete;132 DiagnosticHandlerWrapper(DiagnosticHandlerWrapper &&other) = default;133 DiagnosticHandlerWrapper &134 operator=(const DiagnosticHandlerWrapper &) = delete;135 DiagnosticHandlerWrapper &operator=(DiagnosticHandlerWrapper &&) = default;136 137 /// Verifies the captured "expected-*" diagnostics if required.138 llvm::LogicalResult verify() const {139 if (auto *ptr =140 dyn_cast<mlir::SourceMgrDiagnosticVerifierHandler *>(handler)) {141 return ptr->verify();142 }143 return mlir::success();144 }145 146 /// Destructs the object of the same type as allocated.147 ~DiagnosticHandlerWrapper() {148 if (auto *ptr = dyn_cast<mlir::SourceMgrDiagnosticHandler *>(handler)) {149 delete ptr;150 } else {151 delete cast<mlir::SourceMgrDiagnosticVerifierHandler *>(handler);152 }153 }154 155private:156 /// Internal storage is a type-safe union.157 llvm::PointerUnion<mlir::SourceMgrDiagnosticHandler *,158 mlir::SourceMgrDiagnosticVerifierHandler *>159 handler;160};161 162/// MLIR has deeply rooted expectations that the LLVM source manager contains163/// exactly one buffer, until at least the lexer level. This class wraps164/// multiple LLVM source managers each managing a buffer to match MLIR's165/// expectations while still providing a centralized handling mechanism.166class TransformSourceMgr {167public:168 /// Constructs the source manager indicating whether diagnostic messages will169 /// be verified later on.170 explicit TransformSourceMgr(171 std::optional<mlir::SourceMgrDiagnosticVerifierHandler::Level>172 verifyDiagnostics)173 : verifyDiagnostics(verifyDiagnostics) {}174 175 /// Deconstructs the source manager. Note that `checkResults` must have been176 /// called on this instance before deconstructing it.177 ~TransformSourceMgr() {178 assert(resultChecked && "must check the result of diagnostic handlers by "179 "running TransformSourceMgr::checkResult");180 }181 182 /// Parses the given buffer and creates the top-level operation of the kind183 /// specified as template argument in the given context. Additional parsing184 /// options may be provided.185 template <typename OpTy = mlir::Operation *>186 mlir::OwningOpRef<OpTy> parseBuffer(std::unique_ptr<MemoryBuffer> buffer,187 mlir::MLIRContext &context,188 const mlir::ParserConfig &config) {189 // Create a single-buffer LLVM source manager. Note that `unique_ptr` allows190 // the code below to capture a reference to the source manager in such a way191 // that it is not invalidated when the vector contents is eventually192 // reallocated.193 llvm::SourceMgr &mgr =194 *sourceMgrs.emplace_back(std::make_unique<llvm::SourceMgr>());195 mgr.AddNewSourceBuffer(std::move(buffer), llvm::SMLoc());196 197 // Choose the type of diagnostic handler depending on whether diagnostic198 // verification needs to happen and store it.199 if (verifyDiagnostics) {200 diagHandlers.emplace_back(201 DiagnosticHandlerWrapper::Kind::VerifyDiagnostics, mgr, &context,202 verifyDiagnostics);203 } else {204 diagHandlers.emplace_back(DiagnosticHandlerWrapper::Kind::EmitDiagnostics,205 mgr, &context);206 }207 208 // Defer to MLIR's parser.209 return mlir::parseSourceFile<OpTy>(mgr, config);210 }211 212 /// If diagnostic message verification has been requested upon construction of213 /// this source manager, performs the verification, reports errors and returns214 /// the result of the verification. Otherwise passes through the given value.215 llvm::LogicalResult checkResult(llvm::LogicalResult result) {216 resultChecked = true;217 if (!verifyDiagnostics)218 return result;219 220 return mlir::failure(llvm::any_of(diagHandlers, [](const auto &handler) {221 return mlir::failed(handler.verify());222 }));223 }224 225private:226 /// Indicates whether diagnostic message verification is requested.227 const std::optional<mlir::SourceMgrDiagnosticVerifierHandler::Level>228 verifyDiagnostics;229 230 /// Indicates that diagnostic message verification has taken place, and the231 /// deconstruction is therefore safe.232 bool resultChecked = false;233 234 /// Storage for per-buffer source managers and diagnostic handlers. These are235 /// wrapped into unique pointers in order to make it safe to capture236 /// references to these objects: if the vector is reallocated, the unique237 /// pointer objects are moved by the pointer addresses won't change. Also, for238 /// handlers, this allows to store the pointer to the base class.239 SmallVector<std::unique_ptr<llvm::SourceMgr>> sourceMgrs;240 SmallVector<DiagnosticHandlerWrapper> diagHandlers;241};242} // namespace243 244/// Trivial wrapper around `applyTransforms` that doesn't support extra mapping245/// and doesn't enforce the entry point transform ops being top-level.246static llvm::LogicalResult247applyTransforms(mlir::Operation *payloadRoot,248 mlir::transform::TransformOpInterface transformRoot,249 const mlir::transform::TransformOptions &options) {250 return applyTransforms(payloadRoot, transformRoot, {}, options,251 /*enforceToplevelTransformOp=*/false);252}253 254/// Applies transforms indicated in the transform dialect script to the input255/// buffer. The transform script may be embedded in the input buffer or as a256/// separate buffer. The transform script may have external symbols, the257/// definitions of which must be provided in transform library buffers. If the258/// application is successful, prints the transformed input buffer into the259/// given output stream. Additional configuration options are derived from260/// command-line options.261static llvm::LogicalResult processPayloadBuffer(262 raw_ostream &os, std::unique_ptr<MemoryBuffer> inputBuffer,263 std::unique_ptr<llvm::MemoryBuffer> transformBuffer,264 MutableArrayRef<std::unique_ptr<MemoryBuffer>> transformLibraries,265 mlir::DialectRegistry ®istry) {266 267 // Initialize the MLIR context, and various configurations.268 mlir::MLIRContext context(registry, mlir::MLIRContext::Threading::DISABLED);269 context.allowUnregisteredDialects(clOptions->allowUnregisteredDialects);270 mlir::ParserConfig config(&context);271 TransformSourceMgr sourceMgr(272 /*verifyDiagnostics=*/clOptions->verifyDiagnostics.getNumOccurrences()273 ? std::optional{clOptions->verifyDiagnostics.getValue()}274 : std::nullopt);275 276 // Parse the input buffer that will be used as transform payload.277 mlir::OwningOpRef<mlir::Operation *> payloadRoot =278 sourceMgr.parseBuffer(std::move(inputBuffer), context, config);279 if (!payloadRoot)280 return sourceMgr.checkResult(mlir::failure());281 282 // Identify the module containing the transform script entry point. This may283 // be the same module as the input or a separate module. In the former case,284 // make a copy of the module so it can be modified freely. Modification may285 // happen in the script itself (at which point it could be rewriting itself286 // during interpretation, leading to tricky memory errors) or by embedding287 // library modules in the script.288 mlir::OwningOpRef<mlir::ModuleOp> transformRoot;289 if (transformBuffer) {290 transformRoot = sourceMgr.parseBuffer<mlir::ModuleOp>(291 std::move(transformBuffer), context, config);292 if (!transformRoot)293 return sourceMgr.checkResult(mlir::failure());294 } else {295 transformRoot = cast<mlir::ModuleOp>(payloadRoot->clone());296 }297 298 // Parse and merge the libraries into the main transform module.299 for (auto &&transformLibrary : transformLibraries) {300 mlir::OwningOpRef<mlir::ModuleOp> libraryModule =301 sourceMgr.parseBuffer<mlir::ModuleOp>(std::move(transformLibrary),302 context, config);303 304 if (!libraryModule ||305 mlir::failed(mlir::transform::detail::mergeSymbolsInto(306 *transformRoot, std::move(libraryModule))))307 return sourceMgr.checkResult(mlir::failure());308 }309 310 // If requested, dump the combined transform module.311 if (clOptions->dumpLibraryModule)312 transformRoot->dump();313 314 // Find the entry point symbol. Even if it had originally been in the payload315 // module, it was cloned into the transform module so only look there.316 mlir::transform::TransformOpInterface entryPoint =317 mlir::transform::detail::findTransformEntryPoint(318 *transformRoot, mlir::ModuleOp(), clOptions->transformEntryPoint);319 if (!entryPoint)320 return sourceMgr.checkResult(mlir::failure());321 322 // Apply the requested transformations.323 mlir::transform::TransformOptions transformOptions;324 transformOptions.enableExpensiveChecks(!clOptions->disableExpensiveChecks);325 if (mlir::failed(applyTransforms(*payloadRoot, entryPoint, transformOptions)))326 return sourceMgr.checkResult(mlir::failure());327 328 // Print the transformed result and check the captured diagnostics if329 // requested.330 payloadRoot->print(os);331 return sourceMgr.checkResult(mlir::success());332}333 334/// Tool entry point.335static llvm::LogicalResult runMain(int argc, char **argv) {336 // Register all upstream dialects and extensions. Specific uses are advised337 // not to register all dialects indiscriminately but rather hand-pick what is338 // necessary for their use case.339 mlir::DialectRegistry registry;340 mlir::registerAllDialects(registry);341 mlir::registerAllExtensions(registry);342 mlir::registerAllPasses();343 344 // Explicitly register the transform dialect. This is not strictly necessary345 // since it has been already registered as part of the upstream dialect list,346 // but useful for example purposes for cases when dialects to register are347 // hand-picked. The transform dialect must be registered.348 registry.insert<mlir::transform::TransformDialect>();349 350 // Register various command-line options. Note that the LLVM initializer351 // object is a RAII that ensures correct deconstruction of command-line option352 // objects inside ManagedStatic.353 llvm::InitLLVM y(argc, argv);354 mlir::registerAsmPrinterCLOptions();355 mlir::registerMLIRContextCLOptions();356 registerCLOptions();357 llvm::cl::ParseCommandLineOptions(argc, argv,358 "Minimal Transform dialect driver\n");359 360 // Try opening the main input file.361 std::string errorMessage;362 std::unique_ptr<llvm::MemoryBuffer> payloadFile =363 mlir::openInputFile(clOptions->payloadFilename, &errorMessage);364 if (!payloadFile) {365 llvm::errs() << errorMessage << "\n";366 return mlir::failure();367 }368 369 // Try opening the output file.370 std::unique_ptr<llvm::ToolOutputFile> outputFile =371 mlir::openOutputFile(clOptions->outputFilename, &errorMessage);372 if (!outputFile) {373 llvm::errs() << errorMessage << "\n";374 return mlir::failure();375 }376 377 // Try opening the main transform file if provided.378 std::unique_ptr<llvm::MemoryBuffer> transformRootFile;379 if (!clOptions->transformMainFilename.empty()) {380 if (clOptions->transformMainFilename == clOptions->payloadFilename) {381 llvm::errs() << "warning: " << clOptions->payloadFilename382 << " is provided as both payload and transform file\n";383 } else {384 transformRootFile =385 mlir::openInputFile(clOptions->transformMainFilename, &errorMessage);386 if (!transformRootFile) {387 llvm::errs() << errorMessage << "\n";388 return mlir::failure();389 }390 }391 }392 393 // Try opening transform library files if provided.394 SmallVector<std::unique_ptr<llvm::MemoryBuffer>> transformLibraries;395 transformLibraries.reserve(clOptions->transformLibraryFilenames.size());396 for (llvm::StringRef filename : clOptions->transformLibraryFilenames) {397 transformLibraries.emplace_back(398 mlir::openInputFile(filename, &errorMessage));399 if (!transformLibraries.back()) {400 llvm::errs() << errorMessage << "\n";401 return mlir::failure();402 }403 }404 405 return processPayloadBuffer(outputFile->os(), std::move(payloadFile),406 std::move(transformRootFile), transformLibraries,407 registry);408}409 410int main(int argc, char **argv) {411 return mlir::asMainReturnCode(runMain(argc, argv));412}413