use upstream types

Author: makslevental

examples/mlir_python_extras.ipynb

@@ -59,6 +59,7 @@
     "from mlir.extras.dialects.ext.scf import canonicalizer as scf, range_ as range\n",
     "from mlir.extras.runtime.passes import Pipeline, run_pipeline\n",
     "from mlir.extras.runtime.refbackend import LLVMJITBackend\n",
+    "from mlir.ir import StridedLayoutAttr\n",
     "\n",
     "# you need this to register the memref value caster\n",
     "# noinspection PyUnresolvedReferences\n",
@@ -91,15 +92,15 @@
    "outputs": [],
    "source": [
     "K = 10\n",
-    "memref_i64 = T.memref(K, K, T.i64)\n",
+    "memref_i64 = T.memref(K, K, T.i64())\n",
     "\n",
     "@func(emit=True)\n",
     "@canonicalize(using=scf)\n",
     "def memfoo(A: memref_i64, B: memref_i64, C: memref_i64):\n",
     "    one = constant(1)\n",
     "    two = constant(2)\n",
     "    if one > two:\n",
-    "        C[0, 0] = constant(3, T.i64)\n",
+    "        C[0, 0] = constant(3, T.i64())\n",
     "    else:\n",
     "        for i in range(0, K):\n",
     "            for j in range(0, K):\n",
@@ -447,8 +448,9 @@
     "D = 32\n",
     "\n",
     "F = K // D\n",
-    "ranked_memref_kxk_f32 = T.memref(K, K, T.f32)\n",
-    "ranked_memref_dxd_f32 = T.memref(D, D, T.f32, layout=((K, 1), S))\n",
+    "ranked_memref_kxk_f32 = T.memref(K, K, T.f32())\n",
+    "layout = StridedLayoutAttr.get(S, (K, 1))\n",
+    "ranked_memref_dxd_f32 = T.memref(D, D, T.f32(), layout=layout)\n",
     "\n",
     "@func(emit=True)\n",
     "@canonicalize(using=scf)\n",
@@ -784,8 +786,9 @@
     "ctx_man = mlir_mod_ctx()\n",
     "ctx = ctx_man.__enter__()\n",
     "\n",
-    "ranked_memref_kxk_f32 = T.memref(K, K, T.f32)\n",
-    "ranked_memref_dxd_f32 = T.memref(D, D, T.f32, layout=((K, 1), S))\n",
+    "ranked_memref_kxk_f32 = T.memref(K, K, T.f32())\n",
+    "layout = StridedLayoutAttr.get(S, (K, 1))\n",
+    "ranked_memref_dxd_f32 = T.memref(D, D, T.f32(), layout=layout)\n",
     "\n",
     "from mlir.extras.dialects.ext import linalg\n",
     "\n",

examples/mwe.py

@@ -15,7 +15,7 @@
 
 def setting_memref(ctx: MLIRContext, backend: LLVMJITBackend):
     K = 10
-    memref_i64 = T.memref(K, K, T.i64)
+    memref_i64 = T.memref(K, K, T.i64())
 
     @func
     @canonicalize(using=scf)

mlir/extras/dialects/ext/arith.py

@@ -43,9 +43,8 @@
     FloatAttr,
 )
 
-from ...util import get_user_code_loc
+from ...util import get_user_code_loc, infer_mlir_type, mlir_type_to_np_dtype
 from ...._mlir_libs._mlir import register_value_caster
