EnzymeAD / EnzymeAD/Enzyme-JAX

Slice access analysis to eliminate slices and concats inside loops

Open
#625 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

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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.