refactor coerce and constant

Author: makslevental

mlir_utils/dialects/ext/arith.py

@@ -6,6 +6,7 @@
 
 import numpy as np
 from mlir.dialects import arith as arith_dialect
+from mlir.dialects import complex as complex_dialect
 from mlir.dialects._arith_ops_ext import _is_integer_like_type
 from mlir.dialects._ods_common import get_op_result_or_value
 from mlir.dialects.linalg.opdsl.lang.emitter import (
@@ -19,6 +20,7 @@
     Context,
     DenseElementsAttr,
     IndexType,
+    InsertionPoint,
     IntegerAttr,
     IntegerType,
     Location,
@@ -28,9 +30,20 @@
     Type,
     Value,
     register_attribute_builder,
+    ComplexType,
+    BF16Type,
+    F16Type,
+    F32Type,
+    F64Type,
+    FloatAttr,
 )
 
-from mlir_utils.util import get_result_or_results, maybe_cast, get_user_code_loc
+from mlir_utils.util import (
+    get_result_or_results,
+    maybe_cast,
+    get_user_code_loc,
+    register_value_caster,
+)
 
 try:
     from mlir_utils.dialects.arith import *
@@ -46,7 +59,8 @@ def constant(
     index: Optional[bool] = None,
     *,
     loc: Location = None,
-) -> arith_dialect.ConstantOp:
+    ip: InsertionPoint = None,
+) -> Value:
     """Instantiate arith.constant with value `value`.
 
     Args:
@@ -67,21 +81,62 @@ def constant(
         type = IndexType.get()
     if type is None:
         type = infer_mlir_type(value)
-    elif RankedTensorType.isinstance(type) and isinstance(value, (int, float, bool)):
+
+    assert type is not None
+
+    if _is_complex_type(type):
+        value = complex(value)
+        return maybe_cast(
+            get_result_or_results(
+                complex_dialect.ConstantOp(
+                    type,
+                    list(
+                        map(
+                            lambda x: FloatAttr.get(type.element_type, x),
+                            [value.real, value.imag],
+                        )
+                    ),
+                    loc=loc,
+                    ip=ip,
+                )
+            )
+        )
+
+    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)
-        value = np.ones(
+        value = np.full(
             ranked_tensor_type.shape,
+            value,
             dtype=mlir_type_to_np_dtype(ranked_tensor_type.element_type),
         )
-    assert type is not None
 
     if isinstance(value, np.ndarray):
         value = DenseElementsAttr.get(
             value,
             type=type,
         )
+
     return maybe_cast(
-        get_result_or_results(arith_dialect.ConstantOp(type, value, loc=loc))
+        get_result_or_results(arith_dialect.ConstantOp(type, value, loc=loc, ip=ip))
+    )
+
+
+def index_cast(
+    value: Value,
+    *,
+    to: Type = None,
+    loc: Location = None,
+    ip: InsertionPoint = None,
+) -> Value:
+    if loc is None:
+        loc = get_user_code_loc()
+    if to is None:
+        to = IndexType.get()
+    return maybe_cast(
+        get_result_or_results(arith_dialect.IndexCastOp(to, value, loc=loc, ip=ip))
     )
 
 
@@ -231,6 +286,7 @@ def _binary_op(
     rhs: "ArithValue",
     op: str,
     predicate: str = None,
+    signedness: str = None,
     *,
     loc: Location = None,
 ) -> "ArithValue":
@@ -247,12 +303,15 @@ def _binary_op(
     """
     if loc is None:
         loc = get_user_code_loc()
-    if not isinstance(rhs, lhs.__class__):
+    if (
+        isinstance(rhs, Value)
+        and lhs.type != rhs.type
+        or isinstance(rhs, (float, int, bool, np.ndarray))
+    ):
         lhs, rhs = lhs.coerce(rhs)
-    if lhs.type != rhs.type:
-        raise ValueError(f"{lhs=} {rhs=} must have the same type.")
+    assert lhs.type == rhs.type, f"{lhs=} {rhs=} must have the same type."
 
-    assert op in {"add", "sub", "mul", "cmp", "truediv", "floordiv", "mod"}
+    assert op in {"add", "and", "or", "sub", "mul", "cmp", "truediv", "floordiv", "mod"}
 
     if op == "cmp":
         assert predicate is not None
