[GPU] Teach GPUApplyTilingLevel PartialReduction tiling (#19682)
diff --git a/compiler/src/iree/compiler/Codegen/Common/GPU/GPUApplyTilingLevel.cpp b/compiler/src/iree/compiler/Codegen/Common/GPU/GPUApplyTilingLevel.cpp
index b6f58fe..62a8ac2 100644
--- a/compiler/src/iree/compiler/Codegen/Common/GPU/GPUApplyTilingLevel.cpp
+++ b/compiler/src/iree/compiler/Codegen/Common/GPU/GPUApplyTilingLevel.cpp
@@ -279,6 +279,12 @@
tilingOptions.setMapping(llvm::to_vector(llvm::reverse(mapping)));
}
+ if (tilingLevel == IREE::GPU::TilingLevel::PartialReduction) {
+ tilingOptions.setReductionTilingStrategy(
+ scf::SCFTilingOptions::ReductionTilingStrategy::
+ PartialReductionOuterReduction);
+ }
+
scf::SCFTileAndFuseOptions tileAndFuseOptions;
tileAndFuseOptions.setTilingOptions(tilingOptions);
@@ -288,6 +294,7 @@
-> std::optional<scf::SCFTileAndFuseOptions::ControlFnResult> {
Operation *owner = originalProducer.getOwner();
if (tilingLevel == IREE::GPU::TilingLevel::Reduction ||
+ tilingLevel == IREE::GPU::TilingLevel::PartialReduction ||
tilingLevel == IREE::GPU::TilingLevel::Subgroup) {
// Do not fuse pad in reduction and subgroup tiling. We instead fuse
// pad without zero slice guard as a cleanup pattern.
@@ -298,7 +305,8 @@
bool yieldProducerReplacement = false;
// We dont want this for reduction tiling as it can lead to large tensors
// being yielded.
- if (tilingLevel != IREE::GPU::TilingLevel::Reduction)
+ if (tilingLevel != IREE::GPU::TilingLevel::Reduction &&
+ tilingLevel != IREE::GPU::TilingLevel::PartialReduction)
yieldProducerReplacement = yieldReplacementsFor.contains(owner);
bool shouldFuse = false;
if (auto tilingOwner = dyn_cast<TilingInterface>(owner)) {
@@ -306,7 +314,8 @@
}
// Do not fuse destination operands for reduction tiling.
if (isDestinationOperand &&
- tilingLevel == IREE::GPU::TilingLevel::Reduction) {
+ (tilingLevel == IREE::GPU::TilingLevel::Reduction ||
+ tilingLevel == IREE::GPU::TilingLevel::PartialReduction)) {
shouldFuse = false;
}
if (shouldFuse) {
@@ -395,7 +404,8 @@
if (tilingLevel != IREE::GPU::TilingLevel::Reduction &&
tilingLevel != IREE::GPU::TilingLevel::Thread &&
- tilingLevel != IREE::GPU::TilingLevel::Subgroup) {
+ tilingLevel != IREE::GPU::TilingLevel::Subgroup &&
+ tilingLevel != IREE::GPU::TilingLevel::PartialReduction) {
funcOp.emitError() << "unsupported tiling level: "
<< IREE::GPU::stringifyEnum(tilingLevel) << "\n";
return signalPassFailure();
diff --git a/compiler/src/iree/compiler/Codegen/Common/GPU/Passes.td b/compiler/src/iree/compiler/Codegen/Common/GPU/Passes.td
index 3a71759..7b9f96f 100644
--- a/compiler/src/iree/compiler/Codegen/Common/GPU/Passes.td
+++ b/compiler/src/iree/compiler/Codegen/Common/GPU/Passes.td
@@ -233,6 +233,8 @@
[{llvm::cl::values(
clEnumValN(IREE::GPU::TilingLevel::Reduction, "reduction",
"Tile and fuse all annotated ops to serial loops"),
+ clEnumValN(IREE::GPU::TilingLevel::PartialReduction, "partial_reduction",
+ "Tile and fuse all annotated ops to partial reduuction loops"),
clEnumValN(IREE::GPU::TilingLevel::Thread, "thread",
"Tile and fuse all annotated ops to threads"),
clEnumValN(IREE::GPU::TilingLevel::Subgroup, "subgroup",
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 f47a1ec..ee0c9a2 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
@@ -2,6 +2,7 @@
// RUN: iree-opt --split-input-file --mlir-print-local-scope --pass-pipeline="builtin.module(func.func(iree-codegen-gpu-apply-tiling-level{allow-zero-slices=true}, canonicalize, cse))" %s | FileCheck %s --check-prefix=NOZERO
// 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 --check-prefix=THREAD
// RUN: iree-opt --split-input-file --mlir-print-local-scope --pass-pipeline="builtin.module(func.func(iree-codegen-gpu-apply-tiling-level{tiling-level=subgroup}, canonicalize, cse))" %s | FileCheck %s --check-prefix=SUBGROUP
+// RUN: iree-opt --split-input-file --mlir-print-local-scope --pass-pipeline="builtin.module(func.func(iree-codegen-gpu-apply-tiling-level{tiling-level=partial_reduction}, canonicalize, cse))" %s | FileCheck %s --check-prefix=PARTRED
#config = #iree_gpu.lowering_config<{thread = [2, 16], subgroup = [2, 16]}>
#map = affine_map<(d0, d1) -> (d0, d1)>
@@ -698,3 +699,42 @@
// THREAD: %[[SLICE:.+]] = tensor.extract_slice %{{.*}}[%[[APPLY0]], %[[APPLY1]]] [20, 4] [1, 1]
// THREAD: %[[EXPAND:.+]] = tensor.expand_shape %[[SLICE]] {{\[\[}}0, 1, 2], [3, 4]] output_shape [1, 2, 10, 1, 4]
// THREAD: linalg.exp {{.*}} ins(%[[EXPAND]]
+
+// -----
+
+// Partial reduction tiling tests
+#config = #iree_gpu.lowering_config<{partial_reduction = [0, 8]}>
+#map1 = affine_map<(d0, d1) -> (d0, d1)>
+#map2 = affine_map<(d0, d1) -> (d0)>
+func.func @partial_reduction(%3: tensor<?x?xf32>) -> tensor<?xf32> {
+ %cst = arith.constant 0.000000e+00 : f32
+ %c0 = arith.constant 0 : index
+ %par_dim = tensor.dim %3, %c0 : tensor<?x?xf32>
+ %empty = tensor.empty(%par_dim) : tensor<?xf32>
+ %4 = linalg.fill ins(%cst : f32) outs(%empty : tensor<?xf32>) -> tensor<?xf32>
+ %5 = linalg.generic {
+ indexing_maps = [#map1, #map2],
+ iterator_types = ["parallel", "reduction"]
+ } ins(%3 : tensor<?x?xf32>) outs(%4 : tensor<?xf32>) attrs = {lowering_config = #config} {
+ ^bb0(%in: f32, %out: f32):
+ %7 = arith.addf %in, %out : f32
+ linalg.yield %7 : f32
+ } -> tensor<?xf32>
+ return %5 : tensor<?xf32>
+}
+
+// We only check if the correct tiling implementation was used. We do not
+// check if the tiling implementation itself is correct (it should be tested
+// in the partial tiling unit tests).
+// PARTRED-LABEL: func.func @partial_reduction
+// PARTRED-DAG: %[[DIM0:.+]] = tensor.dim %{{.*}}, %c0
+// PARTRED-DAG: %[[DIM1:.+]] = tensor.dim %{{.*}}, %c1
+// PARTRED-DAG: %[[FULL:.+]] = linalg.fill {{.*}} tensor<?xf32>
+// PARTRED-DAG: %[[PART:.+]] = linalg.fill {{.*}} tensor<?x8xf32>
+// PARTRED: %[[OUT:.+]] = scf.for %{{.*}} = %c0 to %[[DIM1]] step %c8 iter_args(%{{.*}} = %[[PART]])
+// PARTRED: linalg.generic
+// PARTRED-SAME: iterator_types = ["parallel", "parallel"]
+// PARTRED-SAME: ins(%{{.*}} : tensor<?x?xf32>) outs(%{{.*}} : tensor<?x?xf32>)
+// PARTRED: scf.yield
+// PARTRED: linalg.reduce ins(%[[OUT]] : tensor<?x8xf32>)
+// PARTRED-SAME: outs(%[[FULL]] : tensor<?xf32>)
