[mlir][transform] Handle multiple library preloading passes.

The transform dialect stores a "library module" that the preload pass
can populate. Until now, each pass registered an additional module by
simply pushing it to a vector; however, the interpreter only used the
first of them. This commit turns the registration into "loading", i.e.,
each newly added module gets merged into the existing one. This allows
the loading to be split into several passes, and using the library in
the interpreter now takes all of them into account. While this design
avoids repeated merging every time the library is accessed, it requires
that the implementation of merging modules lives in the
`TransformDialect` target (since it at the dialect depend on each
other).

This resolves #69111.

Author: ingomueller-net

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

@@ -78,20 +78,19 @@ def Transform_Dialect : Dialect {
     using ExtensionTypePrintingHook =
         std::function<void(::mlir::Type, ::mlir::AsmPrinter &)>;
 
-    /// Appends the given module as a transform symbol library available to
-    /// all dialect users.
-    void registerLibraryModule(::mlir::OwningOpRef<::mlir::ModuleOp> &&
-                                library) {
-      libraryModules.push_back(std::move(library));
-    }
-
-    /// Returns a range of registered library modules.
-    auto getLibraryModules() const {
-      return ::llvm::map_range(
-          libraryModules,
-          [](const ::mlir::OwningOpRef<::mlir::ModuleOp> &library) {
-        return library.get();
-      });
+    /// Loads the given module into the transform symbol library module.
+    void initializeLibraryModule();
+
+    /// Loads the given module into the transform symbol library module.
+    LogicalResult loadIntoLibraryModule(::mlir::OwningOpRef<::mlir::ModuleOp> &&
+                                        library);
+
+    /// Returns the transform symbol library module available to all dialect
+    /// users.
+    ModuleOp getLibraryModule() const {
+      if (libraryModule)
+        return libraryModule.get();
+      return ModuleOp();
     }
 
   private:
@@ -153,10 +152,9 @@ def Transform_Dialect : Dialect {
     ::llvm::DenseMap<::mlir::TypeID, ExtensionTypePrintingHook>
         typePrintingHooks;
 
-    /// Modules containing symbols, e.g. named sequences, that will be
-    /// resolved by the interpreter when used.
-    ::llvm::SmallVector<::mlir::OwningOpRef<::mlir::ModuleOp>, 2>
-        libraryModules;
+    /// Module containing symbols, e.g. named sequences, that will be resolved
+    /// by the interpreter when used.
+    ::mlir::OwningOpRef<::mlir::ModuleOp> libraryModule;
   }];
 }
 

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

@@ -11,6 +11,7 @@
 #include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
 #include "mlir/Dialect/Transform/IR/TransformOps.h"
 #include "mlir/Dialect/Transform/IR/TransformTypes.h"
+#include "mlir/Dialect/Transform/IR/Utils.h"
 #include "mlir/IR/DialectImplementation.h"
 #include "llvm/ADT/SCCIterator.h"
 
@@ -65,6 +66,7 @@ void transform::TransformDialect::initialize() {
 #include "mlir/Dialect/Transform/IR/TransformOps.cpp.inc"
       >();
   initializeTypes();
+  initializeLibraryModule();
 }
 
 Type transform::TransformDialect::parseType(DialectAsmParser &parser) const {
@@ -89,6 +91,20 @@ void transform::TransformDialect::printType(Type type,
   it->getSecond()(type, printer);
 }
 
+void transform::TransformDialect::initializeLibraryModule() {
+  MLIRContext *context = getContext();
+  auto loc =
+      FileLineColLoc::get(context, "<transform-dialect-library-module>", 0, 0);
+  libraryModule = ModuleOp::create(loc, "__transform_library");
+  libraryModule.get()->setAttr(TransformDialect::kWithNamedSequenceAttrName,
+                               UnitAttr::get(context));
+}
+
+LogicalResult transform::TransformDialect::loadIntoLibraryModule(
+    ::mlir::OwningOpRef<::mlir::ModuleOp> &&library) {
+  return detail::mergeSymbolsInto(getLibraryModule(), std::move(library));
+}
+
 void transform::TransformDialect::reportDuplicateTypeRegistration(
     StringRef mnemonic) {
   std::string buffer;

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

@@ -33,7 +33,9 @@ class PreloadLibraryPass
     // TODO: investigate using a resource blob if some ownership mode allows it.
     auto *dialect =
         getContext().getOrLoadDialect<transform::TransformDialect>();
-    dialect->registerLibraryModule(std::move(mergedParsedLibraries));
+    if (failed(
+            dialect->loadIntoLibraryModule(std::move(mergedParsedLibraries))))
+      signalPassFailure();
   }
 };
 } // namespace

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

@@ -107,12 +107,8 @@ LogicalResult transform::detail::parseTransformModuleFromFile(
 }
 
 ModuleOp transform::detail::getPreloadedTransformModule(MLIRContext *context) {
-  auto preloadedLibraryRange =
-      context->getOrLoadDialect<transform::TransformDialect>()
-          ->getLibraryModules();
-  if (!preloadedLibraryRange.empty())
-    return *preloadedLibraryRange.begin();
-  return ModuleOp();
+  return context->getOrLoadDialect<transform::TransformDialect>()
+      ->getLibraryModule();
 }
 
 transform::TransformOpInterface

mlir/test/Dialect/Transform/preload-library.mlir

@@ -3,6 +3,18 @@
 // RUN:   -transform-interpreter=entry-point=private_helper \
 // RUN:   -split-input-file -verify-diagnostics
 
+// RUN: mlir-opt %s \
+// RUN:   -transform-preload-library=transform-library-paths=%p%{fs-sep}test-interpreter-library/definitions-self-contained.mlir \
+// RUN:   -transform-preload-library=transform-library-paths=%p%{fs-sep}test-interpreter-library/definitions-with-unresolved.mlir \
+// RUN:   -transform-interpreter=entry-point=private_helper \
+// RUN:   -split-input-file -verify-diagnostics
+
+// RUN: mlir-opt %s \
+// RUN:   -transform-preload-library=transform-library-paths=%p%{fs-sep}test-interpreter-library/definitions-with-unresolved.mlir \
+// RUN:   -transform-preload-library=transform-library-paths=%p%{fs-sep}test-interpreter-library/definitions-self-contained.mlir \
+// RUN:   -transform-interpreter=entry-point=private_helper \
+// RUN:   -split-input-file -verify-diagnostics
+
 // expected-remark @below {{message}}
 module {}
 

mlir/unittests/Dialect/Transform/Preload.cpp

@@ -7,6 +7,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Transform/IR/TransformDialect.h"
+#include "mlir/Dialect/Transform/IR/Utils.h"
 #include "mlir/Dialect/Transform/Transforms/TransformInterpreterUtils.h"
 #include "mlir/IR/AsmState.h"
 #include "mlir/IR/DialectRegistry.h"
@@ -67,7 +68,9 @@ TEST(Preload, ContextPreloadConstructedLibrary) {
   OwningOpRef<ModuleOp> transformLibrary =
       parseSourceString<ModuleOp>(library, parserConfig, "<transform-library>");
   EXPECT_TRUE(transformLibrary) << "failed to parse transform module";
-  dialect->registerLibraryModule(std::move(transformLibrary));
+  LogicalResult diag =
+      dialect->loadIntoLibraryModule(std::move(transformLibrary));
+  EXPECT_TRUE(succeeded(diag));
 
   ModuleOp retrievedTransformLibrary =
       transform::detail::getPreloadedTransformModule(&context);