implement div and promotion

Author: makslevental

mlir_utils/ast/canonicalize.py

@@ -4,7 +4,7 @@
 import inspect
 import logging
 import types
-from abc import ABC
+from abc import ABC, abstractmethod
 from dis import findlinestarts
 from opcode import opmap
 from types import CodeType
@@ -101,7 +101,7 @@ class BytecodePatcher(ABC):
     def __init__(self, context=None):
         self.context = context
 
-    @property
+    @abstractmethod
     def patch_bytecode(self, code: ConcreteBytecode, original_f) -> ConcreteBytecode:
         pass
 
@@ -119,10 +119,12 @@ def patch_bytecode(f, patchers: list[type(BytecodePatcher)] = None):
 
 class Canonicalizer(ABC):
     @property
+    @abstractmethod
     def cst_transformers(self) -> list[StrictTransformer]:
         pass
 
     @property
+    @abstractmethod
     def bytecode_patchers(self) -> list[BytecodePatcher]:
         pass
 

mlir_utils/dialects/ext/arith.py

@@ -1,4 +1,5 @@
 import operator
+from abc import abstractmethod
 from copy import deepcopy
 from functools import partialmethod, cached_property
 from typing import Union, Optional
@@ -184,7 +185,7 @@ def _arith_CmpIPredicateAttr(predicate: str | Attribute, context: Context):
     }
     if isinstance(predicate, Attribute):
         return predicate
-    assert predicate in predicates, f"predicate {predicate} not in predicates"
+    assert predicate in predicates, f"{predicate=} not in predicates"
     return IntegerAttr.get(
         IntegerType.get_signless(64, context=context), predicates[predicate]
     )
@@ -219,7 +220,7 @@ def _arith_CmpFPredicateAttr(predicate: str | Attribute, context: Context):
     }
     if isinstance(predicate, Attribute):
         return predicate
-    assert predicate in predicates, f"predicate {predicate} not in predicates"
+    assert predicate in predicates, f"{predicate=} not in predicates"
     return IntegerAttr.get(
         IntegerType.get_signless(64, context=context), predicates[predicate]
     )
@@ -247,13 +248,14 @@ def _binary_op(
     if loc is None:
         loc = get_user_code_loc()
     if not isinstance(rhs, lhs.__class__):
-        rhs = lhs.__class__(rhs, dtype=lhs.type)
+        lhs, rhs = lhs.coerce(rhs)
+    if lhs.type != rhs.type:
+        raise ValueError(f"{lhs=} {rhs=} must have the same type.")
+
+    assert op in {"add", "sub", "mul", "cmp", "truediv", "floordiv", "mod"}
 
-    assert op in {"add", "sub", "mul", "cmp"}
     if op == "cmp":
         assert predicate is not None
-    if lhs.type != rhs.type:
-        raise ValueError(f"{lhs=} {rhs=} must have the same type.")
 
     if lhs.fold() and lhs.fold():
         klass = lhs.__class__
@@ -267,15 +269,30 @@ def _binary_op(
             op = predicate
         op = operator.attrgetter(op)(operator)
         return klass(op(lhs, rhs), fold=True)
+
+    if op == "truediv":
+        op = "div"
+    if op == "mod":
+        op = "rem"
+
+    op = op.capitalize()
+    if _is_floating_point_type(lhs.dtype):
+        if op == "Floordiv":
+            raise ValueError(f"floordiv not supported for {lhs=}")
+        op += "F"
+    elif _is_integer_like_type(lhs.dtype):
+        # TODO(max): this needs to all be regularized
+        if "div" in op.lower() or "rem" in op.lower():
+            if not lhs.dtype.is_signless:
+                raise ValueError(f"{op.lower()}i not supported for {lhs=}")
+            if op == "Floordiv":
+                op = "FloorDiv"
+            op += "S"
+        op += "I"
     else:
-        op = op.capitalize()
-        lhs, rhs = lhs, rhs
-        if _is_floating_point_type(lhs.dtype):
-            op = getattr(arith_dialect, f"{op}FOp")
-        elif _is_integer_like_type(lhs.dtype):
-            op = getattr(arith_dialect, f"{op}IOp")
-        else:
-            raise NotImplementedError(f"Unsupported '{op}' operands: {lhs}, {rhs}")
+        raise NotImplementedError(f"Unsupported '{op}' operands: {lhs}, {rhs}")
+
+    op = getattr(arith_dialect, f"{op}Op")
 
     if predicate is not None:
         if _is_floating_point_type(lhs.dtype):
@@ -315,6 +332,15 @@ def is_constant(self) -> bool:
             self.owner.opview, arith_dialect.ConstantOp
         )
 
+    @property
+    @abstractmethod
+    def literal_value(self):
+        pass
+
+    @abstractmethod
+    def coerce(self, other) -> tuple["ArithValue", "ArithValue"]:
+        pass
+
     def fold(self) -> bool:
         return self.is_constant() and self._fold
 
@@ -329,6 +355,10 @@ def __repr__(self):
     __add__ = partialmethod(_binary_op, op="add")
     __sub__ = partialmethod(_binary_op, op="sub")
     __mul__ = partialmethod(_binary_op, op="mul")
+    __truediv__ = partialmethod(_binary_op, op="truediv")
+    __floordiv__ = partialmethod(_binary_op, op="floordiv")
+    __mod__ = partialmethod(_binary_op, op="mod")
+
     __radd__ = partialmethod(_binary_op, op="add")
     __rsub__ = partialmethod(_binary_op, op="sub")
     __rmul__ = partialmethod(_binary_op, op="mul")
@@ -401,3 +431,10 @@ def __int__(self):
 
     def __float__(self):
         return float(self.literal_value)
+
+    def coerce(self, other) -> tuple["Scalar", "Scalar"]:
+        if isinstance(other, (int, float, bool)):
+            other = Scalar(other, dtype=self.dtype)
+        else:
+            raise ValueError(f"can't coerce {other=} to Scalar")
+        return self, other

mlir_utils/dialects/ext/func.py

@@ -79,7 +79,8 @@ def body_builder_wrapper(self, *call_args):
             loc=self.loc,
             ip=self.ip,
         )
