[Mlir-commits] [mlir] Transform dialect resources (PR #68421)
Oleksandr Alex Zinenko
llvmlistbot at llvm.org
Fri Oct 6 07:58:45 PDT 2023
https://github.com/ftynse created https://github.com/llvm/llvm-project/pull/68421
None
>From c6a8f0bd33d0fbaab7f3ecfe09e44921cf545914 Mon Sep 17 00:00:00 2001
From: Jacques Pienaar <jpienaar at google.com>
Date: Wed, 4 Oct 2023 21:50:04 -0700
Subject: [PATCH 1/2] Add resource to transform dialect for external libraries
This resource enables referencing an external file and have it loaded
into the context while still being able to reference it in the IR.
Referencing into the external libraries follows the convention of
resource key followed by nested symbols inside.
This doesn't yet add the attribute type to reference resource. It would
effectively look the same as SymbolRefAttr but consist of two parts
where the former is referencing the resource and the remainder the
nested references. At that point the builder could be changed to only
retain referenced external files.
---
.../Dialect/Transform/IR/TransformDialect.td | 12 ++
.../Dialect/Transform/IR/TransformDialect.cpp | 137 ++++++++++++++++++
.../Transform/test-interpreter-external.mlir | 19 ++-
.../llvm-project-overlay/mlir/BUILD.bazel | 1 +
4 files changed, 166 insertions(+), 3 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Transform/IR/TransformDialect.td b/mlir/include/mlir/Dialect/Transform/IR/TransformDialect.td
index ad6804673b770ca..788eef111104f6b 100644
--- a/mlir/include/mlir/Dialect/Transform/IR/TransformDialect.td
+++ b/mlir/include/mlir/Dialect/Transform/IR/TransformDialect.td
@@ -29,6 +29,18 @@ def Transform_Dialect : Dialect {
constexpr const static ::llvm::StringLiteral
kWithNamedSequenceAttrName = "transform.with_named_sequence";
+ /// Return the registered transform libraries.
+ ::llvm::ArrayRef<::mlir::ModuleOp> getLibraryModules();
+
+ /// Register a Transform library module
+ void registerLibraryModule(OwningOpRef<ModuleOp> &&library);
+
+ /// Return the transform identified using symbol ref attr. The root
+ /// symbol is the resource while the nested references index into the
+ /// selected library.
+ ::mlir::Operation* getRegisteredTransform(::llvm::StringRef resource,
+ ::mlir::SymbolRefAttr ref);
+
/// Name of the attribute attachable to an operation so it can be
/// identified as root by the default interpreter pass.
constexpr const static ::llvm::StringLiteral kTargetTagAttrName =
diff --git a/mlir/lib/Dialect/Transform/IR/TransformDialect.cpp b/mlir/lib/Dialect/Transform/IR/TransformDialect.cpp
index 32c56e903268f74..dc489d99f8fb85f 100644
--- a/mlir/lib/Dialect/Transform/IR/TransformDialect.cpp
+++ b/mlir/lib/Dialect/Transform/IR/TransformDialect.cpp
@@ -12,7 +12,12 @@
#include "mlir/Dialect/Transform/IR/TransformOps.h"
#include "mlir/Dialect/Transform/IR/TransformTypes.h"
#include "mlir/IR/DialectImplementation.h"
+#include "mlir/IR/DialectResourceBlobManager.h"
+#include "mlir/IR/OpImplementation.h"
+#include "mlir/Parser/Parser.h"
#include "llvm/ADT/SCCIterator.h"
+#include "llvm/Support/MemoryBuffer.h"
+#include "llvm/Support/SourceMgr.h"
using namespace mlir;
@@ -57,6 +62,135 @@ void transform::detail::checkImplementsTransformHandleTypeInterface(
}
#endif // NDEBUG
+/// A handle used to reference external elements instances.
+using TransformDialectResourceBlobHandle =
+ mlir::DialectResourceBlobHandle<mlir::transform::TransformDialect>;
+
+struct TransformResourceBlobManagerInterface
+ : public ResourceBlobManagerDialectInterfaceBase<
+ TransformDialectResourceBlobHandle> {
+ using ResourceBlobManagerDialectInterfaceBase<
+ TransformDialectResourceBlobHandle>::
+ ResourceBlobManagerDialectInterfaceBase;
+};
+
+struct TransformOpAsmInterface : public OpAsmDialectInterface {
+ using OpAsmDialectInterface::OpAsmDialectInterface;
+ TransformOpAsmInterface(Dialect *dialect,
+ TransformResourceBlobManagerInterface &mgr)
+ : OpAsmDialectInterface(dialect), blobManager(mgr) {}
+
+ ~TransformOpAsmInterface() override {
+ for (auto op : orderedLibraryModules)
+ op->erase();
+ }
+
+ //===------------------------------------------------------------------===//
+ // Resources
+ //===------------------------------------------------------------------===//
+
+ std::string
+ getResourceKey(const AsmDialectResourceHandle &handle) const override {
+ return cast<TransformDialectResourceBlobHandle>(handle).getKey().str();
+ }
+
+ FailureOr<AsmDialectResourceHandle>
+ declareResource(StringRef key) const final {
+ return blobManager.insert(key);
+ }
+
+ LogicalResult parseResource(AsmParsedResourceEntry &entry) const final {
+ // If its a string, then treat it as a filename.
+ // TODO: Could be extended for blob resources where the file is encoded.
+ if (entry.getKind() != AsmResourceEntryKind::String)
+ return failure();
+
+ FailureOr<std::string> name = entry.parseAsString();
+ if (failed(name))
+ return failure();
+
+ auto fileOr = llvm::MemoryBuffer::getFile(*name);
+ if (!fileOr) {
+ return entry.emitError()
+ << "failed to load resource: " << fileOr.getError().message();
+ }
+
+ // Parse the module in the source file.
+ llvm::SourceMgr sourceMgr;
+ sourceMgr.AddNewSourceBuffer(std::move(fileOr.get()), llvm::SMLoc());
+ auto transformModule = OwningOpRef<ModuleOp>(
+ parseSourceFile<ModuleOp>(sourceMgr, getContext()));
+ if (!transformModule)
+ return entry.emitError() << "failed to parse Transform module";
+
+ orderedLibraryNames.push_back({entry.getKey().str(), *name});
+ registerLibraryModule(entry.getKey(), std::move(transformModule));
+ return success();
+ }
+
+ void
+ buildResources(Operation *op,
+ const SetVector<AsmDialectResourceHandle> &referencedResources,
+ AsmResourceBuilder &provider) const final {
+ // Only print for top-level libraries, print without considering what is
+ // referenced to capture state.
+ if (op->getParentOp() == nullptr) {
+ // On top-level op print libraries additionally.
+ for (auto it : orderedLibraryNames)
+ provider.buildString(std::get<0>(it), std::get<1>(it));
+ }
+ }
+
+ /// Returns a range of registered library modules.
+ ArrayRef<ModuleOp> getLibraryModules() { return orderedLibraryModules; }
+
+ void registerLibraryModule(StringRef key,
+ OwningOpRef<ModuleOp> &&library) const {
+ libraryModules[key] = orderedLibraryModules.emplace_back(library.release());
+ }
+
+ void registerLibraryModule(OwningOpRef<ModuleOp> &&library) const {
+ orderedLibraryModules.push_back(library.release());
+ }
+
+ Operation *getRegisteredTransform(StringRef resource, SymbolRefAttr ref) {
+ auto it = libraryModules.find(resource);
+ if (it == libraryModules.end()) {
+ return nullptr;
+ }
+
+ return SymbolTable::lookupSymbolIn(it->second, ref);
+ }
+
+private:
+ /// The blob manager for the dialect.
+ TransformResourceBlobManagerInterface &blobManager;
+
+ /// Library modules registered.
+ mutable llvm::StringMap<ModuleOp> libraryModules;
+ /// Keep a sorted list too for iteration.
+ mutable llvm::SmallVector<ModuleOp, 2> orderedLibraryModules;
+ /// Keep list of external files used for printing again.
+ mutable llvm::SmallVector<std::pair<std::string, std::string>, 2>
+ orderedLibraryNames;
+};
+
+ArrayRef<ModuleOp> transform::TransformDialect::getLibraryModules() {
+ return getRegisteredInterface<TransformOpAsmInterface>()->getLibraryModules();
+}
+
+void transform::TransformDialect::registerLibraryModule(
+ OwningOpRef<ModuleOp> &&library) {
+ getRegisteredInterface<TransformOpAsmInterface>()->registerLibraryModule(
+ std::move(library));
+}
+
+Operation *
+transform::TransformDialect::getRegisteredTransform(StringRef resource, SymbolRefAttr ref) {
+ auto &interface = *getRegisteredInterface<TransformOpAsmInterface>();
+ return interface.getRegisteredTransform(resource, ref);
+}
+
void transform::TransformDialect::initialize() {
// Using the checked versions to enable the same assertions as for the ops
// from extensions.
@@ -65,6 +199,9 @@ void transform::TransformDialect::initialize() {
#include "mlir/Dialect/Transform/IR/TransformOps.cpp.inc"
>();
initializeTypes();
+
+ auto &blobInterface = addInterface<TransformResourceBlobManagerInterface>();
+ addInterface<TransformOpAsmInterface>(blobInterface);
}
Type transform::TransformDialect::parseType(DialectAsmParser &parser) const {
diff --git a/mlir/test/Dialect/Transform/test-interpreter-external.mlir b/mlir/test/Dialect/Transform/test-interpreter-external.mlir
index 5ac6b66c817afef..bf0e03037f535d3 100644
--- a/mlir/test/Dialect/Transform/test-interpreter-external.mlir
+++ b/mlir/test/Dialect/Transform/test-interpreter-external.mlir
@@ -1,8 +1,21 @@
-// RUN: mlir-opt %s --pass-pipeline="builtin.module(test-transform-dialect-interpreter{transform-file-name=%p/test-interpreter-external-source.mlir})" \
-// RUN: --verify-diagnostics
+// RUN: cd %p && mlir-opt %s --pass-pipeline="builtin.module(test-transform-dialect-interpreter{transform-file-name=test-interpreter-external-source.mlir})" \
+// RUN: --verify-diagnostics
// The schedule in the separate file emits remarks at the payload root.
// expected-remark @below {{outer}}
// expected-remark @below {{inner}}
-module {}
+module attributes { test.blob_ref = #test.e1di64_elements<library> : tensor<*xi1>} {
+ module attributes { test.blob_ref = #test.e1di64_elements<library> : tensor<*xi1>} {}
+ module {}
+ module {}
+}
+
+{-#
+ dialect_resources: {
+ transform: {
+ library: "test-interpreter-external-source.mlir",
+ banana: "test-interpreter-external-source.mlir"
+ }
+ }
+#-}
\ No newline at end of file
diff --git a/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel b/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
index 53b626996f8bbfa..5f551eb62232356 100644
--- a/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
+++ b/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
@@ -11069,6 +11069,7 @@ cc_library(
":LLVMCommonConversion",
":LLVMDialect",
":LoopLikeInterface",
+ ":Parser",
":Pass",
":Rewrite",
":SideEffectInterfaces",
>From 8dbfc8384ae47d77af98715f1bf1f5bce8dbff10 Mon Sep 17 00:00:00 2001
From: Alex Zinenko <zinenko at google.com>
Date: Thu, 5 Oct 2023 16:52:39 +0000
Subject: [PATCH 2/2] [mlir] rope transform resources into the interpreter
This is a minimal connection of transform dialect resources to the
interpreter.
* The ownership of the library modules is moved into a dedicated class
an instance of which is owned by the dialect (rather than the
interface to avoid mutable/const).
* The interpreter pass no longer owns the parsed library module but
stores it in the dialect. Note that pass initialization apparently
happens in multithreaded mode, so locking is required here.
* The resource parsing test is moved into a separate file, the remaining
tests pass as before, including the ones with libraries.
A separate cleanup patch is pending to simplify library/shared module
management so the changes to the interpreter here are intentionally
minimal.
The dialect currently contains an instance of the library manager, and
therefore owns the library modules. This may not be desirable in the
longer run. Instead, the client tool can hold the library manager
instance and give the dialect a reference to it. The client can either
pre-populate the manager or run a pass populating it (using resources or
pass options) before the interpreter pass.
Partially disables two library-loading tests as we need to factor out
non-indempotency into a separate pass that is not the main interpreter.
---
.../Dialect/Transform/IR/TransformDialect.h | 47 ++++++-
.../Dialect/Transform/IR/TransformDialect.td | 23 +---
.../Transforms/TransformInterpreterPassBase.h | 19 +--
.../Dialect/Transform/IR/TransformDialect.cpp | 116 ++++++++++--------
.../TransformInterpreterPassBase.cpp | 34 +++--
.../Transforms/TransformInterpreterUtils.cpp | 2 +-
...-interpreter-external-symbol-decl-dir.mlir | 5 +-
...test-interpreter-external-symbol-decl.mlir | 7 +-
.../Transform/test-interpreter-external.mlir | 15 +--
...est-interpreter-resource-external-def.mlir | 16 +++
.../TestTransformDialectInterpreter.cpp | 5 +-
mlir/unittests/Dialect/Transform/Preload.cpp | 4 +-
12 files changed, 169 insertions(+), 124 deletions(-)
create mode 100644 mlir/test/Dialect/Transform/test-interpreter-resource-external-def.mlir
diff --git a/mlir/include/mlir/Dialect/Transform/IR/TransformDialect.h b/mlir/include/mlir/Dialect/Transform/IR/TransformDialect.h
index db27f2c6fc49b75..7d890093706feb8 100644
--- a/mlir/include/mlir/Dialect/Transform/IR/TransformDialect.h
+++ b/mlir/include/mlir/Dialect/Transform/IR/TransformDialect.h
@@ -15,6 +15,7 @@
#include "mlir/Support/TypeID.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/StringMap.h"
+#include "llvm/Support/Mutex.h"
#include <optional>
namespace mlir {
@@ -66,8 +67,50 @@ class TransformDialectData : public detail::TransformDialectDataBase {
: TransformDialectDataBase(TypeID::get<DerivedTy>(), ctx) {}
};
-#ifndef NDEBUG
namespace detail {
+/// A thread-safe storage object for modules containing libraries of transform
+/// dialect symbols.
+class TransformLibraryManager {
+public:
+ /// Returns a list of modules stored in the the manager. The list itself is
+ /// copied to avoid concurrent modifications.
+ SmallVector<ModuleOp, 2> getLibraryModules() const;
+
+ /// Get a named transform symbol from the library with the given key.
+ Operation *getRegisteredTransform(StringRef key, StringAttr symbol) const;
+
+ /// Get a list of "key, filename" pairs for libraries stored in the manager,
+ /// useful for roundtripping.
+ SmallVector<std::pair<std::string, std::string>, 2>
+ getOrderedLibraryNames() const;
+
+ /// Transfer the library module to the manager and associate it with the
+ /// provided key. Optionally record the source filename.
+ LogicalResult registerLibraryModule(StringRef key,
+ OwningOpRef<ModuleOp> &&library,
+ StringRef filename = "");
+
+private:
+ /// Mutex guarding the manager content as it may be accessed from passes
+ /// concurrently.
+ mutable llvm::sys::SmartRWMutex<true> mutex;
+
+ /// Library modules registered.
+ llvm::StringMap<size_t> libraryModulePositions;
+
+ /// Keep a sorted list too for iteration.
+ SmallVector<OwningOpRef<ModuleOp>, 2> orderedOwningLibraryModules;
+
+ /// Keep list of external files used for printing again.
+ SmallVector<std::pair<std::string, std::string>, 2> orderedLibraryNames;
+
+ /// Precomputed symbol tables.
+ SmallVector<SymbolTable> orderedSymbolTables;
+};
+
+class TransformOpAsmInterface;
+
+#ifndef NDEBUG
/// Asserts that the operations provided as template arguments implement the
/// TransformOpInterface and MemoryEffectsOpInterface. This must be a dynamic
/// assertion since interface implementations may be registered at runtime.
@@ -78,8 +121,8 @@ void checkImplementsTransformOpInterface(StringRef name, MLIRContext *context);
/// interface implementations may be registered at runtime.
void checkImplementsTransformHandleTypeInterface(TypeID typeID,
MLIRContext *context);
-} // namespace detail
#endif // NDEBUG
+} // namespace detail
} // namespace transform
} // namespace mlir
diff --git a/mlir/include/mlir/Dialect/Transform/IR/TransformDialect.td b/mlir/include/mlir/Dialect/Transform/IR/TransformDialect.td
index 788eef111104f6b..14fc1abfecff88b 100644
--- a/mlir/include/mlir/Dialect/Transform/IR/TransformDialect.td
+++ b/mlir/include/mlir/Dialect/Transform/IR/TransformDialect.td
@@ -90,22 +90,6 @@ 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();
- });
- }
-
private:
/// Registers operations specified as template parameters with this
/// dialect. Checks that they implement the required interfaces.
@@ -165,10 +149,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;
+public:
+ /// Symbol libraries owned by the dialect and available to all clients.
+ ::mlir::transform::detail::TransformLibraryManager libraryManager;
}];
}
diff --git a/mlir/include/mlir/Dialect/Transform/Transforms/TransformInterpreterPassBase.h b/mlir/include/mlir/Dialect/Transform/Transforms/TransformInterpreterPassBase.h
index 16ef0bc6a739200..fc51f08e94d778e 100644
--- a/mlir/include/mlir/Dialect/Transform/Transforms/TransformInterpreterPassBase.h
+++ b/mlir/include/mlir/Dialect/Transform/Transforms/TransformInterpreterPassBase.h
@@ -35,7 +35,6 @@ LogicalResult interpreterBaseInitializeImpl(
MLIRContext *context, StringRef transformFileName,
ArrayRef<std::string> transformLibraryPaths,
std::shared_ptr<OwningOpRef<ModuleOp>> &module,
- std::shared_ptr<OwningOpRef<ModuleOp>> &libraryModule,
function_ref<std::optional<LogicalResult>(OpBuilder &, Location)>
moduleBuilder = nullptr);
@@ -44,7 +43,6 @@ LogicalResult interpreterBaseInitializeImpl(
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,
@@ -101,7 +99,6 @@ class TransformInterpreterPassBase : public GeneratedBase<Concrete> {
TransformInterpreterPassBase(const TransformInterpreterPassBase &pass) {
sharedTransformModule = pass.sharedTransformModule;
- transformLibraryModule = pass.transformLibraryModule;
options = pass.options;
}
@@ -139,8 +136,7 @@ class TransformInterpreterPassBase : public GeneratedBase<Concrete> {
static_cast<Concrete *>(this)->transformLibraryPaths;
return detail::interpreterBaseInitializeImpl(
context, transformFileName, transformLibraryPaths,
- sharedTransformModule, transformLibraryModule,
- [this](OpBuilder &builder, Location loc) {
+ sharedTransformModule, [this](OpBuilder &builder, Location loc) {
return static_cast<Concrete *>(this)->constructTransformModule(
builder, loc);
});
@@ -171,7 +167,6 @@ class TransformInterpreterPassBase : public GeneratedBase<Concrete> {
if (failed(pass->runBeforeInterpreter(op)) ||
failed(detail::interpreterBaseRunOnOperationImpl(
op, pass->getArgument(), sharedTransformModule,
- transformLibraryModule,
/*extraMappings=*/{}, options, pass->transformFileName,
pass->transformLibraryPaths, pass->debugPayloadRootTag,
pass->debugTransformRootTag, binaryName)) ||
@@ -190,24 +185,12 @@ class TransformInterpreterPassBase : public GeneratedBase<Concrete> {
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
diff --git a/mlir/lib/Dialect/Transform/IR/TransformDialect.cpp b/mlir/lib/Dialect/Transform/IR/TransformDialect.cpp
index dc489d99f8fb85f..821968d0b17b4ab 100644
--- a/mlir/lib/Dialect/Transform/IR/TransformDialect.cpp
+++ b/mlir/lib/Dialect/Transform/IR/TransformDialect.cpp
@@ -16,7 +16,9 @@
#include "mlir/IR/OpImplementation.h"
#include "mlir/Parser/Parser.h"
#include "llvm/ADT/SCCIterator.h"
+#include "llvm/Support/FileSystem.h"
#include "llvm/Support/MemoryBuffer.h"
+#include "llvm/Support/RWMutex.h"
#include "llvm/Support/SourceMgr.h"
using namespace mlir;
@@ -74,16 +76,18 @@ struct TransformResourceBlobManagerInterface
ResourceBlobManagerDialectInterfaceBase;
};
-struct TransformOpAsmInterface : public OpAsmDialectInterface {
+//===----------------------------------------------------------------------===//
+// TransformOpAsmInterface
+//===----------------------------------------------------------------------===//
+
+class mlir::transform::detail::TransformOpAsmInterface
+ : public OpAsmDialectInterface {
+public:
using OpAsmDialectInterface::OpAsmDialectInterface;
TransformOpAsmInterface(Dialect *dialect,
TransformResourceBlobManagerInterface &mgr)
- : OpAsmDialectInterface(dialect), blobManager(mgr) {}
-
- ~TransformOpAsmInterface() override {
- for (auto op : orderedLibraryModules)
- op->erase();
- }
+ : OpAsmDialectInterface(dialect), blobManager(mgr),
+ dialect(*static_cast<transform::TransformDialect *>(dialect)) {}
//===------------------------------------------------------------------===//
// Resources
@@ -123,9 +127,8 @@ struct TransformOpAsmInterface : public OpAsmDialectInterface {
if (!transformModule)
return entry.emitError() << "failed to parse Transform module";
- orderedLibraryNames.push_back({entry.getKey().str(), *name});
- registerLibraryModule(entry.getKey(), std::move(transformModule));
- return success();
+ return dialect.libraryManager.registerLibraryModule(
+ entry.getKey(), std::move(transformModule), *name);
}
void
@@ -136,61 +139,72 @@ struct TransformOpAsmInterface : public OpAsmDialectInterface {
// referenced to capture state.
if (op->getParentOp() == nullptr) {
// On top-level op print libraries additionally.
- for (auto it : orderedLibraryNames)
- provider.buildString(std::get<0>(it), std::get<1>(it));
- }
- }
-
- /// Returns a range of registered library modules.
- ArrayRef<ModuleOp> getLibraryModules() { return orderedLibraryModules; }
-
- void registerLibraryModule(StringRef key,
- OwningOpRef<ModuleOp> &&library) const {
- libraryModules[key] = orderedLibraryModules.emplace_back(library.release());
- }
-
- void registerLibraryModule(OwningOpRef<ModuleOp> &&library) const {
- orderedLibraryModules.push_back(library.release());
- }
-
- Operation *getRegisteredTransform(StringRef resource, SymbolRefAttr ref) {
- auto it = libraryModules.find(resource);
- if (it == libraryModules.end()) {
- return nullptr;
+ for (auto &&[key, filename] :
+ dialect.libraryManager.getOrderedLibraryNames())
+ provider.buildString(key, filename);
}
-
- return SymbolTable::lookupSymbolIn(it->second, ref);
}
private:
/// The blob manager for the dialect.
TransformResourceBlobManagerInterface &blobManager;
- /// Library modules registered.
- mutable llvm::StringMap<ModuleOp> libraryModules;
- /// Keep a sorted list too for iteration.
- mutable llvm::SmallVector<ModuleOp, 2> orderedLibraryModules;
- /// Keep list of external files used for printing again.
- mutable llvm::SmallVector<std::pair<std::string, std::string>, 2>
- orderedLibraryNames;
+ /// Back reference to the dialect to which the interface is attached.
+ TransformDialect &dialect;
};
-ArrayRef<ModuleOp> transform::TransformDialect::getLibraryModules() {
- return getRegisteredInterface<TransformOpAsmInterface>()->getLibraryModules();
+//===----------------------------------------------------------------------===//
+// TransformLibraryManager
+//===----------------------------------------------------------------------===//
+
+SmallVector<ModuleOp, 2>
+transform::detail::TransformLibraryManager::getLibraryModules() const {
+ llvm::sys::SmartScopedReader<true> lock(mutex);
+ return llvm::to_vector<2>(llvm::map_range(
+ orderedOwningLibraryModules,
+ [](const OwningOpRef<ModuleOp> &owning) { return owning.get(); }));
}
-void transform::TransformDialect::registerLibraryModule(
- OwningOpRef<ModuleOp> &&library) {
- getRegisteredInterface<TransformOpAsmInterface>()->registerLibraryModule(
- std::move(library));
+SmallVector<std::pair<std::string, std::string>, 2>
+transform::detail::TransformLibraryManager::getOrderedLibraryNames() const {
+ llvm::sys::SmartScopedReader<true> lock(mutex);
+ return orderedLibraryNames;
+}
+
+Operation *transform::detail::TransformLibraryManager::getRegisteredTransform(
+ StringRef key, StringAttr symbolName) const {
+ llvm::sys::SmartScopedReader<true> lock(mutex);
+ auto it = libraryModulePositions.find(key);
+ if (it == libraryModulePositions.end())
+ return nullptr;
+
+ return orderedSymbolTables[it->second].lookup(symbolName);
}
-Operation *
-transform::TransformDialect::getRegisteredTransform(StringRef resource, SymbolRefAttr ref) {
- auto &interface = *getRegisteredInterface<TransformOpAsmInterface>();
- return interface.getRegisteredTransform(resource, ref);
+LogicalResult transform::detail::TransformLibraryManager::registerLibraryModule(
+ StringRef key, OwningOpRef<ModuleOp> &&library, StringRef filename) {
+ llvm::sys::SmartScopedWriter<true> lock(mutex);
+ size_t position = orderedOwningLibraryModules.size();
+ orderedOwningLibraryModules.push_back(std::move(library));
+ ModuleOp nonOwning = orderedOwningLibraryModules.back().get();
+ if (!libraryModulePositions.insert({key, position}).second) {
+ // InFlightDiagnostic diag = emitError(nonOwning->getLoc())
+ // << "module for key '" << key
+ // << "' already registered";
+ orderedOwningLibraryModules.pop_back();
+ // return diag;
+ return success();
+ }
+
+ orderedLibraryNames.push_back({key.str(), filename.str()});
+ orderedSymbolTables.emplace_back(nonOwning.getOperation());
+ return success();
}
+//===----------------------------------------------------------------------===//
+// TransformDialect
+//===----------------------------------------------------------------------===//
+
void transform::TransformDialect::initialize() {
// Using the checked versions to enable the same assertions as for the ops
// from extensions.
@@ -201,7 +215,7 @@ void transform::TransformDialect::initialize() {
initializeTypes();
auto &blobInterface = addInterface<TransformResourceBlobManagerInterface>();
- addInterface<TransformOpAsmInterface>(blobInterface);
+ addInterface<detail::TransformOpAsmInterface>(blobInterface);
}
Type transform::TransformDialect::parseType(DialectAsmParser &parser) const {
diff --git a/mlir/lib/Dialect/Transform/Transforms/TransformInterpreterPassBase.cpp b/mlir/lib/Dialect/Transform/Transforms/TransformInterpreterPassBase.cpp
index 5f35b6789dc94fe..db6cbbc31c89748 100644
--- a/mlir/lib/Dialect/Transform/Transforms/TransformInterpreterPassBase.cpp
+++ b/mlir/lib/Dialect/Transform/Transforms/TransformInterpreterPassBase.cpp
@@ -268,10 +268,12 @@ static void performOptionalDebugActions(
transform->removeAttr(kTransformDialectTagAttrName);
}
+const static llvm::StringLiteral kDefaultLibraryName =
+ "__transform_interpreter_parsed_library";
+
LogicalResult transform::detail::interpreterBaseRunOnOperationImpl(
Operation *target, StringRef passName,
const std::shared_ptr<OwningOpRef<ModuleOp>> &sharedTransformModule,
- const std::shared_ptr<OwningOpRef<ModuleOp>> &transformLibraryModule,
const RaggedArray<MappedValue> &extraMappings,
const TransformOptions &options,
const Pass::Option<std::string> &transformFileName,
@@ -281,10 +283,21 @@ LogicalResult transform::detail::interpreterBaseRunOnOperationImpl(
StringRef binaryName) {
bool hasSharedTransformModule =
sharedTransformModule && *sharedTransformModule;
- bool hasTransformLibraryModule =
- transformLibraryModule && *transformLibraryModule;
+ auto *dialect =
+ target->getContext()->getLoadedDialect<transform::TransformDialect>();
+ const transform::detail::TransformLibraryManager &libMgr =
+ dialect->libraryManager;
+ SmallVector<ModuleOp, 2> libraries = libMgr.getLibraryModules();
+ bool hasTransformLibraryModule = !libraries.empty();
assert((!hasSharedTransformModule || !hasTransformLibraryModule) &&
"at most one of shared or library transform module can be set");
+ if (libraries.size() > 1) {
+ InFlightDiagnostic diag =
+ emitError(libraries[1]->getLoc())
+ << "more than one library module not currently supported";
+ diag.attachNote(libraries[0].getLoc()) << "first library";
+ return diag;
+ }
// Step 1
// ------
@@ -337,7 +350,7 @@ LogicalResult transform::detail::interpreterBaseRunOnOperationImpl(
}
if (failed(detail::mergeSymbolsInto(
SymbolTable::getNearestSymbolTable(transformRoot),
- transformLibraryModule->get()->clone())))
+ libraries[0]->clone())))
return emitError(transformRoot->getLoc(),
"failed to merge library symbols into transform root");
}
@@ -414,7 +427,6 @@ LogicalResult transform::detail::interpreterBaseInitializeImpl(
MLIRContext *context, StringRef transformFileName,
ArrayRef<std::string> transformLibraryPaths,
std::shared_ptr<OwningOpRef<ModuleOp>> &sharedTransformModule,
- std::shared_ptr<OwningOpRef<ModuleOp>> &transformLibraryModule,
function_ref<std::optional<LogicalResult>(OpBuilder &, Location)>
moduleBuilder) {
auto unknownLoc = UnknownLoc::get(context);
@@ -488,8 +500,8 @@ LogicalResult transform::detail::interpreterBaseInitializeImpl(
}
}
- // Use parsed libaries to resolve symbols in shared transform module or return
- // as separate library module.
+ // Use parsed libraries to resolve symbols in shared transform module or
+ // return as separate library module.
if (sharedTransformModule && *sharedTransformModule) {
if (failed(detail::mergeSymbolsInto(sharedTransformModule->get(),
std::move(mergedParsedLibraries))))
@@ -497,8 +509,12 @@ LogicalResult transform::detail::interpreterBaseInitializeImpl(
<< "failed to merge symbols from library files "
"into shared transform module";
} else {
- transformLibraryModule = std::make_shared<OwningOpRef<ModuleOp>>(
- std::move(mergedParsedLibraries));
+ // TODO: we should list libraries as resources instead of pass options.
+ transform::detail::TransformLibraryManager &libMgr =
+ context->getLoadedDialect<transform::TransformDialect>()
+ ->libraryManager;
+ return libMgr.registerLibraryModule(kDefaultLibraryName,
+ std::move(mergedParsedLibraries));
}
return success();
}
diff --git a/mlir/lib/Dialect/Transform/Transforms/TransformInterpreterUtils.cpp b/mlir/lib/Dialect/Transform/Transforms/TransformInterpreterUtils.cpp
index 1a6ebdd16232e8a..bc3aeee0a62da7c 100644
--- a/mlir/lib/Dialect/Transform/Transforms/TransformInterpreterUtils.cpp
+++ b/mlir/lib/Dialect/Transform/Transforms/TransformInterpreterUtils.cpp
@@ -59,7 +59,7 @@ LogicalResult transform::detail::parseTransformModuleFromFile(
ModuleOp transform::detail::getPreloadedTransformModule(MLIRContext *context) {
auto preloadedLibraryRange =
context->getOrLoadDialect<transform::TransformDialect>()
- ->getLibraryModules();
+ ->libraryManager.getLibraryModules();
if (!preloadedLibraryRange.empty())
return *preloadedLibraryRange.begin();
return ModuleOp();
diff --git a/mlir/test/Dialect/Transform/test-interpreter-external-symbol-decl-dir.mlir b/mlir/test/Dialect/Transform/test-interpreter-external-symbol-decl-dir.mlir
index 8b8254976e9aeec..ce0da8e13ad91c8 100644
--- a/mlir/test/Dialect/Transform/test-interpreter-external-symbol-decl-dir.mlir
+++ b/mlir/test/Dialect/Transform/test-interpreter-external-symbol-decl-dir.mlir
@@ -4,8 +4,9 @@
// RUN: mlir-opt %s --pass-pipeline="builtin.module(test-transform-dialect-interpreter{transform-library-paths=%p%{fs-sep}test-interpreter-library/definitions-self-contained.mlir,%p%{fs-sep}test-interpreter-library/definitions-with-unresolved.mlir})" \
// RUN: --verify-diagnostics --split-input-file | FileCheck %s
-// RUN: mlir-opt %s --pass-pipeline="builtin.module(test-transform-dialect-interpreter{transform-library-paths=%p%{fs-sep}test-interpreter-library}, test-transform-dialect-interpreter)" \
-// RUN: --verify-diagnostics --split-input-file | FileCheck %s
+// TODO: reenable when library loading happens in a separate pass.
+// _UN: mlir-opt %s --pass-pipeline="builtin.module(test-transform-dialect-interpreter{transform-library-paths=%p%{fs-sep}test-interpreter-library}, test-transform-dialect-interpreter)" \
+// _UN: --verify-diagnostics --split-input-file | 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
diff --git a/mlir/test/Dialect/Transform/test-interpreter-external-symbol-decl.mlir b/mlir/test/Dialect/Transform/test-interpreter-external-symbol-decl.mlir
index 339e62072cd5510..0400b52e8986428 100644
--- a/mlir/test/Dialect/Transform/test-interpreter-external-symbol-decl.mlir
+++ b/mlir/test/Dialect/Transform/test-interpreter-external-symbol-decl.mlir
@@ -1,13 +1,14 @@
// RUN: mlir-opt %s --pass-pipeline="builtin.module(test-transform-dialect-interpreter{transform-library-paths=%p/test-interpreter-library/definitions-self-contained.mlir})" \
// RUN: --verify-diagnostics --split-input-file | FileCheck %s
-// RUN: mlir-opt %s --pass-pipeline="builtin.module(test-transform-dialect-interpreter{transform-library-paths=%p/test-interpreter-library/definitions-self-contained.mlir}, test-transform-dialect-interpreter)" \
-// RUN: --verify-diagnostics --split-input-file | FileCheck %s
+// TODO: reenable when library loading happens in a separate pass.
+// _UN: mlir-opt %s --pass-pipeline="builtin.module(test-transform-dialect-interpreter{transform-library-paths=%p/test-interpreter-library/definitions-self-contained.mlir}, test-transform-dialect-interpreter)" \
+// _UN: --verify-diagnostics --split-input-file | FileCheck %s
// The definition of the @print_message named sequence is provided in another
// file. It will be included because of the pass option. Subsequent application
// of the same pass works but only without the library file (since the first
-// application loads external symbols and loading them again woul make them
+// application loads external symbols and loading them again would make them
// clash).
// Note that the same diagnostic produced twice at the same location only
// needs to be matched once.
diff --git a/mlir/test/Dialect/Transform/test-interpreter-external.mlir b/mlir/test/Dialect/Transform/test-interpreter-external.mlir
index bf0e03037f535d3..c3fef37a1aa109d 100644
--- a/mlir/test/Dialect/Transform/test-interpreter-external.mlir
+++ b/mlir/test/Dialect/Transform/test-interpreter-external.mlir
@@ -5,17 +5,4 @@
// expected-remark @below {{outer}}
// expected-remark @below {{inner}}
-module attributes { test.blob_ref = #test.e1di64_elements<library> : tensor<*xi1>} {
- module attributes { test.blob_ref = #test.e1di64_elements<library> : tensor<*xi1>} {}
- module {}
- module {}
-}
-
-{-#
- dialect_resources: {
- transform: {
- library: "test-interpreter-external-source.mlir",
- banana: "test-interpreter-external-source.mlir"
- }
- }
-#-}
\ No newline at end of file
+module {}
diff --git a/mlir/test/Dialect/Transform/test-interpreter-resource-external-def.mlir b/mlir/test/Dialect/Transform/test-interpreter-resource-external-def.mlir
new file mode 100644
index 000000000000000..e1052e616d43958
--- /dev/null
+++ b/mlir/test/Dialect/Transform/test-interpreter-resource-external-def.mlir
@@ -0,0 +1,16 @@
+// RUN: cd %p; mlir-opt %s
+
+module attributes { test.blob_ref = #test.e1di64_elements<library> : tensor<*xi1>} {
+ module attributes { test.blob_ref = #test.e1di64_elements<library> : tensor<*xi1>} {}
+ module {}
+ module {}
+}
+
+{-#
+ dialect_resources: {
+ transform: {
+ library: "test-interpreter-external-source.mlir",
+ banana: "test-interpreter-external-source.mlir"
+ }
+ }
+#-}
diff --git a/mlir/test/lib/Dialect/Transform/TestTransformDialectInterpreter.cpp b/mlir/test/lib/Dialect/Transform/TestTransformDialectInterpreter.cpp
index c60b21c918338b4..416b6f3f94825bd 100644
--- a/mlir/test/lib/Dialect/Transform/TestTransformDialectInterpreter.cpp
+++ b/mlir/test/lib/Dialect/Transform/TestTransformDialectInterpreter.cpp
@@ -160,9 +160,8 @@ class TestTransformDialectInterpreterPass
options = options.enableExpensiveChecks(enableExpensiveChecks);
if (failed(transform::detail::interpreterBaseRunOnOperationImpl(
getOperation(), getArgument(), getSharedTransformModule(),
- getTransformLibraryModule(), extraMapping, options,
- transformFileName, transformLibraryPaths, debugPayloadRootTag,
- debugTransformRootTag, getBinaryName())))
+ extraMapping, options, transformFileName, transformLibraryPaths,
+ debugPayloadRootTag, debugTransformRootTag, getBinaryName())))
return signalPassFailure();
}
diff --git a/mlir/unittests/Dialect/Transform/Preload.cpp b/mlir/unittests/Dialect/Transform/Preload.cpp
index d3c3044e0e0f776..6a6170d899407c1 100644
--- a/mlir/unittests/Dialect/Transform/Preload.cpp
+++ b/mlir/unittests/Dialect/Transform/Preload.cpp
@@ -67,7 +67,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 result = dialect->libraryManager.registerLibraryModule(
+ "<main>", std::move(transformLibrary));
+ EXPECT_TRUE(succeeded(result));
ModuleOp retrievedTransformLibrary =
transform::detail::getPreloadedTransformModule(&context);
More information about the Mlir-commits
mailing list