-from ...types import infer_mlir_type, mlir_type_to_np_dtype
 
 
 def constant(

mlir/extras/dialects/ext/gpu.py

@@ -227,7 +227,7 @@ def __init__(
             loc=loc,
             ip=ip,
         )
-        self.regions[0].blocks.append(*[T.index for _ in range(12)])
+        self.regions[0].blocks.append(*[T.index() for _ in range(12)])
 
 
 def launch_(

mlir/extras/dialects/ext/llvm.py

@@ -0,0 +1,5 @@
+from ....ir import Type
+
+
+def llvm_ptr_t():
+    return Type.parse("!llvm.ptr")

mlir/extras/dialects/ext/scf.py

@@ -46,10 +46,17 @@
     IndexType,
     _denseI64ArrayAttr,
     Attribute,
+    OpaqueType,
 )
 
 logger = logging.getLogger(__name__)
 
+opaque = lambda dialect_namespace, buffer: OpaqueType.get(dialect_namespace, buffer)
+
+
+def placeholder_opaque_t():
+    return opaque("scf", "placeholder")
+
 
 def _for(
     start,
@@ -372,7 +379,7 @@ def yield_(*args):
             unpacked_args = list(unpacked_args[0])
 
         for i, r in enumerate(results):
-            if r.type == T.placeholder_opaque():
+            if r.type == placeholder_opaque_t():
                 r.set_type(unpacked_args[i].type)
 
         if len(results) > 1:
@@ -582,7 +589,7 @@ def visit_If(self, updated_node: ast.If) -> ast.With | list[ast.With, ast.With]:
             and isinstance(last_statement.targets[0], ast.Tuple)
             else 0,
         )
-        results = [ast_call(T.placeholder_opaque.__name__) for _ in range(num_results)]
+        results = [ast_call(placeholder_opaque_t.__name__) for _ in range(num_results)]
         results = ast.fix_missing_locations(
             ast.copy_location(ast.Tuple(results, ctx=ast.Load()), test)
         )
@@ -620,7 +627,7 @@ def patch_bytecode(self, code: ConcreteBytecode, f):
         f.__globals__[yield_.__name__] = yield_
         f.__globals__[if_ctx_manager.__name__] = if_ctx_manager
         f.__globals__[else_ctx_manager.__name__] = else_ctx_manager
-        f.__globals__[T.placeholder_opaque.__name__] = T.placeholder_opaque
+        f.__globals__[placeholder_opaque_t.__name__] = placeholder_opaque_t
         return code
 
 

mlir/extras/dialects/ext/transform.py

@@ -13,6 +13,8 @@
     SequenceOp,
     FailurePropagationMode,
     YieldOp,
+    AnyOpType,
+    OperationType,
 )
 from ....dialects.transform.loop import LoopUnrollOp
 from ....dialects.transform import GetParentOp
@@ -23,6 +25,18 @@
     StringAttr,
 )
 from ....dialects._ods_common import get_op_result_or_op_results
+from ....dialects import pdl
+
+
+pdl_operation_t = lambda: pdl.OperationType.get()
+
+
+def transform_any_op_t():
+    return AnyOpType.get()
+
+
+def transform_op_t(name):
+    return OperationType.get(name)
 
 
 def sequence_(
@@ -40,7 +54,7 @@ def sequence_(
     if results_ is None:
         results_ = []
     if target is None:
-        target = T.pdl_operation
+        target = pdl_operation_t()
     # this is a misnomer - it's not about targeting a particular op
     # but about picking which transform sequence runs using
     # transform_dialect_interpreter(debug_transform_root_tag="")
@@ -52,7 +66,7 @@ def sequence_(
         failure_propagation_mode = FailurePropagationMode.Propagate
 
     if isinstance(target, str):
-        target = T.transform_op(target)
+        target = transform_op_t(target)
 
     seq_op = SequenceOp(
         failure_propagation_mode,
@@ -85,7 +99,7 @@ def get_parent(
 
     return get_op_result_or_op_results(
         GetParentOp(
-            T.pdl_operation,
+            pdl_operation_t(),
             target,
             isolated_from_above=isolated_from_above,
             op_name=op_name,
@@ -119,7 +133,7 @@ def match(
         loc = get_user_code_loc()
     return get_op_result_or_op_results(
         MatchOp(
-            T.transform_any_op,
+            transform_any_op_t(),
             target,
             ops=ops,
             interface=interface,

mlir/extras/runtime/refbackend.py

@@ -25,7 +25,7 @@
 from .. import types as T
 from ...dialects.memref import cast
 from ..runtime.passes import Pipeline, run_pipeline
-from ..types import (
+from ..util import (
     memref_type_to_np_dtype,
     mlir_type_to_ctype,
     np_dtype_to_mlir_type,