-        func_op.regions[0].blocks.append(*input_types)
+        arg_locs = [get_user_code_loc()] * len(sig.parameters)
+        func_op.regions[0].blocks.append(*input_types, arg_locs=arg_locs)
         with InsertionPoint(func_op.regions[0].blocks[0]):
             results = get_result_or_results(
                 self.body_builder(

mlir_utils/dialects/ext/tensor.py

@@ -186,6 +186,24 @@ def __setitem__(self, idx, source):
         previous_frame = inspect.currentframe().f_back
         _update_caller_vars(previous_frame, [self], [res])
 
+    def coerce(self, other) -> tuple["Tensor", "Tensor"]:
+        if isinstance(other, np.ndarray):
+            other = Tensor(other)
+            return self, other
+        elif _is_scalar(other):
+            if not self.has_static_shape():
+                raise ValueError(
+                    f"can't coerce {other=} because {self=} doesn't have static shape"
+                )
+            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)
+                return self, other
+
+        raise ValueError(f"can't coerce unknown {other=}")
+
 
 @dataclass(frozen=True)
 class _Indexer:

mlir_utils/types.py

@@ -22,14 +22,27 @@
 
 _index_t = lambda: IndexType.get()
 _bool_t = lambda: IntegerType.get_signless(1)
+
 _i8_t = lambda: IntegerType.get_signless(8)
 _i16_t = lambda: IntegerType.get_signless(16)
 _i32_t = lambda: IntegerType.get_signless(32)
 _i64_t = lambda: IntegerType.get_signless(64)
+
+_si8_t = lambda: IntegerType.get_signed(8)
+_si16_t = lambda: IntegerType.get_signed(16)
+_si32_t = lambda: IntegerType.get_signed(32)
+_si64_t = lambda: IntegerType.get_signed(64)
+
+_ui8_t = lambda: IntegerType.get_unsigned(8)
+_ui16_t = lambda: IntegerType.get_unsigned(16)
+_ui32_t = lambda: IntegerType.get_unsigned(32)
+_ui64_t = lambda: IntegerType.get_unsigned(64)
+
 _f16_t = lambda: F16Type.get()
 _f32_t = lambda: F32Type.get()
 _f64_t = lambda: F64Type.get()
 _bf16_t = lambda: BF16Type.get()
+
 opaque_t = lambda dialect_namespace, buffer: OpaqueType.get(dialect_namespace, buffer)
 
 
@@ -40,10 +53,22 @@ def _placeholder_opaque_t():
 _name_to_type = {
     "index_t": _index_t,
     "bool_t": _bool_t,
+
     "i8_t": _i8_t,
     "i16_t": _i16_t,
     "i32_t": _i32_t,
     "i64_t": _i64_t,
+
+    "si8_t": _si8_t,
+    "si16_t": _si16_t,
+    "si32_t": _si32_t,
+    "si64_t": _si64_t,
+
+    "ui8_t": _ui8_t,
+    "ui16_t": _ui16_t,
+    "ui32_t": _ui32_t,
+    "ui64_t": _ui64_t,
+
     "f16_t": _f16_t,
     "f32_t": _f32_t,
     "f64_t": _f64_t,

mlir_utils/util.py

@@ -5,7 +5,7 @@
 from collections import defaultdict
 from functools import wraps
 from pathlib import Path
-from typing import Callable, Sequence
+from typing import Callable, Sequence, Optional
 
 import mlir
 from mlir.dialects._ods_common import get_op_result_or_value, get_op_results_or_values
@@ -226,18 +226,22 @@ def _update_caller_vars(previous_frame, args: Sequence, replacements: Sequence):
             )
 
 
-def get_user_code_loc():
+def get_user_code_loc(user_base: Optional[Path] = None):
     import mlir_utils
     import mlir
 
     mlir_utis_root_path = Path(mlir_utils.__path__[0])
     mlir_root_path = Path(mlir.__path__[0])
 
-    prev_frame = inspect.currentframe().f_back.f_back
+    prev_frame = inspect.currentframe().f_back
+    if user_base is None:
+        user_base = Path(prev_frame.f_code.co_filename)
+
     while (
         Path(prev_frame.f_code.co_filename).is_relative_to(mlir_utis_root_path)
         or Path(prev_frame.f_code.co_filename).is_relative_to(mlir_root_path)
         or Path(prev_frame.f_code.co_filename).is_relative_to(sys.prefix)
+        or Path(prev_frame.f_code.co_filename).is_relative_to(user_base)
     ):
         prev_frame = prev_frame.f_back
     frame_info = inspect.getframeinfo(prev_frame)

pyproject.toml

@@ -18,7 +18,7 @@ dependencies = [
 test = ["pytest", "mlir-native-tools", "astpretty"]
 torch-mlir = ["torch-mlir-core"]
 jax = ["jax[cpu]", ]
-mlir = ["mlir-python-bindings==17.0.0.2023.7.26.12+25b8433"]
+mlir = ["mlir-python-bindings==17.0.0.2023.7.27.14+70aca7b1"]
 
 [project.scripts]
 configure-mlir-python-utils = "mlir_utils:_configuration.configuration.configure_host_bindings"

tests/test_operator_overloading.py

@@ -13,6 +13,51 @@
 pytest.mark.usefixtures("ctx")
 
 
+def test_arithmetic(ctx: MLIRContext):
+    one = constant(1)
+    two = constant(2)
+    one + two
+    one - two
+    one / two
+    one // two
+    one % two
+
+    one = constant(1.0)
+    two = constant(2.0)
+    one + two
+    one - two
+    one / two
+    try:
+        one // two
+    except ValueError as e:
+        assert str(e) == "floordiv not supported for lhs=Scalar(%cst, f64)"
+    one % two
+
+    ctx.module.operation.verify()
+    filecheck(
+        dedent(
+            """\
+    module {
+      %c1_i64 = arith.constant 1 : i64
+      %c2_i64 = arith.constant 2 : i64
+      %0 = arith.addi %c1_i64, %c2_i64 : i64
+      %1 = arith.subi %c1_i64, %c2_i64 : i64
+      %2 = arith.divsi %c1_i64, %c2_i64 : i64
+      %3 = arith.floordivsi %c1_i64, %c2_i64 : i64
+      %4 = arith.remsi %c1_i64, %c2_i64 : i64
+      %cst = arith.constant 1.000000e+00 : f64
+      %cst_0 = arith.constant 2.000000e+00 : f64
+      %5 = arith.addf %cst, %cst_0 : f64
+      %6 = arith.subf %cst, %cst_0 : f64
+      %7 = arith.divf %cst, %cst_0 : f64
+      %8 = arith.remf %cst, %cst_0 : f64
+    }
+    """
+        ),
+        ctx.module,
+    )
+
+
 def test_tensor_arithmetic(ctx: MLIRContext):
     one = constant(1)
     assert isinstance(one, Scalar)
@@ -117,3 +162,47 @@ def test_arith_cmp(ctx: MLIRContext):
         ),
         ctx.module,
     )
