This is an archive of the discontinued LLVM Phabricator instance.

[mlir][sparse] Support sparse2sparse collapse for dynamic sizes
ClosedPublic

Authored by anlunx on Aug 10 2022, 10:41 AM.

Details

Summary

This patch implements sparse2sparse collapse for operands with dynamic shape.

Diff Detail

Event Timeline

anlunx created this revision.Aug 10 2022, 10:41 AM
anlunx requested review of this revision.Aug 10 2022, 10:41 AM

please use tags in the title, our team usually usesSparse reshaping for dynamic sizes

[mlir][sparse] sparse reshaping for dynamic sizes

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp
457

we already have hasStaticShape() as a utility

The "Index is too large for the dimension"' failure indicates that you are not passing the right (runtime) sizes for the COO data structure. I suspect you pass the -1 of the dynamic size to the kEmptyCOO call.

anlunx updated this revision to Diff 451628.Aug 10 2022, 2:21 PM

use hasStaticShape

anlunx retitled this revision from Sparse reshaping for dynamic sizes to [mlir][sparse] sparse reshaping for dynamic sizes.Aug 10 2022, 2:22 PM
anlunx marked an inline comment as done.
anlunx added inline comments.
mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp
519–527

The "Index is too large for the dimension"' failure indicates that you are not passing the right (runtime) sizes for the COO data structure. I suspect you pass the -1 of the dynamic size to the kEmptyCOO call.

I'm using sizesFromPtr to compute the size of the destination tensor. Is it a valid thing to do? I was thinking that it should compute the runtime size because it generates the "sparseDimSize" function.

aartbik added inline comments.Aug 11 2022, 1:17 PM
mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp
459

This is really very similar to translateIndices, and contains way too much codedup. Can we merge the two and deal with static/dynamic in the logic?

512

llvm style does not want underscores in local vars

519–527

The problem is that you are using

sizesFromPtr(rewriter, dst_sizes, op, encDst, dstTp, src);

to fill the dimensions of the *destination*, but are you are querying the *src* to get the dimension size dynamically!
That way, you set the size of the destination to tensor<1xf64> and the second insertion goes out of bounds.

The problem, of course, is that you have no dst operand to query the sizes from ;-)
You will have to compute the destination format dynamically, using the dynamic sizes of the src to build it.
In this case you will need a computation that does dst.size[0] = src.size[0] * src.size[1]
for lower to higher, the computation will be more elaborate!

anlunx updated this revision to Diff 452074.Aug 11 2022, 9:41 PM

Add function for generate collapse destination shape
Merge translateIndices with translateIndicesDyn

anlunx updated this revision to Diff 452075.Aug 11 2022, 9:49 PM
anlunx marked an inline comment as done.

Rename genDstShape to genReshapeDstShape

anlunx retitled this revision from [mlir][sparse] sparse reshaping for dynamic sizes to [mlir][sparse] Support sparse2sparse collapse for dynamic sizes.Aug 11 2022, 9:55 PM
anlunx edited the summary of this revision. (Show Details)
anlunx marked 2 inline comments as done.
anlunx added inline comments.
mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp
459

This is done. Thanks for the suggestion!

519–527

Thank you for the explanation! Added a function genReshapeDstShape to compute the dynamic size of the destination.

anlunx updated this revision to Diff 452244.Aug 12 2022, 11:29 AM
anlunx marked an inline comment as done.

Fix style using clang-format

aartbik added inline comments.Aug 15 2022, 9:46 AM
mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp
463

Can we make expand of this revision too?

If not, then don't assert, but simply return failure() in the rewriter, i.e. pattern will not kick in.

Peiming added inline comments.Aug 16 2022, 2:52 PM
mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp
428

Is it more easy to follow to always generate the MLIR operation (same below) regardless of whether the size is statically know? We can rely on constant folding to eliminate unnecessary runtime computation.

I am not sure whether it is worth it though :-)

anlunx updated this revision to Diff 453881.Aug 18 2022, 10:07 PM

Implement expand

anlunx marked 2 inline comments as done.Aug 18 2022, 10:14 PM
anlunx added inline comments.
mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp
428

Good point. Removed code for static shapes.

