Add e2e loopnest value test

Author: saulshanabrook

python/egglog/exp/array_api.py

@@ -119,6 +119,9 @@ class Int(Expr, ruleset=array_api_ruleset):
     # https://en.wikipedia.org/wiki/Bottom_type
     NEVER: ClassVar[Int]
 
+    @classmethod
+    def var(cls, name: StringLike) -> Int: ...
+
     def __init__(self, value: i64Like) -> None: ...
 
     def __invert__(self) -> Int: ...

python/egglog/exp/array_api_loopnest.py

@@ -1,8 +1,7 @@
 """
-In progress module
+Example module to replicate behavior expressed in
 
 https://gist.github.com/sklam/5e5737137d48d6e5b816d14a90076f1d
-
 """
 
 # %%

python/egglog/exp/array_api_program_gen.py

@@ -38,7 +38,8 @@ def int_program(x: Int) -> Program: ...
 
 
 @array_api_program_gen_ruleset.register
-def _int_program(i64_: i64, i: Int, j: Int):
+def _int_program(i64_: i64, i: Int, j: Int, s: String):
+    yield rewrite(int_program(Int.var(s))).to(Program(s, True))
     yield rewrite(int_program(Int(i64_))).to(Program(i64_.to_string()))
     yield rewrite(int_program(~i)).to(Program("~") + int_program(i))
     yield rewrite(bool_program(i < j)).to(Program("(") + int_program(i) + " < " + int_program(j) + ")")
