EnzymeAD / EnzymeAD/Enzyme-JAX
Pull repeated ops from conditional branches
- Dominant language
- MLIR
- Stars
- 131
- Forks
- 53
- Avg merge
- 1d 10h
- Merged PRs (30d)
- 193
Description
```mlir
%3:2 = "stablehlo.if"(%2) ({
%14 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor) -> tensor<4x2xi64>
%15 = stablehlo.convert %14 : (tensor<4x2xi64>) -> tensor<4x2xf32>
%16 = stablehlo.divide %1, %15 : tensor<4x2xf32>
stablehlo.return %16, %16 : tensor<4x2xf32>, tensor<4x2xf32>
}, {
%14 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor) -> tensor<4x2xi64>
%15 = stablehlo.convert %14 : (tensor<4x2xi64>) -> tensor<4x2xf32>
%16 = stablehlo.divide %1, %15 : tensor<4x2xf32>
%17 = stablehlo.add %0, %16 : tensor<4x2xf32>
stablehlo.return %17, %17 : tensor<4x2xf32>, tensor<4x2xf32>
}) : (tensor) -> (tensor<4x2xf32>, tensor<4x2xf32>)
```
We should be able to lift
```mlir
%14 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor) -> tensor<4x2xi64>
%15 = stablehlo.convert %14 : (tensor<4x2xi64>) -> tensor<4x2xf32>
%16 = stablehlo.divide %1, %15 : tensor<4x2xf32>
```
into the outer block
Contributor guide
No contributing guide indexed for this repository
Assessment
This issue has not been assessed yet.