[mlir] support external named transform libraries

[ftynse]

Summary

Introduce support for external definitions of named sequences in the
transform dialect by letting the TransformInterpreterPassBase read a
"library" MLIR file. This file is expected to contain definitions for
named sequences that are only declared in the main transformation
script. This allows for sharing non-trivial transform combinations
without duplication.

This patch provides only the minimal plumbing for a single textual IR
file. Further changes are possible to support multiple libraries and
bytecode files.

[nicolasvasilache]

LGTM, thanks!

mlir/lib/Dialect/Transform/IR/TransformOps.cpp

792–793

DCE?

[ftynse]

Address review

[aaronmondal]

I think this is breaking CI.

INFO: From Testing @llvm-project//mlir/test/Dialect:SCF/one-shot-bufferize.mlir.test:
==================== Test output for @llvm-project//mlir/test/Dialect:SCF/one-shot-bufferize.mlir.test:
lit.py: /var/lib/buildkite-agent/builds/llvm-bazel-premerge-buildkite-2-3/llvm-project/premerge-checks/llvm/utils/lit/lit/llvm/subst.py:122: note: Did not find mlir-tblgen in external/llvm-project/mlir:external/llvm-project/llvm

(Temporary) CI log: https://buildkite.com/llvm-project/premerge-checks/builds/143762#018729e9-79e6-4916-a163-e1908bb9fdc2

[aaronmondal]

Naive attempt to fix https://reviews.llvm.org/D147088

[GMNGeoffrey]

In D146961#4228641, @aaronmondal wrote:

I think this is breaking CI.

INFO: From Testing @llvm-project//mlir/test/Dialect:SCF/one-shot-bufferize.mlir.test:
==================== Test output for @llvm-project//mlir/test/Dialect:SCF/one-shot-bufferize.mlir.test:
lit.py: /var/lib/buildkite-agent/builds/llvm-bazel-premerge-buildkite-2-3/llvm-project/premerge-checks/llvm/utils/lit/lit/llvm/subst.py:122: note: Did not find mlir-tblgen in external/llvm-project/mlir:external/llvm-project/llvm

(Temporary) CI log: https://buildkite.com/llvm-project/premerge-checks/builds/143762#018729e9-79e6-4916-a163-e1908bb9fdc2

Note that the Bazel CI is not officially maintained, so in general it isn't the responsibility of the community to maintain it. That said @ftynse has opted in to the Bazel phab project and maintenance of the Bazel build (one of the many perks of working at Google 😛). @ftynse, it looks like this failed on presubmit as well: https://buildkite.com/llvm-project/premerge-checks/builds/143452#01872366-31f6-4d33-b73f-89cbb9af50f3, including the build that finished before review and submission: https://buildkite.com/llvm-project/premerge-checks/builds/143452#01872366-31f6-4d33-b73f-89cbb9af50f3. Could you please make sure not to commit things that break the Bazel build?

[ftynse]

Has this failure been clearly root-caused to this patch? Because I have really hard time seeing how SCF bufferization tests are related to this change.

mlir/include/mlir/Dialect/Transform/IR/TransformOps.td

let description = [{
     Named sequences may include other named sequences via `transform.include`,
     but recursion is *not* allowed.
   }];
   
   let arguments = (ins
     SymbolNameAttr:$sym_name,
     TypeAttrBase<"::mlir::FunctionType",
                  "function type attribute">:$function_type,
+    OptionalAttr<StrAttr>:$sym_visibility,
     OptionalAttr<DictArrayAttr>:$arg_attrs,
     OptionalAttr<DictArrayAttr>:$res_attrs);
