Page MenuHomePhabricator

[mlir][sparse] auto-insertion of conversion to resolve cycles
ClosedPublic

Authored by aartbik on Jun 29 2022, 12:44 PM.

Details

Summary

When the iteration graph is cyclic (even after several attempts using less and less constraints), the current sparse compiler bails out, and no rewriting hapens. However, this revision adds some new logic where the sparse compiler tries to find a single input sparse tensor that breaks the cycle, and then adds a proper sparse conversion operation. This way, more incoming kernels can be handled!

Note, the resulting code is not optimal (although it keeps more or less proper "sparse" complexity), and more improvements should be added (especially when the kernel directly yields without computation, such as the transpose example). However, handling is better than not handling ;-)

Diff Detail

Event Timeline

aartbik created this revision.Jun 29 2022, 12:44 PM
aartbik requested review of this revision.Jun 29 2022, 12:44 PM
aartbik edited the summary of this revision. (Show Details)Jun 29 2022, 12:50 PM
aartbik updated this revision to Diff 441171.Jun 29 2022, 2:34 PM

rebased with main

bixia accepted this revision.Jun 29 2022, 4:25 PM
bixia added inline comments.
mlir/test/Dialect/SparseTensor/sparse_transpose.mlir
2

Maybe we shouldn't even need such a test as the generated code is so complicated and we already have an integration test to show that it works?

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_transpose.mlir
32–49

Would it be better to just replace this old function with the new one without the convert op?

This revision is now accepted and ready to land.Jun 29 2022, 4:25 PM
aartbik marked 2 inline comments as done.Jun 29 2022, 5:33 PM
aartbik added inline comments.
mlir/test/Dialect/SparseTensor/sparse_transpose.mlir
2

This was really to show the IR as part of the review too. Also, it shows that we should really address the TODO, since converting and then yielding is a bit strange. When we make that fix, this test will be triggered as being changed.

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_transpose.mlir
32–49

I wanted to contrast the old with the new, but I agree that the comment is now a bit out of date.

I have updated the comments to show the manual vs. auto insertion.

aartbik updated this revision to Diff 441218.Jun 29 2022, 5:55 PM
aartbik marked 2 inline comments as done.

addressed comments