Page MenuHomePhabricator

[mlir][linalg] Add output tensor args folding for linalg.tiled_loop.
ClosedPublic

Authored by pifon2a on Mar 25 2021, 5:29 AM.

Details

Summary

Folds away TiledLoopOp output tensors when the following conditions are met:

  • result of linalg.tiled_loop has no uses
  • output tensor is the argument of linalg.yield

Example:

%0 = linalg.tiled_loop ...  outs (%out, %out_buf:tensor<...>, memref<...>) {
  ...
  linalg.yield %out : tensor ...
}

Becomes

linalg.tiled_loop ...  outs (%out_buf:memref<...>) {
  ...
  linalg.yield
}

Diff Detail

Event Timeline

pifon2a created this revision.Mar 25 2021, 5:29 AM
pifon2a requested review of this revision.Mar 25 2021, 5:29 AM
pifon2a edited the summary of this revision. (Show Details)Mar 25 2021, 5:30 AM
pifon2a retitled this revision from [mlir][linalg] Add results folding for linalg.tiled_loop. to [mlir][linalg] Add output tensor args folding for linalg.tiled_loop..
pifon2a updated this revision to Diff 333312.Mar 25 2021, 8:19 AM

Simplify.

pifon2a edited the summary of this revision. (Show Details)Mar 25 2021, 8:20 AM
pifon2a edited the summary of this revision. (Show Details)
pifon2a edited the summary of this revision. (Show Details)
mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
1980

can you use llvm::enumerate and drop resultId ?

1982

can you use newOutputOperands ?
I got myself confused thinking these were results.
also, assert memreftype ?

1999

I would just set operands in place with the special inplaceupdateroot thing and the xxxMutable + only clone/erase the terminator.
The rest will DCE later.

mlir/test/Dialect/Linalg/canonicalize.mlir
819

I'd just drop the operations in the body and the memref.tensor_load and instead pass a new function arg: these are mostly noise for this test.
instead you can just have a call to an external function that takes %A, %B and %C for the purpose of simulating internal side-effects.

pifon2a updated this revision to Diff 333331.Mar 25 2021, 9:55 AM
pifon2a marked 4 inline comments as done.

Address the comments

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
1980

unfortunately no, because |outputs()| >= |results|. If we have strictly more outputs (because of memref output operands), then it would not work.

1982

yes, agreed, it's a better name

1999

Is it possible also to erase the results in-place?

LG modulo using the rewriter for IR creation.

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
1999

true, sorry about the noise.

2005

we need to use the rewriter to clone these ops.

This revision is now accepted and ready to land.Mar 25 2021, 10:05 AM
pifon2a updated this revision to Diff 333334.Mar 25 2021, 10:08 AM

Addressed the comments.

pifon2a marked 2 inline comments as done.Mar 25 2021, 10:09 AM
This revision was landed with ongoing or failed builds.Mar 25 2021, 10:12 AM
This revision was automatically updated to reflect the committed changes.