Reusable pattern for folding static dimensions into shape-aware ops.
- Dominant language
- C++
- Stars
- 3.9k
- Forks
- 1k
- Avg merge
- 4d 16h
- Merged PRs (30d)
- 47
Description
Though a majority of the cases where dimensions become static happen early in the pipeline there are some that happen after lowering into `flow` and `stream.tensor.*`. It'd be really nice to be able to propagate this static information before lowering into the stream dialects.
These come from consuming tensors whose shape may be knowable:
```mlir
%0 = tensor.cast %input : tensor<1x2xf32> -> tensor
%dim = tensor.dim %0, %c1 : tensor
%1 = flow.tensor.clone %0 : tensor{%dim}
```
->
```mlir
%0 = tensor.cast %input : tensor<1x2xf32> -> tensor
%1 = flow.tensor.clone %0 : tensor{%c1}
```
And can also happen with consumers:
```mlir
%dim = tensor.dim %input, %c1 : tensor
%0 = flow.tensor.clone %input : tensor{%dim}
%1 = tensor.cast %0 : tensor -> tensor<1x2xf32>
```
I'm not sure there's anything today that replaces the dim if there's a subsequent cast, but that'd be useful!:
```mlir
%0 = flow.tensor.clone %input : tensor{%c1}
%1 = tensor.cast %0 : tensor -> tensor<1x2xf32>
```
There's a few approaches here with the most robust being to make the `Util_ShapeAwareOp` interface support recreating the op with new static shape dimensions. A canonicalization pattern registered on the interface could then check operands/results for cases where more static information is available and recreate the op with that. Another approach would be to expose mutable fields on the interface but that gets messy - there's only a dozen ops and it'd be easier to just rebuild them.
Ideally we'd then end up with something that turned the casts into clones:
```mlir
%0 = tensor.cast %input : tensor<1x2xf32> -> tensor
%1 = flow.tensor.clone %0 : tensor{%c1}
```
->
```mlir
%0 = flow.tensor.clone %input : tensor<1x2xf32> -> tensor<1x2xf32>
%1 = flow.tensor.clone %0 : tensor<1x2xf32>
```
and
```mlir
%0 = flow.tensor.clone %input : tensor{%c1}
%1 = tensor.cast %0 : tensor -> tensor<1x2xf32>
```
->
```mlir
%0 = flow.tensor.clone %input : tensor<1x2xf32>
%1 = flow.tensor.clone %0 : tensor<1x2xf32> -> tensor<1x2xf32>
```
Since in tensor form the cast may be forking an immutable tensor value the clone preserves the behavior (for example if there were two cast ops consuming the same tensor) - clone canonicalizers can then kick in and handle the rest of the cleanup.
Examples of where this can arise and not be caught earlier are any IPO we do (across globals/function/branch boundaries), const eval that ends up introducing more static information, or specialization where we branch off paths that are say size 1 and size N and want to optimize the size 1 half. #8441 would also benefit as we can get more static 0's and kill more ops.
Contributor guide
Research direction
Start by locating the Util_ShapeAwareOp interface and the flow.tensor.clone, stream.tensor.*, tensor.cast, and tensor.dim operations mentioned in the issue. Review existing canonicalizers and determine how static dimensions can be propagated across operands and results. Done means shape-aware operations are recreated with available static dimensions and the illustrated casts, clones, and redundant dynamic dimensions canonicalize as shown.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp
- Domain
- compilers
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 30/100