-  let regions = (region SizedRegion<1>:$body);
+  let regions = (region MaxSizedRegion<1>:$body);
 
   let hasCustomAssemblyFormat = 1;
   let hasVerifier = 1;
 
   let extraClassDeclaration = [{
     ::llvm::ArrayRef<::mlir::Type> getArgumentTypes() {
       return getFunctionType().getInputs();
     }

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

 class Operation;
 template <typename>
 class OwningOpRef;
 class Region;
 
 namespace transform {
 namespace detail {
 /// Template-free implementation of TransformInterpreterPassBase::initialize.
-LogicalResult
-interpreterBaseInitializeImpl(MLIRContext *context, StringRef transformFileName,
-                              std::shared_ptr<OwningOpRef<ModuleOp>> &module);
+LogicalResult interpreterBaseInitializeImpl(
+    MLIRContext *context, StringRef transformFileName,
+    StringRef transformLibraryFileName,
+    std::shared_ptr<OwningOpRef<ModuleOp>> &module,
+    std::shared_ptr<OwningOpRef<ModuleOp>> &libraryModule);
 
 /// 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,
     const TransformOptions &options,
     const Pass::Option<std::string> &transformFileName,
+    const Pass::Option<std::string> &transformLibraryFileName,
     const Pass::Option<std::string> &debugPayloadRootTag,
     const Pass::Option<std::string> &debugTransformRootTag,
     StringRef binaryName);
 } // namespace detail
 
 /// Base class for transform dialect interpreter passes that can consume and
 /// dump transform dialect scripts in separate files. The pass is controlled by
 /// three string options:
 ///
 ///   - transformFileName: if non-empty, the name of the file containing the
 ///     transform script. If empty, `debugTransformRootTag` is considered or the
 ///     pass root operation must contain a single top-level transform op that
 ///     will be interpreted.
+///   - transformLibraryFileName: if non-empty, the name of the file containing
+///     definitions of external symbols referenced in the transform script.
+///     These definitions will be used to replace declarations.
 ///   - debugPayloadRootTag: if non-empty, the value of the attribute named
 ///     `kTransformDialectTagAttrName` indicating the single op that is
 ///     considered the payload root of the transform interpreter; otherwise, the
 ///     root operation of the pass is used.
 ///   - debugTransformRootTag: if non-empty, the value of the attribute named
 ///     `kTransformDialectTagAttrName` indicating the single top-level transform
 ///     op contained in the payload root to be used as the entry point by the
 ///     transform interpreter; mutually exclusive with `transformFileName`.
static_assert(                                                               \
       std::is_same_v<                                                          \
           std::remove_reference_t<decltype(std::declval<Concrete &>().NAME)>,  \
           Pass::Option<std::string>>,                                          \
       "required " #NAME " string pass option is missing")
 
     REQUIRE_PASS_OPTION(transformFileName);
     REQUIRE_PASS_OPTION(debugPayloadRootTag);
     REQUIRE_PASS_OPTION(debugTransformRootTag);
+    REQUIRE_PASS_OPTION(transformLibraryFileName);
 
 #undef REQUIRE_PASS_OPTION
 
     StringRef transformFileName =
         static_cast<Concrete *>(this)->transformFileName;
-    return detail::interpreterBaseInitializeImpl(context, transformFileName,
-                                                 sharedTransformModule);
+    StringRef transformLibraryFileName =
+        static_cast<Concrete *>(this)->transformLibraryFileName;
+    return detail::interpreterBaseInitializeImpl(
+        context, transformFileName, transformLibraryFileName,
+        sharedTransformModule, transformLibraryModule);
   }
 
   /// 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(); }
 
   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,
             /*extraMappings=*/{}, options, pass->transformFileName,
-            pass->debugPayloadRootTag, pass->debugTransformRootTag,
-            binaryName)) ||
+            pass->transformLibraryFileName, pass->debugPayloadRootTag,
+            pass->debugTransformRootTag, binaryName)) ||
         failed(pass->runAfterInterpreter(op))) {
       return pass->signalPassFailure();
     }
   }
 
 protected:
   /// Transform interpreter options.
   TransformOptions options;
 
   /// Returns a read-only reference to shared transform module.
   const std::shared_ptr<OwningOpRef<ModuleOp>> &
   getSharedTransformModule() const {
     return sharedTransformModule;
   }
 
