Revert "[mlir][linalg] Add a new pattern to handle folding unit reduction dims."

[hanchung]

Summary

This reverts commit 6eee66d12ab33f35a37a1514342b51ae93d175e8.

It breaks builds, see https://lab.llvm.org/buildbot/#/builders/61/builds/35742

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

 
 #include "mlir/Dialect/Linalg/Passes.h"
 
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
 #include "mlir/Dialect/Linalg/Utils/Utils.h"
-#include "mlir/Dialect/MemRef/Transforms/Passes.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
-#include "mlir/Dialect/Tensor/Utils/Utils.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/Transforms/FoldUtils.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
-#include "llvm/ADT/SetVector.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 
 namespace mlir {
 #define GEN_PASS_DEF_LINALGFOLDUNITEXTENTDIMS
 #include "mlir/Dialect/Linalg/Passes.h.inc"
 } // namespace mlir
 
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
-/// 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>
-///  %2 = linalg.generic {indexing_maps = [affine_map<(d0) -> (0, d0, 0, 0)>,
-///                                        affine_map<(d0) -> (0, d0)>],
-///                       iterator_types = ["parallel"]}
-///    ins(%arg0 : tensor<1x?x1x1xf32>)
-///    outs(%1 : tensor<1x1xf32>) {
-///  ^bb0(%in: f32, %out: f32):
-///    %3 = arith.addf %in, %out : f32
-///    linalg.yield %3 : f32
-///  } -> tensor<1x1xf32>
-///
-///  into
-///
-///  %0 = tensor.empty() : tensor<1x1xf32>
-///  %1 = linalg.fill
-///    ins(%cst : f32)
-///    outs(%0 : tensor<1x1xf32>) -> tensor<1x1xf32>
-///  %2 = tensor.empty() : tensor<1x1xf32>
-///  %3 = linalg.generic {indexing_maps = [affine_map<(d0) -> (0, d0, 0, 0)>,
-///                                        affine_map<(d0) -> (0, d0)>,
-///                                        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> {
-  using OpRewritePattern<GenericOp>::OpRewritePattern;
-  LogicalResult matchAndRewrite(GenericOp genericOp,
-                                PatternRewriter &rewriter) const override {
-    if (!genericOp.hasTensorSemantics())
-      return failure();
-    if (genericOp.getNumParallelLoops() != genericOp.getNumLoops())
-      return failure();
-
-    auto outputOperands = genericOp.getDpsInitOperands();
-    SetVector<OpOperand *> candidates;
-    for (OpOperand *op : outputOperands) {
-      if (genericOp.getMatchingBlockArgument(op).use_empty())
-        continue;
-      candidates.insert(op);
-    }
-
-    if (candidates.empty())
-      return failure();
-
-    // Compute the modified indexing maps.
-    int64_t origNumInput = genericOp.getNumDpsInputs();
-    SmallVector<Value> newInputOperands = genericOp.getDpsInputOperands();
-    SmallVector<AffineMap> indexingMaps = genericOp.getIndexingMapsArray();
-    SmallVector<AffineMap> newIndexingMaps;
-    newIndexingMaps.append(indexingMaps.begin(),
-                           std::next(indexingMaps.begin(), origNumInput));
-    for (OpOperand *op : candidates) {
-      newInputOperands.push_back(op->get());
-      newIndexingMaps.push_back(genericOp.getMatchingIndexingMap(op));
-    }
-    newIndexingMaps.append(std::next(indexingMaps.begin(), origNumInput),
-                           indexingMaps.end());
-
-    Location loc = genericOp.getLoc();
-    SmallVector<Value> newOutputOperands = outputOperands;
-    for (OpOperand *op : candidates) {
-      OpBuilder::InsertionGuard guard(rewriter);
-      rewriter.setInsertionPointAfterValue(op->get());
-      auto elemType = op->get().getType().cast<ShapedType>().getElementType();
-      auto empty = rewriter.create<tensor::EmptyOp>(
-          loc, tensor::createDimValues(rewriter, loc, op->get()), elemType);
-
-      auto [start, end] = genericOp.getDpsInitsPositionRange();
-      newOutputOperands[op->getOperandNumber() - start] = empty.getResult();
-    }
-
-    auto newOp = rewriter.create<GenericOp>(
-        loc, genericOp.getResultTypes(), newInputOperands, newOutputOperands,
-        newIndexingMaps, genericOp.getIteratorTypesArray(),
-        /*bodyBuild=*/nullptr, linalg::getPrunedAttributeList(genericOp));
-
-    Region &region = newOp.getRegion();
-    Block *block = new Block();
-    region.push_back(block);
-    BlockAndValueMapping mapper;
-    OpBuilder::InsertionGuard guard(rewriter);
-    rewriter.setInsertionPointToStart(block);
-    for (auto bbarg : genericOp.getRegionInputArgs())
-      mapper.map(bbarg, block->addArgument(bbarg.getType(), loc));
-
-    for (OpOperand *op : candidates) {
-      BlockArgument bbarg = genericOp.getMatchingBlockArgument(op);
-      mapper.map(bbarg, block->addArgument(bbarg.getType(), loc));
-    }
-
-    for (OpOperand *op : outputOperands) {
-      BlockArgument bbarg = genericOp.getMatchingBlockArgument(op);
-      if (candidates.count(op))
-        block->addArgument(bbarg.getType(), loc);
-      else
-        mapper.map(bbarg, block->addArgument(bbarg.getType(), loc));
-    }
-
-    for (auto &op : genericOp.getBody()->getOperations()) {
-      rewriter.clone(op, mapper);
-    }
-    rewriter.replaceOp(genericOp, newOp.getResults());
-
-    return success();
-  }
-};
-
 struct UnitExtentReplacementInfo {
   Type type;
   AffineMap indexMap;
   ArrayAttr reassociation;
 };
 } // namespace
 
 /// Utility function for replacing operands/results to a linalg generic

 };
 } // namespace
 
 /// Patterns that are used to canonicalize the use of unit-extent dims for
 /// broadcasting.
 void mlir::linalg::populateFoldUnitExtentDimsPatterns(
     RewritePatternSet &patterns) {
   auto *context = patterns.getContext();
-  patterns.add<FoldUnitDimLoops, AddInitOperandsToInput, ReplaceUnitExtents,
-               RankReducedExtractSliceOp,
+  patterns.add<FoldUnitDimLoops, ReplaceUnitExtents, RankReducedExtractSliceOp,
                RankReducedInsertSliceOp<tensor::InsertSliceOp>,
                RankReducedInsertSliceOp<tensor::ParallelInsertSliceOp>>(
       context);
   linalg::FillOp::getCanonicalizationPatterns(patterns, context);
   tensor::CollapseShapeOp::getCanonicalizationPatterns(patterns, context);
   tensor::EmptyOp::getCanonicalizationPatterns(patterns, context);
   tensor::ExpandShapeOp::getCanonicalizationPatterns(patterns, context);
-  memref::populateResolveRankedShapeTypeResultDimsPatterns(patterns);
-  memref::populateResolveShapedTypeResultDimsPatterns(patterns);
 }
 
 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>(context);
     else
       populateFoldUnitExtentDimsPatterns(patterns);
     (void)applyPatternsAndFoldGreedily(op, std::move(patterns));
   }
 };
 } // namespace
 
 std::unique_ptr<Pass> mlir::createLinalgFoldUnitExtentDimsPass() {
   return std::make_unique<LinalgFoldUnitExtentDimsPass>();
 }

