add RAIIMLIRContextModule (#111)

Author: makslevental

mlir/extras/context.py

@@ -45,8 +45,12 @@ class RAIIMLIRContext:
     context: ir.Context
     location: ir.Location
 
-    def __init__(self, location: Optional[ir.Location] = None):
+    def __init__(
+        self, location: Optional[ir.Location] = None, allow_unregistered_dialects=False
+    ):
         self.context = ir.Context()
+        if allow_unregistered_dialects:
+            self.context.allow_unregistered_dialects = True
         self.context.__enter__()
         if location is None:
             location = ir.Location.unknown()
@@ -61,6 +65,36 @@ def __del__(self):
             assert ir.Context is not self.context
 
 
+class RAIIMLIRContextModule:
+    context: ir.Context
+    location: ir.Location
+    insertion_point: ir.InsertionPoint
+    module: ir.Module
+
+    def __init__(
+        self, location: Optional[ir.Location] = None, allow_unregistered_dialects=False
+    ):
+        self.context = ir.Context()
+        if allow_unregistered_dialects:
+            self.context.allow_unregistered_dialects = True
+        self.context.__enter__()
+        if location is None:
+            location = ir.Location.unknown()
+        self.location = location
+        self.location.__enter__()
+        self.module = ir.Module.create()
+        self.insertion_point = ir.InsertionPoint(self.module.body)
+        self.insertion_point.__enter__()
+
+    def __del__(self):
+        self.insertion_point.__exit__(None, None, None)
+        self.location.__exit__(None, None, None)
+        self.context.__exit__(None, None, None)
+        # i guess the extension gets destroyed before this object sometimes?
+        if ir is not None:
+            assert ir.Context is not self.context
+
+
 class ExplicitlyManagedModule:
     module: ir.Module
     _ip: ir.InsertionPoint

tests/test_func.py

@@ -1,5 +1,6 @@
 import inspect
 import sys
+import threading
 from textwrap import dedent
 from typing import TypeVar
 
@@ -8,7 +9,7 @@
 import mlir.extras.types as T
 
 from mlir.extras.ast.canonicalize import canonicalize
-from mlir.extras.context import mlir_mod_ctx
+from mlir.extras.context import mlir_mod_ctx, RAIIMLIRContextModule
 from mlir.extras.dialects.ext.arith import constant
 from mlir.extras.dialects.ext.func import func
 from mlir.extras.dialects.ext import linalg, arith, scf, memref
@@ -294,3 +295,28 @@ def mat_product_kernel(
     """
     )
     filecheck(correct, ctx.module)
+
+
+def test_raii_mlir_context_module():
+    tls = threading.local()
+    tls.ctx = RAIIMLIRContextModule()
+
+    @func
+    def demo_fun1():
+        one = constant(1)
+        return one
+
+    assert hasattr(demo_fun1, "emit")
+    assert inspect.ismethod(demo_fun1.emit)
+    demo_fun1.emit()
+    correct = dedent(
+        """\
+    module {
+      func.func @demo_fun1() -> i32 {
+        %c1_i32 = arith.constant 1 : i32
+        return %c1_i32 : i32
+      }
+    }
+    """
+    )
+    filecheck(correct, tls.ctx.module)