@@ -145,9 +146,14 @@ def _value_program(i: Int, b: Boolean, f: Float, x: NDArray, v1: Value, v2: Valu
     yield rewrite(value_program(x.to_value())).to(ndarray_program(x))
     yield rewrite(value_program(v1 < v2)).to(Program("(") + value_program(v1) + " < " + value_program(v2) + ")")
     yield rewrite(value_program(v1 / v2)).to(Program("(") + value_program(v1) + " / " + value_program(v2) + ")")
+    yield rewrite(value_program(v1 + v2)).to(Program("(") + value_program(v1) + " + " + value_program(v2) + ")")
+    yield rewrite(value_program(v1 * v2)).to(Program("(") + value_program(v1) + " * " + value_program(v2) + ")")
     yield rewrite(bool_program(v1.to_bool)).to(value_program(v1))
     yield rewrite(int_program(v1.to_int)).to(value_program(v1))
-    yield rewrite(value_program(xs.index(ti))).to(ndarray_program(xs) + "[" + tuple_int_program(ti) + "]")
+    yield rewrite(value_program(xs.index(ti))).to((ndarray_program(xs) + "[" + tuple_int_program(ti) + "]").assign())
+    yield rewrite(value_program(v1.sqrt())).to(Program("np.sqrt(") + value_program(v1) + ")")
+    yield rewrite(value_program(v1.real())).to(Program("np.real(") + value_program(v1) + ")")
+    yield rewrite(value_program(v1.conj())).to(Program("np.conj(") + value_program(v1) + ")")
 
 
 @function

python/egglog/exp/program_gen.py

@@ -45,6 +45,13 @@ def function_two(self, arg1: ProgramLike, arg2: ProgramLike, name: StringLike =
         Returns a new program defining a function with two arguments.
         """
 
+    def function_three(
+        self, arg1: ProgramLike, arg2: ProgramLike, arg3: ProgramLike, name: StringLike = String("__fn")
+    ) -> Program:
+        """
+        Returns a new program defining a function with three arguments.
+        """
+
     def expr_to_statement(self) -> Program:
         """
         Returns a new program with the expression as a statement and the new expression empty.
@@ -142,10 +149,12 @@ def program_gen_ruleset(
     s3: String,
     s4: String,
     s5: String,
+    s6: String,
     p: Program,
     p1: Program,
     p2: Program,
     p3: Program,
+    p4: Program,
     i: i64,
     i2: i64,
     b: Bool,
@@ -357,6 +366,7 @@ def program_gen_ruleset(
 
     ##
     # Function two
+    ##
 
     # When compiling a function, the two args, p2 and p3, should get compiled when we compile p1, and should just be vars.
     fn_two = eq(p).to(p1.function_two(p2, p3, s1))
@@ -387,3 +397,35 @@ def program_gen_ruleset(
         set_(p.next_sym).to(i),
         set_(p.expr).to(s1),
     )
+
+    ##
+    # Function three
+    ##
+
+    fn_three = eq(p).to(p1.function_three(p2, p3, p4, s1))
+    yield rule(fn_three, p.compile(i)).then(
+        set_(p2.parent).to(p),
+        set_(p3.parent).to(p),
+        set_(p1.parent).to(p),
+        set_(p4.parent).to(p),
+        p2.compile(i),
+        p3.compile(i),
+        p1.compile(i),
+        p4.compile(i),
+        set_(p.is_identifer).to(Bool(True)),
+    )
+    yield rule(
+        fn_three,
+        p.compile(i),
+        eq(s2).to(p1.expr),
+        eq(s3).to(p1.statements),
+        eq(s4).to(p2.expr),
+        eq(s5).to(p3.expr),
+        eq(s6).to(p4.expr),
+    ).then(
+        set_(p.statements).to(
+            join("def ", s1, "(", s4, ", ", s5, ", ", s6, "):\n    ", s3.replace("\n", "\n    "), "return ", s2, "\n")
+        ),
+        set_(p.next_sym).to(i),
+        set_(p.expr).to(s1),
+    )

python/egglog/runtime.py

@@ -601,8 +601,10 @@ def resolve_callable(callable: object) -> tuple[CallableRef, Declarations]:
         case RuntimeClass(thunk, tp):
             return InitRef(tp.name), thunk()
         case RuntimeExpr(decl_thunk, expr_thunk):
-            if not isinstance((expr := expr_thunk().expr), CallDecl) or not isinstance(expr.callable, ConstantRef):
-                raise NotImplementedError(f"Can only turn constants into callable refs, not {expr}")
+            if not isinstance((expr := expr_thunk().expr), CallDecl) or not isinstance(
+                expr.callable, ConstantRef | ClassVariableRef
+            ):
+                raise NotImplementedError(f"Can only turn constants or classvars into callable refs, not {expr}")
             return expr.callable, decl_thunk()
         case _:
             raise NotImplementedError(f"Cannot turn {callable} of type {type(callable)} into a callable ref")

python/tests/__snapshots__/test_array_api/TestLoopNest.test_index_codegen[code].py

@@ -0,0 +1,8 @@
+def __fn(X, i, j):
+    _0 = X[(0, 0, i, j, )]
+    _1 = X[(0, 1, i, j, )]
+    _2 = X[(1, 0, i, j, )]
+    _3 = X[(1, 1, i, j, )]
+    _4 = X[(2, 0, i, j, )]
+    _5 = X[(2, 1, i, j, )]
+    return np.sqrt((((((np.real((np.conj(_0) * _0)) + np.real((np.conj(_1) * _1))) + np.real((np.conj(_2) * _2))) + np.real((np.conj(_3) * _3))) + np.real((np.conj(_4) * _4))) + np.real((np.conj(_5) * _5))))

python/tests/__snapshots__/test_array_api/TestLoopNest.test_index_codegen[expr].py

@@ -0,0 +1,13 @@
+_Value_1 = NDArray.var("X").index(TupleInt.from_vec(Vec[Int](Int(0), Int(0), Int.var("i"), Int.var("j"))))
+_Value_2 = NDArray.var("X").index(TupleInt.from_vec(Vec[Int](Int(0), Int(1), Int.var("i"), Int.var("j"))))
+_Value_3 = NDArray.var("X").index(TupleInt.from_vec(Vec[Int](Int(1), Int(0), Int.var("i"), Int.var("j"))))
+_Value_4 = NDArray.var("X").index(TupleInt.from_vec(Vec[Int](Int(1), Int(1), Int.var("i"), Int.var("j"))))
+_Value_5 = NDArray.var("X").index(TupleInt.from_vec(Vec[Int](Int(2), Int(0), Int.var("i"), Int.var("j"))))
+_Value_6 = NDArray.var("X").index(TupleInt.from_vec(Vec[Int](Int(2), Int(1), Int.var("i"), Int.var("j"))))
+(
+    (
+        ((((_Value_1.conj() * _Value_1).real() + (_Value_2.conj() * _Value_2).real()) + (_Value_3.conj() * _Value_3).real()) + (_Value_4.conj() * _Value_4).real())
+        + (_Value_5.conj() * _Value_5).real()
+    )
+    + (_Value_6.conj() * _Value_6).real()
+).sqrt()

python/tests/__snapshots__/test_program_gen/test_to_string_function_three.py

@@ -0,0 +1,8 @@
+def my_fn(x, y, z):
+    _0 = -x
+    assert _0 > 0
+    _1 = _0 + y
+    _2 = _1 + z
+    _3 = _2 + 2
+    _4 = _3 + _2
+    return _4

python/tests/test_array_api.py

@@ -225,29 +225,33 @@ def linalg_norm_v2(X: NDArrayLike, axis: TupleIntLike) -> NDArray:
     )
 
 
-def linalg_val(linalg_fn: Callable[[NDArray, TupleIntLike], NDArray]) -> NDArray:
-    X = NDArray.var("X")
+def linalg_val(X: NDArray, linalg_fn: Callable[[NDArray, TupleIntLike], NDArray]) -> NDArray:
     assume_shape(X, (3, 2, 3, 4))
     return linalg_fn(X, (0, 1))
 
 
 class TestLoopNest:
     @pytest.mark.parametrize("linalg_fn", [linalg_norm, linalg_norm_v2])
     def test_shape(self, linalg_fn):
-        check_eq(linalg_val(linalg_fn).shape, TupleInt.from_vec((3, 4)), array_api_schedule)
+        X = np.random.random((3, 2, 3, 4))
+        expect = np.linalg.norm(X, axis=(0, 1))
+        assert expect.shape == (3, 4)
+
+        check_eq(linalg_val(constant("X", NDArray), linalg_fn).shape, TupleInt.from_vec((3, 4)), array_api_schedule)
 
     @pytest.mark.parametrize("linalg_fn", [linalg_norm, linalg_norm_v2])
-    def test_index(self, linalg_fn):
+    def test_abstract_index(self, linalg_fn):
         i = constant("i", Int)
         j = constant("j", Int)
-        idxed = linalg_val(linalg_fn).index((i, j))
-        _NDArray_1 = NDArray.var("X")
-        _Value_1 = _NDArray_1.index(TupleInt.from_vec(Vec[Int](Int(0), Int(0), i, j)))
-        _Value_2 = _NDArray_1.index(TupleInt.from_vec(Vec[Int](Int(0), Int(1), i, j)))
-        _Value_3 = _NDArray_1.index(TupleInt.from_vec(Vec[Int](Int(1), Int(0), i, j)))
-        _Value_4 = _NDArray_1.index(TupleInt.from_vec(Vec[Int](Int(1), Int(1), i, j)))
-        _Value_5 = _NDArray_1.index(TupleInt.from_vec(Vec[Int](Int(2), Int(0), i, j)))
-        _Value_6 = _NDArray_1.index(TupleInt.from_vec(Vec[Int](Int(2), Int(1), i, j)))
+        X = constant("X", NDArray)
+        idxed = linalg_val(X, linalg_fn).index((i, j))
+
+        _Value_1 = X.index(TupleInt.from_vec(Vec[Int](Int(0), Int(0), i, j)))
+        _Value_2 = X.index(TupleInt.from_vec(Vec[Int](Int(0), Int(1), i, j)))
+        _Value_3 = X.index(TupleInt.from_vec(Vec[Int](Int(1), Int(0), i, j)))
+        _Value_4 = X.index(TupleInt.from_vec(Vec[Int](Int(1), Int(1), i, j)))
+        _Value_5 = X.index(TupleInt.from_vec(Vec[Int](Int(2), Int(0), i, j)))
+        _Value_6 = X.index(TupleInt.from_vec(Vec[Int](Int(2), Int(1), i, j)))
         res = (
             (
                 (
@@ -263,6 +267,36 @@ def test_index(self, linalg_fn):
         ).sqrt()
         check_eq(idxed, res, array_api_schedule)
 
+    def test_index_codegen(self, snapshot_py):
+        X = NDArray.var("X")
+        i = Int.var("i")
+        j = Int.var("j")
+        idxed = linalg_val(X, linalg_norm_v2).index((i, j))
+        simplified_index = simplify(idxed, array_api_schedule)
+        assert str(simplified_index) == snapshot_py(name="expr")
+
+        res = EvalProgram(
+            value_program(simplified_index).function_three(ndarray_program(X), int_program(i), int_program(j)),
+            {"np": np},
+        )
+        egraph = EGraph()
+        egraph.register(res)
+        egraph.run(array_api_program_gen_schedule)
+        print(
+            egraph.extract(
+                value_program(simplified_index).function_three(ndarray_program(X), int_program(i), int_program(j))
+            )
+        )
+        # egraph.display(split_primitive_outputs=True, n_inline_leaves=3, split_functions=[TupleInt.EMPTY, TupleInt.append, Int])
+        assert egraph.eval(res.statements) == snapshot_py(name="code")
+
+        fn_value = egraph.eval(res.py_object)
+        X = np.random.random((3, 2, 3, 4))
+        expect = np.linalg.norm(X, axis=(0, 1))
+
+        for idxs in np.ndindex(*expect.shape):
+            assert np.allclose(fn_value(X, *idxs), expect[idxs], rtol=1e-03)
+
 
 # This test happens in different steps. Each will be benchmarked and saved as a snapshot.
 # The next step will load the old snapshot and run their test on it.

python/tests/test_program_gen.py

@@ -61,6 +61,20 @@ def test_to_string(snapshot_py) -> None:
     assert egraph.eval(fn.statements) == snapshot_py
 
 
+def test_to_string_function_three(snapshot_py) -> None:
+    first = assume_pos(-Math.var("x")) + Math.var("y") + Math.var("z")
+    fn = (first + Math(2) + first).program.function_three(
+        Math.var("x").program, Math.var("y").program, Math.var("z").program, "my_fn"
+    )
+    egraph = EGraph()
+    egraph.register(fn)
+    egraph.register(fn.compile())
+    egraph.run((to_program_ruleset | program_gen_ruleset).saturate())
+    # egraph.display(n_inline_leaves=1)
+    assert egraph.eval(fn.expr) == "my_fn"
+    assert egraph.eval(fn.statements) == snapshot_py
+
+
 def test_py_object():
     x = Math.var("x")
     y = Math.var("y")