+  /// Returns a read-only reference to the transform library module.
+  const std::shared_ptr<OwningOpRef<ModuleOp>> &
+  getTransformLibraryModule() const {
+    return transformLibraryModule;
+  }
+
 private:
   /// The separate transform module to be used for transformations, shared
   /// across multiple instances of the pass if it is applied in parallel to
   /// avoid potentially expensive cloning. MUST NOT be modified after the pass
   /// has been initialized.
   std::shared_ptr<OwningOpRef<ModuleOp>> sharedTransformModule = nullptr;
+
+  /// The transform module containing symbol definitions that become available
+  /// in the transform scripts. Similar to dynamic linking for binaries. This is
+  /// shared across multiple instances of the pass and therefore MUST NOT be
+  /// modified after the pass has been initialized.
+  std::shared_ptr<OwningOpRef<ModuleOp>> transformLibraryModule = nullptr;
 };
 
 } // namespace transform
 } // namespace mlir
 
 #endif // MLIR_DIALECT_TRANSFORM_TRANSFORMS_TRANSFORMINTERPRETERPASSBASE_H

mlir/lib/Dialect/Transform/IR/TransformOps.cpp

 
 DiagnosedSilenceableFailure
 transform::IncludeOp::apply(transform::TransformResults &results,
                             transform::TransformState &state) {
   auto callee = SymbolTable::lookupNearestSymbolFrom<NamedSequenceOp>(
       getOperation(), getTarget());
   assert(callee && "unverified reference to unknown symbol");
 
+  if (callee.isExternal())
+    return emitDefiniteFailure() << "unresolved external named sequence";
+
   // Map operands to block arguments.
   SmallVector<SmallVector<MappedValue>> mappings;
   detail::prepareValueMappings(mappings, getOperands(), state);
   auto scope = state.make_isolated_region_scope(callee.getBody());
   for (auto &&[arg, map] :
        llvm::zip_equal(callee.getBody().front().getArguments(), mappings)) {
     if (failed(state.mapBlockArgument(arg, map)))
       return DiagnosedSilenceableFailure::definiteFailure();
void transform::IncludeOp::getEffects(
   DiagnosedSilenceableFailure earlyVerifierResult =
       verifyNamedSequenceOp(callee);
   if (!earlyVerifierResult.succeeded()) {
     (void)earlyVerifierResult.silence();
     return;
   }
 
   // Carry over effects from the callee.
+  // TODO: external callees must provides attributes annotating the
+  // readonly/consume effects on operands.
+  if (!callee.isExternal())
     remapArgumentEffects(callee.getBody().front(), getOperands(), effects);
 
   // Proper effects.
   onlyReadsHandle(getOperands(), effects);
   producesHandle(getResults(), effects);
 }
 
 template <typename... Tys>
 static bool implementSameInterface(Type t1, Type t2) {
function_interface_impl::printFunctionOp(
       getFunctionTypeAttrName().getValue(), getArgAttrsAttrName(),
       getResAttrsAttrName());
 }
 
 /// Verification of a NamedSequenceOp. This does not report the error
 /// immediately, so it can be used to check for op's well-formedness before the
 /// verifier runs, e.g., during trait verification.
 static DiagnosedSilenceableFailure
 verifyNamedSequenceOp(transform::NamedSequenceOp op) {
-  if (op.isExternal())
-    return emitSilenceableFailure(op) << "cannot be empty";
-
   if (Operation *parent = op->getParentWithTrait<OpTrait::SymbolTable>()) {
     if (!parent->getAttr(
             transform::TransformDialect::kWithNamedSequenceAttrName)) {
       DiagnosedSilenceableFailure diag =
           emitSilenceableFailure(op)
           << "expects the parent symbol table to have the '"
           << transform::TransformDialect::kWithNamedSequenceAttrName
           << "' attribute";
       diag.attachNote(parent->getLoc()) << "symbol table operation";
       return diag;
     }
   }
 
   if (auto parent = op->getParentOfType<transform::TransformOpInterface>()) {
     DiagnosedSilenceableFailure diag =
         emitSilenceableFailure(op)
         << "cannot be defined inside another transform op";
     diag.attachNote(parent.getLoc()) << "ancestor transform op";
     return diag;
   }
 
+  if (op.isExternal() || op.getBody().empty())
+    return DiagnosedSilenceableFailure::success();
+
   if (op.getBody().front().empty())
     return emitSilenceableFailure(op) << "expected a non-empty body block";
 
   Operation *terminator = &op.getBody().front().back();
   if (!isa<transform::YieldOp>(terminator)) {
     DiagnosedSilenceableFailure diag = emitSilenceableFailure(op)
                                        << "expected '"
                                        << transform::YieldOp::getOperationName()

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

 //===- TransformInterpreterPassBase.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
 //
 //===----------------------------------------------------------------------===//
 //
 // Base class with shared implementation for transform dialect interpreter
 // passes.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Transform/Transforms/TransformInterpreterPassBase.h"
 #include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
 #include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/FunctionInterfaces.h"
+#include "mlir/IR/Verifier.h"
 #include "mlir/IR/Visitors.h"
 #include "mlir/Parser/Parser.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Support/FileUtilities.h"
+#include "llvm/ADT/ScopeExit.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/FileSystem.h"
 #include "llvm/Support/FormatVariadic.h"
 #include "llvm/Support/Mutex.h"
 #include "llvm/Support/Path.h"
 #include "llvm/Support/SourceMgr.h"
 #include "llvm/Support/raw_ostream.h"

 
 /// Prints the CLI command running the repro with the current path.
 // TODO: make binary name optional by querying LLVM command line API for the
 // name of the current binary.
 static llvm::raw_ostream &
 printReproCall(llvm::raw_ostream &os, StringRef rootOpName, StringRef passName,
                const Pass::Option<std::string> &debugPayloadRootTag,
                const Pass::Option<std::string> &debugTransformRootTag,
+               const Pass::Option<std::string> &transformLibraryFileName,
                StringRef binaryName) {
+  std::string transformLibraryOption = "";
+  if (!transformLibraryFileName.empty()) {
+    transformLibraryOption =
+        llvm::formatv(" {0}={1}", transformLibraryFileName.getArgStr(),
+                      transformLibraryFileName.getValue())
+            .str();
+  }
   os << llvm::formatv(
-      "{6} --pass-pipeline=\"{0}({1}{{{2}={3} {4}={5}})\"", rootOpName,
+      "{7} --pass-pipeline=\"{0}({1}{{{2}={3} {4}={5}{6}})\"", rootOpName,
       passName, debugPayloadRootTag.getArgStr(),
       debugPayloadRootTag.empty()
           ? StringRef(kTransformDialectTagPayloadRootValue)
           : debugPayloadRootTag,
       debugTransformRootTag.getArgStr(),
       debugTransformRootTag.empty()
           ? StringRef(kTransformDialectTagTransformContainerValue)
           : debugTransformRootTag,
-      binaryName);
+      transformLibraryOption, binaryName);
   return os;
 }
 
 /// Prints the module rooted at `root` to `os` and appends
 /// `transformContainer` if it is not nested in `root`.
 llvm::raw_ostream &printModuleForRepro(llvm::raw_ostream &os, Operation *root,
                                        Operation *transform) {
   root->print(os);
   if (!root->isAncestor(transform))
     transform->print(os);
   return os;
 }
 
 /// Saves the payload and the transform IR into a temporary file and reports
 /// the file name to `os`.
-void saveReproToTempFile(llvm::raw_ostream &os, Operation *target,
-                         Operation *transform, StringRef passName,
-                         const Pass::Option<std::string> &debugPayloadRootTag,
-                         const Pass::Option<std::string> &debugTransformRootTag,
-                         StringRef binaryName) {
+void saveReproToTempFile(
+    llvm::raw_ostream &os, Operation *target, Operation *transform,
+    StringRef passName, const Pass::Option<std::string> &debugPayloadRootTag,
+    const Pass::Option<std::string> &transformLibraryFileName,
   using llvm::sys::fs::TempFile;
   Operation *root = getRootOperation(target);
 
   SmallVector<char, 128> tmpPath;
   llvm::sys::path::system_temp_directory(/*erasedOnReboot=*/true, tmpPath);
   llvm::sys::path::append(tmpPath, "transform_dialect_%%%%%%.mlir");
   llvm::Expected<TempFile> tempFile = TempFile::create(tmpPath);
   if (!tempFile) {
     os << "could not open temporary file to save the repro\n";
     return;
   }
 
   llvm::raw_fd_ostream fout(tempFile->FD, /*shouldClose=*/false);
   printModuleForRepro(fout, root, transform);
   fout.flush();
   std::string filename = tempFile->TmpName;
 
   if (tempFile->keep()) {
     os << "could not preserve the temporary file with the repro\n";
     return;
   }
 
   os << "=== Transform Interpreter Repro ===\n";
   printReproCall(os, root->getName().getStringRef(), passName,
-                 debugPayloadRootTag, debugTransformRootTag, binaryName)
+                 debugPayloadRootTag, debugTransformRootTag,
+                 transformLibraryFileName, binaryName)
       << " " << filename << "\n";
   os << "===================================\n";
 }
 
 // Optionally perform debug actions requested by the user to dump IR and a
 // repro to stderr and/or a file.
 static void performOptionalDebugActions(
     Operation *target, Operation *transform, StringRef passName,
     const Pass::Option<std::string> &debugPayloadRootTag,
     const Pass::Option<std::string> &debugTransformRootTag,
+    const Pass::Option<std::string> &transformLibraryFileName,
     StringRef binaryName) {
   MLIRContext *context = target->getContext();
 
   // If we are not planning to print, bail early.
   bool hasDebugFlags = false;
   DEBUG_WITH_TYPE(DEBUG_TYPE_DUMP_STDERR, { hasDebugFlags = true; });
   DEBUG_WITH_TYPE(DEBUG_TYPE_DUMP_FILE, { hasDebugFlags = true; });
   if (!hasDebugFlags)
transform->setAttr(
         kTransformDialectTagAttrName,
         StringAttr::get(context, kTransformDialectTagTransformContainerValue));
   }
 
   DEBUG_WITH_TYPE(DEBUG_TYPE_DUMP_STDERR, {
     llvm::dbgs() << "=== Transform Interpreter Repro ===\n";
     printReproCall(llvm::dbgs() << "cat <<EOF | ",
                    root->getName().getStringRef(), passName,
-                   debugPayloadRootTag, debugTransformRootTag, binaryName)
+                   debugPayloadRootTag, debugTransformRootTag,
+                   transformLibraryFileName, binaryName)
         << "\n";
     printModuleForRepro(llvm::dbgs(), root, transform);
     llvm::dbgs() << "\nEOF\n";
     llvm::dbgs() << "===================================\n";
   });
   (void)root;
   DEBUG_WITH_TYPE(DEBUG_TYPE_DUMP_FILE, {
     saveReproToTempFile(llvm::dbgs(), target, transform, passName,
-                        debugPayloadRootTag, debugTransformRootTag, binaryName);
+                        debugPayloadRootTag, debugTransformRootTag,
+                        transformLibraryFileName, binaryName);
   });
 }
 