+
+
+def test_scalar_promotion(ctx: MLIRContext):
+    one = constant(1)
+    one + 2
+    one - 2
+    one / 2
+    one // 2
+    one % 2
+
+    one = constant(1.0)
+    one + 2.0
+    one - 2.0
+    one / 2.0
+    one % 2.0
+
+    ctx.module.operation.verify()
+    correct = dedent(
+        """\
+    module {
+      %c1_i64 = arith.constant 1 : i64
+      %c2_i64 = arith.constant 2 : i64
+      %0 = arith.addi %c1_i64, %c2_i64 : i64
+      %c2_i64_0 = arith.constant 2 : i64
+      %1 = arith.subi %c1_i64, %c2_i64_0 : i64
+      %c2_i64_1 = arith.constant 2 : i64
+      %2 = arith.divsi %c1_i64, %c2_i64_1 : i64
+      %c2_i64_2 = arith.constant 2 : i64
+      %3 = arith.floordivsi %c1_i64, %c2_i64_2 : i64
+      %c2_i64_3 = arith.constant 2 : i64
+      %4 = arith.remsi %c1_i64, %c2_i64_3 : i64
+      %cst = arith.constant 1.000000e+00 : f64
+      %cst_4 = arith.constant 2.000000e+00 : f64
+      %5 = arith.addf %cst, %cst_4 : f64
+      %cst_5 = arith.constant 2.000000e+00 : f64
+      %6 = arith.subf %cst, %cst_5 : f64
+      %cst_6 = arith.constant 2.000000e+00 : f64
+      %7 = arith.divf %cst, %cst_6 : f64
+      %cst_7 = arith.constant 2.000000e+00 : f64
+      %8 = arith.remf %cst, %cst_7 : f64
+    }
+    """
+    )
+    filecheck(correct, ctx.module)