EnzymeAD / EnzymeAD/Enzyme-JAX
`slice(transpose(...))` --> `transpose(slice(...))` if transpose is only being used in slice
- Dominant language
- MLIR
- Stars
- 131
- Forks
- 53
- Avg merge
- 1d 10h
- Merged PRs (30d)
- 193
Description
enables optimizations of routines like (if transpose is being used exclusively for slice operations)
```mlir
module {
func.func @main(%arg0: tensor<32x16xf32>, %arg1: tensor<32x8xf32>) -> (tensor<32x8xf32>, tensor<32x8xf32>) {
%0 = stablehlo.transpose %arg0, dims = [1, 0] : (tensor<32x16xf32>) -> tensor<16x32xf32>
%1 = stablehlo.transpose %arg1, dims = [1, 0] : (tensor<32x8xf32>) -> tensor<8x32xf32>
%2 = stablehlo.slice %0 [0:8, 0:32] : (tensor<16x32xf32>) -> tensor<8x32xf32>
%3 = stablehlo.slice %0 [8:16, 0:32] : (tensor<16x32xf32>) -> tensor<8x32xf32>
%4 = stablehlo.exponential %3 : tensor<8x32xf32>
%5 = stablehlo.multiply %1, %4 : tensor<8x32xf32>
%6 = stablehlo.add %5, %2 : tensor<8x32xf32>
%7 = stablehlo.transpose %6, dims = [1, 0] : (tensor<8x32xf32>) -> tensor<32x8xf32>
%8 = stablehlo.transpose %3, dims = [1, 0] : (tensor<8x32xf32>) -> tensor<32x8xf32>
return %7, %8 : tensor<32x8xf32>, tensor<32x8xf32>
}
}
```
Contributor guide
No contributing guide indexed for this repository
Assessment
This issue has not been assessed yet.