From fe100534db1a4dd18ebeb0af5b42c63d57f71b2d Mon Sep 17 00:00:00 2001 From: JoelMathewC Date: Fri, 18 Sep 2026 14:29:01 -0400 Subject: [PATCH 01/20] feat: rexport jit from finch-tensor --- pixi.toml | 2 + pyproject.toml | 2 + src/galley_jl_python/__init__.py | 3 + tests/conftest.py | 16 + .../test_jit_dependent_loop_inserted_code.py | 11 + .../test_jit_if_branch_inserted_code.py | 12 + ...jit_local_module_function_inserted_code.py | 12 + .../test_jit_module_function_inserted_code.py | 12 + .../test_jit_return_expr_inserted_code.py | 6 + .../test_jit_scalar_loop_inserted_code.py | 11 + .../test_jit_straight_line_inserted_code.py | 6 + ...t_jit_two_independent_ops_inserted_code.py | 8 + .../reference/test_jit_while_inserted_code.py | 12 + tests/test_fused.py | 782 ++++++++++++++++++ 14 files changed, 895 insertions(+) create mode 100644 tests/reference/test_jit_dependent_loop_inserted_code.py create mode 100644 tests/reference/test_jit_if_branch_inserted_code.py create mode 100644 tests/reference/test_jit_local_module_function_inserted_code.py create mode 100644 tests/reference/test_jit_module_function_inserted_code.py create mode 100644 tests/reference/test_jit_return_expr_inserted_code.py create mode 100644 tests/reference/test_jit_scalar_loop_inserted_code.py create mode 100644 tests/reference/test_jit_straight_line_inserted_code.py create mode 100644 tests/reference/test_jit_two_independent_ops_inserted_code.py create mode 100644 tests/reference/test_jit_while_inserted_code.py create mode 100644 tests/test_fused.py diff --git a/pixi.toml b/pixi.toml index fcdcdfa..b2778ab 100644 --- a/pixi.toml +++ b/pixi.toml @@ -18,6 +18,7 @@ galley-jl-python = { path = ".", editable = true } juliapkg = ">=0.1.16,<0.2" juliacall = ">=0.9.24,<0.10" numpy = ">=1.19" +finch-tensor = ">=0.7.0,<0.8.0" [feature.test.pypi-dependencies] pytest = "*" @@ -27,6 +28,7 @@ numba = ">=0.61" scipy = "*" numpy = "==2.*" pytest-xdist = ">=3.6.1,<4" +pytest-regressions = ">=2.8.1,<3.0.0" [feature.test.tasks] test = { cmd = "pytest", depends-on = ["compile"] } diff --git a/pyproject.toml b/pyproject.toml index 915143b..518edb9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -11,6 +11,7 @@ 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] @@ -27,6 +28,7 @@ test = [ "hypothesis (>=6.151.9,<7.0.0)", "ndindex (>=1.10.1,<2.0.0)", "pytest-json-report (>=1.5.0,<2.0.0)", + "pytest-regressions>=2.8.1,<3.0.0", ] [build-system] diff --git a/src/galley_jl_python/__init__.py b/src/galley_jl_python/__init__.py index 5658656..beb66e9 100644 --- a/src/galley_jl_python/__init__.py +++ b/src/galley_jl_python/__init__.py @@ -84,6 +84,8 @@ pi as pi, ) +from finch.finch_fused import jit + from . import linalg from ._array_api_info import __array_namespace_info__ from .compiled import ( @@ -292,6 +294,7 @@ "isfinite", "isinf", "isnan", + "jit", "lazy", "less", "less_equal", diff --git a/tests/conftest.py b/tests/conftest.py index 61800c9..dcd5795 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,6 +1,9 @@ +from pathlib import Path + import pytest import numpy as np +from numpy.testing import assert_allclose def pytest_addoption(parser): @@ -47,3 +50,16 @@ def arr3d(): [[0, 0, 0, 0], [1, 5, 0, 3]], ] ) + + +@pytest.fixture(scope="session") +def original_datadir() -> Path: + return Path(__file__).parent / "reference" + + +def finch_assert_allclose(result, expected, **kwargs): + if hasattr(result, "to_numpy"): + result = result.to_numpy() + if hasattr(expected, "to_numpy"): + expected = expected.to_numpy() + assert_allclose(result, expected, **kwargs) diff --git a/tests/reference/test_jit_dependent_loop_inserted_code.py b/tests/reference/test_jit_dependent_loop_inserted_code.py new file mode 100644 index 0000000..6b3c319 --- /dev/null +++ b/tests/reference/test_jit_dependent_loop_inserted_code.py @@ -0,0 +1,11 @@ +def opt_fn(A, n): + A, n = maybedefer((A, n)) + B = A + A, B = compute((A, B)) + for _i in range(getattr(sum(B), 'item')()): + A, B = maybedefer((A, B)) + B = add(B, A) + A, B = compute((A, B)) + B, = defer((B,)) + B, = compute((B,)) + return B diff --git a/tests/reference/test_jit_if_branch_inserted_code.py b/tests/reference/test_jit_if_branch_inserted_code.py new file mode 100644 index 0000000..602c6be --- /dev/null +++ b/tests/reference/test_jit_if_branch_inserted_code.py @@ -0,0 +1,12 @@ +def opt_fn(A, B, use_matmul): + if use_matmul: + A, B = maybedefer((A, B)) + result = matmul(A, B) + result, = compute((result,)) + else: + A, B = maybedefer((A, B)) + result = add(A, B) + result, = compute((result,)) + result, = defer((result,)) + result, = compute((result,)) + return result diff --git a/tests/reference/test_jit_local_module_function_inserted_code.py b/tests/reference/test_jit_local_module_function_inserted_code.py new file mode 100644 index 0000000..8093464 --- /dev/null +++ b/tests/reference/test_jit_local_module_function_inserted_code.py @@ -0,0 +1,12 @@ +def opt_fn(A, B, n): + A, B, n = maybedefer((A, B, n)) + C = A + B, C, n = compute((B, C, n)) + while n > 0: + B, C, n = maybedefer((B, C, n)) + C = getattr(finch.interface, 'add')(C, B) + n = n - 1 + B, C, n = compute((B, C, n)) + C, = defer((C,)) + C, = compute((C,)) + return C diff --git a/tests/reference/test_jit_module_function_inserted_code.py b/tests/reference/test_jit_module_function_inserted_code.py new file mode 100644 index 0000000..595633d --- /dev/null +++ b/tests/reference/test_jit_module_function_inserted_code.py @@ -0,0 +1,12 @@ +def opt_fn(A, B, n): + A, B, n = maybedefer((A, B, n)) + C = A + B, C, n = compute((B, C, n)) + while n > 0: + B, C, n = maybedefer((B, C, n)) + C = getattr(getattr(finch, 'interface'), 'add')(C, B) + n = n - 1 + B, C, n = compute((B, C, n)) + C, = defer((C,)) + C, = compute((C,)) + return C diff --git a/tests/reference/test_jit_return_expr_inserted_code.py b/tests/reference/test_jit_return_expr_inserted_code.py new file mode 100644 index 0000000..ea7d32b --- /dev/null +++ b/tests/reference/test_jit_return_expr_inserted_code.py @@ -0,0 +1,6 @@ +def opt_fn(A, B): + A, B = defer((A, B)) + ret = matmul(A, B) + ret_2 = matmul(A, B) + ret, ret_2 = compute((ret, ret_2)) + return (ret, ret_2) diff --git a/tests/reference/test_jit_scalar_loop_inserted_code.py b/tests/reference/test_jit_scalar_loop_inserted_code.py new file mode 100644 index 0000000..6633b5f --- /dev/null +++ b/tests/reference/test_jit_scalar_loop_inserted_code.py @@ -0,0 +1,11 @@ +def opt_fn(A, n): + A, n = maybedefer((A, n)) + B = A + A, B, n = compute((A, B, n)) + for _i in range(n): + A, B, n = maybedefer((A, B, n)) + B = add(B, A) + A, B, n = compute((A, B, n)) + B, = defer((B,)) + B, = compute((B,)) + return B diff --git a/tests/reference/test_jit_straight_line_inserted_code.py b/tests/reference/test_jit_straight_line_inserted_code.py new file mode 100644 index 0000000..93340be --- /dev/null +++ b/tests/reference/test_jit_straight_line_inserted_code.py @@ -0,0 +1,6 @@ +def opt_fn(A, B): + A, B = maybedefer((A, B)) + C = matmul(A, B) + C, = defer((C,)) + C, = compute((C,)) + return C diff --git a/tests/reference/test_jit_two_independent_ops_inserted_code.py b/tests/reference/test_jit_two_independent_ops_inserted_code.py new file mode 100644 index 0000000..6c70eb0 --- /dev/null +++ b/tests/reference/test_jit_two_independent_ops_inserted_code.py @@ -0,0 +1,8 @@ +def opt_fn(A, B, C): + A, B, C = maybedefer((A, B, C)) + D = matmul(A, B) + E = add(A, C) + F = add(D, E) + F, = defer((F,)) + F, = compute((F,)) + return F diff --git a/tests/reference/test_jit_while_inserted_code.py b/tests/reference/test_jit_while_inserted_code.py new file mode 100644 index 0000000..45f2006 --- /dev/null +++ b/tests/reference/test_jit_while_inserted_code.py @@ -0,0 +1,12 @@ +def opt_fn(A, B, n): + A, B, n = maybedefer((A, B, n)) + C = A + B, C, n = compute((B, C, n)) + while n > 0: + B, C, n = maybedefer((B, C, n)) + C = add(C, B) + n = n - 1 + B, C, n = compute((B, C, n)) + C, = defer((C,)) + C, = compute((C,)) + return C diff --git a/tests/test_fused.py b/tests/test_fused.py new file mode 100644 index 0000000..32b5cee --- /dev/null +++ b/tests/test_fused.py @@ -0,0 +1,782 @@ +import ast +import operator +import textwrap + +import pytest + +import numpy as np + +import finch +from finch.autoschedule import ( + DefaultLogicFormatter, + DefaultLogicOptimizer, + DefaultLoopOrderer, + LogicCapture, + LogicCompiler, + LogicExecutor, + LogicNormalizer, +) +from galley_jl_python import jit +from finch.finch_fused import nodes as fzd +from finch.finch_fused.cfg_builder import ( + fused_build_cfg, + fused_desugar, + number_statements, +) +from finch.finch_fused.dataflow import ( + LivenessAnalysis, + insert_lazy_and_compute, + maybedefer, +) +from finch.finch_fused.parser import ( + fused_function_to_python_ast, + parse_fused_function, +) +from finch.finch_notation.interpreter import NotationInterpreter +from finch.interface import add, asarray, matmul, sum +from finch.interface.lazy import LazyTensor +from finch.tensor.scalar import ConstantScalar, ScalarFType + +from .conftest import finch_assert_allclose + + +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(ConstantScalar(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(ConstantScalar(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(ConstantScalar(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. Numeric literals survive as literals: + the parser leaves them alone, and a constant only becomes a ConstantScalar + later, when `defer` hands it to the lazy layer. + """ + + 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 _transformed_jit_source(fn): + transformed_fn = insert_lazy_and_compute(parse_fused_function(fn)) + assert isinstance(transformed_fn, fzd.Function) + return ast.unparse(fused_function_to_python_ast(transformed_fn)) + "\n" + + +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 test_jit_straight_line(): + """A jit function with no loops should produce the same result as eager.""" + + def simple_fn(A, B): + C = matmul(A, B) + return C # noqa: RET504 + + @jit + def opt_fn(A, B): + C = matmul(A, B) + return C # noqa: RET504 + + A = asarray(np.array([[1, 2], [3, 4]])) + B = asarray(np.array([[5, 6], [7, 8]])) + + finch_assert_allclose(opt_fn(A, B), simple_fn(A, B)) + + +def test_jit_straight_line_inserted_code(file_regression): + def opt_fn(A, B): + C = matmul(A, B) + return C # noqa: RET504 + + file_regression.check(_transformed_jit_source(opt_fn), extension=".py") + + +def test_jit_eager_only_operation_warns_and_computes(): + def simple_fn(A): + return finch.linalg.inv(A) + + @jit + def opt_fn(A): + return finch.linalg.inv(A) + + A = asarray(np.array([[1.0, 2.0], [3.0, 5.0]])) + + with pytest.warns(RuntimeWarning, match="inv"): + result = opt_fn(A) + + finch_assert_allclose(result, simple_fn(A)) + + +def test_jit_return_expr(): + """A jit function with no loops should produce the same result as eager.""" + + def simple_fn(A, B): + return matmul(A, B), matmul(A, B) + + @jit + def opt_fn(A, B): + return matmul(A, B), matmul(A, B) + + A = asarray(np.array([[1, 2], [3, 4]])) + B = asarray(np.array([[5, 6], [7, 8]])) + + finch_assert_allclose(opt_fn(A, B), simple_fn(A, B)) + + +def test_jit_return_expr_inserted_code(file_regression): + def opt_fn(A, B): + return matmul(A, B), matmul(A, B) + + file_regression.check(_transformed_jit_source(opt_fn), extension=".py") + + +def test_jit_two_independent_ops(): + """Two independent tensor ops whose results are both used.""" + + def simple_fn(A, B, C): + D = matmul(A, B) + E = add(A, C) + F = add(D, E) + return F # noqa: RET504 + + @jit + def opt_fn(A, B, C): + D = matmul(A, B) + E = add(A, C) + F = add(D, E) + return F # noqa: RET504 + + A = asarray(np.array([[1, 2], [3, 4]])) + B = asarray(np.array([[1, 0], [0, 1]])) + C = asarray(np.array([[1, 1], [1, 1]])) + + finch_assert_allclose(opt_fn(A, B, C), simple_fn(A, B, C)) + + +def test_jit_two_independent_ops_inserted_code(file_regression): + def opt_fn(A, B, C): + D = matmul(A, B) + E = add(A, C) + F = add(D, E) + return F # noqa: RET504 + + file_regression.check(_transformed_jit_source(opt_fn), extension=".py") + + +@pytest.mark.parametrize( + ("operand", "bound"), + [ + pytest.param(ConstantScalar(2.0), 0, id="constant_scalar"), + pytest.param(1.0, 0, id="specializable_literal"), + pytest.param(2.0, 1, id="runtime_literal"), + ], +) +def test_a_constant_operand_is_not_bound_as_a_tensor(operand, bound): + """A constant costs no runtime binding, however it was written. + + Everything is bound as a table on the way in; `inline_constant_scalars` + then replaces the constants with their values and drops those bindings. A + value the optimizer cannot act on stays a binding, which is the point. + """ + capture = _CollectBindings(LogicCompiler(NotationInterpreter())) + executor = LogicExecutor( + DefaultLogicOptimizer(DefaultLoopOrderer(DefaultLogicFormatter(capture))) + ) + arr = np.arange(3.0) + with finch.with_default_scheduler(LogicNormalizer(executor)): + result = finch.compute(finch.defer(asarray(arr)) + operand) + finch_assert_allclose(result, arr + float(np.asarray(operand))) + scalars = [ + t + for bindings in capture.all_bindings + for t in bindings.values() + if isinstance(t, ScalarFType) + ] + assert len(scalars) == bound + + +def test_maybedefer_defers_every_tensor(): + """`maybedefer` makes no judgement about fills -- it only defers tensors.""" + A = asarray(np.arange(3.0)) + (lazy_A, lazy_c, plain) = maybedefer((A, ConstantScalar(2.0), 2.0)) + assert isinstance(lazy_A, LazyTensor) + assert isinstance(lazy_c, LazyTensor) + assert plain == 2.0 + + +class _CollectBindings(LogicCapture): + """A LogicCapture that keeps the bindings of every lowering, not just the last.""" + + def __init__(self, ctx): + super().__init__(ctx) + self.all_bindings: list[dict] = [] + + def lower(self, prgm, bindings, stats, stats_factory): + self.all_bindings.append(bindings.copy()) + return super().lower(prgm, bindings, stats, stats_factory) + + +def test_jit_inlines_literal_operands(): + """A specializable literal in a jit body reaches the kernel as a constant.""" + + @jit + def opt_fn(A): + return add(A, 1.0) + + capture = _CollectBindings(LogicCompiler(NotationInterpreter())) + executor = LogicExecutor( + DefaultLogicOptimizer(DefaultLoopOrderer(DefaultLogicFormatter(capture))) + ) + + arr = np.arange(3.0) + with finch.with_default_scheduler(LogicNormalizer(executor)): + result = opt_fn(asarray(arr)) + finch_assert_allclose(result, arr + 1.0) + scalars = [ + t + for bindings in capture.all_bindings + for t in bindings.values() + if isinstance(t, ScalarFType) + ] + assert not scalars + + +@pytest.mark.parametrize("trip_counts", [(1, 2, 4, 8, 16)]) +def test_jit_constant_folded_in_a_loop_does_not_grow_the_kernel_cache(trip_counts): + """A constant incremented in a loop must not cost a kernel per iteration. + + `n` folds to a fresh value on every trip. Without the demotion in + `maybedefer` each fresh value is a fresh literal, so the program -- and the + kernel compiled from it -- differs per iteration and the cache grows without + bound in the trip count. + """ + + @jit + def const_loop(A, k): + n = 1 + B = A + for _i in range(k): + n = n + 1 + B = add(B, n) + return B + + arr = np.arange(3.0) + counts = [] + for k in trip_counts: + executor = LogicExecutor( + DefaultLogicOptimizer( + DefaultLoopOrderer( + DefaultLogicFormatter(LogicCompiler(NotationInterpreter())) + ) + ), + cache=True, + ) + with finch.with_default_scheduler(LogicNormalizer(executor)): + result = const_loop(asarray(arr), k) + want = arr.copy() + n = 1 + for _ in range(k): + n += 1 + want = want + n + finch_assert_allclose(result, want) + counts.append(len(executor.cached_kernels)) + + # Flat, not linear: once the loop body has been compiled the cache stops + # growing, however many more trips are added. + assert counts[-1] == counts[-2] == counts[len(counts) // 2], counts + + +def test_jit_scalar_loop(): + """A loop with a scalar iteration count and tensor accumulation.""" + + def simple_fn(A, n): + B = A + for _i in range(n): + B = add(B, A) + return B + + @jit + def opt_fn(A, n): + B = A + for _i in range(n): + B = add(B, A) + return B + + A = asarray(np.array([[1, 0], [0, 1]], dtype=float)) + + finch_assert_allclose(opt_fn(A, 3), simple_fn(A, 3)) + + +def test_jit_scalar_loop_inserted_code(file_regression): + def opt_fn(A, n): + B = A + for _i in range(n): + B = add(B, A) + return B + + file_regression.check(_transformed_jit_source(opt_fn), extension=".py") + + +def test_jit_dependent_loop(): + """A loop with an iterator that depends on a computation.""" + + def simple_fn(A, n): + B = A + for _i in range(sum(B).item()): + B = add(B, A) + return B + + @jit + def opt_fn(A, n): + B = A + for _i in range(sum(B).item()): + B = add(B, A) + return B + + A = asarray(np.array([[1, 0], [0, 1]], dtype=int)) + + finch_assert_allclose(opt_fn(A, 3), simple_fn(A, 3)) + + +def test_jit_dependent_loop_inserted_code(file_regression): + def opt_fn(A, n): + B = A + for _i in range(sum(B).item()): + B = add(B, A) + return B + + file_regression.check(_transformed_jit_source(opt_fn), extension=".py") + + +def test_jit_if_branch(): + """A jit function with an if/else over tensor ops.""" + + def simple_fn(A, B, use_matmul): + if use_matmul: # noqa: SIM108 + result = matmul(A, B) + else: + result = add(A, B) + return result + + @jit + def opt_fn(A, B, use_matmul): + if use_matmul: # noqa: SIM108 + result = matmul(A, B) + else: + result = add(A, B) + return result + + A = asarray(np.array([[1, 2], [3, 4]])) + B = asarray(np.array([[1, 0], [0, 1]])) + + finch_assert_allclose(opt_fn(A, B, True), simple_fn(A, B, True)) + finch_assert_allclose(opt_fn(A, B, False), simple_fn(A, B, False)) + + +def test_jit_if_branch_inserted_code(file_regression): + def opt_fn(A, B, use_matmul): + if use_matmul: # noqa: SIM108 + result = matmul(A, B) + else: + result = add(A, B) + return result + + file_regression.check(_transformed_jit_source(opt_fn), extension=".py") + + +def test_jit_while(): + """A jit function with a while loop.""" + + def simple_fn(A, B, n): + C = A + while n > 0: + C = add(C, B) + n = n - 1 + return C + + @jit + def opt_fn(A, B, n): + C = A + while n > 0: + C = add(C, B) + n = n - 1 + return C + + A = asarray(np.array([[1, 2], [3, 4]])) + B = asarray(np.array([[1, 0], [0, 1]])) + n = 3 + + finch_assert_allclose(opt_fn(A, B, n), simple_fn(A, B, n)) + + +def test_jit_while_inserted_code(file_regression): + def opt_fn(A, B, n): + C = A + while n > 0: + C = add(C, B) + n = n - 1 + return C + + file_regression.check(_transformed_jit_source(opt_fn), extension=".py") + + +def test_jit_module_function(): + """A jit function with a function from a module.""" + + def simple_fn(A, B, n): + C = A + while n > 0: + C = finch.interface.add(C, B) + n = n - 1 + return C + + @jit + def opt_fn(A, B, n): + C = A + while n > 0: + C = finch.interface.add(C, B) + n = n - 1 + return C + + A = asarray(np.array([[1, 2], [3, 4]])) + B = asarray(np.array([[1, 0], [0, 1]])) + n = 3 + + finch_assert_allclose(opt_fn(A, B, n), simple_fn(A, B, n)) + + +def test_jit_module_function_inserted_code(file_regression): + def opt_fn(A, B, n): + C = A + while n > 0: + C = finch.interface.add(C, B) + n = n - 1 + return C + + file_regression.check(_transformed_jit_source(opt_fn), extension=".py") + + +def test_jit_local_module_function(): + """A jit function with a function from a module.""" + + xp = finch.interface + + def simple_fn(A, B, n): + C = A + while n > 0: + C = xp.add(C, B) + n = n - 1 + return C + + @jit + def opt_fn(A, B, n): + C = A + while n > 0: + C = xp.add(C, B) + n = n - 1 + return C + + A = asarray(np.array([[1, 2], [3, 4]])) + B = asarray(np.array([[1, 0], [0, 1]])) + n = 3 + + finch_assert_allclose(simple_fn(A, B, n), opt_fn(A, B, n)) + + +def test_jit_local_module_function_inserted_code(file_regression): + xp = finch.interface + + def opt_fn(A, B, n): + C = A + while n > 0: + C = xp.add(C, B) + n = n - 1 + return C + + file_regression.check(_transformed_jit_source(opt_fn), extension=".py") From 53c8289d4ae41998db72df7ab87ddb62368c4863 Mon Sep 17 00:00:00 2001 From: JoelMathewC Date: Fri, 18 Sep 2026 14:31:52 -0400 Subject: [PATCH 02/20] chore: fix precommit issues --- pyproject.toml | 1 + tests/test_fused.py | 3 ++- 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 518edb9..0b9dfd2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -37,6 +37,7 @@ build-backend = "poetry.core.masonry.api" [tool.ruff] line-length = 88 +extend-exclude = ["tests/reference"] [tool.ruff.lint] select = ["F", "E", "W", "I", "B", "UP", "YTT", "BLE", "C4", "T10", "ISC", "ICN", "PIE", "PYI", "RSE", "RET", "SIM", "PGH", "FLY", "NPY", "PERF", "RUF022"] diff --git a/tests/test_fused.py b/tests/test_fused.py index 32b5cee..a599327 100644 --- a/tests/test_fused.py +++ b/tests/test_fused.py @@ -16,7 +16,6 @@ LogicExecutor, LogicNormalizer, ) -from galley_jl_python import jit from finch.finch_fused import nodes as fzd from finch.finch_fused.cfg_builder import ( fused_build_cfg, @@ -37,6 +36,8 @@ from finch.interface.lazy import LazyTensor from finch.tensor.scalar import ConstantScalar, ScalarFType +from galley_jl_python import jit + from .conftest import finch_assert_allclose From 3426d241c6aa6e59c873d9e1e38eb7e060fce0eb Mon Sep 17 00:00:00 2001 From: JoelMathewC Date: Fri, 18 Sep 2026 15:17:19 -0400 Subject: [PATCH 03/20] chore: removing redundant tests --- pixi.toml | 1 - pyproject.toml | 3 - src/galley_jl_python/__init__.py | 3 +- src/galley_jl_python/finch_fused.py | 33 + tests/conftest.py | 16 - .../test_jit_dependent_loop_inserted_code.py | 11 - .../test_jit_if_branch_inserted_code.py | 12 - ...jit_local_module_function_inserted_code.py | 12 - .../test_jit_module_function_inserted_code.py | 12 - .../test_jit_return_expr_inserted_code.py | 6 - .../test_jit_scalar_loop_inserted_code.py | 11 - .../test_jit_straight_line_inserted_code.py | 6 - ...t_jit_two_independent_ops_inserted_code.py | 8 - .../reference/test_jit_while_inserted_code.py | 12 - tests/test_fused.py | 786 +----------------- 15 files changed, 73 insertions(+), 859 deletions(-) create mode 100644 src/galley_jl_python/finch_fused.py delete mode 100644 tests/reference/test_jit_dependent_loop_inserted_code.py delete mode 100644 tests/reference/test_jit_if_branch_inserted_code.py delete mode 100644 tests/reference/test_jit_local_module_function_inserted_code.py delete mode 100644 tests/reference/test_jit_module_function_inserted_code.py delete mode 100644 tests/reference/test_jit_return_expr_inserted_code.py delete mode 100644 tests/reference/test_jit_scalar_loop_inserted_code.py delete mode 100644 tests/reference/test_jit_straight_line_inserted_code.py delete mode 100644 tests/reference/test_jit_two_independent_ops_inserted_code.py delete mode 100644 tests/reference/test_jit_while_inserted_code.py diff --git a/pixi.toml b/pixi.toml index b2778ab..f19ec57 100644 --- a/pixi.toml +++ b/pixi.toml @@ -28,7 +28,6 @@ numba = ">=0.61" scipy = "*" numpy = "==2.*" pytest-xdist = ">=3.6.1,<4" -pytest-regressions = ">=2.8.1,<3.0.0" [feature.test.tasks] test = { cmd = "pytest", depends-on = ["compile"] } diff --git a/pyproject.toml b/pyproject.toml index 0b9dfd2..da0ad37 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -28,7 +28,6 @@ test = [ "hypothesis (>=6.151.9,<7.0.0)", "ndindex (>=1.10.1,<2.0.0)", "pytest-json-report (>=1.5.0,<2.0.0)", - "pytest-regressions>=2.8.1,<3.0.0", ] [build-system] @@ -37,7 +36,6 @@ build-backend = "poetry.core.masonry.api" [tool.ruff] line-length = 88 -extend-exclude = ["tests/reference"] [tool.ruff.lint] select = ["F", "E", "W", "I", "B", "UP", "YTT", "BLE", "C4", "T10", "ISC", "ICN", "PIE", "PYI", "RSE", "RET", "SIM", "PGH", "FLY", "NPY", "PERF", "RUF022"] @@ -64,4 +62,3 @@ section-order = [ [tool.mypy] ignore_missing_imports = true -exclude = ["tests/reference"] diff --git a/src/galley_jl_python/__init__.py b/src/galley_jl_python/__init__.py index beb66e9..0a19a2c 100644 --- a/src/galley_jl_python/__init__.py +++ b/src/galley_jl_python/__init__.py @@ -84,8 +84,6 @@ pi as pi, ) -from finch.finch_fused import jit - from . import linalg from ._array_api_info import __array_namespace_info__ from .compiled import ( @@ -117,6 +115,7 @@ uint32, uint64, ) +from .finch_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 new file mode 100644 index 0000000..2f64efc --- /dev/null +++ b/src/galley_jl_python/finch_fused.py @@ -0,0 +1,33 @@ +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/tests/conftest.py b/tests/conftest.py index dcd5795..61800c9 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,9 +1,6 @@ -from pathlib import Path - import pytest import numpy as np -from numpy.testing import assert_allclose def pytest_addoption(parser): @@ -50,16 +47,3 @@ def arr3d(): [[0, 0, 0, 0], [1, 5, 0, 3]], ] ) - - -@pytest.fixture(scope="session") -def original_datadir() -> Path: - return Path(__file__).parent / "reference" - - -def finch_assert_allclose(result, expected, **kwargs): - if hasattr(result, "to_numpy"): - result = result.to_numpy() - if hasattr(expected, "to_numpy"): - expected = expected.to_numpy() - assert_allclose(result, expected, **kwargs) diff --git a/tests/reference/test_jit_dependent_loop_inserted_code.py b/tests/reference/test_jit_dependent_loop_inserted_code.py deleted file mode 100644 index 6b3c319..0000000 --- a/tests/reference/test_jit_dependent_loop_inserted_code.py +++ /dev/null @@ -1,11 +0,0 @@ -def opt_fn(A, n): - A, n = maybedefer((A, n)) - B = A - A, B = compute((A, B)) - for _i in range(getattr(sum(B), 'item')()): - A, B = maybedefer((A, B)) - B = add(B, A) - A, B = compute((A, B)) - B, = defer((B,)) - B, = compute((B,)) - return B diff --git a/tests/reference/test_jit_if_branch_inserted_code.py b/tests/reference/test_jit_if_branch_inserted_code.py deleted file mode 100644 index 602c6be..0000000 --- a/tests/reference/test_jit_if_branch_inserted_code.py +++ /dev/null @@ -1,12 +0,0 @@ -def opt_fn(A, B, use_matmul): - if use_matmul: - A, B = maybedefer((A, B)) - result = matmul(A, B) - result, = compute((result,)) - else: - A, B = maybedefer((A, B)) - result = add(A, B) - result, = compute((result,)) - result, = defer((result,)) - result, = compute((result,)) - return result diff --git a/tests/reference/test_jit_local_module_function_inserted_code.py b/tests/reference/test_jit_local_module_function_inserted_code.py deleted file mode 100644 index 8093464..0000000 --- a/tests/reference/test_jit_local_module_function_inserted_code.py +++ /dev/null @@ -1,12 +0,0 @@ -def opt_fn(A, B, n): - A, B, n = maybedefer((A, B, n)) - C = A - B, C, n = compute((B, C, n)) - while n > 0: - B, C, n = maybedefer((B, C, n)) - C = getattr(finch.interface, 'add')(C, B) - n = n - 1 - B, C, n = compute((B, C, n)) - C, = defer((C,)) - C, = compute((C,)) - return C diff --git a/tests/reference/test_jit_module_function_inserted_code.py b/tests/reference/test_jit_module_function_inserted_code.py deleted file mode 100644 index 595633d..0000000 --- a/tests/reference/test_jit_module_function_inserted_code.py +++ /dev/null @@ -1,12 +0,0 @@ -def opt_fn(A, B, n): - A, B, n = maybedefer((A, B, n)) - C = A - B, C, n = compute((B, C, n)) - while n > 0: - B, C, n = maybedefer((B, C, n)) - C = getattr(getattr(finch, 'interface'), 'add')(C, B) - n = n - 1 - B, C, n = compute((B, C, n)) - C, = defer((C,)) - C, = compute((C,)) - return C diff --git a/tests/reference/test_jit_return_expr_inserted_code.py b/tests/reference/test_jit_return_expr_inserted_code.py deleted file mode 100644 index ea7d32b..0000000 --- a/tests/reference/test_jit_return_expr_inserted_code.py +++ /dev/null @@ -1,6 +0,0 @@ -def opt_fn(A, B): - A, B = defer((A, B)) - ret = matmul(A, B) - ret_2 = matmul(A, B) - ret, ret_2 = compute((ret, ret_2)) - return (ret, ret_2) diff --git a/tests/reference/test_jit_scalar_loop_inserted_code.py b/tests/reference/test_jit_scalar_loop_inserted_code.py deleted file mode 100644 index 6633b5f..0000000 --- a/tests/reference/test_jit_scalar_loop_inserted_code.py +++ /dev/null @@ -1,11 +0,0 @@ -def opt_fn(A, n): - A, n = maybedefer((A, n)) - B = A - A, B, n = compute((A, B, n)) - for _i in range(n): - A, B, n = maybedefer((A, B, n)) - B = add(B, A) - A, B, n = compute((A, B, n)) - B, = defer((B,)) - B, = compute((B,)) - return B diff --git a/tests/reference/test_jit_straight_line_inserted_code.py b/tests/reference/test_jit_straight_line_inserted_code.py deleted file mode 100644 index 93340be..0000000 --- a/tests/reference/test_jit_straight_line_inserted_code.py +++ /dev/null @@ -1,6 +0,0 @@ -def opt_fn(A, B): - A, B = maybedefer((A, B)) - C = matmul(A, B) - C, = defer((C,)) - C, = compute((C,)) - return C diff --git a/tests/reference/test_jit_two_independent_ops_inserted_code.py b/tests/reference/test_jit_two_independent_ops_inserted_code.py deleted file mode 100644 index 6c70eb0..0000000 --- a/tests/reference/test_jit_two_independent_ops_inserted_code.py +++ /dev/null @@ -1,8 +0,0 @@ -def opt_fn(A, B, C): - A, B, C = maybedefer((A, B, C)) - D = matmul(A, B) - E = add(A, C) - F = add(D, E) - F, = defer((F,)) - F, = compute((F,)) - return F diff --git a/tests/reference/test_jit_while_inserted_code.py b/tests/reference/test_jit_while_inserted_code.py deleted file mode 100644 index 45f2006..0000000 --- a/tests/reference/test_jit_while_inserted_code.py +++ /dev/null @@ -1,12 +0,0 @@ -def opt_fn(A, B, n): - A, B, n = maybedefer((A, B, n)) - C = A - B, C, n = compute((B, C, n)) - while n > 0: - B, C, n = maybedefer((B, C, n)) - C = add(C, B) - n = n - 1 - B, C, n = compute((B, C, n)) - C, = defer((C,)) - C, = compute((C,)) - return C diff --git a/tests/test_fused.py b/tests/test_fused.py index a599327..b95c5ff 100644 --- a/tests/test_fused.py +++ b/tests/test_fused.py @@ -1,783 +1,75 @@ -import ast -import operator -import textwrap - -import pytest - import numpy as np +from numpy.testing import assert_allclose -import finch -from finch.autoschedule import ( - DefaultLogicFormatter, - DefaultLogicOptimizer, - DefaultLoopOrderer, - LogicCapture, - LogicCompiler, - LogicExecutor, - LogicNormalizer, -) -from finch.finch_fused import nodes as fzd -from finch.finch_fused.cfg_builder import ( - fused_build_cfg, - fused_desugar, - number_statements, -) -from finch.finch_fused.dataflow import ( - LivenessAnalysis, - insert_lazy_and_compute, - maybedefer, -) -from finch.finch_fused.parser import ( - fused_function_to_python_ast, - parse_fused_function, -) -from finch.finch_notation.interpreter import NotationInterpreter -from finch.interface import add, asarray, matmul, sum -from finch.interface.lazy import LazyTensor -from finch.tensor.scalar import ConstantScalar, ScalarFType - -from galley_jl_python import jit - -from .conftest import finch_assert_allclose - - -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(ConstantScalar(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(ConstantScalar(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(ConstantScalar(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. Numeric literals survive as literals: - the parser leaves them alone, and a constant only becomes a ConstantScalar - later, when `defer` hands it to the lazy layer. - """ - - 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 _transformed_jit_source(fn): - transformed_fn = insert_lazy_and_compute(parse_fused_function(fn)) - assert isinstance(transformed_fn, fzd.Function) - return ast.unparse(fused_function_to_python_ast(transformed_fn)) + "\n" - - -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 +import galley_jl_python as galley -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 test_jit_straight_line(): - """A jit function with no loops should produce the same result as eager.""" - - def simple_fn(A, B): - C = matmul(A, B) - return C # noqa: RET504 - - @jit - def opt_fn(A, B): - C = matmul(A, B) - return C # noqa: RET504 - - A = asarray(np.array([[1, 2], [3, 4]])) - B = asarray(np.array([[5, 6], [7, 8]])) - - finch_assert_allclose(opt_fn(A, B), simple_fn(A, B)) - - -def test_jit_straight_line_inserted_code(file_regression): - def opt_fn(A, B): - C = matmul(A, B) - return C # noqa: RET504 - - file_regression.check(_transformed_jit_source(opt_fn), extension=".py") - - -def test_jit_eager_only_operation_warns_and_computes(): - def simple_fn(A): - return finch.linalg.inv(A) - - @jit - def opt_fn(A): - return finch.linalg.inv(A) - - A = asarray(np.array([[1.0, 2.0], [3.0, 5.0]])) - - with pytest.warns(RuntimeWarning, match="inv"): - result = opt_fn(A) - - finch_assert_allclose(result, simple_fn(A)) +def assert_galley_allclose(actual, expected): + assert isinstance(actual, galley.Tensor) + assert_allclose(actual.todense(), expected) -def test_jit_return_expr(): - """A jit function with no loops should produce the same result as eager.""" - - def simple_fn(A, B): - return matmul(A, B), matmul(A, B) - - @jit +def test_jit_straight_line_uses_galley_tensors(): + @galley.jit def opt_fn(A, B): - return matmul(A, B), matmul(A, B) + return galley.matmul(A, B) - A = asarray(np.array([[1, 2], [3, 4]])) - B = asarray(np.array([[5, 6], [7, 8]])) + A = galley.Tensor(np.array([[1, 2], [3, 4]])) + B = galley.Tensor(np.array([[5, 6], [7, 8]])) - finch_assert_allclose(opt_fn(A, B), simple_fn(A, B)) + assert_galley_allclose(opt_fn(A, B), np.matmul(A.todense(), B.todense())) -def test_jit_return_expr_inserted_code(file_regression): +def test_jit_return_expr_uses_galley_tensors(): + @galley.jit def opt_fn(A, B): - return matmul(A, B), matmul(A, B) - - file_regression.check(_transformed_jit_source(opt_fn), extension=".py") - - -def test_jit_two_independent_ops(): - """Two independent tensor ops whose results are both used.""" - - def simple_fn(A, B, C): - D = matmul(A, B) - E = add(A, C) - F = add(D, E) - return F # noqa: RET504 - - @jit - def opt_fn(A, B, C): - D = matmul(A, B) - E = add(A, C) - F = add(D, E) - return F # noqa: RET504 - - A = asarray(np.array([[1, 2], [3, 4]])) - B = asarray(np.array([[1, 0], [0, 1]])) - C = asarray(np.array([[1, 1], [1, 1]])) - - finch_assert_allclose(opt_fn(A, B, C), simple_fn(A, B, C)) - - -def test_jit_two_independent_ops_inserted_code(file_regression): - def opt_fn(A, B, C): - D = matmul(A, B) - E = add(A, C) - F = add(D, E) - return F # noqa: RET504 - - file_regression.check(_transformed_jit_source(opt_fn), extension=".py") - - -@pytest.mark.parametrize( - ("operand", "bound"), - [ - pytest.param(ConstantScalar(2.0), 0, id="constant_scalar"), - pytest.param(1.0, 0, id="specializable_literal"), - pytest.param(2.0, 1, id="runtime_literal"), - ], -) -def test_a_constant_operand_is_not_bound_as_a_tensor(operand, bound): - """A constant costs no runtime binding, however it was written. - - Everything is bound as a table on the way in; `inline_constant_scalars` - then replaces the constants with their values and drops those bindings. A - value the optimizer cannot act on stays a binding, which is the point. - """ - capture = _CollectBindings(LogicCompiler(NotationInterpreter())) - executor = LogicExecutor( - DefaultLogicOptimizer(DefaultLoopOrderer(DefaultLogicFormatter(capture))) - ) - arr = np.arange(3.0) - with finch.with_default_scheduler(LogicNormalizer(executor)): - result = finch.compute(finch.defer(asarray(arr)) + operand) - finch_assert_allclose(result, arr + float(np.asarray(operand))) - scalars = [ - t - for bindings in capture.all_bindings - for t in bindings.values() - if isinstance(t, ScalarFType) - ] - assert len(scalars) == bound - - -def test_maybedefer_defers_every_tensor(): - """`maybedefer` makes no judgement about fills -- it only defers tensors.""" - A = asarray(np.arange(3.0)) - (lazy_A, lazy_c, plain) = maybedefer((A, ConstantScalar(2.0), 2.0)) - assert isinstance(lazy_A, LazyTensor) - assert isinstance(lazy_c, LazyTensor) - assert plain == 2.0 - - -class _CollectBindings(LogicCapture): - """A LogicCapture that keeps the bindings of every lowering, not just the last.""" - - def __init__(self, ctx): - super().__init__(ctx) - self.all_bindings: list[dict] = [] - - def lower(self, prgm, bindings, stats, stats_factory): - self.all_bindings.append(bindings.copy()) - return super().lower(prgm, bindings, stats, stats_factory) - - -def test_jit_inlines_literal_operands(): - """A specializable literal in a jit body reaches the kernel as a constant.""" - - @jit - def opt_fn(A): - return add(A, 1.0) - - capture = _CollectBindings(LogicCompiler(NotationInterpreter())) - executor = LogicExecutor( - DefaultLogicOptimizer(DefaultLoopOrderer(DefaultLogicFormatter(capture))) - ) - - arr = np.arange(3.0) - with finch.with_default_scheduler(LogicNormalizer(executor)): - result = opt_fn(asarray(arr)) - finch_assert_allclose(result, arr + 1.0) - scalars = [ - t - for bindings in capture.all_bindings - for t in bindings.values() - if isinstance(t, ScalarFType) - ] - assert not scalars + return galley.matmul(A, B), galley.add(A, B) + A = galley.Tensor(np.array([[1, 2], [3, 4]])) + B = galley.Tensor(np.array([[5, 6], [7, 8]])) -@pytest.mark.parametrize("trip_counts", [(1, 2, 4, 8, 16)]) -def test_jit_constant_folded_in_a_loop_does_not_grow_the_kernel_cache(trip_counts): - """A constant incremented in a loop must not cost a kernel per iteration. + actual_matmul, actual_add = opt_fn(A, B) + assert_galley_allclose(actual_matmul, np.matmul(A.todense(), B.todense())) + assert_galley_allclose(actual_add, A.todense() + B.todense()) - `n` folds to a fresh value on every trip. Without the demotion in - `maybedefer` each fresh value is a fresh literal, so the program -- and the - kernel compiled from it -- differs per iteration and the cache grows without - bound in the trip count. - """ - @jit - def const_loop(A, k): - n = 1 - B = A - for _i in range(k): - n = n + 1 - B = add(B, n) - return B - - arr = np.arange(3.0) - counts = [] - for k in trip_counts: - executor = LogicExecutor( - DefaultLogicOptimizer( - DefaultLoopOrderer( - DefaultLogicFormatter(LogicCompiler(NotationInterpreter())) - ) - ), - cache=True, - ) - with finch.with_default_scheduler(LogicNormalizer(executor)): - result = const_loop(asarray(arr), k) - want = arr.copy() - n = 1 - for _ in range(k): - n += 1 - want = want + n - finch_assert_allclose(result, want) - counts.append(len(executor.cached_kernels)) - - # Flat, not linear: once the loop body has been compiled the cache stops - # growing, however many more trips are added. - assert counts[-1] == counts[-2] == counts[len(counts) // 2], counts - - -def test_jit_scalar_loop(): - """A loop with a scalar iteration count and tensor accumulation.""" - - def simple_fn(A, n): - B = A - for _i in range(n): - B = add(B, A) - return B - - @jit - def opt_fn(A, n): - B = A - for _i in range(n): - B = add(B, A) - return B - - A = asarray(np.array([[1, 0], [0, 1]], dtype=float)) - - finch_assert_allclose(opt_fn(A, 3), simple_fn(A, 3)) - - -def test_jit_scalar_loop_inserted_code(file_regression): +def test_jit_scalar_loop_uses_galley_tensors(): + @galley.jit def opt_fn(A, n): B = A for _i in range(n): - B = add(B, A) - return B - - file_regression.check(_transformed_jit_source(opt_fn), extension=".py") - - -def test_jit_dependent_loop(): - """A loop with an iterator that depends on a computation.""" - - def simple_fn(A, n): - B = A - for _i in range(sum(B).item()): - B = add(B, A) - return B - - @jit - def opt_fn(A, n): - B = A - for _i in range(sum(B).item()): - B = add(B, A) + B = galley.add(B, A) return B - A = asarray(np.array([[1, 0], [0, 1]], dtype=int)) - - finch_assert_allclose(opt_fn(A, 3), simple_fn(A, 3)) - - -def test_jit_dependent_loop_inserted_code(file_regression): - def opt_fn(A, n): - B = A - for _i in range(sum(B).item()): - B = add(B, A) - return B - - file_regression.check(_transformed_jit_source(opt_fn), extension=".py") - - -def test_jit_if_branch(): - """A jit function with an if/else over tensor ops.""" - - def simple_fn(A, B, use_matmul): - if use_matmul: # noqa: SIM108 - result = matmul(A, B) - else: - result = add(A, B) - return result + A = galley.Tensor(np.array([[1, 0], [0, 1]], dtype=float)) - @jit - def opt_fn(A, B, use_matmul): - if use_matmul: # noqa: SIM108 - result = matmul(A, B) - else: - result = add(A, B) - return result - - A = asarray(np.array([[1, 2], [3, 4]])) - B = asarray(np.array([[1, 0], [0, 1]])) + assert_galley_allclose(opt_fn(A, 3), 4 * A.todense()) - finch_assert_allclose(opt_fn(A, B, True), simple_fn(A, B, True)) - finch_assert_allclose(opt_fn(A, B, False), simple_fn(A, B, False)) - -def test_jit_if_branch_inserted_code(file_regression): +def test_jit_if_branch_uses_galley_tensors(): + @galley.jit def opt_fn(A, B, use_matmul): - if use_matmul: # noqa: SIM108 - result = matmul(A, B) - else: - result = add(A, B) - return result + if use_matmul: + return galley.matmul(A, B) + return galley.add(A, B) - file_regression.check(_transformed_jit_source(opt_fn), extension=".py") + A = galley.Tensor(np.array([[1, 2], [3, 4]])) + B = galley.Tensor(np.array([[1, 0], [0, 1]])) + assert_galley_allclose(opt_fn(A, B, True), np.matmul(A.todense(), B.todense())) + assert_galley_allclose(opt_fn(A, B, False), A.todense() + B.todense()) -def test_jit_while(): - """A jit function with a while loop.""" - def simple_fn(A, B, n): - C = A - while n > 0: - C = add(C, B) - n = n - 1 - return C - - @jit +def test_jit_while_loop_uses_galley_tensors(): + @galley.jit def opt_fn(A, B, n): C = A while n > 0: - C = add(C, B) + C = galley.add(C, B) n = n - 1 return C - A = asarray(np.array([[1, 2], [3, 4]])) - B = asarray(np.array([[1, 0], [0, 1]])) - n = 3 - - finch_assert_allclose(opt_fn(A, B, n), simple_fn(A, B, n)) - - -def test_jit_while_inserted_code(file_regression): - def opt_fn(A, B, n): - C = A - while n > 0: - C = add(C, B) - n = n - 1 - return C - - file_regression.check(_transformed_jit_source(opt_fn), extension=".py") - - -def test_jit_module_function(): - """A jit function with a function from a module.""" - - def simple_fn(A, B, n): - C = A - while n > 0: - C = finch.interface.add(C, B) - n = n - 1 - return C - - @jit - def opt_fn(A, B, n): - C = A - while n > 0: - C = finch.interface.add(C, B) - n = n - 1 - return C - - A = asarray(np.array([[1, 2], [3, 4]])) - B = asarray(np.array([[1, 0], [0, 1]])) - n = 3 - - finch_assert_allclose(opt_fn(A, B, n), simple_fn(A, B, n)) - - -def test_jit_module_function_inserted_code(file_regression): - def opt_fn(A, B, n): - C = A - while n > 0: - C = finch.interface.add(C, B) - n = n - 1 - return C - - file_regression.check(_transformed_jit_source(opt_fn), extension=".py") - - -def test_jit_local_module_function(): - """A jit function with a function from a module.""" - - xp = finch.interface - - def simple_fn(A, B, n): - C = A - while n > 0: - C = xp.add(C, B) - n = n - 1 - return C - - @jit - def opt_fn(A, B, n): - C = A - while n > 0: - C = xp.add(C, B) - n = n - 1 - return C - - A = asarray(np.array([[1, 2], [3, 4]])) - B = asarray(np.array([[1, 0], [0, 1]])) - n = 3 - - finch_assert_allclose(simple_fn(A, B, n), opt_fn(A, B, n)) - - -def test_jit_local_module_function_inserted_code(file_regression): - xp = finch.interface - - def opt_fn(A, B, n): - C = A - while n > 0: - C = xp.add(C, B) - n = n - 1 - return C + A = galley.Tensor(np.array([[1, 2], [3, 4]])) + B = galley.Tensor(np.array([[1, 0], [0, 1]])) - file_regression.check(_transformed_jit_source(opt_fn), extension=".py") + assert_galley_allclose(opt_fn(A, B, 3), A.todense() + 3 * B.todense()) From e51281722025af545b09cdc833e6c10c25b12268 Mon Sep 17 00:00:00 2001 From: JoelMathewC Date: Fri, 18 Sep 2026 15:20:14 -0400 Subject: [PATCH 04/20] chore: rename galley->gl --- tests/test_fused.py | 44 ++++++++++++++++++++++---------------------- 1 file changed, 22 insertions(+), 22 deletions(-) diff --git a/tests/test_fused.py b/tests/test_fused.py index b95c5ff..ac6a456 100644 --- a/tests/test_fused.py +++ b/tests/test_fused.py @@ -1,32 +1,32 @@ import numpy as np from numpy.testing import assert_allclose -import galley_jl_python as galley +import galley_jl_python as gl def assert_galley_allclose(actual, expected): - assert isinstance(actual, galley.Tensor) + assert isinstance(actual, gl.Tensor) assert_allclose(actual.todense(), expected) def test_jit_straight_line_uses_galley_tensors(): - @galley.jit + @gl.jit def opt_fn(A, B): - return galley.matmul(A, B) + return gl.matmul(A, B) - A = galley.Tensor(np.array([[1, 2], [3, 4]])) - B = galley.Tensor(np.array([[5, 6], [7, 8]])) + A = gl.Tensor(np.array([[1, 2], [3, 4]])) + B = gl.Tensor(np.array([[5, 6], [7, 8]])) assert_galley_allclose(opt_fn(A, B), np.matmul(A.todense(), B.todense())) def test_jit_return_expr_uses_galley_tensors(): - @galley.jit + @gl.jit def opt_fn(A, B): - return galley.matmul(A, B), galley.add(A, B) + return gl.matmul(A, B), gl.add(A, B) - A = galley.Tensor(np.array([[1, 2], [3, 4]])) - B = galley.Tensor(np.array([[5, 6], [7, 8]])) + A = gl.Tensor(np.array([[1, 2], [3, 4]])) + B = gl.Tensor(np.array([[5, 6], [7, 8]])) actual_matmul, actual_add = opt_fn(A, B) assert_galley_allclose(actual_matmul, np.matmul(A.todense(), B.todense())) @@ -34,42 +34,42 @@ def opt_fn(A, B): def test_jit_scalar_loop_uses_galley_tensors(): - @galley.jit + @gl.jit def opt_fn(A, n): B = A for _i in range(n): - B = galley.add(B, A) + B = gl.add(B, A) return B - A = galley.Tensor(np.array([[1, 0], [0, 1]], dtype=float)) + A = gl.Tensor(np.array([[1, 0], [0, 1]], dtype=float)) assert_galley_allclose(opt_fn(A, 3), 4 * A.todense()) def test_jit_if_branch_uses_galley_tensors(): - @galley.jit + @gl.jit def opt_fn(A, B, use_matmul): if use_matmul: - return galley.matmul(A, B) - return galley.add(A, B) + return gl.matmul(A, B) + return gl.add(A, B) - A = galley.Tensor(np.array([[1, 2], [3, 4]])) - B = galley.Tensor(np.array([[1, 0], [0, 1]])) + A = gl.Tensor(np.array([[1, 2], [3, 4]])) + B = gl.Tensor(np.array([[1, 0], [0, 1]])) assert_galley_allclose(opt_fn(A, B, True), np.matmul(A.todense(), B.todense())) assert_galley_allclose(opt_fn(A, B, False), A.todense() + B.todense()) def test_jit_while_loop_uses_galley_tensors(): - @galley.jit + @gl.jit def opt_fn(A, B, n): C = A while n > 0: - C = galley.add(C, B) + C = gl.add(C, B) n = n - 1 return C - A = galley.Tensor(np.array([[1, 2], [3, 4]])) - B = galley.Tensor(np.array([[1, 0], [0, 1]])) + A = gl.Tensor(np.array([[1, 2], [3, 4]])) + B = gl.Tensor(np.array([[1, 0], [0, 1]])) assert_galley_allclose(opt_fn(A, B, 3), A.todense() + 3 * B.todense()) From 090f896937b99b7c5cfe17731388ea3812516230 Mon Sep 17 00:00:00 2001 From: kylebd99 Date: Fri, 25 Sep 2026 17:51:49 -0400 Subject: [PATCH 05/20] merge to pyproject.toml --- pixi.toml | 36 ------------------------------------ pyproject.toml | 25 +++++++++++++++++++++++++ 2 files changed, 25 insertions(+), 36 deletions(-) delete mode 100644 pixi.toml diff --git a/pixi.toml b/pixi.toml deleted file mode 100644 index f19ec57..0000000 --- a/pixi.toml +++ /dev/null @@ -1,36 +0,0 @@ -[project] -authors = ["Hameer Abbasi <2190658+hameerabbasi@users.noreply.github.com>"] -channels = ["conda-forge"] -description = "Add a short description here" -name = "galley-jl-python" -platforms = ["osx-arm64"] -version = "0.2.12" - -[tasks] -compile = "python -c 'import galley_jl_python'" - -[dependencies] -python = ">=3.10,<3.13" -juliaup = ">=1.17.10,<2" - -[pypi-dependencies] -galley-jl-python = { path = ".", editable = true } -juliapkg = ">=0.1.16,<0.2" -juliacall = ">=0.9.24,<0.10" -numpy = ">=1.19" -finch-tensor = ">=0.7.0,<0.8.0" - -[feature.test.pypi-dependencies] -pytest = "*" -pytest-cov = "*" -sparse = ">=0.16,<0.17" -numba = ">=0.61" -scipy = "*" -numpy = "==2.*" -pytest-xdist = ">=3.6.1,<4" - -[feature.test.tasks] -test = { cmd = "pytest", depends-on = ["compile"] } - -[environments] -test = ["test"] diff --git a/pyproject.toml b/pyproject.toml index da0ad37..bab48c0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -28,12 +28,37 @@ test = [ "hypothesis (>=6.151.9,<7.0.0)", "ndindex (>=1.10.1,<2.0.0)", "pytest-json-report (>=1.5.0,<2.0.0)", + "pytest-xdist>=3.6.1,<4", ] [build-system] requires = ["poetry-core>=1.0.8"] build-backend = "poetry.core.masonry.api" +[tool.pixi.workspace] +channels = ["conda-forge"] +platforms = [ + { name = "linux-amd64", platform = "linux-64", glibc = "2.34" }, + { name = "linux-arm64", platform = "linux-aarch64", glibc = "2.34" }, + { name = "macos-arm64", platform = "osx-arm64", macos = "15.0" }, + { name = "win-amd64", platform = "win-64" }, +] + +[tool.pixi.dependencies] +python = ">=3.11,<3.13" + +[tool.pixi.pypi-dependencies] +galley-jl-python = { path = ".", editable = true } + +[tool.pixi.tasks] +compile = "python -c 'import galley_jl_python'" + +[tool.pixi.feature.test.tasks] +test = { cmd = "pytest", depends-on = ["compile"] } + +[tool.pixi.environments] +test = { features = ["test"], solve-group = "default" } + [tool.ruff] line-length = 88 From e3759c8645f3e50a83d4275fc1d7434eb4416783 Mon Sep 17 00:00:00 2001 From: kylebd99 Date: Fri, 25 Sep 2026 19:04:58 -0400 Subject: [PATCH 06/20] update --- array-api-skips.txt | 1 + src/galley_jl_python/compiled.py | 3 +++ src/galley_jl_python/tensor.py | 8 +++++++- 3 files changed, 11 insertions(+), 1 deletion(-) diff --git a/array-api-skips.txt b/array-api-skips.txt index 797379d..ade08fe 100644 --- a/array-api-skips.txt +++ b/array-api-skips.txt @@ -249,6 +249,7 @@ array_api_tests/test_indexing_functions.py::test_take # test_linalg # not implemented +array_api_tests/test_linalg.py::test_cross array_api_tests/test_linalg.py::test_matrix_transpose array_api_tests/test_linalg.py::test_vecdot array_api_tests/test_linalg.py::test_eigh diff --git a/src/galley_jl_python/compiled.py b/src/galley_jl_python/compiled.py index af05973..83907ad 100644 --- a/src/galley_jl_python/compiled.py +++ b/src/galley_jl_python/compiled.py @@ -141,6 +141,9 @@ def set_optimizer(opt: AbstractScheduler) -> None: jl.Finch.set_scheduler_b(opt.get_julia_scheduler()) +set_optimizer(GalleyScheduler()) + + def lazy(tensor: Tensor) -> Tensor: from .tensor import Tensor diff --git a/src/galley_jl_python/tensor.py b/src/galley_jl_python/tensor.py index 06195c7..8f98335 100644 --- a/src/galley_jl_python/tensor.py +++ b/src/galley_jl_python/tensor.py @@ -25,6 +25,11 @@ ) from .typing import Device, DType, JuliaObj, OrderType, TupleOf3Arrays, spmatrix +# Builds a 0-d Finch tensor in Julia so the element keeps its Julia eltype. +_scalar_to_0d = jl.seval( + "(T, fv, x) -> Finch.Tensor(Finch.Element{convert(T, fv)}(T[convert(T, x)]))" +) + class SparseArray: """ @@ -279,7 +284,8 @@ def __getitem__(self, key): return Tensor(result) if jl.isa(result, jl.Finch.Tensor): return Tensor(jl.swizzle(result, *range(1, jl.ndims(result) + 1))) - return result + # the Array API requires a 0-d array, not a scalar, for integer-only keys + return Tensor(jl.swizzle(_scalar_to_0d(self.dtype, self.fill_value, result))) @property def dtype(self) -> DType: From c8a05441516b8091a127b5a7a3017949c8a7f18c Mon Sep 17 00:00:00 2001 From: kylebd99 Date: Sat, 26 Sep 2026 11:14:05 -0400 Subject: [PATCH 07/20] Add sysimage workflow --- .github/workflows/sysimage.yml | 53 ++++ .gitignore | 3 + README.md | 30 +++ pyproject.toml | 5 +- scripts/sysimage/_load.py | 15 ++ scripts/sysimage/build.py | 79 ++++++ scripts/sysimage/build_sysimage.jl | 26 ++ scripts/sysimage/fetch.py | 58 +++++ scripts/sysimage/pin_julia_deps.py | 54 +++++ scripts/sysimage/workload.py | 62 +++++ src/galley_jl_python/_sysimage.py | 79 ++++++ src/galley_jl_python/julia.py | 8 +- src/galley_jl_python/juliapkg.json | 374 ++++++++++++++++++++++++++++- 13 files changed, 839 insertions(+), 7 deletions(-) create mode 100644 .github/workflows/sysimage.yml create mode 100644 scripts/sysimage/_load.py create mode 100644 scripts/sysimage/build.py create mode 100644 scripts/sysimage/build_sysimage.jl create mode 100644 scripts/sysimage/fetch.py create mode 100644 scripts/sysimage/pin_julia_deps.py create mode 100644 scripts/sysimage/workload.py create mode 100644 src/galley_jl_python/_sysimage.py diff --git a/.github/workflows/sysimage.yml b/.github/workflows/sysimage.yml new file mode 100644 index 0000000..fa246cb --- /dev/null +++ b/.github/workflows/sysimage.yml @@ -0,0 +1,53 @@ +name: Sysimage + +# Builds the Julia sysimage for each platform and publishes it to the GitHub +# release tagged `sysimage-`, where identifies the exact Julia +# environment. Users download it with `pixi run fetch-sysimage`. +on: + workflow_dispatch: + push: + branches: + - main + paths: + - src/galley_jl_python/juliapkg.json + - scripts/sysimage/** + - pyproject.toml + - pixi.lock + - .github/workflows/sysimage.yml + +permissions: + contents: write + +jobs: + build: + strategy: + matrix: + os: [ubuntu-latest, ubuntu-24.04-arm, macos-latest, windows-latest] + fail-fast: false + runs-on: ${{ matrix.os }} + # GitHub-hosted maximum; the generic multi-CPU build is slow + timeout-minutes: 360 + defaults: + run: + shell: bash + steps: + - name: Checkout Repo + uses: actions/checkout@v7 + - uses: prefix-dev/setup-pixi@v0.10.0 + with: + pixi-version: v0.72.2 + - name: Build sysimage + id: build + run: pixi run build-sysimage + - name: Publish to release + env: + GH_TOKEN: ${{ github.token }} + TAG: ${{ steps.build.outputs.tag }} + IMAGE: ${{ steps.build.outputs.path }} + run: | + # the matrix jobs race to create the release; losing the race is fine + gh release view "$TAG" > /dev/null 2>&1 \ + || gh release create "$TAG" --prerelease --title "$TAG" \ + --notes "Prebuilt Julia sysimages. Fetch with \`pixi run fetch-sysimage\`." \ + || gh release view "$TAG" > /dev/null + gh release upload "$TAG" "$IMAGE" --clobber diff --git a/.gitignore b/.gitignore index 2b2b89f..52ec6f2 100644 --- a/.gitignore +++ b/.gitignore @@ -168,3 +168,6 @@ cython_debug/ .pixi/ *.egg-info pixi.lock + +# PackageCompiler environment used by scripts/sysimage/build_sysimage.jl +scripts/sysimage/.build-env/ diff --git a/README.md b/README.md index 9f01527..bd13d3b 100644 --- a/README.md +++ b/README.md @@ -38,6 +38,36 @@ Options: --path Path to the local copy of Finch.jl [default: ../Finch.jl]. ``` +### Julia sysimage + +Most of Galley's startup time is Julia compiling Finch itself. A prebuilt Julia +sysimage removes it: the first operations of a session drop from minutes to a +few seconds. + +- `pixi run fetch-sysimage` downloads the image for your platform into + `~/.cache/galley-jl-python/` (`GALLEY_JL_PYTHON_CACHE` overrides this). + `pixi run compile` and `pixi run test` run this step first. +- `pixi run build-sysimage` builds the image locally instead. This takes hours. +- `import galley_jl_python` loads a cached image automatically when it matches + the current Julia environment. Otherwise Julia starts without it. Set + `GALLEY_JL_PYTHON_SYSIMAGE=0` to turn this off. + +An image works only with the exact Julia and package versions it was built +from. For this reason `src/galley_jl_python/juliapkg.json` pins Julia and every +Julia package. Each image's name includes a hash of that environment. A local +Finch.jl from `develop.py` therefore runs without the image. + +To update the Julia dependencies: + +1. Loosen the pins you want to change. +2. Resolve with `pixi run compile`. +3. Re-pin with `python scripts/sysimage/pin_julia_deps.py`. + +Pushing the new pins to `main` runs the "Sysimage" GitHub Action. It builds +images for Linux, macOS and Windows and publishes them to a `sysimage-` +GitHub release, where `fetch-sysimage` finds them. The action can also be run +manually from the Actions tab. + ### Publishing The "Publish" GitHub Action is a manual workflow for publishing Python packages to PyPI using Poetry. It handles the version management based on the `pyproject.toml` file and automates tagging and creating GitHub releases. diff --git a/pyproject.toml b/pyproject.toml index bab48c0..1fd8224 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -51,7 +51,10 @@ python = ">=3.11,<3.13" galley-jl-python = { path = ".", editable = true } [tool.pixi.tasks] -compile = "python -c 'import galley_jl_python'" +# pixi has no post-install hook, so the entry-point tasks fetch the sysimage +compile = { cmd = "python -c 'import galley_jl_python'", depends-on = ["fetch-sysimage"] } +build-sysimage = "python scripts/sysimage/build.py" +fetch-sysimage = "python scripts/sysimage/fetch.py" [tool.pixi.feature.test.tasks] test = { cmd = "pytest", depends-on = ["compile"] } diff --git a/scripts/sysimage/_load.py b/scripts/sysimage/_load.py new file mode 100644 index 0000000..cd04e5b --- /dev/null +++ b/scripts/sysimage/_load.py @@ -0,0 +1,15 @@ +"""Load ``galley_jl_python._sysimage`` by path. + +Importing it through the package would run ``galley_jl_python/__init__.py``, +which starts Julia. +""" + +import importlib.util +from pathlib import Path + +_PATH = ( + Path(__file__).resolve().parents[2] / "src" / "galley_jl_python" / "_sysimage.py" +) +_spec = importlib.util.spec_from_file_location("_galley_sysimage", _PATH) +sysimage = importlib.util.module_from_spec(_spec) +_spec.loader.exec_module(sysimage) diff --git a/scripts/sysimage/build.py b/scripts/sysimage/build.py new file mode 100644 index 0000000..950a39c --- /dev/null +++ b/scripts/sysimage/build.py @@ -0,0 +1,79 @@ +"""Build a Julia sysimage that removes most of galley's first-call compile time. + +Traces the Julia methods compiled by ``workload.py``, compiles them into a +sysimage with PackageCompiler, and stores it in the per-user cache under a name +tied to the current Julia environment, where ``galley_jl_python`` finds it. +""" + +import argparse +import os +import shutil +import subprocess +import sys +import tempfile +from pathlib import Path + +import juliapkg + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) + +from _load import sysimage # noqa: E402 + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + parser.add_argument( + "--statements", + type=Path, + help="reuse precompile statements from an earlier trace instead of tracing", + ) + args = parser.parse_args() + + project = juliapkg.project() + key = sysimage.env_key(project) + image = sysimage.image_path(key) + image.parent.mkdir(parents=True, exist_ok=True) + + with tempfile.TemporaryDirectory() as tmp: + statements = args.statements or Path(tmp) / "precompile_statements.jl" + if args.statements is None: + print(f"[sysimage] tracing workload -> {statements}", flush=True) + env = { + k: v for k, v in os.environ.items() if k != "PYTHON_JULIACALL_SYSIMAGE" + } + subprocess.run( + [sys.executable, str(HERE / "workload.py")], + env={ + **env, + "GALLEY_JL_PYTHON_SYSIMAGE": "0", + "PYTHON_JULIACALL_TRACE_COMPILE": str(statements), + }, + check=True, + ) + + # build next to the target and rename, so an interrupted build never + # leaves a partial image where it would be loaded + partial = Path(tmp) / image.name + print(f"[sysimage] building {image.name} (this takes a while)", flush=True) + subprocess.run( + [ + juliapkg.executable(), + "--startup-file=no", + str(HERE / "build_sysimage.jl"), + project, + str(statements), + str(partial), + ], + check=True, + ) + shutil.move(partial, image) + + print(f"[sysimage] wrote {image}", flush=True) + if github_output := os.environ.get("GITHUB_OUTPUT"): + with open(github_output, "a") as f: + f.write(f"key={key}\npath={image}\ntag={sysimage.release_tag(key)}\n") + + +if __name__ == "__main__": + main() diff --git a/scripts/sysimage/build_sysimage.jl b/scripts/sysimage/build_sysimage.jl new file mode 100644 index 0000000..fcd8b59 --- /dev/null +++ b/scripts/sysimage/build_sysimage.jl @@ -0,0 +1,26 @@ +# Usage: julia build_sysimage.jl +# +# PackageCompiler lives in its own environment so it never enters the +# project that juliapkg manages. +using Pkg + +project, statements_file, output = ARGS + +build_env = joinpath(@__DIR__, ".build-env") +Pkg.activate(build_env; io=devnull) +if Base.find_package("PackageCompiler") === nothing + Pkg.add("PackageCompiler") +end +using PackageCompiler + +packages = ["Finch", "HDF5", "NPZ", "TensorMarket", "PythonCall"] + +create_sysimage( + packages; + project=project, + sysimage_path=output, + precompile_statements_file=statements_file, + # a generic target so the image runs on any x86_64 CPU, not just this one + cpu_target=PackageCompiler.default_app_cpu_target(), + include_transitive_dependencies=true, +) diff --git a/scripts/sysimage/fetch.py b/scripts/sysimage/fetch.py new file mode 100644 index 0000000..725569e --- /dev/null +++ b/scripts/sysimage/fetch.py @@ -0,0 +1,58 @@ +"""Download the prebuilt sysimage for this platform and Julia environment. + +CI publishes one image per platform to the GitHub release tagged +``sysimage-``, where the key identifies the exact Julia environment. +""" + +import os +import sys +import tempfile +import urllib.error +import urllib.request +from pathlib import Path + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) + +from _load import sysimage # noqa: E402 + + +def main() -> None: + key = sysimage.current_key() + image = sysimage.image_path(key) + if image.is_file(): + print(f"[sysimage] already present: {image}") + return + + url = f"{sysimage.RELEASES_URL}/{sysimage.release_tag(key)}/{image.name}" + print(f"[sysimage] downloading {url}", flush=True) + image.parent.mkdir(parents=True, exist_ok=True) + fd, partial = tempfile.mkstemp(dir=image.parent, suffix=".part") + try: + with os.fdopen(fd, "wb") as out, urllib.request.urlopen(url) as resp: + while chunk := resp.read(1 << 20): + out.write(chunk) + os.replace(partial, image) + # never fail: other tasks depend on this one, and galley works without an + # image, just with a slower first call + except urllib.error.HTTPError as e: + if e.code != 404: + print(f"[sysimage] download failed ({e}); continuing without it.") + return + print( + f"[sysimage] no prebuilt image for this platform and Julia environment " + f"({image.name}); continuing without it. Build one locally with " + "`pixi run build-sysimage`." + ) + return + except (urllib.error.URLError, OSError) as e: + print(f"[sysimage] download failed ({e}); continuing without it.") + return + finally: + if os.path.exists(partial): + os.remove(partial) + print(f"[sysimage] saved {image}") + + +if __name__ == "__main__": + main() diff --git a/scripts/sysimage/pin_julia_deps.py b/scripts/sysimage/pin_julia_deps.py new file mode 100644 index 0000000..94726b7 --- /dev/null +++ b/scripts/sysimage/pin_julia_deps.py @@ -0,0 +1,54 @@ +"""Pin galley's Julia dependencies to the currently resolved environment. + +Rewrites ``src/galley_jl_python/juliapkg.json`` so that Julia and every +registered package in the resolved ``Manifest.toml`` (transitive dependencies +included) are pinned to exact versions. Everyone then resolves the same +manifest, which the prebuilt sysimage requires. + +To update: loosen the pins you want to move, resolve (e.g. ``pixi run +compile``), run this script, then rebuild with ``pixi run build-sysimage``. +""" + +import argparse +import json +import sys +import tomllib +from pathlib import Path + +ROOT = Path(__file__).resolve().parents[2] +JULIAPKG_JSON = ROOT / "src" / "galley_jl_python" / "juliapkg.json" +# juliacall pins these itself (PythonCall to its own release, OpenSSL_jll to +# match Python's OpenSSL), so they must stay unpinned here +UNPINNED = {"OpenSSL_jll", "PythonCall"} + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + parser.add_argument( + "--manifest", + type=Path, + default=Path(sys.prefix) / "julia_env" / "Manifest.toml", + ) + args = parser.parse_args() + + manifest = tomllib.loads(args.manifest.read_text()) + config = json.loads(JULIAPKG_JSON.read_text()) + config["julia"] = f"={manifest['julia_version']}" + + packages = {} + for name, (entry, *_) in sorted(manifest["deps"].items()): + # stdlibs have no tree hash and ship with Julia itself + if "git-tree-sha1" not in entry or name in UNPINNED: + continue + # drop build metadata (e.g. JLL "+0"): juliapkg can't parse it and Pkg + # compat ignores it anyway + version = entry["version"].split("+")[0] + packages[name] = {"uuid": entry["uuid"], "version": f"={version}"} + config["packages"] = packages + + JULIAPKG_JSON.write_text(json.dumps(config, indent=2) + "\n") + print(f"pinned julia {config['julia']} and {len(packages)} packages") + + +if __name__ == "__main__": + main() diff --git a/scripts/sysimage/workload.py b/scripts/sysimage/workload.py new file mode 100644 index 0000000..030167d --- /dev/null +++ b/scripts/sysimage/workload.py @@ -0,0 +1,62 @@ +"""Representative workload traced to build the Julia sysimage. + +Every Julia method compiled while this runs is recorded (via juliacall's +``trace_compile`` option) and baked into the sysimage by ``build.py``. Add +operations here to cover more of the library ahead of time. +""" + +import numpy as np + +import galley_jl_python as gl + + +def run_ops(A, B, C, v): + gl.matmul(A, B) + gl.sum(gl.matmul(gl.matmul(A, B), C)) + gl.add(A, B) + gl.multiply(A, B) + gl.subtract(A, B) + gl.sum(A, axis=0) + gl.sum(A, axis=1) + gl.max(A) + gl.tensordot(A, v, axes=1) + gl.permute_dims(A, (1, 0)) + A[0] + A[0, 0] + A[1:3, :] + A.todense() + + @gl.jit + def sum_abc(A, B, C): + return gl.sum(gl.matmul(gl.matmul(A, B), C)) + + @gl.jit + def chain(A, B, C): + return gl.add(gl.matmul(A, B), C) + + sum_abc(A, B, C) + chain(A, B, C) + + +def main(): + rng = np.random.default_rng(0) + for scheduler in (gl.GalleyScheduler(), gl.DefaultScheduler()): + gl.set_optimizer(scheduler) + for n in (50, 400): + A, B, C = ( + gl.random((n, n), density=0.02, random_state=s) for s in (1, 2, 3) + ) + v = gl.asarray(rng.random(n)) + run_ops(A, B, C, v) + for fmt in ("coo", "csr", "csc", "dense"): + for dtype in (np.float64, np.int64): + dense = (rng.random((60, 60)) < 0.05).astype(dtype) + X = gl.asarray(dense, format=fmt) + gl.matmul(X, X) + gl.sum(X) + gl.add(X, X) + gl.set_optimizer(gl.GalleyScheduler()) + + +if __name__ == "__main__": + main() diff --git a/src/galley_jl_python/_sysimage.py b/src/galley_jl_python/_sysimage.py new file mode 100644 index 0000000..a96ffbe --- /dev/null +++ b/src/galley_jl_python/_sysimage.py @@ -0,0 +1,79 @@ +"""Locate a prebuilt Julia sysimage and point juliacall at it. + +This module must not import juliacall: it runs before Julia starts. + +A sysimage only works with the exact Julia version and package versions it was +built from, so each image is named after a hash of both (see ``env_key``) and +lives in a per-user cache directory. An image is used only when its name +matches the current environment; otherwise Julia starts normally. + +Get an image with ``pixi run fetch-sysimage`` (downloads the one CI built for +this environment) or ``pixi run build-sysimage`` (builds it locally). Set +``GALLEY_JL_PYTHON_SYSIMAGE=0`` to never load one, or +``PYTHON_JULIACALL_SYSIMAGE`` to use a specific image. +""" + +import hashlib +import json +import os +import platform +import sys +import tomllib +from pathlib import Path + +RELEASES_URL = "https://github.com/finch-tensor/galley-jl-python/releases/download" +_EXTENSIONS = {"linux": "so", "darwin": "dylib", "win32": "dll"} + + +def cache_dir() -> Path: + if custom := os.environ.get("GALLEY_JL_PYTHON_CACHE"): + return Path(custom) + if sys.platform == "win32": + base = Path(os.environ.get("LOCALAPPDATA", Path.home() / "AppData" / "Local")) + elif sys.platform == "darwin": + base = Path.home() / "Library" / "Caches" + else: + base = Path(os.environ.get("XDG_CACHE_HOME", Path.home() / ".cache")) + return base / "galley-jl-python" + + +def env_key(project: str) -> str: + """Hash the Julia version and resolved package set of a juliapkg project.""" + manifest = tomllib.loads((Path(project) / "Manifest.toml").read_text()) + resolved = sorted( + (name, entry.get("uuid"), entry.get("version"), entry.get("git-tree-sha1")) + for name, entries in manifest["deps"].items() + for entry in entries + ) + payload = json.dumps([manifest["julia_version"], resolved]) + return hashlib.sha256(payload.encode()).hexdigest()[:16] + + +def release_tag(key: str) -> str: + return f"sysimage-{key}" + + +def image_name(key: str) -> str: + ext = _EXTENSIONS.get(sys.platform, "so") + return f"galley-{sys.platform}-{platform.machine().lower()}-{key}.{ext}" + + +def image_path(key: str) -> Path: + return cache_dir() / image_name(key) + + +def current_key() -> str: + import juliapkg + + return env_key(juliapkg.project()) + + +def configure() -> None: + """Point juliacall at the cached sysimage for this environment, if any.""" + if os.environ.get("GALLEY_JL_PYTHON_SYSIMAGE") == "0": + return + if "PYTHON_JULIACALL_SYSIMAGE" in os.environ: + return + image = image_path(current_key()) + if image.is_file(): + os.environ["PYTHON_JULIACALL_SYSIMAGE"] = str(image) diff --git a/src/galley_jl_python/julia.py b/src/galley_jl_python/julia.py index c8f2191..b865900 100644 --- a/src/galley_jl_python/julia.py +++ b/src/galley_jl_python/julia.py @@ -1,10 +1,16 @@ import juliapkg # noqa: I001, F401 +from . import _sysimage + +_sysimage.configure() + # To change the version of Finch used, see the documentation for pyjuliapkg here: https://github.com/JuliaPy/pyjuliapkg # Use pyjuliapkg to modify the `juliapkg.json` file in the root of this repo. # You can also run `develop.py` to quickly use a local copy of Finch.jl. # An example development json is found in `juliapkg_dev.json` -import juliacall as jc # noqa: F401 +# Julia and all packages are pinned so the prebuilt sysimage matches; update the +# pins with `scripts/sysimage/pin_julia_deps.py`. +import juliacall as jc # noqa: E402, F401 from juliacall import Main as jl # noqa: E402, F401 jl.seval("using Finch") diff --git a/src/galley_jl_python/juliapkg.json b/src/galley_jl_python/juliapkg.json index ad999ee..77dd492 100644 --- a/src/galley_jl_python/juliapkg.json +++ b/src/galley_jl_python/juliapkg.json @@ -1,21 +1,385 @@ { - "julia": "1.10 - 1.12", + "julia": "=1.12.2", "packages": { + "AbstractTrees": { + "uuid": "1520ce14-60c1-5f80-bbc7-55ef81b5835c", + "version": "=0.4.5" + }, + "AcceleratedKernels": { + "uuid": "6a4ca0a5-0e36-4168-a932-d9be78d558f1", + "version": "=0.4.3" + }, + "Accessors": { + "uuid": "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697", + "version": "=0.1.45" + }, + "Adapt": { + "uuid": "79e6a3ab-5dfb-504d-930d-738a2a938a0e", + "version": "=4.7.1" + }, + "AliasTables": { + "uuid": "66dad0bd-aa9a-41b7-9441-69ab47430ed8", + "version": "=1.1.3" + }, + "ArgCheck": { + "uuid": "dce04be8-c92d-5529-be00-80e4d2c0e197", + "version": "=2.5.0" + }, + "Atomix": { + "uuid": "a9b6321e-bd34-4604-b9c9-b65b8de01458", + "version": "=1.2.1" + }, + "AutoHashEquals": { + "uuid": "15f4f7f2-30c1-5605-9d31-71845cf9641f", + "version": "=2.2.0" + }, + "CommonSolve": { + "uuid": "38540f10-b2f7-11e9-35d8-d573e4eb0ff2", + "version": "=0.2.14" + }, + "Compat": { + "uuid": "34da2185-b29b-5c13-b0c7-acf172513d20", + "version": "=4.18.1" + }, + "CompositionsBase": { + "uuid": "a33af91c-f02d-484b-be07-31d278c5ca2b", + "version": "=0.1.2" + }, + "CondaPkg": { + "uuid": "992eb4ea-22a4-4c89-a5bb-47a3300528ab", + "version": "=0.2.36" + }, + "ConstructionBase": { + "uuid": "187b0558-2788-49d3-abe0-74a17ed4e7c9", + "version": "=1.6.0" + }, + "DataAPI": { + "uuid": "9a962f9c-6df0-11e9-0e5d-c546b8b5ee8a", + "version": "=1.16.0" + }, + "DataStructures": { + "uuid": "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8", + "version": "=0.19.6" + }, + "DataValueInterfaces": { + "uuid": "e2d170a0-9d28-54be-80f0-106bbe20a464", + "version": "=1.0.0" + }, + "Distributions": { + "uuid": "31c24e10-a181-5473-b8eb-7969acd0382f", + "version": "=0.25.131" + }, + "DocStringExtensions": { + "uuid": "ffbed154-4ef7-542d-bbb7-c09d3a79fcae", + "version": "=0.9.5" + }, + "FileIO": { + "uuid": "5789e2e9-d7fb-5bc7-8068-2c6fae9b9549", + "version": "=1.20.0" + }, + "FillArrays": { + "uuid": "1a297f60-69ca-5386-bcde-b61e274b549b", + "version": "=1.17.1" + }, "Finch": { "uuid": "9177782c-1635-4eb9-9bfb-d9dfa25e6bce", - "version": "1.2.10" + "version": "=1.6.0" + }, + "GPUArraysCore": { + "uuid": "46192b85-c4d5-4398-a991-12ede77f4527", + "version": "=0.2.1" + }, + "Gamma": { + "uuid": "a0844989-3bd2-4988-8bea-c9407ab0941b", + "version": "=1.2.0" }, "HDF5": { "uuid": "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f", - "version": "0.17.2" + "version": "=0.17.4" + }, + "HDF5_jll": { + "uuid": "0234f1f7-429e-5d53-9886-15a909be8d59", + "version": "=2.2.2" + }, + "Hwloc_jll": { + "uuid": "e33a78d0-f292-5ffc-b300-72abe9b543c8", + "version": "=2.14.1" + }, + "HypergeometricFunctions": { + "uuid": "34004b35-14d8-5ef3-9330-4cdb6864b03a", + "version": "=0.3.30" + }, + "InverseFunctions": { + "uuid": "3587e190-3f89-42d0-90ee-14403ec27112", + "version": "=0.1.17" + }, + "IrrationalConstants": { + "uuid": "92d709cd-6900-40b7-9082-c6be49f344b6", + "version": "=0.2.6" + }, + "IterTools": { + "uuid": "c8e1da08-722c-5040-9ed9-7db0dc04731e", + "version": "=1.10.0" + }, + "IteratorInterfaceExtensions": { + "uuid": "82899510-4779-5014-852e-03e436cf321d", + "version": "=1.0.0" + }, + "JLLWrappers": { + "uuid": "692b3bcd-3c85-4b1f-b108-f13ce0eb3210", + "version": "=1.8.0" + }, + "JSON": { + "uuid": "682c06a0-de6a-54ab-a142-c8b1cf79cde6", + "version": "=1.9.0" + }, + "KernelAbstractions": { + "uuid": "63c18a36-062a-441e-b654-da1e3ab1ce7c", + "version": "=0.9.42" + }, + "Libiconv_jll": { + "uuid": "94ce4f54-9a6c-5748-9c1c-f9c7231a4531", + "version": "=1.18.0" + }, + "LogExpFunctions": { + "uuid": "2ab3a3ac-af41-5b50-aa03-7779005ae688", + "version": "=1.0.2" + }, + "MPIABI_jll": { + "uuid": "b5ada748-db0f-5fc0-8972-9331c762740c", + "version": "=1.0.0" + }, + "MPICH_jll": { + "uuid": "7cb0a576-ebde-5e09-9194-50597f1243b4", + "version": "=5.0.1" + }, + "MPIPreferences": { + "uuid": "3da0fdf6-3ccc-4f1b-acd9-58baa6c99267", + "version": "=0.1.12" + }, + "MPItrampoline_jll": { + "uuid": "f1f71cc9-e9ae-5b93-9b94-4fe0e1ad3748", + "version": "=5.5.6" + }, + "MacroTools": { + "uuid": "1914dd2f-81c6-5fcd-8719-6d5c9610ff09", + "version": "=0.5.16" + }, + "MicroMamba": { + "uuid": "0b3b1443-0f03-428d-bdfb-f27f9c1191ea", + "version": "=0.1.15" + }, + "MicrosoftMPI_jll": { + "uuid": "9237b28f-5490-5468-be7b-bb81f5f5e6cf", + "version": "=10.1.4" + }, + "Missings": { + "uuid": "e1d29d7a-bbdc-5cf2-9ac0-f12de2c33e28", + "version": "=1.2.0" }, "NPZ": { "uuid": "15e1cf62-19b3-5cfa-8e77-841668bca605", - "version": "0.4.3" + "version": "=0.4.3" + }, + "OpenMPI_jll": { + "uuid": "fe0851c0-eecd-5654-98d4-656369965a5c", + "version": "=5.0.12" + }, + "OpenSpecFun_jll": { + "uuid": "efe28fd5-8261-553b-a9e1-b2916fc3738e", + "version": "=0.5.6" + }, + "OrderedCollections": { + "uuid": "bac558e1-5e72-5ebc-8fee-abe8a469f55d", + "version": "=2.0.1" + }, + "PDMats": { + "uuid": "90014a1f-27ba-587c-ab20-58faa44d9150", + "version": "=0.11.41" + }, + "Parsers": { + "uuid": "69de0a69-1ddd-5017-9359-2bf0b02dc9f0", + "version": "=3.0.0" + }, + "Pidfile": { + "uuid": "fa939f87-e72e-5be4-a000-7fc836dbe307", + "version": "=1.3.0" + }, + "PrecompileTools": { + "uuid": "aea7be01-6a6a-4083-8856-8a6e6704d82a", + "version": "=1.3.4" + }, + "Preferences": { + "uuid": "21216c6a-2e73-6563-6e65-726566657250", + "version": "=1.6.0" + }, + "PtrArrays": { + "uuid": "43287f4e-b6f4-7ad1-bb20-aadabca52c3d", + "version": "=1.4.0" + }, + "QuadGK": { + "uuid": "1fd47b50-473d-5c70-9696-f719f8f3bcdc", + "version": "=2.11.3" + }, + "Reexport": { + "uuid": "189a3867-3050-52da-a836-e630ba90ab69", + "version": "=1.2.2" + }, + "Requires": { + "uuid": "ae029012-a4dd-5104-9daa-d747884805df", + "version": "=1.3.1" + }, + "RewriteTools": { + "uuid": "5969e224-3634-4c61-9f66-721b69e98b8a", + "version": "=0.2.4" + }, + "Rmath": { + "uuid": "79098fc4-a85e-5d69-aa6a-4863f24498fa", + "version": "=0.9.0" + }, + "Rmath_jll": { + "uuid": "f50d1b31-88e8-58de-be2c-1cc44531875f", + "version": "=0.5.2" + }, + "Roots": { + "uuid": "f2b01f46-fcfa-551c-844a-d8ac1e96c665", + "version": "=3.0.8" + }, + "Scratch": { + "uuid": "6c6a2e73-6563-6170-7368-637461726353", + "version": "=1.3.0" + }, + "SortingAlgorithms": { + "uuid": "a2af1166-a08f-5f64-846c-94a0d3cef48c", + "version": "=1.2.3" + }, + "SpecialFunctions": { + "uuid": "276daf66-3868-5448-9aa4-cd146d93841b", + "version": "=2.9.0" + }, + "StaticArrays": { + "uuid": "90137ffa-7385-5640-81b9-e52037218182", + "version": "=1.9.22" + }, + "StaticArraysCore": { + "uuid": "1e83bf80-4336-4d27-bf5d-d5a4f845583c", + "version": "=1.4.4" + }, + "Statistics": { + "uuid": "10745b16-79ce-11e8-11f9-7d13ad32a3b2", + "version": "=1.11.5" + }, + "StatsAPI": { + "uuid": "82ae8749-77ed-4fe6-ae5f-f523153014b0", + "version": "=1.8.0" + }, + "StatsBase": { + "uuid": "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91", + "version": "=0.34.13" + }, + "StatsFuns": { + "uuid": "4c63d2b9-4356-54db-8cca-17b64c39e42c", + "version": "=2.2.1" + }, + "StructUtils": { + "uuid": "ec057cc2-7a8d-4b58-b3b3-92acb9f63b42", + "version": "=2.9.2" + }, + "SyntaxInterface": { + "uuid": "b33eeca9-aacb-4496-a840-e75f1646a4fb", + "version": "=0.2.1" + }, + "TableTraits": { + "uuid": "3783bdb8-4a98-5b6b-af9a-565f29a5fe9c", + "version": "=1.0.1" + }, + "Tables": { + "uuid": "bd369af6-aec1-5ad0-b16a-f7cc5008161c", + "version": "=1.14.0" }, "TensorMarket": { "uuid": "8b7d4fe7-0b45-4d0d-9dd8-5cc9b23b4b77", - "version": "0.2.0" + "version": "=0.2.2" + }, + "UnsafeAtomics": { + "uuid": "013be700-e6cd-48c3-b4a1-df204f14c38f", + "version": "=0.3.2" + }, + "UnsafePointers": { + "uuid": "e17b2a0c-0bdf-430a-bd0c-3a23cae4ff39", + "version": "=1.0.0" + }, + "XML2_jll": { + "uuid": "02c8fc9c-b97f-50b9-bbe4-9be30ff0a78a", + "version": "=2.15.4" + }, + "Xorg_libpciaccess_jll": { + "uuid": "a65dc6b1-eb27-53a1-bb3e-dea574b5389e", + "version": "=0.19.0" + }, + "ZipFile": { + "uuid": "a5390f91-8eb1-5f08-bee0-b1d1ffed6cea", + "version": "=0.10.1" + }, + "aws_c_auth_jll": { + "uuid": "2b3700d1-4306-52e2-a478-c162f0c514be", + "version": "=0.9.6" + }, + "aws_c_cal_jll": { + "uuid": "70f11efc-bab2-57f1-b0f3-22aad4e67c4b", + "version": "=0.9.13" + }, + "aws_c_common_jll": { + "uuid": "73048d1d-b8c4-5092-a58d-866c5e8d1e50", + "version": "=0.12.6" + }, + "aws_c_compression_jll": { + "uuid": "73a04cd5-f3d7-5bac-9290-e8adb709f224", + "version": "=0.3.2" + }, + "aws_c_http_jll": { + "uuid": "3254fc65-9028-534d-aa9d-d76d128babc6", + "version": "=0.10.15" + }, + "aws_c_io_jll": { + "uuid": "13c41daa-f319-5298-b5eb-5754e0170d52", + "version": "=0.26.3" + }, + "aws_c_s3_jll": { + "uuid": "bd1f34fb-993f-5903-a121-aaf302eed6d4", + "version": "=0.11.5" + }, + "aws_c_sdkutils_jll": { + "uuid": "1282aa60-004d-510b-9f52-12498d409daa", + "version": "=0.2.4" + }, + "aws_checksums_jll": { + "uuid": "b2a88e68-78e7-5e94-8c20-c02986ec140e", + "version": "=0.2.10" + }, + "dlfcn_win32_jll": { + "uuid": "c4b69c83-5512-53e3-94e6-de98773c479f", + "version": "=1.4.2" + }, + "libaec_jll": { + "uuid": "477f73a3-ac25-53e9-8cc3-50b2fa2566f0", + "version": "=1.1.7" + }, + "micromamba_jll": { + "uuid": "f8abcde7-e9b7-5caa-b8af-a437887ae8e4", + "version": "=2.3.1" + }, + "mpif_jll": { + "uuid": "9aeb927a-4695-514f-a259-621a69f20ec0", + "version": "=1.0.0" + }, + "pixi_jll": { + "uuid": "4d7b5844-a134-5dcd-ac86-c8f19cd51bed", + "version": "=0.76.2" + }, + "s2n_tls_jll": { + "uuid": "cddc5d3d-934d-5d3a-9747-62fc12ea3f48", + "version": "=1.7.10" } } } From d659c44e47ad285525d6253dbe239cebcb9abf5d Mon Sep 17 00:00:00 2001 From: kylebd99 Date: Sat, 26 Sep 2026 11:18:30 -0400 Subject: [PATCH 08/20] update for pixi --- README.md | 45 +++++++++++++++++++++++++++++++++++---------- 1 file changed, 35 insertions(+), 10 deletions(-) diff --git a/README.md b/README.md index bd13d3b..34f9444 100644 --- a/README.md +++ b/README.md @@ -16,18 +16,35 @@ pip install galley-jl-python ## Contributing -### Packaging +### Development environment -Galley uses [poetry](https://python-poetry.org/) for packaging. +Galley uses [pixi](https://pixi.sh) for development. All configuration lives in +`pyproject.toml`: the package metadata and dependencies are under `[project]`, +and pixi's settings are under `[tool.pixi]`. -To install for development, clone the repository and run: +To set up, clone the repository and run: ```bash -poetry install --with test +pixi install # the default environment +pixi install -e test # adds the test dependencies (the `test` extra) ``` -to install the current project and dev dependencies. + +pixi installs the package in editable mode. Julia is not a pixi dependency: +[juliapkg](https://github.com/JuliaPy/pyjuliapkg) installs the pinned Julia +version and packages the first time `galley_jl_python` is imported. To trigger +that, and to fetch the sysimage (see below), run: +```bash +pixi run compile +``` + +Run any other command inside an environment with `pixi run`, for example +`pixi run -e test python`. + +The package is still built and published with Poetry (see +[Publishing](#publishing)). ### Working with a local copy of Finch.jl -The `develop.py ` script can be used to set up a local copy of Finch.jl for development. +The `develop.py` script can be used to set up a local copy of Finch.jl for +development. Run it with `pixi run python develop.py`. ``` Usage: @@ -93,16 +110,24 @@ On successful execution, the action publishes the package to PyPI and tags the r To add pre-commit hooks, run: ```bash -poetry run pre-commit install +pixi run -e test pre-commit install ``` ### Testing -Finch uses [pytest](https://docs.pytest.org/en/latest/) for testing. To run the +Galley uses [pytest](https://docs.pytest.org/en/latest/) for testing. To run the tests: ```bash -poetry run pytest +pixi run test +``` + +This runs `compile` first and then the full suite, including the Array API +tests described below. To run a subset, call pytest in the test environment +directly: + +```bash +pixi run -e test pytest tests/test_fused.py ``` Array API tests are included in `tests/test_array_api.py`. These tests invoke @@ -112,7 +137,7 @@ To forward `pytest` options to the nested `--array-api-pytest-args`): ```bash -poetry run pytest tests/test_array_api.py \ +pixi run -e test pytest tests/test_array_api.py \ --array-api="-k creation_functions" \ --array-api="-x" ``` From 026d51f589df15a08229159655a8d03386f8f71e Mon Sep 17 00:00:00 2001 From: kylebd99 Date: Sat, 26 Sep 2026 12:02:31 -0400 Subject: [PATCH 09/20] small fixes --- .github/workflows/ci.yml | 32 +++++++++----------------- pyproject.toml | 14 +++++++++++- src/galley_jl_python/finch_patches.jl | 33 +++++++++++++++++++++++++++ src/galley_jl_python/julia.py | 2 ++ 4 files changed, 59 insertions(+), 22 deletions(-) create mode 100644 src/galley_jl_python/finch_patches.jl diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 7f1e00f..fd59ab2 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -3,40 +3,30 @@ jobs: test: defaults: run: - shell: bash -el {0} + shell: bash strategy: matrix: os: [ubuntu-latest] - python: ['3.11', '3.12', '3.13'] + environment: [test-py311, test-py312, test-py313] include: - os: macos-latest - python: '3.11' + environment: test-py311 - os: windows-latest - python: '3.11' + environment: test-py311 fail-fast: false runs-on: ${{ matrix.os }} steps: - name: Checkout Repo - uses: actions/checkout@v3 - - uses: actions/setup-python@v5 + uses: actions/checkout@v7 + - uses: prefix-dev/setup-pixi@v0.10.0 with: - python-version: ${{ matrix.python }} - - name: Install Poetry - uses: snok/install-poetry@v1 - - name: Configure Poetry - run: | - if [[ "${{ runner.os }}" == "Windows" ]]; then - poetry env use python - else - poetry env use python${{ matrix.python }} - fi - - name: Install package - run: | - poetry install --extras test + pixi-version: v0.72.2 + environments: ${{ matrix.environment }} - name: Run tests run: | - poetry run pytest --junit-xml=test-${{ matrix.os }}-Python-${{ matrix.python }}.xml - - uses: codecov/codecov-action@v3 + pixi run -e ${{ matrix.environment }} test \ + --junit-xml=test-${{ matrix.os }}-${{ matrix.environment }}.xml + - uses: codecov/codecov-action@v7 on: # Trigger the workflow on push or pull request, diff --git a/pyproject.toml b/pyproject.toml index 1fd8224..d1e9eb5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -45,7 +45,7 @@ platforms = [ ] [tool.pixi.dependencies] -python = ">=3.11,<3.13" +python = ">=3.11,<3.14" [tool.pixi.pypi-dependencies] galley-jl-python = { path = ".", editable = true } @@ -59,8 +59,20 @@ fetch-sysimage = "python scripts/sysimage/fetch.py" [tool.pixi.feature.test.tasks] test = { cmd = "pytest", depends-on = ["compile"] } +[tool.pixi.feature.py311.dependencies] +python = "3.11.*" + +[tool.pixi.feature.py312.dependencies] +python = "3.12.*" + +[tool.pixi.feature.py313.dependencies] +python = "3.13.*" + [tool.pixi.environments] test = { features = ["test"], solve-group = "default" } +test-py311 = ["test", "py311"] +test-py312 = ["test", "py312"] +test-py313 = ["test", "py313"] [tool.ruff] line-length = 88 diff --git a/src/galley_jl_python/finch_patches.jl b/src/galley_jl_python/finch_patches.jl new file mode 100644 index 0000000..62986d7 --- /dev/null +++ b/src/galley_jl_python/finch_patches.jl @@ -0,0 +1,33 @@ +# Fixes for bugs in the pinned Finch.jl, evaluated into Finch.Galley at import. +# Delete each fix once the pinned Finch includes it. + +# Finch <= 1.6.0 builds `OrderedDict{DCKey}()`, which has no constructor, so +# Galley throws a MethodError when it re-checks a cached plan for tensors with +# 50 or more degree constraints (e.g. einsums over 4-d tensors). +# Upstream fix: https://github.com/finch-tensor/Finch.jl/pull/822 +function issimilar(stat1::DCStats, stat2::DCStats, rel_granularity) + if length(stat1.dcs) < 50 + for dc1 in stat1.dcs + for dc2 in stat2.dcs + if dc1.X == dc2.X && dc1.Y == dc2.Y && + abs(log(rel_granularity, dc1.d) - log(rel_granularity, dc2.d)) > 1 + return false + end + end + end + else + dc_dict = OrderedDict{DCKey,Float64}() + for dc1 in stat1.dcs + dc_dict[get_dc_key(dc1)] = dc1.d + end + for dc2 in stat2.dcs + # like the loop above, only compare DCs present in both stats + d1 = get(dc_dict, get_dc_key(dc2), nothing) + if !isnothing(d1) && + abs(log(rel_granularity, d1) - log(rel_granularity, dc2.d)) > 1 + return false + end + end + end + return true +end diff --git a/src/galley_jl_python/julia.py b/src/galley_jl_python/julia.py index b865900..e22b879 100644 --- a/src/galley_jl_python/julia.py +++ b/src/galley_jl_python/julia.py @@ -1,4 +1,5 @@ import juliapkg # noqa: I001, F401 +from pathlib import Path from . import _sysimage @@ -14,6 +15,7 @@ from juliacall import Main as jl # noqa: E402, F401 jl.seval("using Finch") +jl.Base.include(jl.Finch.Galley, str(Path(__file__).with_name("finch_patches.jl"))) jl.seval("using HDF5") jl.seval("using NPZ") jl.seval("using TensorMarket") From 47da44bc30b63bb6b23cf0e13ce0f73ab4493031 Mon Sep 17 00:00:00 2001 From: JoelMathewC Date: Fri, 18 Sep 2026 14:29:01 -0400 Subject: [PATCH 10/20] feat: rexport jit from finch-tensor --- pixi.toml | 2 + pyproject.toml | 2 + src/galley_jl_python/__init__.py | 3 + tests/conftest.py | 16 + .../test_jit_dependent_loop_inserted_code.py | 11 + .../test_jit_if_branch_inserted_code.py | 12 + ...jit_local_module_function_inserted_code.py | 12 + .../test_jit_module_function_inserted_code.py | 12 + .../test_jit_return_expr_inserted_code.py | 6 + .../test_jit_scalar_loop_inserted_code.py | 11 + .../test_jit_straight_line_inserted_code.py | 6 + ...t_jit_two_independent_ops_inserted_code.py | 8 + .../reference/test_jit_while_inserted_code.py | 12 + tests/test_fused.py | 782 ++++++++++++++++++ 14 files changed, 895 insertions(+) create mode 100644 tests/reference/test_jit_dependent_loop_inserted_code.py create mode 100644 tests/reference/test_jit_if_branch_inserted_code.py create mode 100644 tests/reference/test_jit_local_module_function_inserted_code.py create mode 100644 tests/reference/test_jit_module_function_inserted_code.py create mode 100644 tests/reference/test_jit_return_expr_inserted_code.py create mode 100644 tests/reference/test_jit_scalar_loop_inserted_code.py create mode 100644 tests/reference/test_jit_straight_line_inserted_code.py create mode 100644 tests/reference/test_jit_two_independent_ops_inserted_code.py create mode 100644 tests/reference/test_jit_while_inserted_code.py create mode 100644 tests/test_fused.py diff --git a/pixi.toml b/pixi.toml index fcdcdfa..b2778ab 100644 --- a/pixi.toml +++ b/pixi.toml @@ -18,6 +18,7 @@ galley-jl-python = { path = ".", editable = true } juliapkg = ">=0.1.16,<0.2" juliacall = ">=0.9.24,<0.10" numpy = ">=1.19" +finch-tensor = ">=0.7.0,<0.8.0" [feature.test.pypi-dependencies] pytest = "*" @@ -27,6 +28,7 @@ numba = ">=0.61" scipy = "*" numpy = "==2.*" pytest-xdist = ">=3.6.1,<4" +pytest-regressions = ">=2.8.1,<3.0.0" [feature.test.tasks] test = { cmd = "pytest", depends-on = ["compile"] } diff --git a/pyproject.toml b/pyproject.toml index 915143b..518edb9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -11,6 +11,7 @@ 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] @@ -27,6 +28,7 @@ test = [ "hypothesis (>=6.151.9,<7.0.0)", "ndindex (>=1.10.1,<2.0.0)", "pytest-json-report (>=1.5.0,<2.0.0)", + "pytest-regressions>=2.8.1,<3.0.0", ] [build-system] diff --git a/src/galley_jl_python/__init__.py b/src/galley_jl_python/__init__.py index 5658656..beb66e9 100644 --- a/src/galley_jl_python/__init__.py +++ b/src/galley_jl_python/__init__.py @@ -84,6 +84,8 @@ pi as pi, ) +from finch.finch_fused import jit + from . import linalg from ._array_api_info import __array_namespace_info__ from .compiled import ( @@ -292,6 +294,7 @@ "isfinite", "isinf", "isnan", + "jit", "lazy", "less", "less_equal", diff --git a/tests/conftest.py b/tests/conftest.py index 61800c9..dcd5795 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,6 +1,9 @@ +from pathlib import Path + import pytest import numpy as np +from numpy.testing import assert_allclose def pytest_addoption(parser): @@ -47,3 +50,16 @@ def arr3d(): [[0, 0, 0, 0], [1, 5, 0, 3]], ] ) + + +@pytest.fixture(scope="session") +def original_datadir() -> Path: + return Path(__file__).parent / "reference" + + +def finch_assert_allclose(result, expected, **kwargs): + if hasattr(result, "to_numpy"): + result = result.to_numpy() + if hasattr(expected, "to_numpy"): + expected = expected.to_numpy() + assert_allclose(result, expected, **kwargs) diff --git a/tests/reference/test_jit_dependent_loop_inserted_code.py b/tests/reference/test_jit_dependent_loop_inserted_code.py new file mode 100644 index 0000000..6b3c319 --- /dev/null +++ b/tests/reference/test_jit_dependent_loop_inserted_code.py @@ -0,0 +1,11 @@ +def opt_fn(A, n): + A, n = maybedefer((A, n)) + B = A + A, B = compute((A, B)) + for _i in range(getattr(sum(B), 'item')()): + A, B = maybedefer((A, B)) + B = add(B, A) + A, B = compute((A, B)) + B, = defer((B,)) + B, = compute((B,)) + return B diff --git a/tests/reference/test_jit_if_branch_inserted_code.py b/tests/reference/test_jit_if_branch_inserted_code.py new file mode 100644 index 0000000..602c6be --- /dev/null +++ b/tests/reference/test_jit_if_branch_inserted_code.py @@ -0,0 +1,12 @@ +def opt_fn(A, B, use_matmul): + if use_matmul: + A, B = maybedefer((A, B)) + result = matmul(A, B) + result, = compute((result,)) + else: + A, B = maybedefer((A, B)) + result = add(A, B) + result, = compute((result,)) + result, = defer((result,)) + result, = compute((result,)) + return result diff --git a/tests/reference/test_jit_local_module_function_inserted_code.py b/tests/reference/test_jit_local_module_function_inserted_code.py new file mode 100644 index 0000000..8093464 --- /dev/null +++ b/tests/reference/test_jit_local_module_function_inserted_code.py @@ -0,0 +1,12 @@ +def opt_fn(A, B, n): + A, B, n = maybedefer((A, B, n)) + C = A + B, C, n = compute((B, C, n)) + while n > 0: + B, C, n = maybedefer((B, C, n)) + C = getattr(finch.interface, 'add')(C, B) + n = n - 1 + B, C, n = compute((B, C, n)) + C, = defer((C,)) + C, = compute((C,)) + return C diff --git a/tests/reference/test_jit_module_function_inserted_code.py b/tests/reference/test_jit_module_function_inserted_code.py new file mode 100644 index 0000000..595633d --- /dev/null +++ b/tests/reference/test_jit_module_function_inserted_code.py @@ -0,0 +1,12 @@ +def opt_fn(A, B, n): + A, B, n = maybedefer((A, B, n)) + C = A + B, C, n = compute((B, C, n)) + while n > 0: + B, C, n = maybedefer((B, C, n)) + C = getattr(getattr(finch, 'interface'), 'add')(C, B) + n = n - 1 + B, C, n = compute((B, C, n)) + C, = defer((C,)) + C, = compute((C,)) + return C diff --git a/tests/reference/test_jit_return_expr_inserted_code.py b/tests/reference/test_jit_return_expr_inserted_code.py new file mode 100644 index 0000000..ea7d32b --- /dev/null +++ b/tests/reference/test_jit_return_expr_inserted_code.py @@ -0,0 +1,6 @@ +def opt_fn(A, B): + A, B = defer((A, B)) + ret = matmul(A, B) + ret_2 = matmul(A, B) + ret, ret_2 = compute((ret, ret_2)) + return (ret, ret_2) diff --git a/tests/reference/test_jit_scalar_loop_inserted_code.py b/tests/reference/test_jit_scalar_loop_inserted_code.py new file mode 100644 index 0000000..6633b5f --- /dev/null +++ b/tests/reference/test_jit_scalar_loop_inserted_code.py @@ -0,0 +1,11 @@ +def opt_fn(A, n): + A, n = maybedefer((A, n)) + B = A + A, B, n = compute((A, B, n)) + for _i in range(n): + A, B, n = maybedefer((A, B, n)) + B = add(B, A) + A, B, n = compute((A, B, n)) + B, = defer((B,)) + B, = compute((B,)) + return B diff --git a/tests/reference/test_jit_straight_line_inserted_code.py b/tests/reference/test_jit_straight_line_inserted_code.py new file mode 100644 index 0000000..93340be --- /dev/null +++ b/tests/reference/test_jit_straight_line_inserted_code.py @@ -0,0 +1,6 @@ +def opt_fn(A, B): + A, B = maybedefer((A, B)) + C = matmul(A, B) + C, = defer((C,)) + C, = compute((C,)) + return C diff --git a/tests/reference/test_jit_two_independent_ops_inserted_code.py b/tests/reference/test_jit_two_independent_ops_inserted_code.py new file mode 100644 index 0000000..6c70eb0 --- /dev/null +++ b/tests/reference/test_jit_two_independent_ops_inserted_code.py @@ -0,0 +1,8 @@ +def opt_fn(A, B, C): + A, B, C = maybedefer((A, B, C)) + D = matmul(A, B) + E = add(A, C) + F = add(D, E) + F, = defer((F,)) + F, = compute((F,)) + return F diff --git a/tests/reference/test_jit_while_inserted_code.py b/tests/reference/test_jit_while_inserted_code.py new file mode 100644 index 0000000..45f2006 --- /dev/null +++ b/tests/reference/test_jit_while_inserted_code.py @@ -0,0 +1,12 @@ +def opt_fn(A, B, n): + A, B, n = maybedefer((A, B, n)) + C = A + B, C, n = compute((B, C, n)) + while n > 0: + B, C, n = maybedefer((B, C, n)) + C = add(C, B) + n = n - 1 + B, C, n = compute((B, C, n)) + C, = defer((C,)) + C, = compute((C,)) + return C diff --git a/tests/test_fused.py b/tests/test_fused.py new file mode 100644 index 0000000..32b5cee --- /dev/null +++ b/tests/test_fused.py @@ -0,0 +1,782 @@ +import ast +import operator +import textwrap + +import pytest + +import numpy as np + +import finch +from finch.autoschedule import ( + DefaultLogicFormatter, + DefaultLogicOptimizer, + DefaultLoopOrderer, + LogicCapture, + LogicCompiler, + LogicExecutor, + LogicNormalizer, +) +from galley_jl_python import jit +from finch.finch_fused import nodes as fzd +from finch.finch_fused.cfg_builder import ( + fused_build_cfg, + fused_desugar, + number_statements, +) +from finch.finch_fused.dataflow import ( + LivenessAnalysis, + insert_lazy_and_compute, + maybedefer, +) +from finch.finch_fused.parser import ( + fused_function_to_python_ast, + parse_fused_function, +) +from finch.finch_notation.interpreter import NotationInterpreter +from finch.interface import add, asarray, matmul, sum +from finch.interface.lazy import LazyTensor +from finch.tensor.scalar import ConstantScalar, ScalarFType + +from .conftest import finch_assert_allclose + + +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(ConstantScalar(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(ConstantScalar(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(ConstantScalar(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. Numeric literals survive as literals: + the parser leaves them alone, and a constant only becomes a ConstantScalar + later, when `defer` hands it to the lazy layer. + """ + + 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 _transformed_jit_source(fn): + transformed_fn = insert_lazy_and_compute(parse_fused_function(fn)) + assert isinstance(transformed_fn, fzd.Function) + return ast.unparse(fused_function_to_python_ast(transformed_fn)) + "\n" + + +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 test_jit_straight_line(): + """A jit function with no loops should produce the same result as eager.""" + + def simple_fn(A, B): + C = matmul(A, B) + return C # noqa: RET504 + + @jit + def opt_fn(A, B): + C = matmul(A, B) + return C # noqa: RET504 + + A = asarray(np.array([[1, 2], [3, 4]])) + B = asarray(np.array([[5, 6], [7, 8]])) + + finch_assert_allclose(opt_fn(A, B), simple_fn(A, B)) + + +def test_jit_straight_line_inserted_code(file_regression): + def opt_fn(A, B): + C = matmul(A, B) + return C # noqa: RET504 + + file_regression.check(_transformed_jit_source(opt_fn), extension=".py") + + +def test_jit_eager_only_operation_warns_and_computes(): + def simple_fn(A): + return finch.linalg.inv(A) + + @jit + def opt_fn(A): + return finch.linalg.inv(A) + + A = asarray(np.array([[1.0, 2.0], [3.0, 5.0]])) + + with pytest.warns(RuntimeWarning, match="inv"): + result = opt_fn(A) + + finch_assert_allclose(result, simple_fn(A)) + + +def test_jit_return_expr(): + """A jit function with no loops should produce the same result as eager.""" + + def simple_fn(A, B): + return matmul(A, B), matmul(A, B) + + @jit + def opt_fn(A, B): + return matmul(A, B), matmul(A, B) + + A = asarray(np.array([[1, 2], [3, 4]])) + B = asarray(np.array([[5, 6], [7, 8]])) + + finch_assert_allclose(opt_fn(A, B), simple_fn(A, B)) + + +def test_jit_return_expr_inserted_code(file_regression): + def opt_fn(A, B): + return matmul(A, B), matmul(A, B) + + file_regression.check(_transformed_jit_source(opt_fn), extension=".py") + + +def test_jit_two_independent_ops(): + """Two independent tensor ops whose results are both used.""" + + def simple_fn(A, B, C): + D = matmul(A, B) + E = add(A, C) + F = add(D, E) + return F # noqa: RET504 + + @jit + def opt_fn(A, B, C): + D = matmul(A, B) + E = add(A, C) + F = add(D, E) + return F # noqa: RET504 + + A = asarray(np.array([[1, 2], [3, 4]])) + B = asarray(np.array([[1, 0], [0, 1]])) + C = asarray(np.array([[1, 1], [1, 1]])) + + finch_assert_allclose(opt_fn(A, B, C), simple_fn(A, B, C)) + + +def test_jit_two_independent_ops_inserted_code(file_regression): + def opt_fn(A, B, C): + D = matmul(A, B) + E = add(A, C) + F = add(D, E) + return F # noqa: RET504 + + file_regression.check(_transformed_jit_source(opt_fn), extension=".py") + + +@pytest.mark.parametrize( + ("operand", "bound"), + [ + pytest.param(ConstantScalar(2.0), 0, id="constant_scalar"), + pytest.param(1.0, 0, id="specializable_literal"), + pytest.param(2.0, 1, id="runtime_literal"), + ], +) +def test_a_constant_operand_is_not_bound_as_a_tensor(operand, bound): + """A constant costs no runtime binding, however it was written. + + Everything is bound as a table on the way in; `inline_constant_scalars` + then replaces the constants with their values and drops those bindings. A + value the optimizer cannot act on stays a binding, which is the point. + """ + capture = _CollectBindings(LogicCompiler(NotationInterpreter())) + executor = LogicExecutor( + DefaultLogicOptimizer(DefaultLoopOrderer(DefaultLogicFormatter(capture))) + ) + arr = np.arange(3.0) + with finch.with_default_scheduler(LogicNormalizer(executor)): + result = finch.compute(finch.defer(asarray(arr)) + operand) + finch_assert_allclose(result, arr + float(np.asarray(operand))) + scalars = [ + t + for bindings in capture.all_bindings + for t in bindings.values() + if isinstance(t, ScalarFType) + ] + assert len(scalars) == bound + + +def test_maybedefer_defers_every_tensor(): + """`maybedefer` makes no judgement about fills -- it only defers tensors.""" + A = asarray(np.arange(3.0)) + (lazy_A, lazy_c, plain) = maybedefer((A, ConstantScalar(2.0), 2.0)) + assert isinstance(lazy_A, LazyTensor) + assert isinstance(lazy_c, LazyTensor) + assert plain == 2.0 + + +class _CollectBindings(LogicCapture): + """A LogicCapture that keeps the bindings of every lowering, not just the last.""" + + def __init__(self, ctx): + super().__init__(ctx) + self.all_bindings: list[dict] = [] + + def lower(self, prgm, bindings, stats, stats_factory): + self.all_bindings.append(bindings.copy()) + return super().lower(prgm, bindings, stats, stats_factory) + + +def test_jit_inlines_literal_operands(): + """A specializable literal in a jit body reaches the kernel as a constant.""" + + @jit + def opt_fn(A): + return add(A, 1.0) + + capture = _CollectBindings(LogicCompiler(NotationInterpreter())) + executor = LogicExecutor( + DefaultLogicOptimizer(DefaultLoopOrderer(DefaultLogicFormatter(capture))) + ) + + arr = np.arange(3.0) + with finch.with_default_scheduler(LogicNormalizer(executor)): + result = opt_fn(asarray(arr)) + finch_assert_allclose(result, arr + 1.0) + scalars = [ + t + for bindings in capture.all_bindings + for t in bindings.values() + if isinstance(t, ScalarFType) + ] + assert not scalars + + +@pytest.mark.parametrize("trip_counts", [(1, 2, 4, 8, 16)]) +def test_jit_constant_folded_in_a_loop_does_not_grow_the_kernel_cache(trip_counts): + """A constant incremented in a loop must not cost a kernel per iteration. + + `n` folds to a fresh value on every trip. Without the demotion in + `maybedefer` each fresh value is a fresh literal, so the program -- and the + kernel compiled from it -- differs per iteration and the cache grows without + bound in the trip count. + """ + + @jit + def const_loop(A, k): + n = 1 + B = A + for _i in range(k): + n = n + 1 + B = add(B, n) + return B + + arr = np.arange(3.0) + counts = [] + for k in trip_counts: + executor = LogicExecutor( + DefaultLogicOptimizer( + DefaultLoopOrderer( + DefaultLogicFormatter(LogicCompiler(NotationInterpreter())) + ) + ), + cache=True, + ) + with finch.with_default_scheduler(LogicNormalizer(executor)): + result = const_loop(asarray(arr), k) + want = arr.copy() + n = 1 + for _ in range(k): + n += 1 + want = want + n + finch_assert_allclose(result, want) + counts.append(len(executor.cached_kernels)) + + # Flat, not linear: once the loop body has been compiled the cache stops + # growing, however many more trips are added. + assert counts[-1] == counts[-2] == counts[len(counts) // 2], counts + + +def test_jit_scalar_loop(): + """A loop with a scalar iteration count and tensor accumulation.""" + + def simple_fn(A, n): + B = A + for _i in range(n): + B = add(B, A) + return B + + @jit + def opt_fn(A, n): + B = A + for _i in range(n): + B = add(B, A) + return B + + A = asarray(np.array([[1, 0], [0, 1]], dtype=float)) + + finch_assert_allclose(opt_fn(A, 3), simple_fn(A, 3)) + + +def test_jit_scalar_loop_inserted_code(file_regression): + def opt_fn(A, n): + B = A + for _i in range(n): + B = add(B, A) + return B + + file_regression.check(_transformed_jit_source(opt_fn), extension=".py") + + +def test_jit_dependent_loop(): + """A loop with an iterator that depends on a computation.""" + + def simple_fn(A, n): + B = A + for _i in range(sum(B).item()): + B = add(B, A) + return B + + @jit + def opt_fn(A, n): + B = A + for _i in range(sum(B).item()): + B = add(B, A) + return B + + A = asarray(np.array([[1, 0], [0, 1]], dtype=int)) + + finch_assert_allclose(opt_fn(A, 3), simple_fn(A, 3)) + + +def test_jit_dependent_loop_inserted_code(file_regression): + def opt_fn(A, n): + B = A + for _i in range(sum(B).item()): + B = add(B, A) + return B + + file_regression.check(_transformed_jit_source(opt_fn), extension=".py") + + +def test_jit_if_branch(): + """A jit function with an if/else over tensor ops.""" + + def simple_fn(A, B, use_matmul): + if use_matmul: # noqa: SIM108 + result = matmul(A, B) + else: + result = add(A, B) + return result + + @jit + def opt_fn(A, B, use_matmul): + if use_matmul: # noqa: SIM108 + result = matmul(A, B) + else: + result = add(A, B) + return result + + A = asarray(np.array([[1, 2], [3, 4]])) + B = asarray(np.array([[1, 0], [0, 1]])) + + finch_assert_allclose(opt_fn(A, B, True), simple_fn(A, B, True)) + finch_assert_allclose(opt_fn(A, B, False), simple_fn(A, B, False)) + + +def test_jit_if_branch_inserted_code(file_regression): + def opt_fn(A, B, use_matmul): + if use_matmul: # noqa: SIM108 + result = matmul(A, B) + else: + result = add(A, B) + return result + + file_regression.check(_transformed_jit_source(opt_fn), extension=".py") + + +def test_jit_while(): + """A jit function with a while loop.""" + + def simple_fn(A, B, n): + C = A + while n > 0: + C = add(C, B) + n = n - 1 + return C + + @jit + def opt_fn(A, B, n): + C = A + while n > 0: + C = add(C, B) + n = n - 1 + return C + + A = asarray(np.array([[1, 2], [3, 4]])) + B = asarray(np.array([[1, 0], [0, 1]])) + n = 3 + + finch_assert_allclose(opt_fn(A, B, n), simple_fn(A, B, n)) + + +def test_jit_while_inserted_code(file_regression): + def opt_fn(A, B, n): + C = A + while n > 0: + C = add(C, B) + n = n - 1 + return C + + file_regression.check(_transformed_jit_source(opt_fn), extension=".py") + + +def test_jit_module_function(): + """A jit function with a function from a module.""" + + def simple_fn(A, B, n): + C = A + while n > 0: + C = finch.interface.add(C, B) + n = n - 1 + return C + + @jit + def opt_fn(A, B, n): + C = A + while n > 0: + C = finch.interface.add(C, B) + n = n - 1 + return C + + A = asarray(np.array([[1, 2], [3, 4]])) + B = asarray(np.array([[1, 0], [0, 1]])) + n = 3 + + finch_assert_allclose(opt_fn(A, B, n), simple_fn(A, B, n)) + + +def test_jit_module_function_inserted_code(file_regression): + def opt_fn(A, B, n): + C = A + while n > 0: + C = finch.interface.add(C, B) + n = n - 1 + return C + + file_regression.check(_transformed_jit_source(opt_fn), extension=".py") + + +def test_jit_local_module_function(): + """A jit function with a function from a module.""" + + xp = finch.interface + + def simple_fn(A, B, n): + C = A + while n > 0: + C = xp.add(C, B) + n = n - 1 + return C + + @jit + def opt_fn(A, B, n): + C = A + while n > 0: + C = xp.add(C, B) + n = n - 1 + return C + + A = asarray(np.array([[1, 2], [3, 4]])) + B = asarray(np.array([[1, 0], [0, 1]])) + n = 3 + + finch_assert_allclose(simple_fn(A, B, n), opt_fn(A, B, n)) + + +def test_jit_local_module_function_inserted_code(file_regression): + xp = finch.interface + + def opt_fn(A, B, n): + C = A + while n > 0: + C = xp.add(C, B) + n = n - 1 + return C + + file_regression.check(_transformed_jit_source(opt_fn), extension=".py") From 81f5c8dbc98d90c9cd89b4d34ae4d66ae2a12886 Mon Sep 17 00:00:00 2001 From: JoelMathewC Date: Fri, 18 Sep 2026 14:31:52 -0400 Subject: [PATCH 11/20] chore: fix precommit issues --- pyproject.toml | 1 + tests/test_fused.py | 3 ++- 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 518edb9..0b9dfd2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -37,6 +37,7 @@ build-backend = "poetry.core.masonry.api" [tool.ruff] line-length = 88 +extend-exclude = ["tests/reference"] [tool.ruff.lint] select = ["F", "E", "W", "I", "B", "UP", "YTT", "BLE", "C4", "T10", "ISC", "ICN", "PIE", "PYI", "RSE", "RET", "SIM", "PGH", "FLY", "NPY", "PERF", "RUF022"] diff --git a/tests/test_fused.py b/tests/test_fused.py index 32b5cee..a599327 100644 --- a/tests/test_fused.py +++ b/tests/test_fused.py @@ -16,7 +16,6 @@ LogicExecutor, LogicNormalizer, ) -from galley_jl_python import jit from finch.finch_fused import nodes as fzd from finch.finch_fused.cfg_builder import ( fused_build_cfg, @@ -37,6 +36,8 @@ from finch.interface.lazy import LazyTensor from finch.tensor.scalar import ConstantScalar, ScalarFType +from galley_jl_python import jit + from .conftest import finch_assert_allclose From a54dbeed2d3dc5392274e696a9cc769a7916f29f Mon Sep 17 00:00:00 2001 From: JoelMathewC Date: Fri, 18 Sep 2026 15:17:19 -0400 Subject: [PATCH 12/20] chore: removing redundant tests --- pixi.toml | 1 - pyproject.toml | 3 - src/galley_jl_python/__init__.py | 3 +- src/galley_jl_python/finch_fused.py | 33 + tests/conftest.py | 16 - .../test_jit_dependent_loop_inserted_code.py | 11 - .../test_jit_if_branch_inserted_code.py | 12 - ...jit_local_module_function_inserted_code.py | 12 - .../test_jit_module_function_inserted_code.py | 12 - .../test_jit_return_expr_inserted_code.py | 6 - .../test_jit_scalar_loop_inserted_code.py | 11 - .../test_jit_straight_line_inserted_code.py | 6 - ...t_jit_two_independent_ops_inserted_code.py | 8 - .../reference/test_jit_while_inserted_code.py | 12 - tests/test_fused.py | 786 +----------------- 15 files changed, 73 insertions(+), 859 deletions(-) create mode 100644 src/galley_jl_python/finch_fused.py delete mode 100644 tests/reference/test_jit_dependent_loop_inserted_code.py delete mode 100644 tests/reference/test_jit_if_branch_inserted_code.py delete mode 100644 tests/reference/test_jit_local_module_function_inserted_code.py delete mode 100644 tests/reference/test_jit_module_function_inserted_code.py delete mode 100644 tests/reference/test_jit_return_expr_inserted_code.py delete mode 100644 tests/reference/test_jit_scalar_loop_inserted_code.py delete mode 100644 tests/reference/test_jit_straight_line_inserted_code.py delete mode 100644 tests/reference/test_jit_two_independent_ops_inserted_code.py delete mode 100644 tests/reference/test_jit_while_inserted_code.py diff --git a/pixi.toml b/pixi.toml index b2778ab..f19ec57 100644 --- a/pixi.toml +++ b/pixi.toml @@ -28,7 +28,6 @@ numba = ">=0.61" scipy = "*" numpy = "==2.*" pytest-xdist = ">=3.6.1,<4" -pytest-regressions = ">=2.8.1,<3.0.0" [feature.test.tasks] test = { cmd = "pytest", depends-on = ["compile"] } diff --git a/pyproject.toml b/pyproject.toml index 0b9dfd2..da0ad37 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -28,7 +28,6 @@ test = [ "hypothesis (>=6.151.9,<7.0.0)", "ndindex (>=1.10.1,<2.0.0)", "pytest-json-report (>=1.5.0,<2.0.0)", - "pytest-regressions>=2.8.1,<3.0.0", ] [build-system] @@ -37,7 +36,6 @@ build-backend = "poetry.core.masonry.api" [tool.ruff] line-length = 88 -extend-exclude = ["tests/reference"] [tool.ruff.lint] select = ["F", "E", "W", "I", "B", "UP", "YTT", "BLE", "C4", "T10", "ISC", "ICN", "PIE", "PYI", "RSE", "RET", "SIM", "PGH", "FLY", "NPY", "PERF", "RUF022"] @@ -64,4 +62,3 @@ section-order = [ [tool.mypy] ignore_missing_imports = true -exclude = ["tests/reference"] diff --git a/src/galley_jl_python/__init__.py b/src/galley_jl_python/__init__.py index beb66e9..0a19a2c 100644 --- a/src/galley_jl_python/__init__.py +++ b/src/galley_jl_python/__init__.py @@ -84,8 +84,6 @@ pi as pi, ) -from finch.finch_fused import jit - from . import linalg from ._array_api_info import __array_namespace_info__ from .compiled import ( @@ -117,6 +115,7 @@ uint32, uint64, ) +from .finch_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 new file mode 100644 index 0000000..2f64efc --- /dev/null +++ b/src/galley_jl_python/finch_fused.py @@ -0,0 +1,33 @@ +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/tests/conftest.py b/tests/conftest.py index dcd5795..61800c9 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,9 +1,6 @@ -from pathlib import Path - import pytest import numpy as np -from numpy.testing import assert_allclose def pytest_addoption(parser): @@ -50,16 +47,3 @@ def arr3d(): [[0, 0, 0, 0], [1, 5, 0, 3]], ] ) - - -@pytest.fixture(scope="session") -def original_datadir() -> Path: - return Path(__file__).parent / "reference" - - -def finch_assert_allclose(result, expected, **kwargs): - if hasattr(result, "to_numpy"): - result = result.to_numpy() - if hasattr(expected, "to_numpy"): - expected = expected.to_numpy() - assert_allclose(result, expected, **kwargs) diff --git a/tests/reference/test_jit_dependent_loop_inserted_code.py b/tests/reference/test_jit_dependent_loop_inserted_code.py deleted file mode 100644 index 6b3c319..0000000 --- a/tests/reference/test_jit_dependent_loop_inserted_code.py +++ /dev/null @@ -1,11 +0,0 @@ -def opt_fn(A, n): - A, n = maybedefer((A, n)) - B = A - A, B = compute((A, B)) - for _i in range(getattr(sum(B), 'item')()): - A, B = maybedefer((A, B)) - B = add(B, A) - A, B = compute((A, B)) - B, = defer((B,)) - B, = compute((B,)) - return B diff --git a/tests/reference/test_jit_if_branch_inserted_code.py b/tests/reference/test_jit_if_branch_inserted_code.py deleted file mode 100644 index 602c6be..0000000 --- a/tests/reference/test_jit_if_branch_inserted_code.py +++ /dev/null @@ -1,12 +0,0 @@ -def opt_fn(A, B, use_matmul): - if use_matmul: - A, B = maybedefer((A, B)) - result = matmul(A, B) - result, = compute((result,)) - else: - A, B = maybedefer((A, B)) - result = add(A, B) - result, = compute((result,)) - result, = defer((result,)) - result, = compute((result,)) - return result diff --git a/tests/reference/test_jit_local_module_function_inserted_code.py b/tests/reference/test_jit_local_module_function_inserted_code.py deleted file mode 100644 index 8093464..0000000 --- a/tests/reference/test_jit_local_module_function_inserted_code.py +++ /dev/null @@ -1,12 +0,0 @@ -def opt_fn(A, B, n): - A, B, n = maybedefer((A, B, n)) - C = A - B, C, n = compute((B, C, n)) - while n > 0: - B, C, n = maybedefer((B, C, n)) - C = getattr(finch.interface, 'add')(C, B) - n = n - 1 - B, C, n = compute((B, C, n)) - C, = defer((C,)) - C, = compute((C,)) - return C diff --git a/tests/reference/test_jit_module_function_inserted_code.py b/tests/reference/test_jit_module_function_inserted_code.py deleted file mode 100644 index 595633d..0000000 --- a/tests/reference/test_jit_module_function_inserted_code.py +++ /dev/null @@ -1,12 +0,0 @@ -def opt_fn(A, B, n): - A, B, n = maybedefer((A, B, n)) - C = A - B, C, n = compute((B, C, n)) - while n > 0: - B, C, n = maybedefer((B, C, n)) - C = getattr(getattr(finch, 'interface'), 'add')(C, B) - n = n - 1 - B, C, n = compute((B, C, n)) - C, = defer((C,)) - C, = compute((C,)) - return C diff --git a/tests/reference/test_jit_return_expr_inserted_code.py b/tests/reference/test_jit_return_expr_inserted_code.py deleted file mode 100644 index ea7d32b..0000000 --- a/tests/reference/test_jit_return_expr_inserted_code.py +++ /dev/null @@ -1,6 +0,0 @@ -def opt_fn(A, B): - A, B = defer((A, B)) - ret = matmul(A, B) - ret_2 = matmul(A, B) - ret, ret_2 = compute((ret, ret_2)) - return (ret, ret_2) diff --git a/tests/reference/test_jit_scalar_loop_inserted_code.py b/tests/reference/test_jit_scalar_loop_inserted_code.py deleted file mode 100644 index 6633b5f..0000000 --- a/tests/reference/test_jit_scalar_loop_inserted_code.py +++ /dev/null @@ -1,11 +0,0 @@ -def opt_fn(A, n): - A, n = maybedefer((A, n)) - B = A - A, B, n = compute((A, B, n)) - for _i in range(n): - A, B, n = maybedefer((A, B, n)) - B = add(B, A) - A, B, n = compute((A, B, n)) - B, = defer((B,)) - B, = compute((B,)) - return B diff --git a/tests/reference/test_jit_straight_line_inserted_code.py b/tests/reference/test_jit_straight_line_inserted_code.py deleted file mode 100644 index 93340be..0000000 --- a/tests/reference/test_jit_straight_line_inserted_code.py +++ /dev/null @@ -1,6 +0,0 @@ -def opt_fn(A, B): - A, B = maybedefer((A, B)) - C = matmul(A, B) - C, = defer((C,)) - C, = compute((C,)) - return C diff --git a/tests/reference/test_jit_two_independent_ops_inserted_code.py b/tests/reference/test_jit_two_independent_ops_inserted_code.py deleted file mode 100644 index 6c70eb0..0000000 --- a/tests/reference/test_jit_two_independent_ops_inserted_code.py +++ /dev/null @@ -1,8 +0,0 @@ -def opt_fn(A, B, C): - A, B, C = maybedefer((A, B, C)) - D = matmul(A, B) - E = add(A, C) - F = add(D, E) - F, = defer((F,)) - F, = compute((F,)) - return F diff --git a/tests/reference/test_jit_while_inserted_code.py b/tests/reference/test_jit_while_inserted_code.py deleted file mode 100644 index 45f2006..0000000 --- a/tests/reference/test_jit_while_inserted_code.py +++ /dev/null @@ -1,12 +0,0 @@ -def opt_fn(A, B, n): - A, B, n = maybedefer((A, B, n)) - C = A - B, C, n = compute((B, C, n)) - while n > 0: - B, C, n = maybedefer((B, C, n)) - C = add(C, B) - n = n - 1 - B, C, n = compute((B, C, n)) - C, = defer((C,)) - C, = compute((C,)) - return C diff --git a/tests/test_fused.py b/tests/test_fused.py index a599327..b95c5ff 100644 --- a/tests/test_fused.py +++ b/tests/test_fused.py @@ -1,783 +1,75 @@ -import ast -import operator -import textwrap - -import pytest - import numpy as np +from numpy.testing import assert_allclose -import finch -from finch.autoschedule import ( - DefaultLogicFormatter, - DefaultLogicOptimizer, - DefaultLoopOrderer, - LogicCapture, - LogicCompiler, - LogicExecutor, - LogicNormalizer, -) -from finch.finch_fused import nodes as fzd -from finch.finch_fused.cfg_builder import ( - fused_build_cfg, - fused_desugar, - number_statements, -) -from finch.finch_fused.dataflow import ( - LivenessAnalysis, - insert_lazy_and_compute, - maybedefer, -) -from finch.finch_fused.parser import ( - fused_function_to_python_ast, - parse_fused_function, -) -from finch.finch_notation.interpreter import NotationInterpreter -from finch.interface import add, asarray, matmul, sum -from finch.interface.lazy import LazyTensor -from finch.tensor.scalar import ConstantScalar, ScalarFType - -from galley_jl_python import jit - -from .conftest import finch_assert_allclose - - -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(ConstantScalar(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(ConstantScalar(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(ConstantScalar(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. Numeric literals survive as literals: - the parser leaves them alone, and a constant only becomes a ConstantScalar - later, when `defer` hands it to the lazy layer. - """ - - 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 _transformed_jit_source(fn): - transformed_fn = insert_lazy_and_compute(parse_fused_function(fn)) - assert isinstance(transformed_fn, fzd.Function) - return ast.unparse(fused_function_to_python_ast(transformed_fn)) + "\n" - - -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 +import galley_jl_python as galley -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 test_jit_straight_line(): - """A jit function with no loops should produce the same result as eager.""" - - def simple_fn(A, B): - C = matmul(A, B) - return C # noqa: RET504 - - @jit - def opt_fn(A, B): - C = matmul(A, B) - return C # noqa: RET504 - - A = asarray(np.array([[1, 2], [3, 4]])) - B = asarray(np.array([[5, 6], [7, 8]])) - - finch_assert_allclose(opt_fn(A, B), simple_fn(A, B)) - - -def test_jit_straight_line_inserted_code(file_regression): - def opt_fn(A, B): - C = matmul(A, B) - return C # noqa: RET504 - - file_regression.check(_transformed_jit_source(opt_fn), extension=".py") - - -def test_jit_eager_only_operation_warns_and_computes(): - def simple_fn(A): - return finch.linalg.inv(A) - - @jit - def opt_fn(A): - return finch.linalg.inv(A) - - A = asarray(np.array([[1.0, 2.0], [3.0, 5.0]])) - - with pytest.warns(RuntimeWarning, match="inv"): - result = opt_fn(A) - - finch_assert_allclose(result, simple_fn(A)) +def assert_galley_allclose(actual, expected): + assert isinstance(actual, galley.Tensor) + assert_allclose(actual.todense(), expected) -def test_jit_return_expr(): - """A jit function with no loops should produce the same result as eager.""" - - def simple_fn(A, B): - return matmul(A, B), matmul(A, B) - - @jit +def test_jit_straight_line_uses_galley_tensors(): + @galley.jit def opt_fn(A, B): - return matmul(A, B), matmul(A, B) + return galley.matmul(A, B) - A = asarray(np.array([[1, 2], [3, 4]])) - B = asarray(np.array([[5, 6], [7, 8]])) + A = galley.Tensor(np.array([[1, 2], [3, 4]])) + B = galley.Tensor(np.array([[5, 6], [7, 8]])) - finch_assert_allclose(opt_fn(A, B), simple_fn(A, B)) + assert_galley_allclose(opt_fn(A, B), np.matmul(A.todense(), B.todense())) -def test_jit_return_expr_inserted_code(file_regression): +def test_jit_return_expr_uses_galley_tensors(): + @galley.jit def opt_fn(A, B): - return matmul(A, B), matmul(A, B) - - file_regression.check(_transformed_jit_source(opt_fn), extension=".py") - - -def test_jit_two_independent_ops(): - """Two independent tensor ops whose results are both used.""" - - def simple_fn(A, B, C): - D = matmul(A, B) - E = add(A, C) - F = add(D, E) - return F # noqa: RET504 - - @jit - def opt_fn(A, B, C): - D = matmul(A, B) - E = add(A, C) - F = add(D, E) - return F # noqa: RET504 - - A = asarray(np.array([[1, 2], [3, 4]])) - B = asarray(np.array([[1, 0], [0, 1]])) - C = asarray(np.array([[1, 1], [1, 1]])) - - finch_assert_allclose(opt_fn(A, B, C), simple_fn(A, B, C)) - - -def test_jit_two_independent_ops_inserted_code(file_regression): - def opt_fn(A, B, C): - D = matmul(A, B) - E = add(A, C) - F = add(D, E) - return F # noqa: RET504 - - file_regression.check(_transformed_jit_source(opt_fn), extension=".py") - - -@pytest.mark.parametrize( - ("operand", "bound"), - [ - pytest.param(ConstantScalar(2.0), 0, id="constant_scalar"), - pytest.param(1.0, 0, id="specializable_literal"), - pytest.param(2.0, 1, id="runtime_literal"), - ], -) -def test_a_constant_operand_is_not_bound_as_a_tensor(operand, bound): - """A constant costs no runtime binding, however it was written. - - Everything is bound as a table on the way in; `inline_constant_scalars` - then replaces the constants with their values and drops those bindings. A - value the optimizer cannot act on stays a binding, which is the point. - """ - capture = _CollectBindings(LogicCompiler(NotationInterpreter())) - executor = LogicExecutor( - DefaultLogicOptimizer(DefaultLoopOrderer(DefaultLogicFormatter(capture))) - ) - arr = np.arange(3.0) - with finch.with_default_scheduler(LogicNormalizer(executor)): - result = finch.compute(finch.defer(asarray(arr)) + operand) - finch_assert_allclose(result, arr + float(np.asarray(operand))) - scalars = [ - t - for bindings in capture.all_bindings - for t in bindings.values() - if isinstance(t, ScalarFType) - ] - assert len(scalars) == bound - - -def test_maybedefer_defers_every_tensor(): - """`maybedefer` makes no judgement about fills -- it only defers tensors.""" - A = asarray(np.arange(3.0)) - (lazy_A, lazy_c, plain) = maybedefer((A, ConstantScalar(2.0), 2.0)) - assert isinstance(lazy_A, LazyTensor) - assert isinstance(lazy_c, LazyTensor) - assert plain == 2.0 - - -class _CollectBindings(LogicCapture): - """A LogicCapture that keeps the bindings of every lowering, not just the last.""" - - def __init__(self, ctx): - super().__init__(ctx) - self.all_bindings: list[dict] = [] - - def lower(self, prgm, bindings, stats, stats_factory): - self.all_bindings.append(bindings.copy()) - return super().lower(prgm, bindings, stats, stats_factory) - - -def test_jit_inlines_literal_operands(): - """A specializable literal in a jit body reaches the kernel as a constant.""" - - @jit - def opt_fn(A): - return add(A, 1.0) - - capture = _CollectBindings(LogicCompiler(NotationInterpreter())) - executor = LogicExecutor( - DefaultLogicOptimizer(DefaultLoopOrderer(DefaultLogicFormatter(capture))) - ) - - arr = np.arange(3.0) - with finch.with_default_scheduler(LogicNormalizer(executor)): - result = opt_fn(asarray(arr)) - finch_assert_allclose(result, arr + 1.0) - scalars = [ - t - for bindings in capture.all_bindings - for t in bindings.values() - if isinstance(t, ScalarFType) - ] - assert not scalars + return galley.matmul(A, B), galley.add(A, B) + A = galley.Tensor(np.array([[1, 2], [3, 4]])) + B = galley.Tensor(np.array([[5, 6], [7, 8]])) -@pytest.mark.parametrize("trip_counts", [(1, 2, 4, 8, 16)]) -def test_jit_constant_folded_in_a_loop_does_not_grow_the_kernel_cache(trip_counts): - """A constant incremented in a loop must not cost a kernel per iteration. + actual_matmul, actual_add = opt_fn(A, B) + assert_galley_allclose(actual_matmul, np.matmul(A.todense(), B.todense())) + assert_galley_allclose(actual_add, A.todense() + B.todense()) - `n` folds to a fresh value on every trip. Without the demotion in - `maybedefer` each fresh value is a fresh literal, so the program -- and the - kernel compiled from it -- differs per iteration and the cache grows without - bound in the trip count. - """ - @jit - def const_loop(A, k): - n = 1 - B = A - for _i in range(k): - n = n + 1 - B = add(B, n) - return B - - arr = np.arange(3.0) - counts = [] - for k in trip_counts: - executor = LogicExecutor( - DefaultLogicOptimizer( - DefaultLoopOrderer( - DefaultLogicFormatter(LogicCompiler(NotationInterpreter())) - ) - ), - cache=True, - ) - with finch.with_default_scheduler(LogicNormalizer(executor)): - result = const_loop(asarray(arr), k) - want = arr.copy() - n = 1 - for _ in range(k): - n += 1 - want = want + n - finch_assert_allclose(result, want) - counts.append(len(executor.cached_kernels)) - - # Flat, not linear: once the loop body has been compiled the cache stops - # growing, however many more trips are added. - assert counts[-1] == counts[-2] == counts[len(counts) // 2], counts - - -def test_jit_scalar_loop(): - """A loop with a scalar iteration count and tensor accumulation.""" - - def simple_fn(A, n): - B = A - for _i in range(n): - B = add(B, A) - return B - - @jit - def opt_fn(A, n): - B = A - for _i in range(n): - B = add(B, A) - return B - - A = asarray(np.array([[1, 0], [0, 1]], dtype=float)) - - finch_assert_allclose(opt_fn(A, 3), simple_fn(A, 3)) - - -def test_jit_scalar_loop_inserted_code(file_regression): +def test_jit_scalar_loop_uses_galley_tensors(): + @galley.jit def opt_fn(A, n): B = A for _i in range(n): - B = add(B, A) - return B - - file_regression.check(_transformed_jit_source(opt_fn), extension=".py") - - -def test_jit_dependent_loop(): - """A loop with an iterator that depends on a computation.""" - - def simple_fn(A, n): - B = A - for _i in range(sum(B).item()): - B = add(B, A) - return B - - @jit - def opt_fn(A, n): - B = A - for _i in range(sum(B).item()): - B = add(B, A) + B = galley.add(B, A) return B - A = asarray(np.array([[1, 0], [0, 1]], dtype=int)) - - finch_assert_allclose(opt_fn(A, 3), simple_fn(A, 3)) - - -def test_jit_dependent_loop_inserted_code(file_regression): - def opt_fn(A, n): - B = A - for _i in range(sum(B).item()): - B = add(B, A) - return B - - file_regression.check(_transformed_jit_source(opt_fn), extension=".py") - - -def test_jit_if_branch(): - """A jit function with an if/else over tensor ops.""" - - def simple_fn(A, B, use_matmul): - if use_matmul: # noqa: SIM108 - result = matmul(A, B) - else: - result = add(A, B) - return result + A = galley.Tensor(np.array([[1, 0], [0, 1]], dtype=float)) - @jit - def opt_fn(A, B, use_matmul): - if use_matmul: # noqa: SIM108 - result = matmul(A, B) - else: - result = add(A, B) - return result - - A = asarray(np.array([[1, 2], [3, 4]])) - B = asarray(np.array([[1, 0], [0, 1]])) + assert_galley_allclose(opt_fn(A, 3), 4 * A.todense()) - finch_assert_allclose(opt_fn(A, B, True), simple_fn(A, B, True)) - finch_assert_allclose(opt_fn(A, B, False), simple_fn(A, B, False)) - -def test_jit_if_branch_inserted_code(file_regression): +def test_jit_if_branch_uses_galley_tensors(): + @galley.jit def opt_fn(A, B, use_matmul): - if use_matmul: # noqa: SIM108 - result = matmul(A, B) - else: - result = add(A, B) - return result + if use_matmul: + return galley.matmul(A, B) + return galley.add(A, B) - file_regression.check(_transformed_jit_source(opt_fn), extension=".py") + A = galley.Tensor(np.array([[1, 2], [3, 4]])) + B = galley.Tensor(np.array([[1, 0], [0, 1]])) + assert_galley_allclose(opt_fn(A, B, True), np.matmul(A.todense(), B.todense())) + assert_galley_allclose(opt_fn(A, B, False), A.todense() + B.todense()) -def test_jit_while(): - """A jit function with a while loop.""" - def simple_fn(A, B, n): - C = A - while n > 0: - C = add(C, B) - n = n - 1 - return C - - @jit +def test_jit_while_loop_uses_galley_tensors(): + @galley.jit def opt_fn(A, B, n): C = A while n > 0: - C = add(C, B) + C = galley.add(C, B) n = n - 1 return C - A = asarray(np.array([[1, 2], [3, 4]])) - B = asarray(np.array([[1, 0], [0, 1]])) - n = 3 - - finch_assert_allclose(opt_fn(A, B, n), simple_fn(A, B, n)) - - -def test_jit_while_inserted_code(file_regression): - def opt_fn(A, B, n): - C = A - while n > 0: - C = add(C, B) - n = n - 1 - return C - - file_regression.check(_transformed_jit_source(opt_fn), extension=".py") - - -def test_jit_module_function(): - """A jit function with a function from a module.""" - - def simple_fn(A, B, n): - C = A - while n > 0: - C = finch.interface.add(C, B) - n = n - 1 - return C - - @jit - def opt_fn(A, B, n): - C = A - while n > 0: - C = finch.interface.add(C, B) - n = n - 1 - return C - - A = asarray(np.array([[1, 2], [3, 4]])) - B = asarray(np.array([[1, 0], [0, 1]])) - n = 3 - - finch_assert_allclose(opt_fn(A, B, n), simple_fn(A, B, n)) - - -def test_jit_module_function_inserted_code(file_regression): - def opt_fn(A, B, n): - C = A - while n > 0: - C = finch.interface.add(C, B) - n = n - 1 - return C - - file_regression.check(_transformed_jit_source(opt_fn), extension=".py") - - -def test_jit_local_module_function(): - """A jit function with a function from a module.""" - - xp = finch.interface - - def simple_fn(A, B, n): - C = A - while n > 0: - C = xp.add(C, B) - n = n - 1 - return C - - @jit - def opt_fn(A, B, n): - C = A - while n > 0: - C = xp.add(C, B) - n = n - 1 - return C - - A = asarray(np.array([[1, 2], [3, 4]])) - B = asarray(np.array([[1, 0], [0, 1]])) - n = 3 - - finch_assert_allclose(simple_fn(A, B, n), opt_fn(A, B, n)) - - -def test_jit_local_module_function_inserted_code(file_regression): - xp = finch.interface - - def opt_fn(A, B, n): - C = A - while n > 0: - C = xp.add(C, B) - n = n - 1 - return C + A = galley.Tensor(np.array([[1, 2], [3, 4]])) + B = galley.Tensor(np.array([[1, 0], [0, 1]])) - file_regression.check(_transformed_jit_source(opt_fn), extension=".py") + assert_galley_allclose(opt_fn(A, B, 3), A.todense() + 3 * B.todense()) From 44d3df013840deecf4e6b25c8132b9ada3a657ce Mon Sep 17 00:00:00 2001 From: JoelMathewC Date: Fri, 18 Sep 2026 15:20:14 -0400 Subject: [PATCH 13/20] chore: rename galley->gl --- tests/test_fused.py | 44 ++++++++++++++++++++++---------------------- 1 file changed, 22 insertions(+), 22 deletions(-) diff --git a/tests/test_fused.py b/tests/test_fused.py index b95c5ff..ac6a456 100644 --- a/tests/test_fused.py +++ b/tests/test_fused.py @@ -1,32 +1,32 @@ import numpy as np from numpy.testing import assert_allclose -import galley_jl_python as galley +import galley_jl_python as gl def assert_galley_allclose(actual, expected): - assert isinstance(actual, galley.Tensor) + assert isinstance(actual, gl.Tensor) assert_allclose(actual.todense(), expected) def test_jit_straight_line_uses_galley_tensors(): - @galley.jit + @gl.jit def opt_fn(A, B): - return galley.matmul(A, B) + return gl.matmul(A, B) - A = galley.Tensor(np.array([[1, 2], [3, 4]])) - B = galley.Tensor(np.array([[5, 6], [7, 8]])) + A = gl.Tensor(np.array([[1, 2], [3, 4]])) + B = gl.Tensor(np.array([[5, 6], [7, 8]])) assert_galley_allclose(opt_fn(A, B), np.matmul(A.todense(), B.todense())) def test_jit_return_expr_uses_galley_tensors(): - @galley.jit + @gl.jit def opt_fn(A, B): - return galley.matmul(A, B), galley.add(A, B) + return gl.matmul(A, B), gl.add(A, B) - A = galley.Tensor(np.array([[1, 2], [3, 4]])) - B = galley.Tensor(np.array([[5, 6], [7, 8]])) + A = gl.Tensor(np.array([[1, 2], [3, 4]])) + B = gl.Tensor(np.array([[5, 6], [7, 8]])) actual_matmul, actual_add = opt_fn(A, B) assert_galley_allclose(actual_matmul, np.matmul(A.todense(), B.todense())) @@ -34,42 +34,42 @@ def opt_fn(A, B): def test_jit_scalar_loop_uses_galley_tensors(): - @galley.jit + @gl.jit def opt_fn(A, n): B = A for _i in range(n): - B = galley.add(B, A) + B = gl.add(B, A) return B - A = galley.Tensor(np.array([[1, 0], [0, 1]], dtype=float)) + A = gl.Tensor(np.array([[1, 0], [0, 1]], dtype=float)) assert_galley_allclose(opt_fn(A, 3), 4 * A.todense()) def test_jit_if_branch_uses_galley_tensors(): - @galley.jit + @gl.jit def opt_fn(A, B, use_matmul): if use_matmul: - return galley.matmul(A, B) - return galley.add(A, B) + return gl.matmul(A, B) + return gl.add(A, B) - A = galley.Tensor(np.array([[1, 2], [3, 4]])) - B = galley.Tensor(np.array([[1, 0], [0, 1]])) + A = gl.Tensor(np.array([[1, 2], [3, 4]])) + B = gl.Tensor(np.array([[1, 0], [0, 1]])) assert_galley_allclose(opt_fn(A, B, True), np.matmul(A.todense(), B.todense())) assert_galley_allclose(opt_fn(A, B, False), A.todense() + B.todense()) def test_jit_while_loop_uses_galley_tensors(): - @galley.jit + @gl.jit def opt_fn(A, B, n): C = A while n > 0: - C = galley.add(C, B) + C = gl.add(C, B) n = n - 1 return C - A = galley.Tensor(np.array([[1, 2], [3, 4]])) - B = galley.Tensor(np.array([[1, 0], [0, 1]])) + A = gl.Tensor(np.array([[1, 2], [3, 4]])) + B = gl.Tensor(np.array([[1, 0], [0, 1]])) assert_galley_allclose(opt_fn(A, B, 3), A.todense() + 3 * B.todense()) From 0ab8e39e84cabe71ee36c989adf68aef87b8d8e3 Mon Sep 17 00:00:00 2001 From: kylebd99 Date: Fri, 25 Sep 2026 17:51:49 -0400 Subject: [PATCH 14/20] merge to pyproject.toml --- pixi.toml | 36 ------------------------------------ pyproject.toml | 25 +++++++++++++++++++++++++ 2 files changed, 25 insertions(+), 36 deletions(-) delete mode 100644 pixi.toml diff --git a/pixi.toml b/pixi.toml deleted file mode 100644 index f19ec57..0000000 --- a/pixi.toml +++ /dev/null @@ -1,36 +0,0 @@ -[project] -authors = ["Hameer Abbasi <2190658+hameerabbasi@users.noreply.github.com>"] -channels = ["conda-forge"] -description = "Add a short description here" -name = "galley-jl-python" -platforms = ["osx-arm64"] -version = "0.2.12" - -[tasks] -compile = "python -c 'import galley_jl_python'" - -[dependencies] -python = ">=3.10,<3.13" -juliaup = ">=1.17.10,<2" - -[pypi-dependencies] -galley-jl-python = { path = ".", editable = true } -juliapkg = ">=0.1.16,<0.2" -juliacall = ">=0.9.24,<0.10" -numpy = ">=1.19" -finch-tensor = ">=0.7.0,<0.8.0" - -[feature.test.pypi-dependencies] -pytest = "*" -pytest-cov = "*" -sparse = ">=0.16,<0.17" -numba = ">=0.61" -scipy = "*" -numpy = "==2.*" -pytest-xdist = ">=3.6.1,<4" - -[feature.test.tasks] -test = { cmd = "pytest", depends-on = ["compile"] } - -[environments] -test = ["test"] diff --git a/pyproject.toml b/pyproject.toml index da0ad37..bab48c0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -28,12 +28,37 @@ test = [ "hypothesis (>=6.151.9,<7.0.0)", "ndindex (>=1.10.1,<2.0.0)", "pytest-json-report (>=1.5.0,<2.0.0)", + "pytest-xdist>=3.6.1,<4", ] [build-system] requires = ["poetry-core>=1.0.8"] build-backend = "poetry.core.masonry.api" +[tool.pixi.workspace] +channels = ["conda-forge"] +platforms = [ + { name = "linux-amd64", platform = "linux-64", glibc = "2.34" }, + { name = "linux-arm64", platform = "linux-aarch64", glibc = "2.34" }, + { name = "macos-arm64", platform = "osx-arm64", macos = "15.0" }, + { name = "win-amd64", platform = "win-64" }, +] + +[tool.pixi.dependencies] +python = ">=3.11,<3.13" + +[tool.pixi.pypi-dependencies] +galley-jl-python = { path = ".", editable = true } + +[tool.pixi.tasks] +compile = "python -c 'import galley_jl_python'" + +[tool.pixi.feature.test.tasks] +test = { cmd = "pytest", depends-on = ["compile"] } + +[tool.pixi.environments] +test = { features = ["test"], solve-group = "default" } + [tool.ruff] line-length = 88 From 54e79bbdf6d9f028a11eac9e894c700f77b24526 Mon Sep 17 00:00:00 2001 From: kylebd99 Date: Fri, 25 Sep 2026 19:04:58 -0400 Subject: [PATCH 15/20] update --- array-api-skips.txt | 1 + src/galley_jl_python/compiled.py | 3 +++ src/galley_jl_python/tensor.py | 8 +++++++- 3 files changed, 11 insertions(+), 1 deletion(-) diff --git a/array-api-skips.txt b/array-api-skips.txt index 797379d..ade08fe 100644 --- a/array-api-skips.txt +++ b/array-api-skips.txt @@ -249,6 +249,7 @@ array_api_tests/test_indexing_functions.py::test_take # test_linalg # not implemented +array_api_tests/test_linalg.py::test_cross array_api_tests/test_linalg.py::test_matrix_transpose array_api_tests/test_linalg.py::test_vecdot array_api_tests/test_linalg.py::test_eigh diff --git a/src/galley_jl_python/compiled.py b/src/galley_jl_python/compiled.py index af05973..83907ad 100644 --- a/src/galley_jl_python/compiled.py +++ b/src/galley_jl_python/compiled.py @@ -141,6 +141,9 @@ def set_optimizer(opt: AbstractScheduler) -> None: jl.Finch.set_scheduler_b(opt.get_julia_scheduler()) +set_optimizer(GalleyScheduler()) + + def lazy(tensor: Tensor) -> Tensor: from .tensor import Tensor diff --git a/src/galley_jl_python/tensor.py b/src/galley_jl_python/tensor.py index 06195c7..8f98335 100644 --- a/src/galley_jl_python/tensor.py +++ b/src/galley_jl_python/tensor.py @@ -25,6 +25,11 @@ ) from .typing import Device, DType, JuliaObj, OrderType, TupleOf3Arrays, spmatrix +# Builds a 0-d Finch tensor in Julia so the element keeps its Julia eltype. +_scalar_to_0d = jl.seval( + "(T, fv, x) -> Finch.Tensor(Finch.Element{convert(T, fv)}(T[convert(T, x)]))" +) + class SparseArray: """ @@ -279,7 +284,8 @@ def __getitem__(self, key): return Tensor(result) if jl.isa(result, jl.Finch.Tensor): return Tensor(jl.swizzle(result, *range(1, jl.ndims(result) + 1))) - return result + # the Array API requires a 0-d array, not a scalar, for integer-only keys + return Tensor(jl.swizzle(_scalar_to_0d(self.dtype, self.fill_value, result))) @property def dtype(self) -> DType: From e8921d1736f5bb53cac053e7a9143dd79ff69490 Mon Sep 17 00:00:00 2001 From: kylebd99 Date: Sat, 26 Sep 2026 11:14:05 -0400 Subject: [PATCH 16/20] Add sysimage workflow --- .github/workflows/sysimage.yml | 53 ++++ .gitignore | 3 + README.md | 30 +++ pyproject.toml | 5 +- scripts/sysimage/_load.py | 15 ++ scripts/sysimage/build.py | 79 ++++++ scripts/sysimage/build_sysimage.jl | 26 ++ scripts/sysimage/fetch.py | 58 +++++ scripts/sysimage/pin_julia_deps.py | 54 +++++ scripts/sysimage/workload.py | 62 +++++ src/galley_jl_python/_sysimage.py | 79 ++++++ src/galley_jl_python/julia.py | 8 +- src/galley_jl_python/juliapkg.json | 374 ++++++++++++++++++++++++++++- 13 files changed, 839 insertions(+), 7 deletions(-) create mode 100644 .github/workflows/sysimage.yml create mode 100644 scripts/sysimage/_load.py create mode 100644 scripts/sysimage/build.py create mode 100644 scripts/sysimage/build_sysimage.jl create mode 100644 scripts/sysimage/fetch.py create mode 100644 scripts/sysimage/pin_julia_deps.py create mode 100644 scripts/sysimage/workload.py create mode 100644 src/galley_jl_python/_sysimage.py diff --git a/.github/workflows/sysimage.yml b/.github/workflows/sysimage.yml new file mode 100644 index 0000000..fa246cb --- /dev/null +++ b/.github/workflows/sysimage.yml @@ -0,0 +1,53 @@ +name: Sysimage + +# Builds the Julia sysimage for each platform and publishes it to the GitHub +# release tagged `sysimage-`, where identifies the exact Julia +# environment. Users download it with `pixi run fetch-sysimage`. +on: + workflow_dispatch: + push: + branches: + - main + paths: + - src/galley_jl_python/juliapkg.json + - scripts/sysimage/** + - pyproject.toml + - pixi.lock + - .github/workflows/sysimage.yml + +permissions: + contents: write + +jobs: + build: + strategy: + matrix: + os: [ubuntu-latest, ubuntu-24.04-arm, macos-latest, windows-latest] + fail-fast: false + runs-on: ${{ matrix.os }} + # GitHub-hosted maximum; the generic multi-CPU build is slow + timeout-minutes: 360 + defaults: + run: + shell: bash + steps: + - name: Checkout Repo + uses: actions/checkout@v7 + - uses: prefix-dev/setup-pixi@v0.10.0 + with: + pixi-version: v0.72.2 + - name: Build sysimage + id: build + run: pixi run build-sysimage + - name: Publish to release + env: + GH_TOKEN: ${{ github.token }} + TAG: ${{ steps.build.outputs.tag }} + IMAGE: ${{ steps.build.outputs.path }} + run: | + # the matrix jobs race to create the release; losing the race is fine + gh release view "$TAG" > /dev/null 2>&1 \ + || gh release create "$TAG" --prerelease --title "$TAG" \ + --notes "Prebuilt Julia sysimages. Fetch with \`pixi run fetch-sysimage\`." \ + || gh release view "$TAG" > /dev/null + gh release upload "$TAG" "$IMAGE" --clobber diff --git a/.gitignore b/.gitignore index 2b2b89f..52ec6f2 100644 --- a/.gitignore +++ b/.gitignore @@ -168,3 +168,6 @@ cython_debug/ .pixi/ *.egg-info pixi.lock + +# PackageCompiler environment used by scripts/sysimage/build_sysimage.jl +scripts/sysimage/.build-env/ diff --git a/README.md b/README.md index 9f01527..bd13d3b 100644 --- a/README.md +++ b/README.md @@ -38,6 +38,36 @@ Options: --path Path to the local copy of Finch.jl [default: ../Finch.jl]. ``` +### Julia sysimage + +Most of Galley's startup time is Julia compiling Finch itself. A prebuilt Julia +sysimage removes it: the first operations of a session drop from minutes to a +few seconds. + +- `pixi run fetch-sysimage` downloads the image for your platform into + `~/.cache/galley-jl-python/` (`GALLEY_JL_PYTHON_CACHE` overrides this). + `pixi run compile` and `pixi run test` run this step first. +- `pixi run build-sysimage` builds the image locally instead. This takes hours. +- `import galley_jl_python` loads a cached image automatically when it matches + the current Julia environment. Otherwise Julia starts without it. Set + `GALLEY_JL_PYTHON_SYSIMAGE=0` to turn this off. + +An image works only with the exact Julia and package versions it was built +from. For this reason `src/galley_jl_python/juliapkg.json` pins Julia and every +Julia package. Each image's name includes a hash of that environment. A local +Finch.jl from `develop.py` therefore runs without the image. + +To update the Julia dependencies: + +1. Loosen the pins you want to change. +2. Resolve with `pixi run compile`. +3. Re-pin with `python scripts/sysimage/pin_julia_deps.py`. + +Pushing the new pins to `main` runs the "Sysimage" GitHub Action. It builds +images for Linux, macOS and Windows and publishes them to a `sysimage-` +GitHub release, where `fetch-sysimage` finds them. The action can also be run +manually from the Actions tab. + ### Publishing The "Publish" GitHub Action is a manual workflow for publishing Python packages to PyPI using Poetry. It handles the version management based on the `pyproject.toml` file and automates tagging and creating GitHub releases. diff --git a/pyproject.toml b/pyproject.toml index bab48c0..1fd8224 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -51,7 +51,10 @@ python = ">=3.11,<3.13" galley-jl-python = { path = ".", editable = true } [tool.pixi.tasks] -compile = "python -c 'import galley_jl_python'" +# pixi has no post-install hook, so the entry-point tasks fetch the sysimage +compile = { cmd = "python -c 'import galley_jl_python'", depends-on = ["fetch-sysimage"] } +build-sysimage = "python scripts/sysimage/build.py" +fetch-sysimage = "python scripts/sysimage/fetch.py" [tool.pixi.feature.test.tasks] test = { cmd = "pytest", depends-on = ["compile"] } diff --git a/scripts/sysimage/_load.py b/scripts/sysimage/_load.py new file mode 100644 index 0000000..cd04e5b --- /dev/null +++ b/scripts/sysimage/_load.py @@ -0,0 +1,15 @@ +"""Load ``galley_jl_python._sysimage`` by path. + +Importing it through the package would run ``galley_jl_python/__init__.py``, +which starts Julia. +""" + +import importlib.util +from pathlib import Path + +_PATH = ( + Path(__file__).resolve().parents[2] / "src" / "galley_jl_python" / "_sysimage.py" +) +_spec = importlib.util.spec_from_file_location("_galley_sysimage", _PATH) +sysimage = importlib.util.module_from_spec(_spec) +_spec.loader.exec_module(sysimage) diff --git a/scripts/sysimage/build.py b/scripts/sysimage/build.py new file mode 100644 index 0000000..950a39c --- /dev/null +++ b/scripts/sysimage/build.py @@ -0,0 +1,79 @@ +"""Build a Julia sysimage that removes most of galley's first-call compile time. + +Traces the Julia methods compiled by ``workload.py``, compiles them into a +sysimage with PackageCompiler, and stores it in the per-user cache under a name +tied to the current Julia environment, where ``galley_jl_python`` finds it. +""" + +import argparse +import os +import shutil +import subprocess +import sys +import tempfile +from pathlib import Path + +import juliapkg + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) + +from _load import sysimage # noqa: E402 + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + parser.add_argument( + "--statements", + type=Path, + help="reuse precompile statements from an earlier trace instead of tracing", + ) + args = parser.parse_args() + + project = juliapkg.project() + key = sysimage.env_key(project) + image = sysimage.image_path(key) + image.parent.mkdir(parents=True, exist_ok=True) + + with tempfile.TemporaryDirectory() as tmp: + statements = args.statements or Path(tmp) / "precompile_statements.jl" + if args.statements is None: + print(f"[sysimage] tracing workload -> {statements}", flush=True) + env = { + k: v for k, v in os.environ.items() if k != "PYTHON_JULIACALL_SYSIMAGE" + } + subprocess.run( + [sys.executable, str(HERE / "workload.py")], + env={ + **env, + "GALLEY_JL_PYTHON_SYSIMAGE": "0", + "PYTHON_JULIACALL_TRACE_COMPILE": str(statements), + }, + check=True, + ) + + # build next to the target and rename, so an interrupted build never + # leaves a partial image where it would be loaded + partial = Path(tmp) / image.name + print(f"[sysimage] building {image.name} (this takes a while)", flush=True) + subprocess.run( + [ + juliapkg.executable(), + "--startup-file=no", + str(HERE / "build_sysimage.jl"), + project, + str(statements), + str(partial), + ], + check=True, + ) + shutil.move(partial, image) + + print(f"[sysimage] wrote {image}", flush=True) + if github_output := os.environ.get("GITHUB_OUTPUT"): + with open(github_output, "a") as f: + f.write(f"key={key}\npath={image}\ntag={sysimage.release_tag(key)}\n") + + +if __name__ == "__main__": + main() diff --git a/scripts/sysimage/build_sysimage.jl b/scripts/sysimage/build_sysimage.jl new file mode 100644 index 0000000..fcd8b59 --- /dev/null +++ b/scripts/sysimage/build_sysimage.jl @@ -0,0 +1,26 @@ +# Usage: julia build_sysimage.jl +# +# PackageCompiler lives in its own environment so it never enters the +# project that juliapkg manages. +using Pkg + +project, statements_file, output = ARGS + +build_env = joinpath(@__DIR__, ".build-env") +Pkg.activate(build_env; io=devnull) +if Base.find_package("PackageCompiler") === nothing + Pkg.add("PackageCompiler") +end +using PackageCompiler + +packages = ["Finch", "HDF5", "NPZ", "TensorMarket", "PythonCall"] + +create_sysimage( + packages; + project=project, + sysimage_path=output, + precompile_statements_file=statements_file, + # a generic target so the image runs on any x86_64 CPU, not just this one + cpu_target=PackageCompiler.default_app_cpu_target(), + include_transitive_dependencies=true, +) diff --git a/scripts/sysimage/fetch.py b/scripts/sysimage/fetch.py new file mode 100644 index 0000000..725569e --- /dev/null +++ b/scripts/sysimage/fetch.py @@ -0,0 +1,58 @@ +"""Download the prebuilt sysimage for this platform and Julia environment. + +CI publishes one image per platform to the GitHub release tagged +``sysimage-``, where the key identifies the exact Julia environment. +""" + +import os +import sys +import tempfile +import urllib.error +import urllib.request +from pathlib import Path + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) + +from _load import sysimage # noqa: E402 + + +def main() -> None: + key = sysimage.current_key() + image = sysimage.image_path(key) + if image.is_file(): + print(f"[sysimage] already present: {image}") + return + + url = f"{sysimage.RELEASES_URL}/{sysimage.release_tag(key)}/{image.name}" + print(f"[sysimage] downloading {url}", flush=True) + image.parent.mkdir(parents=True, exist_ok=True) + fd, partial = tempfile.mkstemp(dir=image.parent, suffix=".part") + try: + with os.fdopen(fd, "wb") as out, urllib.request.urlopen(url) as resp: + while chunk := resp.read(1 << 20): + out.write(chunk) + os.replace(partial, image) + # never fail: other tasks depend on this one, and galley works without an + # image, just with a slower first call + except urllib.error.HTTPError as e: + if e.code != 404: + print(f"[sysimage] download failed ({e}); continuing without it.") + return + print( + f"[sysimage] no prebuilt image for this platform and Julia environment " + f"({image.name}); continuing without it. Build one locally with " + "`pixi run build-sysimage`." + ) + return + except (urllib.error.URLError, OSError) as e: + print(f"[sysimage] download failed ({e}); continuing without it.") + return + finally: + if os.path.exists(partial): + os.remove(partial) + print(f"[sysimage] saved {image}") + + +if __name__ == "__main__": + main() diff --git a/scripts/sysimage/pin_julia_deps.py b/scripts/sysimage/pin_julia_deps.py new file mode 100644 index 0000000..94726b7 --- /dev/null +++ b/scripts/sysimage/pin_julia_deps.py @@ -0,0 +1,54 @@ +"""Pin galley's Julia dependencies to the currently resolved environment. + +Rewrites ``src/galley_jl_python/juliapkg.json`` so that Julia and every +registered package in the resolved ``Manifest.toml`` (transitive dependencies +included) are pinned to exact versions. Everyone then resolves the same +manifest, which the prebuilt sysimage requires. + +To update: loosen the pins you want to move, resolve (e.g. ``pixi run +compile``), run this script, then rebuild with ``pixi run build-sysimage``. +""" + +import argparse +import json +import sys +import tomllib +from pathlib import Path + +ROOT = Path(__file__).resolve().parents[2] +JULIAPKG_JSON = ROOT / "src" / "galley_jl_python" / "juliapkg.json" +# juliacall pins these itself (PythonCall to its own release, OpenSSL_jll to +# match Python's OpenSSL), so they must stay unpinned here +UNPINNED = {"OpenSSL_jll", "PythonCall"} + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + parser.add_argument( + "--manifest", + type=Path, + default=Path(sys.prefix) / "julia_env" / "Manifest.toml", + ) + args = parser.parse_args() + + manifest = tomllib.loads(args.manifest.read_text()) + config = json.loads(JULIAPKG_JSON.read_text()) + config["julia"] = f"={manifest['julia_version']}" + + packages = {} + for name, (entry, *_) in sorted(manifest["deps"].items()): + # stdlibs have no tree hash and ship with Julia itself + if "git-tree-sha1" not in entry or name in UNPINNED: + continue + # drop build metadata (e.g. JLL "+0"): juliapkg can't parse it and Pkg + # compat ignores it anyway + version = entry["version"].split("+")[0] + packages[name] = {"uuid": entry["uuid"], "version": f"={version}"} + config["packages"] = packages + + JULIAPKG_JSON.write_text(json.dumps(config, indent=2) + "\n") + print(f"pinned julia {config['julia']} and {len(packages)} packages") + + +if __name__ == "__main__": + main() diff --git a/scripts/sysimage/workload.py b/scripts/sysimage/workload.py new file mode 100644 index 0000000..030167d --- /dev/null +++ b/scripts/sysimage/workload.py @@ -0,0 +1,62 @@ +"""Representative workload traced to build the Julia sysimage. + +Every Julia method compiled while this runs is recorded (via juliacall's +``trace_compile`` option) and baked into the sysimage by ``build.py``. Add +operations here to cover more of the library ahead of time. +""" + +import numpy as np + +import galley_jl_python as gl + + +def run_ops(A, B, C, v): + gl.matmul(A, B) + gl.sum(gl.matmul(gl.matmul(A, B), C)) + gl.add(A, B) + gl.multiply(A, B) + gl.subtract(A, B) + gl.sum(A, axis=0) + gl.sum(A, axis=1) + gl.max(A) + gl.tensordot(A, v, axes=1) + gl.permute_dims(A, (1, 0)) + A[0] + A[0, 0] + A[1:3, :] + A.todense() + + @gl.jit + def sum_abc(A, B, C): + return gl.sum(gl.matmul(gl.matmul(A, B), C)) + + @gl.jit + def chain(A, B, C): + return gl.add(gl.matmul(A, B), C) + + sum_abc(A, B, C) + chain(A, B, C) + + +def main(): + rng = np.random.default_rng(0) + for scheduler in (gl.GalleyScheduler(), gl.DefaultScheduler()): + gl.set_optimizer(scheduler) + for n in (50, 400): + A, B, C = ( + gl.random((n, n), density=0.02, random_state=s) for s in (1, 2, 3) + ) + v = gl.asarray(rng.random(n)) + run_ops(A, B, C, v) + for fmt in ("coo", "csr", "csc", "dense"): + for dtype in (np.float64, np.int64): + dense = (rng.random((60, 60)) < 0.05).astype(dtype) + X = gl.asarray(dense, format=fmt) + gl.matmul(X, X) + gl.sum(X) + gl.add(X, X) + gl.set_optimizer(gl.GalleyScheduler()) + + +if __name__ == "__main__": + main() diff --git a/src/galley_jl_python/_sysimage.py b/src/galley_jl_python/_sysimage.py new file mode 100644 index 0000000..a96ffbe --- /dev/null +++ b/src/galley_jl_python/_sysimage.py @@ -0,0 +1,79 @@ +"""Locate a prebuilt Julia sysimage and point juliacall at it. + +This module must not import juliacall: it runs before Julia starts. + +A sysimage only works with the exact Julia version and package versions it was +built from, so each image is named after a hash of both (see ``env_key``) and +lives in a per-user cache directory. An image is used only when its name +matches the current environment; otherwise Julia starts normally. + +Get an image with ``pixi run fetch-sysimage`` (downloads the one CI built for +this environment) or ``pixi run build-sysimage`` (builds it locally). Set +``GALLEY_JL_PYTHON_SYSIMAGE=0`` to never load one, or +``PYTHON_JULIACALL_SYSIMAGE`` to use a specific image. +""" + +import hashlib +import json +import os +import platform +import sys +import tomllib +from pathlib import Path + +RELEASES_URL = "https://github.com/finch-tensor/galley-jl-python/releases/download" +_EXTENSIONS = {"linux": "so", "darwin": "dylib", "win32": "dll"} + + +def cache_dir() -> Path: + if custom := os.environ.get("GALLEY_JL_PYTHON_CACHE"): + return Path(custom) + if sys.platform == "win32": + base = Path(os.environ.get("LOCALAPPDATA", Path.home() / "AppData" / "Local")) + elif sys.platform == "darwin": + base = Path.home() / "Library" / "Caches" + else: + base = Path(os.environ.get("XDG_CACHE_HOME", Path.home() / ".cache")) + return base / "galley-jl-python" + + +def env_key(project: str) -> str: + """Hash the Julia version and resolved package set of a juliapkg project.""" + manifest = tomllib.loads((Path(project) / "Manifest.toml").read_text()) + resolved = sorted( + (name, entry.get("uuid"), entry.get("version"), entry.get("git-tree-sha1")) + for name, entries in manifest["deps"].items() + for entry in entries + ) + payload = json.dumps([manifest["julia_version"], resolved]) + return hashlib.sha256(payload.encode()).hexdigest()[:16] + + +def release_tag(key: str) -> str: + return f"sysimage-{key}" + + +def image_name(key: str) -> str: + ext = _EXTENSIONS.get(sys.platform, "so") + return f"galley-{sys.platform}-{platform.machine().lower()}-{key}.{ext}" + + +def image_path(key: str) -> Path: + return cache_dir() / image_name(key) + + +def current_key() -> str: + import juliapkg + + return env_key(juliapkg.project()) + + +def configure() -> None: + """Point juliacall at the cached sysimage for this environment, if any.""" + if os.environ.get("GALLEY_JL_PYTHON_SYSIMAGE") == "0": + return + if "PYTHON_JULIACALL_SYSIMAGE" in os.environ: + return + image = image_path(current_key()) + if image.is_file(): + os.environ["PYTHON_JULIACALL_SYSIMAGE"] = str(image) diff --git a/src/galley_jl_python/julia.py b/src/galley_jl_python/julia.py index c8f2191..b865900 100644 --- a/src/galley_jl_python/julia.py +++ b/src/galley_jl_python/julia.py @@ -1,10 +1,16 @@ import juliapkg # noqa: I001, F401 +from . import _sysimage + +_sysimage.configure() + # To change the version of Finch used, see the documentation for pyjuliapkg here: https://github.com/JuliaPy/pyjuliapkg # Use pyjuliapkg to modify the `juliapkg.json` file in the root of this repo. # You can also run `develop.py` to quickly use a local copy of Finch.jl. # An example development json is found in `juliapkg_dev.json` -import juliacall as jc # noqa: F401 +# Julia and all packages are pinned so the prebuilt sysimage matches; update the +# pins with `scripts/sysimage/pin_julia_deps.py`. +import juliacall as jc # noqa: E402, F401 from juliacall import Main as jl # noqa: E402, F401 jl.seval("using Finch") diff --git a/src/galley_jl_python/juliapkg.json b/src/galley_jl_python/juliapkg.json index ad999ee..77dd492 100644 --- a/src/galley_jl_python/juliapkg.json +++ b/src/galley_jl_python/juliapkg.json @@ -1,21 +1,385 @@ { - "julia": "1.10 - 1.12", + "julia": "=1.12.2", "packages": { + "AbstractTrees": { + "uuid": "1520ce14-60c1-5f80-bbc7-55ef81b5835c", + "version": "=0.4.5" + }, + "AcceleratedKernels": { + "uuid": "6a4ca0a5-0e36-4168-a932-d9be78d558f1", + "version": "=0.4.3" + }, + "Accessors": { + "uuid": "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697", + "version": "=0.1.45" + }, + "Adapt": { + "uuid": "79e6a3ab-5dfb-504d-930d-738a2a938a0e", + "version": "=4.7.1" + }, + "AliasTables": { + "uuid": "66dad0bd-aa9a-41b7-9441-69ab47430ed8", + "version": "=1.1.3" + }, + "ArgCheck": { + "uuid": "dce04be8-c92d-5529-be00-80e4d2c0e197", + "version": "=2.5.0" + }, + "Atomix": { + "uuid": "a9b6321e-bd34-4604-b9c9-b65b8de01458", + "version": "=1.2.1" + }, + "AutoHashEquals": { + "uuid": "15f4f7f2-30c1-5605-9d31-71845cf9641f", + "version": "=2.2.0" + }, + "CommonSolve": { + "uuid": "38540f10-b2f7-11e9-35d8-d573e4eb0ff2", + "version": "=0.2.14" + }, + "Compat": { + "uuid": "34da2185-b29b-5c13-b0c7-acf172513d20", + "version": "=4.18.1" + }, + "CompositionsBase": { + "uuid": "a33af91c-f02d-484b-be07-31d278c5ca2b", + "version": "=0.1.2" + }, + "CondaPkg": { + "uuid": "992eb4ea-22a4-4c89-a5bb-47a3300528ab", + "version": "=0.2.36" + }, + "ConstructionBase": { + "uuid": "187b0558-2788-49d3-abe0-74a17ed4e7c9", + "version": "=1.6.0" + }, + "DataAPI": { + "uuid": "9a962f9c-6df0-11e9-0e5d-c546b8b5ee8a", + "version": "=1.16.0" + }, + "DataStructures": { + "uuid": "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8", + "version": "=0.19.6" + }, + "DataValueInterfaces": { + "uuid": "e2d170a0-9d28-54be-80f0-106bbe20a464", + "version": "=1.0.0" + }, + "Distributions": { + "uuid": "31c24e10-a181-5473-b8eb-7969acd0382f", + "version": "=0.25.131" + }, + "DocStringExtensions": { + "uuid": "ffbed154-4ef7-542d-bbb7-c09d3a79fcae", + "version": "=0.9.5" + }, + "FileIO": { + "uuid": "5789e2e9-d7fb-5bc7-8068-2c6fae9b9549", + "version": "=1.20.0" + }, + "FillArrays": { + "uuid": "1a297f60-69ca-5386-bcde-b61e274b549b", + "version": "=1.17.1" + }, "Finch": { "uuid": "9177782c-1635-4eb9-9bfb-d9dfa25e6bce", - "version": "1.2.10" + "version": "=1.6.0" + }, + "GPUArraysCore": { + "uuid": "46192b85-c4d5-4398-a991-12ede77f4527", + "version": "=0.2.1" + }, + "Gamma": { + "uuid": "a0844989-3bd2-4988-8bea-c9407ab0941b", + "version": "=1.2.0" }, "HDF5": { "uuid": "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f", - "version": "0.17.2" + "version": "=0.17.4" + }, + "HDF5_jll": { + "uuid": "0234f1f7-429e-5d53-9886-15a909be8d59", + "version": "=2.2.2" + }, + "Hwloc_jll": { + "uuid": "e33a78d0-f292-5ffc-b300-72abe9b543c8", + "version": "=2.14.1" + }, + "HypergeometricFunctions": { + "uuid": "34004b35-14d8-5ef3-9330-4cdb6864b03a", + "version": "=0.3.30" + }, + "InverseFunctions": { + "uuid": "3587e190-3f89-42d0-90ee-14403ec27112", + "version": "=0.1.17" + }, + "IrrationalConstants": { + "uuid": "92d709cd-6900-40b7-9082-c6be49f344b6", + "version": "=0.2.6" + }, + "IterTools": { + "uuid": "c8e1da08-722c-5040-9ed9-7db0dc04731e", + "version": "=1.10.0" + }, + "IteratorInterfaceExtensions": { + "uuid": "82899510-4779-5014-852e-03e436cf321d", + "version": "=1.0.0" + }, + "JLLWrappers": { + "uuid": "692b3bcd-3c85-4b1f-b108-f13ce0eb3210", + "version": "=1.8.0" + }, + "JSON": { + "uuid": "682c06a0-de6a-54ab-a142-c8b1cf79cde6", + "version": "=1.9.0" + }, + "KernelAbstractions": { + "uuid": "63c18a36-062a-441e-b654-da1e3ab1ce7c", + "version": "=0.9.42" + }, + "Libiconv_jll": { + "uuid": "94ce4f54-9a6c-5748-9c1c-f9c7231a4531", + "version": "=1.18.0" + }, + "LogExpFunctions": { + "uuid": "2ab3a3ac-af41-5b50-aa03-7779005ae688", + "version": "=1.0.2" + }, + "MPIABI_jll": { + "uuid": "b5ada748-db0f-5fc0-8972-9331c762740c", + "version": "=1.0.0" + }, + "MPICH_jll": { + "uuid": "7cb0a576-ebde-5e09-9194-50597f1243b4", + "version": "=5.0.1" + }, + "MPIPreferences": { + "uuid": "3da0fdf6-3ccc-4f1b-acd9-58baa6c99267", + "version": "=0.1.12" + }, + "MPItrampoline_jll": { + "uuid": "f1f71cc9-e9ae-5b93-9b94-4fe0e1ad3748", + "version": "=5.5.6" + }, + "MacroTools": { + "uuid": "1914dd2f-81c6-5fcd-8719-6d5c9610ff09", + "version": "=0.5.16" + }, + "MicroMamba": { + "uuid": "0b3b1443-0f03-428d-bdfb-f27f9c1191ea", + "version": "=0.1.15" + }, + "MicrosoftMPI_jll": { + "uuid": "9237b28f-5490-5468-be7b-bb81f5f5e6cf", + "version": "=10.1.4" + }, + "Missings": { + "uuid": "e1d29d7a-bbdc-5cf2-9ac0-f12de2c33e28", + "version": "=1.2.0" }, "NPZ": { "uuid": "15e1cf62-19b3-5cfa-8e77-841668bca605", - "version": "0.4.3" + "version": "=0.4.3" + }, + "OpenMPI_jll": { + "uuid": "fe0851c0-eecd-5654-98d4-656369965a5c", + "version": "=5.0.12" + }, + "OpenSpecFun_jll": { + "uuid": "efe28fd5-8261-553b-a9e1-b2916fc3738e", + "version": "=0.5.6" + }, + "OrderedCollections": { + "uuid": "bac558e1-5e72-5ebc-8fee-abe8a469f55d", + "version": "=2.0.1" + }, + "PDMats": { + "uuid": "90014a1f-27ba-587c-ab20-58faa44d9150", + "version": "=0.11.41" + }, + "Parsers": { + "uuid": "69de0a69-1ddd-5017-9359-2bf0b02dc9f0", + "version": "=3.0.0" + }, + "Pidfile": { + "uuid": "fa939f87-e72e-5be4-a000-7fc836dbe307", + "version": "=1.3.0" + }, + "PrecompileTools": { + "uuid": "aea7be01-6a6a-4083-8856-8a6e6704d82a", + "version": "=1.3.4" + }, + "Preferences": { + "uuid": "21216c6a-2e73-6563-6e65-726566657250", + "version": "=1.6.0" + }, + "PtrArrays": { + "uuid": "43287f4e-b6f4-7ad1-bb20-aadabca52c3d", + "version": "=1.4.0" + }, + "QuadGK": { + "uuid": "1fd47b50-473d-5c70-9696-f719f8f3bcdc", + "version": "=2.11.3" + }, + "Reexport": { + "uuid": "189a3867-3050-52da-a836-e630ba90ab69", + "version": "=1.2.2" + }, + "Requires": { + "uuid": "ae029012-a4dd-5104-9daa-d747884805df", + "version": "=1.3.1" + }, + "RewriteTools": { + "uuid": "5969e224-3634-4c61-9f66-721b69e98b8a", + "version": "=0.2.4" + }, + "Rmath": { + "uuid": "79098fc4-a85e-5d69-aa6a-4863f24498fa", + "version": "=0.9.0" + }, + "Rmath_jll": { + "uuid": "f50d1b31-88e8-58de-be2c-1cc44531875f", + "version": "=0.5.2" + }, + "Roots": { + "uuid": "f2b01f46-fcfa-551c-844a-d8ac1e96c665", + "version": "=3.0.8" + }, + "Scratch": { + "uuid": "6c6a2e73-6563-6170-7368-637461726353", + "version": "=1.3.0" + }, + "SortingAlgorithms": { + "uuid": "a2af1166-a08f-5f64-846c-94a0d3cef48c", + "version": "=1.2.3" + }, + "SpecialFunctions": { + "uuid": "276daf66-3868-5448-9aa4-cd146d93841b", + "version": "=2.9.0" + }, + "StaticArrays": { + "uuid": "90137ffa-7385-5640-81b9-e52037218182", + "version": "=1.9.22" + }, + "StaticArraysCore": { + "uuid": "1e83bf80-4336-4d27-bf5d-d5a4f845583c", + "version": "=1.4.4" + }, + "Statistics": { + "uuid": "10745b16-79ce-11e8-11f9-7d13ad32a3b2", + "version": "=1.11.5" + }, + "StatsAPI": { + "uuid": "82ae8749-77ed-4fe6-ae5f-f523153014b0", + "version": "=1.8.0" + }, + "StatsBase": { + "uuid": "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91", + "version": "=0.34.13" + }, + "StatsFuns": { + "uuid": "4c63d2b9-4356-54db-8cca-17b64c39e42c", + "version": "=2.2.1" + }, + "StructUtils": { + "uuid": "ec057cc2-7a8d-4b58-b3b3-92acb9f63b42", + "version": "=2.9.2" + }, + "SyntaxInterface": { + "uuid": "b33eeca9-aacb-4496-a840-e75f1646a4fb", + "version": "=0.2.1" + }, + "TableTraits": { + "uuid": "3783bdb8-4a98-5b6b-af9a-565f29a5fe9c", + "version": "=1.0.1" + }, + "Tables": { + "uuid": "bd369af6-aec1-5ad0-b16a-f7cc5008161c", + "version": "=1.14.0" }, "TensorMarket": { "uuid": "8b7d4fe7-0b45-4d0d-9dd8-5cc9b23b4b77", - "version": "0.2.0" + "version": "=0.2.2" + }, + "UnsafeAtomics": { + "uuid": "013be700-e6cd-48c3-b4a1-df204f14c38f", + "version": "=0.3.2" + }, + "UnsafePointers": { + "uuid": "e17b2a0c-0bdf-430a-bd0c-3a23cae4ff39", + "version": "=1.0.0" + }, + "XML2_jll": { + "uuid": "02c8fc9c-b97f-50b9-bbe4-9be30ff0a78a", + "version": "=2.15.4" + }, + "Xorg_libpciaccess_jll": { + "uuid": "a65dc6b1-eb27-53a1-bb3e-dea574b5389e", + "version": "=0.19.0" + }, + "ZipFile": { + "uuid": "a5390f91-8eb1-5f08-bee0-b1d1ffed6cea", + "version": "=0.10.1" + }, + "aws_c_auth_jll": { + "uuid": "2b3700d1-4306-52e2-a478-c162f0c514be", + "version": "=0.9.6" + }, + "aws_c_cal_jll": { + "uuid": "70f11efc-bab2-57f1-b0f3-22aad4e67c4b", + "version": "=0.9.13" + }, + "aws_c_common_jll": { + "uuid": "73048d1d-b8c4-5092-a58d-866c5e8d1e50", + "version": "=0.12.6" + }, + "aws_c_compression_jll": { + "uuid": "73a04cd5-f3d7-5bac-9290-e8adb709f224", + "version": "=0.3.2" + }, + "aws_c_http_jll": { + "uuid": "3254fc65-9028-534d-aa9d-d76d128babc6", + "version": "=0.10.15" + }, + "aws_c_io_jll": { + "uuid": "13c41daa-f319-5298-b5eb-5754e0170d52", + "version": "=0.26.3" + }, + "aws_c_s3_jll": { + "uuid": "bd1f34fb-993f-5903-a121-aaf302eed6d4", + "version": "=0.11.5" + }, + "aws_c_sdkutils_jll": { + "uuid": "1282aa60-004d-510b-9f52-12498d409daa", + "version": "=0.2.4" + }, + "aws_checksums_jll": { + "uuid": "b2a88e68-78e7-5e94-8c20-c02986ec140e", + "version": "=0.2.10" + }, + "dlfcn_win32_jll": { + "uuid": "c4b69c83-5512-53e3-94e6-de98773c479f", + "version": "=1.4.2" + }, + "libaec_jll": { + "uuid": "477f73a3-ac25-53e9-8cc3-50b2fa2566f0", + "version": "=1.1.7" + }, + "micromamba_jll": { + "uuid": "f8abcde7-e9b7-5caa-b8af-a437887ae8e4", + "version": "=2.3.1" + }, + "mpif_jll": { + "uuid": "9aeb927a-4695-514f-a259-621a69f20ec0", + "version": "=1.0.0" + }, + "pixi_jll": { + "uuid": "4d7b5844-a134-5dcd-ac86-c8f19cd51bed", + "version": "=0.76.2" + }, + "s2n_tls_jll": { + "uuid": "cddc5d3d-934d-5d3a-9747-62fc12ea3f48", + "version": "=1.7.10" } } } From 5a8b338a54c23f91eee8673abb3117f5a3e6ecc7 Mon Sep 17 00:00:00 2001 From: kylebd99 Date: Sat, 26 Sep 2026 11:18:30 -0400 Subject: [PATCH 17/20] update for pixi --- README.md | 45 +++++++++++++++++++++++++++++++++++---------- 1 file changed, 35 insertions(+), 10 deletions(-) diff --git a/README.md b/README.md index bd13d3b..34f9444 100644 --- a/README.md +++ b/README.md @@ -16,18 +16,35 @@ pip install galley-jl-python ## Contributing -### Packaging +### Development environment -Galley uses [poetry](https://python-poetry.org/) for packaging. +Galley uses [pixi](https://pixi.sh) for development. All configuration lives in +`pyproject.toml`: the package metadata and dependencies are under `[project]`, +and pixi's settings are under `[tool.pixi]`. -To install for development, clone the repository and run: +To set up, clone the repository and run: ```bash -poetry install --with test +pixi install # the default environment +pixi install -e test # adds the test dependencies (the `test` extra) ``` -to install the current project and dev dependencies. + +pixi installs the package in editable mode. Julia is not a pixi dependency: +[juliapkg](https://github.com/JuliaPy/pyjuliapkg) installs the pinned Julia +version and packages the first time `galley_jl_python` is imported. To trigger +that, and to fetch the sysimage (see below), run: +```bash +pixi run compile +``` + +Run any other command inside an environment with `pixi run`, for example +`pixi run -e test python`. + +The package is still built and published with Poetry (see +[Publishing](#publishing)). ### Working with a local copy of Finch.jl -The `develop.py ` script can be used to set up a local copy of Finch.jl for development. +The `develop.py` script can be used to set up a local copy of Finch.jl for +development. Run it with `pixi run python develop.py`. ``` Usage: @@ -93,16 +110,24 @@ On successful execution, the action publishes the package to PyPI and tags the r To add pre-commit hooks, run: ```bash -poetry run pre-commit install +pixi run -e test pre-commit install ``` ### Testing -Finch uses [pytest](https://docs.pytest.org/en/latest/) for testing. To run the +Galley uses [pytest](https://docs.pytest.org/en/latest/) for testing. To run the tests: ```bash -poetry run pytest +pixi run test +``` + +This runs `compile` first and then the full suite, including the Array API +tests described below. To run a subset, call pytest in the test environment +directly: + +```bash +pixi run -e test pytest tests/test_fused.py ``` Array API tests are included in `tests/test_array_api.py`. These tests invoke @@ -112,7 +137,7 @@ To forward `pytest` options to the nested `--array-api-pytest-args`): ```bash -poetry run pytest tests/test_array_api.py \ +pixi run -e test pytest tests/test_array_api.py \ --array-api="-k creation_functions" \ --array-api="-x" ``` From 3e4506cf713dd09ca10a334da3f03e6fcf86a60e Mon Sep 17 00:00:00 2001 From: kylebd99 Date: Sat, 26 Sep 2026 12:02:31 -0400 Subject: [PATCH 18/20] small fixes --- .github/workflows/ci.yml | 32 +++++++++----------------- pyproject.toml | 14 +++++++++++- src/galley_jl_python/finch_patches.jl | 33 +++++++++++++++++++++++++++ src/galley_jl_python/julia.py | 2 ++ 4 files changed, 59 insertions(+), 22 deletions(-) create mode 100644 src/galley_jl_python/finch_patches.jl diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 7f1e00f..fd59ab2 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -3,40 +3,30 @@ jobs: test: defaults: run: - shell: bash -el {0} + shell: bash strategy: matrix: os: [ubuntu-latest] - python: ['3.11', '3.12', '3.13'] + environment: [test-py311, test-py312, test-py313] include: - os: macos-latest - python: '3.11' + environment: test-py311 - os: windows-latest - python: '3.11' + environment: test-py311 fail-fast: false runs-on: ${{ matrix.os }} steps: - name: Checkout Repo - uses: actions/checkout@v3 - - uses: actions/setup-python@v5 + uses: actions/checkout@v7 + - uses: prefix-dev/setup-pixi@v0.10.0 with: - python-version: ${{ matrix.python }} - - name: Install Poetry - uses: snok/install-poetry@v1 - - name: Configure Poetry - run: | - if [[ "${{ runner.os }}" == "Windows" ]]; then - poetry env use python - else - poetry env use python${{ matrix.python }} - fi - - name: Install package - run: | - poetry install --extras test + pixi-version: v0.72.2 + environments: ${{ matrix.environment }} - name: Run tests run: | - poetry run pytest --junit-xml=test-${{ matrix.os }}-Python-${{ matrix.python }}.xml - - uses: codecov/codecov-action@v3 + pixi run -e ${{ matrix.environment }} test \ + --junit-xml=test-${{ matrix.os }}-${{ matrix.environment }}.xml + - uses: codecov/codecov-action@v7 on: # Trigger the workflow on push or pull request, diff --git a/pyproject.toml b/pyproject.toml index 1fd8224..d1e9eb5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -45,7 +45,7 @@ platforms = [ ] [tool.pixi.dependencies] -python = ">=3.11,<3.13" +python = ">=3.11,<3.14" [tool.pixi.pypi-dependencies] galley-jl-python = { path = ".", editable = true } @@ -59,8 +59,20 @@ fetch-sysimage = "python scripts/sysimage/fetch.py" [tool.pixi.feature.test.tasks] test = { cmd = "pytest", depends-on = ["compile"] } +[tool.pixi.feature.py311.dependencies] +python = "3.11.*" + +[tool.pixi.feature.py312.dependencies] +python = "3.12.*" + +[tool.pixi.feature.py313.dependencies] +python = "3.13.*" + [tool.pixi.environments] test = { features = ["test"], solve-group = "default" } +test-py311 = ["test", "py311"] +test-py312 = ["test", "py312"] +test-py313 = ["test", "py313"] [tool.ruff] line-length = 88 diff --git a/src/galley_jl_python/finch_patches.jl b/src/galley_jl_python/finch_patches.jl new file mode 100644 index 0000000..62986d7 --- /dev/null +++ b/src/galley_jl_python/finch_patches.jl @@ -0,0 +1,33 @@ +# Fixes for bugs in the pinned Finch.jl, evaluated into Finch.Galley at import. +# Delete each fix once the pinned Finch includes it. + +# Finch <= 1.6.0 builds `OrderedDict{DCKey}()`, which has no constructor, so +# Galley throws a MethodError when it re-checks a cached plan for tensors with +# 50 or more degree constraints (e.g. einsums over 4-d tensors). +# Upstream fix: https://github.com/finch-tensor/Finch.jl/pull/822 +function issimilar(stat1::DCStats, stat2::DCStats, rel_granularity) + if length(stat1.dcs) < 50 + for dc1 in stat1.dcs + for dc2 in stat2.dcs + if dc1.X == dc2.X && dc1.Y == dc2.Y && + abs(log(rel_granularity, dc1.d) - log(rel_granularity, dc2.d)) > 1 + return false + end + end + end + else + dc_dict = OrderedDict{DCKey,Float64}() + for dc1 in stat1.dcs + dc_dict[get_dc_key(dc1)] = dc1.d + end + for dc2 in stat2.dcs + # like the loop above, only compare DCs present in both stats + d1 = get(dc_dict, get_dc_key(dc2), nothing) + if !isnothing(d1) && + abs(log(rel_granularity, d1) - log(rel_granularity, dc2.d)) > 1 + return false + end + end + end + return true +end diff --git a/src/galley_jl_python/julia.py b/src/galley_jl_python/julia.py index b865900..e22b879 100644 --- a/src/galley_jl_python/julia.py +++ b/src/galley_jl_python/julia.py @@ -1,4 +1,5 @@ import juliapkg # noqa: I001, F401 +from pathlib import Path from . import _sysimage @@ -14,6 +15,7 @@ from juliacall import Main as jl # noqa: E402, F401 jl.seval("using Finch") +jl.Base.include(jl.Finch.Galley, str(Path(__file__).with_name("finch_patches.jl"))) jl.seval("using HDF5") jl.seval("using NPZ") jl.seval("using TensorMarket") From b047961498254350ca4f5deadac7e3b0d39b0965 Mon Sep 17 00:00:00 2001 From: kylebd99 Date: Sat, 26 Sep 2026 18:23:04 -0400 Subject: [PATCH 19/20] update finch version & parallelize tests --- .github/workflows/ci.yml | 5 ++-- README.md | 16 +++++++++---- pyproject.toml | 28 +++++++++++++++++++++-- scripts/sysimage/pin_julia_deps.py | 14 +++++++++++- src/galley_jl_python/finch_patches.jl | 33 --------------------------- src/galley_jl_python/julia.py | 5 ++-- src/galley_jl_python/juliapkg.json | 3 ++- tests/conftest.py | 6 +++++ 8 files changed, 63 insertions(+), 47 deletions(-) delete mode 100644 src/galley_jl_python/finch_patches.jl diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index fd59ab2..040baf2 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -23,9 +23,8 @@ jobs: pixi-version: v0.72.2 environments: ${{ matrix.environment }} - name: Run tests - run: | - pixi run -e ${{ matrix.environment }} test \ - --junit-xml=test-${{ matrix.os }}-${{ matrix.environment }}.xml + # junit and coverage reports go where pytest.ini says (junit/, coverage.xml) + run: pixi run -e ${{ matrix.environment }} test - uses: codecov/codecov-action@v7 on: diff --git a/README.md b/README.md index 34f9444..01af67a 100644 --- a/README.md +++ b/README.md @@ -119,12 +119,20 @@ Galley uses [pytest](https://docs.pytest.org/en/latest/) for testing. To run the tests: ```bash -pixi run test +pixi run test # one pytest-xdist worker per CPU +pixi run test 4 # or a fixed number of workers ``` -This runs `compile` first and then the full suite, including the Array API -tests described below. To run a subset, call pytest in the test environment -directly: +This runs `compile` first, then two suites one after the other, each spread +over the workers: + +- `pixi run test-unit` runs the unit tests. +- `pixi run test-array-api` runs the Array API tests described below. + +Each worker is a separate Julia process that uses 1–2 GB of memory, so lower +the worker count on machines with little memory. + +To run a subset, call pytest in the test environment directly: ```bash pixi run -e test pytest tests/test_fused.py diff --git a/pyproject.toml b/pyproject.toml index d1e9eb5..bce2d4c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -56,8 +56,32 @@ compile = { cmd = "python -c 'import galley_jl_python'", depends-on = ["fetch-sy build-sysimage = "python scripts/sysimage/build.py" fetch-sysimage = "python scripts/sysimage/fetch.py" -[tool.pixi.feature.test.tasks] -test = { cmd = "pytest", depends-on = ["compile"] } +# `pixi run test [njobs]` runs both suites, one after the other, each spread +# over `njobs` pytest-xdist workers (default: one per CPU). +[tool.pixi.feature.test.tasks.test] +args = [{ arg = "njobs", default = "auto" }] +# a no-op cmd, since `default-environment` is only allowed on tasks with one +cmd = "echo all test suites passed" +default-environment = "test" +depends-on = [ + { task = "test-unit", args = ["{{ njobs }}"] }, + { task = "test-array-api", args = ["{{ njobs }}"] }, +] + +[tool.pixi.feature.test.tasks.test-unit] +args = [{ arg = "njobs", default = "auto" }] +default-environment = "test" +cmd = "pytest -n {{ njobs }} --ignore=tests/test_array_api.py" +depends-on = ["compile"] + +# The Array API wrapper is a single test that runs array-api-tests in a nested +# pytest, so the workers go to the nested run. It skips coverage and writes its +# own junit file so it doesn't overwrite test-unit's reports. +[tool.pixi.feature.test.tasks.test-array-api] +args = [{ arg = "njobs", default = "auto" }] +default-environment = "test" +cmd = "pytest tests/test_array_api.py --no-cov --junitxml=junit/array-api-results.xml --array-api='-n {{ njobs }}'" +depends-on = ["compile"] [tool.pixi.feature.py311.dependencies] python = "3.11.*" diff --git a/scripts/sysimage/pin_julia_deps.py b/scripts/sysimage/pin_julia_deps.py index 94726b7..0a2f8d7 100644 --- a/scripts/sysimage/pin_julia_deps.py +++ b/scripts/sysimage/pin_julia_deps.py @@ -5,6 +5,9 @@ included) are pinned to exact versions. Everyone then resolves the same manifest, which the prebuilt sysimage requires. +Packages pinned to a git commit (``url`` and ``rev`` in juliapkg.json) are +left as they are; change their ``rev`` by hand. + To update: loosen the pins you want to move, resolve (e.g. ``pixi run compile``), run this script, then rebuild with ``pixi run build-sysimage``. """ @@ -35,8 +38,14 @@ def main() -> None: config = json.loads(JULIAPKG_JSON.read_text()) config["julia"] = f"={manifest['julia_version']}" + git_pinned = { + name: spec for name, spec in config["packages"].items() if "rev" in spec + } packages = {} for name, (entry, *_) in sorted(manifest["deps"].items()): + if name in git_pinned: + packages[name] = git_pinned[name] + continue # stdlibs have no tree hash and ship with Julia itself if "git-tree-sha1" not in entry or name in UNPINNED: continue @@ -47,7 +56,10 @@ def main() -> None: config["packages"] = packages JULIAPKG_JSON.write_text(json.dumps(config, indent=2) + "\n") - print(f"pinned julia {config['julia']} and {len(packages)} packages") + print( + f"pinned julia {config['julia']} and {len(packages)} packages " + f"({len(git_pinned)} to git commits: {', '.join(git_pinned) or 'none'})" + ) if __name__ == "__main__": diff --git a/src/galley_jl_python/finch_patches.jl b/src/galley_jl_python/finch_patches.jl deleted file mode 100644 index 62986d7..0000000 --- a/src/galley_jl_python/finch_patches.jl +++ /dev/null @@ -1,33 +0,0 @@ -# Fixes for bugs in the pinned Finch.jl, evaluated into Finch.Galley at import. -# Delete each fix once the pinned Finch includes it. - -# Finch <= 1.6.0 builds `OrderedDict{DCKey}()`, which has no constructor, so -# Galley throws a MethodError when it re-checks a cached plan for tensors with -# 50 or more degree constraints (e.g. einsums over 4-d tensors). -# Upstream fix: https://github.com/finch-tensor/Finch.jl/pull/822 -function issimilar(stat1::DCStats, stat2::DCStats, rel_granularity) - if length(stat1.dcs) < 50 - for dc1 in stat1.dcs - for dc2 in stat2.dcs - if dc1.X == dc2.X && dc1.Y == dc2.Y && - abs(log(rel_granularity, dc1.d) - log(rel_granularity, dc2.d)) > 1 - return false - end - end - end - else - dc_dict = OrderedDict{DCKey,Float64}() - for dc1 in stat1.dcs - dc_dict[get_dc_key(dc1)] = dc1.d - end - for dc2 in stat2.dcs - # like the loop above, only compare DCs present in both stats - d1 = get(dc_dict, get_dc_key(dc2), nothing) - if !isnothing(d1) && - abs(log(rel_granularity, d1) - log(rel_granularity, dc2.d)) > 1 - return false - end - end - end - return true -end diff --git a/src/galley_jl_python/julia.py b/src/galley_jl_python/julia.py index e22b879..7ac7934 100644 --- a/src/galley_jl_python/julia.py +++ b/src/galley_jl_python/julia.py @@ -1,5 +1,4 @@ import juliapkg # noqa: I001, F401 -from pathlib import Path from . import _sysimage @@ -10,12 +9,12 @@ # You can also run `develop.py` to quickly use a local copy of Finch.jl. # An example development json is found in `juliapkg_dev.json` # Julia and all packages are pinned so the prebuilt sysimage matches; update the -# pins with `scripts/sysimage/pin_julia_deps.py`. +# pins with `scripts/sysimage/pin_julia_deps.py`. Finch is pinned to a git +# commit (`rev`), which that script leaves alone. import juliacall as jc # noqa: E402, F401 from juliacall import Main as jl # noqa: E402, F401 jl.seval("using Finch") -jl.Base.include(jl.Finch.Galley, str(Path(__file__).with_name("finch_patches.jl"))) jl.seval("using HDF5") jl.seval("using NPZ") jl.seval("using TensorMarket") diff --git a/src/galley_jl_python/juliapkg.json b/src/galley_jl_python/juliapkg.json index 77dd492..70fdfea 100644 --- a/src/galley_jl_python/juliapkg.json +++ b/src/galley_jl_python/juliapkg.json @@ -83,7 +83,8 @@ }, "Finch": { "uuid": "9177782c-1635-4eb9-9bfb-d9dfa25e6bce", - "version": "=1.6.0" + "url": "https://github.com/finch-tensor/Finch.jl", + "rev": "368602490c2a106778afc9ca9d567a2201393012" }, "GPUArraysCore": { "uuid": "46192b85-c4d5-4398-a991-12ede77f4527", diff --git a/tests/conftest.py b/tests/conftest.py index 61800c9..14c22c7 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -2,6 +2,12 @@ import numpy as np +# Start Julia here, on the main thread, in every pytest process. Otherwise the +# pytest-xdist controller first imports galley when it rebuilds a +# PerformanceWarning sent by a worker, which starts Julia on execnet's receiver +# thread, and Julia's exit hook then deadlocks at interpreter shutdown. +import galley_jl_python # noqa: E402, F401 + def pytest_addoption(parser): parser.addoption( From 7b457b44492549b7d0081e0d5f652bcf9d7fcfd3 Mon Sep 17 00:00:00 2001 From: kylebd99 Date: Sat, 26 Sep 2026 18:37:47 -0400 Subject: [PATCH 20/20] add pre-commit + array api skips --- README.md | 11 +++++++++-- array-api-skips.txt | 6 ++++++ pyproject.toml | 16 +++++++++++++++- scripts/sysimage/pin_julia_deps.py | 2 +- src/galley_jl_python/juliapkg.json | 4 ++-- 5 files changed, 33 insertions(+), 6 deletions(-) diff --git a/README.md b/README.md index 01af67a..90e80e5 100644 --- a/README.md +++ b/README.md @@ -108,11 +108,18 @@ On successful execution, the action publishes the package to PyPI and tags the r ### Pre-commit hooks -To add pre-commit hooks, run: +The hooks run in their own `pre-commit` environment, which doesn't install +galley's dependencies or Julia. + ```bash -pixi run -e test pre-commit install +pixi run pre-commit-install # run the hooks on every `git commit` +pixi run pre-commit # run every hook on every file now +pixi run pre-commit ruff # the same, skipping the listed hook ids ``` +To run a single hook, use `pixi shell -e pre-commit` and then +`pre-commit run -a`. + ### Testing Galley uses [pytest](https://docs.pytest.org/en/latest/) for testing. To run the diff --git a/array-api-skips.txt b/array-api-skips.txt index ade08fe..29a70d4 100644 --- a/array-api-skips.txt +++ b/array-api-skips.txt @@ -6,6 +6,12 @@ array_api_tests/test_special_cases.py::test_unary[sign((x_i is -0 or x_i == +0)) array_api_tests/test_searching_functions.py::test_where # `test_solve` is not defined in Finch, hangs as xfail array_api_tests/test_linalg.py::test_solve +# Galley drops the array's shape for `x & False`, `x * 0` and other annihilating +# constants (https://github.com/finch-tensor/Finch.jl/issues/820), and binary +# functions don't accept a scalar as the first argument. Remove these once both +# are fixed. +array_api_tests/test_operators_and_elementwise_functions.py::test_binary_with_scalars_bool[logical_and] +array_api_tests/test_operators_and_elementwise_functions.py::test_binary_with_scalars_bitwise[bitwise_and] # test_signatures diff --git a/pyproject.toml b/pyproject.toml index bce2d4c..ad0cddb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -21,7 +21,6 @@ requires-poetry = ">=2.0,<3.0" [project.optional-dependencies] test = [ "pytest>=7.4.4,<7.6", - "pre-commit>=3.6.0,<3.9", "pytest-cov>=4.1.0,<4.2", "sparse>=0.17.0,<0.17.1", "scipy>=1.7,<1.17", @@ -83,6 +82,20 @@ default-environment = "test" cmd = "pytest tests/test_array_api.py --no-cov --junitxml=junit/array-api-results.xml --array-api='-n {{ njobs }}'" depends-on = ["compile"] +# Hooks run in their own environment, without galley's dependencies or Julia. +[tool.pixi.feature.pre-commit.dependencies] +pre-commit = "*" + +[tool.pixi.feature.pre-commit.tasks] +pre-commit-install = "pre-commit install" + +# `pixi run pre-commit [skip]` runs every hook on every file; `skip` is a +# comma-separated list of hook ids to leave out. +[tool.pixi.feature.pre-commit.tasks.pre-commit] +args = [{ arg = "skip", default = "" }] +env = { SKIP = "{{ skip }}" } +cmd = "pre-commit run -a" + [tool.pixi.feature.py311.dependencies] python = "3.11.*" @@ -97,6 +110,7 @@ test = { features = ["test"], solve-group = "default" } test-py311 = ["test", "py311"] test-py312 = ["test", "py312"] test-py313 = ["test", "py313"] +pre-commit = { features = ["pre-commit"], no-default-feature = true } [tool.ruff] line-length = 88 diff --git a/scripts/sysimage/pin_julia_deps.py b/scripts/sysimage/pin_julia_deps.py index 0a2f8d7..3a0c25d 100644 --- a/scripts/sysimage/pin_julia_deps.py +++ b/scripts/sysimage/pin_julia_deps.py @@ -55,7 +55,7 @@ def main() -> None: packages[name] = {"uuid": entry["uuid"], "version": f"={version}"} config["packages"] = packages - JULIAPKG_JSON.write_text(json.dumps(config, indent=2) + "\n") + JULIAPKG_JSON.write_text(json.dumps(config, indent=2, sort_keys=True) + "\n") print( f"pinned julia {config['julia']} and {len(packages)} packages " f"({len(git_pinned)} to git commits: {', '.join(git_pinned) or 'none'})" diff --git a/src/galley_jl_python/juliapkg.json b/src/galley_jl_python/juliapkg.json index 70fdfea..84c5be3 100644 --- a/src/galley_jl_python/juliapkg.json +++ b/src/galley_jl_python/juliapkg.json @@ -82,9 +82,9 @@ "version": "=1.17.1" }, "Finch": { - "uuid": "9177782c-1635-4eb9-9bfb-d9dfa25e6bce", + "rev": "368602490c2a106778afc9ca9d567a2201393012", "url": "https://github.com/finch-tensor/Finch.jl", - "rev": "368602490c2a106778afc9ca9d567a2201393012" + "uuid": "9177782c-1635-4eb9-9bfb-d9dfa25e6bce" }, "GPUArraysCore": { "uuid": "46192b85-c4d5-4398-a991-12ede77f4527",