[Mlir-commits] [mlir] [mlir][sparse] Reuse input sparsity for materialized outputs (PR #215579)

[Kohei Yamaguchi]

sott0n wrote:

@aartbik Thanks for checking it. The intended abstraction is output-support containment rather than an SDDMM-specific rewrite.

## Problem

The sparsifier currently treats a sparse output initialized by `tensor.empty` as having an unknown structure. It therefore constructs the output coordinates dynamically through insertion or `expand`/`compress`.

For some expressions, however, the result can only be nonzero at coordinates stored by one of the sparse inputs. Value-weighted SDDMM is one example:

```
    O(i, j) = S(i, j) * sum_k A(i, k) * B(k, j)
```

Since `S(i, j) = 0` implies `O(i, j) = 0`, the output support is contained in the support of `S`.

## Proposed analysis

For each sparse input with the same tensor type and indexing map as the output, the tensor expression is analyzed to determine whether:

```
    support(output) ⊆ stored_coordinates(input)
```

The analysis follows the expression tree:

- multiplication/intersection requires the condition on either operand;
- addition/union requires it on both operands;
- zero-preserving unary operations propagate the condition;
- unsupported or ambiguous operations conservatively fail.

This applies to other operations such as sparse-dense elementwise multiplication and sparse-sparse intersection, not only SDDMM.

## Lowering

When containment is proven, the output is materialized with a copy of the input's coordinate structure and a separate zero-initialized values buffer. The existing sparsifier and loop emitter then compute the output values over that structure.

Coordinates whose computed value is zero may remain explicitly stored. This is consistent with using the input's stored coordinates as an upper bound rather than requiring exact equality of mathematical nonzero supports.

This avoids dynamic sparse-output construction and allows the existing parallel-loop lowering to be used.

## Initial scope

The expression analysis is format-independent, but the initial materialization is limited to static rank-2 identity, ordered, unique CSR with functional emission. Supporting additional formats requires generalizing the format-dependent structure-copying step; it does not require changing the support-containment rule.

https://github.com/llvm/llvm-project/pull/215579