Page MenuHomePhabricator

[mlir][linalg] Fusion on tensors.
ClosedPublic

Authored by gysit on Sep 14 2021, 9:22 AM.

Details

Summary

Add a new version of fusion on tensors that supports the following scenarios:

  • support input and output operand fusion
  • fuse a producer result passed in via tile loop iteration arguments (update the tile loop iteration arguments)
  • supports only linalg operations on tensors
  • supports only scf::for
  • cannot add an output to the tile loop nest

The LinalgTileAndFuseOnTensors pass tiles the root operation and fuses its producers.

Diff Detail

Event Timeline

gysit created this revision.Sep 14 2021, 9:22 AM
gysit requested review of this revision.Sep 14 2021, 9:22 AM
nicolasvasilache accepted this revision.Sep 17 2021, 7:00 AM

Thanks for pushing this !

mlir/lib/Dialect/Linalg/Transforms/FusionOnTensors.cpp
222

InsertionGuard + setIP please, don't assume cross boundary correctness with IP

242

can this similarly be replaced by some applyPermutation for better readability?

275

can this be a failure someplace earlier rather than a hard stop ?

353

"except when/if" ?

354

early exit to avoid increasing the nesting level wherever possible plz

378

early exit to avoid increasing the nesting level wherever possible plz

393

use/refactor applyPermutationToVector so that you can compose existing things in a more idiomatic fashion plz

438

Please mark clearly that this is the heuristic.

459

This looks quite adhoc and should be generally reusable.
It is related to the concept of "projected permutation", can you refactor in a way that uses AffineMap and composes / extends existing AffineMap abstractions?

465

this check for instance is part of isProjectedPermutation IIRC

This revision is now accepted and ready to land.Sep 17 2021, 7:00 AM
mravishankar accepted this revision.Sep 17 2021, 9:06 AM

THis mostly looks OK to me. Just a few style comments.

mlir/include/mlir/Dialect/Linalg/Utils/Utils.h
180

Pure style comments: Its probably better to have a struct that is just fields, and have helper methods within a file. Unless these methods are required outside of the tiling algorithm, I dont think there is a need to expose these as methods.

mlir/lib/Dialect/Linalg/Transforms/FusionOnTensors.cpp
42

Not sure I understand this part. Is this just checking that all the ops are scf::ForOp, cuase it is putting them back into an ArrayRef<Operation *> which means before and after tileLoopOps is unmodified.

If this is just a check, use llvm::any_of(...)

212

I know that the tiled_loop abstraction is to be revisited, but I think that helps in this case cause you dont need to follow iter-args. Just a note.

434

Nit: please add { } around multiple textual lines.

gysit updated this revision to Diff 373496.Sep 20 2021, 12:10 AM
gysit marked 12 inline comments as done.

Address comments and fix issue with repeated fusion.

gysit updated this revision to Diff 373582.Sep 20 2021, 7:14 AM

Address additional comments.

nicolasvasilache accepted this revision.Sep 20 2021, 7:41 AM
This revision was landed with ongoing or failed builds.Sep 20 2021, 7:46 AM
This revision was automatically updated to reflect the committed changes.