This is an archive of the discontinued LLVM Phabricator instance.

[TOSA] Loosen folding restrictions for tosa.add,tosa.sub, tosa.mul
ClosedPublic

Authored by sjw36 on Mar 9 2023, 4:12 PM.

Details

Summary

Allow folding of different tensor types when the constant tensor is broadcast. Removed redundant and incorrect AddZero and MulOne canonical optimizations.

Diff Detail

Event Timeline

sjw36 created this revision.Mar 9 2023, 4:12 PM
Herald added a project: Restricted Project. · View Herald TranscriptMar 9 2023, 4:12 PM
sjw36 requested review of this revision.Mar 9 2023, 4:12 PM
sjw36 updated this revision to Diff 504165.Mar 10 2023, 8:49 AM

update format

eric-k256 added a subscriber: rsuderman.

Adding @rsuderman, the author of the existing fold code as a more appropriate reviewer.

TOSA operators don't allow implicit rank changes, they need to be done with an explicit tosa.reshape. Implicit broadcast within a rank is allowed, so this could be changed to allow broadcast, but not dimension differences.

mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
1

This looks like an accidental inclusion?

The change looks good, just make the fix the @eric-k256 pointed out.

mgehre-amd added inline comments.Mar 26 2023, 11:34 PM
mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
677

I wonder if we can generalize the two branches into a single, non-templated isSplatZero(rhsAttr) && lhsTy == resultTy branch

sjw36 updated this revision to Diff 508882.Mar 27 2023, 8:57 PM
  • simplified logic
  • removed irrelevant change
sjw36 marked an inline comment as done.Mar 27 2023, 8:59 PM

@mgehre-amd thanks, simplified the logic.
@rsuderman cleaned up PR.

eric-k256 requested changes to this revision.Mar 28 2023, 11:59 AM

I like the cleanup and the overall simplification of the code, but the broadcast issue still remains. Not specifically called out, but in addition to the marked issues, updating the commit message would also be good.

mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
618

Same comment on the new patch, this should check against a broadcastable shape, ranks must match, and dimension sizes must be equal, or 1.
See resolveBroadcastShape in TosaOps.cpp for a reference.

mlir/test/Dialect/Tosa/canonicalize.mlir
24

This isn't a valid TOSA add, the rank of both tensors needs to match. So tensor<1x1x1xi32> would be valid (would broadcast along dimensions of size 1).

193

Same issue here.

This revision now requires changes to proceed.Mar 28 2023, 11:59 AM
sjw36 marked an inline comment as done.Mar 28 2023, 4:13 PM

@eric-k256 thanks. I didn't see a hard requirement in the spec and was confused by the test above for add_zero_different_shape. Also turns out there is a AddZeroOptimization and MulOneOptimization that also allow mixed rank folding.. also kind of redundant. Shall I remove those as well?

@eric-k256 thanks. I didn't see a hard requirement in the spec and was confused by the test above for add_zero_different_shape. Also turns out there is a AddZeroOptimization and MulOneOptimization that also allow mixed rank folding.. also kind of redundant. Shall I remove those as well?

No problem. Section 1.9.4 covers broadcasting, it's possible if you were only looking at ADD/MUL it wouldn't be as obvious, but it says: TOSA broadcast requires the rank of both tensors to be the same. https://www.mlplatform.org/tosa/tosa_spec.html#_broadcasting

Yes, to reduce confusion, it would be good if you could also remove those. I must not have caught that when they went in.

sjw36 updated this revision to Diff 509188.EditedMar 28 2023, 7:25 PM
  • added rank checks for the folders
  • further simplified
  • also removed redundant and incorrect pattern rewriters for AddZero and MulOne
mgehre-amd added inline comments.Mar 29 2023, 12:07 AM
mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
671–672

Why do we need to check this here?

From what I understand, a sub is invalid when the rank of lhs and rhs differs. The verifier would catch this.
So passes/transforms can assume that they get valid IR as input, and thus don't need to each
verify the validity by themselves.

sjw36 marked 2 inline comments as done.Mar 29 2023, 7:56 AM
sjw36 added inline comments.
mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
671–672

Agreed. The canonicalize test below has these bad IR cases but the verifier doesn't seem to catch them in any case, @eric-k256 perhaps file a ticket?

Good catch @mgehre-amd. I had been thinking along a similar path to what you described, your comment helped solidify my thoughts.
If we track fixing the verifier, then agree that we can drop the check here. In terms of this change, we should change the tests so that if we fix the verifier, we don't then have to change these tests. (We're likely to find other tests that fail, but let's make our future work easier 😄 ). While doing that, changing the commit message to drop the mention of lower dimension helps avoid confusion.

mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
671–672

Yes, a ticket sounds reasonable to me, with the caveat that I'm not sure how well those get monitored for MLIR overall, and less for TOSA. I expect that I can't be assigned to the issue. I don't know if you wanted to file it, or would like me to. If you file it, just let me know the number, and I'll subscribe. It seems like an important issue to fix, so I'll try to find someone here to give it a shot, unless you're already planning on working on it.

I just scanned through the code, and I think I understand the issue. The TOSA ops have the ResultsBroadcastableShape trait on them, which does do checking of the operands, however it uses a broader compatibility check than what TOSA allows, so the mismatch of ranks passes verification. It appears that we need to switch to a TOSA specific verification of broadcasting to catch this case.

sjw36 edited the summary of this revision. (Show Details)Mar 29 2023, 10:37 AM

@eric-k256 I guess I will just remove the _different_shape tests and keep the new _bcast_ tests since that should be sufficient coverage. Is that acceptable?

sjw36 updated this revision to Diff 509422.Mar 29 2023, 11:07 AM
  • removed rank checks
  • removed invalid canonicalize tests (_different_shape)
sjw36 marked an inline comment as done.Mar 29 2023, 3:08 PM
sjw36 added inline comments.
mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
671–672

I unfortunately won't have time to work on that in the near future. Perhaps if you file the ticket you can assign it? I don't know the system here either.

eric-k256 accepted this revision.Mar 29 2023, 3:51 PM

Thanks for making all of the changes. Issue if you want to follow along: https://github.com/llvm/llvm-project/issues/61822. Not assigned to me, but if you're not working on it I'm less worried that someone else will pick it up without me knowing.

I may have overstepped on removing the existing optimization/canonicalization passes in favor of only the folder. @rsuderman, you'll probably want to make sure that's okay.
But from my perspective it looks ready to go.

This revision is now accepted and ready to land.Mar 29 2023, 3:51 PM
rsuderman accepted this revision.Mar 30 2023, 11:08 AM