mlir/test/Dialect/Linalg/drop-unit-extent-dims.mlir

func.func @unit_dim_for_both_reduction(%arg0: tensor<1x?x1x1xf32>) -> tensor<1x1xf32> {
   return %3 : tensor<1x1xf32>
 }
 //  CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0) -> (d0)>
 //      CHECK: func @unit_dim_for_both_reduction
 // CHECK-SAME:   %[[ARG0:.+]]: tensor<1x?x1x1xf32>
 //  CHECK-DAG:   %[[RESHAPE:.+]] = tensor.collapse_shape %[[ARG0]] {{\[}}[0, 1, 2, 3]
 //      CHECK:   %[[INIT:.+]] = tensor.empty() : tensor<1xf32>
 //      CHECK:   %[[FILL:.+]] = linalg.fill ins(%{{.+}}{{.*}}outs(%[[INIT]]
-//      CHECK:   %[[INIT2:.+]] = tensor.empty() : tensor<1xf32>
 //      CHECK:   %[[RESULT:.+]] = linalg.generic
-// CHECK-SAME:     indexing_maps = [#[[MAP2]], #[[MAP2]], #[[MAP2]]]
+// CHECK-SAME:     indexing_maps = [#[[MAP2]], #[[MAP2]]]
 // CHECK-SAME:     iterator_types = ["parallel"]
-// CHECK-SAME:     ins(%[[RESHAPE]], %[[FILL]] : tensor<?xf32>, tensor<1xf32>)
-// CHECK-SAME:     outs(%[[INIT2]] : tensor<1xf32>)
+// CHECK-SAME:     ins(%[[RESHAPE]] : tensor<?xf32>)
+// CHECK-SAME:     outs(%[[FILL]] : tensor<1xf32>)
 //      CHECK:   %[[RESULT_RESHAPE:.+]] = tensor.expand_shape %[[RESULT]] {{\[}}[0, 1]]
 //      CHECK:   return %[[RESULT_RESHAPE]]
 
 // -----
 
 func.func @unit_dim_for_reduction_inner(%arg0: tensor<?x1x?x1xf32>) -> tensor<?x1xf32> {
   %cst = arith.constant 1.000000e+00 : f32
   %c2 = arith.constant 2 : index

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

deps = [
         ":IR",
         ":LinalgAnalysis",
         ":LinalgDialect",
         ":LinalgPassIncGen",
         ":LinalgStructuredOpsIncGen",
         ":LinalgUtils",
         ":MathDialect",
         ":MemRefDialect",
-        ":MemRefTransforms",
         ":Pass",
         ":SCFDialect",
         ":SCFTransforms",
         ":SCFUtils",
         ":SparseTensorDialect",
         ":Support",
         ":TensorDialect",
         ":TensorTilingInterfaceImpl",