[LinalgExt] Argsort leads to dynamic-sized stack allocation
bug 🐞
### What happened?
I was trying to compile a model containing the [torch.argsort](https://docs.pytorch.org/docs/2.13/generated/torch.argsort.html) operator, applied to a tensor of dynamic shape. This resulted in the following error:
> error: 'memref.alloca' op expected no unbounded stack allocations
### Steps to reproduce your issue
1) Save this reproducer as `argsort.mlir`
```mlir
#map0 = affine_map<(d0) -> (d0)>
func.func @argsort(%arg0: tensor<?xi64>) -> tensor<?xi64> {
%c0 = arith.constant 0 : index
%dim = tensor.dim %arg0, %c0 : tensor<?xi64>
%idx_init = tensor.empty(%dim) : tensor<?xi64>
%idx = linalg.generic {indexing_maps = [#map0], iterator_types = ["parallel"]} outs(%idx_init : tensor<?xi64>) {
^bb0(%out: i64):
%i = linalg.index 0 : index
%iv = arith.index_cast %i : index to i64
linalg.yield %iv : i64
} -> tensor<?xi64>
%sorted, %indices = iree_linalg_ext.sort dimension(0) outs(%arg0, %idx : tensor<?xi64>, tensor<?xi64>) {
^bb0(%a0: i64, %a1: i64, %b0: i64, %b1: i64):
%cmp = arith.cmpi sle, %a0, %a1 : i64
iree_linalg_ext.yield %cmp : i1
} -> tensor<?xi64>, tensor<?xi64>
return %indices : tensor<?xi64>
}
```
2) Try to compile via:
```sh
iree-compile ./path/to/argsort.mlir \
--iree-hal-target-backends=llvm-cpu \
--iree-llvmcpu-target-cpu=generic \
-o /tmp/argsort.vmfb \
--dump-compilation-phases-to=/tmp/argsort
```
### What component(s) does this issue relate to?
Compiler
### Version information
The issue has been present for a long time; confirmed it's still not solved as of `iree-3.12.0rc20260721` (today).
### Additional context
This is due to the first argument of sort being passed as readonly, since it is not used outside:
```mlir
%3 = flow.dispatch.workgroups[%0](%1, %2, %0) : (tensor<?xi64>{%0}, tensor<?xi64>{%0}, index) -> %2{%0} =
(%arg1: !iree_tensor_ext.dispatch.tensor<readonly:tensor<?xi64>>, %arg2: !iree_tensor_ext.dispatch.tensor<readwrite:tensor<?xi64>>, %arg3: index) {
%5 = iree_tensor_ext.dispatch.workload.ordinal %arg3, 0 : index
%6 = iree_tensor_ext.dispatch.tensor.load %arg1, offsets = [0], sizes = [%5], strides = [1] : !iree_tensor_ext.dispatch.tensor<readonly:tensor<?xi64>>{%5} -> tensor<?xi64>
%7 = iree_tensor_ext.dispatch.tensor.load %arg2, offsets = [0], sizes = [%5], strides = [1] : !iree_tensor_ext.dispatch.tensor<readwrite:tensor<?xi64>>{%5} -> tensor<?xi64>
%8:2 = iree_linalg_ext.sort dimension(0) outs(%6, %7 : tensor<?xi64>, tensor<?xi64>) {
^bb0(%arg4: i64, %arg5: i64, %arg6: i64, %arg7: i64):
%9 = arith.cmpi sle, %arg4, %arg5 : i64
iree_linalg_ext.yield %9 : i1
} -> tensor<?xi64>, tensor<?xi64>
iree_tensor_ext.dispatch.tensor.store %8#1, %arg2, offsets = [0], sizes = [%5], strides = [1] : tensor<?xi64> -> !iree_tensor_ext.dispatch.tensor<readwrite:tensor<?xi64>>{%5}
flow.return
} count(%arg1: index) -> (index, index, index) {
%x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_slice(%arg1)
flow.return %x, %y, %z : index, index, index
}
```
This causes a temporary copy to be allocated in codegen, since the sort op needs to modify all arguments in-place.
0 条评论