Initial draft

Author: mtsokol

sparse/__init__.py

@@ -2,6 +2,7 @@
 import warnings
 from enum import Enum
 
+from . import scheduler  # noqa: F401
 from ._version import __version__, __version_tuple__  # noqa: F401
 
 __array_api_version__ = "2022.12"

sparse/scheduler/__init__.py

@@ -0,0 +1,40 @@
+from .finch_logic import (
+    Aggregate,
+    Alias,
+    Deferred,
+    Field,
+    Immediate,
+    MapJoin,
+    Plan,
+    Produces,
+    Query,
+    Reformat,
+    Relabel,
+    Reorder,
+    Subquery,
+    Table,
+)
+from .optimize import optimize, propagate_map_queries
+from .rewrite_tools import PostOrderDFS, PostWalk, PreWalk
+
+__all__ = [
+    "Aggregate",
+    "Alias",
+    "Deferred",
+    "Field",
+    "Immediate",
+    "MapJoin",
+    "Plan",
+    "Produces",
+    "Query",
+    "Reformat",
+    "Relabel",
+    "Reorder",
+    "Subquery",
+    "Table",
+    "optimize",
+    "propagate_map_queries",
+    "PostOrderDFS",
+    "PostWalk",
+    "PreWalk",
+]

sparse/scheduler/compiler.py

@@ -0,0 +1,79 @@
+from textwrap import dedent
+
+from .finch_logic import Alias, Deferred, Field, Immediate, LogicNode, MapJoin, Query, Reformat, Relabel, Reorder, Table
+
+
+class PointwiseLowerer:
+    def __init__(self):
+        self.bound_idxs = []
+
+    def __call__(self, ex):
+        match ex:
+            case MapJoin(op, args) if isinstance(op, Immediate):
+                return f":({op.val}({','.join([self(arg) for arg in args])}))"
+            case Reorder(Relabel(arg, idxs_1), idxs_2) if isinstance(arg, Alias):
+                self.bound_idxs.append(idxs_1)
+                return f":({arg.name}[{','.join([idx.name if idx in idxs_2 else 1 for idx in idxs_1])}])"
+            case Reorder(arg, _) if isinstance(arg, Immediate):
+                return arg.val
+            case Immediate(val):
+                return val
+            case _:
+                raise Exception(f"Unrecognized logic: {ex}")
+
+
+def compile_pointwise_logic(ex: LogicNode) -> tuple:
+    ctx = PointwiseLowerer()
+    code = ctx(ex)
+    return (code, ctx.bound_idxs)
+
+
+def compile_logic_constant(ex):
+    match ex:
+        case Immediate(val):
+            return val
+        case Deferred(ex, type_):
+            return f":({ex}::{type_})"
+        case _:
+            raise Exception(f"Invalid constant: {ex}")
+
+
+class LogicLowerer:
+    def __init__(self, mode: str = "fast"):
+        self.mode = mode
+
+    def __call__(self, ex):
+        match ex:
+            case Query(lhs, Table(tns, _)) if isinstance(lhs, Alias):
+                return f":({lhs.name} = {compile_logic_constant(tns)})"
+
+            case Query(lhs, Reformat(tns, Reorder(Relabel(arg, idxs_1), idxs_2))) if isinstance(
+                lhs, Alias
+            ) and isinstance(arg, Alias):
+                loop_idxs = [idx.name for idx in withsubsequence(intersect(idxs_1, idxs_2), idxs_2)]  # noqa: F821
+                lhs_idxs = [idx.name for idx in idxs_2]
+                (rhs, rhs_idxs) = compile_pointwise_logic(Reorder(Relabel(arg, idxs_1), idxs_2))
+                body = f":({lhs.name}[{','.join(lhs_idxs)}] = {rhs})"
+                for idx in loop_idxs:
+                    if Field(idx) in rhs_idxs:
+                        body = f":(for {idx} = _ \n {body} end)"
+                    elif idx in lhs_idxs:
+                        body = f":(for {idx} = 1:1 \n {body} end)"
+
+                result = f"""\
+                    quote
+                        {lhs.name} = {compile_logic_constant(tns)}
+                        @finch mode = {self.mode} begin
+                            {lhs.name} .= {tns.fill_value}
+                            {body}
+                            return {lhs.name}
+                        end
+                    end
+                    """
+                return dedent(result)
+
+
+class LogicCompiler:
+    def __call__(self, prgm):
+        prgm = format_queries(prgm, True)  # noqa: F821
+        return LogicLowerer()(prgm)

sparse/scheduler/executor.py

@@ -0,0 +1,27 @@
+from .compiler import LogicCompiler
+
+
+class LogicExecutor:
+    def __init__(self, ctx, verbose=False):
+        self.ctx: LogicCompiler = ctx
+        self.codes = {}
+        self.verbose = verbose
+
+    def __call__(self, prgm):
+        prgm_structure = prgm
+        if prgm_structure not in self.codes:
+            thunk = logic_executor_code(self.ctx, prgm)
+            self.codes[prgm_structure] = eval(thunk), thunk
+
+        f, code = self.codes[prgm_structure]
+        if self.verbose:
+            print(code)
+        return f(prgm)
+
+
+def get_structure():
+    pass
+
+
+def logic_executor_code(ctx, prgm):
+    pass