EnzymeAD / EnzymeAD/Enzyme-JAX

Transpose Elementwise always if all branches can remove transpose

Open
#1,094 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

```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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.