Repository: drhagen/tensora

Issue:
Don't emit computation when doing assembly
Right now, assembly does the exact same thing as evaluate. It should not emit any of the computation.

Discussion:
codecov-commenter: ## [Codecov](https://app.codecov.io/gh/drhagen/tensora/pull/27?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=David+Hagen) Report
Attention: `2 lines` in your changes are missing coverage. Please review.
> Comparison is base [(`f2d3098`)](https://app.codecov.io/gh/drhagen/tensora/commit/f2d30985c01a50a013543deccbcb148a3620f7b8?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=David+Hagen) 88.59% compared to head [(`c97a736`)](https://app.codecov.io/gh/drhagen/tensora/pull/27?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=David+Hagen) 88.55%.

| [Files](https://app.codecov.io/gh/drhagen/tensora/pull/27?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=David+Hagen) | Patch % | Lines |
|---|---|---|
| [src/tensora/iteration\_graph/generate\_ir.py](https://app.codecov.io/gh/drhagen/tensora/pull/27?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=David+Hagen#diff-c3JjL3RlbnNvcmEvaXRlcmF0aW9uX2dyYXBoL2dlbmVyYXRlX2lyLnB5) | 84.61% | [0 Missing and 2 partials :warning: ](https://app.codecov.io/gh/drhagen/tensora/pull/27?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=David+Hagen) |

:exclamation: Your organization needs to install the [Codecov GitHub app](https://github.com/apps/codecov/installations/select_target) to enable full functionality.

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##           master      #27      +/-   ##
==========================================
- Coverage   88.59%   88.55%   -0.04%     
==========================================
  Files          50       50              
  Lines        2911     2919       +8     
  Branches      516      519       +3     
==========================================
+ Hits         2579     2585       +6     
  Misses        252      252              
- Partials       80       82       +2     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/drhagen/tensora/pull/27?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=David+Hagen).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=David+Hagen).


Patch:
diff --git a/README.md b/README.md
index c7fda23..b13fd6d 100644
--- a/README.md
+++ b/README.md
@@ -166,32 +166,36 @@ tensora 'y(i) = A(i,j) * x(j)' -f A:ds -t compute -o kernel.c
 Here is the output of `tensora --help` for reference:
 
 ```
- Usage: tensora [OPTIONS] ASSIGNMENT                                                                     
-                                                                                                         
-╭─ Arguments ───────────────────────────────────────────────────────────────────────────────────────────╮
-│ *    assignment      TEXT  The assignment for which to generate code, e.g. y(i) = A(i,j) * x(j).      │
-│                            [required]                                                                 │
-╰───────────────────────────────────────────────────────────────────────────────────────────────────────╯
-╭─ Options ─────────────────────────────────────────────────────────────────────────────────────────────╮
-│ --format              -f      TEXT                             A tensor and its format separated by a │
-│                                                                colon, e.g. A:d1s0 for CSC matrix.     │
-│                                                                Unmentioned tensors are be assumed to  │
-│                                                                be all dense.                          │
-│ --type                -t      [assembly|compute|evaluate]      The type of kernel that will be        │
-│                                                                generated. Can be mentioned multiple   │
-│                                                                times.                                 │
-│                                                                [default: KernelType.compute]          │
-│ --output              -o      PATH                             The file to which the kernel will be   │
-│                                                                written. If not specified, prints to   │
-│                                                                standard out.                          │
-│                                                                [default: None]                        │
-│ --install-completion          [bash|zsh|fish|powershell|pwsh]  Install completion for the specified   │
-│                                                                shell.                                 │
-│                                                                [default: None]                        │
-│ --show-completion             [bash|zsh|fish|powershell|pwsh]  Show completion for the specified      │
-│                                                                shell, to copy it or customize the     │
-│                                                                installation.                          │
-│                                                                [default: None]                        │
-│ --help                                                         Show this message and exit.            │
-╰───────────────────────────────────────────────────────────────────────────────────────────────────────╯
-``````
\ No newline at end of file
+ Usage: tensora [OPTIONS] ASSIGNMENT
+
+╭─ Arguments ──────────────────────────────────────────────────────────────────────────────╮
+│ *    assignment      TEXT  The assignment for which to generate code, e.g. y(i) = A(i,j) │
+│                            * x(j).                                                       │
+│                            [required]                                                    │
+╰──────────────────────────────────────────────────────────────────────────────────────────╯
+╭─ Options ────────────────────────────────────────────────────────────────────────────────╮
+│ --format              -f      TEXT                         A tensor and its format       │
+│                                                            separated by a colon, e.g.    │
+│                                                            A:d1s0 for CSC matrix.        │
+│                                                            Unmentioned tensors are be    │
+│                                                            assumed to be all dense.      │
+│ --type                -t      [assemble|compute|evaluate]  The type of kernel that will  │
+│                                                            be generated. Can be          │
+│                                                            mentioned multiple times.     │
+│                                                            [default: compute]            │
+│ --output              -o      PATH                         The file to which the kernel  │
+│                                                            will be written. If not       │
+│                                                            specified, prints to standard │
+│                                                            out.                          │
+│                                                            [default: None]               │
+│ --install-completion          [bash|zsh|fish|powershell|p  Install completion for the    │
+│                               wsh]                         specified shell.              │
+│                                                            [default: None]               │
+│ --show-completion             [bash|zsh|fish|powershell|p  Show completion for the       │
+│                               wsh]                         specified shell, to copy it   │
+│                                                            or customize the              │
+│                                                            installation.                 │
+│                                                            [default: None]               │
+│ --help                                                     Show this message and exit.   │
+╰──────────────────────────────────────────────────────────────────────────────────────────╯
+```
\ No newline at end of file
diff --git a/src/tensora/iteration_graph/generate_ir.py b/src/tensora/iteration_graph/generate_ir.py
index 1ec3055..c071dce 100644
--- a/src/tensora/iteration_graph/generate_ir.py
+++ b/src/tensora/iteration_graph/generate_ir.py
@@ -24,8 +24,7 @@ from ..kernel_type import KernelType
 from .definition import Definition
 from .identifiable_expression import to_ir
 from .identifiable_expression.tensor_layer import TensorLayer
-from .iteration_graph import IterationGraph, IterationVariable, TerminalExpression
-from .iteration_graph import Sum as GraphAdd
+from .iteration_graph import IterationGraph, IterationVariable, Sum, TerminalExpression
 from .names import crd_name, dimension_name, pos_name, vals_name
 from .outputs import AppendOutput, Output
 from .write_sparse_ir import (
@@ -49,7 +48,8 @@ def to_ir_iteration_graph(
 def to_ir_terminal_expression(self: TerminalExpression, output: Output, kernel_type: KernelType):
     source = SourceBuilder("*** Computation of expression ***")
 
-    source.append(output.write_assignment(to_ir(self.expression), kernel_type))
+    if kernel_type.is_compute():
+        source.append(output.write_assignment(to_ir(self.expression), kernel_type))
 
     return source
 
@@ -86,6 +86,10 @@ def generate_subgraphs(graph: IterationVariable) -> list[IterationVariable]:
 def to_ir_iteration_variable(self: IterationVariable, output: Output, kernel_type: KernelType):
     source = SourceBuilder(f"*** Iteration over {self.index_variable} ***")
 
+    if not kernel_type.is_compute() and not self.has_assemble():
+        # If not doing compute and no assembly is left, return early
+        return source
+
     loop_variable = Variable(self.index_variable)
 
     # If this node is sparse, then we can use efficient sparse iteration. Sparse graphs produce
@@ -321,17 +325,19 @@ def to_ir_iteration_variable(self: IterationVariable, output: Output, kernel_typ
     return source
 
 
-@to_ir_iteration_graph.register(GraphAdd)
-def to_ir_add(self: GraphAdd, output: Output, kernel_type: KernelType):
-    source = SourceBuilder("*** Add ***")
+@to_ir_iteration_graph.register(Sum)
+def to_ir_sum(self: Sum, output: Output, kernel_type: KernelType):
+    source = SourceBuilder("*** Sum ***")
 
-    next_output, next_output_declarations, next_output_cleanup = output.next_output(
-        None, kernel_type
-    )
-    source.append(next_output_declarations)
+    if kernel_type.is_compute():
+        # No assembly is currently allowed downstream of a Sum node
+        next_output, next_output_declarations, next_output_cleanup = output.next_output(
+            None, kernel_type
+        )
+        source.append(next_output_declarations)
 
-    for term in self.terms:
-        source.append(to_ir_iteration_graph(term, next_output, kernel_type))
+        for term in self.terms:
+            source.append(to_ir_iteration_graph(term, next_output, kernel_type))
 
     return source
 
diff --git a/src/tensora/iteration_graph/identifiable_expression/extract_context.py b/src/tensora/iteration_graph/identifiable_expression/extract_context.py
index 787a9e5..652811b 100644
--- a/src/tensora/iteration_graph/identifiable_expression/extract_context.py
+++ b/src/tensora/iteration_graph/identifiable_expression/extract_context.py
@@ -17,6 +17,7 @@ class Context:
     dense_leaves: list[TensorLayer] = field(default_factory=list)
     indexes: frozenset[str] = frozenset()
     has_output: bool = False
+    has_assemble: bool = False
 
     def add(self, other: Context) -> Context:
         return Context(
@@ -25,6 +26,7 @@ class Context:
             dense_leaves=self.dense_leaves + other.dense_leaves,
             indexes=self.indexes | other.indexes,
             has_output=self.has_output or other.has_output,
+            has_assemble=self.has_assemble or other.has_assemble,
         )
 
     def multiply(self, other: Context) -> Context:
@@ -34,6 +36,7 @@ class Context:
             dense_leaves=self.dense_leaves + other.dense_leaves,
             indexes=self.indexes | other.indexes,
             has_output=self.has_output or other.has_output,
+            has_assemble=self.has_assemble or other.has_assemble,
         )
 
 
diff --git a/src/tensora/iteration_graph/iteration_graph.py b/src/tensora/iteration_graph/iteration_graph.py
index 64371a0..70fd59a 100644
--- a/src/tensora/iteration_graph/iteration_graph.py
+++ b/src/tensora/iteration_graph/iteration_graph.py
@@ -85,6 +85,7 @@ class IterationVariable(IterationGraph):
             next_context,
             indexes=next_context.indexes | frozenset([self.index_variable]),
             has_output=next_context.has_output or self.output is not None,
+            has_assemble=next_context.has_assemble or self.is_sparse_output(),
         )
 
     def exhaust_tensor(self, tensor: TensorLeaf) -> IterationGraph:
@@ -116,6 +117,9 @@ class IterationVariable(IterationGraph):
     def has_output(self) -> bool:
         return self.context.has_output
 
+    def has_assemble(self) -> bool:
+        return self.context.has_assemble
+
 
 @dataclass(frozen=True)
 class Sum(IterationGraph):
diff --git a/src/tensora/iteration_graph/outputs/append.py b/src/tensora/iteration_graph/outputs/append.py
index 811baa6..766cd29 100644
--- a/src/tensora/iteration_graph/outputs/append.py
+++ b/src/tensora/iteration_graph/outputs/append.py
@@ -50,7 +50,7 @@ class AppendOutput(Output):
                     if all_dense:
                         # If the previous dimensions were all dense, then the size of pos in this dimension is fixed
                         pos_size = Multiply.join(
-                            [output_tensor.attr("dimensions").idx(i_prev) for i_prev in range(i)]
+                            [dimension_name(self.output.indexes[i_prev]) for i_prev in range(i)]
                         ).plus(1)
                     else:
                         # Otherwise, the value will change, so provide a good default
diff --git a/src/tensora/kernel_type.py b/src/tensora/kernel_type.py
index fa931f6..4c3a803 100644
--- a/src/tensora/kernel_type.py
+++ b/src/tensora/kernel_type.py
@@ -5,9 +5,15 @@ from enum import Enum
 
 class KernelType(str, Enum):
     # Python 3.10 does not support StrEnum, so do it manually
-    assembly = "assembly"
+    assemble = "assemble"
     compute = "compute"
     evaluate = "evaluate"
 
     def is_assembly(self):
-        return self == KernelType.assembly or self == KernelType.evaluate
+        return self == KernelType.assemble or self == KernelType.evaluate
+
+    def is_compute(self):
+        return self == KernelType.compute or self == KernelType.evaluate
+
+    def __str__(self) -> str:
+        return self.name


Test patch:
diff --git a/tests/test_cli.py b/tests/test_cli.py
index 2bbb460..99fae0f 100644
--- a/tests/test_cli.py
+++ b/tests/test_cli.py
@@ -25,13 +25,13 @@ def test_cli():
 
 def test_multiple_kernels():
     result = runner.invoke(
-        app, ["y(i) = A(i,j) * x(j)", "-t", "compute", "-t", "evaluate", "-t", "assembly"]
+        app, ["y(i) = A(i,j) * x(j)", "-t", "compute", "-t", "evaluate", "-t", "assemble"]
     )
 
     assert result.exit_code == 0
     assert "compute" in result.stdout
     assert "evaluate" in result.stdout
-    assert "assembly" in result.stdout
+    assert "assemble" in result.stdout
 
 
 def test_write_to_file():