Revert "[SPIRV] Reduce the number of warps used by subgroup reduce"

This reverts commit 31e7635.
nod-ai · Aug 30, 2023 · 974c895 · 974c895
1 parent d2d37b9
commit 974c895
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 6 deletions.
diff --git a/compiler/src/iree/compiler/Codegen/SPIRV/KernelConfig.cpp b/compiler/src/iree/compiler/Codegen/SPIRV/KernelConfig.cpp
@@ -40,7 +40,6 @@ using llvm::APIntOps::GreatestCommonDivisor;
 constexpr unsigned numTilesPerSubgroupDimK = 2;
 
 constexpr int kMaxVectorNumBits = 128;
-constexpr int kPreferredReductionVectorUnrollAmount = 8;
 
 namespace mlir {
 namespace iree_compiler {
@@ -1281,8 +1280,7 @@ static LogicalResult setReductionConfig(const spirv::TargetEnv &targetEnv,
     return failure();
 
   // Let each thread handle `vectorSize` elements.
-  unsigned vectorSize =
-      kPreferredReductionVectorUnrollAmount * kMaxVectorNumBits / bitWidth;
+  unsigned vectorSize = kMaxVectorNumBits / bitWidth;
   while ((dimSize / vectorSize) % subgroupSize != 0)
     vectorSize /= 2;
 

diff --git a/compiler/src/iree/compiler/Codegen/SPIRV/test/config_default_matvec.mlir b/compiler/src/iree/compiler/Codegen/SPIRV/test/config_default_matvec.mlir
@@ -140,7 +140,7 @@ hal.executable @i4_dequant_matvec_f32 {
 //   CHECK-DAG: #[[$TRANSLATION:.+]] = #iree_codegen.translation_info<SPIRVSubgroupReduce>
 // CHECK-LABEL: hal.executable.export public @i4_dequant_matvec_f32
 //  CHECK-SAME:   translation_info = #[[$TRANSLATION]]
-//  CHECK-SAME:   workgroup_size = [128 : index, 1 : index, 1 : index]
+//  CHECK-SAME:   workgroup_size = [1024 : index, 1 : index, 1 : index]
 //       CHECK: func.func @i4_dequant_matvec_f32()
 //       CHECK:   linalg.generic
 //  CHECK-SAME:     lowering_config = #[[$CONFIG]]

diff --git a/compiler/src/iree/compiler/Codegen/SPIRV/test/config_default_reduction.mlir b/compiler/src/iree/compiler/Codegen/SPIRV/test/config_default_reduction.mlir
@@ -47,7 +47,7 @@ hal.executable private @subgroup_reduce_f32 {
 //  CHECK-DAG: #[[TRANSLATION:.+]] = #iree_codegen.translation_info<SPIRVSubgroupReduce>
 //      CHECK: hal.executable.export public @subgroup_reduce_f32
 // CHECK-SAME:   translation_info = #[[TRANSLATION]]
-// CHECK-SAME:   workgroup_size = [16 : index, 1 : index, 1 : index]
+// CHECK-SAME:   workgroup_size = [128 : index, 1 : index, 1 : index]
 //      CHECK: func.func @subgroup_reduce_f32()
 //      CHECK:   linalg.generic
 // CHECK-SAME:     lowering_config = #[[CONFIG]]
@@ -105,7 +105,7 @@ hal.executable private @subgroup_reduce_f16 {
 //  CHECK-DAG: #[[TRANSLATION:.+]] = #iree_codegen.translation_info<SPIRVSubgroupReduce>
 //      CHECK: hal.executable.export public @subgroup_reduce_f16
 // CHECK-SAME:   translation_info = #[[TRANSLATION]]
-// CHECK-SAME:   workgroup_size = [64 : index, 1 : index, 1 : index]
+// CHECK-SAME:   workgroup_size = [512 : index, 1 : index, 1 : index]
 //      CHECK: func.func @subgroup_reduce_f16()
 //      CHECK:   linalg.generic
 // CHECK-SAME:     lowering_config = #[[CONFIG]]