This is an archive of the discontinued LLVM Phabricator instance.

[mlir][sparse] enable SDDMM-flavored fusion
ClosedPublic

Authored by aartbik on Aug 2 2022, 10:06 AM.

Details

Summary

This rewriting was no longer functional after recent migration to one shot
bufferization. However, this revision makes it work again, with a CHECK test
to ensure fusion happens. Note that functionality is tested by several
integration tests.

Diff Detail

Event Timeline

aartbik created this revision.Aug 2 2022, 10:06 AM
Herald added a project: Restricted Project. · View Herald TranscriptAug 2 2022, 10:06 AM
aartbik requested review of this revision.Aug 2 2022, 10:06 AM
Peiming added inline comments.Aug 2 2022, 10:36 AM
mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp
149

Just out of curiosity, why does the producer has to be zero?

Even if it is non-zero
T(i,j) = SUM(k, A(i,j,k) * B(i,j,k) * ... )
can be transformed to
T0 = alloc_tensor<0>; T0(i,j) = SUM(T(i, j), SUM(k, A(i,j,k) * B(i,j,k) * ... ))

aartbik marked an inline comment as done.Aug 2 2022, 11:54 AM
aartbik added inline comments.
mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp
149

I am perhaps a bit overly conservative here, but for all bufferization patterns generated, zero init will probably be the norm. We can relax later if we find we are missing real-world cases.

Peiming accepted this revision.Aug 2 2022, 12:07 PM
This revision is now accepted and ready to land.Aug 2 2022, 12:07 PM
aartbik updated this revision to Diff 449391.Aug 2 2022, 12:39 PM
aartbik marked an inline comment as done.

rebased against main

This revision was landed with ongoing or failed builds.Aug 2 2022, 12:40 PM
This revision was automatically updated to reflect the committed changes.