register vmath desugar as post inference and remove pass (#573)

- vmath desugar rewrite (#568) is now registered as a vmath dialect
post-inference rewrite rule
- remove vmath/passes.py as it is no longer needed

These changes are separate from (#568) since they are not compatible
with kirin 0.17 (in case we want to backport).

Author: neelay893

src/kirin/dialects/vmath/__init__.py

@@ -6,6 +6,7 @@
 
 from . import stmts as stmts, interp as interp
 from ._dialect import dialect as dialect
+from .rewrites import desugar as desugar
 
 pi = pymath.pi
 e = pymath.e

src/kirin/dialects/vmath/passes.py

@@ -6,6 +6,7 @@
 from kirin.dialects.ilist import IListType
 
 from ..stmts import add as vadd, div as vdiv, sub as vsub, mult as vmult
+from .._dialect import dialect
 
 
 class DesugarBinOp(RewriteRule):
@@ -52,6 +53,7 @@ def replace_binop(self, node: ir.Statement) -> RewriteResult:
                 return RewriteResult()
 
 
+@dialect.post_inference
 class WalkDesugarBinop(RewriteRule):
     """
     Walks DesugarBinop. Needed for correct behavior when

src/kirin/dialects/vmath/rewrites/desugar.py

@@ -4,7 +4,6 @@
 
 from kirin.prelude import basic
 from kirin.dialects import vmath
-from kirin.dialects.vmath.passes import VMathDesugar
 from kirin.dialects.ilist.runtime import IList
 
 
@@ -25,7 +24,6 @@ def add_scalar_lhs():
 
 def test_add_scalar_lhs():
     # out = add_scalar_lhs()
-    VMathDesugar(add_scalar_lhs.dialects).unsafe_run(add_scalar_lhs)
     add_scalar_lhs.print()
     res = add_scalar_lhs()
     assert isinstance(res, IList)
@@ -34,51 +32,46 @@ def test_add_scalar_lhs():
 
 
 def test_typed_kernel_add():
-    VMathDesugar(add_scalar_rhs_typed.dialects).unsafe_run(add_scalar_rhs_typed)
     add_scalar_rhs_typed.print()
     res = add_scalar_rhs_typed(IList([0, 1, 2]), 3.1)
     assert np.allclose(np.asarray(res), np.asarray([3.1, 4.1, 5.1]))
 
 
-@basic.union([vmath])
+@basic.union([vmath])(typeinfer=True)
 def add_two_lists():
     return add_kernel(x=[0, 1, 2], y=[3, 4, 5])
 
 
 def test_add_lists():
-    VMathDesugar(add_two_lists.dialects).unsafe_run(add_two_lists)
     res = add_two_lists()
     assert np.allclose(np.asarray(res), np.array([0, 1, 2, 3, 4, 5]))
 
 
-@basic.union([vmath])
+@basic.union([vmath])(typeinfer=True)
 def sub_scalar_rhs_typed(x: IList[float, Any], y: float):
     return x - y
 
 
 def test_sub_scalar_typed():
-    VMathDesugar(sub_scalar_rhs_typed.dialects).unsafe_run(sub_scalar_rhs_typed)
     res = sub_scalar_rhs_typed(IList([0, 1, 2]), 3.1)
     assert np.allclose(np.asarray(res), np.asarray([-3.1, -2.1, -1.1]))
 
 
-@basic.union([vmath])
+@basic.union([vmath])(typeinfer=True)
 def mult_scalar_lhs_typed(x: float, y: IList[float, Any]):
     return x * y
 
 
 def test_mult_scalar_typed():
-    VMathDesugar(mult_scalar_lhs_typed.dialects).unsafe_run(mult_scalar_lhs_typed)
     res = mult_scalar_lhs_typed(3, IList([0, 1, 2]))
     assert np.allclose(np.asarray(res), np.asarray([0, 3, 6]))
 
 
-@basic.union([vmath])
+@basic.union([vmath])(typeinfer=True)
 def div_scalar_lhs_typed(x: float, y: IList[float, Any]):
     return x / y
 
 
 def test_div_scalar_typed():
-    VMathDesugar(div_scalar_lhs_typed.dialects).unsafe_run(div_scalar_lhs_typed)
     res = div_scalar_lhs_typed(3, IList([1, 1.5, 2]))
     assert np.allclose(np.asarray(res), np.asarray([3, 2, 1.5]))