[mlir][linalg] Remove RemoveOutsDependency pattern

[springerm]

Summary

This pattern does not work well with bufferization because it destroys tensor SSA use-def chains.

Depends On D129354

[mravishankar]

This should not affect bufferization. It breaks unnecessary dependency. Can you elaborate more about the failure case here.

[springerm]

Bufferization analyzes tensor use-def chains to decide what should bufferize in-place. A linalg.generic may have a certain tensor as an out operand even if it is not used in the computation. This may be on purpose to ensure that a certain buffer is used for in-place bufferization. This pattern would replace such as an out tensor with an init_tensor, thus breaking the use-def chain.

Removing this patterns breaks an IREE test case, that's why I put it on "changes planned" for the moment. In what way does IREE depend on this pattern?

I stumbled upon this pattern when looking at the sparse compiler pipeline, where it causes an issue with bufferization.

[mravishankar]

In D129355#3644375, @springerm wrote:

Bufferization analyzes tensor use-def chains to decide what should bufferize in-place. A linalg.generic may have a certain tensor as an out operand even if it is not used in the computation. This may be on purpose to ensure that a certain buffer is used for in-place bufferization. This pattern would replace such as an out tensor with an init_tensor, thus breaking the use-def chain.

Removing this patterns breaks an IREE test case, that's why I put it on "changes planned" for the moment. In what way does IREE depend on this pattern?

I stumbled upon this pattern when looking at the sparse compiler pipeline, where it causes an issue with bufferization.

I think you are implicitly expecting every program to be in destination passing style. I realize that this is what bufferization needs, but this is not the form of the program to start with during compilation of full programs (unless the destination passing style is lifted all the way into user code which no framework does today). So stepping out of the destination passing style requirements, adding a dependence between an op to another then makes it hard to disambiguate what is a real dependence vs what is a dependence that is introduced by bufferization requirements. For things that want to do producer-consumer fusion (like these patterns), its much easier to have to deal with real dependence as opposed to dependence due to bufferization.

