[mlir][transform] generate transform module on-the-fly

[ftynse]

Summary

Add a TransformInterpreterPassBase capability to generate the (shared)
module containing the transform script during the pass initialization.
This is helpful to programmatically generate the script as opposed to
parsing it from the textual module.

mlir/include/mlir/Dialect/Transform/Transforms/TransformInterpreterPassBase.h

 
 namespace transform {
 namespace detail {
 /// Template-free implementation of TransformInterpreterPassBase::initialize.
 LogicalResult interpreterBaseInitializeImpl(
     MLIRContext *context, StringRef transformFileName,
     StringRef transformLibraryFileName,
     std::shared_ptr<OwningOpRef<ModuleOp>> &module,
-    std::shared_ptr<OwningOpRef<ModuleOp>> &libraryModule);
+    std::shared_ptr<OwningOpRef<ModuleOp>> &libraryModule,
+    function_ref<std::optional<LogicalResult>(OpBuilder &, Location)>
+        moduleBuilder = nullptr);
 
 /// Template-free implementation of
 /// TransformInterpreterPassBase::runOnOperation.
 LogicalResult interpreterBaseRunOnOperationImpl(
     Operation *target, StringRef passName,
     const std::shared_ptr<OwningOpRef<ModuleOp>> &sharedTransformModule,
     const std::shared_ptr<OwningOpRef<ModuleOp>> &libraryModule,
     const RaggedArray<MappedValue> &extraMappings,

 #undef REQUIRE_PASS_OPTION
 
     StringRef transformFileName =
         static_cast<Concrete *>(this)->transformFileName;
     StringRef transformLibraryFileName =
         static_cast<Concrete *>(this)->transformLibraryFileName;
     return detail::interpreterBaseInitializeImpl(
         context, transformFileName, transformLibraryFileName,
-        sharedTransformModule, transformLibraryModule);
+        sharedTransformModule, transformLibraryModule,
+        [this](OpBuilder &builder, Location loc) {
+          return static_cast<Concrete *>(this)->constructTransformModule(
+              builder, loc);
+        });
   }
 
   /// Hook for passes to run additional logic in the pass before the
   /// interpreter. If failure is returned, the pass fails and the interpreter is
   /// not run.
   LogicalResult runBeforeInterpreter(Operation *) { return success(); }
 
   /// Hook for passes to run additional logic in the pass after the interpreter.
   /// Only runs if everything succeeded before. If failure is returned, the pass
   /// fails.
   LogicalResult runAfterInterpreter(Operation *) { return success(); }
 
