EnzymeAD / EnzymeAD/Enzyme-JAX
Two broadcast optimizations
- Dominant language
- MLIR
- Stars
- 131
- Forks
- 53
- Avg merge
- 1d 10h
- Merged PRs (30d)
- 193
Description
```
// AUTOBATCH-NEXT: %[[a5:.+]] = stablehlo.reshape %[[a4]] : (tensor<14x1xf32>) -> tensor<14x1x1xf32>
// AUTOBATCH-NEXT: %[[a6:.+]] = stablehlo.broadcast_in_dim %[[a4]], dims = [0, 2, 1] : (tensor<14x1x1xf32>) -> tensor<14x5x3xf32>
```
From here:
First, a broadcast of singleton dims which are permuted can be sorted (thus in the case above it can be dims = [0, 1, 2]).
Second, we should extend broadcastreshape to support inserted dims (and thus fold the reshape into the broadcast in dim, reducing the number of ops).
Contributor guide
No contributing guide indexed for this repository
Research direction
Start by locating the tests and implementation for stablehlo.reshape and stablehlo.broadcast_in_dim, using the AUTOBATCH-NEXT example as the initial reproduction. Verify the singleton-dimension permutation case and the inserted-dimension broadcastreshape case; done means both optimizations are covered and the resulting operation sequence has fewer or canonicalized broadcasts.
Written by the indexing model from the issue text.
Assessment
- Domain
- compilers
- Issue type
- Feature
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100