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

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-linalg

Author: Tuomas Kärnä (tkarna)

<details>
<summary>Changes</summary>

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

For motivation, see https://github.com/llvm/llvm-project/pull/212974#issuecomment-5176932196

---
Full diff: https://github.com/llvm/llvm-project/pull/213907.diff


3 Files Affected:

- (modified) mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td (+11) 
- (modified) mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp (+5) 
- (modified) mlir/test/python/dialects/transform.py (+14) 


``````````diff
diff --git a/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td b/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td
index 09a6c4c4d652f..0f8d45779bb43 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 8c5bd50bd6f8b..e6bfebef3e100 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/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())

``````````

</details>


https://github.com/llvm/llvm-project/pull/213907