mlir/test/python/ir/context_lifecycle.py

# RUN: %PYTHON %s
# Standalone sanity check of context life-cycle.
import gc
import mlir.ir

assert mlir.ir.Context._get_live_count() == 0

# Create first context.
print("CREATE C1")
c1 = mlir.ir.Context()
assert mlir.ir.Context._get_live_count() == 1
c1_repr = repr(c1)
print("C1 = ", c1_repr)

print("GETTING AGAIN...")
c2 = c1._get_context_again()
c2_repr = repr(c2)
assert mlir.ir.Context._get_live_count() == 1
assert c1_repr == c2_repr

print("C2 =", c2)

# Make sure new contexts on constructor.
print("CREATE C3")
c3 = mlir.ir.Context()
assert mlir.ir.Context._get_live_count() == 2
c3_repr = repr(c3)
print("C3 =", c3)
assert c3_repr != c1_repr
print("FREE C3")
c3 = None
gc.collect()
assert mlir.ir.Context._get_live_count() == 1

print("Free C1")
c1 = None
gc.collect()
assert mlir.ir.Context._get_live_count() == 1
print("Free C2")
c2 = None
gc.collect()
assert mlir.ir.Context._get_live_count() == 0

# Create a context, get its capsule and create from capsule.
c4 = mlir.ir.Context()
c4_capsule = c4._CAPIPtr
assert '"mlir.ir.Context._CAPIPtr"' in repr(c4_capsule)
c5 = mlir.ir.Context._CAPICreate(c4_capsule)
assert c4 is c5
c4 = None
c5 = None
gc.collect()

# Create a global threadpool and use it in two contexts
tp = mlir.ir.ThreadPool()
assert tp.get_max_concurrency() > 0
c5 = mlir.ir.Context()
c5.set_thread_pool(tp)
assert c5.get_num_threads() == tp.get_max_concurrency()
assert c5._mlir_thread_pool_ptr() == tp._mlir_thread_pool_ptr()
c6 = mlir.ir.Context()
c6.set_thread_pool(tp)
assert c6.get_num_threads() == tp.get_max_concurrency()
assert c6._mlir_thread_pool_ptr() == tp._mlir_thread_pool_ptr()
c7 = mlir.ir.Context(thread_pool=tp)
assert c7.get_num_threads() == tp.get_max_concurrency()
assert c7._mlir_thread_pool_ptr() == tp._mlir_thread_pool_ptr()
assert mlir.ir.Context._get_live_count() == 3
c5 = None
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