[Mlir-commits] [mlir] [mlir][tosa] Harden folds/canonicalizations for unranked and dynamic shapes (PR #188188)

[Hocky Yudhiono]

hockyy wrote:


## Changes

### `TosaCanonicalizations.cpp`

- Add `foldToInputIfTypeMatches`: central guard for folds that would otherwise return the input value unchanged.
- `ResizeOp::fold`, `TransposeOp::fold`, `NegateOp::fold`, `AbsOp::fold`: Use the same type-match rule for no-op / algebraic folds.
- Revamp `ReverseOp::fold` a bit, changed the overall algorithm
- `binaryFolder` / `unaryFolder`: Early-out if the result `ShapedType` is missing or not statically shaped, it's working with `DenseElementsAttr`, so adding this check would make sense I believe.
- `CastOp::fold`, `ReciprocalOp::fold`: Skip folds that would need a static result shape for dense materialization.
- `GreaterOp` / `GreaterEqualOp` / `EqualOp::fold`: Use `ShapedType` for the result type where appropriate (aligned with static-shape checks elsewhere), logic check bail out is now checked at `binaryFolder`.
- `ClampIsNoOp`: Use `ShapedType` for the input tensor type, NFC.
- `SliceDynamicSizeCanonicalization`: `notifyMatchFailure` if the slice result is not ranked (avoids bad assumptions on unranked outputs).

### `TosaFolders.cpp`

- `TosaFoldConstantTranspose`: Do not fold when the transpose result type is not statically shaped.
- `TosaFoldConstantReciprocal`: Require input type == reciprocal result type before rewriting to a new constant.

### Tests

- `canonicalize.mlir`: Identity resize fold (and no-fold when result is unranked / asymmetric scale); unranked clamp and slice must not canonicalize unsafely; reverse no-fold cases (unranked operand, unranked result, splat type mismatch). Adjust existing resize test to use BILINEAR where appropriate.
- `constant_folding.mlir`: Ops with unranked results should not fold away incorrectly.
- `tosa-layerwise-constant-fold.mlir`: Transpose and reciprocal constant folds skipped for unranked results.

@sahas3 @lhutton1 ^ summary


https://github.com/llvm/llvm-project/pull/188188