EnzymeAD / EnzymeAD/Enzyme-JAX
Slice access analysis to eliminate slices and concats inside loops
- Dominant language
- MLIR
- Stars
- 131
- Forks
- 53
- Avg merge
- 1d 10h
- Merged PRs (30d)
- 193
Description
### case I
```mlir
%235 = stablehlo.slice %iterArg_101 [0:1, 1:2033, 0:4080] {mhlo.sharding = "{devices=[1,4,8]<=[8,4]T(1,0)}"} : (tensor<1x2034x4080xf64>) -> tensor<1x2032x4080xf64> loc(#loc1095) // <-- this is just %240
....
%244 = stablehlo.concatenate %242, %240, %243, dim = 1 {mhlo.sharding = "{devices=[1,4,8]<=[8,4]T(1,0)}"} : (tensor<1x1x4080xf64>, tensor<1x2032x4080xf64>, tensor<1x1x4080xf64>) -> tensor<1x2034x4080xf64> loc(#loc1167)
return %iterArg_101 = %244
```
### case II
```mlir
%209 = stablehlo.slice %iterArg_102 [1:129, 1:2033, 0:4080] {mhlo.sharding = "{devices=[1,4,8]<=[8,4]T(1,0)}"} : (tensor<130x2034x4080xf64>) -> tensor<128x2032x4080xf64> loc(#loc1093) // same as slicing %254 along dim 0
%571 = stablehlo.slice %iterArg_102 [0:1, 0:1, 0:4080] {mhlo.sharding = "{devices=[1,4,8]<=[8,4]T(1,0)}"} : (tensor<130x2034x4080xf64>) -> tensor<1x1x4080xf64> loc(#loc) // <-- not needed
%572 = stablehlo.slice %iterArg_102 [129:130, 0:1, 0:4080] {mhlo.sharding = "{devices=[1,4,8]<=[8,4]T(1,0)}"} : (tensor<130x2034x4080xf64>) -> tensor<1x1x4080xf64> loc(#loc) // <-- not needed
%574 = stablehlo.slice %iterArg_102 [0:1, 2033:2034, 0:4080] {mhlo.sharding = "{devices=[1,4,8]<=[8,4]T(1,0)}"} : (tensor<130x2034x4080xf64>) -> tensor<1x1x4080xf64> loc(#loc) // <-- not needed
%575 = stablehlo.slice %iterArg_102 [129:130, 2033:2034, 0:4080] {mhlo.sharding = "{devices=[1,4,8]<=[8,4]T(1,0)}"} : (tensor<130x2034x4080xf64>) -> tensor<1x1x4080xf64> loc(#loc) // <-- not needed
%254 = stablehlo.concatenate %246, %214, %248, dim = 0 {mhlo.sharding = "{devices=[1,4,8]<=[8,4]T(1,0)}"} : (tensor<1x2032x4080xf64>, tensor<128x2032x4080xf64>, tensor<1x2032x4080xf64>) -> tensor<130x2032x4080xf64> loc(#loc1194)
%573 = stablehlo.concatenate %571, %260, %572, dim = 0 {mhlo.sharding = "{devices=[1,4,8]<=[8,4]T(1,0)}"} : (tensor<1x1x4080xf64>, tensor<128x1x4080xf64>, tensor<1x1x4080xf64>) -> tensor<130x1x4080xf64> loc(#loc1190) // <-- not needed
%576 = stablehlo.concatenate %574, %262, %575, dim = 0 {mhlo.sharding = "{devices=[1,4,8]<=[8,4]T(1,0)}"} : (tensor<1x1x4080xf64>, tensor<128x1x4080xf64>, tensor<1x1x4080xf64>) -> tensor<130x1x4080xf64> loc(#loc1189) // <-- not needed
%577 = stablehlo.concatenate %573, %254, %576, dim = 1 {mhlo.sharding = "{devices=[1,4,8]<=[8,4]T(1,0)}"} : (tensor<130x1x4080xf64>, tensor<130x2032x4080xf64>, tensor<130x1x4080xf64>) -> tensor<130x2034x4080xf64> loc(#loc1162)
return %iterArg_102 = %577
```
Contributor guide
No contributing guide indexed for this repository
Research direction
Start with the two StableHLO examples in the issue and trace how the loop-carried tensors produce the shown slice and concatenate operations. Determine how the access analysis should recognize equivalent regions and remove redundant operations; done means the loop no longer contains the unnecessary slices and concats while preserving the demonstrated tensor shapes and results.
Written by the indexing model from the issue text.
Assessment
- Domain
- compilers, performance
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 25/100