[Codegen][DispatchCreation] Separate dispatch for iree_linalg_ext.scan producers (#24651)
This fix is necessitated by lack of fusion support for LinalgExt Scan -
since the default LLVMCPU tiling does not respect the fusability of
producer ops, nor the limit for stack-bound allocations, we force this
non-fusable "cumulative reduction"-style operation to be dispatched
separately from non-trivial linalg operations.
---------
Co-authored-by: Artem Gindinson <gindinson@roofline.ai>
diff --git a/compiler/src/iree/compiler/Dialect/Flow/Transforms/RegionOpUtils.cpp b/compiler/src/iree/compiler/Dialect/Flow/Transforms/RegionOpUtils.cpp
index be39b20..3ab05b9 100644
--- a/compiler/src/iree/compiler/Dialect/Flow/Transforms/RegionOpUtils.cpp
+++ b/compiler/src/iree/compiler/Dialect/Flow/Transforms/RegionOpUtils.cpp
@@ -779,6 +779,10 @@
return false;
}
+static bool hasExplicitNonFusableUsers(Operation *op) {
+ return llvm::any_of(op->getUsers(), llvm::IsaPred<IREE::LinalgExt::ScanOp>);
+}
+
/// Operations that are cloned into dispatch regions formed with other
/// operations as roots.
bool isCloneableIntoDispatchOp(Operation *op,
@@ -795,6 +799,11 @@
complex::CreateOp, IREE::Encoding::UnsetEncodingOp>(op)) {
return true;
}
+ // TODO: Tune the cases excluded through hasExplicitNonFusableUsers
+ // condition in a more targeted manner, then remove the condition.
+ if (isa<linalg::LinalgOp>(op) && hasExplicitNonFusableUsers(op)) {
+ return false;
+ }
if (LinalgExt::isBitExtendOp(op)) {
return true;
}
diff --git a/compiler/src/iree/compiler/DispatchCreation/FormDispatchRegions.cpp b/compiler/src/iree/compiler/DispatchCreation/FormDispatchRegions.cpp
index 91f4b9c..cf4f535 100644
--- a/compiler/src/iree/compiler/DispatchCreation/FormDispatchRegions.cpp
+++ b/compiler/src/iree/compiler/DispatchCreation/FormDispatchRegions.cpp
@@ -888,6 +888,13 @@
// Once all root linalg ops have been tagged, put all remaining generic ops
// into their own dispatches.
+ // TODO: when identyfing root operations, there are multiple
+ // isFusableWith[Producer|Consumer] checks being invoked, however we only
+ // access the positive results of those checks (i.e. explicit assignment to
+ // the fusion group). We should consider assigning *all* operations to fusion
+ // groups during the fuseRootsWith...() phase, creating new groups for more
+ // negative fusion cases, and potentially eliding the second iteration over
+ // the region altogether.
for (Block &block : region) {
SmallVector<Operation *> roots;
for (Operation &op : llvm::reverse(block)) {
diff --git a/compiler/src/iree/compiler/DispatchCreation/test/dispatch_linalg_ext_fusion.mlir b/compiler/src/iree/compiler/DispatchCreation/test/dispatch_linalg_ext_fusion.mlir
index aa25951..0f50ebe 100644
--- a/compiler/src/iree/compiler/DispatchCreation/test/dispatch_linalg_ext_fusion.mlir
+++ b/compiler/src/iree/compiler/DispatchCreation/test/dispatch_linalg_ext_fusion.mlir
@@ -1,4 +1,7 @@
-// RUN: iree-opt --split-input-file --pass-pipeline="builtin.module(util.func(iree-dispatch-creation-form-dispatch-regions{aggressive-fusion=true}, iree-dispatch-creation-clone-producers-into-dispatch-regions{aggressive=true}), cse, canonicalize, cse)" %s | FileCheck %s
+// RUN: iree-opt %s --split-input-file \
+// RUN: --pass-pipeline="builtin.module(util.func(iree-dispatch-creation-form-dispatch-regions{aggressive-fusion=true}, \
+// RUN: iree-dispatch-creation-clone-producers-into-dispatch-regions{aggressive=true}), cse, canonicalize, cse)" \
+// RUN: | FileCheck %s
#map = affine_map<(d0, d1) -> (d0, d1)>
#map1 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>
@@ -175,3 +178,42 @@
// CHECK: %[[GEN2:.+]] = linalg.generic
// CHECK-SAME: ins(%[[ATTN]]
// CHECK: flow.return %[[GEN2]]
+
+// -----
+
+// Check that a simple elementwise bit extend producer is assigned to a separate dispatch
+// (until fusion is supported)
+#map = affine_map<(d0, d1) -> (d0, d1)>
+util.func public @linalgext_scan_inclusive_dispatch_non_fusable(%arg0: tensor<8x32xi32>) -> tensor<8x32xi64> {
+ %c0_i64 = arith.constant 0 : i64
+ %0 = tensor.empty() : tensor<8x32xi64>
+ %1 = linalg.generic {indexing_maps = [#map, #map], iterator_types = ["parallel", "parallel"]} ins(%arg0 : tensor<8x32xi32>) outs(%0 : tensor<8x32xi64>) {
+ ^bb0(%in: i32, %out: i64):
+ %6 = arith.extsi %in : i32 to i64
+ linalg.yield %6 : i64
+ } -> tensor<8x32xi64>
+ %2 = linalg.fill ins(%c0_i64 : i64) outs(%0 : tensor<8x32xi64>) -> tensor<8x32xi64>
+ %3 = tensor.empty() : tensor<8xi64>
+ %4 = linalg.fill ins(%c0_i64 : i64) outs(%3 : tensor<8xi64>) -> tensor<8xi64>
+ %5:2 = iree_linalg_ext.scan dimension(1) inclusive(true) ins(%1 : tensor<8x32xi64>) outs(%2, %4 : tensor<8x32xi64>, tensor<8xi64>) {
+ ^bb0(%arg3: i64, %arg4: i64):
+ %6 = arith.addi %arg3, %arg4 : i64
+ iree_linalg_ext.yield %6 : i64
+ } -> tensor<8x32xi64>, tensor<8xi64>
+ util.return %5#0 : tensor<8x32xi64>
+}
+
+// CHECK-LABEL: util.func public @linalgext_scan_inclusive_dispatch_non_fusable(
+// CHECK-SAME: %[[ARG:.+]]: tensor<8x32xi32>) -> tensor<8x32xi64>
+// CHECK: %[[ZERO_CONST:.+]] = arith.constant 0 : i64
+// CHECK: %[[PRODUCER_REGION:.+]] = flow.dispatch.region -> (tensor<8x32xi64>)
+// CHECK: %[[EMPTY:.+]] = tensor.empty()
+// CHECK: %[[PRODUCER:.+]] = linalg.generic {{.+}} ins(%[[ARG]] : tensor<8x32xi32>) outs(%[[EMPTY]] : tensor<8x32xi64>)
+// CHECK: flow.return %[[PRODUCER]]
+// CHECK: %[[LINALGEXT_REGION:.+]] = flow.dispatch.region -> (tensor<8x32xi64>)
+// CHECK: %[[CUMULATIVE_FILL:.+]] = linalg.fill ins(%[[ZERO_CONST]] : i64) outs(%{{.+}} : tensor<8x32xi64>)
+// CHECK: %[[REDUCED_FILL:.+]] = linalg.fill ins(%[[ZERO_CONST]] : i64) outs(%{{.+}} : tensor<8xi64>)
+// CHECK: %[[SCAN_RESULT:.+]]:2 = iree_linalg_ext.scan dimension(1) inclusive(true)
+// CHECK-SAME: ins(%[[PRODUCER_REGION]] : tensor<8x32xi64>) outs(%[[CUMULATIVE_FILL]], %[[REDUCED_FILL]] : {{.+}}) {
+// CHECK: flow.return %[[SCAN_RESULT]]#0
+// CHECK: util.return %[[LINALGEXT_REGION]] : tensor<8x32xi64>