@@ -301,15 +360,20 @@ def _binary_op(
         elif _is_integer_like_type(lhs.dtype):
             # eq, ne signs don't matter
             if predicate not in {"eq", "ne"}:
-                if lhs.dtype.is_signed:
-                    predicate = "s" + predicate
+                if signedness is not None:
+                    predicate = signedness + predicate
                 else:
-                    predicate = "u" + predicate
+                    if lhs.dtype.is_signed:
+                        predicate = "s" + predicate
+                    else:
+                        predicate = "u" + predicate
         return lhs.__class__(op(predicate, lhs, rhs, loc=loc), dtype=lhs.dtype)
     else:
         return lhs.__class__(op(lhs, rhs, loc=loc), dtype=lhs.dtype)
 
 
+# TODO(max): these could be generic in the dtype
+# TODO(max): hit .verify() before constructing (maybe)
 class ArithValue(Value, metaclass=ArithValueMeta):
     """Class for functionality shared by Value subclasses that support
     arithmetic operations.
@@ -363,6 +427,9 @@ def __repr__(self):
     __rsub__ = partialmethod(_binary_op, op="sub")
     __rmul__ = partialmethod(_binary_op, op="mul")
 
+    __and__ = partialmethod(_binary_op, op="and")
+    __or__ = partialmethod(_binary_op, op="or")
+
     def __eq__(self, other):
         if not isinstance(other, self.__class__):
             try:
@@ -435,6 +502,14 @@ def __float__(self):
     def coerce(self, other) -> tuple["Scalar", "Scalar"]:
         if isinstance(other, (int, float, bool)):
             other = Scalar(other, dtype=self.dtype)
+        elif isinstance(other, Scalar) and _is_index_type(self.type):
+            other = index_cast(other)
+        elif isinstance(other, Scalar) and _is_index_type(other.type):
+            other = index_cast(other, to=self.type)
         else:
-            raise ValueError(f"can't coerce {other=} to Scalar")
+            raise ValueError(f"can't coerce {other=} to {self=}")
         return self, other
+
+
+for t in [BF16Type, F16Type, F32Type, F64Type, IndexType, IntegerType, ComplexType]:
+    register_value_caster(t.static_typeid)(Scalar)

mlir_utils/dialects/ext/scf.py

@@ -5,6 +5,7 @@
 from typing import Optional, Sequence
 
 from bytecode import ConcreteBytecode, ConcreteInstr
+from mlir.dialects.linalg.opdsl.lang.emitter import _is_index_type
 from mlir.dialects.scf import IfOp, ForOp
 from mlir.ir import InsertionPoint, Value, OpResultList, OpResult
 
@@ -16,7 +17,7 @@
     OpCode,
 )
 from mlir_utils.ast.util import ast_call, set_lineno
-from mlir_utils.dialects.ext.arith import constant
+from mlir_utils.dialects.ext.arith import constant, index_cast
 from mlir_utils.dialects.scf import yield_ as yield__
 from mlir_utils.util import (
     region_op,
@@ -43,15 +44,17 @@ def _for(
     if stop is None:
         stop = start
         start = 0
-    if isinstance(start, int):
-        start = constant(start, index=True)
-    if isinstance(stop, int):
-        stop = constant(stop, index=True)
-    if isinstance(step, int):
-        step = constant(step, index=True)
+    params = [start, stop, step]
+    for i, p in enumerate(params):
+        if isinstance(p, int):
+            p = constant(p, index=True)
+        if not _is_index_type(p.type):
+            p = index_cast(p)
+        params[i] = p
+
     if loc is None:
         loc = get_user_code_loc()
-    return ForOp(start, stop, step, iter_args, loc=loc, ip=ip)
+    return ForOp(*params, iter_args, loc=loc, ip=ip)
 
 
 for_ = region_op(_for, terminator=yield__)

mlir_utils/dialects/ext/tensor.py

@@ -140,7 +140,7 @@ def __getitem__(self, idx: tuple) -> "Tensor":
         if isinstance(idx, tuple) and all(i == slice(None) for i in idx):
             return self
         if idx is None:
-            return _expand_dims(self, (0,))
+            return expand_dims(self, (0,))
 
         idx = list((idx,) if isinstance(idx, int) else idx)
         for i, d in enumerate(idx):
@@ -198,8 +198,10 @@ def coerce(self, other) -> tuple["Tensor", "Tensor"]:
             if isinstance(other, (int, float)):
                 other = Tensor(np.full(self.shape, other), dtype=self.dtype)
                 return self, other
-            elif _is_scalar(other):
-                other = tensor.splat(self.type, other)
+            elif isinstance(other, Scalar):
+                other = tensor.splat(
+                    RankedTensorType.get(self.shape, other.dtype), other
+                )
                 return self, other
 
         raise ValueError(f"can't coerce unknown {other=}")
@@ -256,7 +258,7 @@ def static_strides(self):
         return tuple(strides)
 
 
-def _expand_dims(inp, newaxis_dims) -> Tensor:
+def expand_dims(inp, newaxis_dims) -> Tensor:
     """Expand the shape of a tensor.
 
     Insert a new axis that will appear at the `axis` position in the expanded
@@ -514,7 +516,7 @@ def _extract_slice(
         raise ValueError(f"non-constant indices not supported {indexer}")
 
     # This adds newaxis/None dimensions.
-    return _expand_dims(out, indexer.newaxis_dims)
+    return expand_dims(out, indexer.newaxis_dims)
 
 
 def _insert_slice(
@@ -523,7 +525,7 @@ def _insert_slice(
     idx,
 ):
     if isinstance(source, Scalar):
-        source = _expand_dims(source, (0,))
+        source = expand_dims(source, (0,))
 
     indexer = _indices_to_indexer(idx, dest.shape)