[Mlir-commits] [mlir] [mlir][transform] Add linalg.fold_reshape_ops_by_expansion transform op (PR #233732)
Tuomas Kärnä
llvmlistbot at llvm.org
Tue Oct 20 03:55:52 PDT 2026
https://github.com/tkarna updated https://github.com/llvm/llvm-project/pull/233732
>From 59a19c43e57c06adca7a32ce36dcffa8ee6b32a6 Mon Sep 17 00:00:00 2001
From: Tuomas Karna <tuomas.karna at intel.com>
Date: Mon, 19 Oct 2026 14:16:58 +0300
Subject: [PATCH 1/2] add linalg.fold_reshape_ops_by_expansion transform op
---
.../Linalg/TransformOps/LinalgTransformOps.td | 17 ++++
.../TransformOps/LinalgTransformOps.cpp | 12 +++
mlir/test/Dialect/Linalg/reshape_fusion.mlir | 3 +
.../td/fold-reshape-ops-by-expansion.mlir | 25 +++++
...form-op-fold-reshape-ops-by-expansion.mlir | 92 +++++++++++++++++++
mlir/test/python/dialects/transform.py | 21 +++++
6 files changed, 170 insertions(+)
create mode 100644 mlir/test/Dialect/Linalg/td/fold-reshape-ops-by-expansion.mlir
create mode 100644 mlir/test/Dialect/Linalg/transform-op-fold-reshape-ops-by-expansion.mlir
diff --git a/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td b/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td
index c31cd399032986..857b2236334bbc 100644
--- a/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td
+++ b/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td
@@ -187,6 +187,23 @@ def ApplyEraseUnusedOperandsAndResultsPatternsOp : Op<Transform_Dialect,
let assemblyFormat = "attr-dict";
}
+def ApplyFoldReshapeOpsByExpansionPatternsOp : Op<Transform_Dialect,
+ "apply_patterns.linalg.fold_reshape_ops_by_expansion",
+ [DeclareOpInterfaceMethods<PatternDescriptorOpInterface>]> {
+ let description = [{
+ Collects patterns that fold reshape operations into their producers or
+ consumers by expanding dimensions, including patterns for Linalg and
+ tensor.pad operations.
+
+ If `fold_single_use_only` is true, folding is restricted to cases where the
+ defining producer operation of the operand being fused has exactly one use.
+ Otherwise, no use-count restriction is imposed.
+ }];
+
+ let arguments = (ins DefaultValuedAttr<BoolAttr, "false">:$fold_single_use_only);
+ let assemblyFormat = "oilist(`fold_single_use_only` `=` $fold_single_use_only) attr-dict";
+}
+
//===----------------------------------------------------------------------===//
// BufferizeToAllocationOp
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
index 816d7499845789..ad33c9c99b158c 100644
--- a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
+++ b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
@@ -276,6 +276,18 @@ void transform::ApplyEraseUnusedOperandsAndResultsPatternsOp::populatePatterns(
linalg::populateEraseUnusedOperandsAndResultsPatterns(patterns);
}
+void transform::ApplyFoldReshapeOpsByExpansionPatternsOp::populatePatterns(
+ RewritePatternSet &patterns) {
+ linalg::ControlFusionFn controlFn =
+ [foldSingleUseOnly = getFoldSingleUseOnly()](OpOperand *fusedOperand) {
+ if (!foldSingleUseOnly)
+ return true;
+ Operation *producer = fusedOperand->get().getDefiningOp();
+ return producer && producer->hasOneUse();
+ };
+ linalg::populateFoldReshapeOpsByExpansionPatterns(patterns, controlFn);
+}
+
//===----------------------------------------------------------------------===//
// BufferizeToAllocationOp
//===----------------------------------------------------------------------===//
diff --git a/mlir/test/Dialect/Linalg/reshape_fusion.mlir b/mlir/test/Dialect/Linalg/reshape_fusion.mlir
index 874de9e14714aa..4a1ea5cf39c192 100644
--- a/mlir/test/Dialect/Linalg/reshape_fusion.mlir
+++ b/mlir/test/Dialect/Linalg/reshape_fusion.mlir
@@ -1,4 +1,7 @@
// RUN: mlir-opt %s -test-linalg-elementwise-fusion-patterns=fuse-with-reshape-by-expansion -split-input-file | FileCheck %s
+// RUN: mlir-opt %s -split-input-file \
+// RUN: -transform-preload-library='transform-library-paths=%p/td/fold-reshape-ops-by-expansion.mlir' \
+// RUN: -transform-interpreter=entry-point=fold_reshape_ops_by_expansion | FileCheck %s
#map0 = affine_map<(d0, d1, d2) -> (d2, d0, d1)>
#map1 = affine_map<(d0, d1, d2) -> (d1, d2, d0)>
diff --git a/mlir/test/Dialect/Linalg/td/fold-reshape-ops-by-expansion.mlir b/mlir/test/Dialect/Linalg/td/fold-reshape-ops-by-expansion.mlir
new file mode 100644
index 00000000000000..16fda19a1926d0
--- /dev/null
+++ b/mlir/test/Dialect/Linalg/td/fold-reshape-ops-by-expansion.mlir
@@ -0,0 +1,25 @@
+module @transforms attributes { transform.with_named_sequence } {
+ transform.named_sequence @fold_reshape_ops_by_expansion(%module: !transform.any_op {transform.readonly}) {
+ %func = transform.structured.match ops{["func.func"]} in %module : (!transform.any_op) -> !transform.any_op
+ transform.apply_patterns to %func {
+ transform.apply_patterns.linalg.fold_reshape_ops_by_expansion
+ } : !transform.any_op
+ transform.yield
+ }
+
+ transform.named_sequence @fold_reshape_ops_by_expansion_unrestricted(%module: !transform.any_op {transform.readonly}) {
+ %func = transform.structured.match ops{["func.func"]} in %module : (!transform.any_op) -> !transform.any_op
+ transform.apply_patterns to %func {
+ transform.apply_patterns.linalg.fold_reshape_ops_by_expansion fold_single_use_only = false
+ } : !transform.any_op
+ transform.yield
+ }
+
+ transform.named_sequence @fold_reshape_ops_by_expansion_single_use(%module: !transform.any_op {transform.readonly}) {
+ %func = transform.structured.match ops{["func.func"]} in %module : (!transform.any_op) -> !transform.any_op
+ transform.apply_patterns to %func {
+ transform.apply_patterns.linalg.fold_reshape_ops_by_expansion fold_single_use_only = true
+ } : !transform.any_op
+ transform.yield
+ }
+}
diff --git a/mlir/test/Dialect/Linalg/transform-op-fold-reshape-ops-by-expansion.mlir b/mlir/test/Dialect/Linalg/transform-op-fold-reshape-ops-by-expansion.mlir
new file mode 100644
index 00000000000000..e3c9d781b49a2a
--- /dev/null
+++ b/mlir/test/Dialect/Linalg/transform-op-fold-reshape-ops-by-expansion.mlir
@@ -0,0 +1,92 @@
+// RUN: mlir-opt %s \
+// RUN: -transform-preload-library='transform-library-paths=%p/td/fold-reshape-ops-by-expansion.mlir' \
+// RUN: -transform-interpreter=entry-point=fold_reshape_ops_by_expansion \
+// RUN: | FileCheck %s --check-prefixes=CHECK,UNRESTRICTED
+// RUN: mlir-opt %s \
+// RUN: -transform-preload-library='transform-library-paths=%p/td/fold-reshape-ops-by-expansion.mlir' \
+// RUN: -transform-interpreter=entry-point=fold_reshape_ops_by_expansion_unrestricted \
+// RUN: | FileCheck %s --check-prefixes=CHECK,UNRESTRICTED
+// RUN: mlir-opt %s \
+// RUN: -transform-preload-library='transform-library-paths=%p/td/fold-reshape-ops-by-expansion.mlir' \
+// RUN: -transform-interpreter=entry-point=fold_reshape_ops_by_expansion_single_use \
+// RUN: | FileCheck %s --check-prefixes=CHECK,SINGLE
+
+// CHECK-LABEL: func.func @collapse_single_use
+// CHECK: %[[COLLAPSE_SINGLE_FUSED:.*]] = linalg.generic
+// CHECK-SAME: iterator_types = ["parallel", "parallel"]
+// CHECK: %[[COLLAPSE_SINGLE_RESULT:.*]] = tensor.collapse_shape %[[COLLAPSE_SINGLE_FUSED]]
+// CHECK: return %[[COLLAPSE_SINGLE_RESULT]]
+func.func @collapse_single_use(%arg: tensor<2x3xf32>, %init: tensor<6xf32>) -> tensor<6xf32> {
+ %collapsed = tensor.collapse_shape %arg [[0, 1]] : tensor<2x3xf32> into tensor<6xf32>
+ %result = linalg.generic {
+ indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>],
+ iterator_types = ["parallel"]
+ } ins(%collapsed : tensor<6xf32>) outs(%init : tensor<6xf32>) {
+ ^bb0(%in: f32, %out: f32):
+ %neg = arith.negf %in : f32
+ linalg.yield %neg : f32
+ } -> tensor<6xf32>
+ return %result : tensor<6xf32>
+}
+
+// CHECK-LABEL: func.func @collapse_multi_use
+// CHECK: %[[COLLAPSED:.*]] = tensor.collapse_shape
+// UNRESTRICTED: %[[COLLAPSE_MULTI_FUSED:.*]] = linalg.generic
+// UNRESTRICTED-SAME: iterator_types = ["parallel", "parallel"]
+// UNRESTRICTED: %[[COLLAPSE_MULTI_RESULT:.*]] = tensor.collapse_shape %[[COLLAPSE_MULTI_FUSED]]
+// UNRESTRICTED: return %[[COLLAPSED]], %[[COLLAPSE_MULTI_RESULT]]
+// SINGLE: %[[COLLAPSE_MULTI_ORIGINAL:.*]] = linalg.generic
+// SINGLE-SAME: iterator_types = ["parallel"]
+// SINGLE-SAME: ins(%[[COLLAPSED]] : tensor<6xf32>)
+// SINGLE: return %[[COLLAPSED]], %[[COLLAPSE_MULTI_ORIGINAL]]
+func.func @collapse_multi_use(%arg: tensor<2x3xf32>, %init: tensor<6xf32>) -> (tensor<6xf32>, tensor<6xf32>) {
+ %collapsed = tensor.collapse_shape %arg [[0, 1]] : tensor<2x3xf32> into tensor<6xf32>
+ %result = linalg.generic {
+ indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>],
+ iterator_types = ["parallel"]
+ } ins(%collapsed : tensor<6xf32>) outs(%init : tensor<6xf32>) {
+ ^bb0(%in: f32, %out: f32):
+ %neg = arith.negf %in : f32
+ linalg.yield %neg : f32
+ } -> tensor<6xf32>
+ return %collapsed, %result : tensor<6xf32>, tensor<6xf32>
+}
+
+// CHECK-LABEL: func.func @expand_single_use
+// CHECK: %[[EXPAND_SINGLE_FUSED:.*]] = linalg.generic
+// CHECK-SAME: iterator_types = ["parallel", "parallel"]
+// CHECK: return %[[EXPAND_SINGLE_FUSED]]
+func.func @expand_single_use(%arg: tensor<6xf32>, %init: tensor<6xf32>) -> tensor<2x3xf32> {
+ %result = linalg.generic {
+ indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>],
+ iterator_types = ["parallel"]
+ } ins(%arg : tensor<6xf32>) outs(%init : tensor<6xf32>) {
+ ^bb0(%in: f32, %out: f32):
+ %neg = arith.negf %in : f32
+ linalg.yield %neg : f32
+ } -> tensor<6xf32>
+ %expanded = tensor.expand_shape %result [[0, 1]] output_shape [2, 3] : tensor<6xf32> into tensor<2x3xf32>
+ return %expanded : tensor<2x3xf32>
+}
+
+// CHECK-LABEL: func.func @expand_multi_use
+// UNRESTRICTED: %[[EXPAND_MULTI_FUSED:.*]] = linalg.generic
+// UNRESTRICTED-SAME: iterator_types = ["parallel", "parallel"]
+// UNRESTRICTED: %[[EXPAND_MULTI_COLLAPSED:.*]] = tensor.collapse_shape %[[EXPAND_MULTI_FUSED]]
+// UNRESTRICTED: return %[[EXPAND_MULTI_COLLAPSED]], %[[EXPAND_MULTI_FUSED]]
+// SINGLE: %[[EXPAND_MULTI_ORIGINAL:.*]] = linalg.generic
+// SINGLE-SAME: iterator_types = ["parallel"]
+// SINGLE: %[[EXPAND_MULTI_RESULT:.*]] = tensor.expand_shape %[[EXPAND_MULTI_ORIGINAL]]
+// SINGLE: return %[[EXPAND_MULTI_ORIGINAL]], %[[EXPAND_MULTI_RESULT]]
+func.func @expand_multi_use(%arg: tensor<6xf32>, %init: tensor<6xf32>) -> (tensor<6xf32>, tensor<2x3xf32>) {
+ %result = linalg.generic {
+ indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>],
+ iterator_types = ["parallel"]
+ } ins(%arg : tensor<6xf32>) outs(%init : tensor<6xf32>) {
+ ^bb0(%in: f32, %out: f32):
+ %neg = arith.negf %in : f32
+ linalg.yield %neg : f32
+ } -> tensor<6xf32>
+ %expanded = tensor.expand_shape %result [[0, 1]] output_shape [2, 3] : tensor<6xf32> into tensor<2x3xf32>
+ return %result, %expanded : tensor<6xf32>, tensor<2x3xf32>
+}
diff --git a/mlir/test/python/dialects/transform.py b/mlir/test/python/dialects/transform.py
index 55643f2bba870a..f586afb4cba078 100644
--- a/mlir/test/python/dialects/transform.py
+++ b/mlir/test/python/dialects/transform.py
@@ -328,6 +328,27 @@ def testApplyLinalgEraseUnusedOperandsAndResultsPattern(module: Module):
# CHECK: transform.apply_patterns.linalg.erase_unused_operands_and_results
+ at run
+def testApplyLinalgFoldReshapeOpsByExpansionPattern(module: Module):
+ sequence = transform.SequenceOp(
+ transform.FailurePropagationMode.Propagate, [], transform.AnyOpType.get()
+ )
+ with InsertionPoint(sequence.body):
+ with InsertionPoint(transform.ApplyPatternsOp(sequence.bodyTarget).patterns):
+ structured.apply_patterns_linalg_fold_reshape_ops_by_expansion()
+ structured.apply_patterns_linalg_fold_reshape_ops_by_expansion(
+ fold_single_use_only=False
+ )
+ structured.apply_patterns_linalg_fold_reshape_ops_by_expansion(
+ fold_single_use_only=True
+ )
+ transform.YieldOp()
+ # CHECK-LABEL: TEST: testApplyLinalgFoldReshapeOpsByExpansionPattern
+ # CHECK: transform.apply_patterns.linalg.fold_reshape_ops_by_expansion{{$}}
+ # CHECK-NEXT: transform.apply_patterns.linalg.fold_reshape_ops_by_expansion{{$}}
+ # CHECK-NEXT: transform.apply_patterns.linalg.fold_reshape_ops_by_expansion fold_single_use_only = true
+
+
@run
def testReplicateOp(module: Module):
with_pdl = transform_pdl.WithPDLPatternsOp(transform.AnyOpType.get())
>From cc4d5ad8183b3bc80664d29bf9222de0fb2e849e Mon Sep 17 00:00:00 2001
From: Tuomas Karna <tuomas.karna at intel.com>
Date: Tue, 20 Oct 2026 13:54:50 +0300
Subject: [PATCH 2/2] update docstring
---
.../mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td | 3 +--
1 file changed, 1 insertion(+), 2 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td b/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td
index 857b2236334bbc..035bb6a016e782 100644
--- a/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td
+++ b/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td
@@ -192,8 +192,7 @@ def ApplyFoldReshapeOpsByExpansionPatternsOp : Op<Transform_Dialect,
[DeclareOpInterfaceMethods<PatternDescriptorOpInterface>]> {
let description = [{
Collects patterns that fold reshape operations into their producers or
- consumers by expanding dimensions, including patterns for Linalg and
- tensor.pad operations.
+ consumers by expanding dimensions.
If `fold_single_use_only` is true, folding is restricted to cases where the
defining producer operation of the operand being fused has exactly one use.
More information about the Mlir-commits
mailing list