fix wrong type and interp for pow and add missing stmts for vmath (#476)

Author: kaihsin

src/kirin/dialects/vmath/__init__.py

@@ -125,9 +125,7 @@ def log2(x: ilist.IList[float, ListLen]) -> ilist.IList[float, ListLen]: ...
 
 
 @lowering.wraps(stmts.pow)
-def pow(
-    x: ilist.IList[float, ListLen], y: ilist.IList[float, ListLen]
-) -> ilist.IList[float, ListLen]: ...
+def pow(x: ilist.IList[float, ListLen], y: float) -> ilist.IList[float, ListLen]: ...
 
 
 @lowering.wraps(stmts.radians)
@@ -162,3 +160,15 @@ def tanh(x: ilist.IList[float, ListLen]) -> ilist.IList[float, ListLen]: ...
 
 @lowering.wraps(stmts.trunc)
 def trunc(x: ilist.IList[float, ListLen]) -> ilist.IList[float, ListLen]: ...
+
+
+@lowering.wraps(stmts.scale)
+def scale(
+    value: float, x: ilist.IList[float, ListLen]
+) -> ilist.IList[float, ListLen]: ...
+
+
+@lowering.wraps(stmts.offset)
+def offset(
+    value: float, x: ilist.IList[float, ListLen]
+) -> ilist.IList[float, ListLen]: ...

src/kirin/dialects/vmath/interp.py

@@ -187,10 +187,11 @@ def log2(self, interp, frame: Frame, stmt: stmts.log2):
 
     @impl(stmts.pow)
     def pow(self, interp, frame: Frame, stmt: stmts.pow):
-        values = frame.get_values(stmt.args)
+        x = frame.get(stmt.x)
+        y = frame.get(stmt.y)
         return (
             ilist.IList(
-                np.pow(np.asarray(values[0]), np.asarray(values[1])).tolist(),
+                np.pow(np.asarray(x), y).tolist(),
                 elem=types.Float,
             ),
         )
@@ -243,3 +244,15 @@ def trunc(self, interp, frame: Frame, stmt: stmts.trunc):
         return (
             ilist.IList(np.trunc(np.asarray(values[0])).tolist(), elem=types.Float),
         )
+
+    @impl(stmts.scale)
+    def scale(self, interp, frame: Frame, stmt: stmts.scale):
+        a = frame.get(stmt.value)
+        x = frame.get(stmt.x)
+        return (ilist.IList((np.asarray(x) * a).tolist(), elem=types.Float),)
+
+    @impl(stmts.offset)
+    def offset(self, interp, frame: Frame, stmt: stmts.offset):
+        a = frame.get(stmt.value)
+        x = frame.get(stmt.x)
+        return (ilist.IList((np.asarray(x) + a).tolist(), elem=types.Float),)

src/kirin/dialects/vmath/stmts.py

@@ -360,3 +360,25 @@ class trunc(ir.Statement):
     traits = frozenset({ir.Pure(), lowering.FromPythonCall()})
     x: ir.SSAValue = info.argument(ilist.IListType[types.Float, ListLen])
     result: ir.ResultValue = info.result(ilist.IListType[types.Float, ListLen])
+
+
+@statement(dialect=dialect)
+class scale(ir.Statement):
+    """scale with a scalar statement"""
+
+    name = "scale"
+    traits = frozenset({ir.Pure(), lowering.FromPythonCall()})
+    x: ir.SSAValue = info.argument(ilist.IListType[types.Float, ListLen])
+    value: ir.SSAValue = info.argument(types.Float)
+    result: ir.ResultValue = info.result(ilist.IListType[types.Float, ListLen])
+
+
+@statement(dialect=dialect)
+class offset(ir.Statement):
+    """offset with a scalar statement"""
+
+    name = "offset"
+    traits = frozenset({ir.Pure(), lowering.FromPythonCall()})
+    x: ir.SSAValue = info.argument(ilist.IListType[types.Float, ListLen])
+    value: ir.SSAValue = info.argument(types.Float)
+    result: ir.ResultValue = info.result(ilist.IListType[types.Float, ListLen])

test/dialects/vmath/test_basic.py

@@ -370,11 +370,11 @@ def pow_func(x, y):
 def test_pow():
     truth = np.pow(
         ilist.IList([0.42, 0.87, 0.32], elem=types.Float),
-        ilist.IList([0.42, 0.87, 0.32], elem=types.Float),
+        3.33,
     )
     out = pow_func(
         ilist.IList([0.42, 0.87, 0.32], elem=types.Float),
-        ilist.IList([0.42, 0.87, 0.32], elem=types.Float),
+        3.33,
     )
     assert isinstance(out, ilist.IList)
     assert out.elem == types.Float
@@ -489,3 +489,31 @@ def test_trunc():
     assert isinstance(out, ilist.IList)
     assert out.elem == types.Float
     assert np.allclose(out, truth)
+
+
+@basic.union([vmath])
+def scale_func(x, y):
+    return vmath.scale(value=y, x=x)
+
+
+def test_scale():
+    a = 3.3
+    truth = np.array(ilist.IList([0.42, 0.87, 0.32], elem=types.Float)) * a
+    out = scale_func(ilist.IList([0.42, 0.87, 0.32], elem=types.Float), a)
+    assert isinstance(out, ilist.IList)
+    assert out.elem == types.Float
+    assert np.allclose(out, truth)
+
+
+@basic.union([vmath])
+def offset_func(x, y):
+    return vmath.offset(value=y, x=x)
+
+
+def test_offset():
+    a = 3.3
+    truth = np.array(ilist.IList([0.42, 0.87, 0.32], elem=types.Float)) + a
+    out = offset_func(ilist.IList([0.42, 0.87, 0.32], elem=types.Float), a)
+    assert isinstance(out, ilist.IList)
+    assert out.elem == types.Float
+    assert np.allclose(out, truth)