[Mlir-commits] [mlir] [mlir][linalg] Expose transform.fuse_into_containing_op helpers: NFC (PR #72473)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-linalg

@llvm/pr-subscribers-mlir

Author: Quinn Dawkins (qedawkins)

<details>
<summary>Changes</summary>

This allows use of the various fusion helpers called by
`transform.fuse_into_containing_op` in other contexts.

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


2 Files Affected:

- (modified) mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.h (+27) 
- (modified) mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp (+11-20) 


``````````diff
diff --git a/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.h b/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.h
index 12923663b3fb6ce..f9f96ae0b7b6494 100644
--- a/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.h
+++ b/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.h
@@ -59,6 +59,33 @@ tileToForallOpImpl(RewriterBase &rewriter, transform::TransformState &state,
                    std::optional<ArrayAttr> mapping,
                    linalg::ForallTilingResult &tilingResult);
 
+/// Find the first "extract" user of `producerOp` and tile it right before its
+/// use. The tiled op is fused under the `containingOp`.
+/// Return this fused op on success or nullptr if anything fails.
+/// If tiled op has uses that are dominated by `containingOp`, return
+/// a new `containingOp` with results of the fused op appended to
+/// results of the `containingOp` or nullptr if there are no dominated uses.
+std::tuple<SmallVector<Operation *>, Operation *>
+tileAndFuseFirstExtractUse(RewriterBase &rewriter, Diagnostic &diag,
+                           Operation *producerOp, Operation *containingOp);
+
+/// First, find the first "scf::ForallOp" user of `producerOp` and ensure
+/// it is exactly the `containingOp`, otherwise bail.
+/// Then, find the first "extract" user of the tied block argument and tile it
+/// right before its "extract" use. The tiled op is fused under the
+/// `containingOp`.
+/// Return this fused op on success or nullptr if anything fails.
+SmallVector<Operation *>
+tileAndFuseFirstExtractUseThroughContainingOpBlockArgument(
+    RewriterBase &rewriter, Diagnostic &diag, Operation *producerOp,
+    Operation *containingOp);
+
+/// Find the first use of `producerOp` inside `containingOp` and fuse into
+/// the containing op by cloning the producer. Return nullptr if no such
+/// fusion opportunity exists.
+Operation *cloneAndFuseFirstUse(RewriterBase &rewriter, Diagnostic &diag,
+                                Operation *producerOp, Operation *containingOp);
+
 } // namespace transform
 } // namespace mlir
 
diff --git a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
index de4965f937162ea..fe98dfbf287a8ad 100644
--- a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
+++ b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
@@ -628,15 +628,11 @@ static Operation *replaceForAllWithNewSignature(
   return newforallOp;
 }
 
-/// Find the first "extract" user of `producerOp` and tile it right before its
-/// use. The tiled op is fused under the `containingOp`.
-/// Return this fused op on success or nullptr if anything fails.
-/// If tiled op has uses that are dominated by `containingOp`, return
-/// a new `containingOp` with results of the fused op appended to
-/// results of the `containingOp` or nullptr if there are no dominated uses.
-static std::tuple<SmallVector<Operation *>, Operation *>
-tileAndFuseFirstExtractUse(RewriterBase &rewriter, Diagnostic &diag,
-                           Operation *producerOp, Operation *containingOp) {
+std::tuple<SmallVector<Operation *>, Operation *>
+mlir::transform::tileAndFuseFirstExtractUse(RewriterBase &rewriter,
+                                            Diagnostic &diag,
+                                            Operation *producerOp,
+                                            Operation *containingOp) {
   LLVM_DEBUG(DBGS() << "Try to fuse a direct extract use\n");
   auto tileableProducer = dyn_cast<TilingInterface>(producerOp);
   if (!tileableProducer) {
@@ -710,14 +706,8 @@ tileAndFuseFirstExtractUse(RewriterBase &rewriter, Diagnostic &diag,
   return std::make_tuple(tileAndFuseResult->tiledOps, newContainingOp);
 }
 
-/// First, find the first "scf::ForallOp" user of `producerOp` and ensure
-/// it is exactly the `containingOp`, otherwise bail.
-/// Then, find the first "extract" user of the tied block argument and tile it
-/// right before its "extract" use. The tiled op is fused under the
-/// `containingOp`.
-/// Return this fused op on success or nullptr if anything fails.
-static SmallVector<Operation *>
-tileAndFuseFirstExtractUseThroughContainingOpBlockArgument(
+SmallVector<Operation *>
+mlir::transform::tileAndFuseFirstExtractUseThroughContainingOpBlockArgument(
     RewriterBase &rewriter, Diagnostic &diag, Operation *producerOp,
     Operation *containingOp) {
   LLVM_DEBUG(DBGS() << "Try to fuse an extract use through block argument\n");
@@ -819,9 +809,10 @@ tileAndFuseFirstExtractUseThroughContainingOpBlockArgument(
   return tileAndFuseResult->tiledOps;
 }
 
-static Operation *cloneAndFuseFirstUse(RewriterBase &rewriter, Diagnostic &diag,
-                                       Operation *producerOp,
-                                       Operation *containingOp) {
+Operation *mlir::transform::cloneAndFuseFirstUse(RewriterBase &rewriter,
+                                                 Diagnostic &diag,
+                                                 Operation *producerOp,
+                                                 Operation *containingOp) {
   LLVM_DEBUG(DBGS() << "Try to fuse an use by cloning\n");
 
   // Gather all uses inside the containing op.

``````````

</details>


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