This is an archive of the discontinued LLVM Phabricator instance.

[mlir][linalg] Add a new pattern to handle folding unit reduction dims.
ClosedPublic

Authored by hanchung on Nov 17 2022, 6:11 PM.

Details

Summary

The output operands will be added to input operands if the generic op (on tensors)
becomes an elementwise operation. The outputs of the generic op is still the same.
They will be cleaned up by ReplaceWithEmptyTensorIfUnused pattern.

Diff Detail

Event Timeline

hanchung created this revision.Nov 17 2022, 6:11 PM
hanchung requested review of this revision.Nov 17 2022, 6:11 PM
hanchung updated this revision to Diff 476304.Nov 17 2022, 6:22 PM

update a comment

hanchung added inline comments.Nov 17 2022, 6:26 PM
mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp
239–269

Maybe I should move these to StaticValueUtils.h. They are also used in IREE LinalgExt dialect...

mravishankar requested changes to this revision.Nov 17 2022, 7:36 PM
mravishankar added inline comments.
mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp
152

I think you can replace all the uses of this with rewriter.getAffineMapArrayAttr https://github.com/llvm/llvm-project/blob/d5844685810980399397a4310b943532361790ef/mlir/include/mlir/IR/Builders.h#L154

162

Change to Block seems unnecessary?

239–269
318

Instead of adding a separate pattern for reduction (and making the existing pattern handle only parallel iterator types), it might be better to move the code to switch the operand from outs to ins in a separate pattern (if the operation has tensor semantics). Sorry, I should have been more explicit in my recommendation that this the pattern I was thinking about.

This revision now requires changes to proceed.Nov 17 2022, 7:36 PM
hanchung updated this revision to Diff 476562.Nov 18 2022, 12:20 PM

address comments

hanchung marked 3 inline comments as done.Nov 18 2022, 12:22 PM
hanchung added inline comments.
mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp
152

reverted as following the other suggestion.

239–269

thanks for the pointer, very useful!

hanchung edited the summary of this revision. (Show Details)Nov 18 2022, 12:23 PM
hanchung marked an inline comment as done.
hanchung added inline comments.
mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp
318

After reviewing the lit cases and offline discussion, this makes more sense to me now. Thanks for the suggestion!

mravishankar requested changes to this revision.Nov 18 2022, 6:03 PM

Just a few more clarifications.

mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp
243

Actually I am not really sure why we need this separate pattern. SHouldnt it be easier to just add a tensor.empty as the value for the new outs operand added to AddInitOperandsToInput?

254

Instead of using this, use getDpsInits and genericOp.getMatchingBlockArgument to get the argument.

267

I am not sure we need this....

282

Maybe add some IR to explain this better. Also init have uses should be clarified to say blockArgument corresponding to init is used in the region.

This revision now requires changes to proceed.Nov 18 2022, 6:03 PM
hanchung updated this revision to Diff 477025.Nov 21 2022, 4:12 PM
hanchung marked 3 inline comments as done.

address comments

remove unneeded pattern

use DPS/StructuredOp methods

mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp
243

I don't remember what the issues I hit, but it works now. Removed the pattern.

254

thanks, I haven't touch DSP stuff for a while. They are all replaced with DSP/StructuredOp methods, and look cleaner. Thanks for the suggestion!

267

removed.

mravishankar accepted this revision.Nov 21 2022, 11:15 PM
mravishankar added inline comments.
mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp
762–765

LLVM style I think is to not use { } here...

This revision is now accepted and ready to land.Nov 21 2022, 11:15 PM
hanchung added inline comments.Nov 22 2022, 6:02 PM
mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp
762–765

good point, it was two lines when having three patterns, so I added braces. I forgot to remove them...