add shaped type wrappers

Author: makslevental

mlir_utils/types.py

@@ -1,14 +1,20 @@
+from functools import partial
 from typing import Union
 
 import numpy as np
 from mlir.ir import (
-    IntegerType,
-    F64Type,
-    RankedTensorType,
-    IndexType,
+    Attribute,
     F16Type,
     F32Type,
+    F64Type,
+    IndexType,
+    IntegerType,
+    MemRefType,
+    RankedTensorType,
     Type,
+    UnrankedMemRefType,
+    UnrankedTensorType,
+    VectorType,
 )
 
 index_t = IndexType.get()
@@ -66,15 +72,55 @@ def infer_mlir_type(
         )
 
 
-def tensor_t(*args, element_type: Type = None):
-    if (element_type is None and not isinstance(args[-1], Type)) or (
-        isinstance(args[-1], Type) and element_type is not None
+def shaped_t(*args, element_type: Type = None, type_constructor=None):
+    if type_constructor is None:
+        raise ValueError("shaped_t is an abstract base class - cannot be constructed")
+    if (element_type is None and args and not isinstance(args[-1], Type)) or (
+        args and isinstance(args[-1], Type) and element_type is not None
     ):
         raise ValueError(
             f"either element_type must be provided explicitly XOR last arg to tensor type constructor must be the element type"
         )
     if element_type is not None:
         type = element_type
+        sizes = args
     else:
         type = args[-1]
-    return RankedTensorType.get(args[:-1], type)
+        sizes = args[:-1]
+    if sizes:
+        return type_constructor(sizes, type)
+    else:
+        return type_constructor(type)
+
+
+def vector_t(*args, element_type: Type = None):
+    return shaped_t(*args, element_type=element_type, type_constructor=VectorType.get)
+
+
+def tensor_t(*args, element_type: Type = None):
+    if not len(args) or len(args) == 1 and isinstance(args[-1], Type):
+        return shaped_t(
+            *args, element_type=element_type, type_constructor=UnrankedTensorType.get
+        )
+    else:
+        return shaped_t(
+            *args, element_type=element_type, type_constructor=RankedTensorType.get
+        )
+
+
+def memref_t(*args, element_type: Type = None, memory_space: int = None):
+    if memory_space is None:
+        memory_space = 0
+    memory_space = Attribute.parse(str(memory_space))
+    if not len(args) or len(args) == 1 and isinstance(args[-1], Type):
+        return shaped_t(
+            *args,
+            element_type=element_type,
+            type_constructor=partial(UnrankedMemRefType.get, memory_space=memory_space),
+        )
+    else:
+        return shaped_t(
+            *args,
+            element_type=element_type,
+            type_constructor=partial(MemRefType.get, memory_space=memory_space),
+        )

tests/test_types.py

@@ -0,0 +1,33 @@
+import pytest
+
+from mlir_utils.dialects.ext.tensor import S
+
+# noinspection PyUnresolvedReferences
+from mlir_utils.testing import mlir_ctx as ctx, filecheck, MLIRContext
+from mlir_utils.types import f64_t, tensor_t, memref_t, vector_t
+
+# needed since the fix isn't defined here nor conftest.py
+pytest.mark.usefixtures("ctx")
+
+
+def test_shaped_types(ctx: MLIRContext):
+    t = tensor_t(S, 3, S, f64_t)
+    assert repr(t) == "RankedTensorType(tensor<?x3x?xf64>)"
+    ut = tensor_t(f64_t)
+    assert repr(ut) == "UnrankedTensorType(tensor<*xf64>)"
+    t = tensor_t(S, 3, S, element_type=f64_t)
+    assert repr(t) == "RankedTensorType(tensor<?x3x?xf64>)"
+    ut = tensor_t(element_type=f64_t)
+    assert repr(ut) == "UnrankedTensorType(tensor<*xf64>)"
+
+    m = memref_t(S, 3, S, f64_t)
+    assert repr(m) == "MemRefType(memref<?x3x?xf64>)"
+    um = memref_t(f64_t)
+    assert repr(um) == "UnrankedMemRefType(memref<*xf64>)"
+    m = memref_t(S, 3, S, element_type=f64_t)
+    assert repr(m) == "MemRefType(memref<?x3x?xf64>)"
+    um = memref_t(element_type=f64_t)
+    assert repr(um) == "UnrankedMemRefType(memref<*xf64>)"
+
+    v = vector_t(3, 3, 3, f64_t)
+    assert repr(v) == "VectorType(vector<3x3x3xf64>)"