EnzymeAD / EnzymeAD/Enzyme-JAX
Transpose Elementwise always if all branches can remove transpose
- Dominant language
- MLIR
- Stars
- 131
- Forks
- 53
- Avg merge
- 1d 10h
- Merged PRs (30d)
- 193
Description
```mlir
%21 = stablehlo.convolution(%20, %arg4) dim_numbers = [0, 1, f, b]x[o, i, 1, 0]->[0, 1, f, b], window = {stride = [1, 1], pad = [[0, 0], [0, 0]], rhs_dilate = [1, 1]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64, precision_config = [#stablehlo, #stablehlo]} : (tensor<64x64x128x256xf32>, tensor<64x128x1x1xf32>) -> tensor<64x64x64x256xf32>
%22 = stablehlo.broadcast_in_dim %arg5, dims = [2] : (tensor<64xf32>) -> tensor<64x64x64x256xf32>
%23 = stablehlo.add %21, %22 : tensor<64x64x64x256xf32>
%24 = stablehlo.tanh %23 : tensor<64x64x64x256xf32>
%25 = stablehlo.transpose %24, dims = [3, 2, 1, 0] : (tensor<64x64x64x256xf32>) -> tensor<256x64x64x64xf32>
%51 = stablehlo.broadcast_in_dim %arg12, dims = [2, 1, 0] : (tensor<64x1x1xf32>) -> tensor<64x64x64x256xf32>
%52 = stablehlo.multiply %23, %51 : tensor<64x64x64x256xf32>
```
Contributor guide
No contributing guide indexed for this repository
Research direction
Start from the StableHLO convolution, elementwise operations, and transpose shown in the issue, and trace the compiler pass that decides whether to transpose elementwise branches. Check how the branches are analyzed for removable transposes. Done means the intended transpose behavior is applied when all branches can remove their transpose, with coverage for this operation pattern.
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