remove default context

Author: makslevental

mlir_utils/__init__.py

@@ -1,22 +1,3 @@
-import atexit
-
 from ._configuration.configuration import alias_upstream_bindings
 
-if alias_upstream_bindings():
-    from mlir import ir
-
-    DefaultContext = ir.Context()
-    # Push a default context onto the context stack at import time.
-    DefaultContext.__enter__()
-    DefaultContext.allow_unregistered_dialects = False
-
-    DefaultLocation = ir.Location.unknown()
-    DefaultLocation.__enter__()
-
-    @atexit.register
-    def __exit_ctxt():
-        DefaultContext.__exit__(None, None, None)
-
-    @atexit.register
-    def __exit_loc():
-        DefaultLocation.__exit__(None, None, None)
+alias_upstream_bindings()

mlir_utils/context.py

@@ -4,8 +4,6 @@
 
 import mlir.ir
 
-from mlir_utils import DefaultContext
-
 
 @dataclass
 class MLIRContext:
@@ -19,7 +17,7 @@ def __str__(self):
 @contextmanager
 def mlir_mod_ctx(
     src: Optional[str] = None,
-    context: mlir.ir.Context = DefaultContext,
+    context: mlir.ir.Context = None,
     location: mlir.ir.Location = None,
     allow_unregistered_dialects=False,
 ) -> MLIRContext:

mlir_utils/dialects/ext/arith.py

@@ -35,7 +35,7 @@
 except ModuleNotFoundError:
     pass
 
