[GPU] Add thread tile size inference for scatter (#19694)
This adds basic logic for picking thread distribution tile sizes for
scatter ops. This allows fallback distribution to kick in for scatter.
Additionally splits derived_thread_config tests into their own file.
The existing test for GPUApplyTilingLevel has grown rather large. The
tests for derived_thread_config within that file aren't really testing
the pass but the logic in how we pick tile sizes. As a result it makes
more sense to split out those tests into a separate file. In the future
the logic for derived_thread_configs will probably need to move
elsewhere but this is good cleanup to start.
diff --git a/compiler/src/iree/compiler/Codegen/Common/GPU/test/BUILD.bazel b/compiler/src/iree/compiler/Codegen/Common/GPU/test/BUILD.bazel
index a43b450..473ca95 100644
--- a/compiler/src/iree/compiler/Codegen/Common/GPU/test/BUILD.bazel
+++ b/compiler/src/iree/compiler/Codegen/Common/GPU/test/BUILD.bazel
@@ -18,6 +18,7 @@
name = "lit",
srcs = enforce_glob(
[
+ "gpu_apply_derived_thread_config.mlir",
"gpu_apply_tiling_level.mlir",
"gpu_check_resource_usage.mlir",
"gpu_create_fast_slow_path.mlir",
diff --git a/compiler/src/iree/compiler/Codegen/Common/GPU/test/CMakeLists.txt b/compiler/src/iree/compiler/Codegen/Common/GPU/test/CMakeLists.txt
index 8efb734..cf1cf03 100644
--- a/compiler/src/iree/compiler/Codegen/Common/GPU/test/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Codegen/Common/GPU/test/CMakeLists.txt
@@ -14,6 +14,7 @@
NAME
lit
SRCS
+ "gpu_apply_derived_thread_config.mlir"
"gpu_apply_tiling_level.mlir"
"gpu_check_resource_usage.mlir"
"gpu_combine_value_barriers.mlir"
diff --git a/compiler/src/iree/compiler/Codegen/Common/GPU/test/gpu_apply_derived_thread_config.mlir b/compiler/src/iree/compiler/Codegen/Common/GPU/test/gpu_apply_derived_thread_config.mlir
new file mode 100644
index 0000000..dbaa3c2
--- /dev/null
+++ b/compiler/src/iree/compiler/Codegen/Common/GPU/test/gpu_apply_derived_thread_config.mlir
@@ -0,0 +1,214 @@
+// RUN: iree-opt --split-input-file --mlir-print-local-scope --pass-pipeline="builtin.module(func.func(iree-codegen-gpu-apply-tiling-level{tiling-level=thread}, canonicalize, cse))" %s | FileCheck %s
+
+#config = #iree_gpu.derived_thread_config
+#map = affine_map<(d0, d1) -> (d0, d1)>
+module {
+ func.func @inferred_add_tensor(%3: tensor<64x256xf32>, %4: tensor<64x256xf32>, %5: tensor<64x256xf32>) -> tensor<64x256xf32>
+ attributes {
+ translation_info = #iree_codegen.translation_info<pipeline = LLVMGPUTileAndFuse workgroup_size = [16, 32, 1] subgroup_size = 64, {}>
+ } {
+ %6 = linalg.generic {
+ indexing_maps = [#map, #map, #map],
+ iterator_types = ["parallel", "parallel"]
+ } ins(%3, %4 : tensor<64x256xf32>, tensor<64x256xf32>) outs(%5 : tensor<64x256xf32>) attrs = {lowering_config = #config} {
+ ^bb0(%in: f32, %in_0: f32, %out: f32):
+ %7 = arith.addf %in, %in_0 : f32
+ linalg.yield %7 : f32
+ } -> tensor<64x256xf32>
+ return %6 : tensor<64x256xf32>
+ }
+}
+
+// CHECK-LABEL: func.func @inferred_add_tensor
+// CHECK: scf.forall ({{.*}}) = (0, 0) to (64, 256) step (1, 4)
+// CHECK: linalg.generic {{.*}} ins(%{{.*}}: tensor<1x4xf32>, tensor<1x4xf32>)
+// CHECK: scf.forall.in_parallel
+// CHECK: mapping = [#gpu.thread<linear_dim_1>, #gpu.thread<linear_dim_0>]
+
+// -----
+
+#config = #iree_gpu.derived_thread_config
+#map = affine_map<(d0, d1) -> (d0, d1)>
+module {
+ func.func @inferred_dynamic(%3: tensor<?x?xf32>, %4: tensor<?x?xf32>, %5: tensor<?x?xf32>) -> tensor<?x?xf32>
+ attributes {
+ translation_info = #iree_codegen.translation_info<pipeline = LLVMGPUTileAndFuse workgroup_size = [16, 32, 1] subgroup_size = 64, {}>
+ } {
+ %6 = linalg.generic {
+ indexing_maps = [#map, #map, #map],
+ iterator_types = ["parallel", "parallel"]
+ } ins(%3, %4 : tensor<?x?xf32>, tensor<?x?xf32>) outs(%5 : tensor<?x?xf32>) attrs = {lowering_config = #config} {
+ ^bb0(%in: f32, %in_0: f32, %out: f32):
+ %7 = arith.addf %in, %in_0 : f32
+ linalg.yield %7 : f32
+ } -> tensor<?x?xf32>
+ return %6 : tensor<?x?xf32>
+ }
+}
+
+// CHECK-LABEL: func.func @inferred_dynamic
+// CHECK-SAME: %[[A:[A-Za-z0-9]+]]: tensor<?x?xf32>
+// CHECK-DAG: %[[DIM0:.+]] = tensor.dim %[[A]], %c0 : tensor<?x?xf32>
+// CHECK-DAG: %[[DIM1:.+]] = tensor.dim %[[A]], %c1 : tensor<?x?xf32>
+// CHECK: scf.forall ({{.*}}) = (0, 0) to (%[[DIM0]], %[[DIM1]]) step (1, 4)
+// CHECK: linalg.generic {{.*}} ins(%{{.*}}: tensor<1x?xf32>, tensor<1x?xf32>)
+// CHECK: scf.forall.in_parallel
+// CHECK: mapping = [#gpu.thread<linear_dim_1>, #gpu.thread<linear_dim_0>]
+
+// -----
+
+#config = #iree_gpu.derived_thread_config
+#map = affine_map<(d0, d1) -> (d0, d1)>
+module {
+ func.func @inferred_small_inner_dim(%3: tensor<8x2xf32>, %4: tensor<8x2xf32>, %5: tensor<8x2xf32>) -> tensor<8x2xf32>
+ attributes {
+ translation_info = #iree_codegen.translation_info<pipeline = LLVMGPUTileAndFuse workgroup_size = [16, 32, 1] subgroup_size = 64, {}>
+ } {
+ %6 = linalg.generic {
+ indexing_maps = [#map, #map, #map],
+ iterator_types = ["parallel", "parallel"]
+ } ins(%3, %4 : tensor<8x2xf32>, tensor<8x2xf32>) outs(%5 : tensor<8x2xf32>) attrs = {lowering_config = #config} {
+ ^bb0(%in: f32, %in_0: f32, %out: f32):
+ %7 = arith.addf %in, %in_0 : f32
+ linalg.yield %7 : f32
+ } -> tensor<8x2xf32>
+ return %6 : tensor<8x2xf32>
+ }
+}
+
+// CHECK-LABEL: func.func @inferred_small_inner_dim
+// CHECK: scf.forall ({{.*}}) = (0, 0) to (8, 2) step (1, 2)
+// CHECK: linalg.generic {{.*}} ins(%{{.*}}: tensor<1x2xf32>, tensor<1x2xf32>)
+// CHECK: scf.forall.in_parallel
+// CHECK: mapping = [#gpu.thread<linear_dim_1>, #gpu.thread<linear_dim_0>]
+
+// -----
+
+#map = affine_map<(d0, d1) -> (d0, d1)>
+module {
+ func.func @inferred_small_inner_dim_fill_vector_sizes(%0: tensor<4x16x8x4x16x2x4xf16>, %1: tensor<4x16x8x4x16x2x4xf16>) -> tensor<4x16x8x4x16x2x4xf16>
+ attributes {
+ translation_info = #iree_codegen.translation_info<pipeline = LLVMGPUTileAndFuse workgroup_size = [256, 1, 1] subgroup_size = 64, {}>
+ } {
+ %2 = linalg.copy {lowering_config = #iree_gpu.derived_thread_config}
+ ins(%0 : tensor<4x16x8x4x16x2x4xf16>)
+ outs(%1 : tensor<4x16x8x4x16x2x4xf16>) -> tensor<4x16x8x4x16x2x4xf16>
+ return %2 : tensor<4x16x8x4x16x2x4xf16>
+ }
+}
+
+// CHECK-LABEL: func.func @inferred_small_inner_dim_fill_vector_sizes
+// CHECK: scf.forall ({{.*}}) = (0, 0, 0, 0, 0, 0, 0) to (4, 16, 8, 4, 16, 2, 4) step (1, 1, 1, 1, 1, 2, 4)
+// CHECK: linalg.copy{{.*}}ins({{.*}} : tensor<1x1x1x1x1x2x4xf16>) outs({{.*}} : tensor<1x1x1x1x1x2x4xf16>)
+// CHECK: scf.forall.in_parallel
+// CHECK: mapping = [#gpu.thread<linear_dim_6>, {{.*}}, #gpu.thread<linear_dim_0>]
+
+// -----
+
+#map = affine_map<(d0, d1) -> (d0, d1)>
+module {
+ func.func @inferred_small_inner_dim_dont_fill_non_contiguous(
+ %0: tensor<4x16x4x4xf16>, %1: tensor<4x16x4x4xf16>) -> tensor<4x16x4x4xf16>
+ attributes {
+ translation_info = #iree_codegen.translation_info<pipeline = LLVMGPUTileAndFuse workgroup_size = [64, 1, 1] subgroup_size = 64, {}>
+ } {
+ %2 = linalg.copy {lowering_config = #iree_gpu.derived_thread_config}
+ ins(%0 : tensor<4x16x4x4xf16>)
+ outs(%1 : tensor<4x16x4x4xf16>) -> tensor<4x16x4x4xf16>
+ return %2 : tensor<4x16x4x4xf16>
+ }
+}
+
+// CHECK-LABEL: func.func @inferred_small_inner_dim_dont_fill_non_contiguous
+// CHECK: scf.forall ({{.*}}) = (0, 0, 0, 0) to (4, 16, 4, 4) step (1, 1, 1, 4)
+// CHECK: linalg.copy{{.*}}ins({{.*}} : tensor<1x1x1x4xf16>) outs({{.*}} : tensor<1x1x1x4xf16>)
+// CHECK: scf.forall.in_parallel
+// CHECK: mapping = [#gpu.thread<linear_dim_3>, {{.*}}, #gpu.thread<linear_dim_0>]
+
+// -----
+
+#map = affine_map<(d0, d1) -> (d0, d1)>
+module {
+ func.func @inferred_unaligned(%0: tensor<70xf16>, %1: tensor<70xf16>) -> tensor<70xf16>
+ attributes {
+ translation_info = #iree_codegen.translation_info<pipeline = LLVMGPUTileAndFuse workgroup_size = [64, 1, 1] subgroup_size = 64, {}>
+ } {
+ %2 = linalg.copy {lowering_config = #iree_gpu.derived_thread_config}
+ ins(%0 : tensor<70xf16>)
+ outs(%1 : tensor<70xf16>) -> tensor<70xf16>
+ return %2 : tensor<70xf16>
+ }
+}
+
+// CHECK-LABEL: func.func @inferred_unaligned
+// CHECK: scf.forall ({{.*}}) = (0) to (70) step (8)
+// CHECK: linalg.copy{{.*}}ins({{.*}} : tensor<?xf16>) outs({{.*}} : tensor<?xf16>)
+// CHECK: scf.forall.in_parallel
+// CHECK: mapping = [#gpu.thread<linear_dim_0>]
+
+// -----
+
+#map = affine_map<(d0, d1) -> (d0, d1)>
+module {
+ func.func @inferred_smaller_load(%0: tensor<128xf16>, %1: tensor<128xf16>) -> tensor<128xf16>
+ attributes {
+ translation_info = #iree_codegen.translation_info<pipeline = LLVMGPUTileAndFuse workgroup_size = [64, 1, 1] subgroup_size = 64, {}>
+ } {
+ %2 = linalg.copy {lowering_config = #iree_gpu.derived_thread_config}
+ ins(%0 : tensor<128xf16>)
+ outs(%1 : tensor<128xf16>) -> tensor<128xf16>
+ return %2 : tensor<128xf16>
+ }
+}
+
+// CHECK-LABEL: func.func @inferred_smaller_load
+// CHECK: scf.forall ({{.*}}) = (0) to (128) step (2)
+// CHECK: linalg.copy{{.*}}ins({{.*}} : tensor<2xf16>) outs({{.*}} : tensor<2xf16>)
+// CHECK: scf.forall.in_parallel
+// CHECK: mapping = [#gpu.thread<linear_dim_0>]
+
+// -----
+
+#config = #iree_gpu.derived_thread_config
+module {
+ func.func @inferred_im2col(%2: tensor<2x34x34x128xf16>, %3: tensor<2x128x8xf16>) -> tensor<2x128x8xf16>
+ attributes {
+ translation_info = #iree_codegen.translation_info<pipeline = LLVMGPUTileAndFuse workgroup_size = [16, 32, 1] subgroup_size = 64, {}>
+ } {
+ %4 = iree_linalg_ext.im2col {lowering_config = #config}
+ strides = [1, 1] dilations = [1, 1] kernel_size = [3, 3]
+ m_offset = [0] * [1] k_offset = [0] * [1]
+ batch_pos = [0] m_pos = [2, 3] k_pos = [1]
+ ins(%2 : tensor<2x34x34x128xf16>)
+ outs(%3 : tensor<2x128x8xf16>) -> tensor<2x128x8xf16>
+ return %4 : tensor<2x128x8xf16>
+ }
+}
+
+// CHECK-LABEL: func.func @inferred_im2col
+// CHECK: scf.forall ({{.*}}) = (0, 0, 0) to (2, 128, 8) step (1, 1, 4)
+// CHECK: iree_linalg_ext.im2col {{.*}} ins(%{{.*}}: tensor<1x34x34x128xf16>) outs({{.*}}: tensor<1x1x4xf16>)
+// CHECK: scf.forall.in_parallel
+// CHECK: mapping = [#gpu.thread<linear_dim_2>, #gpu.thread<linear_dim_1>, #gpu.thread<linear_dim_0>]
+
+// -----
+
+#config = #iree_gpu.derived_thread_config
+func.func @scatter(%arg0: tensor<3x32x16xf32>, %arg1: tensor<3x1xi32>) -> tensor<3x32x16xf32>
+ attributes {
+ translation_info = #iree_codegen.translation_info<pipeline = LLVMGPUTileAndFuse workgroup_size = [64, 1, 1] subgroup_size = 64, {}>
+ } {
+ %cst = arith.constant 0.000000e+00 : f32
+ %0 = tensor.empty() : tensor<3x32x16xf32>
+ %1 = iree_linalg_ext.scatter {lowering_config = #config} dimension_map = [0] unique_indices(true)
+ ins(%arg0, %arg1 : tensor<3x32x16xf32>, tensor<3x1xi32>) outs(%0 : tensor<3x32x16xf32>) {
+ ^bb0(%arg2: f32, %arg3: f32):
+ iree_linalg_ext.yield %arg2 : f32
+ } -> tensor<3x32x16xf32>
+ return %1 : tensor<3x32x16xf32>
+}
+
+// CHECK-LABEL: func.func @scatter
+// CHECK: scf.forall ({{.*}}) = (0, 0, 0) to (3, 32, 16) step (1, 1, 4)
+// CHECK: linalg_ext.scatter
+// CHECK: scf.forall.in_parallel
diff --git a/compiler/src/iree/compiler/Codegen/Common/GPU/test/gpu_apply_tiling_level.mlir b/compiler/src/iree/compiler/Codegen/Common/GPU/test/gpu_apply_tiling_level.mlir
index ee0c9a2..99a46be 100644
--- a/compiler/src/iree/compiler/Codegen/Common/GPU/test/gpu_apply_tiling_level.mlir
+++ b/compiler/src/iree/compiler/Codegen/Common/GPU/test/gpu_apply_tiling_level.mlir
@@ -206,211 +206,6 @@
// -----
-#config = #iree_gpu.derived_thread_config
-#map = affine_map<(d0, d1) -> (d0, d1)>
-module {
- func.func @inferred_add_tensor(%3: tensor<64x256xf32>, %4: tensor<64x256xf32>, %5: tensor<64x256xf32>) -> tensor<64x256xf32>
- attributes {
- translation_info = #iree_codegen.translation_info<pipeline = LLVMGPUTileAndFuse workgroup_size = [16, 32, 1] subgroup_size = 64, {}>
- } {
- %6 = linalg.generic {
- indexing_maps = [#map, #map, #map],
- iterator_types = ["parallel", "parallel"]
- } ins(%3, %4 : tensor<64x256xf32>, tensor<64x256xf32>) outs(%5 : tensor<64x256xf32>) attrs = {lowering_config = #config} {
- ^bb0(%in: f32, %in_0: f32, %out: f32):
- %7 = arith.addf %in, %in_0 : f32
- linalg.yield %7 : f32
- } -> tensor<64x256xf32>
- return %6 : tensor<64x256xf32>
- }
-}
-
-// Verify that no loops are generated without a reduction configuration.
-// CHECK-LABEL: func.func @inferred_add_tensor
-// CHECK-NOT: scf.for
-
-// THREAD-LABEL: func.func @inferred_add_tensor
-// THREAD: scf.forall ({{.*}}) = (0, 0) to (64, 256) step (1, 4)
-// THREAD: linalg.generic {{.*}} ins(%{{.*}}: tensor<1x4xf32>, tensor<1x4xf32>)
-// THREAD: scf.forall.in_parallel
-// THREAD: mapping = [#gpu.thread<linear_dim_1>, #gpu.thread<linear_dim_0>]
-
-// -----
-
-#config = #iree_gpu.derived_thread_config
-#map = affine_map<(d0, d1) -> (d0, d1)>
-module {
- func.func @inferred_dynamic(%3: tensor<?x?xf32>, %4: tensor<?x?xf32>, %5: tensor<?x?xf32>) -> tensor<?x?xf32>
- attributes {
- translation_info = #iree_codegen.translation_info<pipeline = LLVMGPUTileAndFuse workgroup_size = [16, 32, 1] subgroup_size = 64, {}>
- } {
- %6 = linalg.generic {
- indexing_maps = [#map, #map, #map],
- iterator_types = ["parallel", "parallel"]
- } ins(%3, %4 : tensor<?x?xf32>, tensor<?x?xf32>) outs(%5 : tensor<?x?xf32>) attrs = {lowering_config = #config} {
- ^bb0(%in: f32, %in_0: f32, %out: f32):
- %7 = arith.addf %in, %in_0 : f32
- linalg.yield %7 : f32
- } -> tensor<?x?xf32>
- return %6 : tensor<?x?xf32>
- }
-}
-
-// THREAD-LABEL: func.func @inferred_dynamic
-// THREAD-SAME: %[[A:[A-Za-z0-9]+]]: tensor<?x?xf32>
-// THREAD-DAG: %[[DIM0:.+]] = tensor.dim %[[A]], %c0 : tensor<?x?xf32>
-// THREAD-DAG: %[[DIM1:.+]] = tensor.dim %[[A]], %c1 : tensor<?x?xf32>
-// THREAD: scf.forall ({{.*}}) = (0, 0) to (%[[DIM0]], %[[DIM1]]) step (1, 4)
-// THREAD: linalg.generic {{.*}} ins(%{{.*}}: tensor<1x?xf32>, tensor<1x?xf32>)
-// THREAD: scf.forall.in_parallel
-// THREAD: mapping = [#gpu.thread<linear_dim_1>, #gpu.thread<linear_dim_0>]
-
-// -----
-
-#config = #iree_gpu.derived_thread_config
-#map = affine_map<(d0, d1) -> (d0, d1)>
-module {
- func.func @inferred_small_inner_dim(%3: tensor<8x2xf32>, %4: tensor<8x2xf32>, %5: tensor<8x2xf32>) -> tensor<8x2xf32>
- attributes {
- translation_info = #iree_codegen.translation_info<pipeline = LLVMGPUTileAndFuse workgroup_size = [16, 32, 1] subgroup_size = 64, {}>
- } {
- %6 = linalg.generic {
- indexing_maps = [#map, #map, #map],
- iterator_types = ["parallel", "parallel"]
- } ins(%3, %4 : tensor<8x2xf32>, tensor<8x2xf32>) outs(%5 : tensor<8x2xf32>) attrs = {lowering_config = #config} {
- ^bb0(%in: f32, %in_0: f32, %out: f32):
- %7 = arith.addf %in, %in_0 : f32
- linalg.yield %7 : f32
- } -> tensor<8x2xf32>
- return %6 : tensor<8x2xf32>
- }
-}
-
-// THREAD-LABEL: func.func @inferred_small_inner_dim
-// THREAD: scf.forall ({{.*}}) = (0, 0) to (8, 2) step (1, 2)
-// THREAD: linalg.generic {{.*}} ins(%{{.*}}: tensor<1x2xf32>, tensor<1x2xf32>)
-// THREAD: scf.forall.in_parallel
-// THREAD: mapping = [#gpu.thread<linear_dim_1>, #gpu.thread<linear_dim_0>]
-
-// -----
-
-#config = #iree_gpu.derived_thread_config
-#map = affine_map<(d0, d1) -> (d0, d1)>
-module {
- func.func @inferred_small_inner_dim_fill_vector_sizes(%0: tensor<4x16x8x4x16x2x4xf16>, %1: tensor<4x16x8x4x16x2x4xf16>) -> tensor<4x16x8x4x16x2x4xf16>
- attributes {
- translation_info = #iree_codegen.translation_info<pipeline = LLVMGPUTileAndFuse workgroup_size = [256, 1, 1] subgroup_size = 64, {}>
- } {
- %2 = linalg.copy {lowering_config = #iree_gpu.derived_thread_config}
- ins(%0 : tensor<4x16x8x4x16x2x4xf16>)
- outs(%1 : tensor<4x16x8x4x16x2x4xf16>) -> tensor<4x16x8x4x16x2x4xf16>
- return %2 : tensor<4x16x8x4x16x2x4xf16>
- }
-}
-
-// THREAD-LABEL: func.func @inferred_small_inner_dim_fill_vector_sizes
-// THREAD: scf.forall ({{.*}}) = (0, 0, 0, 0, 0, 0, 0) to (4, 16, 8, 4, 16, 2, 4) step (1, 1, 1, 1, 1, 2, 4)
-// THREAD: linalg.copy{{.*}}ins({{.*}} : tensor<1x1x1x1x1x2x4xf16>) outs({{.*}} : tensor<1x1x1x1x1x2x4xf16>)
-// THREAD: scf.forall.in_parallel
-// THREAD: mapping = [#gpu.thread<linear_dim_6>, {{.*}}, #gpu.thread<linear_dim_0>]
-
-// -----
-
-#config = #iree_gpu.derived_thread_config
-#map = affine_map<(d0, d1) -> (d0, d1)>
-module {
- func.func @inferred_small_inner_dim_dont_fill_non_contiguous(
- %0: tensor<4x16x4x4xf16>, %1: tensor<4x16x4x4xf16>) -> tensor<4x16x4x4xf16>
- attributes {
- translation_info = #iree_codegen.translation_info<pipeline = LLVMGPUTileAndFuse workgroup_size = [64, 1, 1] subgroup_size = 64, {}>
- } {
- %2 = linalg.copy {lowering_config = #iree_gpu.derived_thread_config}
- ins(%0 : tensor<4x16x4x4xf16>)
- outs(%1 : tensor<4x16x4x4xf16>) -> tensor<4x16x4x4xf16>
- return %2 : tensor<4x16x4x4xf16>
- }
-}
-
-// THREAD-LABEL: func.func @inferred_small_inner_dim_dont_fill_non_contiguous
-// THREAD: scf.forall ({{.*}}) = (0, 0, 0, 0) to (4, 16, 4, 4) step (1, 1, 1, 4)
-// THREAD: linalg.copy{{.*}}ins({{.*}} : tensor<1x1x1x4xf16>) outs({{.*}} : tensor<1x1x1x4xf16>)
-// THREAD: scf.forall.in_parallel
-// THREAD: mapping = [#gpu.thread<linear_dim_3>, {{.*}}, #gpu.thread<linear_dim_0>]
-
-// -----
-
-#config = #iree_gpu.derived_thread_config
-#map = affine_map<(d0, d1) -> (d0, d1)>
-module {
- func.func @inferred_unaligned(%0: tensor<70xf16>, %1: tensor<70xf16>) -> tensor<70xf16>
- attributes {
- translation_info = #iree_codegen.translation_info<pipeline = LLVMGPUTileAndFuse workgroup_size = [64, 1, 1] subgroup_size = 64, {}>
- } {
- %2 = linalg.copy {lowering_config = #iree_gpu.derived_thread_config}
- ins(%0 : tensor<70xf16>)
- outs(%1 : tensor<70xf16>) -> tensor<70xf16>
- return %2 : tensor<70xf16>
- }
-}
-
-// THREAD-LABEL: func.func @inferred_unaligned
-// THREAD: scf.forall ({{.*}}) = (0) to (70) step (8)
-// THREAD: linalg.copy{{.*}}ins({{.*}} : tensor<?xf16>) outs({{.*}} : tensor<?xf16>)
-// THREAD: scf.forall.in_parallel
-// THREAD: mapping = [#gpu.thread<linear_dim_0>]
-
-// -----
-
-#config = #iree_gpu.derived_thread_config
-#map = affine_map<(d0, d1) -> (d0, d1)>
-module {
- func.func @inferred_smaller_load(%0: tensor<128xf16>, %1: tensor<128xf16>) -> tensor<128xf16>
- attributes {
- translation_info = #iree_codegen.translation_info<pipeline = LLVMGPUTileAndFuse workgroup_size = [64, 1, 1] subgroup_size = 64, {}>
- } {
- %2 = linalg.copy {lowering_config = #iree_gpu.derived_thread_config}
- ins(%0 : tensor<128xf16>)
- outs(%1 : tensor<128xf16>) -> tensor<128xf16>
- return %2 : tensor<128xf16>
- }
-}
-
-// THREAD-LABEL: func.func @inferred_smaller_load
-// THREAD: scf.forall ({{.*}}) = (0) to (128) step (2)
-// THREAD: linalg.copy{{.*}}ins({{.*}} : tensor<2xf16>) outs({{.*}} : tensor<2xf16>)
-// THREAD: scf.forall.in_parallel
-// THREAD: mapping = [#gpu.thread<linear_dim_0>]
-
-// -----
-
-#config = #iree_gpu.derived_thread_config
-module {
- func.func @inferred_im2col(%2: tensor<2x34x34x128xf16>, %3: tensor<2x128x8xf16>) -> tensor<2x128x8xf16>
- attributes {
- translation_info = #iree_codegen.translation_info<pipeline = LLVMGPUTileAndFuse workgroup_size = [16, 32, 1] subgroup_size = 64, {}>
- } {
- %4 = iree_linalg_ext.im2col {lowering_config = #config}
- strides = [1, 1] dilations = [1, 1] kernel_size = [3, 3]
- m_offset = [0] * [1] k_offset = [0] * [1]
- batch_pos = [0] m_pos = [2, 3] k_pos = [1]
- ins(%2 : tensor<2x34x34x128xf16>)
- outs(%3 : tensor<2x128x8xf16>) -> tensor<2x128x8xf16>
- return %4 : tensor<2x128x8xf16>
- }
-}
-
-// Verify that no loops are generated without a reduction configuration.
-// CHECK-LABEL: func.func @inferred_im2col
-// CHECK-NOT: scf.for
-
-// THREAD-LABEL: func.func @inferred_im2col
-// THREAD: scf.forall ({{.*}}) = (0, 0, 0) to (2, 128, 8) step (1, 1, 4)
-// THREAD: iree_linalg_ext.im2col {{.*}} ins(%{{.*}}: tensor<1x34x34x128xf16>) outs({{.*}}: tensor<1x1x4xf16>)
-// THREAD: scf.forall.in_parallel
-// THREAD: mapping = [#gpu.thread<linear_dim_2>, #gpu.thread<linear_dim_1>, #gpu.thread<linear_dim_0>]
-
-// -----
-
#config = #iree_gpu.lowering_config<{reduction = [0, 0, 8], subgroup = [2, 4]}>
#contraction_accesses = [
affine_map<(i, j, k) -> (i, k)>,
diff --git a/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/DerivedConfigUtils.cpp b/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/DerivedConfigUtils.cpp
index 6f7b75f..249ba7a 100644
--- a/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/DerivedConfigUtils.cpp
+++ b/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/DerivedConfigUtils.cpp
@@ -153,7 +153,17 @@
.Case([&](IREE::LinalgExt::Im2colOp im2colOp) -> SmallVector<int64_t> {
return deriveIm2colOpThreadTileSizes(im2colOp, numThreads);
})
- .Default([](Operation *op) -> SmallVector<int64_t> { return {}; });
+ .Case([&](IREE::LinalgExt::ScatterOp scatterOp) -> SmallVector<int64_t> {
+ int64_t loopDepth = scatterOp.getLoopIteratorTypes().size();
+ SmallVector<int64_t> loopBounds =
+ scatterOp.getStaticLoopRanges().value_or(
+ SmallVector<int64_t>(loopDepth, ShapedType::kDynamic));
+ int64_t elemBits = scatterOp.getOriginalType().getElementTypeBitWidth();
+ int64_t vectorSize = kPreferredCopyNumBits / elemBits;
+ return getVectorTileSizesFromLoopRanges(loopBounds, numThreads,
+ vectorSize);
+ })
+ .Default([&](Operation *op) -> SmallVector<int64_t> { return {}; });
}
} // namespace mlir::iree_compiler::IREE::GPU
