From b676697f48035d83ca4d36da476763e81c02d717 Mon Sep 17 00:00:00 2001 From: David Hagen Date: Thu, 16 Jul 2026 12:44:51 -0400 Subject: [PATCH] Do a full iteration graph merge when addition terms have no contractions Stop trying to distribute multiplication over a SumNode when generating iteration graphs. --- src/tensora/desugar/_to_iteration_graphs.py | 54 ++++++++++-- tests/test_addition_in_multiplication.py | 91 +++++++++++++++++++++ 2 files changed, 139 insertions(+), 6 deletions(-) create mode 100644 tests/test_addition_in_multiplication.py diff --git a/src/tensora/desugar/_to_iteration_graphs.py b/src/tensora/desugar/_to_iteration_graphs.py index 839b9ae..a837429 100644 --- a/src/tensora/desugar/_to_iteration_graphs.py +++ b/src/tensora/desugar/_to_iteration_graphs.py @@ -96,6 +96,42 @@ def to_iteration_graphs_tensor( yield graph +def contains_contraction(self: ast.Expression) -> bool: + """Whether an expression contains a contraction.""" + match self: + case ast.Contract(): + return True + case ast.Add() | ast.Multiply(): + return contains_contraction(self.left) or contains_contraction(self.right) + case _: + return False + + +def merge_add(left: ig.IterationGraph, right: ig.IterationGraph) -> Iterator[ig.IterationGraph]: + """Produce all possible ways to co-iterate two iteration graphs under addition.""" + match (left, right): + case (ig.TerminalNode(), ig.TerminalNode()): + yield ig.TerminalNode(id.Add(left.expression, right.expression)) + case (ig.IterationNode(), ig.TerminalNode()): + for tail in merge_add(left.next, right): + yield replace(left, next=tail) + case (ig.TerminalNode(), ig.IterationNode()): + for tail in merge_add(left, right.next): + yield replace(right, next=tail) + case (ig.IterationNode(), ig.IterationNode()): + if left.index_variable == right.index_variable: + for tail in merge_add(left.next, right.next): + yield replace(left, next=tail) + else: + if left.index_variable not in right.next.later_indexes(): + for tail in merge_add(left.next, right): + yield replace(left, next=tail) + + if right.index_variable not in left.next.later_indexes(): + for tail in merge_add(left, right.next): + yield replace(right, next=tail) + + def simplify_add(graph: ig.SumNode) -> ig.IterationGraph: """Simplify an Add by combining terms with the same index variable. @@ -148,6 +184,13 @@ def to_iteration_graphs_add( formats: dict[str, Format], counter: Iterator[int], ) -> Iterator[ig.IterationGraph]: + if not (contains_contraction(self.left) or contains_contraction(self.right)): + # A contraction-free sum can be merged into a single iteration node. + for left in to_iteration_graphs_expression(self.left, formats, counter): + for right in to_iteration_graphs_expression(self.right, formats, counter): + yield from merge_add(left, right) + return + name = f"sum_{next(counter)}" for left in to_iteration_graphs_expression(self.left, formats, counter): for right in to_iteration_graphs_expression(self.right, formats, counter): @@ -189,12 +232,11 @@ def merge_multiply( if right.index_variable not in left.next.later_indexes(): for tail in merge_multiply(left, right.next): yield replace(right, next=tail) - case (ig.SumNode(), _): - for terms in product(*[merge_multiply(term, right) for term in left.terms]): - yield ig.SumNode(left.name, list(terms)) - case (_, ig.SumNode()): - for terms in product(*[merge_multiply(left, term) for term in right.terms]): - yield ig.SumNode(right.name, list(terms)) + case (ig.SumNode(), _) | (_, ig.SumNode()): + # A surviving SumNode contains a contraction (contraction-free sums are merged into + # iteration nodes or terminal nodes by `to_iteration_graphs_add`). Multiplying it would + # require a workspace, which is not yet implemented. + return @to_iteration_graphs_expression.register(ast.Multiply) diff --git a/tests/test_addition_in_multiplication.py b/tests/test_addition_in_multiplication.py new file mode 100644 index 0000000..346e774 --- /dev/null +++ b/tests/test_addition_in_multiplication.py @@ -0,0 +1,91 @@ +import pytest + +from tensora.compile import TensorMethod +from tensora.desugar import NoKernelFoundError +from tensora.expression import parse_assignment +from tensora.format import parse_format +from tensora.problem import make_problem + +from .assert_same_as_dense import assert_same_as_dense + + +@pytest.mark.parametrize("dense_b", [[0, 2, 0, 3], [0, 0, 0, 0]]) +@pytest.mark.parametrize("dense_c", [[5, 0, 0, 7], [0, 0, 0, 0]]) +@pytest.mark.parametrize("format_b", ["s", "d"]) +@pytest.mark.parametrize("format_c", ["s", "d"]) +@pytest.mark.parametrize("format_out", ["s", "d"]) +def test_broadcast_scalar_in_product(dense_b, dense_c, format_b, format_c, format_out): + assert_same_as_dense( + "a(i) = b(i) * (c(i) + 1)", + format_out, + b=(dense_b, format_b), + c=(dense_c, format_c), + ) + + +@pytest.mark.parametrize("format_b", ["dd", "ds", "sd", "ss"]) +@pytest.mark.parametrize("format_c", ["s", "d"]) +@pytest.mark.parametrize("format_d", ["s", "d"]) +@pytest.mark.parametrize("format_out", ["dd", "ds", "sd", "ss"]) +def test_different_free_indexes_in_product(format_b, format_c, format_d, format_out): + # c(i) and d(j) broadcast over different free indexes; folding must co-iterate them without + # duplicating b. + assert_same_as_dense( + "a(i,j) = b(i,j) * (c(i) + d(j))", + format_out, + b=([[1, 2], [3, 4]], format_b), + c=([10, 20], format_c), + d=([100, 200], format_d), + ) + + +@pytest.mark.parametrize("format_b", ["s", "d"]) +@pytest.mark.parametrize("format_c", ["dd", "ds", "sd", "ss"]) +@pytest.mark.parametrize("format_out", ["dd", "ds", "sd", "ss"]) +def test_nested_broadcast_scalar_in_product(format_b, format_c, format_out): + assert_same_as_dense( + "a(i,j) = b(i) * (c(i,j) + 1)", + format_out, + b=([1, 2], format_b), + c=([[3, 0], [0, 4]], format_c), + ) + + +@pytest.mark.parametrize("format_c", ["s", "d"]) +@pytest.mark.parametrize("format_d", ["s", "d"]) +@pytest.mark.parametrize("format_out", ["s", "d"]) +def test_product_of_sums(format_c, format_d, format_out): + assert_same_as_dense( + "a(i) = (c(i) + 1) * (d(i) + 1)", + format_out, + c=([4, 0, 6], format_c), + d=([1, 2, 0], format_d), + ) + + +@pytest.mark.parametrize("format_b", ["s", "d"]) +@pytest.mark.parametrize("format_c", ["dd", "ds", "sd", "ss"]) +@pytest.mark.parametrize("format_d", ["s", "d"]) +@pytest.mark.parametrize("format_out", ["dd", "ds", "sd", "ss"]) +def test_three_term_sum_in_product(format_b, format_c, format_d, format_out): + assert_same_as_dense( + "a(i,j) = b(i) * (c(i,j) + d(j) + 1)", + format_out, + b=([1, 2], format_b), + c=([[3, 0], [0, 4]], format_c), + d=([5, 6], format_d), + ) + + +def test_contraction_in_product_is_not_yet_supported(): + # A contraction (`j`) inside the parentheses cannot be merged into a single expression because + # a reduction is a loop, not a term. Until workspaces are implemented, no kernel exists for it. + assignment = parse_assignment("a(i) = b(i) * (c(i) + d(j,i))").unwrap() + formats = { + name: parse_format(fmt).unwrap() + for name, fmt in [("a", "d"), ("b", "d"), ("c", "d"), ("d", "dd")] + } + problem = make_problem(assignment, formats).unwrap() + + with pytest.raises(NoKernelFoundError): + TensorMethod(problem)