add vmath.add

Author: neelay893

src/kirin/dialects/vmath/__init__.py

@@ -14,6 +14,13 @@
 ListLen = TypeVar("ListLen")
 
 
+@lowering.wraps(stmts.add)
+def add(
+    lhs: ilist.IList[float, ListLen] | float,
+    rhs: ilist.IList[float, ListLen] | float,
+) -> ilist.IList[float, ListLen]: ...
+
+
 @lowering.wraps(stmts.acos)
 def acos(x: ilist.IList[float, ListLen]) -> ilist.IList[float, ListLen]: ...
 

src/kirin/dialects/vmath/interp.py

@@ -12,6 +12,17 @@
 @dialect.register
 class MathMethodTable(MethodTable):
 
+    @impl(stmts.add)
+    def add(self, interp, frame: Frame, stmt: stmts.add):
+        lhs = frame.get(stmt.lhs)
+        rhs = frame.get(stmt.rhs)
+        if isinstance(lhs, ilist.IList):
+            lhs = np.asarray(lhs)
+        if isinstance(rhs, ilist.IList):
+            rhs = np.asarray(rhs)
+        arraysum = lhs + rhs
+        return (ilist.IList(arraysum.tolist(), elem=types.Float),)
+
     @impl(stmts.acos)
     def acos(self, interp, frame: Frame, stmt: stmts.acos):
         values = frame.get_values(stmt.args)

src/kirin/dialects/vmath/stmts.py

@@ -7,6 +7,21 @@
 ListLen = types.TypeVar("ListLen")
 
 
+@statement(dialect=dialect)
+class add(ir.Statement):
+    """Addition statement"""
+
+    name = "Add"
+    traits = frozenset({ir.Pure(), lowering.FromPythonCall()})
+    lhs: ir.SSAValue = info.argument(
+        ilist.IListType[types.Float, ListLen] | types.Float
+    )
+    rhs: ir.SSAValue = info.argument(
+        ilist.IListType[types.Float, ListLen] | types.Float
+    )
+    result: ir.ResultValue = info.result(types.Any)
+
+
 @statement(dialect=dialect)
 class acos(ir.Statement):
     """acos statement, wrapping the math.acos function"""

test/dialects/vmath/test_basic.py

@@ -6,6 +6,31 @@
 from kirin.dialects import ilist, vmath
 
 
+@basic.union([vmath])
+def add_kernel(x, y):
+    return vmath.add(x, y)
+
+
+def test_add_lists():
+    a = ilist.IList([0.0, 1.0, 2.0], elem=types.Float)
+    b = ilist.IList([3.0, 4.0, 5.0], elem=types.Float)
+    truth = np.array([0.0, 1.0, 2.0]) + np.array([3.0, 4.0, 5.0])
+    out = add_kernel(a, b)
+    assert isinstance(out, ilist.IList)
+    assert out.elem == types.Float
+    assert np.allclose(out, truth)
+
+
+def test_add_scalar_list():
+    a = 2.0
+    b = ilist.IList([3.0, 4.0, 5.0], elem=types.Float)
+    truth = 2.0 + np.array([3.0, 4.0, 5.0])
+    out = add_kernel(a, b)
+    assert isinstance(out, ilist.IList)
+    assert out.elem == types.Float
+    assert np.allclose(out, truth)
+
+
 @basic.union([vmath])
 def acos_func(x):
     return vmath.acos(x)