[mlir][sparse] Add a new reduce operation in sparse_tensor dialect

[jim22k]

Summary

Custom reduction is a binary operators, but differs from the existing
sparse_tensor.binary operation in significant ways.

1. It does not deal with sparsity overlaps. It is given a finite set of non-empty values and must collapse them to a single value
2. The nature of reduction requires a starting value (or identity value). This is normally zero for summation, but for the product needs to be 1.0. The custom code needs to explicitly pass an identity value.

[aartbik]

mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td

554

The documentation can use a bit more explanation on why we added this.

For example, you can show that "regular" reductions can be expressed "the easy way", but this makes the sparse specifications more expressive (I want to avoid that newcomers think *all* reductions need to be defined this way)

560

empty line after Example:

mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp

378–379

since we don't use the result of the cast, why not simply use a "isa"?

mlir/lib/Dialect/SparseTensor/Utils/Merger.cpp

959 ↗

(On Diff #437689)

can't we start a new == 3 section then for ternary operations?

mlir/test/Dialect/SparseTensor/invalid.mlir

344

looks like we never had an invalid test for "must end with sparse_tensor.yield", for this or another one, so perhaps add this as part of this revision to get full coverage

[jim22k]

• Merge with latest changes (D128000)
• Update based on feedback

[jim22k]

Fix formatting

[aartbik]

Perhaps for the sake of progress, can we split this into two revisions? Once that introduces the new sparse dialect op, and then one that provides the right implementation?
I feel that the former is almost ready to go in, but the latter needs some work, so splitting this up into smaller pieces may make it easier.

mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td

570

This example is missing the surrounding linalg.generic, making it also harder to see what %x and %y are exactly. Also, in the examples above you use %a,%b,%c for the top level, and then %arg0 again in the block inside the semi-ring op. Please make the syntax in all examples similar, so it is easier to read and recognize common components.

mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp

745 ↗

(On Diff #440630)

I am puzzled we need to have new code in the merger for this, can't we just handle that in the sparse codegen file as before. Also, please note that we need to deal with vector values during vectorization. Probably not for the semi ring op, but your ReducID seems to deal with the other cases too.

Since you are using so many if (codegen.redKind == kNoReduc) tests, it feels like you are entering the wrong methods, and should have deal with reductions at a higher level (see updateReduc() and friends)

[jim22k]

Add reduce operation definition and validation tests

[jim22k]

I split the PR in two. This is now just the .td definition and validation tests. Hopefully this makes it easier to approve, and we can discuss implementation in a subsequent PR.

[aartbik]

In D128004#3652267, @jim22k wrote:

I split the PR in two. This is now just the .td definition and validation tests. Hopefully this makes it easier to approve, and we can discuss implementation in a subsequent PR.

Thanks for splitting this up! This part is directly good to go, and we can start working on the implementation with a smaller revision!