[Mlir-commits] [mlir] [mlir][IR] Add transient scope support for resettable MLIRContext. (PR #217320)
Jacques Pienaar
llvmlistbot at llvm.org
Wed Aug 19 05:18:35 PDT 2026
https://github.com/jpienaar created https://github.com/llvm/llvm-project/pull/217320
We have cases where we have many context being used to compile independent modules. Creatign a context per ends up being rather expensive and the only reason to not reuse contexts more is bloat over time. This changes adds a "transient" scope which can be used for such cases. Initially I had tried a "cheap fork" approach for contexts, this is complimentary to that (that one allows for independent contexts cheaply), while this one is less invasive.
Introduce an overlay/layered state architecture in StorageUniquer and MLIRContext to enable efficient scoping and rollback of transient types, attributes, affine expressions, and distinct attributes added since entering a transient scope.
When executing repeated compilation passes or running in long-lived compiler services (e.g. JITs, servers, REPLs), allocating and destroying an MLIRContext incurs significant overhead (~3.7 ms per context) to re-register and load dialects. This change allows freezing an initialized context as a base state and rolling back transient types/attributes in O(1) time without re-instantiating dialects or builtins.
Assisted-By: Gemini
>From 8634b7c4e954190a9e4b0d4c63140234b81affd4 Mon Sep 17 00:00:00 2001
From: Jacques Pienaar <jpienaar at google.com>
Date: Wed, 19 Aug 2026 12:17:14 +0000
Subject: [PATCH] [mlir][IR] Add transient scope support for resettable
MLIRContext.
We have cases where we have many context being used to compile
independent modules. Creatign a context per ends up being rather
expensive and the only reason to not reuse contexts more is bloat over
time. This changes adds a "transient" scope which can be used for such
cases. Initially I had tried a "cheap fork" approach for contexts, this
is complimentary to that (that one allows for independent contexts
cheaply), while this one is less invasive.
Introduce an overlay/layered state architecture in StorageUniquer and
MLIRContext to enable efficient scoping and rollback of transient types,
attributes, affine expressions, and distinct attributes added since
entering a transient scope.
When executing repeated compilation passes or running in long-lived
compiler services (e.g. JITs, servers, REPLs), allocating and destroying
an MLIRContext incurs significant overhead (~3.7 ms per context) to
re-register and load dialects. This change allows freezing an
initialized context as a base state and rolling back transient
types/attributes in O(1) time without re-instantiating dialects or
builtins.
Assisted-By: Gemini
---
mlir/include/mlir-c/IR.h | 14 ++
mlir/include/mlir/IR/MLIRContext.h | 43 ++++
mlir/include/mlir/Support/StorageUniquer.h | 19 ++
mlir/include/mlir/Support/ThreadLocalCache.h | 3 +
mlir/lib/Bindings/Python/IRCore.cpp | 32 ++-
mlir/lib/CAPI/IR/IR.cpp | 12 +
mlir/lib/IR/AttributeDetail.h | 37 ++-
mlir/lib/IR/MLIRContext.cpp | 107 +++++++++
mlir/lib/Support/StorageUniquer.cpp | 212 +++++++++++++++++-
mlir/python/mlir/_mlir_libs/__init__.py | 16 ++
mlir/test/CAPI/ir.c | 50 +++++
mlir/test/python/ir/context_lifecycle.py | 68 ++++++
mlir/test/python/ir/context_managers.py | 29 +++
mlir/unittests/IR/MLIRContextResetTest.cpp | 192 ++++++++++++++++
mlir/unittests/Support/StorageUniquerTest.cpp | 100 +++++++++
15 files changed, 917 insertions(+), 17 deletions(-)
create mode 100644 mlir/unittests/IR/MLIRContextResetTest.cpp
diff --git a/mlir/include/mlir-c/IR.h b/mlir/include/mlir-c/IR.h
index ac8ffbaf028b0..866d90e621384 100644
--- a/mlir/include/mlir-c/IR.h
+++ b/mlir/include/mlir-c/IR.h
@@ -179,6 +179,20 @@ MLIR_CAPI_EXPORTED unsigned mlirContextGetNumThreads(MlirContext context);
MLIR_CAPI_EXPORTED MlirLlvmThreadPool
mlirContextGetThreadPool(MlirContext context);
+/// Begins a transient scope on the context, freezing the base layer (loaded
+/// dialects, registered operations, interface models, and existing
+/// types/attributes).
+/// Precondition: The context must not already be in a transient scope.
+MLIR_CAPI_EXPORTED void mlirContextBeginTransientScope(MlirContext context);
+
+/// Ends the transient scope and resets the context to the base state, pruning
+/// transient types, attributes, affine expressions, distinct attributes, and
+/// unregistered operations added during the transient scope.
+MLIR_CAPI_EXPORTED void mlirContextEndTransientScope(MlirContext context);
+
+/// Returns whether the context is currently in a transient scope.
+MLIR_CAPI_EXPORTED bool mlirContextIsInTransientScope(MlirContext context);
+
//===----------------------------------------------------------------------===//
// Dialect API.
//===----------------------------------------------------------------------===//
diff --git a/mlir/include/mlir/IR/MLIRContext.h b/mlir/include/mlir/IR/MLIRContext.h
index 12cd00b60215e..ce195518066d9 100644
--- a/mlir/include/mlir/IR/MLIRContext.h
+++ b/mlir/include/mlir/IR/MLIRContext.h
@@ -145,6 +145,49 @@ class MLIRContext {
/// operations) without being caught by assertions or other means.
void allowUnregisteredDialects(bool allow = true);
+ /// Begins a transient scope on the context, freezing the current state (all
+ /// loaded dialects, registered operations, types, attributes, affine
+ /// expressions, and singletons) as the base state. Subsequent types,
+ /// attributes, and expressions allocated will belong to the transient layer.
+ ///
+ /// Preconditions:
+ /// - The context must not already be in a transient scope.
+ /// - Must be called from a single-threaded execution context.
+ /// - Loading dialects, modifying dialect registries, or mutating base
+ /// storage instances is not supported while in a transient scope.
+ void beginTransientScope();
+
+ /// Ends the transient scope and resets the context to the base state, pruning
+ /// all types, attributes, affine expressions, distinct attributes, and
+ /// unregistered operations created during the transient scope.
+ ///
+ /// Preconditions:
+ /// - The context must be in a transient scope.
+ /// - There must be no remaining IR (operations, blocks, regions,
+ /// values) referencing the transient types/attributes.
+ /// - Must be called from a single-threaded execution context.
+ void endTransientScope();
+
+ /// Returns true if the context is currently in a transient scope.
+ bool isInTransientScope() const;
+
+ /// RAII scope guard that calls beginTransientScope() on construction and
+ /// endTransientScope() on destruction. Provides exception-safe and
+ /// forgetting-proof transient scope management.
+ class TransientScope {
+ public:
+ explicit TransientScope(MLIRContext &ctx) : ctx(ctx) {
+ ctx.beginTransientScope();
+ }
+ ~TransientScope() { ctx.endTransientScope(); }
+
+ TransientScope(const TransientScope &) = delete;
+ TransientScope &operator=(const TransientScope &) = delete;
+
+ private:
+ MLIRContext &ctx;
+ };
+
/// Return true if multi-threading is enabled by the context.
bool isMultithreadingEnabled();
diff --git a/mlir/include/mlir/Support/StorageUniquer.h b/mlir/include/mlir/Support/StorageUniquer.h
index b5f4df680ac69..39018ebba106f 100644
--- a/mlir/include/mlir/Support/StorageUniquer.h
+++ b/mlir/include/mlir/Support/StorageUniquer.h
@@ -246,8 +246,27 @@ class StorageUniquer {
/// is initialized when a dialect is loaded.
bool isParametricStorageInitialized(TypeID id);
+ /// Begins a transient scope. All subsequent allocations and uniquing
+ /// registrations will be recorded in a transient layer.
+ /// Precondition: The uniquer must not already be in a transient scope.
+ void beginTransientScope();
+
+ /// Ends the transient scope and resets back to the base state, freeing
+ /// all transiently allocated storage instances and clearing transient lookup
+ /// tables.
+ /// Precondition: The uniquer must be in a transient scope.
+ void endTransientScope();
+
+ /// Returns true if the uniquer is currently in a transient scope.
+ bool isInTransientScope() const;
+
/// Changes the mutable component of 'storage' by forwarding the trailing
/// arguments to the 'mutate' function of the derived class.
+ /// Note: When in a transient scope, any storage memory allocated during
+ /// mutation will come from the transient allocator. Mutating base storage
+ /// instances (allocated before the transient scope) while in a transient
+ /// scope is not supported as that transient memory will be freed on scope
+ /// exit.
template <typename Storage, typename... Args>
LogicalResult mutate(TypeID id, Storage *storage, Args &&...args) {
auto mutationFn = [&](StorageAllocator &allocator) -> LogicalResult {
diff --git a/mlir/include/mlir/Support/ThreadLocalCache.h b/mlir/include/mlir/Support/ThreadLocalCache.h
index a1f88ec609e86..66041638b89a2 100644
--- a/mlir/include/mlir/Support/ThreadLocalCache.h
+++ b/mlir/include/mlir/Support/ThreadLocalCache.h
@@ -140,6 +140,9 @@ class ThreadLocalCache {
// scope and invalidate the weak pointers held by the thread_local caches.
}
+ /// Invalidate and clear all thread-local cache entries for this instance.
+ void clear() { perInstanceState = std::make_shared<PerInstanceState>(); }
+
/// Return an instance of the value type for the current thread.
ValueT &get() {
// Check for an already existing instance for this thread.
diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index 75cfd2a0a1c0b..f80114c578a20 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -3536,7 +3536,37 @@ void populateIRCore(nb::module_ &m) {
Loads all dialects available in the registry into the context.
This eagerly loads all dialects that have been registered, making them
- immediately available for use.)");
+ immediately available for use.)")
+ .def(
+ "begin_transient_scope",
+ [](PyMlirContext &self) {
+ if (mlirContextIsInTransientScope(self.get()))
+ throw nb::value_error("Context is already in a transient scope");
+ mlirContextBeginTransientScope(self.get());
+ },
+ R"(
+ Begins a transient scope on the context, freezing the base layer.
+
+ All subsequently allocated types, attributes, and unregistered operations
+ are treated as transient and will be deallocated with end_transient_scope().
+ Raises a ValueError if the context is already in a transient scope.)")
+ .def(
+ "end_transient_scope",
+ [](PyMlirContext &self) { mlirContextEndTransientScope(self.get()); },
+ R"(
+ Ends the transient scope and resets the context to the base state.
+
+ Prunes all transient types, attributes, affine expressions, distinct
+ attributes, and unregistered operations added during the transient scope.
+
+ Note: Any Python objects referencing transient IR entities become invalid
+ after this call and must not be accessed.)")
+ .def_prop_ro(
+ "is_in_transient_scope",
+ [](PyMlirContext &self) -> bool {
+ return mlirContextIsInTransientScope(self.get());
+ },
+ "Returns whether the context is currently in a transient scope.");
//----------------------------------------------------------------------------
// Mapping of PyDialectDescriptor
diff --git a/mlir/lib/CAPI/IR/IR.cpp b/mlir/lib/CAPI/IR/IR.cpp
index ef730b26cdd5d..f862efa420daa 100644
--- a/mlir/lib/CAPI/IR/IR.cpp
+++ b/mlir/lib/CAPI/IR/IR.cpp
@@ -127,6 +127,18 @@ MlirLlvmThreadPool mlirContextGetThreadPool(MlirContext context) {
return wrap(&unwrap(context)->getThreadPool());
}
+void mlirContextBeginTransientScope(MlirContext context) {
+ unwrap(context)->beginTransientScope();
+}
+
+void mlirContextEndTransientScope(MlirContext context) {
+ unwrap(context)->endTransientScope();
+}
+
+bool mlirContextIsInTransientScope(MlirContext context) {
+ return unwrap(context)->isInTransientScope();
+}
+
//===----------------------------------------------------------------------===//
// Dialect API.
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/IR/AttributeDetail.h b/mlir/lib/IR/AttributeDetail.h
index 4b7e7dd5677e8..58a9106c65707 100644
--- a/mlir/lib/IR/AttributeDetail.h
+++ b/mlir/lib/IR/AttributeDetail.h
@@ -22,6 +22,7 @@
#include "mlir/IR/MLIRContext.h"
#include "llvm/ADT/APFloat.h"
#include "llvm/Support/Allocator.h"
+#include "llvm/Support/RWMutex.h"
#include <mutex>
namespace mlir {
@@ -321,18 +322,40 @@ class DistinctAttributeAllocator final {
operator=(const DistinctAttributeAllocator &) = delete;
DistinctAttrStorage *allocate(Attribute referencedAttr) {
- std::scoped_lock<std::mutex> guard(allocatorMutex);
- return new (allocator.Allocate<DistinctAttrStorage>())
+ llvm::sys::SmartScopedWriter<true> guard(allocatorMutex);
+ llvm::BumpPtrAllocator &alloc =
+ transientAllocator ? *transientAllocator : allocator;
+ return new (alloc.Allocate<DistinctAttrStorage>())
DistinctAttrStorage(referencedAttr);
- };
+ }
+
+ void beginTransientScope() {
+ llvm::sys::SmartScopedWriter<true> guard(allocatorMutex);
+ assert(!transientAllocator &&
+ "distinct attribute allocator is already in a transient scope");
+ transientAllocator = std::make_unique<llvm::BumpPtrAllocator>();
+ }
+
+ void endTransientScope() {
+ llvm::sys::SmartScopedWriter<true> guard(allocatorMutex);
+ transientAllocator.reset();
+ }
+
+ bool isInTransientScope() const {
+ llvm::sys::SmartScopedReader<true> guard(allocatorMutex);
+ return transientAllocator != nullptr;
+ }
private:
- /// Used to allocate distict attribute storages. The managed memory is freed
- /// automatically when the allocator instance is destroyed.
+ /// Used to allocate distinct attribute storages in base layer.
llvm::BumpPtrAllocator allocator;
- /// Used to lock access to the allocator.
- std::mutex allocatorMutex;
+ /// Used to allocate distinct attribute storages in transient layer.
+ std::unique_ptr<llvm::BumpPtrAllocator> transientAllocator;
+
+ /// Used to synchronize access to the allocator. Uses a RW mutex so that
+ /// isInTransientScope() can take a cheap reader lock.
+ mutable llvm::sys::SmartRWMutex<true> allocatorMutex;
};
} // namespace detail
} // namespace mlir
diff --git a/mlir/lib/IR/MLIRContext.cpp b/mlir/lib/IR/MLIRContext.cpp
index da891a7e6e014..06da14504a293 100644
--- a/mlir/lib/IR/MLIRContext.cpp
+++ b/mlir/lib/IR/MLIRContext.cpp
@@ -269,6 +269,17 @@ class MLIRContextImpl {
/// destruction.
DistinctAttributeAllocator distinctAttributeAllocator;
+ /// Bundled state dynamically allocated when in a transient scope.
+ struct TransientScopeState {
+ /// Set of operation names in `operations` at snapshot time.
+ llvm::DenseSet<StringRef> baseOperations;
+
+ /// Number of entries in `dialectReferencingStrAttrs` per dialect at
+ /// snapshot time.
+ llvm::DenseMap<StringRef, size_t> baseDialectReferencingStrAttrCounts;
+ };
+ std::unique_ptr<TransientScopeState> transientState;
+
public:
MLIRContextImpl(bool threadingIsEnabled)
: threadingIsEnabled(threadingIsEnabled) {
@@ -428,6 +439,9 @@ void MLIRContext::appendDialectRegistry(const DialectRegistry ®istry) {
assert(impl->multiThreadedExecutionContext == 0 &&
"appending to the MLIRContext dialect registry while in a "
"multi-threaded execution context");
+ assert(!impl->transientState &&
+ "cannot append to dialect registry while in a transient scope");
+
registry.appendTo(impl->dialectsRegistry);
// For the already loaded dialects, apply any possible extensions immediately.
@@ -486,6 +500,8 @@ MLIRContext::getOrLoadDialect(StringRef dialectNamespace, TypeID dialectID,
if (dialectIt.second) {
LDBG() << "Load new dialect in Context " << dialectNamespace;
+ assert(!impl.transientState &&
+ "cannot load new dialects while in a transient scope");
#ifndef NDEBUG
if (impl.multiThreadedExecutionContext != 0)
llvm::report_fatal_error(
@@ -679,6 +695,97 @@ void MLIRContext::exitMultiThreadedExecution() {
#endif
}
+void MLIRContext::beginTransientScope() {
+ MLIRContextImpl &ctxImpl = getImpl();
+ assert(ctxImpl.multiThreadedExecutionContext == 0 &&
+ "Beginning a transient scope while in a multi-threaded execution "
+ "context");
+ assert(!ctxImpl.transientState && "context is already in a transient scope");
+ ctxImpl.transientState =
+ std::make_unique<MLIRContextImpl::TransientScopeState>();
+
+ // Begin transient scope in the uniquers.
+ ctxImpl.typeUniquer.beginTransientScope();
+ ctxImpl.attributeUniquer.beginTransientScope();
+ ctxImpl.affineUniquer.beginTransientScope();
+ ctxImpl.distinctAttributeAllocator.beginTransientScope();
+
+ // Record base operations in operations map.
+ {
+ llvm::sys::SmartScopedReader<true> contextLock(ctxImpl.operationInfoMutex);
+ for (const auto &entry : ctxImpl.operations)
+ ctxImpl.transientState->baseOperations.insert(entry.first());
+ }
+
+ // Record dialect referencing string attribute counts.
+ {
+ llvm::sys::SmartScopedLock<true> lock(ctxImpl.dialectRefStrAttrMutex);
+ for (const auto &entry : ctxImpl.dialectReferencingStrAttrs)
+ ctxImpl.transientState->baseDialectReferencingStrAttrCounts[entry.first] =
+ entry.second.size();
+ }
+}
+
+void MLIRContext::endTransientScope() {
+ MLIRContextImpl &ctxImpl = getImpl();
+ assert(ctxImpl.transientState && "context is not in a transient scope");
+ assert(ctxImpl.multiThreadedExecutionContext == 0 &&
+ "Ending a transient scope while in a multi-threaded execution "
+ "context");
+ if (!ctxImpl.transientState)
+ return;
+
+ // Prune unregistered operations created during transient scope before
+ // destroying the attribute uniquer that holds their string attribute names.
+ {
+ llvm::sys::SmartScopedWriter<true> contextLock(ctxImpl.operationInfoMutex);
+ SmallVector<StringRef> opsToErase;
+ for (const auto &entry : ctxImpl.operations) {
+ if (!entry.second->isRegistered() &&
+ !ctxImpl.transientState->baseOperations.contains(entry.first()))
+ opsToErase.push_back(entry.first());
+ }
+ for (StringRef op : opsToErase)
+ ctxImpl.operations.erase(op);
+ }
+
+ // Restore dialect referencing string attributes before destroying the
+ // attribute uniquer that holds the underlying StringAttrStorage pointers.
+ // Note: Transient entries are always appended to the end of each dialect's
+ // vector, so truncating via resize() to the base count safely removes only
+ // transient entries while preserving base entries.
+ {
+ llvm::sys::SmartScopedLock<true> lock(ctxImpl.dialectRefStrAttrMutex);
+ SmallVector<StringRef> dialectsToErase;
+ for (auto &entry : ctxImpl.dialectReferencingStrAttrs) {
+ auto countIt =
+ ctxImpl.transientState->baseDialectReferencingStrAttrCounts.find(
+ entry.first);
+ if (countIt ==
+ ctxImpl.transientState->baseDialectReferencingStrAttrCounts.end()) {
+ dialectsToErase.push_back(entry.first);
+ } else {
+ entry.second.resize(countIt->second);
+ }
+ }
+ for (StringRef dialect : dialectsToErase)
+ ctxImpl.dialectReferencingStrAttrs.erase(dialect);
+ }
+
+ // End transient scope in the uniquers now that all referencing structures
+ // are cleaned up.
+ ctxImpl.typeUniquer.endTransientScope();
+ ctxImpl.attributeUniquer.endTransientScope();
+ ctxImpl.affineUniquer.endTransientScope();
+ ctxImpl.distinctAttributeAllocator.endTransientScope();
+
+ ctxImpl.transientState.reset();
+}
+
+bool MLIRContext::isInTransientScope() const {
+ return getImpl().transientState != nullptr;
+}
+
/// Return true if we should attach the operation to diagnostics emitted via
/// Operation::emit.
bool MLIRContext::shouldPrintOpOnDiagnostic() {
diff --git a/mlir/lib/Support/StorageUniquer.cpp b/mlir/lib/Support/StorageUniquer.cpp
index ef7f19b9918d9..dd9de846f9d76 100644
--- a/mlir/lib/Support/StorageUniquer.cpp
+++ b/mlir/lib/Support/StorageUniquer.cpp
@@ -11,6 +11,7 @@
#include "mlir/Support/LLVM.h"
#include "mlir/Support/ThreadLocalCache.h"
#include "mlir/Support/TypeID.h"
+#include "llvm/ADT/Bitfields.h"
#include "llvm/Support/RWMutex.h"
using namespace mlir;
@@ -66,13 +67,20 @@ class ParametricStorageUniquer {
};
using StorageTypeSet = DenseSet<HashedStorage, StorageKeyInfo>;
+ using InTransientScope = llvm::Bitfield::Element<bool, 0, 1>;
+ using NumShards = llvm::Bitfield::Element<unsigned, 1, 31>;
+
/// This class represents a single shard of the uniquer. The uniquer uses a
/// set of shards to allow for multiple threads to create instances with less
/// lock contention.
struct Shard {
- /// The set containing the allocated storage instances.
+ /// The set containing the allocated storage instances in the base layer.
StorageTypeSet instances;
+ /// The set containing the allocated storage instances in the transient
+ /// layer, lazily instantiated on the first transient allocation.
+ std::unique_ptr<StorageTypeSet> transientInstances;
+
#if LLVM_ENABLE_THREADS != 0
/// A mutex to keep uniquing thread-safe.
llvm::sys::SmartRWMutex<true> mutex;
@@ -83,6 +91,24 @@ class ParametricStorageUniquer {
/// fashion.
BaseStorage *getOrCreateUnsafe(Shard &shard, LookupKey &key,
function_ref<BaseStorage *()> ctorFn) {
+ if (isInTransientScope()) {
+ if (shard.transientInstances) {
+ auto transientIt = shard.transientInstances->find_as(key);
+ if (transientIt != shard.transientInstances->end())
+ return transientIt->storage;
+ }
+ auto baseIt = shard.instances.find_as(key);
+ if (baseIt != shard.instances.end())
+ return baseIt->storage;
+ if (!shard.transientInstances)
+ shard.transientInstances = std::make_unique<StorageTypeSet>();
+ auto existing = shard.transientInstances->insert_as({key.hashValue}, key);
+ BaseStorage *&storage = existing.first->storage;
+ if (existing.second)
+ storage = ctorFn();
+ return storage;
+ }
+
auto existing = shard.instances.insert_as({key.hashValue}, key);
BaseStorage *&storage = existing.first->storage;
if (existing.second)
@@ -96,6 +122,10 @@ class ParametricStorageUniquer {
return;
for (HashedStorage &instance : shard.instances)
destructorFn(instance.storage);
+ if (shard.transientInstances) {
+ for (HashedStorage &instance : *shard.transientInstances)
+ destructorFn(instance.storage);
+ }
}
public:
@@ -105,15 +135,17 @@ class ParametricStorageUniquer {
/// destructor function is used to destroy any allocated storage instances.
ParametricStorageUniquer(function_ref<void(BaseStorage *)> destructorFn,
size_t numShards = 8)
- : shards(new std::atomic<Shard *>[numShards]), numShards(numShards),
+ : shards(new std::atomic<Shard *>[numShards]),
destructorFn(destructorFn) {
assert(llvm::isPowerOf2_64(numShards) &&
"the number of shards is required to be a power of 2");
+ llvm::Bitfield::set<NumShards>(numShardsAndScope, numShards);
for (size_t i = 0; i < numShards; i++)
shards[i].store(nullptr, std::memory_order_relaxed);
}
~ParametricStorageUniquer() {
// Free all of the allocated shards.
+ size_t numShards = getNumShards();
for (size_t i = 0; i != numShards; ++i) {
if (Shard *shard = shards[i].load()) {
destroyShardInstances(*shard);
@@ -139,6 +171,11 @@ class ParametricStorageUniquer {
// Check for an existing instance in read-only mode.
{
llvm::sys::SmartScopedReader<true> typeLock(shard.mutex);
+ if (isInTransientScope() && shard.transientInstances) {
+ auto it = shard.transientInstances->find_as(lookupKey);
+ if (it != shard.transientInstances->end())
+ return localInst = it->storage;
+ }
auto it = shard.instances.find_as(lookupKey);
if (it != shard.instances.end())
return localInst = it->storage;
@@ -165,11 +202,48 @@ class ParametricStorageUniquer {
return mutationFn();
}
+ void beginTransientScope() {
+ assert(!isInTransientScope() &&
+ "parametric storage uniquer is already in a transient scope");
+ llvm::Bitfield::set<InTransientScope>(numShardsAndScope, true);
+ }
+
+ void endTransientScope() {
+ assert(isInTransientScope() &&
+ "parametric storage uniquer is not in a transient scope");
+ if (!isInTransientScope())
+ return;
+
+ size_t numShards = getNumShards();
+ for (size_t i = 0; i != numShards; ++i) {
+ if (Shard *shard = shards[i].load()) {
+ llvm::sys::SmartScopedWriter<true> typeLock(shard->mutex);
+ if (shard->transientInstances) {
+ if (destructorFn) {
+ for (HashedStorage &instance : *shard->transientInstances)
+ destructorFn(instance.storage);
+ }
+ shard->transientInstances.reset();
+ }
+ }
+ }
+ localCache.clear();
+ llvm::Bitfield::set<InTransientScope>(numShardsAndScope, false);
+ }
+
+ bool isInTransientScope() const {
+ return llvm::Bitfield::get<InTransientScope>(numShardsAndScope);
+ }
+
+ size_t getNumShards() const {
+ return llvm::Bitfield::get<NumShards>(numShardsAndScope);
+ }
+
private:
/// Return the shard used for the given hash value.
Shard &getShard(unsigned hashValue) {
// Get a shard number from the provided hashvalue.
- unsigned shardNum = hashValue & (numShards - 1);
+ unsigned shardNum = hashValue & (getNumShards() - 1);
// Try to acquire an already initialized shard.
Shard *shard = shards[shardNum].load(std::memory_order_acquire);
@@ -195,8 +269,9 @@ class ParametricStorageUniquer {
/// the overhead when only a small amount of shards are in use.
std::unique_ptr<std::atomic<Shard *>[]> shards;
- /// The number of available shards.
- size_t numShards;
+ /// Packed field storing numShards in bits 1..31 and inTransientScope in bit
+ /// 0.
+ uint32_t numShardsAndScope = 0;
/// Function to used to destruct any allocated storage instances.
function_ref<void(BaseStorage *)> destructorFn;
@@ -226,12 +301,38 @@ class ParametricStorageUniquer {
return mutationFn();
}
+ void beginTransientScope() {
+ assert(!inTransientScope &&
+ "parametric storage uniquer is already in a transient scope");
+ inTransientScope = true;
+ }
+
+ void endTransientScope() {
+ assert(inTransientScope &&
+ "parametric storage uniquer is not in a transient scope");
+ if (!inTransientScope)
+ return;
+ if (shard.transientInstances) {
+ if (destructorFn) {
+ for (HashedStorage &instance : *shard.transientInstances)
+ destructorFn(instance.storage);
+ }
+ shard.transientInstances.reset();
+ }
+ inTransientScope = false;
+ }
+
+ bool isInTransientScope() const { return inTransientScope; }
+
private:
/// The main uniquer shard that is used for allocating storage instances.
Shard shard;
/// Function to used to destruct any allocated storage instances.
function_ref<void(BaseStorage *)> destructorFn;
+
+ /// Flag indicating if the uniquer is currently in a transient scope.
+ bool inTransientScope = false;
#endif
};
} // namespace
@@ -243,6 +344,26 @@ struct StorageUniquerImpl {
using BaseStorage = StorageUniquer::BaseStorage;
using StorageAllocator = StorageUniquer::StorageAllocator;
+ /// Bundled state dynamically allocated when entering a transient scope.
+ struct TransientState {
+#if LLVM_ENABLE_THREADS != 0
+ /// Transient thread local set of allocators used when in a transient scope.
+ ThreadLocalCache<StorageAllocator *> threadSafeAllocator;
+
+ /// Transient allocators created during transient scope.
+ std::vector<std::unique_ptr<StorageAllocator>> threadAllocators;
+
+ /// A mutex used for safely adding a new transient thread allocator.
+ llvm::sys::SmartMutex<true> threadAllocatorMutex;
+#endif
+
+ /// Single-threaded allocator used during transient scope.
+ std::unique_ptr<StorageAllocator> allocator;
+
+ /// Transient singleton instances registered during transient scope.
+ DenseMap<TypeID, BaseStorage *> singletonInstances;
+ };
+
//===--------------------------------------------------------------------===//
// Parametric Storage
//===--------------------------------------------------------------------===//
@@ -280,8 +401,27 @@ struct StorageUniquerImpl {
/// current thread.
StorageAllocator &getThreadSafeAllocator() {
#if LLVM_ENABLE_THREADS != 0
- if (!threadingIsEnabled)
+ if (!threadingIsEnabled) {
+ if (transientState) {
+ if (!transientState->allocator)
+ transientState->allocator = std::make_unique<StorageAllocator>();
+ return *transientState->allocator;
+ }
return allocator;
+ }
+
+ if (transientState) {
+ StorageAllocator *&threadAllocator =
+ transientState->threadSafeAllocator.get();
+ if (!threadAllocator) {
+ threadAllocator = new StorageAllocator();
+ llvm::sys::SmartScopedLock<true> lock(
+ transientState->threadAllocatorMutex);
+ transientState->threadAllocators.push_back(
+ std::unique_ptr<StorageAllocator>(threadAllocator));
+ }
+ return *threadAllocator;
+ }
// If the allocator has not been initialized, create a new one.
StorageAllocator *&threadAllocator = threadSafeAllocator.get();
@@ -297,23 +437,56 @@ struct StorageUniquerImpl {
return *threadAllocator;
#else
+ if (transientState) {
+ if (!transientState->allocator)
+ transientState->allocator = std::make_unique<StorageAllocator>();
+ return *transientState->allocator;
+ }
return allocator;
#endif
}
+ void beginTransientScope() {
+ assert(!transientState &&
+ "storage uniquer is already in a transient scope");
+ transientState = std::make_unique<TransientState>();
+ for (auto &entry : parametricUniquers)
+ entry.second->beginTransientScope();
+ }
+
+ void endTransientScope() {
+ assert(transientState && "storage uniquer is not in a transient scope");
+ if (!transientState)
+ return;
+ for (auto &entry : parametricUniquers)
+ entry.second->endTransientScope();
+ transientState.reset();
+ }
+
+ bool isInTransientScope() const { return transientState != nullptr; }
+
//===--------------------------------------------------------------------===//
// Singleton Storage
//===--------------------------------------------------------------------===//
/// Get or create an instance of a singleton storage class.
BaseStorage *getSingleton(TypeID id) {
+ if (transientState) {
+ auto it = transientState->singletonInstances.find(id);
+ if (it != transientState->singletonInstances.end())
+ return it->second;
+ }
BaseStorage *singletonInstance = singletonInstances[id];
assert(singletonInstance && "expected singleton instance to exist");
return singletonInstance;
}
/// Check if an instance of a singleton storage class exists.
- bool hasSingleton(TypeID id) const { return singletonInstances.count(id); }
+ bool hasSingleton(TypeID id) const {
+ if (transientState && transientState->singletonInstances.count(id))
+ return true;
+ return singletonInstances.count(id);
+ }
//===--------------------------------------------------------------------===//
// Instance Storage
@@ -335,6 +508,9 @@ struct StorageUniquerImpl {
/// thread safety is guaranteed.
StorageAllocator allocator;
+ /// Transient state bundled into a unique pointer (nullptr when inactive).
+ std::unique_ptr<TransientState> transientState;
+
/// Map of type ids to the storage uniquer to use for registered objects.
DenseMap<TypeID, std::unique_ptr<ParametricStorageUniquer>>
parametricUniquers;
@@ -357,6 +533,14 @@ void StorageUniquer::disableMultithreading(bool disable) {
impl->threadingIsEnabled = !disable;
}
+void StorageUniquer::beginTransientScope() { impl->beginTransientScope(); }
+
+void StorageUniquer::endTransientScope() { impl->endTransientScope(); }
+
+bool StorageUniquer::isInTransientScope() const {
+ return impl->isInTransientScope();
+}
+
/// Implementation for getting/creating an instance of a derived type with
/// parametric storage.
auto StorageUniquer::getParametricStorageTypeImpl(
@@ -370,8 +554,10 @@ auto StorageUniquer::getParametricStorageTypeImpl(
/// parametric storage.
void StorageUniquer::registerParametricStorageTypeImpl(
TypeID id, function_ref<void(BaseStorage *)> destructorFn) {
- impl->parametricUniquers.try_emplace(
- id, std::make_unique<ParametricStorageUniquer>(destructorFn));
+ auto uniquer = std::make_unique<ParametricStorageUniquer>(destructorFn);
+ if (impl->isInTransientScope())
+ uniquer->beginTransientScope();
+ impl->parametricUniquers.try_emplace(id, std::move(uniquer));
}
/// Implementation for getting an instance of a derived type with default
@@ -394,6 +580,14 @@ bool StorageUniquer::isParametricStorageInitialized(TypeID id) {
/// storage.
void StorageUniquer::registerSingletonImpl(
TypeID id, function_ref<BaseStorage *(StorageAllocator &)> ctorFn) {
+ if (impl->transientState) {
+ assert(!impl->transientState->singletonInstances.count(id) &&
+ !impl->singletonInstances.count(id) &&
+ "storage class already registered");
+ impl->transientState->singletonInstances.try_emplace(
+ id, ctorFn(impl->getThreadSafeAllocator()));
+ return;
+ }
assert(!impl->singletonInstances.count(id) &&
"storage class already registered");
impl->singletonInstances.try_emplace(id, ctorFn(impl->allocator));
diff --git a/mlir/python/mlir/_mlir_libs/__init__.py b/mlir/python/mlir/_mlir_libs/__init__.py
index 886f38a6bf793..a8b28a1d20678 100644
--- a/mlir/python/mlir/_mlir_libs/__init__.py
+++ b/mlir/python/mlir/_mlir_libs/__init__.py
@@ -2,6 +2,7 @@
# See https://llvm.org/LICENSE.txt for license information.
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+from contextlib import contextmanager
from typing import Any, Mapping, Sequence
import os
@@ -198,6 +199,21 @@ def __init__(
)
init_module.register_llvm_translations(self)
+ @contextmanager
+ def transient_scope(self):
+ """Context manager that begins a transient scope and ends it on exit.
+
+ All types, attributes, affine expressions, and unregistered operations
+ allocated within this scope will be rolled back upon exit. Any Python
+ objects referencing transient IR entities must not be accessed after
+ exiting this context manager.
+ """
+ self.begin_transient_scope()
+ try:
+ yield self
+ finally:
+ self.end_transient_scope()
+
ir.Context = Context
# Register containers as Sequences, so they can be used with `match`.
diff --git a/mlir/test/CAPI/ir.c b/mlir/test/CAPI/ir.c
index bd9c8e0f50778..eed0335d6b053 100644
--- a/mlir/test/CAPI/ir.c
+++ b/mlir/test/CAPI/ir.c
@@ -2378,6 +2378,55 @@ void testExplicitThreadPools(void) {
mlirLlvmThreadPoolDestroy(threadPool);
}
+void testContextTransientScope(void) {
+ MlirContext ctx = mlirContextCreate();
+ fprintf(stderr, "@test_context_transient_scope\n");
+
+ // CHECK-LABEL: @test_context_transient_scope
+ // CHECK: is_in_transient_scope before: 0
+ fprintf(stderr, "is_in_transient_scope before: %d\n",
+ mlirContextIsInTransientScope(ctx));
+
+ MlirType i32Type =
+ mlirTypeParseGet(ctx, mlirStringRefCreateFromCString("i32"));
+
+ mlirContextBeginTransientScope(ctx);
+
+ // CHECK: is_in_transient_scope during: 1
+ fprintf(stderr, "is_in_transient_scope during: %d\n",
+ mlirContextIsInTransientScope(ctx));
+
+ MlirType transientVectorType =
+ mlirTypeParseGet(ctx, mlirStringRefCreateFromCString("vector<4xi32>"));
+ MlirAttribute transientStrAttr =
+ mlirStringAttrGet(ctx, mlirStringRefCreateFromCString("transient_str"));
+
+ // CHECK: transient vector valid: 1
+ fprintf(stderr, "transient vector valid: %d\n",
+ !mlirTypeIsNull(transientVectorType));
+ // CHECK: transient str attr valid: 1
+ fprintf(stderr, "transient str attr valid: %d\n",
+ !mlirAttributeIsNull(transientStrAttr));
+
+ mlirContextEndTransientScope(ctx);
+
+ // CHECK: is_in_transient_scope after: 0
+ fprintf(stderr, "is_in_transient_scope after: %d\n",
+ mlirContextIsInTransientScope(ctx));
+
+ MlirType postResetI32 =
+ mlirTypeParseGet(ctx, mlirStringRefCreateFromCString("i32"));
+ // CHECK: base i32 equal: 1
+ fprintf(stderr, "base i32 equal: %d\n", mlirTypeEqual(i32Type, postResetI32));
+
+ MlirType newVectorType =
+ mlirTypeParseGet(ctx, mlirStringRefCreateFromCString("vector<4xi32>"));
+ // CHECK: new vector valid: 1
+ fprintf(stderr, "new vector valid: %d\n", !mlirTypeIsNull(newVectorType));
+
+ mlirContextDestroy(ctx);
+}
+
void testLocation(void) {
MlirContext ctx = mlirContextCreate();
fprintf(stderr, "@test_location\n");
@@ -3259,6 +3308,7 @@ int main(void) {
return 16;
testExplicitThreadPools();
+ testContextTransientScope();
testLocation();
testDiagnostics();
diff --git a/mlir/test/python/ir/context_lifecycle.py b/mlir/test/python/ir/context_lifecycle.py
index 230db8277c8e7..e664ca0556f85 100644
--- a/mlir/test/python/ir/context_lifecycle.py
+++ b/mlir/test/python/ir/context_lifecycle.py
@@ -70,3 +70,71 @@
c6 = None
c7 = None
gc.collect()
+assert mlir.ir.Context._get_live_count() == 0
+
+# Test begin_transient_scope and end_transient_scope APIs
+print("TEST TRANSIENT SCOPE")
+ctx = mlir.ir.Context()
+assert not ctx.is_in_transient_scope
+
+with ctx, mlir.ir.Location.unknown(ctx):
+ i32 = mlir.ir.IntegerType.get_signless(32)
+
+ ctx.begin_transient_scope()
+ assert ctx.is_in_transient_scope
+
+ # Verify exception when attempting to enter again while already active
+ try:
+ ctx.begin_transient_scope()
+ except ValueError as e:
+ assert "Context is already in a transient scope" in str(e)
+ else:
+ assert False, "Expected ValueError when entering transient scope twice"
+
+ # Create transient types and attributes
+ vec_type = mlir.ir.VectorType.get([4], i32)
+ str_attr = mlir.ir.StringAttr.get("transient_ident")
+ assert str_attr.value == "transient_ident"
+
+ # End transient scope back to base
+ ctx.end_transient_scope()
+ assert not ctx.is_in_transient_scope
+
+ # Base type is intact
+ post_i32 = mlir.ir.IntegerType.get_signless(32)
+ assert post_i32 == i32
+
+ # Re-creating types post reset succeeds
+ new_vec_type = mlir.ir.VectorType.get([4], i32)
+ assert new_vec_type is not None
+
+ # Test transient_scope context manager
+ with ctx.transient_scope():
+ assert ctx.is_in_transient_scope
+ transient_f32 = mlir.ir.F32Type.get()
+ transient_vec = mlir.ir.VectorType.get([2, 2], transient_f32)
+ assert transient_vec is not None
+
+ # Verify exception on nested transient_scope
+ try:
+ with ctx.transient_scope():
+ pass
+ except ValueError as e:
+ assert "Context is already in a transient scope" in str(e)
+ else:
+ assert False, "Expected ValueError on nested transient_scope"
+
+ assert not ctx.is_in_transient_scope
+
+ctx = None
+i32 = None
+post_i32 = None
+vec_type = None
+new_vec_type = None
+str_attr = None
+transient_f32 = None
+transient_vec = None
+gc.collect()
+assert mlir.ir.Context._get_live_count() == 0
+
+
diff --git a/mlir/test/python/ir/context_managers.py b/mlir/test/python/ir/context_managers.py
index 5d9f9ceee97f3..3a7410ef293ab 100644
--- a/mlir/test/python/ir/context_managers.py
+++ b/mlir/test/python/ir/context_managers.py
@@ -88,3 +88,32 @@ def testInsertionPointEnterExit():
run(testInsertionPointEnterExit)
+
+
+# CHECK-LABEL: TEST: testTransientScope
+def testTransientScope():
+ ctx = Context()
+ with ctx, Location.unknown(ctx):
+ i32 = IntegerType.get_signless(32)
+ assert not ctx.is_in_transient_scope
+ with ctx.transient_scope():
+ assert ctx.is_in_transient_scope
+ v = VectorType.get([4], i32)
+ assert v is not None
+
+ # Test exception on entering nested transient scope
+ try:
+ with ctx.transient_scope():
+ pass
+ except ValueError as e:
+ # CHECK: Context is already in a transient scope
+ print(e)
+ else:
+ assert False, "Expected ValueError on nested transient scope"
+
+ assert not ctx.is_in_transient_scope
+ # Base type remains intact after scope exit
+ assert IntegerType.get_signless(32) == i32
+
+
+run(testTransientScope)
diff --git a/mlir/unittests/IR/MLIRContextResetTest.cpp b/mlir/unittests/IR/MLIRContextResetTest.cpp
new file mode 100644
index 0000000000000..442ad54ab2725
--- /dev/null
+++ b/mlir/unittests/IR/MLIRContextResetTest.cpp
@@ -0,0 +1,192 @@
+//===- MLIRContextResetTest.cpp - Tests for transient MLIRContext ---------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/IR/AffineExpr.h"
+#include "mlir/IR/AffineMap.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/BuiltinDialect.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/MLIRContext.h"
+#include "mlir/IR/OperationSupport.h"
+#include "llvm/Support/ThreadPool.h"
+#include "gtest/gtest.h"
+
+using namespace mlir;
+
+namespace {
+
+TEST(MLIRContextResetTest, BasicTypeTransientScopeAndReset) {
+ MLIRContext ctx;
+ OpBuilder builder(&ctx);
+
+ // Base types created before transient scope.
+ Type i32Type = builder.getI32Type();
+ Type f32Type = builder.getF32Type();
+ Type baseVectorType = VectorType::get({2, 2}, f32Type);
+
+ EXPECT_FALSE(ctx.isInTransientScope());
+ ctx.beginTransientScope();
+ EXPECT_TRUE(ctx.isInTransientScope());
+
+ // Base types remain identical under transient scope.
+ EXPECT_EQ(i32Type, builder.getI32Type());
+ EXPECT_EQ(f32Type, builder.getF32Type());
+ EXPECT_EQ(baseVectorType, VectorType::get({2, 2}, f32Type));
+
+ // Create transient types in transient scope.
+ Type transientVectorType = VectorType::get({4, 8}, i32Type);
+ Type transientTupleType = TupleType::get(&ctx, {i32Type, baseVectorType});
+
+ // Verify uniquing in transient state.
+ EXPECT_EQ(transientVectorType, VectorType::get({4, 8}, i32Type));
+ EXPECT_EQ(transientTupleType,
+ TupleType::get(&ctx, {i32Type, baseVectorType}));
+
+ // End transient scope back to base state.
+ ctx.endTransientScope();
+ EXPECT_FALSE(ctx.isInTransientScope());
+
+ // Verify base types are completely preserved.
+ EXPECT_EQ(i32Type, builder.getI32Type());
+ EXPECT_EQ(f32Type, builder.getF32Type());
+ EXPECT_EQ(baseVectorType, VectorType::get({2, 2}, f32Type));
+
+ // Types can be cleanly recreated post-reset.
+ Type postResetVectorType = VectorType::get({4, 8}, i32Type);
+ EXPECT_NE(postResetVectorType, Type());
+ EXPECT_EQ(postResetVectorType, VectorType::get({4, 8}, i32Type));
+}
+
+TEST(MLIRContextResetTest, AttributeTransientScopeAndReset) {
+ MLIRContext ctx;
+ OpBuilder builder(&ctx);
+
+ // Base attributes.
+ StringAttr baseStr = builder.getStringAttr("base_identifier");
+ UnitAttr baseUnit = builder.getUnitAttr();
+ IntegerAttr baseInt = builder.getI32IntegerAttr(42);
+ DistinctAttr baseDistinct = DistinctAttr::create(baseUnit);
+
+ ctx.beginTransientScope();
+ EXPECT_TRUE(ctx.isInTransientScope());
+
+ // Transient attributes.
+ StringAttr transientStr = builder.getStringAttr("transient_identifier");
+ IntegerAttr transientInt = builder.getI32IntegerAttr(100);
+ ArrayAttr transientArray = builder.getArrayAttr({baseInt, transientInt});
+ DictionaryAttr transientDict =
+ builder.getDictionaryAttr({builder.getNamedAttr("key", transientStr)});
+ DistinctAttr transientDistinct = DistinctAttr::create(transientInt);
+
+ EXPECT_EQ(transientStr, builder.getStringAttr("transient_identifier"));
+ EXPECT_EQ(transientArray, builder.getArrayAttr({baseInt, transientInt}));
+ EXPECT_EQ(transientDict, builder.getDictionaryAttr(
+ {builder.getNamedAttr("key", transientStr)}));
+
+ // End transient scope.
+ ctx.endTransientScope();
+ EXPECT_FALSE(ctx.isInTransientScope());
+
+ // Verify base attributes.
+ EXPECT_EQ(baseStr, builder.getStringAttr("base_identifier"));
+ EXPECT_EQ(baseUnit, builder.getUnitAttr());
+ EXPECT_EQ(baseInt, builder.getI32IntegerAttr(42));
+ EXPECT_NE(baseDistinct, DistinctAttr());
+
+ // Re-create attributes post-reset.
+ StringAttr postResetStr = builder.getStringAttr("transient_identifier");
+ EXPECT_EQ(postResetStr.getValue(), "transient_identifier");
+}
+
+TEST(MLIRContextResetTest, AffineTransientScopeAndReset) {
+ MLIRContext ctx;
+
+ // Base affine expression and map.
+ AffineExpr d0 = getAffineDimExpr(0, &ctx);
+ AffineExpr d1 = getAffineDimExpr(1, &ctx);
+ AffineMap baseMap = AffineMap::get(2, 0, {d0 + d1}, &ctx);
+
+ ctx.beginTransientScope();
+
+ // Transient affine expressions and maps.
+ AffineExpr c42 = getAffineConstantExpr(42, &ctx);
+ AffineMap transientMap = AffineMap::get(2, 0, {d0 * 4 + d1 + c42}, &ctx);
+
+ EXPECT_EQ(transientMap, AffineMap::get(2, 0, {d0 * 4 + d1 + c42}, &ctx));
+
+ ctx.endTransientScope();
+
+ // Verify base map.
+ EXPECT_EQ(baseMap, AffineMap::get(2, 0, {d0 + d1}, &ctx));
+
+ // Recreate map post reset.
+ AffineExpr newC42 = getAffineConstantExpr(42, &ctx);
+ AffineMap newMap = AffineMap::get(2, 0, {d0 * 4 + d1 + newC42}, &ctx);
+ EXPECT_EQ(newMap.getNumResults(), 1u);
+}
+
+TEST(MLIRContextResetTest, UnregisteredOperationPruning) {
+ MLIRContext ctx;
+ ctx.allowUnregisteredDialects(true);
+
+ // Base unregistered op.
+ OperationName baseOpName("custom_base.op", &ctx);
+
+ ctx.beginTransientScope();
+
+ // Transient unregistered op.
+ OperationName transientOpName("custom_transient.op", &ctx);
+ EXPECT_EQ(transientOpName.getStringRef(), "custom_transient.op");
+
+ ctx.endTransientScope();
+
+ // Base op is still valid and lookup succeeds.
+ OperationName baseOpNameAfter("custom_base.op", &ctx);
+ EXPECT_EQ(baseOpName, baseOpNameAfter);
+}
+
+TEST(MLIRContextResetTest, MultithreadedTransientScopeAndReset) {
+ MLIRContext ctx;
+ OpBuilder builder(&ctx);
+
+ Type i32Type = builder.getI32Type();
+
+ ctx.beginTransientScope();
+
+ // Allocate in parallel across multiple threads during transient scope.
+ llvm::DefaultThreadPool pool;
+ for (int i = 0; i < 20; ++i) {
+ pool.async([&ctx, i, i32Type]() {
+ for (int j = 0; j < 50; ++j) {
+ (void)VectorType::get({i + 1, j + 1}, i32Type);
+ (void)StringAttr::get(&ctx, "thread_str_" + std::to_string(i) + "_" +
+ std::to_string(j));
+ }
+ });
+ }
+ pool.wait();
+
+ // End transient scope back to base.
+ ctx.endTransientScope();
+
+ // Allocate again in parallel across multiple threads.
+ for (int i = 0; i < 20; ++i) {
+ pool.async([&ctx, i, i32Type]() {
+ for (int j = 0; j < 50; ++j) {
+ Type ty = VectorType::get({i + 1, j + 1}, i32Type);
+ EXPECT_TRUE(isa<VectorType>(ty));
+ }
+ });
+ }
+ pool.wait();
+
+ EXPECT_EQ(i32Type, builder.getI32Type());
+}
+
+} // namespace
diff --git a/mlir/unittests/Support/StorageUniquerTest.cpp b/mlir/unittests/Support/StorageUniquerTest.cpp
index 6db6783bb89f9..209e18cbedddb 100644
--- a/mlir/unittests/Support/StorageUniquerTest.cpp
+++ b/mlir/unittests/Support/StorageUniquerTest.cpp
@@ -58,3 +58,103 @@ TEST(StorageUniquerTest, NonTrivialDestructor) {
EXPECT_TRUE(wasDestructed);
}
+
+TEST(StorageUniquerTest, TransientScopeAndReset) {
+ struct IntStorage : public SimpleStorage<IntStorage, int> {
+ using Base::Base;
+ };
+
+ StorageUniquer uniquer;
+ uniquer.registerParametricStorageType<IntStorage>();
+
+ // Allocate in base layer.
+ IntStorage *base1 = IntStorage::get(uniquer, 1);
+ IntStorage *base2 = IntStorage::get(uniquer, 2);
+ EXPECT_EQ(base1, IntStorage::get(uniquer, 1));
+ EXPECT_EQ(base2, IntStorage::get(uniquer, 2));
+
+ // Begin transient scope.
+ EXPECT_FALSE(uniquer.isInTransientScope());
+ uniquer.beginTransientScope();
+ EXPECT_TRUE(uniquer.isInTransientScope());
+
+ // Base instances are still found and have same pointer.
+ EXPECT_EQ(base1, IntStorage::get(uniquer, 1));
+ EXPECT_EQ(base2, IntStorage::get(uniquer, 2));
+
+ // Allocate in transient layer.
+ IntStorage *transient3 = IntStorage::get(uniquer, 3);
+ IntStorage *transient4 = IntStorage::get(uniquer, 4);
+ EXPECT_EQ(transient3, IntStorage::get(uniquer, 3));
+ EXPECT_EQ(transient4, IntStorage::get(uniquer, 4));
+
+ // End transient scope.
+ uniquer.endTransientScope();
+ EXPECT_FALSE(uniquer.isInTransientScope());
+
+ // Base instances are still intact!
+ EXPECT_EQ(base1, IntStorage::get(uniquer, 1));
+ EXPECT_EQ(base2, IntStorage::get(uniquer, 2));
+
+ // Re-allocating transient key 3 now produces a valid object.
+ IntStorage *new3 = IntStorage::get(uniquer, 3);
+ EXPECT_NE(new3, nullptr);
+ EXPECT_EQ(std::get<0>(new3->key), 3);
+}
+
+TEST(StorageUniquerTest, DestructorOnEndTransientScope) {
+ struct NonTrivialStorage : public SimpleStorage<NonTrivialStorage, bool *> {
+ using Base::Base;
+ ~NonTrivialStorage() {
+ bool *wasDestructed = std::get<0>(key);
+ *wasDestructed = true;
+ }
+ };
+
+ StorageUniquer uniquer;
+ uniquer.registerParametricStorageType<NonTrivialStorage>();
+
+ bool baseDestructed = false;
+ NonTrivialStorage::get(uniquer, &baseDestructed);
+
+ uniquer.beginTransientScope();
+
+ bool transientDestructed = false;
+ NonTrivialStorage::get(uniquer, &transientDestructed);
+
+ EXPECT_FALSE(baseDestructed);
+ EXPECT_FALSE(transientDestructed);
+
+ // Ending transient scope should destroy only the transient instance.
+ uniquer.endTransientScope();
+
+ EXPECT_FALSE(baseDestructed);
+ EXPECT_TRUE(transientDestructed);
+}
+
+TEST(StorageUniquerTest, MutableStorageInTransientScope) {
+ struct MutableStorage : public SimpleStorage<MutableStorage, int> {
+ using Base::Base;
+ LogicalResult mutate(StorageUniquer::StorageAllocator &alloc, int newVal) {
+ std::get<0>(key) = newVal;
+ return success();
+ }
+ };
+
+ StorageUniquer uniquer;
+ uniquer.registerParametricStorageType<MutableStorage>();
+
+ uniquer.beginTransientScope();
+ MutableStorage *storage = MutableStorage::get(uniquer, 10);
+ EXPECT_EQ(std::get<0>(storage->key), 10);
+
+ EXPECT_TRUE(
+ succeeded(uniquer.mutate(TypeID::get<MutableStorage>(), storage, 20)));
+ EXPECT_EQ(std::get<0>(storage->key), 20);
+
+ uniquer.endTransientScope();
+
+ // Re-allocating after reset
+ MutableStorage *newStorage = MutableStorage::get(uniquer, 10);
+ EXPECT_EQ(std::get<0>(newStorage->key), 10);
+}
More information about the Mlir-commits
mailing list