[LinalgExt] Implement CustomOp::getStaticLoopRanges (#24734)
CustomOp implements LinalgFusionInterface but does not provide static
loop ranges, so dispatch formation has to exclude it from several fusion
checks.
Compute the ranges from the indexing maps and operand shapes, matching
the existing iteration-domain calculation. Remove the CustomOp-specific
exclusions and cover both static and dynamic shapes.
Fixes #23345
Signed-off-by: zhangpan <98025960+zhangpan2001@users.noreply.github.com>
Co-authored-by: OpenAI <noreply@openai.com>
diff --git a/compiler/src/iree/compiler/Dialect/LinalgExt/IR/LinalgExtInterfaces.td b/compiler/src/iree/compiler/Dialect/LinalgExt/IR/LinalgExtInterfaces.td
index e6a91dd..fb748f0 100644
--- a/compiler/src/iree/compiler/Dialect/LinalgExt/IR/LinalgExtInterfaces.td
+++ b/compiler/src/iree/compiler/Dialect/LinalgExt/IR/LinalgExtInterfaces.td
@@ -133,7 +133,8 @@
>,
InterfaceMethod<
/*desc=*/[{
- Return the static loop ranges.
+ Return one range per logical loop. Returns `ShapedType::kDynamic` for
+ ranges that cannot be statically recovered.
}],
/*retTy=*/"SmallVector<int64_t>",
/*methodName=*/"getStaticLoopRanges",
diff --git a/compiler/src/iree/compiler/Dialect/LinalgExt/IR/LinalgExtOps.cpp b/compiler/src/iree/compiler/Dialect/LinalgExt/IR/LinalgExtOps.cpp
index 0016d82..0a9b7d0 100644
--- a/compiler/src/iree/compiler/Dialect/LinalgExt/IR/LinalgExtOps.cpp
+++ b/compiler/src/iree/compiler/Dialect/LinalgExt/IR/LinalgExtOps.cpp
@@ -3311,6 +3311,39 @@
[](Attribute attr) { return cast<AffineMapAttr>(attr).getValue(); });
}
+SmallVector<int64_t> CustomOp::getStaticLoopRanges() {
+ SmallVector<AffineMap> maps = getIndexingMapsArray();
+ SmallVector<int64_t> loopRanges(getNumLoops(), ShapedType::kDynamic);
+ SmallVector<bool> conflictingRanges(getNumLoops(), false);
+
+ for (auto [operand, map] :
+ llvm::zip_equal(getOperation()->getOpOperands(), maps)) {
+ if (map.isEmpty()) {
+ continue;
+ }
+
+ for (auto [size, expr] :
+ llvm::zip_equal(getShape(&operand), map.getResults())) {
+ auto dimExpr = dyn_cast<AffineDimExpr>(expr);
+ if (!dimExpr || ShapedType::isDynamic(size)) {
+ continue;
+ }
+
+ unsigned position = dimExpr.getPosition();
+ if (conflictingRanges[position]) {
+ continue;
+ }
+ if (ShapedType::isDynamic(loopRanges[position])) {
+ loopRanges[position] = size;
+ } else if (loopRanges[position] != size) {
+ loopRanges[position] = ShapedType::kDynamic;
+ conflictingRanges[position] = true;
+ }
+ }
+ }
+ return loopRanges;
+}
+
/// End `LinalgFusionInterface` implementation
/// Start `ReifyRankedShapedTypeOpInterface` implementation
diff --git a/compiler/src/iree/compiler/Dialect/LinalgExt/IR/LinalgExtOps.td b/compiler/src/iree/compiler/Dialect/LinalgExt/IR/LinalgExtOps.td
index c0fc2db..7f6ec1c 100644
--- a/compiler/src/iree/compiler/Dialect/LinalgExt/IR/LinalgExtOps.td
+++ b/compiler/src/iree/compiler/Dialect/LinalgExt/IR/LinalgExtOps.td
@@ -1994,7 +1994,7 @@
AttrSizedOperandSegments,
DeclareOpInterfaceMethods<AggregatedOpInterface, [
"decomposeOperation"]>,
- DeclareOpInterfaceMethods<LinalgFusionInterface>,
+ DeclareOpInterfaceMethods<LinalgFusionInterface, ["getStaticLoopRanges"]>,
DeclareOpInterfaceMethods<ReifyRankedShapedTypeOpInterface, ["reifyResultShapes"]>,
DeclareOpInterfaceMethods<TilingInterface,
["getIterationDomain",
diff --git a/compiler/src/iree/compiler/Dialect/LinalgExt/IR/unittests/BUILD.bazel b/compiler/src/iree/compiler/Dialect/LinalgExt/IR/unittests/BUILD.bazel
new file mode 100644
index 0000000..7fd48fe
--- /dev/null
+++ b/compiler/src/iree/compiler/Dialect/LinalgExt/IR/unittests/BUILD.bazel
@@ -0,0 +1,26 @@
+# Copyright 2026 The IREE Authors
+#
+# Licensed under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+load("//build_tools/bazel:build_defs.oss.bzl", "iree_compiler_cc_test")
+
+package(
+ features = ["layering_check"],
+ licenses = ["notice"], # Apache 2.0
+)
+
+iree_compiler_cc_test(
+ name = "LinalgExtOpsTest",
+ testonly = True,
+ srcs = ["LinalgExtOpsTest.cpp"],
+ deps = [
+ "//compiler/src/iree/compiler/Dialect/LinalgExt/IR",
+ "//compiler/src/iree/testing:gtest_main",
+ "@com_google_googletest//:gtest",
+ "@llvm-project//mlir:FuncDialect",
+ "@llvm-project//mlir:IR",
+ "@llvm-project//mlir:Parser",
+ ],
+)
diff --git a/compiler/src/iree/compiler/Dialect/LinalgExt/IR/unittests/CMakeLists.txt b/compiler/src/iree/compiler/Dialect/LinalgExt/IR/unittests/CMakeLists.txt
new file mode 100644
index 0000000..f258a6f
--- /dev/null
+++ b/compiler/src/iree/compiler/Dialect/LinalgExt/IR/unittests/CMakeLists.txt
@@ -0,0 +1,28 @@
+################################################################################
+# Autogenerated by build_tools/bazel_to_cmake/bazel_to_cmake.py from #
+# compiler/src/iree/compiler/Dialect/LinalgExt/IR/unittests/BUILD.bazel #
+# #
+# Use iree_cmake_extra_content from iree/build_defs.oss.bzl to add arbitrary #
+# CMake-only content. #
+# #
+# To disable autogeneration for this file entirely, delete this header. #
+################################################################################
+
+iree_add_all_subdirs()
+
+iree_cc_test(
+ NAME
+ LinalgExtOpsTest
+ SRCS
+ "LinalgExtOpsTest.cpp"
+ DEPS
+ MLIRFuncDialect
+ MLIRIR
+ MLIRParser
+ gmock
+ gtest
+ iree::compiler::Dialect::LinalgExt::IR
+ iree::testing::gtest_main
+)
+
+### BAZEL_TO_CMAKE_PRESERVES_ALL_CONTENT_BELOW_THIS_LINE ###
diff --git a/compiler/src/iree/compiler/Dialect/LinalgExt/IR/unittests/LinalgExtOpsTest.cpp b/compiler/src/iree/compiler/Dialect/LinalgExt/IR/unittests/LinalgExtOpsTest.cpp
new file mode 100644
index 0000000..8889af5
--- /dev/null
+++ b/compiler/src/iree/compiler/Dialect/LinalgExt/IR/unittests/LinalgExtOpsTest.cpp
@@ -0,0 +1,143 @@
+// Copyright 2026 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+
+#include "iree/compiler/Dialect/LinalgExt/IR/LinalgExtDialect.h"
+#include "iree/compiler/Dialect/LinalgExt/IR/LinalgExtOps.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/Verifier.h"
+#include "mlir/Parser/Parser.h"
+
+namespace mlir::iree_compiler::IREE::LinalgExt {
+namespace {
+
+using ::testing::ElementsAre;
+
+class LinalgExtOpsTest : public ::testing::Test {
+protected:
+ LinalgExtOpsTest() {
+ registry.insert<IREELinalgExtDialect, func::FuncDialect>();
+ context.appendDialectRegistry(registry);
+ context.loadAllAvailableDialects();
+ }
+
+ MLIRContext *getContext() { return &context; }
+
+private:
+ MLIRContext context;
+ DialectRegistry registry;
+};
+
+static CustomOp findCustomOp(ModuleOp module, StringRef functionName) {
+ func::FuncOp function = module.lookupSymbol<func::FuncOp>(functionName);
+ if (!function) {
+ return {};
+ }
+ auto customOps = function.getOps<CustomOp>();
+ return customOps.empty() ? CustomOp{} : *customOps.begin();
+}
+
+TEST_F(LinalgExtOpsTest, GetStaticLoopRanges) {
+ OwningOpRef<ModuleOp> module = parseSourceString<ModuleOp>(
+ R"mlir(
+module {
+ func.func @symbol_dimensions(
+ %input : tensor<4x16xf32>, %output : tensor<4x16xf32>)
+ -> tensor<4x16xf32> {
+ %0 = iree_linalg_ext.custom_op {
+ indexing_maps = [affine_map<(d0)[s0] -> (d0, s0)>,
+ affine_map<(d0)[s0] -> (d0, s0)>],
+ iterator_types = [#iree_linalg_ext.iterator_type<parallel>]}
+ ins(%input : tensor<4x16xf32>)
+ outs(%output : tensor<4x16xf32>) {
+ ^bb0(%input_slice : tensor<?x?xf32>,
+ %output_slice : tensor<?x?xf32>):
+ iree_linalg_ext.yield %output_slice : tensor<?x?xf32>
+ } -> tensor<4x16xf32>
+ return %0 : tensor<4x16xf32>
+ }
+
+ func.func @partial_static_loop_ranges(
+ %input : tensor<7xf32>, %output : tensor<4xf32>) -> tensor<4xf32> {
+ %0 = iree_linalg_ext.custom_op {
+ indexing_maps = [affine_map<(d0, d1) -> (d0 + d1)>,
+ affine_map<(d0, d1) -> (d0)>],
+ iterator_types = [#iree_linalg_ext.iterator_type<parallel>,
+ #iree_linalg_ext.iterator_type<reduction>]}
+ ins(%input : tensor<7xf32>) outs(%output : tensor<4xf32>) {
+ ^bb0(%input_slice : tensor<?xf32>, %output_slice : tensor<?xf32>):
+ iree_linalg_ext.yield %output_slice : tensor<?xf32>
+ } -> tensor<4xf32>
+ return %0 : tensor<4xf32>
+ }
+
+ func.func @dynamic_and_static(
+ %input : tensor<?xf32>, %output : tensor<4xf32>) -> tensor<4xf32> {
+ %0 = iree_linalg_ext.custom_op {
+ indexing_maps = [affine_map<(d0) -> (d0)>,
+ affine_map<(d0) -> (d0)>],
+ iterator_types = [#iree_linalg_ext.iterator_type<parallel>]}
+ ins(%input : tensor<?xf32>) outs(%output : tensor<4xf32>) {
+ ^bb0(%input_slice : tensor<?xf32>, %output_slice : tensor<?xf32>):
+ iree_linalg_ext.yield %output_slice : tensor<?xf32>
+ } -> tensor<4xf32>
+ return %0 : tensor<4xf32>
+ }
+
+ func.func @conflicting_static(
+ %input : tensor<4xf32>, %output : tensor<8xf32>) -> tensor<8xf32> {
+ %0 = iree_linalg_ext.custom_op {
+ indexing_maps = [affine_map<(d0) -> (d0)>,
+ affine_map<(d0) -> (d0)>],
+ iterator_types = [#iree_linalg_ext.iterator_type<parallel>]}
+ ins(%input : tensor<4xf32>) outs(%output : tensor<8xf32>) {
+ ^bb0(%input_slice : tensor<?xf32>, %output_slice : tensor<?xf32>):
+ iree_linalg_ext.yield %output_slice : tensor<?xf32>
+ } -> tensor<8xf32>
+ return %0 : tensor<8xf32>
+ }
+
+ func.func @empty_map(%input : tensor<4xf32>,
+ %ignored : tensor<5x6xf32>, %output : tensor<4xf32>)
+ -> tensor<4xf32> {
+ %0 = iree_linalg_ext.custom_op {
+ indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<() -> ()>,
+ affine_map<(d0) -> (d0)>],
+ iterator_types = [#iree_linalg_ext.iterator_type<parallel>]}
+ ins(%input, %ignored : tensor<4xf32>, tensor<5x6xf32>)
+ outs(%output : tensor<4xf32>) {
+ ^bb0(%input_slice : tensor<?xf32>,
+ %ignored_slice : tensor<?x?xf32>,
+ %output_slice : tensor<?xf32>):
+ iree_linalg_ext.yield %output_slice : tensor<?xf32>
+ } -> tensor<4xf32>
+ return %0 : tensor<4xf32>
+ }
+}
+)mlir",
+ getContext());
+
+ ASSERT_TRUE(module);
+ ASSERT_TRUE(succeeded(verify(*module)));
+
+ EXPECT_THAT(findCustomOp(*module, "symbol_dimensions").getStaticLoopRanges(),
+ ElementsAre(4));
+ EXPECT_THAT(
+ findCustomOp(*module, "partial_static_loop_ranges").getStaticLoopRanges(),
+ ElementsAre(4, ShapedType::kDynamic));
+ EXPECT_THAT(findCustomOp(*module, "dynamic_and_static").getStaticLoopRanges(),
+ ElementsAre(4));
+ EXPECT_THAT(findCustomOp(*module, "conflicting_static").getStaticLoopRanges(),
+ ElementsAre(ShapedType::kDynamic));
+ EXPECT_THAT(findCustomOp(*module, "empty_map").getStaticLoopRanges(),
+ ElementsAre(4));
+}
+
+} // namespace
+} // namespace mlir::iree_compiler::IREE::LinalgExt
diff --git a/compiler/src/iree/compiler/DispatchCreation/FormDispatchRegions.cpp b/compiler/src/iree/compiler/DispatchCreation/FormDispatchRegions.cpp
index cf4f535..20d3a06 100644
--- a/compiler/src/iree/compiler/DispatchCreation/FormDispatchRegions.cpp
+++ b/compiler/src/iree/compiler/DispatchCreation/FormDispatchRegions.cpp
@@ -580,11 +580,6 @@
// TODO(#12664): This is unnecessary requirement, but we need a better config
// to tile the consumer with a larger iteration space.
if (!options.aggressiveFusion) {
- // FIXME: Implement getStaticLoopRanges for LinalgExt::CustomOp.
- if (isa<IREE::LinalgExt::CustomOp>(producer)) {
- return false;
- }
-
SmallVector<int64_t> producerIterationSpace =
producerFusionOp.getStaticLoopRanges();
SmallVector<int64_t> consumerIterationSpace =
@@ -601,7 +596,7 @@
Operation *rootOp = tracker.getFusionGroup(producer).getRoot();
if (auto rootFusionOp =
dyn_cast<IREE::LinalgExt::LinalgFusionOpInterface>(rootOp);
- rootFusionOp && !isa<IREE::LinalgExt::CustomOp>(rootOp)) {
+ rootFusionOp) {
SmallVector<int64_t> rootLoopRanges = rootFusionOp.getStaticLoopRanges();
SmallVector<int64_t> consumerLoopRanges =
consumerFusionOp.getStaticLoopRanges();
diff --git a/compiler/src/iree/compiler/DispatchCreation/test/form_dispatch_regions.mlir b/compiler/src/iree/compiler/DispatchCreation/test/form_dispatch_regions.mlir
index ee358f0..8e4eb51 100644
--- a/compiler/src/iree/compiler/DispatchCreation/test/form_dispatch_regions.mlir
+++ b/compiler/src/iree/compiler/DispatchCreation/test/form_dispatch_regions.mlir
@@ -807,6 +807,147 @@
// CHECK-SAME: ins(%[[CUSTOM_OP]] :
// CHECK: flow.return %[[GENERIC]]
// CHECK: util.return %[[DISPATCH]]
+// DEFAULT-LABEL: func public @custom_op_consumer_fusion
+// DEFAULT: %[[DISPATCH:.+]] = flow.dispatch.region
+// DEFAULT: %[[CUSTOM_OP:.+]] = iree_linalg_ext.custom_op
+// DEFAULT: linalg.generic
+// DEFAULT: %[[GENERIC:.+]] = linalg.generic
+// DEFAULT-SAME: ins(%[[CUSTOM_OP]] :
+// DEFAULT: flow.return %[[GENERIC]]
+// DEFAULT-NOT: flow.dispatch.region
+// DEFAULT: util.return %[[DISPATCH]]
+
+// -----
+
+util.func @custom_op_static_consumer_fusion(%arg0 : tensor<4xf32>,
+ %arg1 : tensor<4xf32>, %arg2 : tensor<4xf32>) -> tensor<4xf32> {
+ %0 = iree_linalg_ext.custom_op {
+ indexing_maps = [affine_map<(d0) -> (d0)>,
+ affine_map<(d0) -> (d0)>],
+ iterator_types = [#iree_linalg_ext.iterator_type<parallel>]}
+ ins(%arg0 : tensor<4xf32>) outs(%arg1 : tensor<4xf32>) {
+ ^bb0(%b0 : tensor<?xf32>, %b1 : tensor<?xf32>):
+ %1 = linalg.generic {
+ indexing_maps = [affine_map<(d0) -> (d0)>,
+ affine_map<(d0) -> (d0)>],
+ iterator_types = ["parallel"]}
+ ins(%b0 : tensor<?xf32>) outs(%b1 : tensor<?xf32>) {
+ ^bb1(%bb0 : f32, %bb1 : f32):
+ %2 = arith.addf %bb0, %bb1 : f32
+ linalg.yield %2 : f32
+ } -> tensor<?xf32>
+ iree_linalg_ext.yield %1 : tensor<?xf32>
+ } -> tensor<4xf32>
+ %3 = linalg.generic {
+ indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>],
+ iterator_types = ["parallel"]}
+ ins(%0 : tensor<4xf32>) outs(%arg2 : tensor<4xf32>) {
+ ^bb0(%b0 : f32, %b1 : f32):
+ %4 = arith.mulf %b0, %b0 : f32
+ linalg.yield %4 : f32
+ } -> tensor<4xf32>
+ util.return %3 : tensor<4xf32>
+}
+// CHECK-LABEL: func public @custom_op_static_consumer_fusion
+// CHECK: %[[DISPATCH:.+]] = flow.dispatch.region
+// CHECK: %[[CUSTOM_OP:.+]] = iree_linalg_ext.custom_op
+// CHECK: linalg.generic
+// CHECK: %[[GENERIC:.+]] = linalg.generic
+// CHECK-SAME: ins(%[[CUSTOM_OP]] :
+// CHECK: flow.return %[[GENERIC]]
+// CHECK: util.return %[[DISPATCH]]
+// DEFAULT-LABEL: func public @custom_op_static_consumer_fusion
+// DEFAULT: %[[DISPATCH:.+]] = flow.dispatch.region
+// DEFAULT: %[[CUSTOM_OP:.+]] = iree_linalg_ext.custom_op
+// DEFAULT: linalg.generic
+// DEFAULT: %[[GENERIC:.+]] = linalg.generic
+// DEFAULT-SAME: ins(%[[CUSTOM_OP]] :
+// DEFAULT: flow.return %[[GENERIC]]
+// DEFAULT-NOT: flow.dispatch.region
+// DEFAULT: util.return %[[DISPATCH]]
+
+// -----
+
+util.func @custom_op_dynamic_loop_range_with_static_consumer_fusion(
+ %arg0 : tensor<7xf32>, %arg1 : tensor<4xf32>,
+ %arg2 : tensor<4xf32>) -> tensor<4xf32> {
+ %0 = iree_linalg_ext.custom_op {
+ indexing_maps = [affine_map<(d0, d1) -> (d0 + d1)>,
+ affine_map<(d0, d1) -> (d0)>],
+ iterator_types = [#iree_linalg_ext.iterator_type<parallel>,
+ #iree_linalg_ext.iterator_type<reduction>]}
+ ins(%arg0 : tensor<7xf32>) outs(%arg1 : tensor<4xf32>) {
+ ^bb0(%b0 : tensor<?xf32>, %b1 : tensor<?xf32>):
+ iree_linalg_ext.yield %b1 : tensor<?xf32>
+ } -> tensor<4xf32>
+ %1 = linalg.generic {
+ indexing_maps = [affine_map<(d0) -> (d0)>,
+ affine_map<(d0) -> (d0)>],
+ iterator_types = ["parallel"]}
+ ins(%0 : tensor<4xf32>) outs(%arg2 : tensor<4xf32>) {
+ ^bb0(%b0 : f32, %b1 : f32):
+ %2 = arith.mulf %b0, %b0 : f32
+ linalg.yield %2 : f32
+ } -> tensor<4xf32>
+ util.return %1 : tensor<4xf32>
+}
+// CHECK-LABEL: func public @custom_op_dynamic_loop_range_with_static_consumer_fusion
+// CHECK: %[[DISPATCH:.+]] = flow.dispatch.region
+// CHECK: %[[CUSTOM_OP:.+]] = iree_linalg_ext.custom_op
+// CHECK: %[[GENERIC:.+]] = linalg.generic
+// CHECK-SAME: ins(%[[CUSTOM_OP]] :
+// CHECK: flow.return %[[GENERIC]]
+// CHECK: util.return %[[DISPATCH]]
+// DEFAULT-LABEL: func public @custom_op_dynamic_loop_range_with_static_consumer_fusion
+// DEFAULT: %[[DISPATCH:.+]] = flow.dispatch.region
+// DEFAULT: %[[CUSTOM_OP:.+]] = iree_linalg_ext.custom_op
+// DEFAULT: %[[GENERIC:.+]] = linalg.generic
+// DEFAULT-SAME: ins(%[[CUSTOM_OP]] :
+// DEFAULT: flow.return %[[GENERIC]]
+// DEFAULT-NOT: flow.dispatch.region
+// DEFAULT: util.return %[[DISPATCH]]
+
+// -----
+
+util.func @custom_op_consumer_fusion_larger_consumer_loop_count(
+ %arg0 : tensor<4xf32>, %arg1 : tensor<4xf32>,
+ %arg2 : tensor<4x1xf32>) -> tensor<4x1xf32> {
+ %0 = iree_linalg_ext.custom_op {
+ indexing_maps = [affine_map<(d0) -> (d0)>,
+ affine_map<(d0) -> (d0)>],
+ iterator_types = [#iree_linalg_ext.iterator_type<parallel>]}
+ ins(%arg0 : tensor<4xf32>) outs(%arg1 : tensor<4xf32>) {
+ ^bb0(%b0 : tensor<?xf32>, %b1 : tensor<?xf32>):
+ iree_linalg_ext.yield %b1 : tensor<?xf32>
+ } -> tensor<4xf32>
+ %1 = linalg.generic {
+ indexing_maps = [affine_map<(d0, d1) -> (d0)>,
+ affine_map<(d0, d1) -> (d0, d1)>],
+ iterator_types = ["parallel", "parallel"]}
+ ins(%0 : tensor<4xf32>) outs(%arg2 : tensor<4x1xf32>) {
+ ^bb0(%b0 : f32, %b1 : f32):
+ %2 = arith.mulf %b0, %b0 : f32
+ linalg.yield %2 : f32
+ } -> tensor<4x1xf32>
+ util.return %1 : tensor<4x1xf32>
+}
+// CHECK-LABEL: func public @custom_op_consumer_fusion_larger_consumer_loop_count
+// CHECK: %[[DISPATCH:.+]] = flow.dispatch.region
+// CHECK: %[[CUSTOM_OP:.+]] = iree_linalg_ext.custom_op
+// CHECK: %[[GENERIC:.+]] = linalg.generic
+// CHECK-SAME: ins(%[[CUSTOM_OP]] :
+// CHECK: flow.return %[[GENERIC]]
+// CHECK-NOT: flow.dispatch.region
+// CHECK: util.return %[[DISPATCH]]
+// DEFAULT-LABEL: func public @custom_op_consumer_fusion_larger_consumer_loop_count
+// DEFAULT: %[[CUSTOM_DISPATCH:.+]] = flow.dispatch.region
+// DEFAULT: %[[CUSTOM_OP:.+]] = iree_linalg_ext.custom_op
+// DEFAULT: flow.return %[[CUSTOM_OP]]
+// DEFAULT: %[[GENERIC_DISPATCH:.+]] = flow.dispatch.region
+// DEFAULT: %[[GENERIC:.+]] = linalg.generic
+// DEFAULT-SAME: ins(%[[CUSTOM_DISPATCH]] :
+// DEFAULT: flow.return %[[GENERIC]]
+// DEFAULT: util.return %[[GENERIC_DISPATCH]]
// -----
