[Mlir-commits] [mlir] 1a083de - [mlir][transform] Add ApplySwapExtractSliceWithFillPatternsOp transform op (#213907)

[llvmlistbot at llvm.org]

Author: Tuomas Kärnä
Date: 2026-08-25T17:46:32+03:00
New Revision: 1a083de492e092a7334a1dd0209262296949566c

URL: https://github.com/llvm/llvm-project/commit/1a083de492e092a7334a1dd0209262296949566c
DIFF: https://github.com/llvm/llvm-project/commit/1a083de492e092a7334a1dd0209262296949566c.diff

LOG: [mlir][transform] Add ApplySwapExtractSliceWithFillPatternsOp transform op (#213907)

Adds transform op for the existing linalg `SwapExtractSliceOfFill`
pattern as `apply_patterns.linalg.swap_extract_slice_with_fill`.

Assisted-by: GPT-5.3-Codex

Added: 
    mlir/test/Dialect/Linalg/transform-op-swap-extract-slice-with-fill.mlir

Modified: 
    mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td
    mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
    mlir/test/python/dialects/transform.py

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td b/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td
index 9a6796f999931..ed6ab90c65401 100644
--- a/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td
+++ b/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td
@@ -165,6 +165,17 @@ def ApplyExtractSliceSinkingPatternsOp : Op<Transform_Dialect,
   let assemblyFormat = "attr-dict";
 }
 
+def ApplySwapExtractSliceWithFillPatternsOp : Op<Transform_Dialect,
+    "apply_patterns.linalg.swap_extract_slice_with_fill",
+    [DeclareOpInterfaceMethods<PatternDescriptorOpInterface>]> {
+  let description = [{
+    Patterns to swap `tensor.extract_slice(linalg.fill(...))` with
+    `linalg.fill(tensor.extract_slice(...))`.
+  }];
+
+  let assemblyFormat = "attr-dict";
+}
+
 //===----------------------------------------------------------------------===//
 // BufferizeToAllocationOp
 //===----------------------------------------------------------------------===//

diff  --git a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
index 22724e3c31121..af50aa79bd491 100644
--- a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
+++ b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
@@ -292,6 +292,11 @@ void transform::ApplyExtractSliceSinkingPatternsOp::populatePatterns(
   linalg::populateExtractSliceSinkingPatterns(patterns, defaultControlFn);
 }
 
+void transform::ApplySwapExtractSliceWithFillPatternsOp::populatePatterns(
+    RewritePatternSet &patterns) {
+  linalg::populateSwapExtractSliceWithFillPatterns(patterns);
+}
+
 //===----------------------------------------------------------------------===//
 // BufferizeToAllocationOp
 //===----------------------------------------------------------------------===//

diff  --git a/mlir/test/Dialect/Linalg/transform-op-swap-extract-slice-with-fill.mlir b/mlir/test/Dialect/Linalg/transform-op-swap-extract-slice-with-fill.mlir
new file mode 100644
index 0000000000000..c4d2cec5e5034
--- /dev/null
+++ b/mlir/test/Dialect/Linalg/transform-op-swap-extract-slice-with-fill.mlir
@@ -0,0 +1,25 @@
+// RUN: mlir-opt %s -transform-interpreter | FileCheck %s
+
+func.func @swap_fill_extract_slice(%init : tensor<?x?x?xf32>, %offset0: index, %size1: index) -> tensor<?x6xf32> {
+  %f0 = arith.constant 0.000000e+00 : f32
+  %0 = linalg.fill ins(%f0 : f32) outs(%init : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
+  %1 = tensor.extract_slice %0[%offset0, 8, 4] [1, %size1, 6] [1, 3, 1]
+    : tensor<?x?x?xf32> to tensor<?x6xf32>
+  return %1: tensor<?x6xf32>
+}
+
+module attributes {transform.with_named_sequence} {
+  transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {
+    %f = transform.structured.match ops{["func.func"]} in %arg0 : (!transform.any_op) -> !transform.any_op
+    transform.apply_patterns to %f {
+      transform.apply_patterns.linalg.swap_extract_slice_with_fill
+    } : !transform.any_op
+    transform.yield
+  }
+}
+
+// CHECK-LABEL: func.func @swap_fill_extract_slice
+// CHECK: %[[F0:.*]] = arith.constant 0.000000e+00 : f32
+// CHECK: %[[EXT:.*]] = tensor.extract_slice %{{.*}}[%{{.*}}, 8, 4] [1, %{{.*}}, 6] [1, 3, 1]
+// CHECK: %[[FILL:.*]] = linalg.fill ins(%[[F0]] : f32) outs(%[[EXT]] : tensor<?x6xf32>) -> tensor<?x6xf32>
+// CHECK: return %[[FILL]] : tensor<?x6xf32>

diff  --git a/mlir/test/python/dialects/transform.py b/mlir/test/python/dialects/transform.py
index dfcc890b83ffc..0d5b0d26ac443 100644
--- a/mlir/test/python/dialects/transform.py
+++ b/mlir/test/python/dialects/transform.py
@@ -2,6 +2,7 @@
 
 from mlir.ir import *
 from mlir.dialects import transform
+from mlir.dialects.transform import structured
 from mlir.dialects.transform import pdl as transform_pdl
 
 
@@ -301,6 +302,19 @@ def testApplyPatternsOpWithType(module: Module):
         # CHECK: !transform.op<"test.dummy">
 
 
+ at run
+def testApplyLinalgSwapExtractSliceWithFillPattern(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_swap_extract_slice_with_fill()
+        transform.YieldOp()
+    # CHECK-LABEL: TEST: testApplyLinalgSwapExtractSliceWithFillPattern
+    # CHECK: transform.apply_patterns.linalg.swap_extract_slice_with_fill
+
+
 @run
 def testReplicateOp(module: Module):
     with_pdl = transform_pdl.WithPDLPatternsOp(transform.AnyOpType.get())