[mlir] Export `collapseGenericOpIterationDims` (NFC)

[guraypp]

Summary

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.

[guraypp]

rebase

mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h

 ///
 ///  return %4 : tensor<16x32xf32>
 /// ```
 FailureOr<SplitReductionResult>
 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
 
 #endif // MLIR_DIALECT_LINALG_TRANSFORMS_TRANSFORMS_H

mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp

static Value getCollapsedOpOperand(Location loc, GenericOp genericOp,
   return reshapeOp.getResult();
 }
 
 /// Modify the `linalg.index` operations in the original generic op, to its
 /// value in the collapsed operation.
 void generateCollapsedIndexingRegion(Location loc, Block *block,
                                      const CollapsingInfo &collapsingInfo,
                                      ValueRange loopRange,
-                                     PatternRewriter &rewriter) {
+                                     RewriterBase &rewriter) {
   OpBuilder::InsertionGuard g(rewriter);
   rewriter.setInsertionPointToStart(block);
 
   // Collect all the original index ops.
   auto indexOps = llvm::to_vector(block->getOps<linalg::IndexOp>());
 
   // For each folded dimension list resolve the original induction variable
   // values in terms of the folded dimension induction variable.
void generateCollapsedIndexingRegion(Location loc, Block *block,
 
   for (auto indexOp : indexOps) {
     auto dim = indexOp.getDim();
     rewriter.replaceOp(indexOp, indexReplacementVals[dim]);
   }
 }
 
 /// 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) {
         return foldedDims.size() <= 1;
       }))
     return failure();
 
   CollapsingInfo collapsingInfo;
SmallVector<ReassociationIndices> collapsableIterationDims =
         controlCollapseDimension(genericOp);
     if (collapsableIterationDims.empty())
       return failure();
 
     Optional<SmallVector<Value>> replacements = collapseGenericOpIterationDims(
         genericOp, collapsableIterationDims, rewriter);
     if (!replacements) {
       return rewriter.notifyMatchFailure(genericOp,
-                                         "failed to collpase dimensions");
+                                         "failed to collapse dimensions");
     }
     rewriter.replaceOp(genericOp, *replacements);
     return success();
   }
 
 private:
   GetCollapsableDimensionsFn controlCollapseDimension;
 };