incorporate comments

Author: makslevental

mlir/lib/Bindings/Python/IRCore.cpp

@@ -2538,8 +2538,8 @@ void mlir::python::populateIRCore(py::module &m) {
           [](py::object & /*class*/) {
             auto *context = PyThreadContextEntry::getDefaultContext();
             if (!context)
-              throw py::value_error("No current Context");
-            return context;
+              return py::none().cast<py::object>();
+            return py::cast(context);
           },
           "Gets the Context bound to the current thread or raises ValueError")
       .def_property_readonly(

mlir/python/CMakeLists.txt

@@ -21,7 +21,7 @@ declare_mlir_python_sources(MLIRPythonSources.Core.Python
     _mlir_libs/__init__.py
     ir.py
     passmanager.py
-    types.py
+    extras/types.py
     dialects/_ods_common.py
 
     # The main _mlir module has submodules: include stubs from each.

mlir/python/mlir/extras/__init__.py

@@ -5,11 +5,10 @@
 from functools import partial
 from typing import Optional, List
 
-from .ir import (
+from ..ir import (
     Attribute,
     BF16Type,
     ComplexType,
-    Context,
     F16Type,
     F32Type,
     F64Type,
@@ -32,55 +31,54 @@
     VectorType,
 )
 
-__all__ = []
+index = lambda: IndexType.get()
 
-_index = lambda: IndexType.get()
-_bool = lambda: IntegerType.get_signless(1)
 
-_i8 = lambda: IntegerType.get_signless(8)
-_i16 = lambda: IntegerType.get_signless(16)
-_i32 = lambda: IntegerType.get_signless(32)
-_i64 = lambda: IntegerType.get_signless(64)
+def i(width):
+    return IntegerType.get_signless(width)
 
-_si8 = lambda: IntegerType.get_signed(8)
-_si16 = lambda: IntegerType.get_signed(16)
-_si32 = lambda: IntegerType.get_signed(32)
-_si64 = lambda: IntegerType.get_signed(64)
 
-_ui8 = lambda: IntegerType.get_unsigned(8)
-_ui16 = lambda: IntegerType.get_unsigned(16)
-_ui32 = lambda: IntegerType.get_unsigned(32)
-_ui64 = lambda: IntegerType.get_unsigned(64)
+def si(width):
+    return IntegerType.get_signed(width)
 
-_f16 = lambda: F16Type.get()
-_f32 = lambda: F32Type.get()
-_f64 = lambda: F64Type.get()
-_bf16 = lambda: BF16Type.get()
 
-_f8e5m2 = lambda: Float8E5M2Type.get()
-_f8e4m3 = lambda: Float8E4M3FNType.get()
-_f8e4m3b11fnuz = lambda: Float8E4M3B11FNUZType.get()
+def ui(width):
+    return IntegerType.get_unsigned(width)
 
-_none = lambda: NoneType.get()
 
+bool = lambda: i(1)
+i8 = lambda: i(8)
+i16 = lambda: i(16)
+i32 = lambda: i(32)
+i64 = lambda: i(64)
 
-def _i(width):
-    return IntegerType.get_signless(width)
+si8 = lambda: si(8)
+si16 = lambda: si(16)
+si32 = lambda: si(32)
+si64 = lambda: si(64)
 
+ui8 = lambda: ui(8)
+ui16 = lambda: ui(16)
+ui32 = lambda: ui(32)
+ui64 = lambda: ui(64)
 
-def _si(width):
-    return IntegerType.get_signed(width)
+f16 = lambda: F16Type.get()
+f32 = lambda: F32Type.get()
+f64 = lambda: F64Type.get()
+bf16 = lambda: BF16Type.get()
 
+f8E5M2 = lambda: Float8E5M2Type.get()
+f8E4M3 = lambda: Float8E4M3FNType.get()
+f8E4M3B11FNUZ = lambda: Float8E4M3B11FNUZType.get()
 
-def _ui(width):
-    return IntegerType.get_unsigned(width)
+none = lambda: NoneType.get()
 
 
-def _complex(type):
+def complex(type):
     return ComplexType.get(type)
 
 
-def _opaque(dialect_namespace, type_data):
+def opaque(dialect_namespace, type_data):
     return OpaqueType.get(dialect_namespace, type_data)
 
 
@@ -105,7 +103,7 @@ def _shaped(*shape, element_type: Type = None, type_constructor=None):
         return type_constructor(type)
 
 
-def _vector(
+def vector(
     *shape,
     element_type: Type = None,
     scalable: Optional[List[bool]] = None,
@@ -120,7 +118,7 @@ def _vector(
     )
 
 
-def _tensor(*shape, element_type: Type = None, encoding: Optional[str] = None):
+def tensor(*shape, element_type: Type = None, encoding: Optional[str] = None):
     if encoding is not None:
         encoding = StringAttr.get(encoding)
     if not shape or (len(shape) == 1 and isinstance(shape[-1], Type)):
@@ -136,7 +134,7 @@ def _tensor(*shape, element_type: Type = None, encoding: Optional[str] = None):
     )
 
 
-def _memref(
+def memref(
     *shape,
     element_type: Type = None,
     memory_space: Optional[int] = None,
@@ -159,31 +157,9 @@ def _memref(
     )
 
 
-def _tuple(*elements):
+def tuple(*elements):
     return TupleType.get_tuple(elements)
 
 
-def _function(*, inputs, results):
+def function(*, inputs, results):
     return FunctionType.get(inputs, results)
-
-
-def __getattr__(name):
-    if name == "__path__":
-        # https://docs.python.org/3/reference/import.html#path__
-        # If a module is a package (either regular or namespace), the module object’s __path__ attribute must be set.
-        # This module is NOT a package and so this must be None (rather than throw the RuntimeError below).
-        return None
-    try:
-        Context.current
-    except ValueError:
-        raise RuntimeError("Types can only be instantiated under an active context.")
-
-    if f"_{name}" in globals():
-        builder = globals()[f"_{name}"]
-        if (
-            isinstance(builder, type(lambda: None))
-            and builder.__name__ == (lambda: None).__name__
-        ):
-            return builder()
-        return builder
-    raise RuntimeError(f"{name} is not a legal type.")

mlir/python/mlir/types.py renamed to mlir/python/mlir/extras/types.py

@@ -3,7 +3,7 @@
 import gc
 from mlir.ir import *
 from mlir.dialects import arith, tensor, func, memref
-import mlir.types as T
+import mlir.extras.types as T
 
 
 def run(f):
@@ -778,109 +778,96 @@ def testCustomTypeTypeCaster():
 # CHECK-LABEL: TEST: testTypeWrappers
 @run
 def testTypeWrappers():
-    try:
-        from mlir.types import i32
-    except RuntimeError as e:
-        assert e.args[0] == "Types can only be instantiated under an active context."
-
-    try:
-        from mlir.types import tensor
-    except RuntimeError as e:
-        assert e.args[0] == "Types can only be instantiated under an active context."
-
     def stride(strides, offset=0):
         return StridedLayoutAttr.get(offset, strides)
 
     with Context(), Location.unknown():
-        try:
-            from mlir.types import non_existent_type
-        except RuntimeError as e:
-            assert e.args[0] == "non_existent_type is not a legal type."
-
         ia = T.i(5)
         sia = T.si(6)
         uia = T.ui(7)
         assert repr(ia) == "IntegerType(i5)"
         assert repr(sia) == "IntegerType(si6)"
         assert repr(uia) == "IntegerType(ui7)"
 
-        assert T.i(16) == T.i16
-        assert T.si(16) == T.si16
-        assert T.ui(16) == T.ui16
+        assert T.i(16) == T.i16()
+        assert T.si(16) == T.si16()
+        assert T.ui(16) == T.ui16()
 
-        c1 = T.complex(T.f16)
-        c2 = T.complex(T.i32)
+        c1 = T.complex(T.f16())
+        c2 = T.complex(T.i32())
         assert repr(c1) == "ComplexType(complex<f16>)"
         assert repr(c2) == "ComplexType(complex<i32>)"
 
-        vec_1 = T.vector(2, 3, T.f32)
-        vec_2 = T.vector(2, 3, 4, T.f32)
+        vec_1 = T.vector(2, 3, T.f32())
+        vec_2 = T.vector(2, 3, 4, T.f32())
         assert repr(vec_1) == "VectorType(vector<2x3xf32>)"
         assert repr(vec_2) == "VectorType(vector<2x3x4xf32>)"
 
-        m1 = T.memref(2, 3, 4, T.f64)
+        m1 = T.memref(2, 3, 4, T.f64())
         assert repr(m1) == "MemRefType(memref<2x3x4xf64>)"
 
-        m2 = T.memref(2, 3, 4, T.f64, memory_space=1)
+        m2 = T.memref(2, 3, 4, T.f64(), memory_space=1)
         assert repr(m2) == "MemRefType(memref<2x3x4xf64, 1>)"
 
-        m3 = T.memref(2, 3, 4, T.f64, memory_space=1, layout=stride([5, 7, 13]))
+        m3 = T.memref(2, 3, 4, T.f64(), memory_space=1, layout=stride([5, 7, 13]))
         assert repr(m3) == "MemRefType(memref<2x3x4xf64, strided<[5, 7, 13]>, 1>)"
 
-        m4 = T.memref(2, 3, 4, T.f64, memory_space=1, layout=stride([5, 7, 13], 42))
+        m4 = T.memref(2, 3, 4, T.f64(), memory_space=1, layout=stride([5, 7, 13], 42))
         assert (
             repr(m4)
             == "MemRefType(memref<2x3x4xf64, strided<[5, 7, 13], offset: 42>, 1>)"
         )
 
         S = ShapedType.get_dynamic_size()
 
-        t1 = T.tensor(S, 3, S, T.f64)
+        t1 = T.tensor(S, 3, S, T.f64())
         assert repr(t1) == "RankedTensorType(tensor<?x3x?xf64>)"
-        ut1 = T.tensor(T.f64)
+        ut1 = T.tensor(T.f64())
         assert repr(ut1) == "UnrankedTensorType(tensor<*xf64>)"
-        t2 = T.tensor(S, 3, S, element_type=T.f64)
+        t2 = T.tensor(S, 3, S, element_type=T.f64())
         assert repr(t2) == "RankedTensorType(tensor<?x3x?xf64>)"
-        ut2 = T.tensor(element_type=T.f64)
+        ut2 = T.tensor(element_type=T.f64())
         assert repr(ut2) == "UnrankedTensorType(tensor<*xf64>)"
 
-        t3 = T.tensor(S, 3, S, T.f64, encoding="encoding")
+        t3 = T.tensor(S, 3, S, T.f64(), encoding="encoding")
         assert repr(t3) == 'RankedTensorType(tensor<?x3x?xf64, "encoding">)'
 
-        v = T.vector(3, 3, 3, T.f64)
+        v = T.vector(3, 3, 3, T.f64())
         assert repr(v) == "VectorType(vector<3x3x3xf64>)"
 
-        m5 = T.memref(S, 3, S, T.f64)
+        m5 = T.memref(S, 3, S, T.f64())
         assert repr(m5) == "MemRefType(memref<?x3x?xf64>)"
-        um1 = T.memref(T.f64)
+        um1 = T.memref(T.f64())
         assert repr(um1) == "UnrankedMemRefType(memref<*xf64>)"
-        m6 = T.memref(S, 3, S, element_type=T.f64)
+        m6 = T.memref(S, 3, S, element_type=T.f64())
         assert repr(m6) == "MemRefType(memref<?x3x?xf64>)"
-        um2 = T.memref(element_type=T.f64)
+        um2 = T.memref(element_type=T.f64())
         assert repr(um2) == "UnrankedMemRefType(memref<*xf64>)"
 
-        m7 = T.memref(S, 3, S, T.f64)
+        m7 = T.memref(S, 3, S, T.f64())
         assert repr(m7) == "MemRefType(memref<?x3x?xf64>)"
-        um3 = T.memref(T.f64)
+        um3 = T.memref(T.f64())
         assert repr(um3) == "UnrankedMemRefType(memref<*xf64>)"
 
-        scalable_1 = T.vector(2, 3, T.f32, scalable=[False, True])
-        scalable_2 = T.vector(2, 3, 4, T.f32, scalable=[True, False, True])
+        scalable_1 = T.vector(2, 3, T.f32(), scalable=[False, True])
+        scalable_2 = T.vector(2, 3, 4, T.f32(), scalable=[True, False, True])
         assert repr(scalable_1) == "VectorType(vector<2x[3]xf32>)"
         assert repr(scalable_2) == "VectorType(vector<[2]x3x[4]xf32>)"
 
-        scalable_3 = T.vector(2, 3, T.f32, scalable_dims=[1])
-        scalable_4 = T.vector(2, 3, 4, T.f32, scalable_dims=[0, 2])
+        scalable_3 = T.vector(2, 3, T.f32(), scalable_dims=[1])
+        scalable_4 = T.vector(2, 3, 4, T.f32(), scalable_dims=[0, 2])
         assert scalable_3 == scalable_1
         assert scalable_4 == scalable_2
 
         opaq = T.opaque("scf", "placeholder")
         assert repr(opaq) == "OpaqueType(!scf.placeholder)"
 
-        tup1 = T.tuple(T.i16, T.i32, T.i64)
-        tup2 = T.tuple(T.f16, T.f32, T.f64)
+        tup1 = T.tuple(T.i16(), T.i32(), T.i64())
+        tup2 = T.tuple(T.f16(), T.f32(), T.f64())
         assert repr(tup1) == "TupleType(tuple<i16, i32, i64>)"
         assert repr(tup2) == "TupleType(tuple<f16, f32, f64>)"
 
-        func = T.function(inputs=(T.i16, T.i32, T.i64), results=(T.f16, T.f32, T.f64))
+        func = T.function(
+            inputs=(T.i16(), T.i32(), T.i64()), results=(T.f16(), T.f32(), T.f64())
+        )
         assert repr(func) == "FunctionType((i16, i32, i64) -> (f16, f32, f64))"

mlir/test/python/ir/builtin_types.py

@@ -15,13 +15,7 @@ def run(f):
 def testContextEnterExit():
     with Context() as ctx:
         assert Context.current is ctx
-    try:
-        _ = Context.current
-    except ValueError as e:
-        # CHECK: No current Context
-        print(e)
-    else:
-        assert False, "Expected exception"
+    assert Context.current is None
 
 
 run(testContextEnterExit)