ITADN

[LinalgExt] Argsort leads to dynamic-sized stack allocation

#24735Openmbknust 创建于 28 天前
bug 🐞
M
mbknustcommented
### 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 条评论