Investigate folding tensor slices into linalg ops to reduce work/memory use.
- Dominant language
- C++
- Stars
- 3.9k
- Forks
- 1k
- Avg merge
- 4d 16h
- Merged PRs (30d)
- 47
Description
Today we can end up with sequences of linalg ops -> slices. In some of these cases the unsliced output from linalg is discarded after the slice meaning that there was a bunch of extra work performed and memory allocated.
Example:
```mlir
%126 = linalg.generic {indexing_maps = [#map1, #map2, #map1, #map0, #map2, #map2, #map0, #map1], iterator_types = ["parallel", "parallel"]} ins(%119, %__iree_flow___sm_node79__m.layer-3.layer-11._output_dense.bias, %114, %125, %__iree_flow___sm_node80__m.layer-3.layer-11._output_layer_norm.gamma, %__iree_flow___sm_node81__m.layer-3.layer-11._output_layer_norm.beta, %122 : tensor<512x768xf32>, tensor<768xf32>, tensor<512x768xf32>, tensor<512xf32>, tensor<768xf32>, tensor<768xf32>, tensor<512xf32>) outs(%3 : tensor<512x768xf32>) {
^bb0(%arg4: f32, %arg5: f32, %arg6: f32, %arg7: f32, %arg8: f32, %arg9: f32, %arg10: f32, %arg11: f32): // no predecessors
%498 = divf %arg7, %cst_18 : f32
%499 = addf %498, %cst_13 : f32
%500 = math.rsqrt %499 : f32
%501 = mulf %500, %arg8 : f32
%502 = divf %arg10, %cst_18 : f32
%503 = mulf %502, %501 : f32
%504 = subf %arg9, %503 : f32
%505 = linalg.index 0 : index
%506 = linalg.index 1 : index
%507 = index_cast %505 : index to i32
%508 = addi %507, %c12345_i32 : i32
%509 = muli %508, %c1103515245_i32 : i32
%510 = addi %509, %c12345_i32 : i32
%511 = index_cast %506 : index to i32
%512 = addi %511, %510 : i32
%513 = muli %512, %c1103515245_i32 : i32
%514 = addi %513, %c12345_i32 : i32
%515 = uitofp %514 : i32 to f32
%516 = mulf %515, %cst_16 : f32
%517 = addf %516, %cst_15 : f32
%518 = cmpf oge, %517, %cst_17 : f32
%519 = uitofp %518 : i1 to f32
%520 = addf %arg4, %arg5 : f32
%521 = mulf %520, %cst_14 : f32
%522 = mulf %521, %519 : f32
%523 = addf %522, %arg6 : f32
%524 = mulf %523, %501 : f32
%525 = addf %524, %504 : f32
linalg.yield %525 : f32
} -> tensor<512x768xf32>
%127 = linalg.tensor_expand_shape %126 [[0, 1], [2]] : tensor<512x768xf32> into tensor<1x512x768xf32>
%128 = tensor.extract_slice %127[0, 0, 0] [1, 1, 768] [1, 1, 1] : tensor<1x512x768xf32> to tensor<768xf32>
```
Here if the linalg.generic was updated to produce only the range needed by the subsequent slice (1x1x768) it would avoid 512x the work and transient memory. Note that the reshape may confuse things a bit, but propagating slices up to producers through such ops seems like something linalg would excel at :)
Contributor guide
Assessment
This issue has not been assessed yet.