Having said that though, it does seem that this particular pattern breaks programs that are already in destination passing style.... If you remove this pattern, there is something else needed in the fusion patterns that would be needed to avoid fusing across these dependences.... At the very least for your use case you can create a separate populate* method that doesnt have this pattern.

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

 
 } // namespace
 
 //===---------------------------------------------------------------------===//
 // Miscellaneous patterns that help fusion.
 //===---------------------------------------------------------------------===//
 
 namespace {
-/// Forces `outs` operands of linalg operations to use `linalg.init_tensor` if
-/// the value of the `outs` operand is not used within the op.  This is only
-/// implemented for `linalg.generic` operations for now, but should hold for all
-/// linalg structured ops.
-struct RemoveOutsDependency : public OpRewritePattern<GenericOp> {
-  using OpRewritePattern<GenericOp>::OpRewritePattern;
-
-  LogicalResult matchAndRewrite(GenericOp op,
-                                PatternRewriter &rewriter) const override {
-    rewriter.startRootUpdate(op);
-    bool modifiedOutput = false;
-    Location loc = op.getLoc();
-    for (OpOperand *opOperand : op.getOutputOperands()) {
-      if (!op.payloadUsesValueFromOperand(opOperand)) {
-        Value operandVal = opOperand->get();
-        auto operandType = operandVal.getType().dyn_cast<RankedTensorType>();
-        if (!operandType)
-          continue;
-
-        // If outs is sparse, leave it to the sparse compiler.
-        if (sparse_tensor::getSparseTensorEncoding(operandVal.getType()))
-          continue;
-
-        // If outs is already an `init_tensor` operation, nothing to do.
-        auto definingOp = operandVal.getDefiningOp<InitTensorOp>();
-        if (definingOp)
-          continue;
-        modifiedOutput = true;
-        SmallVector<Value> dynamicDims;
-        for (const auto &dim : llvm::enumerate(operandType.getShape())) {
-          if (dim.value() != ShapedType::kDynamicSize)
-            continue;
-          dynamicDims.push_back(rewriter.createOrFold<tensor::DimOp>(
-              loc, operandVal, dim.index()));
-        }
-        Value initTensor = rewriter.create<InitTensorOp>(
-            loc, dynamicDims, operandType.getShape(),
-            operandType.getElementType());
-        op->setOperand(opOperand->getOperandNumber(), initTensor);
-      }
-    }
-    if (!modifiedOutput) {
-      rewriter.cancelRootUpdate(op);
-      return failure();
-    }
-    rewriter.finalizeRootUpdate(op);
-    return success();
-  }
-};
-
 /// Fold linalg.fill into linalg.generic
 struct FoldFillWithGenericOp : public OpRewritePattern<GenericOp> {
   using OpRewritePattern<GenericOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(GenericOp genericOp,
                                 PatternRewriter &rewriter) const override {
     if (!genericOp.hasTensorSemantics())
       return failure();
patterns.add<FoldWithProducerReshapeOpByCollapsing>(patterns.getContext(),
                                                       controlFoldingReshapes);
 }
 
 void mlir::linalg::populateElementwiseOpsFusionPatterns(
     RewritePatternSet &patterns,
     const ControlFusionFn &controlElementwiseOpsFusion) {
   auto *context = patterns.getContext();
   patterns.add<FuseElementwiseOps>(context, controlElementwiseOpsFusion);
-  patterns.add<FoldFillWithGenericOp, FoldScalarOrSplatConstant,
-               RemoveOutsDependency>(context);
+  patterns.add<FoldFillWithGenericOp, FoldScalarOrSplatConstant>(context);
 }
 
 //===---------------------------------------------------------------------===//
 // Passes
 //===---------------------------------------------------------------------===//
 
 namespace {
 

mlir/test/Dialect/Linalg/fusion-elementwise-ops.mlir

^bb0(%arg0 : f32, %arg1 : f32):
         %0 = arith.addf %arg0, %arg1 : f32
         linalg.yield %0 : f32
   } -> tensor<3xf32>
   return %result : tensor<3xf32>
 }
 
 // -----
 
-#map = affine_map<(d0, d1) -> (d0, d1)>
-#trait = {
-  indexing_maps = [#map, #map],
-  iterator_types = ["parallel", "parallel"]
-}
-func.func @break_outs_dependency(%arg0 : tensor<?x?xf32>) -> tensor<?x?xf32>
-{
-  %0 = linalg.generic #trait ins(%arg0 : tensor<?x?xf32>) outs(%arg0 : tensor<?x?xf32>) {
-       ^bb0(%arg1 : f32, %arg2 : f32) :
-         %1 = arith.addf %arg1, %arg1 : f32
-         linalg.yield %1 : f32
-       } -> tensor<?x?xf32>
-  %2 = linalg.generic #trait ins(%0 : tensor<?x?xf32>) outs(%0 : tensor<?x?xf32>) {
-       ^bb0(%arg1 : f32, %arg2 : f32) :
-         %3 = arith.mulf %arg1, %arg1 : f32
-         linalg.yield %3 : f32
-       } -> tensor<?x?xf32>
-  return %2 : tensor<?x?xf32>
-}
-//      CHECK: func @break_outs_dependency(
-// CHECK-SAME:   %[[ARG0:.+]]: tensor<?x?xf32>)
-//  CHECK-DAG:   %[[C0:.+]] = arith.constant 0 : index
-//  CHECK-DAG:   %[[C1:.+]] = arith.constant 1 : index
-//  CHECK-DAG:   %[[D0:.+]] = tensor.dim %[[ARG0]], %[[C0]]
-//  CHECK-DAG:   %[[D1:.+]] = tensor.dim %[[ARG0]], %[[C1]]
-//  CHECK-DAG:   %[[INIT:.+]] = linalg.init_tensor [%[[D0]], %[[D1]]]
-//      CHECK:   %[[GENERIC1:.+]] = linalg.generic
-// CHECK-SAME:     outs(%[[INIT]] : tensor<?x?xf32>)
-//  CHECK-DAG:   %[[D0:.+]] = tensor.dim %[[GENERIC1]], %[[C0]]
-//  CHECK-DAG:   %[[D1:.+]] = tensor.dim %[[GENERIC1]], %[[C1]]
-//  CHECK-DAG:   %[[INIT:.+]] = linalg.init_tensor [%[[D0]], %[[D1]]]
-//      CHECK:   %[[RESULT:.+]] = linalg.generic
-// CHECK-SAME:     outs(%[[INIT]] : tensor<?x?xf32>)
-
-// -----
-
 func.func @fuse_scalar_constant(%arg0 : tensor<?x?xf32>) -> (tensor<?x?xf32>, tensor<?x?xi32>) {
   %cst = arith.constant 4.0 : f32
   %c42 = arith.constant 42 : i32
   %c0 = arith.constant 0 : index
   %c1 = arith.constant 1 : index
   %d0 = tensor.dim %arg0, %c0 : tensor<?x?xf32>
   %d1 = tensor.dim %arg0, %c1 : tensor<?x?xf32>
   %0 = linalg.init_tensor[%d0, %d1] : tensor<?x?xf32>