iree-org / iree-org/iree

Investigate folding tensor slices into linalg ops to reduce work/memory use.

Open
#6,901 1 comment 0 reactions 0 assignees View on GitHub
codegen performance ⚡
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.