fix func and add module/context convenience classes (#52)

Author: makslevental

mlir/extras/context.py

@@ -41,6 +41,43 @@ def mlir_mod_ctx(
     context._clear_live_operations()
 
 
+class RAIIMLIRContext:
+    context: ir.Context
+    location: ir.Location
+
+    def __init__(self, location: Optional[ir.Location] = None):
+        self.context = ir.Context()
+        self.context.__enter__()
+        if location is None:
+            location = ir.Location.unknown()
+        self.location = location
+        self.location.__enter__()
+
+    def __del__(self):
+        self.location.__exit__(None, None, None)
+        self.context.__exit__(None, None, None)
+        assert ir.Context is None
+
+
+class ExplicitlyManagedModule:
+    module: ir.Module
+    _ip: ir.InsertionPoint
+
+    def __init__(self, src: Optional[str] = None):
+        if src is not None:
+            self.module = ir.Module.parse(src)
+        else:
+            self.module = ir.Module.create()
+        self._ip = ir.InsertionPoint(self.module.body)
+        self._ip.__enter__()
+
+    def finish(self):
+        self._ip.__exit__(None, None, None)
+
+    def __str__(self):
+        return str(self.module)
+
+
 @contextlib.contextmanager
 def enable_multithreading(context=None):
     from ..ir import Context

mlir/extras/dialects/ext/func.py

@@ -182,7 +182,7 @@ def emit(self, *call_args, decl=False, force=False) -> FuncOp:
                 input_types = self.input_types[:]
                 for i, v in enumerate(input_types):
                     if isinstance(v, str):
-                        input_types[i] = Type(eval(v, {"T": T}))
+                        input_types[i] = Type(eval(v, self.body_builder.__globals__))
                     elif isalambda(v):
                         input_types[i] = v()
             else:

tests/test_regions.py

@@ -521,6 +521,42 @@ def mod():
     filecheck(correct, ctx.module)
 
 
+M, K, N = 64, 32, 64
+
+
+@func(emit=False, sym_visibility="private")
+def matmul_i16_i16(
+    A: "T.memref(M, K, T.i16())",
+    B: "T.memref(K, N, T.i16())",
+    C: "T.memref(M, N, T.i16())",
+):
+    linalg.matmul(A, B, C)
+
+
+def test_defer_emit_3(ctx: MLIRContext):
+
+    matmul_i16_i16.emit(force=True)
+
+    @module
+    def mod():
+        matmul_i16_i16.emit(decl=True)
+
+    correct = dedent(
+        """\
+    module {
+      func.func private @matmul_i16_i16(%arg0: memref<64x32xi16>, %arg1: memref<32x64xi16>, %arg2: memref<64x64xi16>) {
+        linalg.matmul {cast = #linalg.type_fn<cast_signed>} ins(%arg0, %arg1 : memref<64x32xi16>, memref<32x64xi16>) outs(%arg2 : memref<64x64xi16>)
+        return
+      }
+      module {
+        func.func private @matmul_i16_i16(memref<64x32xi16>, memref<32x64xi16>, memref<64x64xi16>)
+      }
+    }
+    """
+    )
+    filecheck(correct, ctx.module)
+
+
 def test_successor_ctx_manager(ctx: MLIRContext):
     @func
     def foo1():

tests/test_runtime.py

@@ -5,11 +5,12 @@
 
 import numpy as np
 import pytest
-from mlir.ir import UnitAttr, Module, StridedLayoutAttr
+from mlir.ir import UnitAttr, Module, StridedLayoutAttr, InsertionPoint, Context
 from mlir.runtime import get_unranked_memref_descriptor, get_ranked_memref_descriptor
 
 import mlir.extras.types as T
 from mlir.extras.ast.canonicalize import canonicalize
+from mlir.extras.context import RAIIMLIRContext, ExplicitlyManagedModule
 from mlir.extras.dialects.ext import linalg
 from mlir.dialects.arith import sitofp, index_cast
 from mlir.extras.dialects.ext.arith import constant
@@ -872,3 +873,50 @@ def tenfoo(
     #
     # invoker.memfoo_capi_wrapper(AA, BB, CC)
     # # assert np.array_equal(A + B, C)
+
+
+def test_raii_context():
+    def foo():
+        ctx = RAIIMLIRContext()
+        mod = Module.create()
+        with InsertionPoint(mod.body):
+
+            @func(emit=True)
+            def foo(x: T.i32()):
+                return x
+
+        correct = dedent(
+            """\
+        module {
+          func.func @foo(%arg0: i32) -> i32 {
+            return %arg0 : i32
+          }
+        }
+        """
+        )
+        filecheck(correct, mod)
+
+    foo()
+    assert Context.current is None
+
+
+def test_explicit_module():
+    ctx = RAIIMLIRContext()
+    mod = ExplicitlyManagedModule()
+
+    @func(emit=True)
+    def foo(x: T.i32()):
+        return x
+
+    mod.finish()
+
+    correct = dedent(
+        """\
+    module {
+      func.func @foo(%arg0: i32) -> i32 {
+        return %arg0 : i32
+      }
+    }
+    """
+    )
+    filecheck(correct, mod)