-from mlir_utils.types import infer_mlir_type, MLIR_TYPE_TO_NP_DTYPE
+from mlir_utils.types import infer_mlir_type, mlir_type_to_np_dtype
 
 
 def constant(
@@ -69,7 +69,7 @@ def constant(
         ranked_tensor_type = RankedTensorType(type)
         value = np.ones(
             ranked_tensor_type.shape,
-            dtype=MLIR_TYPE_TO_NP_DTYPE()[ranked_tensor_type.element_type],
+            dtype=mlir_type_to_np_dtype(ranked_tensor_type.element_type),
         )
     assert type is not None
 

mlir_utils/dialects/ext/scf.py

@@ -9,6 +9,7 @@
 from mlir.dialects.scf import IfOp, ForOp
 from mlir.ir import InsertionPoint, Value, OpResultList, OpResult
 
+import mlir_utils.types as T
 from mlir_utils.ast.canonicalize import (
     StrictTransformer,
     Canonicalizer,
@@ -18,7 +19,6 @@
 from mlir_utils.ast.util import ast_call
 from mlir_utils.dialects.ext.arith import constant
 from mlir_utils.dialects.scf import yield_ as yield__
-from mlir_utils.types import opaque_t
 from mlir_utils.util import (
     region_op,
     maybe_cast,
@@ -101,8 +101,6 @@ def _if(cond, results_=None, *, has_else=False, loc=None, ip=None):
 
 if_ = region_op(_if, terminator=yield__)
 
-_placeholder_opaque_t = opaque_t("scf", "placeholder")
-
 
 class IfStack:
     __current_if_op: list[IfOp] = []
@@ -175,7 +173,7 @@ def yield_(*args):
 
         assert len(results) == len(unpacked_args), f"{results=}, {unpacked_args=}"
         for i, r in enumerate(results):
-            if r.type == _placeholder_opaque_t:
+            if r.type == T._placeholder_opaque_t():
                 r.set_type(unpacked_args[i].type)
 
         yield_(*args)
@@ -325,13 +323,13 @@ def insert_with_results(
         ), f"conditional with := must explicitly yield on last line"
         yield_expr = last_statement.body[0]
         if m.matches(yield_expr.value, m.Call(func=m.Name(stack_yield.__name__))):
-            results = [cst.Element(cst.Name("_placeholder_opaque_t"))] * len(
+            results = [cst.Element(ast_call(T._placeholder_opaque_t.__name__))] * len(
                 yield_expr.value.args
             )
         elif m.matches(yield_expr.value.value, m.Name()):
-            results = [cst.Element(cst.Name("_placeholder_opaque_t"))]
+            results = [cst.Element(ast_call(T._placeholder_opaque_t.__name__))]
         elif m.matches(yield_expr.value.value, m.Tuple()):
-            results = [cst.Element(cst.Name("_placeholder_opaque_t"))] * len(
+            results = [cst.Element(ast_call(T._placeholder_opaque_t.__name__))] * len(
                 yield_expr.value.value.elements
             )
         results = cst.Tuple(results)
@@ -422,14 +420,13 @@ def patch_bytecode(self, code: ConcreteBytecode, f):
                 str(OpCode.NOP), lineno=c.lineno, location=c.location
             )
 
-        # TODO(max): this is bad
         f.__globals__[else_.__name__] = else_
         f.__globals__[end_branch.__name__] = end_branch
         f.__globals__[end_if.__name__] = end_if
         f.__globals__[stack_if.__name__] = stack_if
         f.__globals__[stack_yield.__name__] = stack_yield
         f.__globals__[yield_.__name__] = yield_
-        f.__globals__["_placeholder_opaque_t"] = _placeholder_opaque_t
+        f.__globals__[T._placeholder_opaque_t.__name__] = T._placeholder_opaque_t
         return code
 
 

mlir_utils/types.py

@@ -1,3 +1,4 @@
+import sys
 from functools import partial
 from typing import Union
 
@@ -19,33 +20,70 @@
     OpaqueType,
 )
 
-index_t = IndexType.get()
-bool_t = IntegerType.get_signless(1)
-i8_t = IntegerType.get_signless(8)
-i16_t = IntegerType.get_signless(16)
-i32_t = IntegerType.get_signless(32)
-i64_t = IntegerType.get_signless(64)
-f16_t = F16Type.get()
-f32_t = F32Type.get()
-f64_t = F64Type.get()
-bf16_t = BF16Type.get()
+_index_t = lambda: IndexType.get()
+_bool_t = lambda: IntegerType.get_signless(1)
+_i8_t = lambda: IntegerType.get_signless(8)
+_i16_t = lambda: IntegerType.get_signless(16)
+_i32_t = lambda: IntegerType.get_signless(32)
+_i64_t = lambda: IntegerType.get_signless(64)
+_f16_t = lambda: F16Type.get()
+_f32_t = lambda: F32Type.get()
+_f64_t = lambda: F64Type.get()
+_bf16_t = lambda: BF16Type.get()
 opaque_t = lambda dialect_namespace, buffer: OpaqueType.get(dialect_namespace, buffer)
 
-NP_DTYPE_TO_MLIR_TYPE = lambda: {
-    np.int8: i8_t,
-    np.int16: i16_t,
-    np.int32: i32_t,
-    np.int64: i64_t,
-    # this is techincally wrong i guess but numpy by default casts python scalars to this
-    # so to support passing lists of ints we map this to index type
-    np.longlong: index_t,
-    np.uintp: index_t,
-    np.float16: f16_t,
-    np.float32: f32_t,
-    np.float64: f64_t,
+
+def _placeholder_opaque_t():
+    return opaque_t("scf", "placeholder")
+
+
+_name_to_type = {
+    "index_t": _index_t,
+    "bool_t": _bool_t,
+    "i8_t": _i8_t,
+    "i16_t": _i16_t,
+    "i32_t": _i32_t,
+    "i64_t": _i64_t,
+    "f16_t": _f16_t,
+    "f32_t": _f32_t,
+    "f64_t": _f64_t,
+    "bf16_t": _bf16_t,
+}
+
+
+def __getattr__(name):
+    if name in _name_to_type:
+        return _name_to_type[name]()
+    # this kicks it to the default module attribute lookup (i.e., functions defined below and such)
+    return None
+
+
+_np_dtype_to_mlir_type_ctor = {
+    np.int8: _i8_t,
+    np.int16: _i16_t,
+    np.int32: _i32_t,
+    np.int64: _i64_t,
+    # is technically wrong i guess but numpy by default casts python scalars to this
+    # so to support passing lists of ints we map to index type
+    np.longlong: _index_t,
+    np.uintp: _index_t,
+    np.float16: _f16_t,
+    np.float32: _f32_t,
+    np.float64: _f64_t,
 }
 
-MLIR_TYPE_TO_NP_DTYPE = lambda: {v: k for k, v in NP_DTYPE_TO_MLIR_TYPE().items()}
+_mlir_type_ctor_to_np_dtype = lambda: {
+    v: k for k, v in _np_dtype_to_mlir_type_ctor.items()
+}
+
+
+def np_dtype_to_mlir_type(np_dtype):
+    return _np_dtype_to_mlir_type_ctor[np_dtype]()
+
+
+def mlir_type_to_np_dtype(mlir_type):
+    _mlir_type_to_np_dtype = {v(): k for k, v in _np_dtype_to_mlir_type_ctor.items()}
+    return _mlir_type_to_np_dtype[mlir_type]
 
 
 def infer_mlir_type(
@@ -62,13 +100,13 @@ def infer_mlir_type(
       MLIR type corresponding to py_val.
     """
     if isinstance(py_val, bool):
-        return bool_t
+        return _bool_t()
     elif isinstance(py_val, int):
-        return i64_t
+        return _i64_t()
     elif isinstance(py_val, float):
-        return f64_t
+        return _f64_t()
     elif isinstance(py_val, np.ndarray):
-        dtype = NP_DTYPE_TO_MLIR_TYPE()[py_val.dtype.type]
+        dtype = np_dtype_to_mlir_type(py_val.dtype.type)
         return RankedTensorType.get(py_val.shape, dtype)
     else:
         raise NotImplementedError(

tests/test_location_tracking.py

@@ -1,8 +1,10 @@
+from os import sep
 from pathlib import Path
 from textwrap import dedent
-from os import sep
+
 import pytest
 
+import mlir_utils.types as T
 from mlir_utils.ast.canonicalize import canonicalize
 from mlir_utils.dialects.ext.arith import constant
 from mlir_utils.dialects.ext.scf import (
@@ -14,12 +16,10 @@
 
 # noinspection PyUnresolvedReferences
 from mlir_utils.testing import mlir_ctx as ctx, filecheck, MLIRContext
-from mlir_utils.types import f64_t, index_t, tensor_t
 
 # needed since the fix isn't defined here nor conftest.py
 pytest.mark.usefixtures("ctx")
 
-
 THIS_DIR = str(Path(__file__).parent.absolute())
 
 
@@ -34,7 +34,7 @@ def test_if_replace_yield_5(ctx: MLIRContext):
     def iffoo():
         one = constant(1.0)
         two = constant(2.0)
-        if res := stack_if(one < two, (f64_t, f64_t, f64_t)):
+        if res := stack_if(one < two, (T.f64_t, T.f64_t, T.f64_t)):
             three = constant(3.0)
             yield three, three, three
         else:
@@ -74,11 +74,11 @@ def iffoo():
 
 
 def test_block_args(ctx: MLIRContext):
-    one = constant(1, index_t)
-    two = constant(2, index_t)
+    one = constant(1, T.index_t)
+    two = constant(2, T.index_t)
 
-    @generate(tensor_t(S, 3, S, f64_t), dynamic_extents=[one, two])
-    def demo_fun1(i: index_t, j: index_t, k: index_t):
+    @generate(T.tensor_t(S, 3, S, T.f64_t), dynamic_extents=[one, two])
+    def demo_fun1(i: T.index_t, j: T.index_t, k: T.index_t):
         one = constant(1.0)
         tensor_yield(one)
 

tests/test_operator_overloading.py

@@ -2,12 +2,12 @@
 
 import pytest
 
+import mlir_utils.types as T
 from mlir_utils.dialects.ext.arith import constant, Scalar
 from mlir_utils.dialects.ext.tensor import Tensor, empty
 
 # noinspection PyUnresolvedReferences
 from mlir_utils.testing import mlir_ctx as ctx, filecheck, MLIRContext
-from mlir_utils.types import f64_t
 
 # needed since the fix isn't defined here nor conftest.py
 pytest.mark.usefixtures("ctx")
@@ -21,9 +21,9 @@ def test_tensor_arithmetic(ctx: MLIRContext):
     three = one + two
     assert isinstance(three, Scalar)
 
-    ten1 = empty((10, 10, 10), f64_t)
+    ten1 = empty((10, 10, 10), T.f64_t)
     assert isinstance(ten1, Tensor)
-    ten2 = empty((10, 10, 10), f64_t)
+    ten2 = empty((10, 10, 10), T.f64_t)
     assert isinstance(ten2, Tensor)
     ten3 = ten1 + ten2
     assert isinstance(ten3, Tensor)

tests/test_regions.py

@@ -2,6 +2,7 @@
 
 import pytest
 
+import mlir_utils.types as T
 from mlir_utils.dialects.ext.arith import constant
 from mlir_utils.dialects.ext.func import func
 from mlir_utils.dialects.ext.tensor import S, rank
@@ -11,7 +12,7 @@
 
 # noinspection PyUnresolvedReferences
 from mlir_utils.testing import mlir_ctx as ctx, filecheck, MLIRContext
-from mlir_utils.types import f64_t, index_t, tensor_t
+from mlir_utils.types import tensor_t
 
 # needed since the fix isn't defined here nor conftest.py
 pytest.mark.usefixtures("ctx")
@@ -93,11 +94,11 @@ def demo_fun1():
 
 
 def test_block_args(ctx: MLIRContext):
-    one = constant(1, index_t)
-    two = constant(2, index_t)
+    one = constant(1, T.index_t)
+    two = constant(2, T.index_t)
 
-    @generate(tensor_t(S, 3, S, f64_t), dynamic_extents=[one, two])
-    def demo_fun1(i: index_t, j: index_t, k: index_t):
+    @generate(tensor_t(S, 3, S, T.f64_t), dynamic_extents=[one, two])
+    def demo_fun1(i: T.index_t, j: T.index_t, k: T.index_t):
         one = constant(1.0)
         tensor_yield(one)