[AOT][primitives] Implemented __add__ support for IrScalar (#66)

Author: raikonenfnu

python/shark_turbine/aot/support/ir_utils.py

@@ -280,3 +280,12 @@ def build_tensor_dim_value(
 ) -> Value:
     dim_value = build_index_value(dim, constant_cache=constant_cache)
     return tensor_d.DimOp(t, dim_value).result
+
+
+# API name  inspired by mlir/python/mlir/dialects/_arith_ops_ext.py
+def _is_float_type(type):
+    return isinstance(type, (BF16Type, F16Type, F32Type, F64Type))
+
+
+def _is_integer_like_type(type):
+    return isinstance(type, (IntegerType, IndexType))

python/shark_turbine/aot/support/procedural/primitives.py

@@ -20,13 +20,17 @@
 import torch
 
 from ..ir_imports import (
+    F32Type,
     IrType,
     RankedTensorType,
     Value,
+    arith_d,
 )
 
 from ..ir_utils import (
     build_tensor_dim_value,
+    _is_float_type,
+    _is_integer_like_type,
 )
 
 from ..utils import (
@@ -38,6 +42,7 @@
     Intrinsic,
     IrTrace,
     ShapedTypeDynamicSizeSentinel,
+    current_ir_trace,
 )
 
 ###############################################################################
@@ -58,6 +63,36 @@ class IrScalar(Intrinsic):
     def __init__(self, ir_type: IrType):
         self.ir_type = ir_type
 
+    def __add__(self, other):
+        t = current_ir_trace()
+        with t.ip, t.loc:
+            # Type check and promotion.
+            # TODO: Add more comprehensive type promotion hiearchy as seen in
+            # https://jax.readthedocs.io/en/latest/jep/9407-type-promotion.html
+            lhs = self.ir_value
+            if isinstance(other, IrScalar):
+                # Assumes when both are Value, they have same type.
+                rhs = other.ir_value
+            elif isinstance(other, (int, bool)):
+                rhs = arith_d.ConstantOp(lhs.type, other).result
+            elif isinstance(other, float) and _is_integer_like_type(self.ir_type):
+                lhs = arith_d.SIToFPOp(F32Type.get(), lhs).result
+                rhs = arith_d.ConstantOp(F32Type.get(), other).result
+
+            #  Checks that lhs and rhs has same type.
+            if lhs.type != rhs.type:
+                raise ValueError("Mismatch type between lhs and rhs.")
+
+            # Emit computation.
+            if _is_integer_like_type(lhs.type):
+                return arith_d.AddIOp(lhs, rhs).result
+            elif _is_float_type(lhs.type):
+                return arith_d.AddFOp(lhs, rhs).result
+            else:
+                raise ValueError(
+                    f"Expected operand to be either Int or Float but got {self.ir_type} instead."
+                )
+
 
 class IrImmediateScalar(IrScalar):
     """Represents an IR scalar value."""

tests/aot/iree_procedural_test.py

@@ -186,6 +186,45 @@ def foobar(self, x=AbstractIndex, y=AbstractIndex):
             module_str,
         )
 
+    def testScalarAddInt(self):
+        class ArithModule(CompiledModule):
+            def foobar(self, a=AbstractI32, b=AbstractI32):
+                return a + b
+
+        inst = ArithModule(context=Context())
+        module_str = str(CompiledModule.get_mlir_module(inst))
+        self.assertIn("arith.addi %arg0, %arg1 : i32", module_str)
+
+    def testScalarAddFloat(self):
+        class ArithModule(CompiledModule):
+            def foobar(self, a=AbstractF32, b=AbstractF32):
+                return a + b
+
+        inst = ArithModule(context=Context())
+        module_str = str(CompiledModule.get_mlir_module(inst))
+        self.assertIn("arith.addf %arg0, %arg1 : f32", module_str)
+
+    def testScalarAddLiteral(self):
+        class ArithModule(CompiledModule):
+            def foobar(self, a=AbstractI32):
+                return a + 1
+
+        inst = ArithModule(context=Context())
+        module_str = str(CompiledModule.get_mlir_module(inst))
+        self.assertIn("%c1_i32 = arith.constant 1 : i32", module_str)
+        self.assertIn("arith.addi %arg0, %c1_i32 : i32", module_str)
+
+    def testScalarAddLiteralMixedType(self):
+        class ArithModule(CompiledModule):
+            def foobar(self, a=AbstractI32):
+                return a + 3.23
+
+        inst = ArithModule(context=Context())
+        module_str = str(CompiledModule.get_mlir_module(inst))
+        self.assertIn("%0 = arith.sitofp %arg0 : i32 to f32", module_str)
+        self.assertIn("%cst = arith.constant 3.230000e+00 : f32", module_str)
+        self.assertIn("arith.addf %0, %cst : f32", module_str)
+
 
 if __name__ == "__main__":
     logging.basicConfig(level=logging.DEBUG)