Allow folding of different tensor types when the constant tensor is broadcast. Removed redundant and incorrect AddZero and MulOne canonical optimizations.
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.
This looks like an accidental inclusion?
I wonder if we can generalize the two branches into a single, non-templated isSplatZero(rhsAttr) && lhsTy == resultTy branch
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.
Same comment on the new patch, this should check against a broadcastable shape, ranks must match, and dimension sizes must be equal, or 1.
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).
Same issue here.
@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.
- added rank checks for the folders
- further simplified
- also removed redundant and incorrect pattern rewriters for AddZero and MulOne
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.
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.
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.
@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?
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.
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.