463

Added expand

anlunx updated this revision to Diff 456805.Aug 30 2022, 3:23 PM
anlunx marked 2 inline comments as done.

Fix sparse_reshape round-trip test

apologies for the long review time this time around....

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp
391–392

you have to use ArrayRef<> in the parameter; SmallVector<> is only used at the concrete implementation side (and as written, you copy contents into the call!)

399–400

same, when assigning, only work on references

403

perhaps good to keep this assert since you use the literal

428

Can you check that. In the CHECK test below, we just use the unrelated --cse as extra parameter, and you see that constants are not folded. It is okay to add more flags there, but I would like to see evidence that the code folds indeed.

460

document the "output" parameter, use ArrayRef for all others

470

assert on not dynamic?

482

I typically use the idiom

for (unsigned i = 0, size = srcShape.size(); i < size; i++)

to avoid the potential call in loop (depending on how smart compiler is)

490

Is this skip over intended? Your code can definitely use some examples for each case

mlir/test/Dialect/SparseTensor/sparse_reshape.mlir
27 ↗(On Diff #456805)

see above, okay to add an extra flag, but I would like to keep what is on the left

71 ↗(On Diff #456805)

as here

93 ↗(On Diff #456805)

why did you not add any CHECK tests for the dynamic case

you only rely on integration test now, but I would like to see a complex example generating the right code.

Peiming added a comment.EditedSep 8 2022, 1:59 PM
This comment has been deleted.
mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp
520–521

srcTp or dstTp here?

anlunx updated this revision to Diff 460567.Sep 15 2022, 5:26 PM
anlunx marked 9 inline comments as done.

Address comments

anlunx added inline comments.Sep 15 2022, 5:30 PM
mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp
403

I think I'm not using compile-time constants here. The shape vector here is computed at runtime, so we don't need this assert anymore.

428

Added the --canonicalize flag which folds constants.

470

IIUC srcShape is available no matter if the source shape is dynamic or static? It is computed by calling sizesFromPtr

490

Yes, it is intended. For example, if the srcShape is [8], and the staticDstShape is [2 x ? x 2], then we compute the unknown dimension by dividing 8 by 2 x 2. Therefore we skip the unknown dimension when computing the denominator. Added detailed comments for the algorithm.

520–521

I think they are equivalent since dstTp has static shape if and only if srcTp has static shape. But I agree that dstTp makes more sense here. Changed it to dstTp.

mlir/test/Dialect/SparseTensor/sparse_reshape.mlir
93 ↗(On Diff #456805)

Added codegen tests for dynamic reshape

anlunx updated this revision to Diff 460571.Sep 15 2022, 5:35 PM

fix format

aartbik added inline comments.Sep 20 2022, 2:22 PM
mlir/test/Dialect/SparseTensor/sparse_reshape.mlir
25 ↗(On Diff #460571)

hardcoding %21 is way too brittle (and granted, it was wrong in the original test

this should be either a captured variable, as in

%[[C:.*]] = ....
scf.condition(%[[C]])

or simply

scf.condition

64 ↗(On Diff #460571)

same

84 ↗(On Diff #460571)

extra above
roundtrip:

anlunx updated this revision to Diff 462206.Sep 22 2022, 9:29 AM
anlunx marked 3 inline comments as done.

Fix test format

anlunx added inline comments.Sep 22 2022, 9:29 AM
mlir/test/Dialect/SparseTensor/sparse_reshape.mlir
25 ↗(On Diff #460571)

Removed hardcoded variable

aartbik accepted this revision.Sep 26 2022, 1:04 PM

one last nit, but good to go

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp
403

Yeah, but since you use the shape[j], I would like to make sure we never multiply it with -1 ;-)

I realize we just constructed it, but defensive asserts like this ensure that if later somebody changes that logic to allow for dynamic, we will not forget about this one.

This revision is now accepted and ready to land.Sep 26 2022, 1:04 PM
anlunx added inline comments.Sep 26 2022, 3:30 PM
mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp
403

I still don't understand why the assert is necessary. This change already allows reshaping with dynamic sizes, so the shape variable here should contain actual dynamic sizes, and can never contain -1?