Skip to content

Commit 8223e1f

Browse files
hanhanWhanhanW
authored
[compiler] Retire DispatchWorkgroupCountFromDagRootOp. (#24036)
The op is mostly legacy and there are no active uses for a long time. The codegen does not handle the case properly anyway, since it bails out when FromSlice op is not found in ResolveWorkgroupCountHintsPass. Thus, the revision deletes the op for reducing maintenance burden, and we can always bring it back when it becomes useful. Signed-off-by: hanhanW <hanhan0912@gmail.com>
1 parent 4aa6196 commit 8223e1f

15 files changed

Lines changed: 57 additions & 251 deletions

‎compiler/src/iree/compiler/Codegen/LLVMGPU/test/NVVM/pipeline_vector_distribute_sm80.mlir‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@
1616
hal.executable @matmul_256x256x256_f16_f32 {
1717
hal.executable.variant @cuda target(<"cuda", "cuda-nvptx-fb">) {
1818
hal.executable.export public @matmul_256x256x256_f16_f32 layout(#pipeline_layout) count(%arg0: !hal.device, %arg1: index, %arg2 : index) -> (index, index, index) {
19-
%x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_dag_root(%arg1, %arg2)
19+
%x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_slice(%arg1, %arg2)
2020
hal.return %x, %y, %z : index, index, index
2121
}
2222
builtin.module {
@@ -63,7 +63,7 @@ hal.executable.variant @cuda target(<"cuda", "cuda-nvptx-fb">) {
6363
hal.executable @matmul_256x256x256_f16_f16 {
6464
hal.executable.variant @cuda target(<"cuda", "cuda-nvptx-fb">) {
6565
hal.executable.export public @matmul_256x256x256_f16_f16 layout(#pipeline_layout_f16) count(%arg0: !hal.device, %arg1: index, %arg2 : index) -> (index, index, index) {
66-
%x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_dag_root(%arg1, %arg2)
66+
%x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_slice(%arg1, %arg2)
6767
hal.return %x, %y, %z : index, index, index
6868
}
6969
builtin.module {

‎compiler/src/iree/compiler/Codegen/LLVMGPU/test/ROCDL/pipeline_vector_distribute_gfx1100.mlir‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@
1414
hal.executable @matmul_256x256x256_f16_f32 {
1515
hal.executable.variant @rocm target(<"rocm", "rocm-hsaco-fb">) {
1616
hal.executable.export public @matmul_256x256x256_f16_f32 layout(#pipeline_layout) count(%arg0: !hal.device, %arg1: index, %arg2 : index) -> (index, index, index) {
17-
%x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_dag_root(%arg1, %arg2)
17+
%x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_slice(%arg1, %arg2)
1818
hal.return %x, %y, %z : index, index, index
1919
}
2020
builtin.module {
@@ -60,7 +60,7 @@ hal.executable.variant @rocm target(<"rocm", "rocm-hsaco-fb">) {
6060
hal.executable @matmul_256x256x256_f16_f16 {
6161
hal.executable.variant @rocm target(<"rocm", "rocm-hsaco-fb">) {
6262
hal.executable.export public @matmul_256x256x256_f16_f16 layout(#pipeline_layout) count(%arg0: !hal.device, %arg1: index, %arg2 : index) -> (index, index, index) {
63-
%x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_dag_root(%arg1, %arg2)
63+
%x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_slice(%arg1, %arg2)
6464
hal.return %x, %y, %z : index, index, index
6565
}
6666
builtin.module {

‎compiler/src/iree/compiler/Codegen/LLVMGPU/test/ROCDL/pipeline_vector_distribute_gfx942.mlir‎

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@
1919
hal.executable @matmul_256x256x256_f16_f32 {
2020
hal.executable.variant @rocm target(<"rocm", "rocm-hsaco-fb">) {
2121
hal.executable.export public @matmul_256x256x256_f16_f32 layout(#pipeline_layout) count(%arg0: !hal.device, %arg1: index, %arg2 : index) -> (index, index, index) {
22-
%x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_dag_root(%arg1, %arg2)
22+
%x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_slice(%arg1, %arg2)
2323
hal.return %x, %y, %z : index, index, index
2424
}
2525
builtin.module {
@@ -64,7 +64,7 @@ hal.executable.variant @rocm target(<"rocm", "rocm-hsaco-fb">) {
6464
hal.executable @matmul_256x256x256_f16_f32 {
6565
hal.executable.variant @rocm target(<"rocm", "rocm-hsaco-fb">) {
6666
hal.executable.export public @matmul_256x256x256_f16_f32 layout(#pipeline_layout) count(%arg0: !hal.device, %arg1: index, %arg2 : index) -> (index, index, index) {
67-
%x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_dag_root(%arg1, %arg2)
67+
%x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_slice(%arg1, %arg2)
6868
hal.return %x, %y, %z : index, index, index
6969
}
7070
builtin.module {
@@ -222,7 +222,7 @@ hal.executable @matmul_multiple_k {
222222
hal.executable @matmul_256x256x256_16x16x32_f8_f32 {
223223
hal.executable.variant @rocm target(<"rocm", "rocm-hsaco-fb">) {
224224
hal.executable.export public @matmul_256x256x256_16x16x32_f8_f32 layout(#pipeline_layout) count(%arg0: !hal.device, %arg1: index, %arg2 : index) -> (index, index, index) {
225-
%x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_dag_root(%arg1, %arg2)
225+
%x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_slice(%arg1, %arg2)
226226
hal.return %x, %y, %z : index, index, index
227227
}
228228
builtin.module {
@@ -267,7 +267,7 @@ hal.executable.variant @rocm target(<"rocm", "rocm-hsaco-fb">) {
267267
hal.executable @matmul_256x256x256_i8_i32 {
268268
hal.executable.variant @rocm target(<"rocm", "rocm-hsaco-fb">) {
269269
hal.executable.export public @matmul_256x256x256_i8_i32 layout(#pipeline_layout) count(%arg0: !hal.device, %arg1: index, %arg2 : index) -> (index, index, index) {
270-
%x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_dag_root(%arg1, %arg2)
270+
%x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_slice(%arg1, %arg2)
271271
hal.return %x, %y, %z : index, index, index
272272
}
273273
builtin.module {
@@ -312,7 +312,7 @@ hal.executable.variant @rocm target(<"rocm", "rocm-hsaco-fb">) {
312312
hal.executable @matmul_256x256x256_32x32x16_f8_f32 {
313313
hal.executable.variant @rocm target(<"rocm", "rocm-hsaco-fb">) {
314314
hal.executable.export public @matmul_256x256x256_32x32x16_f8_f32 layout(#pipeline_layout) count(%arg0: !hal.device, %arg1: index, %arg2 : index) -> (index, index, index) {
315-
%x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_dag_root(%arg1, %arg2)
315+
%x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_slice(%arg1, %arg2)
316316
hal.return %x, %y, %z : index, index, index
317317
}
318318
builtin.module {
@@ -357,7 +357,7 @@ hal.executable.variant @rocm target(<"rocm", "rocm-hsaco-fb">) {
357357
hal.executable @matmul_transpose_b_256x256x256_i8_i32 {
358358
hal.executable.variant @rocm target(<"rocm", "rocm-hsaco-fb">) {
359359
hal.executable.export public @matmul_transpose_b_256x256x256_i8_i32 layout(#pipeline_layout) count(%arg0: !hal.device, %arg1: index, %arg2 : index) -> (index, index, index) {
360-
%x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_dag_root(%arg1, %arg2)
360+
%x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_slice(%arg1, %arg2)
361361
hal.return %x, %y, %z : index, index, index
362362
}
363363
builtin.module {
@@ -576,7 +576,7 @@ hal.executable public @contract_schedule_considering_read_layout {
576576
hal.executable @virtual_intrinsic_256x256x256_16x16x32xf8E4M3FNUZ_f32 {
577577
hal.executable.variant @rocm target(<"rocm", "rocm-hsaco-fb">) {
578578
hal.executable.export public @virtual_intrinsic_256x256x256_16x16x32xf8E4M3FNUZ_f32 layout(#pipeline_layout) count(%arg0: !hal.device, %arg1: index, %arg2 : index) -> (index, index, index) {
579-
%x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_dag_root(%arg1, %arg2)
579+
%x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_slice(%arg1, %arg2)
580580
hal.return %x, %y, %z : index, index, index
581581
}
582582
builtin.module {
@@ -636,7 +636,7 @@ hal.executable.variant @rocm target(<"rocm", "rocm-hsaco-fb">) {
636636
hal.executable @virtual_intrinsic_256x256x256_32x32x16_f8E4M3FNUZ_f32 {
637637
hal.executable.variant @rocm target(<"rocm", "rocm-hsaco-fb">) {
638638
hal.executable.export public @virtual_intrinsic_256x256x256_32x32x16_f8E4M3FNUZ_f32 layout(#pipeline_layout) count(%arg0: !hal.device, %arg1: index, %arg2 : index) -> (index, index, index) {
639-
%x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_dag_root(%arg1, %arg2)
639+
%x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_slice(%arg1, %arg2)
640640
hal.return %x, %y, %z : index, index, index
641641
}
642642
builtin.module {
@@ -683,7 +683,7 @@ hal.executable.variant @rocm target(<"rocm", "rocm-hsaco-fb">) {
683683
hal.executable @virtual_intrinsic_256x256x256_16x16x32xf8E4M3FNUZ_f32 {
684684
hal.executable.variant @rocm target(<"rocm", "rocm-hsaco-fb">) {
685685
hal.executable.export public @virtual_intrinsic_256x256x256_16x16x32xf8E4M3FNUZ_f32 layout(#pipeline_layout) count(%arg0: !hal.device, %arg1: index, %arg2 : index) -> (index, index, index) {
686-
%x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_dag_root(%arg1, %arg2)
686+
%x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_slice(%arg1, %arg2)
687687
hal.return %x, %y, %z : index, index, index
688688
}
689689
builtin.module {
@@ -1216,7 +1216,7 @@ hal.executable private @matvec_dispatch_0 {
12161216
hal.executable public @matmul_map_store {
12171217
hal.executable.variant public @rocm target(<"rocm", "rocm-hsaco-fb">) {
12181218
hal.executable.export public @matmul_map_store layout(#pipeline_layout) count(%arg0: !hal.device, %arg1: index, %arg2: index) -> (index, index, index) {
1219-
%x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_dag_root(%arg1, %arg2)
1219+
%x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_slice(%arg1, %arg2)
12201220
hal.return %x, %y, %z : index, index, index
12211221
}
12221222
builtin.module {

‎compiler/src/iree/compiler/Codegen/LLVMGPU/test/ROCDL/pipeline_vector_distribute_gfx950.mlir‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@
1818
hal.executable @matmul_256x256x512_f16_f32 {
1919
hal.executable.variant @rocm target(<"rocm", "rocm-hsaco-fb">) {
2020
hal.executable.export @matmul_256x256x512_f16_f32 layout(#pipeline_layout) count(%arg0: !hal.device, %arg1: index, %arg2 : index) -> (index, index, index) {
21-
%x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_dag_root(%arg1, %arg2)
21+
%x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_slice(%arg1, %arg2)
2222
hal.return %x, %y, %z : index, index, index
2323
}
2424
builtin.module {
@@ -68,7 +68,7 @@ hal.executable.variant @rocm target(<"rocm", "rocm-hsaco-fb">) {
6868
hal.executable @matmul_256x256x512_f16_f32 {
6969
hal.executable.variant @rocm target(<"rocm", "rocm-hsaco-fb">) {
7070
hal.executable.export @matmul_256x256x512_f16_f32 layout(#pipeline_layout) count(%arg0: !hal.device, %arg1: index, %arg2 : index) -> (index, index, index) {
71-
%x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_dag_root(%arg1, %arg2)
71+
%x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_slice(%arg1, %arg2)
7272
hal.return %x, %y, %z : index, index, index
7373
}
7474
builtin.module {
@@ -241,7 +241,7 @@ hal.executable @matmul_multiple_k {
241241
hal.executable @matmul_256x256x256_16x16x128_f8_f32 {
242242
hal.executable.variant @rocm target(<"rocm", "rocm-hsaco-fb">) {
243243
hal.executable.export @matmul_256x256x256_16x16x32_f8_f32 layout(#pipeline_layout) count(%arg0: !hal.device, %arg1: index, %arg2 : index) -> (index, index, index) {
244-
%x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_dag_root(%arg1, %arg2)
244+
%x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_slice(%arg1, %arg2)
245245
hal.return %x, %y, %z : index, index, index
246246
}
247247
builtin.module {
@@ -292,7 +292,7 @@ hal.executable.variant @rocm target(<"rocm", "rocm-hsaco-fb">) {
292292
hal.executable @matmul_256x256x256_i8_i32 {
293293
hal.executable.variant @rocm target(<"rocm", "rocm-hsaco-fb">) {
294294
hal.executable.export @matmul_256x256x256_i8_i32 layout(#pipeline_layout) count(%arg0: !hal.device, %arg1: index, %arg2 : index) -> (index, index, index) {
295-
%x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_dag_root(%arg1, %arg2)
295+
%x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_slice(%arg1, %arg2)
296296
hal.return %x, %y, %z : index, index, index
297297
}
298298
builtin.module {
@@ -343,7 +343,7 @@ hal.executable.variant @rocm target(<"rocm", "rocm-hsaco-fb">) {
343343
hal.executable @matmul_256x256x256_32x32x64_f8_f32 {
344344
hal.executable.variant @rocm target(<"rocm", "rocm-hsaco-fb">) {
345345
hal.executable.export @matmul_256x256x256_32x32x64_f8_f32 layout(#pipeline_layout) count(%arg0: !hal.device, %arg1: index, %arg2 : index) -> (index, index, index) {
346-
%x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_dag_root(%arg1, %arg2)
346+
%x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_slice(%arg1, %arg2)
347347
hal.return %x, %y, %z : index, index, index
348348
}
349349
builtin.module {
@@ -394,7 +394,7 @@ hal.executable.variant @rocm target(<"rocm", "rocm-hsaco-fb">) {
394394
hal.executable @matmul_transpose_b_256x256x256_i8_i32 {
395395
hal.executable.variant @rocm target(<"rocm", "rocm-hsaco-fb">) {
396396
hal.executable.export @matmul_transpose_b_256x256x256_i8_i32 layout(#pipeline_layout) count(%arg0: !hal.device, %arg1: index, %arg2 : index) -> (index, index, index) {
397-
%x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_dag_root(%arg1, %arg2)
397+
%x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_slice(%arg1, %arg2)
398398
hal.return %x, %y, %z : index, index, index
399399
}
400400
builtin.module {

‎compiler/src/iree/compiler/Codegen/LLVMGPU/test/conv_pipeline_test_cuda.mlir‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@
1111
hal.executable private @conv2d_1x230x230x3_7x7x3x64_dispatch_0 {
1212
hal.executable.variant public @cuda_nvptx_fb target(#executable_target_cuda_nvptx_fb) {
1313
hal.executable.export public @conv2d_1x230x230x3_7x7x3x64 ordinal(0) layout(#pipeline_layout) count(%arg0: !hal.device, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: index, %arg7: index) -> (index, index, index) {
14-
%x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_dag_root(%arg1, %arg2, %arg3, %arg4, %arg5, %arg6, %arg7)
14+
%x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_slice(%arg1, %arg2, %arg3, %arg4, %arg5, %arg6, %arg7)
1515
hal.return %x, %y, %z : index, index, index
1616
}
1717
builtin.module {
@@ -55,7 +55,7 @@ hal.executable private @conv2d_1x230x230x3_7x7x3x64_dispatch_0 {
5555
hal.executable private @conv_nchw_dispatch_0 {
5656
hal.executable.variant public @cuda_nvptx_fb target(#executable_target_cuda_nvptx_fb) {
5757
hal.executable.export public @conv_nchw ordinal(0) layout(#pipeline_layout) count(%arg0: !hal.device, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: index, %arg7: index) -> (index, index, index) {
58-
%x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_dag_root(%arg1, %arg2, %arg3, %arg4, %arg5, %arg6, %arg7)
58+
%x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_slice(%arg1, %arg2, %arg3, %arg4, %arg5, %arg6, %arg7)
5959
hal.return %x, %y, %z : index, index, index
6060
}
6161
builtin.module {

‎compiler/src/iree/compiler/Codegen/LLVMGPU/test/reduction_pipeline_cuda.mlir‎

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@
77
hal.executable @warp_reduction_dispatch {
88
hal.executable.variant @cuda target(<"cuda", "cuda-nvptx-fb">) {
99
hal.executable.export public @warp_reduction_dispatch layout(#pipeline_layout) count(%arg0: !hal.device, %arg1: index, %arg2 : index) -> (index, index, index) {
10-
%x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_dag_root(%arg1, %arg2)
10+
%x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_slice(%arg1, %arg2)
1111
hal.return %x, %y, %z : index, index, index
1212
}
1313
builtin.module {
@@ -61,7 +61,7 @@ hal.executable.variant @cuda target(<"cuda", "cuda-nvptx-fb">) {
6161
hal.executable @warp_reduction_broadcast_dispatch {
6262
hal.executable.variant @cuda target(<"cuda", "cuda-nvptx-fb">) {
6363
hal.executable.export public @warp_reduction_broadcast_dispatch layout(#pipeline_layout) count(%arg0: !hal.device, %arg1: index, %arg2 : index) -> (index, index, index) {
64-
%x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_dag_root(%arg1, %arg2)
64+
%x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_slice(%arg1, %arg2)
6565
hal.return %x, %y, %z : index, index, index
6666
}
6767
builtin.module {
@@ -121,7 +121,7 @@ hal.executable.variant @cuda target(<"cuda", "cuda-nvptx-fb">) {
121121
hal.executable @softmax {
122122
hal.executable.variant @cuda target(<"cuda", "cuda-nvptx-fb">) {
123123
hal.executable.export public @softmax layout(#pipeline_layout) count(%arg0: !hal.device, %arg1: index, %arg2 : index) -> (index, index, index) {
124-
%x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_dag_root(%arg1, %arg2)
124+
%x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_slice(%arg1, %arg2)
125125
hal.return %x, %y, %z : index, index, index
126126
}
127127
builtin.module {
@@ -178,7 +178,7 @@ hal.executable.variant @cuda target(<"cuda", "cuda-nvptx-fb">) {
178178
hal.executable @softmax_singlesubgroup {
179179
hal.executable.variant @cuda target(<"cuda", "cuda-nvptx-fb">) {
180180
hal.executable.export public @softmax_singlesubgroup layout(#pipeline_layout) count(%arg0: !hal.device, %arg1: index, %arg2 : index) -> (index, index, index) {
181-
%x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_dag_root(%arg1, %arg2)
181+
%x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_slice(%arg1, %arg2)
182182
hal.return %x, %y, %z : index, index, index
183183
}
184184
builtin.module {
@@ -233,7 +233,7 @@ hal.executable.variant @cuda target(<"cuda", "cuda-nvptx-fb">) {
233233
hal.executable @small_reduction {
234234
hal.executable.variant public @cuda_nvptx_fb target(<"cuda", "cuda-nvptx-fb">) {
235235
hal.executable.export public @small_reduction ordinal(0) layout(#pipeline_layout) count(%arg0: !hal.device, %arg1: index, %arg2: index) -> (index, index, index) {
236-
%x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_dag_root(%arg1, %arg2)
236+
%x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_slice(%arg1, %arg2)
237237
hal.return %x, %y, %z : index, index, index
238238
}
239239
builtin.module {
@@ -271,7 +271,7 @@ hal.executable.variant public @cuda_nvptx_fb target(<"cuda", "cuda-nvptx-fb">) {
271271
hal.executable @group_reduction {
272272
hal.executable.variant public @cuda_nvptx_fb target(<"cuda", "cuda-nvptx-fb">) {
273273
hal.executable.export public @group_reduction ordinal(0) layout(#pipeline_layout) count(%arg0: !hal.device, %arg1: index, %arg2: index) -> (index, index, index) {
274-
%x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_dag_root(%arg1, %arg2)
274+
%x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_slice(%arg1, %arg2)
275275
hal.return %x, %y, %z : index, index, index
276276
}
277277
builtin.module {
@@ -307,7 +307,7 @@ hal.executable.variant public @cuda_nvptx_fb target(<"cuda", "cuda-nvptx-fb">) {
307307
hal.executable @group_elementwise_reduction_elementwise {
308308
hal.executable.variant public @cuda_nvptx_fb target(<"cuda", "cuda-nvptx-fb">) {
309309
hal.executable.export public @group_elementwise_reduction_elementwise ordinal(0) layout(#pipeline_layout) count(%arg0: !hal.device, %arg1: index) -> (index, index, index) {
310-
%x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_dag_root(%arg1)
310+
%x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_slice(%arg1)
311311
hal.return %x, %y, %z : index, index, index
312312
}
313313
builtin.module {
@@ -350,7 +350,7 @@ hal.executable.variant public @cuda_nvptx_fb target(<"cuda", "cuda-nvptx-fb">) {
350350
hal.executable @group_reduction_larger {
351351
hal.executable.variant public @cuda_nvptx_fb target(<"cuda", "cuda-nvptx-fb">) {
352352
hal.executable.export public @group_reduction_larger ordinal(0) layout(#pipeline_layout) count(%arg0: !hal.device, %arg1: index, %arg2: index) -> (index, index, index) {
353-
%x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_dag_root(%arg1, %arg2)
353+
%x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_slice(%arg1, %arg2)
354354
hal.return %x, %y, %z : index, index, index
355355
}
356356
builtin.module {
@@ -387,7 +387,7 @@ hal.executable.variant public @cuda_nvptx_fb target(<"cuda", "cuda-nvptx-fb">) {
387387
hal.executable @group_reduction_1d {
388388
hal.executable.variant public @cuda_nvptx_fb target(<"cuda", "cuda-nvptx-fb">) {
389389
hal.executable.export public @group_reduction_1d ordinal(0) layout(#pipeline_layout) count(%arg0: !hal.device, %arg1: index, %arg2: index) -> (index, index, index) {
390-
%x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_dag_root(%arg1, %arg2)
390+
%x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_slice(%arg1, %arg2)
391391
hal.return %x, %y, %z : index, index, index
392392
}
393393
builtin.module {
@@ -423,7 +423,7 @@ hal.executable.variant public @cuda_nvptx_fb target(<"cuda", "cuda-nvptx-fb">) {
423423
hal.executable @group_elementwise_reduction_elementwise_4d {
424424
hal.executable.variant public @cuda_nvptx_fb target(<"cuda", "cuda-nvptx-fb">) {
425425
hal.executable.export public @group_elementwise_reduction_elementwise_4d ordinal(0) layout(#pipeline_layout) count(%arg0: !hal.device, %arg1: index, %arg2: index, %arg3: index) -> (index, index, index) {
426-
%x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_dag_root(%arg1, %arg2, %arg3)
426+
%x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_slice(%arg1, %arg2, %arg3)
427427
hal.return %x, %y, %z : index, index, index
428428
}
429429
builtin.module {

0 commit comments

Comments
 (0)