+  /// Hook for passes to run custom logic to construct the transform module.
+  /// This will run during initialization. If the external script is provided,
+  /// it overrides the construction, which will not be called.
+  std::optional<LogicalResult> constructTransformModule(OpBuilder &builder,
+                                                        Location loc) {
+    return std::nullopt;
+  }
+
   void runOnOperation() override {
     auto *pass = static_cast<Concrete *>(this);
     Operation *op = pass->getOperation();
     StringRef binaryName = Concrete::getBinaryName();
     if (failed(pass->runBeforeInterpreter(op)) ||
         failed(detail::interpreterBaseRunOnOperationImpl(
             op, pass->getArgument(), sharedTransformModule,
             transformLibraryModule,

mlir/lib/Dialect/Transform/Transforms/TransformInterpreterPassBase.cpp

LogicalResult transform::detail::interpreterBaseRunOnOperationImpl(
   return applyTransforms(payloadRoot, cast<TransformOpInterface>(transformRoot),
                          extraMappings, options);
 }
 
 LogicalResult transform::detail::interpreterBaseInitializeImpl(
     MLIRContext *context, StringRef transformFileName,
     StringRef transformLibraryFileName,
     std::shared_ptr<OwningOpRef<ModuleOp>> &module,
-    std::shared_ptr<OwningOpRef<ModuleOp>> &libraryModule) {
+    std::shared_ptr<OwningOpRef<ModuleOp>> &libraryModule,
+    function_ref<std::optional<LogicalResult>(OpBuilder &, Location)>
+        moduleBuilder) {
   OwningOpRef<ModuleOp> parsed;
   if (failed(parseTransformModuleFromFile(context, transformFileName, parsed)))
     return failure();
   if (parsed && failed(mlir::verify(*parsed)))
     return failure();
 
   OwningOpRef<ModuleOp> parsedLibrary;
   if (failed(parseTransformModuleFromFile(context, transformLibraryFileName,
                                           parsedLibrary)))
     return failure();
   if (parsedLibrary && failed(mlir::verify(*parsedLibrary)))
     return failure();
 
+  if (parsed) {
     module = std::make_shared<OwningOpRef<ModuleOp>>(std::move(parsed));
+  } else if (moduleBuilder) {
+    // TODO: better location story.
+    auto location = UnknownLoc::get(context);
+    auto localModule = std::make_shared<OwningOpRef<ModuleOp>>(
+        ModuleOp::create(location, "__transform"));
+
+    OpBuilder b(context);
+    b.setInsertionPointToEnd(localModule->get().getBody());
+    if (std::optional<LogicalResult> result = moduleBuilder(b, location)) {
+      if (failed(*result))
+        return failure();
+      module = std::move(localModule);
+    }
+  }
+
   if (!parsedLibrary || !*parsedLibrary)
     return success();
 
   if (module && *module) {
     if (failed(defineDeclaredSymbols(*module->get().getBody(),
                                      parsedLibrary.get())))
       return failure();
   } else {
     libraryModule =
         std::make_shared<OwningOpRef<ModuleOp>>(std::move(parsedLibrary));
   }
   return success();
 }

mlir/test/Dialect/Transform/test-interpreter-module-generation.mlir

+// RUN: mlir-opt %s --test-transform-dialect-interpreter=test-module-generation=1 --verify-diagnostics
+
+// expected-remark @below {{remark from generated}}
+module {}

mlir/test/lib/Dialect/Transform/TestTransformDialectInterpreter.cpp

 //===- TestTransformDialectInterpreter.cpp --------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file defines a test pass that interprets Transform dialect operations in
 // the module.
 //
 //===----------------------------------------------------------------------===//
 
+#include "TestTransformDialectExtension.h"
 #include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
+#include "mlir/Dialect/Transform/IR/TransformOps.h"
 #include "mlir/Dialect/Transform/Transforms/TransformInterpreterPassBase.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/Pass/Pass.h"
 
 using namespace mlir;
 
 namespace {
public:
   StringRef getArgument() const override {
     return "test-transform-dialect-interpreter";
   }
 
   StringRef getDescription() const override {
     return "apply transform dialect operations one by one";
   }
 
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<transform::TransformDialect>();
+  }
+
   void findOperationsByName(Operation *root, StringRef name,
                             SmallVectorImpl<Operation *> &operations) {
     root->walk([&](Operation *op) {
       if (op->getName().getStringRef() == name) {
         operations.push_back(op);
       }
     });
   }
unsigned numberOfSetOptions(const Option<std::string> &ops,
                               const Option<std::string> &values) {
     unsigned numSetValues = 0;
     numSetValues += !ops.empty();
     numSetValues += !params.empty();
     numSetValues += !values.empty();
     return numSetValues;
   }
 
+  std::optional<LogicalResult> constructTransformModule(OpBuilder &builder,
+                                                        Location loc) {
+    if (!testModuleGeneration)
+      return std::nullopt;
+
+    builder.create<transform::SequenceOp>(
+        loc, TypeRange(), transform::FailurePropagationMode::Propagate,
+        builder.getType<transform::AnyOpType>(),
+        [](OpBuilder &b, Location nested, Value rootH) {
+          b.create<mlir::test::TestPrintRemarkAtOperandOp>(
+              nested, rootH, "remark from generated");
+          b.create<transform::YieldOp>(nested, ValueRange());
+        });
+    return success();
+  }
+
   void runOnOperation() override {
     unsigned firstSetOptions =
         numberOfSetOptions(bindFirstExtraToOps, bindFirstExtraToParams,
                            bindFirstExtraToResultsOfOps);
     unsigned secondSetOptions =
         numberOfSetOptions(bindSecondExtraToOps, bindSecondExtraToParams,
                            bindSecondExtraToResultsOfOps);
     auto loc = UnknownLoc::get(&getContext());
Option<std::string> debugTransformRootTag{
           "the given value as container IR for top-level transform ops. This "
           "allows user control on what transformation to apply. If empty, "
           "select the container of the top-level transform op.")};
   Option<std::string> transformLibraryFileName{
       *this, "transform-library-file-name", llvm::cl::init(""),
       llvm::cl::desc(
           "Optional name of the file containing transform dialect symbol "
           "definitions to be injected into the transform module.")};
+
+  Option<bool> testModuleGeneration{
+      *this, "test-module-generation", llvm::cl::init(false),
+      llvm::cl::desc("test the generation of the transform module during pass "
+                     "initialization, overridden by parsing")};
 };
 
 struct TestTransformDialectEraseSchedulePass
     : public PassWrapper<TestTransformDialectEraseSchedulePass,
                          OperationPass<ModuleOp>> {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(
       TestTransformDialectEraseSchedulePass)