iree-org / iree-org/iree

Reusable pattern for folding static dimensions into shape-aware ops.

Open
#8,441 1 comment 0 reactions 0 assignees View on GitHub
compiler/dialects good first issue 🌱
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.