This is an archive of the discontinued LLVM Phabricator instance.

[mlir][tosa] Add shape inference for tosa.while
ClosedPublic

Authored by rsuderman on Aug 26 2021, 7:24 PM.

Details

Summary

Tosa.while shape inference requires repeatedly running shape inference across
the body of the loop until the types become static as we do not know the number
of iterations required by the loop body. Once the least specific arguments are
known they are propagated to both regions.

To determine the final end type, the least restrictive types are determined
from all yields.

Diff Detail

Event Timeline

rsuderman created this revision.Aug 26 2021, 7:24 PM
rsuderman requested review of this revision.Aug 26 2021, 7:24 PM

Nice, some questions to help my understanding here

mlir/lib/Dialect/Tosa/IR/TosaOps.cpp
1459

Add comment for failure

1466

So these would require that the operand type has been updated?

1475

Mmm, would

if (auto meet = ...)

resultKnowledge[index] = meet;

result in difficult to read if? (removed need for continue and reduces lifetime scope of variable)

[I needed to check meet here to make sure it was doing what I'd expect for while where it could get less refined every time through]

1485

Do we expect to need this pattern more times? (I see it is repeated 2x, so perhaps a helper method to go from resultKnowledge to shapes could help)

mlir/lib/Dialect/Tosa/Transforms/TosaInferShapes.cpp
38

Both here and below, an array would have worked as well, is the map just easier to populate? (I htink only accessed once in each, so cost of indexing probably the same)

48

Nit: hoist getNumOperands() to document it doesn't change during iteration.

77

Is tosa:: needed? seems like you are using mlir::tosa above.

81

expected ?

82

consistent

82

And this is a property of tosa while?

96

originalTypeMap ? (backup feels like something used in case of failure)

121

OOC, why not yieldOp ?

154

Could you expand on this? I thought the types would be more refined now, not sure why it would need to be reset

167

Ditto

171

It feels like if propagating to regions took a list of shapes, then we could avoid the setting above (but I don't get the resetting part yet, which is either helped by this/can be removed, or shows why it is needed to set)

rsuderman updated this revision to Diff 370643.Sep 3 2021, 12:17 PM

Updated for jpienaar comments.

rsuderman updated this revision to Diff 370644.Sep 3 2021, 12:22 PM
rsuderman marked 7 inline comments as done.

Removed extraneous tosa::

rsuderman updated this revision to Diff 370645.Sep 3 2021, 12:25 PM
rsuderman marked 7 inline comments as done.

Few more tweaks to updates.

rsuderman added inline comments.Sep 3 2021, 12:26 PM
mlir/lib/Dialect/Tosa/IR/TosaOps.cpp
1466

Kinda? It more uses the initial operand type as the starting place for performing shape inference. The idea with while loop propagation is that the block args for cond/while must be compatible with ever invocation. Therefore we start with the initial types (specified by the operand types) and iterate each cycle.

E.g. if we don't know much about the operand types we can't infer as much about the types in the cond/body blocks.

1485

Added a helper to ValueKnowledge.

mlir/lib/Dialect/Tosa/Transforms/TosaInferShapes.cpp
38

This is true however the DenseMap is already defined and it would be the same number of dense map lookups to determine the operand types. So better to just pass the DenseMap.

82

I added more details but overall, yes. Because we execute the regions multiple times, we need their types compatible with every possible iteration.

121

No particular reason. Changed to YieldOp.

171

Yeah, we could go store a map of shapes instead and propagate rather than mutating the IR every iteration. It felt more succinct to accept mutation rather than work around it. Trying to store a type change map introduced a level of indirection that felt unnecessary.

rsuderman accepted this revision.Sep 3 2021, 12:27 PM
rsuderman marked an inline comment as done.
This revision is now accepted and ready to land.Sep 3 2021, 12:27 PM
rsuderman resigned from this revision.Sep 3 2021, 12:27 PM
This revision now requires review to proceed.Sep 3 2021, 12:27 PM

Accidentally reviewed it myself..... woo reflexes

jpienaar accepted this revision.Sep 10 2021, 10:32 AM

Looks good thanks

mlir/lib/Dialect/Tosa/IR/TosaOps.cpp
1466

Sure, I was more referring to operands.getShape which could give a more exact shape here

mlir/lib/Dialect/Tosa/Transforms/TosaInferShapes.cpp
128

Micro nit: should TOSA be used when referring to the opset and tosa.X when referring to op X in opset TOSA? (e.g., is there a capitalization convention to follow here)

145

Perhaps op.emitWarning() ? If this should never happen, then we are only silently failing at the moment.

This revision is now accepted and ready to land.Sep 10 2021, 10:32 AM

Fixed up for last comments.

rsuderman marked 3 inline comments as done.Sep 10 2021, 12:34 PM
This revision was automatically updated to reflect the committed changes.