Page MenuHomePhabricator

[MLIR] Added support for dynamic shaped allocas to promote-buffers-to-stack pass.
ClosedPublic

Authored by dfki-jugr on Nov 23 2020, 7:42 AM.

Details

Summary

Extended promote buffers to stack pass to support dynamically shaped allocas.
The conversion is limited by the rank of the underlying tensor.
An option is added to the pass to adjust the given rank.

Diff Detail

Event Timeline

dfki-jugr created this revision.Nov 23 2020, 7:42 AM
dfki-jugr requested review of this revision.Nov 23 2020, 7:42 AM
pifon2a requested changes to this revision.Nov 24 2020, 5:03 AM
pifon2a added inline comments.
mlir/include/mlir/Transforms/Passes.td
218

I would remove "Define the" from all these 3 options. Just

Maximal size in bytes to promote allocations to stack.
Bitwidth of the index type. Used for size estimation.
Maximal tensor rank to promote dynamic buffers.

mlir/lib/Transforms/BufferOptimizations.cpp
39–40

nit: please, write this in one line:
return type.getRank() < maxRankOfTensors; you might need a static_cast though

322

would it work with Operation* instead of auto? then you wouldn't need to write getOperation() below.

This revision now requires changes to proceed.Nov 24 2020, 5:03 AM
dfki-jugr updated this revision to Diff 307355.Nov 24 2020, 7:39 AM
dfki-jugr marked 3 inline comments as done.

Addressed comments.
Introduce transformation if the alloc operands are RankOps.

Thanks. Just some nits.

mlir/include/mlir/Transforms/Passes.td
224

Can you describe this in more detail? Not all buffers that have less than this rank will be promoted but only "small" ones.

mlir/lib/Transforms/BufferOptimizations.cpp
35

I think we we should limit this to alloc where the defining op is an AllocOp as otherwise we likely do not want to do this transformation.

44

Would return operand.getDefiningOp<RankOp>() also work?

mlir/test/Transforms/promote-buffers-to-stack.mlir
112

Can you add a check where this is more than 3?

dfki-jugr updated this revision to Diff 307607.Nov 25 2020, 7:17 AM
dfki-jugr marked 3 inline comments as done.

Addressed comments.
Changed default max rank to 1 instead of 2 and changed condition to less or equal.
This way is more intuitive.

dfki-jugr added inline comments.Nov 25 2020, 7:18 AM
mlir/test/Transforms/promote-buffers-to-stack.mlir
112

I added another test pass that checks this behaviour. An extra test for this is redundant and skipped.

herhut added inline comments.Nov 30 2020, 2:14 AM
mlir/include/mlir/Transforms/Passes.td
222

Maybe max-rank-of-allocated-memref?

mlir/lib/Transforms/BufferOptimizations.cpp
36

Should this be an ||?

silvas added a subscriber: silvas.Nov 30 2020, 4:48 PM
silvas added inline comments.
mlir/lib/Transforms/BufferOptimizations.cpp
39

nit: I feel like this comment should read "if the dynamic dimension of the alloc is produced by a RankOp, then we know it is likely to be small. Also limit this to maxRankOfTensors to prevent multiple small values produced by RankOp from multiplying to an excessively large value"

silvas added inline comments.Nov 30 2020, 4:53 PM
mlir/test/Transforms/promote-buffers-to-stack.mlir
104

nit: why does this test need control flow? this patch seems like a very localized change unrelated to control flow.

dfki-jugr updated this revision to Diff 308618.Dec 1 2020, 4:31 AM
dfki-jugr marked 4 inline comments as done.

Addressed comments.

herhut accepted this revision.Dec 2 2020, 1:12 AM
This revision was not accepted when it landed; it landed in state Needs Review.Dec 3 2020, 2:48 AM
This revision was automatically updated to reflect the committed changes.