diff --git a/pyproject.toml b/pyproject.toml index ad0cddb..7263744 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -11,7 +11,6 @@ dependencies = [ "numba>=0.61,<0.63.1", "juliacall (>=0.9.24,<0.10.0)", "lark (>=1.3.0,<2.0.0)", - "finch-tensor>=0.7.0,<0.8.0", ] [tool.poetry] diff --git a/src/galley_jl_python/__init__.py b/src/galley_jl_python/__init__.py index 0a19a2c..9c15183 100644 --- a/src/galley_jl_python/__init__.py +++ b/src/galley_jl_python/__init__.py @@ -115,7 +115,7 @@ uint32, uint64, ) -from .finch_fused import jit +from .fused import jit from .io import ( read, write, diff --git a/src/galley_jl_python/finch_fused.py b/src/galley_jl_python/finch_fused.py deleted file mode 100644 index 2f64efc..0000000 --- a/src/galley_jl_python/finch_fused.py +++ /dev/null @@ -1,33 +0,0 @@ -from collections.abc import Callable -from typing import Any - -from finch.finch_fused import jit as _finch_jit - -from .compiled import compute, lazy -from .tensor import Tensor - - -def _defer(value: Any) -> Any: - if isinstance(value, tuple): - return tuple(_defer(item) for item in value) - return lazy(value) if isinstance(value, Tensor) else value - - -def _compute(value: Any) -> Any: - if isinstance(value, tuple): - return tuple(_compute(item) for item in value) - return compute(value) if isinstance(value, Tensor) else value - - -def _maybedefer(values: tuple[Any, ...]) -> tuple[Any, ...]: - return tuple(_defer(value) for value in values) - - -def jit(f: Callable[..., Any], /, ctx: Any = None) -> Callable[..., Any]: - transformed = _finch_jit(f, ctx=ctx) - transformed.__globals__.update( - defer=_defer, - compute=_compute, - maybedefer=_maybedefer, - ) - return transformed diff --git a/src/galley_jl_python/fused/__init__.py b/src/galley_jl_python/fused/__init__.py new file mode 100644 index 0000000..bfea25d --- /dev/null +++ b/src/galley_jl_python/fused/__init__.py @@ -0,0 +1,10 @@ +""" +A JIT that inserts `defer` and `compute` calls into Python functions so that +Galley fuses the tensor operations between control flow. + +Adapted from `finch.finch_fused` in finch-tensor (MIT licensed). +""" + +from .jit import jit + +__all__ = ["jit"] diff --git a/src/galley_jl_python/fused/cfg_builder.py b/src/galley_jl_python/fused/cfg_builder.py new file mode 100644 index 0000000..a92b543 --- /dev/null +++ b/src/galley_jl_python/fused/cfg_builder.py @@ -0,0 +1,277 @@ +from dataclasses import dataclass +from typing import cast + +from .nodes import ( + Assign, + Block, + Break, + ExprStmt, + For, + Function, + FusedNode, + FusedStatement, + FusedTree, + If, + Module, + Return, + While, +) +from .symbolic import BasicBlock, ControlFlowGraph + + +@dataclass(eq=True, frozen=True) +class NumberedStatement(FusedTree, FusedStatement): + """ + Wrapper for AssemblyStatement that assigns a unique id to each statement + for easier tracking in the CFG. + + Attributes: + stmt: The original FusedStatement being wrapped. + sid: A unique integer identifier for the statement. + """ + + stmt: FusedStatement + sid: int + + def __str__(self) -> str: + return str(self.stmt) + + @property + def children(self): + return (self.stmt, self.sid) + + +def number_statements(node: FusedNode, sid: int = 0) -> tuple[FusedNode, int]: + """ + Recursively traverse the AST and wrap each statement in a NumberedStatement + with a unique id. + + Args: + node: The root node of the AST to number. + sid: The starting statement id for numbering. Defaults to 0. + + Returns: + A tuple containing the new AST with numbered statements and the next + available statement id. + """ + + def go(node: FusedNode) -> FusedNode: + match node: + case Module(funcs): + return Module(cast(tuple[Function, ...], tuple(go(f) for f in funcs))) + case Function(name, args, body): + return Function(name, args, cast(Block, go(body))) + case Block(bodies): + new_bodies: list[FusedStatement] = [] + for b in bodies: + b2 = go(b) + if isinstance(b2, Block): + new_bodies.extend(b2.body) + else: + new_bodies.append(cast(FusedStatement, b2)) + return Block(tuple(new_bodies)) + case If(cond, body, else_body): + new_else = cast(Block, go(else_body)) if else_body is not None else None + return If(cond, cast(Block, go(body)), new_else) + case While(cond, body): + return While(cond, cast(Block, go(body))) + case For(target, iter, body): + return For(target, iter, cast(Block, go(body))) + case node: + if isinstance(node, (Assign, ExprStmt, Return, Break)): + nonlocal sid + s = NumberedStatement(node, sid) + sid += 1 + return s + return node + + return go(node), sid + + +def fused_desugar(node: FusedNode) -> FusedNode: + """ + Lower surface syntax to a core AST shape before CFG construction. + + - make function parameters explicit in the function body so they get statement ids + and can be referenced in the CFG (e.g. desugar `def f(x): return x + 1` to + `def f(x): x_1 = x; return x_1 + 1` so that the parameter `x` gets a statement + id and can be referenced in the CFG) + """ + + def go(node: FusedNode) -> FusedNode: + """Recursively desugar the AST.""" + match node: + case Module(funcs): + return Module(cast(tuple[Function, ...], tuple(go(f) for f in funcs))) + case Function(name, args, body): + body_2 = cast(Block, go(body)) + + # Make argument definitions explicit so they get statement ids. + func_prologue = tuple((Assign(arg, arg)) for arg in args) + return Function(name, args, Block((*func_prologue, *body_2.body))) + case Block(bodies): + new_bodies: list[FusedStatement] = [] + for b in bodies: + b2 = go(b) + new_bodies.append(cast(FusedStatement, b2)) + return Block(tuple(new_bodies)) + case If(cond, body, else_body): + new_else = cast(Block, go(else_body)) if else_body is not None else None + return If(cond, cast(Block, go(body)), new_else) + case While(cond, body): + return While(cond, cast(Block, go(body))) + case For(target, iter, body): + return For(target, iter, cast(Block, go(body))) + case node: + return node + + return go(node) + + +class FusedCFGBuilder: + """Incrementally builds control-flow graph for Finch Fused IR.""" + + def __init__(self) -> None: + self.cfg: ControlFlowGraph = ControlFlowGraph() + self.current_block: BasicBlock = self.cfg.entry_block + + def emit(self, stmt) -> None: + """Add a statement to the current block.""" + self.current_block.add_statement(stmt) + + def build(self, node: FusedNode) -> ControlFlowGraph: + return self(node) + + def __call__( + self, + node: FusedNode, + break_block: BasicBlock | None = None, + return_block: BasicBlock | None = None, + ) -> ControlFlowGraph: + match node: + case Assign(): + self.emit(node) + case NumberedStatement(stmt, _): + match stmt: + case Return(_): + self.emit(node) + assert return_block + self.current_block.add_successor(return_block) + unreachable_block = self.cfg.new_block() + self.current_block = unreachable_block + case Break(): + self.emit(node) + assert break_block + self.current_block.add_successor(break_block) + unreachable_block = self.cfg.new_block() + self.current_block = unreachable_block + case _: + self.emit(node) + case Block(bodies): + for body in bodies: + self(body, break_block, return_block) + case If(cond, body, else_body): + before_block = self.current_block + self.emit(cond) + + # create blocks for if, else, and after + if_block = self.cfg.new_block() + else_block = self.cfg.new_block() + after_block = self.cfg.new_block() + + # connect before block to if and else blocks + before_block.add_successor(if_block) + before_block.add_successor(else_block) + + # fill in the if block + self.current_block = if_block + self(body, break_block, return_block) + self.current_block.add_successor(after_block) + + # fill in the else block + self.current_block = else_block + if else_body is not None: + self(else_body, break_block, return_block) + self.current_block.add_successor(after_block) + + # continue building after the if-else + self.current_block = after_block + case While(cond, body): + before_block = self.current_block + self.emit(cond) + + # create blocks for the loop body and the code after the loop + body_block = self.cfg.new_block() + after_block = self.cfg.new_block() + + # connect before block to the loop body and the after block + before_block.add_successor(body_block) + before_block.add_successor(after_block) + + # fill in the loop body + self.current_block = body_block + self(body, after_block, return_block) + self.emit(cond) + + # connect the end of loop body back to the beginning to form the loop + self.current_block.add_successor(body_block) + self.current_block.add_successor(after_block) + self.current_block = after_block + case For(target, iter, body): + before_block = self.current_block + body_block = self.cfg.new_block() + after_block = self.cfg.new_block() + + # connect before block to the loop body and the after block + before_block.add_successor(body_block) + before_block.add_successor(after_block) + + # fill in the loop body + # The loop target is (re)defined from the iterable on every iteration. + self.current_block = body_block + self.emit(Assign(target, iter)) + self(body, after_block, return_block) + + # connect the end of loop body back to the beginning to form the loop + self.current_block.add_successor(body_block) + self.current_block.add_successor(after_block) + self.current_block = after_block + + case Function(func_name, _, body): + # set block names to the function name + self.cfg.block_name = str(func_name.val) + + # create entry/exit block for the function + func_entry_block = self.cfg.new_block() + func_exit_block = self.cfg.new_block() + + # connect CFG entry block to the function's entry block + self.cfg.entry_block.add_successor(func_entry_block) + + # dive into the body of the function + self.current_block = func_entry_block + self(body, break_block, func_exit_block) + + # connect last block in the function to the + # exit block of the function + self.current_block.add_successor(func_exit_block) + + # connect function exit block to the exit block of the CFG + func_exit_block.add_successor(self.cfg.exit_block) + + case node: + raise NotImplementedError(node, "FusedCFGBuilder") + + return self.cfg + + +def fused_build_cfg(node: FusedNode) -> ControlFlowGraph: + """ + Build control-flow graph for a FinchAssembly node. + Args: + node: Root FinchAssembly node to build CFG for. + Returns: + ControlFlowGraph: The constructed control-flow graph. + """ + ctx = FusedCFGBuilder() + return ctx.build(node) diff --git a/src/galley_jl_python/fused/dataflow.py b/src/galley_jl_python/fused/dataflow.py new file mode 100644 index 0000000..20f629e --- /dev/null +++ b/src/galley_jl_python/fused/dataflow.py @@ -0,0 +1,225 @@ +from typing import Any, cast + +from ..compiled import compute as compute_tensor +from ..compiled import lazy +from ..tensor import Tensor +from .cfg_builder import ( + NumberedStatement, + fused_build_cfg, + fused_desugar, + number_statements, +) +from .nodes import ( + Assign, + Block, + Break, + Call, + Function, + FusedNode, + FusedStatement, + Literal, + Return, + Variable, +) +from .symbolic import Chain, DataFlowAnalysis, Namespace, PostWalk, Rewrite + + +def get_variables_in_stmt(stmt: FusedNode) -> set[Variable]: + var_set = set() + + def _var_gatherer(node: FusedNode) -> FusedNode: + match node: + case Variable() as var: + var_set.add(var) + return node + case node: + return node + + Rewrite(PostWalk(_var_gatherer))(stmt) + return var_set + + +class LivenessAnalysis(DataFlowAnalysis): + def stmt_str(self, stmt: FusedNode, state: dict) -> str: + str_state = ", ".join(f"{var}" for var in state) + return f"Live vars: {{{str_state}}} | Stmt: {stmt}" + + def transfer(self, stmts, state: dict) -> dict: + # Walk through the statements in reverse to compute + # liveness + new_state = state.copy() + for stmt in reversed(stmts): + match stmt: + case NumberedStatement(Assign(lhs, rhs), _): + if lhs in new_state: + del new_state[lhs] + for var in get_variables_in_stmt(rhs): + new_state[var] = True + case Assign(lhs, rhs): + if lhs in new_state: + del new_state[lhs] + for var in get_variables_in_stmt(rhs): + new_state[var] = True + case stmt: + for var in get_variables_in_stmt(stmt): + new_state[var] = True + return new_state + + def join(self, state_1: dict, state_2: dict) -> dict: + return state_1 | state_2 + + def direction(self) -> str: + return "backward" + + +# Find the first and last statement in the block to determine +# where to insert lazy and compute calls +def _get_stmt_bounds(stmts: list[FusedNode]) -> tuple[int, int]: + if not stmts: + return -1, -1 + + top_id = -1 + bot_id = -1 + for stmt in stmts: + if isinstance(stmt, NumberedStatement): + if top_id == -1: + top_id = stmt.sid + bot_id = stmt.sid + return top_id, bot_id + + +def _insert_compute( + prgm: FusedNode, compute_sid, vars: set[Variable], nspc: Namespace +) -> FusedNode: + def _visitor(node): + match node: + # In the case of returns, we need to assign the expressions, + # compute them, and then return the computed variables. + case NumberedStatement(Return(ret_expr), sid) if sid == compute_sid: + for expr in ret_expr: + for var in get_variables_in_stmt(expr): + vars.add(var) + lazy_vars = tuple(sorted(vars, key=lambda var: var.name)) + lazy_vars_tuple = Call(Literal(tuple), lazy_vars) + lazies = ( + Assign(lazy_vars_tuple, Call(Literal(defer), (lazy_vars_tuple,))), + ) + exprs_to_compute = () + return_vars = () + for i, expr in enumerate(ret_expr): + if isinstance(expr, Variable) and expr in vars: + return_vars += (expr,) + else: + new_var = Variable(nspc.freshen("ret_" + str(i))) + exprs_to_compute += (Assign(new_var, expr),) + return_vars += (new_var,) + return_vars_tuple = Call(Literal(tuple), return_vars) + exprs_to_compute += ( + Assign( + return_vars_tuple, + Call(Literal(compute), (return_vars_tuple,)), + ), + ) + return Block(lazies + exprs_to_compute + (Return(return_vars),)) + case NumberedStatement(stmt, sid) if sid == compute_sid: + compute_vars = tuple(sorted(vars, key=lambda var: var.name)) + compute_vars_tuple = Call(Literal(tuple), compute_vars) + computes = ( + Assign( + compute_vars_tuple, + Call(Literal(compute), (compute_vars_tuple,)), + ), + ) + match stmt: + case Return(): + return Block(computes + (node,)) + case Break(): + return Block(computes + (node,)) + case stmt: + return Block((node,) + computes) + case node: + return node + + return Rewrite(PostWalk(_visitor))(prgm) + + +def defer(value: Any) -> Any: + if isinstance(value, tuple): + return tuple(defer(item) for item in value) + return lazy(value) if isinstance(value, Tensor) else value + + +def compute(value: Any) -> Any: + if isinstance(value, tuple): + return tuple(compute(item) for item in value) + return compute_tensor(value) if isinstance(value, Tensor) else value + + +def maybedefer(values: tuple[Any, ...]) -> tuple[Any, ...]: + return tuple(defer(value) for value in values) + + +def _insert_lazy(prgm: FusedNode, lazy_sid: int, vars: set[Variable]) -> FusedNode: + def _visitor(node): + match node: + case NumberedStatement(stmt, sid) if sid == lazy_sid and not isinstance( + stmt, (Return, Break) + ): + lazy_vars = tuple(sorted(vars, key=lambda var: var.name)) + lazy_vars_tuple = Call(Literal(tuple), lazy_vars) + lazies = ( + Assign( + lazy_vars_tuple, Call(Literal(maybedefer), (lazy_vars_tuple,)) + ), + ) + return Block(lazies + (node,)) + case node: + return node + + return Rewrite(PostWalk(_visitor))(prgm) + + +def _unwrap_numbered_stmt(node: FusedNode) -> FusedNode: + match node: + case NumberedStatement(stmt, _): + return stmt + case node: + return node + + +def _unnest_block(node: FusedNode) -> FusedNode: + match node: + case Block(bodies): + new_bodies: list[FusedStatement] = [] + for b in bodies: + b2 = _unnest_block(b) + if isinstance(b2, Block): + new_bodies.extend(b2.body) + else: + new_bodies.append(cast(FusedStatement, b2)) + return Block(tuple(new_bodies)) + case node: + return node + + +def insert_lazy_and_compute(prgm: Function) -> Function: + # desugar the input name and number additional statements for CFG construction + nspc = Namespace(prgm) + numbered_prgm, _ = number_statements(prgm) + desugared_prgm = fused_desugar(numbered_prgm) + cfg = fused_build_cfg(desugared_prgm) + liveness = LivenessAnalysis(cfg) + liveness.analyze() + for block in cfg.blocks.values(): + live_outputs = set( + liveness.input_states[block.id].keys() + ) # Backwards analysis, so live outputs are the input state of the block + live_inputs = set( + liveness.output_states[block.id].keys() + ) # Backwards analysis, so live inputs are the output state of the block + min_id, max_id = _get_stmt_bounds(block.statements) + numbered_prgm = _insert_lazy(numbered_prgm, min_id, live_inputs) + numbered_prgm = _insert_compute(numbered_prgm, max_id, live_outputs, nspc) + return Rewrite(PostWalk(Chain([_unwrap_numbered_stmt, _unnest_block])))( + numbered_prgm + ) diff --git a/src/galley_jl_python/fused/jit.py b/src/galley_jl_python/fused/jit.py new file mode 100644 index 0000000..77bd99d --- /dev/null +++ b/src/galley_jl_python/fused/jit.py @@ -0,0 +1,41 @@ +from collections.abc import Callable +from typing import Any + +from .dataflow import insert_lazy_and_compute +from .parser import fused_function_to_python_function, parse_fused_function + + +def jit(f: Callable[..., Any], /, ctx: Any = None) -> Callable[..., Any]: + """ + A decorator that compiles a function to defer and compute Galley tensors, so + that tensor operations between control flow are fused into one computation. + + Parameters: + - f: The function to compile. It can use basic python control flow and + operations (e.g. while, for, if), but not generators, classes or recursion. + - ctx: Unused; accepted for compatibility. + + Example usage: + @jit + def my_function(A, B, C): + D = A @ B + while some_condition(D): + D = D + C + return D + + is transformed into + + def my_function(A, B, C): + A, B = maybedefer((A, B)) + D = A @ B + D, = compute((D,)) + while some_condition(D): + C, D = maybedefer((C, D)) + D = D + C + D, = compute((D,)) + D, = compute((D,)) + return D + """ + fused_fn = parse_fused_function(f) + transformed_fn = insert_lazy_and_compute(fused_fn) + return fused_function_to_python_function(transformed_fn) diff --git a/src/galley_jl_python/fused/nodes.py b/src/galley_jl_python/fused/nodes.py new file mode 100644 index 0000000..58a9c66 --- /dev/null +++ b/src/galley_jl_python/fused/nodes.py @@ -0,0 +1,378 @@ +from __future__ import annotations + +from abc import ABC, abstractmethod +from dataclasses import dataclass +from typing import Any, Self + +from .symbolic import Context, NamedTerm, Term, TermTree, literal_repr, qual_repr + +""" +Finch Fused IR draft. + +This language models simple Python functions with imperative control flow: +- Assignment +- Calls +- If / else +- While +- For + +Nested function declarations are intentionally not representable in this IR. +""" + + +@dataclass(eq=True, frozen=True) +class FusedNode(Term, ABC): + @classmethod + def head(cls): + return cls + + @classmethod + def make_term(cls, head: Any, *children: Term) -> Self: + return head.from_children(*children) + + @classmethod + def from_children(cls, *children: Term) -> Self: + return cls(*children) + + def __str__(self): + ctx = FusedPrinterContext() + res = ctx(self) + return res if res is not None else ctx.emit() + + +@dataclass(eq=True, frozen=True) +class FusedTree(FusedNode, TermTree, ABC): + @property + @abstractmethod + def children(self) -> list[FusedNode]: # type: ignore[override] + ... + + +class FusedExpression(FusedNode, ABC): + pass + + +class FusedStatement(FusedNode, ABC): + pass + + +@dataclass(eq=True, frozen=True) +class Literal(FusedExpression): + val: Any + + def __repr__(self) -> str: + return literal_repr(type(self).__name__, {"val": self.val}) + + +@dataclass(eq=True, frozen=True) +class Variable(FusedExpression, NamedTerm): + name: str + type_: Any = None + + @property + def symbol(self) -> str: + return self.name + + def __repr__(self) -> str: + return literal_repr( + type(self).__name__, {"name": self.name, "type_": self.type_} + ) + + +@dataclass(eq=True, frozen=True) +class Keyword(FusedTree, FusedNode): + arg: str + value: FusedExpression + + @property + def children(self): + return [Literal(self.arg), self.value] + + @classmethod + def from_children(cls, arg, value): + match arg: + case Literal(val=name): + return cls(name, value) + case str(name): + return cls(name, value) + case _: + return cls(str(arg), value) + + +@dataclass(eq=True, frozen=True) +class Call(FusedTree, FusedExpression): + fn: FusedExpression + args: tuple[FusedExpression, ...] = () + kwargs: tuple[Keyword, ...] = () + + @property + def children(self): + return [self.fn, *self.args, *self.kwargs] + + @classmethod + def from_children(cls, fn, *args_and_kwargs): + args: list[FusedExpression] = [] + kwargs: list[Keyword] = [] + for child in args_and_kwargs: + match child: + case Keyword(): + kwargs.append(child) + case _: + args.append(child) + return cls(fn, tuple(args), tuple(kwargs)) + + +@dataclass(eq=True, frozen=True) +class Compare(FusedTree, FusedExpression): + left: FusedExpression + op: Literal + right: FusedExpression + + @property + def children(self): + return [self.left, self.op, self.right] + + +@dataclass(eq=True, frozen=True) +class BinaryOp(FusedTree, FusedExpression): + left: FusedExpression + op: Literal + right: FusedExpression + + @property + def children(self): + return [self.left, self.op, self.right] + + +@dataclass(eq=True, frozen=True) +class Block(FusedTree, FusedStatement): + body: tuple[FusedStatement, ...] = () + + @property + def children(self): + return [*self.body] + + @classmethod + def from_children(cls, *children): + return cls(tuple(children)) + + +@dataclass(eq=True, frozen=True) +class Assign(FusedTree, FusedStatement): + lhs: FusedExpression + rhs: FusedExpression + + @property + def children(self): + return [self.lhs, self.rhs] + + +@dataclass(eq=True, frozen=True) +class ExprStmt(FusedTree, FusedStatement): + value: FusedExpression + + @property + def children(self): + return [self.value] + + +@dataclass(eq=True, frozen=True) +class If(FusedTree, FusedStatement): + cond: FusedExpression + then_body: Block + else_body: Block | None = None + + @property + def children(self): + if self.else_body is None: + return [self.cond, self.then_body] + return [self.cond, self.then_body, self.else_body] + + @classmethod + def from_children(cls, cond, then_body, else_body=None): + return cls(cond, then_body, else_body) + + +@dataclass(eq=True, frozen=True) +class While(FusedTree, FusedStatement): + cond: FusedExpression + body: Block + + @property + def children(self): + return [self.cond, self.body] + + +@dataclass(eq=True, frozen=True) +class For(FusedTree, FusedStatement): + target: Variable + iterable: FusedExpression + body: Block + + @property + def children(self): + return [self.target, self.iterable, self.body] + + +@dataclass(eq=True, frozen=True) +class Return(FusedTree, FusedStatement): + values: tuple[FusedExpression, ...] = () + + @property + def children(self): + return [*self.values] + + @classmethod + def from_children(cls, *values): + return cls(tuple(values)) + + +@dataclass(eq=True, frozen=True) +class Break(FusedTree, FusedStatement): + @property + def children(self): + return [] + + +@dataclass(eq=True, frozen=True) +class Function(FusedTree): + name: Literal + params: tuple[Variable, ...] + body: Block + + @property + def children(self): + return [self.name, *self.params, self.body] + + @classmethod + def from_children(cls, name, *children): + if len(children) == 0: + raise ValueError("Function requires at least a body") + *params, body = children + return cls(name, tuple(params), body) + + +@dataclass(eq=True, frozen=True) +class Module(FusedTree): + functions: tuple[Function, ...] + + @property + def children(self): + return [*self.functions] + + @classmethod + def from_children(cls, *functions): + return cls(tuple(functions)) + + +class FusedPrinterContext(Context): + def __init__(self, tab=" ", indent=0): + super().__init__() + self.tab = tab + self.indent = indent + + @property + def feed(self) -> str: + return self.tab * self.indent + + def emit(self): + return "\n".join([*self.preamble, *self.epilogue]) + + def block(self): + blk = super().block() + blk.indent = self.indent + blk.tab = self.tab + return blk + + def subblock(self): + blk = self.block() + blk.indent = self.indent + 1 + return blk + + def __call__(self, prgm: FusedNode): + feed = self.feed + match prgm: + case Literal(value): + return qual_repr(value).replace("\n", "") + case Variable(name, _): + return name + case Call(fn, args, kwargs): + if isinstance(fn, Literal) and fn.val is tuple: + if len(args) == 1: + return f"({self(args[0])},)" + return f"({', '.join(self(arg) for arg in args)})" + all_args = [self(arg) for arg in args] + [self(kw) for kw in kwargs] + return f"{self(fn)}({', '.join(all_args)})" + case Keyword(arg, value): + return f"{arg}={self(value)}" + case Compare(left, op, right): + return f"({self(left)} {self(op)} {self(right)})" + case BinaryOp(left, op, right): + return f"({self(left)} {self(op)} {self(right)})" + case Block(body): + for stmt in body: + self(stmt) + return None + case Assign(lhs, rhs): + self.exec(f"{feed}{self(lhs)} = {self(rhs)}") + return None + case ExprStmt(value): + self.exec(f"{feed}{self(value)}") + return None + case If(cond, then_body, else_body): + self.exec(f"{feed}if {self(cond)}:") + then_ctx = self.subblock() + then_ctx(then_body) + if len(then_ctx.preamble) == 0: + then_ctx.exec(f"{then_ctx.feed}pass") + self.exec(then_ctx.emit()) + if else_body is not None: + self.exec(f"{feed}else:") + else_ctx = self.subblock() + else_ctx(else_body) + if len(else_ctx.preamble) == 0: + else_ctx.exec(f"{else_ctx.feed}pass") + self.exec(else_ctx.emit()) + return None + case While(cond, body): + self.exec(f"{feed}while {self(cond)}:") + body_ctx = self.subblock() + body_ctx(body) + if len(body_ctx.preamble) == 0: + body_ctx.exec(f"{body_ctx.feed}pass") + self.exec(body_ctx.emit()) + return None + case For(target, iterable, body): + self.exec(f"{feed}for {self(target)} in {self(iterable)}:") + body_ctx = self.subblock() + body_ctx(body) + if len(body_ctx.preamble) == 0: + body_ctx.exec(f"{body_ctx.feed}pass") + self.exec(body_ctx.emit()) + return None + case Return(values): + if len(values) == 0: + self.exec(f"{feed}return") + elif len(values) == 1: + self.exec(f"{feed}return {self(values[0])}") + else: + self.exec(f"{feed}return ({', '.join(self(v) for v in values)})") + return None + case Function(name, params, body): + param_str = ", ".join(self(param) for param in params) + self.exec(f"{feed}def {name.val}({param_str}):") + body_ctx = self.subblock() + body_ctx(body) + if len(body_ctx.preamble) == 0: + body_ctx.exec(f"{body_ctx.feed}pass") + self.exec(body_ctx.emit()) + return None + case Module(functions): + for i, func in enumerate(functions): + if i > 0: + self.exec("") + self(func) + return None + case prgm: + self.exec(f"{self.feed}{str(prgm)}") + return None diff --git a/src/galley_jl_python/fused/parser.py b/src/galley_jl_python/fused/parser.py new file mode 100644 index 0000000..179cda8 --- /dev/null +++ b/src/galley_jl_python/fused/parser.py @@ -0,0 +1,505 @@ +from __future__ import annotations + +import ast +import builtins +import inspect +import operator +import textwrap +import types +from collections.abc import Callable +from typing import Any + +from . import nodes as fzd + +_BIN_OPS = { + ast.Add: operator.add, + ast.Sub: operator.sub, + ast.Mult: operator.mul, + ast.Div: operator.truediv, + ast.FloorDiv: operator.floordiv, + ast.Mod: operator.mod, + ast.Pow: operator.pow, + ast.MatMult: operator.matmul, + ast.BitAnd: operator.and_, + ast.BitOr: operator.or_, + ast.BitXor: operator.xor, + ast.LShift: operator.lshift, + ast.RShift: operator.rshift, +} + +_CMP_OPS = { + ast.Eq: operator.eq, + ast.NotEq: operator.ne, + ast.Lt: operator.lt, + ast.LtE: operator.le, + ast.Gt: operator.gt, + ast.GtE: operator.ge, + ast.Is: operator.is_, + ast.IsNot: operator.is_not, + ast.In: lambda x, y: x in y, + ast.NotIn: lambda x, y: x not in y, +} + +_UNARY_OPS = { + ast.UAdd: operator.pos, + ast.USub: operator.neg, + ast.Invert: operator.invert, + ast.Not: operator.not_, +} + + +def _logical_and(a, b): + return a and b + + +def _logical_or(a, b): + return a or b + + +def _if_expr(cond, then_val, else_val): + return then_val if cond else else_val + + +# Distinct from `operator.and_`/`operator.or_` so that `&`/`|` are not lowered to +# `and`/`or`. +_BOOL_OPS = { + ast.And: _logical_and, + ast.Or: _logical_or, +} + +_REV_BIN_OPS = {fn: op for op, fn in _BIN_OPS.items()} +_REV_CMP_OPS = {fn: op for op, fn in _CMP_OPS.items()} +_REV_UNARY_OPS = {fn: op for op, fn in _UNARY_OPS.items()} +_REV_BOOL_OPS = {fn: op for op, fn in _BOOL_OPS.items()} + + +class _FusedFunctionParser: + def __init__(self, fn: types.FunctionType, fn_def: ast.FunctionDef): + self.fn = fn + self.fn_def = fn_def + self.globals = getattr(fn, "__globals__", {}) + self.closurevars: dict[str, Any] = {} + if hasattr(fn, "__code__") and hasattr(fn, "__closure__") and fn.__closure__: + for name, cell in zip(fn.__code__.co_freevars, fn.__closure__, strict=True): + self.closurevars[name] = cell.cell_contents + self.locals: set[str] = set() + + def parse(self) -> fzd.Function: + if len(self.fn_def.decorator_list) > 1: # The @jit decorator is in this list + raise self._unsupported( + self.fn_def, + "Decorated functions are not supported in the fused JIT parser.", + ) + + params = tuple(self._parse_parameter(arg) for arg in self.fn_def.args.args) + body = self._parse_block(self.fn_def.body) + return fzd.Function(fzd.Literal(self.fn_def.name), params, body) + + def _parse_parameter(self, arg: ast.arg) -> fzd.Variable: + if arg.annotation is not None: + raise self._unsupported(arg, "Annotated parameters are not supported.") + self.locals.add(arg.arg) + return fzd.Variable(arg.arg) + + def _parse_block(self, stmts: list[ast.stmt]) -> fzd.Block: + return fzd.Block(tuple(self._parse_stmt(stmt) for stmt in stmts)) + + def _parse_stmt(self, stmt: ast.stmt) -> fzd.FusedStatement: + match stmt: + case ast.Assign(targets=[ast.Name(id=name)], value=value): + rhs = self._parse_expr(value) + self.locals.add(name) + return fzd.Assign(fzd.Variable(name), rhs) + case ast.AugAssign(target=ast.Name(id=name), op=op, value=value): + lhs = fzd.Variable(name) + rhs = fzd.BinaryOp(lhs, self._parse_op(op), self._parse_expr(value)) + self.locals.add(name) + return fzd.Assign(lhs, rhs) + case ast.Expr(value=value): + return fzd.ExprStmt(self._parse_expr(value)) + case ast.If(test=test, body=body, orelse=orelse): + else_block = self._parse_block(orelse) if orelse else None + return fzd.If( + self._parse_expr(test), self._parse_block(body), else_block + ) + case ast.While(test=test, body=body, orelse=[]): + return fzd.While(self._parse_expr(test), self._parse_block(body)) + case ast.For(target=ast.Name(id=name), iter=iterable, body=body, orelse=[]): + self.locals.add(name) + return fzd.For( + fzd.Variable(name), + self._parse_expr(iterable), + self._parse_block(body), + ) + case ast.Break(): + return fzd.Break() + case ast.Return(value=None): + return fzd.Return(()) + case ast.Return(value=value): + assert value is not None + match value: + case ast.Tuple(elts=elts): + return fzd.Return(tuple(self._parse_expr(elt) for elt in elts)) + case _: + return fzd.Return((self._parse_expr(value),)) + case ast.FunctionDef() | ast.AsyncFunctionDef(): + raise self._unsupported( + stmt, + "Local functions are not supported in the fused JIT parser.", + ) + case ast.For(orelse=orelse) if orelse: + raise self._unsupported(stmt, "For-else blocks are not supported.") + case ast.While(orelse=orelse) if orelse: + raise self._unsupported(stmt, "While-else blocks are not supported.") + case _: + raise self._unsupported( + stmt, + f"Unsupported statement type: {type(stmt).__name__}", + ) + + def _parse_expr(self, expr: ast.expr) -> fzd.FusedExpression: + match expr: + case ast.Constant(value=value): + return fzd.Literal(value) + case ast.Name(id=name): + return self._parse_name(name) + case ast.Call(func=func, args=args, keywords=keywords): + return fzd.Call( + self._parse_expr(func), + tuple(self._parse_expr(arg) for arg in args), + tuple(self._parse_keyword(kw) for kw in keywords), + ) + case ast.BinOp(left=left, op=op, right=right): + return fzd.BinaryOp( + self._parse_expr(left), + self._parse_op(op), + self._parse_expr(right), + ) + case ast.Compare(left=left, ops=[op], comparators=[right]): + return fzd.Compare( + self._parse_expr(left), + self._parse_op(op), + self._parse_expr(right), + ) + case ast.Compare() as cmp if len(cmp.ops) > 1: + return self._parse_chained_compare(cmp) + case ast.BoolOp(op=op, values=values) if len(values) >= 2: + return self._parse_bool_op(op, values) + case ast.UnaryOp(op=op, operand=operand): + return fzd.Call(self._parse_op(op), (self._parse_expr(operand),)) + case ast.Tuple(elts=elts): + return fzd.Call( + fzd.Literal(tuple), + tuple(self._parse_expr(elt) for elt in elts), + ) + case ast.IfExp(test=test, body=body, orelse=orelse): + return fzd.Call( + fzd.Literal(_if_expr), + ( + self._parse_expr(test), + self._parse_expr(body), + self._parse_expr(orelse), + ), + ) + case ast.Attribute(value=value, attr=attr): + base = self._parse_expr(value) + return fzd.Call(fzd.Literal(getattr), (base, fzd.Literal(attr))) + case _: + raise self._unsupported( + expr, + f"Unsupported expression type: {type(expr).__name__}", + ) + + def _parse_name(self, name: str) -> fzd.FusedExpression: + if name in self.locals: + return fzd.Variable(name) + if name in self.closurevars: + return fzd.Literal(self.closurevars[name]) + if name in self.globals: + return fzd.Literal(self.globals[name]) + if hasattr(builtins, name): + return fzd.Literal(getattr(builtins, name)) + return fzd.Variable(name) + + def _parse_op( + self, op: ast.operator | ast.unaryop | ast.boolop | ast.cmpop + ) -> fzd.Literal: + op_type: Any = type(op) + for table in (_BIN_OPS, _UNARY_OPS, _BOOL_OPS, _CMP_OPS): + fn = table.get(op_type) + if fn is not None: + return fzd.Literal(fn) + raise self._unsupported(op, f"Unsupported operator type: {type(op).__name__}") + + def _parse_chained_compare(self, cmp: ast.Compare) -> fzd.FusedExpression: + first = fzd.Compare( + self._parse_expr(cmp.left), + self._parse_op(cmp.ops[0]), + self._parse_expr(cmp.comparators[0]), + ) + expr: fzd.FusedExpression = first + for i in range(1, len(cmp.ops)): + next_cmp = fzd.Compare( + self._parse_expr(cmp.comparators[i - 1]), + self._parse_op(cmp.ops[i]), + self._parse_expr(cmp.comparators[i]), + ) + expr = fzd.BinaryOp(expr, fzd.Literal(_logical_and), next_cmp) + return expr + + def _parse_bool_op( + self, op: ast.boolop, values: list[ast.expr] + ) -> fzd.FusedExpression: + expr = self._parse_expr(values[0]) + op_lit = self._parse_op(op) + for value in values[1:]: + expr = fzd.BinaryOp(expr, op_lit, self._parse_expr(value)) + return expr + + def _parse_keyword(self, kw: ast.keyword) -> fzd.Keyword: + match kw.arg: + case None: + raise self._unsupported(kw, "Unsupported **kwargs unpacking in call") + case str(name): + return fzd.Keyword(name, self._parse_expr(kw.value)) + + def _unsupported(self, node: ast.AST, message: str) -> ValueError: + lineno = getattr(node, "lineno", "?") + return ValueError(f"{message} (line {lineno})") + + +def parse_fused_function(fn: types.FunctionType) -> fzd.Function: + source = textwrap.dedent(inspect.getsource(fn)) + tree = ast.parse(source) + fn_name = getattr(fn, "__name__", None) + + for node in tree.body: + match node: + case ast.FunctionDef(name=name) if fn_name is None or name == fn_name: + return _FusedFunctionParser(fn, node).parse() + case ast.AsyncFunctionDef(name=name) if fn_name is None or name == fn_name: + raise ValueError( + "Async functions are not supported in the fused JIT parser." + ) + + raise ValueError("Expected a top-level Python function definition.") + + +def parse_fused_module(*functions: types.FunctionType) -> fzd.Module: + return fzd.Module(tuple(parse_fused_function(fn) for fn in functions)) + + +def parse_function(fn: types.FunctionType) -> fzd.Function: + return parse_fused_function(fn) + + +class _FusedToPythonAST: + def __init__(self) -> None: + self._bool_fns = tuple(_REV_BOOL_OPS) + self._extra_globals: dict[str, Any] = {} + + def parse_function(self, function: fzd.Function) -> ast.FunctionDef: + if not isinstance(function.name.val, str): + raise ValueError("Fused function name must be a string literal.") + + args = ast.arguments( + posonlyargs=[], + args=[ + ast.arg(arg=param.name, annotation=None) for param in function.params + ], + kwonlyargs=[], + kw_defaults=[], + defaults=[], + vararg=None, + kwarg=None, + ) + + body = [self._stmt_to_ast(stmt) for stmt in function.body.body] + if len(body) == 0: + body = [ast.Pass()] + + return ast.FunctionDef( # noqa: PGH003 + name=function.name.val, + args=args, + body=body, + decorator_list=[], + returns=None, + type_comment=None, + type_params=[], + ) + + def parse_module(self, module: fzd.Module) -> ast.Module: + return ast.Module( + body=[self.parse_function(function) for function in module.functions], + type_ignores=[], + ) + + def _stmt_to_ast(self, stmt: fzd.FusedStatement) -> ast.stmt: + match stmt: + case fzd.Assign(lhs=lhs, rhs=rhs): + return ast.Assign( + targets=[self._target_to_ast(lhs)], + value=self._expr_to_ast(rhs), + ) + case fzd.ExprStmt(value=value): + return ast.Expr(value=self._expr_to_ast(value)) + case fzd.If(cond=cond, then_body=then_body, else_body=else_body): + return ast.If( + test=self._expr_to_ast(cond), + body=self._block_to_ast(then_body), + orelse=self._block_to_ast(else_body) + if else_body is not None + else [], + ) + case fzd.While(cond=cond, body=body): + return ast.While( + test=self._expr_to_ast(cond), + body=self._block_to_ast(body), + orelse=[], + ) + case fzd.For(target=fzd.Variable(name=name), iterable=iterable, body=body): + return ast.For( + target=ast.Name(id=name, ctx=ast.Store()), + iter=self._expr_to_ast(iterable), + body=self._block_to_ast(body), + orelse=[], + type_comment=None, + ) + case fzd.Return(values=values): + if len(values) == 0: + return ast.Return(value=None) + if len(values) == 1: + return ast.Return(value=self._expr_to_ast(values[0])) + return ast.Return( + value=ast.Tuple( + elts=[self._expr_to_ast(value) for value in values], + ctx=ast.Load(), + ) + ) + case fzd.Break(): + return ast.Break() + case _: + raise ValueError( + f"Unsupported fused statement type: {type(stmt).__name__}" + ) + + def _target_to_ast(self, target: fzd.FusedExpression) -> ast.expr: + match target: + case fzd.Variable(name=name): + return ast.Name(id=name, ctx=ast.Store()) + case fzd.Call(fn=fzd.Literal(val=fn), args=args) if fn is tuple: + return ast.Tuple( + elts=[self._target_to_ast(arg) for arg in args], + ctx=ast.Store(), + ) + case _: + raise ValueError( + f"Unsupported fused assignment target: {type(target).__name__}" + ) + + def _block_to_ast(self, block: fzd.Block) -> list[ast.stmt]: + body = [self._stmt_to_ast(stmt) for stmt in block.body] + return body if len(body) > 0 else [ast.Pass()] + + def _expr_to_ast(self, expr: fzd.FusedExpression) -> ast.expr: + match expr: + case fzd.Variable(name=name): + return ast.Name(id=name, ctx=ast.Load()) + case fzd.Literal(val=value): + return self._literal_to_expr(value) + case fzd.Call(fn=fn, args=args, kwargs=kwargs): + if isinstance(fn, fzd.Literal) and fn.val is tuple: + return ast.Tuple( + elts=[self._expr_to_ast(arg) for arg in args], + ctx=ast.Load(), + ) + return ast.Call( + func=self._expr_to_ast(fn), + args=[self._expr_to_ast(arg) for arg in args], + keywords=[ + ast.keyword(arg=kw.arg, value=self._expr_to_ast(kw.value)) + for kw in kwargs + ], + ) + case fzd.Compare(left=left, op=fzd.Literal(val=op_fn), right=right): + cmp_op = _REV_CMP_OPS.get(op_fn) + if cmp_op is None: + raise ValueError(f"Unsupported fused compare operator: {op_fn!r}") + return ast.Compare( + left=self._expr_to_ast(left), + ops=[cmp_op()], + comparators=[self._expr_to_ast(right)], + ) + case fzd.BinaryOp(left=left, op=fzd.Literal(val=op_fn), right=right): + bool_op = _REV_BOOL_OPS.get(op_fn) + if bool_op is not None: + return ast.BoolOp( + op=bool_op(), + values=[self._expr_to_ast(left), self._expr_to_ast(right)], + ) + bin_op = _REV_BIN_OPS.get(op_fn) + if bin_op is None: + raise ValueError(f"Unsupported fused binary operator: {op_fn!r}") + return ast.BinOp( + left=self._expr_to_ast(left), + op=bin_op(), + right=self._expr_to_ast(right), + ) + case _: + raise ValueError( + f"Unsupported fused expression type: {type(expr).__name__}" + ) + + def _global_name(self, name: str, value: Any) -> ast.expr: + # Freshen the name so that distinct objects sharing a name (e.g. `np.add` + # and `operator.add`) do not clobber each other. + fresh_name = name + counter = 1 + while self._extra_globals.get(fresh_name, value) is not value: + fresh_name = f"{name}_{counter}" + counter += 1 + self._extra_globals[fresh_name] = value + return ast.Name(id=fresh_name, ctx=ast.Load()) + + def _literal_to_expr(self, value: Any) -> ast.expr: + if value is None or isinstance( + value, str | bytes | int | float | complex | bool + ): + return ast.Constant(value=value) + + if callable(value): + if getattr(builtins, getattr(value, "__name__", ""), None) is value: + return ast.Name(id=value.__name__, ctx=ast.Load()) # type: ignore[union-attr] + + name = getattr(value, "__name__", None) + if name is not None and name.isidentifier(): + return self._global_name(name, value) + + if isinstance(value, types.ModuleType): + return self._global_name(value.__name__, value) + + raise ValueError(f"Literal cannot be represented in Python AST: {value!r}") + + +def fused_function_to_python_ast(function: fzd.Function) -> ast.FunctionDef: + node = _FusedToPythonAST().parse_function(function) + return ast.fix_missing_locations(node) + + +def fused_function_to_python_function(function: fzd.Function) -> Callable: + converter = _FusedToPythonAST() + func_def = ast.fix_missing_locations(converter.parse_function(function)) + module_node = ast.Module(body=[func_def], type_ignores=[]) + code = compile(module_node, "", "exec") + globals_dict: dict[str, Any] = { + "__builtins__": builtins, + "operator": operator, + **converter._extra_globals, + } + exec(code, globals_dict) # noqa: S102 + return globals_dict[function.name.val] + + +def fused_module_to_python_ast(module: fzd.Module) -> ast.Module: + node = _FusedToPythonAST().parse_module(module) + return ast.fix_missing_locations(node) diff --git a/src/galley_jl_python/fused/symbolic.py b/src/galley_jl_python/fused/symbolic.py new file mode 100644 index 0000000..89e631c --- /dev/null +++ b/src/galley_jl_python/fused/symbolic.py @@ -0,0 +1,301 @@ +""" +Term, rewriting and dataflow utilities used by the fused JIT. + +Adapted from `finch.symbolic` in finch-tensor (MIT licensed), keeping only what +the fused JIT needs. +""" + +import re +from abc import ABC, abstractmethod +from collections import defaultdict +from collections.abc import Callable, Iterable, Iterator +from dataclasses import dataclass +from inspect import isbuiltin, isclass, isfunction +from typing import Any, Generic, Self, TypeVar + + +class Term: + @abstractmethod + def head(self) -> Callable[..., Self]: + """Return the head type of the S-expression.""" + ... + + @classmethod + @abstractmethod + def make_term(cls, head: Callable[..., Self], *children: "Term") -> Self: + """ + Construct a new term in the same family of terms with the given head type and + children. This function should satisfy + `x == x.make_term(x.head(), *x.children)` + """ + ... + + +@dataclass(frozen=True, eq=True) +class TermTree(Term, ABC): + @property + @abstractmethod + def children(self) -> list[Term]: + """Return the children (AKA tail) of the S-expression.""" + ... + + +class NamedTerm(Term, ABC): + """ + A named program node, e.g. a variable, function, or index. + """ + + @property + @abstractmethod + def symbol(self) -> str: ... + + +def qual_repr(val: Any) -> str: + if isbuiltin(val) or isclass(val) or isfunction(val): + return f"{val.__module__}.{val.__qualname__}" + return repr(val) + + +def literal_repr(name: str, fields: dict[str, Any]) -> str: + return ( + name + "(" + ", ".join([f"{k}={qual_repr(v)}" for k, v in fields.items()]) + ")" + ) + + +def PostOrderDFS(node: Term) -> Iterator[Term]: + if isinstance(node, TermTree): + for arg in node.children: + yield from PostOrderDFS(arg) + yield node + + +T = TypeVar("T", bound=Term) + +RwCallable = Callable[[T], T | None] + + +def default_rewrite(x: T | None, y: T) -> T: + return x if x is not None else y + + +class Rewrite(Generic[T]): + """ + A rewriter which returns the original argument even if `rw` returns nothing. + """ + + def __init__(self, rw: RwCallable): + self.rw = rw + + def __call__(self, x: T) -> T: + return default_rewrite(self.rw(x), x) + + +class PostWalk: + """ + A rewriter which recursively rewrites the arguments of each node using + `rw`, then rewrites the resulting node. If all rewriters return `nothing`, + returns `nothing`. + """ + + def __init__(self, rw: RwCallable): + self.rw = rw + + def __call__(self, x: T) -> T | None: + if isinstance(x, TermTree): + args = x.children + new_args = list(map(self, args)) + if all(arg is None for arg in new_args): + return self.rw(x) + y = x.make_term( + x.head(), *map(lambda x1, x2: default_rewrite(x1, x2), new_args, args) + ) + assert isinstance(y, type(x)) + return default_rewrite(self.rw(y), y) + return self.rw(x) + + +class Chain(Generic[T]): + """ + A rewriter which rewrites using each rewriter in `rws`. If all rewriters + return `nothing`, return `nothing`. + """ + + def __init__(self, rws: Iterable[RwCallable]): + self.rws = rws + + def __call__(self, x: T) -> T | None: + is_success = False + for rw in self.rws: + y = rw(x) + if y is not None: + is_success = True + x = y + if is_success: + return x + return None + + +class Namespace: + """ + A namespace for managing variable names and aesthetic fresh variable + generation. You can construct a namespace from an existing tree of `Term`s + to avoid name collisions. + """ + + def __init__(self, root=None): + self.counts = defaultdict(int) + if root is not None: + for node in PostOrderDFS(root): + if isinstance(node, NamedTerm): + self.freshen(node.symbol) + + def freshen(self, *tags) -> str: + """ + Generate a fresh variable name based on the provided tags. + """ + name = "_".join(str(tag) for tag in tags) + m = re.match(r"^(.*)_(\d+)$", name) + if m is None: + tag = name + n = 1 + else: + tag = m.group(1) + n = int(m.group(2)) + n = max(self.counts[tag] + 1, n) + self.counts[tag] = n + if n == 1: + return tag + return f"{tag}_{n}" + + +class Context(ABC): + def __init__(self, namespace=None, preamble=None, epilogue=None): + self.namespace = namespace if namespace is not None else Namespace() + self.preamble = preamble if preamble is not None else [] + self.epilogue = epilogue if epilogue is not None else [] + + def exec(self, thunk: Any): + self.preamble.append(thunk) + + def block(self): + """ + Create a new block. Preambles and epilogues will stay within this block. + """ + blk = self.__class__() + blk.namespace = self.namespace + blk.preamble = [] + blk.epilogue = [] + return blk + + @abstractmethod + def emit(self): + """ + Emit the code in this context. + """ + ... + + +class BasicBlock: + """Linear sequence of statements with a single entry and exit.""" + + def __init__(self, id: str) -> None: + self.id = id + self.statements: list = [] + self.successors: list[BasicBlock] = [] + self.predecessors: list[BasicBlock] = [] + + def add_statement(self, statement) -> None: + self.statements.append(statement) + + def add_successor(self, successor: "BasicBlock") -> None: + if successor not in self.successors: + self.successors.append(successor) + + if self not in successor.predecessors: + successor.predecessors.append(self) + + def __str__(self) -> str: + succ_names = [succ.id for succ in self.successors] + lines = [f"{self.id}: #succs=[{', '.join(succ_names)}]"] + lines.extend(f" {stmt}" for stmt in self.statements) + return "\n".join(lines) + + +class ControlFlowGraph: + """Collection of BasicBlocks plus explicit ENTRY and EXIT nodes.""" + + def __init__(self) -> None: + self.block_counter = 0 + self.block_name = "" + self.blocks: dict[str, BasicBlock] = {} + + self.entry_block = self.new_block_custom("ENTRY") + self.exit_block = self.new_block_custom("EXIT") + + def new_block(self) -> BasicBlock: + return self.new_block_custom(self.block_name) + + def new_block_custom(self, name: str) -> BasicBlock: + bid = f"{name}_{self.block_counter}" + self.block_counter += 1 + block = BasicBlock(bid) + self.blocks[bid] = block + return block + + def __str__(self) -> str: + return "\n\n".join(str(block) for block in self.blocks.values()) + + +class DataFlowAnalysis(ABC): + """ + Base class for performing data flow analyses over a ControlFlowGraph with a + work-list algorithm, in either a "forward" or "backward" direction. + + `input_states` and `output_states` map each basic block id to the lattice + state at the entry and exit of the block. + """ + + def __init__(self, cfg: ControlFlowGraph): + self.cfg: ControlFlowGraph = cfg + self.input_states: dict[str, dict] = { + block.id: {} for block in cfg.blocks.values() + } + self.output_states: dict[str, dict] = { + block.id: {} for block in cfg.blocks.values() + } + + @abstractmethod + def stmt_str(self, stmt, state: dict) -> str: ... + + @abstractmethod + def transfer(self, stmts, state: dict) -> dict: ... + + @abstractmethod + def join(self, state_1: dict, state_2: dict) -> dict: ... + + @abstractmethod + def direction(self) -> str: ... + + def analyze(self) -> None: + forward = self.direction() == "forward" + work_list: list[BasicBlock] = list(self.cfg.blocks.values()) + while work_list: + block = work_list.pop(0) + neighbors = block.predecessors if forward else block.successors + if not neighbors: + input_state = {} + else: + neighbor_outputs = [ + self.output_states.get(neighbor.id, {}) for neighbor in neighbors + ] + input_state = neighbor_outputs[0].copy() + for neighbor_output in neighbor_outputs[1:]: + input_state = self.join(input_state, neighbor_output) + + self.input_states[block.id] = input_state + output_state = self.transfer(block.statements, input_state) + if output_state != self.output_states.get(block.id, {}): + self.output_states[block.id] = output_state + for dependent in block.successors if forward else block.predecessors: + if dependent not in work_list: + work_list.append(dependent) diff --git a/src/galley_jl_python/tensor.py b/src/galley_jl_python/tensor.py index 8f98335..89161e1 100644 --- a/src/galley_jl_python/tensor.py +++ b/src/galley_jl_python/tensor.py @@ -289,7 +289,8 @@ def __getitem__(self, key): @property def dtype(self) -> DType: - return jl.eltype(self._obj.body) + # `eltype` of the swizzle (not its body) also works for lazy tensors. + return jl.eltype(self._obj) @property def ndim(self) -> int: diff --git a/tests/test_fused.py b/tests/test_fused.py index ac6a456..4baa613 100644 --- a/tests/test_fused.py +++ b/tests/test_fused.py @@ -73,3 +73,13 @@ def opt_fn(A, B, n): B = gl.Tensor(np.array([[1, 0], [0, 1]])) assert_galley_allclose(opt_fn(A, B, 3), A.todense() + 3 * B.todense()) + + +def test_jit_scalar_operand_uses_galley_tensors(): + @gl.jit + def opt_fn(A): + return A * 2 + A + + A = gl.Tensor(np.array([[1, 2], [3, 4]])) + + assert_galley_allclose(opt_fn(A), 3 * A.todense()) diff --git a/tests/test_fused_compiler.py b/tests/test_fused_compiler.py new file mode 100644 index 0000000..a1d26a3 --- /dev/null +++ b/tests/test_fused_compiler.py @@ -0,0 +1,406 @@ +import ast +import operator +import sys +import textwrap + +import pytest + +import numpy as np + +import galley_jl_python as gl +from galley_jl_python import add, matmul +from galley_jl_python.fused import nodes as fzd +from galley_jl_python.fused.cfg_builder import ( + fused_build_cfg, + fused_desugar, + number_statements, +) +from galley_jl_python.fused.dataflow import LivenessAnalysis +from galley_jl_python.fused.parser import ( + fused_function_to_python_ast, + parse_fused_function, +) + + +def test_jit_does_not_depend_on_finch_tensor(): + @gl.jit + def opt_fn(A, B): + return gl.add(A, B) + + A = gl.Tensor(np.array([[1, 2], [3, 4]])) + opt_fn(A, A) + + assert not any(name == "finch" or name.startswith("finch.") for name in sys.modules) + + +def test_parse_simple_function_with_control_flow_and_calls(): + def simple_fn(fn, n): + total = 0 + for i in range(n): + if i < n: # noqa: SIM108 + total = fn(total, i, scale=2) + else: + total = total - 1 + while total < n: + total = total + 1 + return total + + result = parse_fused_function(simple_fn) + + expected = fzd.Function( + fzd.Literal("simple_fn"), + (fzd.Variable("fn"), fzd.Variable("n")), + fzd.Block( + ( + fzd.Assign(fzd.Variable("total"), fzd.Literal(0)), + fzd.For( + fzd.Variable("i"), + fzd.Call(fzd.Literal(range), (fzd.Variable("n"),)), + fzd.Block( + ( + fzd.If( + fzd.Compare( + fzd.Variable("i"), + fzd.Literal(operator.lt), + fzd.Variable("n"), + ), + fzd.Block( + ( + fzd.Assign( + fzd.Variable("total"), + fzd.Call( + fzd.Variable("fn"), + ( + fzd.Variable("total"), + fzd.Variable("i"), + ), + ( + fzd.Keyword( + "scale", + fzd.Literal(2), + ), + ), + ), + ), + ) + ), + fzd.Block( + ( + fzd.Assign( + fzd.Variable("total"), + fzd.BinaryOp( + fzd.Variable("total"), + fzd.Literal(operator.sub), + fzd.Literal(1), + ), + ), + ) + ), + ), + ) + ), + ), + fzd.While( + fzd.Compare( + fzd.Variable("total"), + fzd.Literal(operator.lt), + fzd.Variable("n"), + ), + fzd.Block( + ( + fzd.Assign( + fzd.Variable("total"), + fzd.BinaryOp( + fzd.Variable("total"), + fzd.Literal(operator.add), + fzd.Literal(1), + ), + ), + ) + ), + ), + fzd.Return((fzd.Variable("total"),)), + ) + ), + ) + + assert result == expected + + +def test_parse_rejects_local_function_definitions(): + def with_local_fn(x): + def inner(y): + return y + 1 + + return inner(x) + + with pytest.raises(ValueError, match="Local functions are not supported"): + parse_fused_function(with_local_fn) + + +def test_parse_rejects_for_else_blocks(): + def with_for_else(n): + for i in range(n): + n = n + i + else: + n = n + 1 + return n + + with pytest.raises(ValueError, match="For-else blocks are not supported"): + parse_fused_function(with_for_else) + + +def test_parse_rejects_while_else_blocks(): + def with_while_else(n): + while n < 3: + n = n + 1 + else: + n = n + 2 + return n + + with pytest.raises(ValueError, match="While-else blocks are not supported"): + parse_fused_function(with_while_else) + + +def test_parse_reverse_parse_is_lossless_on_supported_subset(): + """ + Round-tripping recovers the source, with numeric literals left as literals. + """ + + def roundtrip_fn(n): + total = 0 + for i in range(n): + if i < n: # noqa: SIM108 + total = total + i + else: + total = total - 1 + while total < n: + total = total + 1 + return total + + expected_source = textwrap.dedent("""\ + def roundtrip_fn(n): + total = 0 + for i in range(n): + if i < n: + total = total + i + else: + total = total - 1 + while total < n: + total = total + 1 + return total + """) + expected_fn = ast.parse(expected_source).body[0] + + fused_fn = parse_fused_function(roundtrip_fn) + roundtrip_fn_ast = fused_function_to_python_ast(fused_fn) + + assert ast.dump(expected_fn, include_attributes=False) == ast.dump( + roundtrip_fn_ast, + include_attributes=False, + ) + + +def test_cfg_builder(): + def simple_fn(fn, n): + total = 0 + for i in range(n): + if i < n: # noqa: SIM108 + total = fn(total, i) + else: + total = total - 1 + while total < n: + total = total + 1 + return total + + fused_fn = parse_fused_function(simple_fn) + numbered_fn, _ = number_statements(fused_fn) + desugared_fn = fused_desugar(numbered_fn) + cfg = fused_build_cfg(desugared_fn) + + # We won't assert on the exact structure of the CFG here, but we can at least + # check that it has the expected number of blocks. The exact number of blocks + # may depend on how the CFG builder handles certain constructs, so this is a + # somewhat loose check. + assert ( + len(cfg.blocks) >= 5 + ) # Entry block, for loop block, if block, while block, return block + + +def _build_liveness(fn): + """Helper: parse, number, desugar, build CFG, run liveness.""" + fused_fn = parse_fused_function(fn) + numbered_fn, _ = number_statements(fused_fn) + desugared_fn = fused_desugar(numbered_fn) + cfg = fused_build_cfg(desugared_fn) + liveness = LivenessAnalysis(cfg) + liveness.analyze() + return liveness, cfg + + +def _all_live_names(liveness, cfg): + """Union of all live variable names across all blocks (in and out).""" + names = set() + for block in cfg.blocks.values(): + names |= {v.name for v in liveness.output_states[block.id]} + names |= {v.name for v in liveness.input_states[block.id]} + return names + + +def test_liveness_straight_line(): + """Parameters must appear live at the function entry block.""" + + def fn(a, b): + c = add(a, b) + return c # noqa: RET504 + + liveness, cfg = _build_liveness(fn) + + # After desugaring, the function body is a single block. + # live-IN (output_states) of that block = {a, b}, since both are used. + # c is defined and consumed in the same block so it never crosses a block + # boundary and does not appear in any block-boundary state. + all_live_in = set() + for block in cfg.blocks.values(): + all_live_in |= {v.name for v in liveness.output_states[block.id]} + + assert "a" in all_live_in + assert "b" in all_live_in + assert "c" not in all_live_in + + +def test_liveness_dead_variable(): + """A variable assigned but never used afterwards must not be live after.""" + + def fn(a, b): + unused = add(a, b) # noqa: F841 + c = matmul(a, b) + return c # noqa: RET504 + + liveness, cfg = _build_liveness(fn) + + exit_block = list(cfg.blocks.values())[-1] + live_at_exit = {v.name for v in liveness.input_states[exit_block.id]} + assert "unused" not in live_at_exit + + +def test_liveness_loop_carried(): + """Loop-carried variables must be live at the top of the loop body.""" + + def fn(n): + total = 0 + for _i in range(n): + total = total + 1 + return total + + liveness, cfg = _build_liveness(fn) + names = _all_live_names(liveness, cfg) + + assert "total" in names + assert "n" in names + + +def test_liveness_multi_loop_carried(): + """Multiple loop-carried variables must all be live inside the loop.""" + + def fn(A, B, C, n): + D = matmul(A, B) + E = add(A, C) + for _i in range(n): + D = add(D, E) + return D + + liveness, cfg = _build_liveness(fn) + names = _all_live_names(liveness, cfg) + + assert "D" in names + assert "E" in names + assert "n" in names + + +def test_liveness_if_branch_merges(): + """Variables used in either branch must be live before the if.""" + + def fn(cond, a, b): + if cond: # noqa: SIM108 + result = add(a, b) + else: + result = matmul(a, b) + return result + + liveness, cfg = _build_liveness(fn) + names = _all_live_names(liveness, cfg) + + assert "a" in names + assert "b" in names + assert "result" in names + + +def _increment(x): + return x + 1 + + +def _decrement(x): + return x - 1 + + +_decrement.__name__ = "_increment" + + +def test_jit_keeps_bitwise_and_logical_operators_apart(): + @gl.jit + def opt_fn(a, b): + return a | b, a & b, a or b, a and b + + assert opt_fn(1, 2) == (3, 0, 1, 2) + + +def test_jit_loop_variable_is_defined_by_the_loop(): + @gl.jit + def opt_fn(n): + total = 0 + for i in range(n): + total = total + i + return total + + assert opt_fn(4) == 6 + + +def test_jit_break(): + @gl.jit + def opt_fn(limit): + i = 0 + while i < 10: + i = i + 1 + if i > limit: + break + return i + + assert opt_fn(3) == 4 + + +def test_jit_expression_statement(capsys): + @gl.jit + def opt_fn(x): + print(x) + return x + + assert opt_fn(5) == 5 + assert capsys.readouterr().out == "5\n" + + +def test_jit_conditional_expression(): + @gl.jit + def opt_fn(c): + return 1 if c else 2 + + assert opt_fn(True) == 1 + assert opt_fn(False) == 2 + + +def test_jit_distinct_callables_with_the_same_name(): + @gl.jit + def opt_fn(x): + return _increment(x), _decrement(x) + + assert opt_fn(5) == (6, 4)