+/// Replaces external symbols in `block` with their (non-external) definitions
+/// from the given module.
+static LogicalResult defineDeclaredSymbols(Block &block, ModuleOp definitions) {
+  for (Operation &op : llvm::make_early_inc_range(block)) {
+    LLVM_DEBUG(DBGS() << op << "\n");
+    auto symbol = dyn_cast<SymbolOpInterface>(op);
+    if (!symbol)
+      continue;
+    if (symbol->getNumRegions() == 1 && !symbol->getRegion(0).empty())
+      continue;
+
+    LLVM_DEBUG(DBGS() << "looking for definition of symbol "
+                      << symbol.getNameAttr() << ":");
+    SymbolTable symbolTable(definitions);
+    Operation *externalSymbol = symbolTable.lookup(symbol.getNameAttr());
+    if (!externalSymbol || externalSymbol->getNumRegions() != 1 ||
+        externalSymbol->getRegion(0).empty()) {
+      LLVM_DEBUG(llvm::dbgs() << "not found\n");
+      continue;
+    }
+
+    auto symbolFunc = dyn_cast<FunctionOpInterface>(op);
+    auto externalSymbolFunc = dyn_cast<FunctionOpInterface>(externalSymbol);
+    if (!symbolFunc || !externalSymbolFunc) {
+      LLVM_DEBUG(llvm::dbgs() << "cannot compare types\n");
+      continue;
+    }
+
+    LLVM_DEBUG(llvm::dbgs() << "found @" << externalSymbol << "\n");
+    if (symbolFunc.getFunctionType() != externalSymbolFunc.getFunctionType()) {
+      return symbolFunc.emitError()
+             << "external definition has a mismatching signature ("
+             << externalSymbolFunc.getFunctionType() << ")";
+    }
+
+    OpBuilder builder(&op);
+    builder.setInsertionPoint(&op);
+    builder.clone(*externalSymbol);
+    symbol->erase();
+  }
+
+  return success();
+}
+
 LogicalResult transform::detail::interpreterBaseRunOnOperationImpl(
     Operation *target, StringRef passName,
     const std::shared_ptr<OwningOpRef<ModuleOp>> &sharedTransformModule,
+    const std::shared_ptr<OwningOpRef<ModuleOp>> &libraryModule,
     const RaggedArray<MappedValue> &extraMappings,
     const TransformOptions &options,
     const Pass::Option<std::string> &transformFileName,
+    const Pass::Option<std::string> &transformLibraryFileName,
     const Pass::Option<std::string> &debugPayloadRootTag,
     const Pass::Option<std::string> &debugTransformRootTag,
     StringRef binaryName) {
 
   // Step 1
   // ------
   // If debugPayloadRootTag was passed, then we are in user-specified selection
   // of the transformed IR. This corresponds to REPL debug mode. Otherwise, just
LogicalResult transform::detail::interpreterBaseRunOnOperationImpl(
   if (!transformRoot->hasTrait<PossibleTopLevelTransformOpTrait>()) {
     return emitError(transformRoot->getLoc())
            << "expected the transform entry point to be a top-level transform "
               "op";
   }
 
   // Step 3
   // ------
+  // Copy external defintions for symbols if provided. Be aware of potential
+  // concurrent execution (normally, the error shouldn't be triggered unless the
+  // transform IR modifies itself in a pass, which is also forbidden elsewhere).
+  if (!sharedTransform && libraryModule && *libraryModule) {
+    if (!target->isProperAncestor(transformRoot)) {
+      InFlightDiagnostic diag =
+          transformRoot->emitError()
+          << "cannot inject transform definitions next to pass anchor op";
+      diag.attachNote(target->getLoc()) << "pass anchor op";
+      return diag;
+    }
+    if (failed(defineDeclaredSymbols(*transformRoot->getBlock(),
+                                     libraryModule->get())))
+      return failure();
+  }
+
+  // Step 4
+  // ------
   // Optionally perform debug actions requested by the user to dump IR and a
   // repro to stderr and/or a file.
   performOptionalDebugActions(target, transformRoot, passName,
                               debugPayloadRootTag, debugTransformRootTag,
-                              binaryName);
+                              transformLibraryFileName, binaryName);
 
-  // Step 4
+  // Step 5
   // ------
   // Apply the transform to the IR
   return applyTransforms(payloadRoot, cast<TransformOpInterface>(transformRoot),
                          extraMappings, options);
 }
 
 LogicalResult transform::detail::interpreterBaseInitializeImpl(
     MLIRContext *context, StringRef transformFileName,
-    std::shared_ptr<OwningOpRef<ModuleOp>> &module) {
+    StringRef transformLibraryFileName,
+    std::shared_ptr<OwningOpRef<ModuleOp>> &module,
+    std::shared_ptr<OwningOpRef<ModuleOp>> &libraryModule) {
   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();
 
   module = std::make_shared<OwningOpRef<ModuleOp>>(std::move(parsed));
+  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/ops-invalid.mlir

 // -----
 
 // expected-error @below {{attribute can only be attached to operations with symbol tables}}
 "test.unknown_container"() { transform.with_named_sequence } : () -> ()
 
 // -----
 
 module attributes { transform.with_named_sequence } {
-  // expected-error @below {{failed to verify constraint: region with 1 blocks}}
-  "transform.named_sequence"() ({}) { sym_name = "external_named_sequence", function_type = () -> () } : () -> ()
+  // expected-error @below {{expected a non-empty body block}}
+  "transform.named_sequence"() ({
+  ^bb0:
+  }) { sym_name = "external_named_sequence", function_type = () -> () } : () -> ()
 
   transform.sequence failures(propagate) {
   ^bb0(%arg0: !transform.any_op):
     transform.include @external_named_sequence failures(propagate) () : () -> ()
   }
 }
 
 // -----

mlir/test/Dialect/Transform/test-interpreter-external-symbol-decl-and-schedule.mlir

+// RUN: mlir-opt %s --pass-pipeline="builtin.module(test-transform-dialect-interpreter{transform-file-name=%p/test-interpreter-external-symbol-decl.mlir transform-library-file-name=%p/test-interpreter-external-symbol-def.mlir})" \
+// RUN:             --verify-diagnostics
+
+// RUN: mlir-opt %s --pass-pipeline="builtin.module(test-transform-dialect-interpreter{transform-file-name=%p/test-interpreter-external-symbol-decl.mlir transform-library-file-name=%p/test-interpreter-external-symbol-def.mlir}, test-transform-dialect-interpreter{transform-file-name=%p/test-interpreter-external-symbol-decl.mlir transform-library-file-name=%p/test-interpreter-external-symbol-def.mlir})" \
+// RUN:             --verify-diagnostics
+
+// The external transform script has a declaration to the named sequence @foo,
+// the definition of which is provided in another file. Repeated application
+// of the same pass should not be a problem. Note that the same diagnostic
+// produced twice at the same location only needs to be matched once.
+
+// expected-remark @below {{message}}
+module {}

mlir/test/Dialect/Transform/test-interpreter-external-symbol-decl-invalid.mlir

+// RUN: mlir-opt %s --pass-pipeline="builtin.module(test-transform-dialect-interpreter{transform-library-file-name=%p/test-interpreter-external-symbol-def.mlir}, test-transform-dialect-interpreter)" \
+// RUN:             --verify-diagnostics --split-input-file
+
+// The definition of the @foo named sequence is provided in another file. It
+// will be included because of the pass option.
+
+module attributes {transform.with_named_sequence} {
+  // expected-error @below {{external definition has a mismatching signature}}
+  transform.named_sequence private @foo(!transform.op<"builtin.module">)
+
+  transform.sequence failures(propagate) {
+  ^bb0(%arg0: !transform.op<"builtin.module">):
+    include @foo failures(propagate) (%arg0) : (!transform.op<"builtin.module">) -> ()
+  }
+}
+
+// -----
+
+module attributes {transform.with_named_sequence} {
+  transform.named_sequence private @undefined_sequence()
+
+  transform.sequence failures(suppress) {
+  ^bb0(%arg0: !transform.any_op):
+    // expected-error @below {{unresolved external named sequence}}
+    include @undefined_sequence failures(suppress) () : () -> ()
+  }
+}

mlir/test/Dialect/Transform/test-interpreter-external-symbol-decl.mlir

+// RUN: mlir-opt %s --pass-pipeline="builtin.module(test-transform-dialect-interpreter{transform-library-file-name=%p/test-interpreter-external-symbol-def.mlir})" \
+// RUN:             --verify-diagnostics | FileCheck %s
+
+// RUN: mlir-opt %s --pass-pipeline="builtin.module(test-transform-dialect-interpreter{transform-library-file-name=%p/test-interpreter-external-symbol-def.mlir}, test-transform-dialect-interpreter)" \
+// RUN:             --verify-diagnostics | FileCheck %s
+
+// RUN: mlir-opt %s --pass-pipeline="builtin.module(test-transform-dialect-interpreter{transform-library-file-name=%p/test-interpreter-external-symbol-def.mlir}, test-transform-dialect-interpreter{transform-library-file-name=%p/test-interpreter-external-symbol-def.mlir})" \
+// RUN:             --verify-diagnostics | FileCheck %s
+
+// The definition of the @foo named sequence is provided in another file. It
+// will be included because of the pass option. Repeated application of the
+// same pass, with or without the library option, should not be a problem.
+// Note that the same diagnostic produced twice at the same location only
+// needs to be matched once.
+
+// expected-remark @below {{message}}
+module attributes {transform.with_named_sequence} {
+  // CHECK: transform.named_sequence @foo
+  // CHECK: test_print_remark_at_operand %{{.*}}, "message"
+  transform.named_sequence private @foo(!transform.any_op)
+
+  transform.sequence failures(propagate) {
+  ^bb0(%arg0: !transform.any_op):
+    include @foo failures(propagate) (%arg0) : (!transform.any_op) -> ()
+  }
+}

mlir/test/Dialect/Transform/test-interpreter-external-symbol-def.mlir

+// RUN: mlir-opt %s
+
+module attributes {transform.with_named_sequence} {
+  transform.named_sequence @foo(%arg0: !transform.any_op) {
+    transform.test_print_remark_at_operand %arg0, "message" : !transform.any_op
+    transform.yield
+  }
+}

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

void runOnOperation() override {
     } else if (!bindSecondExtraToResultsOfOps.empty()) {
       createOpResultMapping(getOperation(), bindSecondExtraToResultsOfOps,
                             extraMapping);
     }
 
     options = options.enableExpensiveChecks(enableExpensiveChecks);
     if (failed(transform::detail::interpreterBaseRunOnOperationImpl(
             getOperation(), getArgument(), getSharedTransformModule(),
-            extraMapping, options, transformFileName, debugPayloadRootTag,
+            getTransformLibraryModule(), extraMapping, options,
+            transformFileName, transformLibraryFileName, debugPayloadRootTag,
             debugTransformRootTag, getBinaryName())))
       return signalPassFailure();
   }
 
   Option<bool> enableExpensiveChecks{
       *this, "enable-expensive-checks", llvm::cl::init(false),
       llvm::cl::desc("perform expensive checks to better report errors in the "
                      "transform IR")};
