[Mlir-commits] [mlir] 135977c - [mlir] Export `collapseGenericOpIterationDims` (NFC)

[Guray Ozen]

Author: Guray Ozen
Date: 2022-11-28T13:54:26+01:00
New Revision: 135977c92a173b716d76c026972a677681f5c6bd

URL: https://github.com/llvm/llvm-project/commit/135977c92a173b716d76c026972a677681f5c6bd
DIFF: https://github.com/llvm/llvm-project/commit/135977c92a173b716d76c026972a677681f5c6bd.diff

LOG: [mlir] Export `collapseGenericOpIterationDims` (NFC)

This revision exports `collapseGenericOpIterationDims` to a header so it can be used outside of the pattern.  We have use-case where we want to call this function directly.

Reviewed By: springerm

Differential Revision: https://reviews.llvm.org/D138697

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
    mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h b/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
index 40bf31c3ab474..29072c2443f6a 100644
--- a/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
+++ b/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
@@ -1069,6 +1069,12 @@ splitReductionByScaling(PatternRewriter &b, LinalgOp op,
                         const ControlSplitReductionFn &controlSplitReductionFn,
                         bool useAlloc = false);
 
+/// Collapses dimensions of linalg.generic operation. It also collapses inputs
+/// before the op and expands outputs after the op.
+FailureOr<SmallVector<Value>> collapseGenericOpIterationDims(
+    GenericOp genericOp, ArrayRef<ReassociationIndices> foldedIterationDims,
+    RewriterBase &rewriter);
+
 } // namespace linalg
 } // namespace mlir
 

diff  --git a/mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp b/mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp
index 11eed11bbabf9..26636bc053653 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp
@@ -1353,7 +1353,7 @@ static Value getCollapsedOpOperand(Location loc, GenericOp genericOp,
 void generateCollapsedIndexingRegion(Location loc, Block *block,
                                      const CollapsingInfo &collapsingInfo,
                                      ValueRange loopRange,
-                                     PatternRewriter &rewriter) {
+                                     RewriterBase &rewriter) {
   OpBuilder::InsertionGuard g(rewriter);
   rewriter.setInsertionPointToStart(block);
 
@@ -1389,9 +1389,9 @@ void generateCollapsedIndexingRegion(Location loc, Block *block,
 }
 
 /// Implementation of fusion with reshape operation by collapsing dimensions.
-static FailureOr<SmallVector<Value>> collapseGenericOpIterationDims(
+FailureOr<SmallVector<Value>> mlir::linalg::collapseGenericOpIterationDims(
     GenericOp genericOp, ArrayRef<ReassociationIndices> foldedIterationDims,
-    PatternRewriter &rewriter) {
+    RewriterBase &rewriter) {
   // Bail on trivial no-op cases.
   if (genericOp.getNumLoops() <= 1 || foldedIterationDims.empty() ||
       llvm::all_of(foldedIterationDims, [](ReassociationIndicesRef foldedDims) {
@@ -1570,7 +1570,7 @@ class CollapseLinalgDimensions : public OpRewritePattern<GenericOp> {
         genericOp, collapsableIterationDims, rewriter);
     if (!replacements) {
       return rewriter.notifyMatchFailure(genericOp,
-                                         "failed to collpase dimensions");
+                                         "failed to collapse dimensions");
     }
     rewriter.replaceOp(genericOp, *replacements);
     return success();