[Mlir-commits] [mlir] [mlir][python] add type wrappers (PR #71218)
Maksim Levental
llvmlistbot at llvm.org
Tue Nov 14 08:46:14 PST 2023
https://github.com/makslevental updated https://github.com/llvm/llvm-project/pull/71218
>From 992283313f3c32219fa306ec0bc899a741b12239 Mon Sep 17 00:00:00 2001
From: max <maksim.levental at gmail.com>
Date: Thu, 2 Nov 2023 11:15:45 -0500
Subject: [PATCH] [mlir][python] add type wrappers
---
mlir/lib/Bindings/Python/IRTypes.cpp | 24 ++--
mlir/python/CMakeLists.txt | 1 +
mlir/python/mlir/types.py | 176 +++++++++++++++++++++++++++
mlir/test/python/ir/builtin_types.py | 101 +++++++++++++++
4 files changed, 286 insertions(+), 16 deletions(-)
create mode 100644 mlir/python/mlir/types.py
diff --git a/mlir/lib/Bindings/Python/IRTypes.cpp b/mlir/lib/Bindings/Python/IRTypes.cpp
index 483db673f989e6b..56e895d3053796e 100644
--- a/mlir/lib/Bindings/Python/IRTypes.cpp
+++ b/mlir/lib/Bindings/Python/IRTypes.cpp
@@ -463,7 +463,7 @@ class PyVectorType : public PyConcreteType<PyVectorType, PyShapedType> {
static void bindDerived(ClassTy &c) {
c.def_static("get", &PyVectorType::get, py::arg("shape"),
- py::arg("elementType"), py::kw_only(),
+ py::arg("element_type"), py::kw_only(),
py::arg("scalable") = py::none(),
py::arg("scalable_dims") = py::none(),
py::arg("loc") = py::none(), "Create a vector type")
@@ -689,13 +689,9 @@ class PyTupleType : public PyConcreteType<PyTupleType> {
static void bindDerived(ClassTy &c) {
c.def_static(
"get_tuple",
- [](py::list elementList, DefaultingPyMlirContext context) {
- intptr_t num = py::len(elementList);
- // Mapping py::list to SmallVector.
- SmallVector<MlirType, 4> elements;
- for (auto element : elementList)
- elements.push_back(element.cast<PyType>());
- MlirType t = mlirTupleTypeGet(context->get(), num, elements.data());
+ [](std::vector<MlirType> elements, DefaultingPyMlirContext context) {
+ MlirType t = mlirTupleTypeGet(context->get(), elements.size(),
+ elements.data());
return PyTupleType(context->getRef(), t);
},
py::arg("elements"), py::arg("context") = py::none(),
@@ -727,13 +723,11 @@ class PyFunctionType : public PyConcreteType<PyFunctionType> {
static void bindDerived(ClassTy &c) {
c.def_static(
"get",
- [](std::vector<PyType> inputs, std::vector<PyType> results,
+ [](std::vector<MlirType> inputs, std::vector<MlirType> results,
DefaultingPyMlirContext context) {
- SmallVector<MlirType, 4> inputsRaw(inputs.begin(), inputs.end());
- SmallVector<MlirType, 4> resultsRaw(results.begin(), results.end());
- MlirType t = mlirFunctionTypeGet(context->get(), inputsRaw.size(),
- inputsRaw.data(), resultsRaw.size(),
- resultsRaw.data());
+ MlirType t =
+ mlirFunctionTypeGet(context->get(), inputs.size(), inputs.data(),
+ results.size(), results.data());
return PyFunctionType(context->getRef(), t);
},
py::arg("inputs"), py::arg("results"), py::arg("context") = py::none(),
@@ -742,7 +736,6 @@ class PyFunctionType : public PyConcreteType<PyFunctionType> {
"inputs",
[](PyFunctionType &self) {
MlirType t = self;
- auto contextRef = self.getContext();
py::list types;
for (intptr_t i = 0, e = mlirFunctionTypeGetNumInputs(self); i < e;
++i) {
@@ -754,7 +747,6 @@ class PyFunctionType : public PyConcreteType<PyFunctionType> {
c.def_property_readonly(
"results",
[](PyFunctionType &self) {
- auto contextRef = self.getContext();
py::list types;
for (intptr_t i = 0, e = mlirFunctionTypeGetNumResults(self); i < e;
++i) {
diff --git a/mlir/python/CMakeLists.txt b/mlir/python/CMakeLists.txt
index 971ad2dd214a15f..12e2dab60f3011b 100644
--- a/mlir/python/CMakeLists.txt
+++ b/mlir/python/CMakeLists.txt
@@ -21,6 +21,7 @@ declare_mlir_python_sources(MLIRPythonSources.Core.Python
_mlir_libs/__init__.py
ir.py
passmanager.py
+ types.py
dialects/_ods_common.py
# The main _mlir module has submodules: include stubs from each.
diff --git a/mlir/python/mlir/types.py b/mlir/python/mlir/types.py
new file mode 100644
index 000000000000000..c57bef924078abe
--- /dev/null
+++ b/mlir/python/mlir/types.py
@@ -0,0 +1,176 @@
+from functools import partial
+from typing import Optional, List
+
+from .ir import (
+ Attribute,
+ BF16Type,
+ ComplexType,
+ Context,
+ F16Type,
+ F32Type,
+ F64Type,
+ Float8E4M3B11FNUZType,
+ Float8E4M3FNType,
+ Float8E5M2Type,
+ FunctionType,
+ IndexType,
+ IntegerType,
+ MemRefType,
+ NoneType,
+ OpaqueType,
+ RankedTensorType,
+ StridedLayoutAttr,
+ StringAttr,
+ TupleType,
+ Type,
+ UnrankedMemRefType,
+ UnrankedTensorType,
+ VectorType,
+)
+
+__all__ = []
+
+_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)
+
+_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)
+
+_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()
+
+_cmp16 = lambda: ComplexType.get(_f16())
+_cmp32 = lambda: ComplexType.get(_f32())
+_cmp64 = lambda: ComplexType.get(_f64())
+
+_none = lambda: NoneType.get()
+
+
+def _opaque(dialect_namespace, type_data):
+ return OpaqueType.get(dialect_namespace, type_data)
+
+
+def _shaped(*shape, element_type: Type = None, type_constructor=None):
+ if type_constructor is None:
+ raise ValueError("shaped is an abstract base class - cannot be constructed.")
+ if (element_type is None and shape and not isinstance(shape[-1], Type)) or (
+ shape and isinstance(shape[-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 = shape
+ else:
+ type = shape[-1]
+ sizes = shape[:-1]
+ if sizes:
+ return type_constructor(sizes, type)
+ else:
+ return type_constructor(type)
+
+
+def _vector(
+ *shape,
+ element_type: Type = None,
+ scalable: Optional[List[bool]] = None,
+ scalable_dims: Optional[List[int]] = None,
+):
+ return _shaped(
+ *shape,
+ element_type=element_type,
+ type_constructor=partial(
+ VectorType.get, scalable=scalable, scalable_dims=scalable_dims
+ ),
+ )
+
+
+def _tensor(*shape, element_type: Type = None, encoding: Optional[str] = None):
+ if encoding is not None:
+ encoding = StringAttr.get(encoding)
+ if not len(shape) or len(shape) == 1 and isinstance(shape[-1], Type):
+ if encoding is not None:
+ raise ValueError("UnrankedTensorType does not support encoding.")
+ return _shaped(
+ *shape, element_type=element_type, type_constructor=UnrankedTensorType.get
+ )
+ else:
+ return _shaped(
+ *shape,
+ element_type=element_type,
+ type_constructor=partial(RankedTensorType.get, encoding=encoding),
+ )
+
+
+def _memref(
+ *shape,
+ element_type: Type = None,
+ memory_space: Optional[int] = None,
+ layout: Optional[StridedLayoutAttr] = None,
+):
+ if memory_space is not None:
+ memory_space = Attribute.parse(str(memory_space))
+ if not len(shape) or len(shape) == 1 and isinstance(shape[-1], Type):
+ return _shaped(
+ *shape,
+ element_type=element_type,
+ type_constructor=partial(UnrankedMemRefType.get, memory_space=memory_space),
+ )
+ return _shaped(
+ *shape,
+ element_type=element_type,
+ type_constructor=partial(
+ MemRefType.get, memory_space=memory_space, layout=layout
+ ),
+ )
+
+
+def _tuple(*elements):
+ return TupleType.get_tuple(elements)
+
+
+def _function(*, inputs, results):
+ return FunctionType.get(inputs, results)
+
+
+def _isa_lambda(v):
+ LAMBDA = lambda: 0
+ return isinstance(v, type(LAMBDA)) and v.__name__ == LAMBDA.__name__
+
+
+def __getattr__(name):
+ if name == "__path__":
+ # If the 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).
+ # https://docs.python.org/3/reference/import.html#path__
+ 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 _isa_lambda(builder):
+ return builder()
+ return builder
+ raise RuntimeError(f"{name} is not a legal type.")
diff --git a/mlir/test/python/ir/builtin_types.py b/mlir/test/python/ir/builtin_types.py
index d4fed86b4f135ee..ec3c8641efb26c5 100644
--- a/mlir/test/python/ir/builtin_types.py
+++ b/mlir/test/python/ir/builtin_types.py
@@ -3,6 +3,7 @@
import gc
from mlir.ir import *
from mlir.dialects import arith, tensor, func, memref
+import mlir.types as T
def run(f):
@@ -772,3 +773,103 @@ def testCustomTypeTypeCaster():
print(t)
# CHECK: OperationType(!transform.op<"foo.bar">)
print(repr(t))
+
+
+# CHECK-LABEL: TEST: testTypeWrappers
+ at 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."
+
+ c1 = T.cmp16
+ c2 = T.cmp32
+ assert repr(c1) == "ComplexType(complex<f16>)"
+ assert repr(c2) == "ComplexType(complex<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)
+ assert repr(m1) == "MemRefType(memref<2x3x4xf64>)"
+
+ 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]))
+ 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))
+ 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)
+ assert repr(t1) == "RankedTensorType(tensor<?x3x?xf64>)"
+ ut1 = T.tensor(T.f64)
+ assert repr(ut1) == "UnrankedTensorType(tensor<*xf64>)"
+ t2 = T.tensor(S, 3, S, element_type=T.f64)
+ assert repr(t2) == "RankedTensorType(tensor<?x3x?xf64>)"
+ ut2 = T.tensor(element_type=T.f64)
+ assert repr(ut2) == "UnrankedTensorType(tensor<*xf64>)"
+
+ 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)
+ assert repr(v) == "VectorType(vector<3x3x3xf64>)"
+
+ m5 = T.memref(S, 3, S, T.f64)
+ assert repr(m5) == "MemRefType(memref<?x3x?xf64>)"
+ um1 = T.memref(T.f64)
+ assert repr(um1) == "UnrankedMemRefType(memref<*xf64>)"
+ m6 = T.memref(S, 3, S, element_type=T.f64)
+ assert repr(m6) == "MemRefType(memref<?x3x?xf64>)"
+ um2 = T.memref(element_type=T.f64)
+ assert repr(um2) == "UnrankedMemRefType(memref<*xf64>)"
+
+ m7 = T.memref(S, 3, S, T.f64)
+ assert repr(m7) == "MemRefType(memref<?x3x?xf64>)"
+ 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])
+ 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])
+ 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)
+ 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))
+ assert repr(func) == "FunctionType((i16, i32, i64) -> (f16, f32, f64))"
More information about the Mlir-commits
mailing list