Option<std::string> debugPayloadRootTag{
           "operation as the payload IR root.")};
   Option<std::string> debugTransformRootTag{
       *this, "debug-transform-root-tag", llvm::cl::init(""),
       llvm::cl::desc(
           "Select the operation with 'transform.target_tag' attribute having "
           "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.")};
 };
 
 struct TestTransformDialectEraseSchedulePass
     : public PassWrapper<TestTransformDialectEraseSchedulePass,
                          OperationPass<ModuleOp>> {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(
       TestTransformDialectEraseSchedulePass)
 

utils/bazel/llvm-project-overlay/mlir/test/Dialect/BUILD.bazel

lit_test(
         srcs = [src],
         data = [
             "//llvm:llvm-symbolizer",
             "//mlir:mlir-opt",
             "//mlir:mlir-translate",
             "//mlir/test:lit_data",
         ] + glob([
             "Transform/*-source.mlir",
+            "Transform/*-symbol-def.mlir",
         ])
     )
     for src in glob(
         include=["**/*.mlir"],
-        exclude=["Transform/*-source.mlir"]
+        exclude=[
+            "Transform/*-source.mlir",
+            "Transform/*-symbol-def.mlir",
+            "Transform/*-symbol-decl-and-schedule.mlir",
+        ]
     )
 ]