arith constant vec (#65)

Author: makslevental

examples/mwe.py

@@ -11,7 +11,7 @@
 # noinspection PyUnresolvedReferences
 import mlir.extras.dialects.ext.memref
 from mlir.extras.context import RAIIMLIRContext, ExplicitlyManagedModule
-from mlir.extras.dialects.ext.bufferization import LayoutMapOption
+from mlir.dialects.bufferization import LayoutMapOption
 from mlir.dialects.transform.vector import (
     VectorContractLowering,
     VectorMultiReductionLowering,

examples/vectorization_e2e.ipynb

@@ -57,9 +57,10 @@
     "from mlir.ir import StringAttr, UnitAttr\n",
     "\n",
     "# you need this to register the memref value caster\n",
+    "# noinspection PyUnresolvedReferences\n",
     "import mlir.extras.dialects.ext.memref\n",
     "from mlir.extras.context import RAIIMLIRContext, ExplicitlyManagedModule\n",
-    "from mlir.extras.dialects.ext.bufferization import LayoutMapOption\n",
+    "from mlir.dialects.bufferization import LayoutMapOption\n",
     "from mlir.dialects.transform.vector import (\n",
     "    VectorContractLowering,\n",
     "    VectorMultiReductionLowering,\n",

mlir/extras/context.py

@@ -1,5 +1,4 @@
 import contextlib
-import warnings
 from contextlib import ExitStack, contextmanager
 from dataclasses import dataclass
 from typing import Optional

mlir/extras/dialects/ext/_shaped_value.py

@@ -0,0 +1,40 @@
+from functools import cached_property, reduce
+from typing import Tuple
+
+import numpy as np
+
+from ....ir import DenseElementsAttr, ShapedType, Type
+
+S = ShapedType.get_dynamic_size()
+
+
+# mixin that requires `is_constant`
+class ShapedValue:
+    @cached_property
+    def literal_value(self) -> np.ndarray:
+        if not self.is_constant:
+            raise ValueError("Can't build literal from non-constant value")
+        return np.array(DenseElementsAttr(self.owner.opview.value), copy=False)
+
+    @cached_property
+    def _shaped_type(self) -> ShapedType:
+        return ShapedType(self.type)
+
+    def has_static_shape(self) -> bool:
+        return self._shaped_type.has_static_shape
+
+    def has_rank(self) -> bool:
+        return self._shaped_type.has_rank
+
+    @cached_property
+    def shape(self) -> Tuple[int, ...]:
+        return tuple(self._shaped_type.shape)
+
+    @cached_property
+    def n_elements(self) -> int:
+        assert self.has_static_shape()
+        return reduce(lambda acc, v: acc * v, self._shaped_type.shape, 1)
+
+    @cached_property
+    def dtype(self) -> Type:
+        return self._shaped_type.element_type

mlir/extras/dialects/ext/arith.py

@@ -1,50 +1,44 @@
 import operator
 from abc import abstractmethod
 from copy import deepcopy
-from functools import partialmethod, cached_property
-from typing import Union, Optional, Tuple
-
-import numpy as np
+from functools import cached_property, partialmethod
+from typing import Optional, Tuple
 
 from ...util import get_user_code_loc, infer_mlir_type, mlir_type_to_np_dtype
 from ...._mlir_libs._mlir import register_value_caster
-from ....dialects import arith as arith_dialect
-from ....dialects import complex as complex_dialect
+from ....dialects import arith as arith_dialect, complex as complex_dialect
 from ....dialects._arith_enum_gen import (
     _arith_cmpfpredicateattr,
-    CmpFPredicate,
-    CmpIPredicate,
     _arith_cmpipredicateattr,
 )
-from ....dialects._ods_common import get_op_result_or_value, get_op_result_or_op_results
+from ....dialects._ods_common import get_op_result_or_op_results, get_op_result_or_value
 from ....dialects.arith import *
 from ....dialects.arith import _is_integer_like_type
 from ....dialects.linalg.opdsl.lang.emitter import (
-    _is_floating_point_type,
-    _is_integer_type,
     _is_complex_type,
+    _is_floating_point_type,
     _is_index_type,
 )
 from ....ir import (
     Attribute,
+    BF16Type,
+    ComplexType,
     Context,
     DenseElementsAttr,
+    F16Type,
+    F32Type,
+    F64Type,
+    FloatAttr,
     IndexType,
     InsertionPoint,
     IntegerType,
     Location,
     OpView,
     Operation,
-    RankedTensorType,
+    ShapedType,
     Type,
     Value,
     register_attribute_builder,
-    ComplexType,
-    BF16Type,
-    F16Type,
-    F32Type,
-    F64Type,
-    FloatAttr,
 )
 
 
@@ -53,6 +47,7 @@ def constant(
     type: Optional[Type] = None,
     index: Optional[bool] = None,
     *,
+    vector: Optional[bool] = False,
     loc: Location = None,
     ip: InsertionPoint = None,
 ) -> Value:
@@ -75,7 +70,7 @@ def constant(
     if index is not None and index:
         type = IndexType.get()
     if type is None:
-        type = infer_mlir_type(value)
+        type = infer_mlir_type(value, vector=vector)
 
     assert type is not None
 
@@ -98,8 +93,8 @@ def constant(
     if _is_floating_point_type(type) and not isinstance(value, np.ndarray):
         value = float(value)
 
-    if RankedTensorType.isinstance(type) and isinstance(value, (int, float, bool)):
-        ranked_tensor_type = RankedTensorType(type)
+    if ShapedType.isinstance(type) and isinstance(value, (int, float, bool)):
+        ranked_tensor_type = ShapedType(type)
         value = np.full(
             ranked_tensor_type.shape,
             value,
@@ -403,7 +398,7 @@ def __str__(self):
         return f"{self.__class__.__name__}({self.get_name()}, {self.type})"
 
     def __repr__(self):
-        return str(self)
+        return str(Value(self)).replace("Value", self.__class__.__name__)
 
     # partialmethod differs from partial in that it also binds the object instance
     # to the first arg (i.e., self)
@@ -468,16 +463,6 @@ class Scalar(ArithValue):
     def dtype(self) -> Type:
         return self.type
 
-    @staticmethod
-    def isinstance(other: Value):
-        return (
-            isinstance(other, Value)
-            and _is_integer_type(other.type)
-            or _is_floating_point_type(other.type)
-            or _is_index_type(other.type)
-            or _is_complex_type(other.type)
-        )
-
     @cached_property
     def literal_value(self) -> Union[int, float, bool]:
         if not self.is_constant():

mlir/extras/dialects/ext/bufferization.py

@@ -1,12 +1,12 @@
-from typing import Union, List
+from typing import List, Union
 
-from ...util import get_user_code_loc, Successor
+from ...util import Successor, get_user_code_loc
 from ....dialects._cf_ops_gen import _Dialect
 from ....dialects._ods_common import (
     _cext,
 )
 from ....dialects.cf import *
-from ....ir import Value, InsertionPoint, Block
+from ....ir import Block, InsertionPoint, Value
 
 
 @_cext.register_operation(_Dialect, replace=True)

mlir/extras/dialects/ext/cf.py

@@ -3,17 +3,17 @@
 
 from ...meta import op_region_builder
 from ...util import get_user_code_loc, make_maybe_no_args_decorator
-from ....dialects.func import *
 from ....dialects._ods_common import get_op_result_or_op_results
+from ....dialects.func import *
 from ....ir import (
     FlatSymbolRefAttr,
     FunctionType,
     InsertionPoint,
+    OpView,
+    Operation,
     Type,
     TypeAttr,
     Value,
-    Operation,
-    OpView,
 )
 
 

mlir/extras/dialects/ext/func.py

@@ -1,28 +1,31 @@
 import inspect
 from functools import partial
-from typing import Optional, Any, List, Tuple
+from typing import Any, List, Optional, Tuple
 
 from .arith import constant
 from .func import FuncBase
 from ... import types as T
 from ...meta import (
     region_op,
 )
-from ...util import get_user_code_loc, make_maybe_no_args_decorator, ModuleMeta
+from ...util import ModuleMeta, get_user_code_loc, make_maybe_no_args_decorator
 from ....dialects._gpu_ops_gen import _Dialect
-from ....dialects._ods_common import get_default_loc_context, _cext
-from ....dialects._ods_common import get_op_result_or_op_results
+from ....dialects._ods_common import (
+    _cext,
+    get_default_loc_context,
+    get_op_result_or_op_results,
+)
 from ....dialects.gpu import *
 from ....ir import (
-    Type,
-    Attribute,
+    ArrayAttr,
     AttrBuilder,
-    UnitAttr,
-    register_attribute_builder,
+    Attribute,
     Context,
-    ArrayAttr,
     InsertionPoint,
+    Type,
+    UnitAttr,
     Value,
+    register_attribute_builder,
 )
 
 

mlir/extras/dialects/ext/gpu.py

@@ -1,4 +1,7 @@
 from ...util import get_user_code_loc
+
+# noinspection PyUnresolvedReferences
+from ....dialects.linalg import *
 from ....dialects import linalg