TEST: Test propagate_map_queries pass

Author: mtsokol

pyproject.toml

@@ -7,7 +7,7 @@ readme = "README.md"
 packages = [{include = "finch", from = "src"}]
 
 [tool.poetry.dependencies]
-python = "^3.10"
+python = "^3.11"
 numpy = ">=1.19"
 
 [tool.poetry.group.test.dependencies]

src/finch/autoschedule/__init__.py

@@ -1,4 +1,4 @@
-from .finch_logic import (
+from ..finch_logic import (
     Aggregate,
     Alias,
     Deferred,
@@ -15,7 +15,7 @@
     Table,
 )
 from .optimize import optimize, propagate_map_queries
-from .rewrite_tools import PostOrderDFS, PostWalk, PreWalk
+from ..symbolic import PostOrderDFS, PostWalk, PreWalk
 
 __all__ = [
     "Aggregate",

src/finch/autoschedule/compiler.py

@@ -2,7 +2,7 @@
 from textwrap import dedent
 from typing import Any
 
-from .finch_logic import (
+from ..finch_logic import (
     Alias,
     Deferred,
     Field,

src/finch/autoschedule/executor.py

@@ -1,5 +1,5 @@
 from .compiler import LogicCompiler
-from .rewrite_tools import gensym
+from ..symbolic import gensym
 
 
 class LogicExecutor:

src/finch/autoschedule/optimize.py

@@ -1,6 +1,6 @@
 from .compiler import LogicCompiler
-from .finch_logic import Aggregate, Alias, LogicNode, MapJoin, Plan, Produces, Query
-from .rewrite_tools import Chain, PostOrderDFS, PostWalk, PreWalk, Rewrite
+from ..finch_logic import Aggregate, Alias, LogicNode, MapJoin, Plan, Produces, Query
+from ..symbolic import Chain, PostOrderDFS, PostWalk, PreWalk, Rewrite
 
 
 def optimize(prgm: LogicNode) -> LogicNode:
@@ -52,4 +52,4 @@ def __init__(self, ctx: LogicCompiler):
 
     def __call__(self, prgm: LogicNode):
         prgm = optimize(prgm)
-        return self.ctx(prgm)
+        return self.ctx(prgm)

src/finch/finch_logic/nodes.py

@@ -27,9 +27,9 @@ def head(cls):
         return cls
 
     @classmethod
-    def make_term(cls, head, *args):
-        """Creates a term with the given head and arguments."""
-        return head(*args)
+    def make_term(cls, *args):
+        """Creates a term of the given cls type and with arguments."""
+        return cls(*args)
 
 
 @dataclass(eq=True, frozen=True)
@@ -191,6 +191,10 @@ def children(self):
         """Returns the children of the node."""
         return [self.op, *self.args]
 
+    @classmethod
+    def make_term(cls, op, *args):
+        return cls(op, args)
+
 
 @dataclass(eq=True, frozen=True)
 class Aggregate(LogicNode):
@@ -361,6 +365,10 @@ def children(self):
         """Returns the children of the node."""
         return [self.lhs, self.rhs]
 
+    @classmethod
+    def make_term(cls, lhs, rhs):
+        return cls(lhs, rhs)
+
 
 @dataclass(eq=True, frozen=True)
 class Produces(LogicNode):
@@ -412,3 +420,7 @@ def is_stateful():
     def children(self):
         """Returns the children of the node."""
         return [*self.bodies]
+
+    @classmethod
+    def make_term(cls, *val):
+        return cls(val)

src/finch/symbolic/rewriters.py

@@ -58,13 +58,13 @@ def __call__(self, x: Term) -> Term | None:
         if y is not None:
             if y.is_expr():
                 args = y.children()
-                return y.make_term(y.head(), *[default_rewrite(self(arg), arg) for arg in args])
+                return y.make_term(*[default_rewrite(self(arg), arg) for arg in args])
             return y
         if x.is_expr():
             args = x.children()
             new_args = list(map(self, args))
             if not all(arg is None for arg in new_args):
-                return x.make_term(x.head(), *map(lambda x1, x2: default_rewrite(x1, x2), new_args, args))
+                return x.make_term(*map(lambda x1, x2: default_rewrite(x1, x2), new_args, args))
             return None
         return None
 

tests/test_scheduler.py

@@ -0,0 +1,21 @@
+from finch.autoschedule import propagate_map_queries
+from finch.finch_logic import *
+
+
+def test_propagate_map_queries_simple():
+    plan = Plan(
+        (
+            Query(Alias("A10"), Aggregate(Immediate("+"), Immediate(0), Immediate("[1,2,3]"), ())),
+            Query(Alias("A11"), Alias("A10")),
+            Produces((Alias("11"),)),
+        )
+    )
+    expected = Plan(
+        (
+            Query(Alias("A11"), MapJoin(Immediate("+"), (Immediate(0), Immediate("[1,2,3]")))),
+            Produces((Alias("11"),)),
+        )
+    )
+
+    result = propagate_map_queries(plan)
+    assert result == expected