context fixture

Author: makslevental

mlir/extras/context.py

@@ -15,6 +15,24 @@ def __str__(self):
         return str(self.module)
 
 
+@contextmanager
+def mlir_mod(
+    src: Optional[str] = None,
+    location: ir.Location = None,
+) -> ir.Module:
+    with ExitStack() as stack:
+        if location is None:
+            location = ir.Location.unknown()
+        stack.enter_context(location)
+        if src is not None:
+            module = ir.Module.parse(src)
+        else:
+            module = ir.Module.create()
+        ip = ir.InsertionPoint(module.body)
+        stack.enter_context(ip)
+        yield module
+
+
 @contextmanager
 def mlir_mod_ctx(
     src: Optional[str] = None,
@@ -112,24 +130,6 @@ def __str__(self):
         return str(self.module)
 
 
-@contextmanager
-def mlir_mod(
-    src: Optional[str] = None,
-    location: ir.Location = None,
-) -> ir.Module:
-    with ExitStack() as stack:
-        if location is None:
-            location = ir.Location.unknown()
-        stack.enter_context(location)
-        if src is not None:
-            module = ir.Module.parse(src)
-        else:
-            module = ir.Module.create()
-        ip = ir.InsertionPoint(module.body)
-        stack.enter_context(ip)
-        yield module
-
-
 @contextlib.contextmanager
 def enable_multithreading(context=None):
     from ..ir import Context

mlir/extras/testing/testing.py

@@ -123,12 +123,12 @@ def filecheck_with_comments(module):
             raise ValueError(f"\n{err}")
 
 
-@pytest.fixture
+@pytest.fixture(scope="function")
 def mlir_ctx() -> MLIRContext:
     with mlir_mod_ctx(allow_unregistered_dialects=True) as ctx:
         yield ctx
 
 
-@pytest.fixture
+@pytest.fixture(scope="function")
 def backend() -> LLVMJITBackend:
     return LLVMJITBackend()

tests/test_gpu.py

@@ -1220,8 +1220,6 @@ def smol_matmul(
     def gpu_module():
         smol_matmul.emit()
 
-    print(gpu_module)
-
     lowered_module = run_pipeline(
         gpu_module,
         Pipeline()

tests/test_nvgpu_nvvm.py

@@ -413,7 +413,6 @@ def main(module: any_op_t()):
     # CHECK:    }
     # CHECK:  }
 
-    print(mod)
     filecheck_with_comments(mod)
 
 

tests/test_runtime.py

@@ -522,7 +522,6 @@ def memfoo(mem: ranked_memref_kxk_f32):
             mem[i, i] = mem[i, i] + mem[i, i] * sitofp(T.f32(), index_cast(T.i32(), i))
 
     memfoo.emit()
-    print(ctx.module)
 
     module = backend.compile(
         ctx.module,

tests/test_vector.py

@@ -120,8 +120,6 @@ def pats():
         ),
     )
 
-    print(vectorized_module)
-
     compiled_module = backend.compile(
         find_ops(
             vectorized_module.operation,