EnzymeAD / EnzymeAD/Enzyme-JAX

Two broadcast optimizations

Open
#2,038 1 comment 0 reactions 0 assignees View on GitHub
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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.