This is an archive of the discontinued LLVM Phabricator instance.

[mlir][sparse] replace zero yield generic op with copy in allocation
ClosedPublic

Authored by aartbik on Aug 3 2022, 3:52 PM.

Details

Summary

This prepares patterns that sometimes are generated by the front-end
and would prohibit fusion of SDDMM flavored kernels.

Diff Detail

Event Timeline

aartbik created this revision.Aug 3 2022, 3:52 PM
Herald added a project: Restricted Project. · View Herald TranscriptAug 3 2022, 3:52 PM
aartbik requested review of this revision.Aug 3 2022, 3:52 PM
Peiming added inline comments.Aug 3 2022, 4:32 PM
mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp
120

Why only fold zero?

aartbik added inline comments.Aug 3 2022, 5:22 PM
mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp
120

For now, the only situation that appears from automatically generated code.

springerm added inline comments.Aug 4 2022, 12:45 AM
mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp
102

What if the op has multiple outputs?

102

In addition, you could also check if yield.getOperand(0) is a zero constant.

120

initial

120

converts

134–139

Is this safe? What if a already has a copy operand and a has multiple users?

138

Should be wrapped in rewriter.updateRootInPlace.

139

Can you use zero here? arith.constant bufferize to a memref.get_global and is "non-writable". I.e., if some operation tries to write to it, the bufferization will automatically insert an alloc_tensor op.

aartbik marked 6 inline comments as done.Aug 4 2022, 8:10 AM
aartbik added inline comments.
mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp
102

That is taken care of in L125. Or do you mean you want an assert here?

102

Extended this for direct value. Like Peiming suggested, for now I restrict this to zero (since that is the only useful case in the set of rewriting, but we may expand other constants too in the future).

134–139

We guarantee there is no copy at L127 (isZero == false). And even if there are several uses, the contents will be filled either way, right?

138

I have changed it, also below. See if this is the right idiom?

139

But that is why we replace the copy in the already existing bufferization allocation, right?

aartbik updated this revision to Diff 449988.Aug 4 2022, 8:12 AM
aartbik marked 4 inline comments as done.

addressed comments

springerm accepted this revision.Aug 4 2022, 8:39 AM
This revision is now accepted and ready to land.Aug 4 2022, 8:39 AM
aartbik updated this revision to Diff 450007.Aug 4 2022, 8:58 AM

more precise update root call