Do a full iteration graph merge when addition terms have no contractions

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)

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)