[mlir][linalg][NFC] Split populateFoldUnitExtentDimsViaReshapesPatterns

[springerm]

Summary

MoveInitOperandsToInput is put into a separate populate... function because it can interfere with certain transformations.

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

 /// Patterns to fold unit-extent dimensions in operands/results of linalg ops on
 /// tensors via reassociative reshape ops.
 void populateFoldUnitExtentDimsViaReshapesPatterns(RewritePatternSet &patterns);
 
 /// Patterns to fold unit-extent dimensions in operands/results of linalg ops on
 /// tensors via rank-reducing slices.
 void populateFoldUnitExtentDimsViaSlicesPatterns(RewritePatternSet &patterns);
 
+/// A pattern that converts init operands to input operands.
+void populateMoveInitOperandsToInputPattern(RewritePatternSet &patterns);
+
 /// Patterns that are used to inline constant operands into linalg generic ops.
 void populateInlineConstantOperandsPatterns(RewritePatternSet &patterns);
 
 /// Patterns that are used to bubble up extract slice op above linalg op.
 void populateBubbleUpExtractSliceOpPatterns(RewritePatternSet &patterns);
 
 /// Adds patterns that waps tensor.extract_slice(linalg.fill(%cst, %init)) into
 /// linalg.fill(%cst, tensor.extract_slice(%init)).

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

LogicalResult matchAndRewrite(GenericOp genericOp,
     genericOp.setIndexingMapsAttr(newIndexingMapAttr);
     genericOp.setIteratorTypesAttr(ArrayAttr::get(context, newIteratorTypes));
     replaceUnitDimIndexOps(genericOp, unitDims, rewriter);
     rewriter.finalizeRootUpdate(genericOp);
     return success();
   }
 };
 
-/// Pattern to add init operands to ins when all the loops are parallel and
+/// Pattern to move init operands to ins when all the loops are parallel and
 /// blockArgument corresponding to init is used in the region. This is a fix-up
 /// when unit reduction dimensions are all folded away. In this context, it
 /// becomes a elementwise generic op. E.g., it converts
 ///
 ///  %0 = tensor.empty() : tensor<1x1xf32>
 ///  %1 = linalg.fill
 ///    ins(%cst : f32)
 ///    outs(%0 : tensor<1x1xf32>) -> tensor<1x1xf32>

 ///                                        affine_map<(d0) -> (0, d0)>],
 ///                       iterator_types = ["parallel"]}
 ///   ins(%arg0, %1 : tensor<1x?x1x1xf32>, tensor<1x1xf32>)
 ///   outs(%2 : tensor<1x1xf32>) {
 ///  ^bb0(%in: f32, %in_0: f32, %out: f32):
 ///    %4 = arith.addf %in, %in_0 : f32
 ///    linalg.yield %4 : f32
 ///  } -> tensor<1x1xf32>
-struct AddInitOperandsToInput : public OpRewritePattern<GenericOp> {
+struct MoveInitOperandsToInput : public OpRewritePattern<GenericOp> {
   using OpRewritePattern<GenericOp>::OpRewritePattern;
   LogicalResult matchAndRewrite(GenericOp genericOp,
                                 PatternRewriter &rewriter) const override {
     if (!genericOp.hasTensorSemantics())
       return failure();
     if (genericOp.getNumParallelLoops() != genericOp.getNumLoops())
       return failure();
 

 /// Patterns that are used to canonicalize the use of unit-extent dims for
 /// broadcasting.
 void mlir::linalg::populateFoldUnitExtentDimsViaReshapesPatterns(
     RewritePatternSet &patterns) {
   auto *context = patterns.getContext();
   patterns.add<ReplaceUnitExtents>(context,
                                    RankReductionStrategy::ReassociativeReshape);
   // TODO: Patterns unrelated to unit dim folding should be factored out.
-  patterns
-      .add<FoldUnitDimLoops, AddInitOperandsToInput, RankReducedExtractSliceOp,
+  patterns.add<FoldUnitDimLoops, RankReducedExtractSliceOp,
                RankReducedInsertSliceOp<tensor::InsertSliceOp>,
-           RankReducedInsertSliceOp<tensor::ParallelInsertSliceOp>>(context);
+               RankReducedInsertSliceOp<tensor::ParallelInsertSliceOp>>(
+      context);
   linalg::FillOp::getCanonicalizationPatterns(patterns, context);
   tensor::CollapseShapeOp::getCanonicalizationPatterns(patterns, context);
   tensor::EmptyOp::getCanonicalizationPatterns(patterns, context);
   tensor::ExpandShapeOp::getCanonicalizationPatterns(patterns, context);
   tensor::populateFoldTensorEmptyPatterns(patterns);
   memref::populateResolveRankedShapeTypeResultDimsPatterns(patterns);
   memref::populateResolveShapedTypeResultDimsPatterns(patterns);
 }
 
 void mlir::linalg::populateFoldUnitExtentDimsViaSlicesPatterns(
     RewritePatternSet &patterns) {
   auto *context = patterns.getContext();
   patterns.add<ReplaceUnitExtents>(context,
                                    RankReductionStrategy::ExtractInsertSlice);
   patterns.add<FoldUnitDimLoops>(context);
 }
 
+void mlir::linalg::populateMoveInitOperandsToInputPattern(
+    RewritePatternSet &patterns) {
+  patterns.add<MoveInitOperandsToInput>(patterns.getContext());
+}
+
 namespace {
 /// Pass that removes unit-extent dims within generic ops.
 struct LinalgFoldUnitExtentDimsPass
     : public impl::LinalgFoldUnitExtentDimsBase<LinalgFoldUnitExtentDimsPass> {
   void runOnOperation() override {
     Operation *op = getOperation();
     MLIRContext *context = op->getContext();
     RewritePatternSet patterns(context);
     if (foldOneTripLoopsOnly) {
-      patterns.add<FoldUnitDimLoops, AddInitOperandsToInput>(context);
+      patterns.add<FoldUnitDimLoops, MoveInitOperandsToInput>(context);
     } else if (useRankReducingSlices) {
       populateFoldUnitExtentDimsViaSlicesPatterns(patterns);
+      populateMoveInitOperandsToInputPattern(patterns);
     } else {
       populateFoldUnitExtentDimsViaReshapesPatterns(patterns);
+      populateMoveInitOperandsToInputPattern(patterns);
     }
     (void)applyPatternsAndFoldGreedily(op, std::move(patterns));
   }
 };
 } // namespace
 
 std::unique_ptr<Pass> mlir::createLinalgFoldUnitExtentDimsPass() {
   return std::make_unique<LinalgFoldUnitExtentDimsPass>();
 }