[mlir][bufferization] Support multiple leaves in EmptyTensorElimination

[springerm]

Summary

Support cases where a source tensor can be traced back to multiple possible tensor.empty ops.

Example:

%0 = scf.if ... {
  %1 = tensor.empty()
  %2 = "some_op"(%1)
  scf.yield %2
} else {
  ...
}
%3 = tensor.insert_slice %0 into %t

The tensor.empty inside the scf.if can now be rewritten to a tensor.extract_slice, allowing it to bufferize without an allocation.

Depends On: D142129

[springerm]

rebase

[springerm]

rebase

[springerm]

rebase

[springerm]

simplify impl

[springerm]

update

[mravishankar]

Have one comment here, more for my understanding of the nuances. Change itself looks good!

mlir/lib/Dialect/Bufferization/Transforms/EmptyTensorElimination.cpp

116

Comment needs to be updated? Now it seems like you can return multiple tensor.empty... I dont know when that happens though

mlir/lib/Dialect/Bufferization/Transforms/EmptyTensorElimination.cpp

if (!insertionPointDominatesUses(domInfo, insertionPoint, emptyTensorOp))
       continue;
     return insertionPoint;
   }
 
   // No suitable insertion point was found.
   return nullptr;
 }
 
+namespace {
+/// A replacement for a tensor.empty op.
+struct EmptyTensorReplacement {
+  OpOperand *operand;
+  SmallVector<Value> neededValues;
+  bool isEquivalent;
+};
+} // namespace
+
 /// Try to eliminate tensor::EmptyOps inside `op`. A tensor::EmptyOp is replaced
 /// with the result of `rewriteFunc` if it is anchored on a matching
 /// OpOperand. "Anchored" means that there is a path on the reverse SSA use-def
 /// chain, starting from the OpOperand and always following the aliasing
 /// OpOperand, that eventually ends at a single tensor::EmptyOp.
 LogicalResult mlir::bufferization::eliminateEmptyTensors(
     RewriterBase &rewriter, Operation *op, AnalysisState &state,
     AnchorMatchFn anchorMatchFunc, RewriteFn rewriteFunc) {
   OpBuilder::InsertionGuard g(rewriter);
+  DenseMap<Operation *, EmptyTensorReplacement> emptyTensors;
 
-  WalkResult status = op->walk([&](Operation *op) {
+  op->walk([&](Operation *op) {
     for (OpOperand &operand : op->getOpOperands()) {
       // Skip operands that do not bufferize inplace.
       if (!state.isInPlace(operand))
         continue;
       // All values that are needed to create the replacement op.
       SmallVector<Value> neededValues;
       // Is this a matching OpOperand?
       if (!anchorMatchFunc(operand, neededValues))
         continue;
-      SetVector<Value> maybeEmptyTensor = state.findValueInReverseUseDefChain(
+      SetVector<Value> maybeEmptyTensors = state.findValueInReverseUseDefChain(
           operand.get(), /*condition=*/[&](Value val) { return false; },
           /*followEquivalentOnly=*/true);
 
+      for (Value v : maybeEmptyTensors) {
         // Replace only if the reverse use-def chain ends at exactly one
         // tensor::EmptyOp.
-      if (maybeEmptyTensor.size() != 1 ||
-          !maybeEmptyTensor.front().getDefiningOp<tensor::EmptyOp>())
+        auto emptyTensor = v.getDefiningOp<tensor::EmptyOp>();
+        if (!emptyTensor)
           continue;
-      Value emptyTensor = maybeEmptyTensor.front();
 
         // Replace only if the types match.
         // TODO: This could be extended to support IR such as:
         // %0 = tensor.empty() : tensor<128xf32>
         // %1 = "some_op"(%0) : (tensor<128xf32>) -> (tensor<128xf32>)
         // %2 = tensor.expand_shape %1 ...
         // %3 = tensor.insert_slice %2 into ...
-      if (emptyTensor.getType() != operand.get().getType())
+        if (v.getType() != operand.get().getType())
+          continue;
+
+        // Check if there is already a good replacement for this op.
+        auto it = emptyTensors.find(emptyTensor.getOperation());
+        if (it != emptyTensors.end() && it->second.isEquivalent)
           continue;
 
         // Find a suitable insertion point.
         Operation *insertionPoint =
-          findValidInsertionPoint(emptyTensor.getDefiningOp(), neededValues);
-      if (!insertionPoint)
-        continue;
+            findValidInsertionPoint(emptyTensor.getOperation(), neededValues);
+        if (!insertionPoint)
+          continue;
+
+        bool isEquivalent =
+            state.areEquivalentBufferizedValues(v, operand.get());
+        emptyTensors[emptyTensor.getOperation()] =
+            EmptyTensorReplacement{&operand, neededValues, isEquivalent};
+      }
+    }
+  });
+
+  for (const auto &it : emptyTensors) {
+    // Find a suitable insertion point.
+    Operation *insertionPoint =
+        findValidInsertionPoint(it.first, it.second.neededValues);
 
-      // Create a replacement for the tensor::EmptyOp.
-      rewriter.setInsertionPoint(insertionPoint);
-      Value replacement = rewriteFunc(rewriter, emptyTensor.getLoc(), operand);
-      if (!replacement)
-        continue;
+    Value replacement =
+        rewriteFunc(rewriter, it.first->getLoc(), *it.second.operand);
 
-      // Replace the tensor::EmptyOp.
-      rewriter.replaceOp(emptyTensor.getDefiningOp(), replacement);
-    }
+    rewriter.replaceOp(it.first, replacement);
 
-    // Advance to the next operation.
-    return WalkResult::advance();
-  });
-
-  return failure(status.wasInterrupted());
+  return success();
 }
 
 /// Try to eliminate tensor::EmptyOps inside `op`. An tensor::EmptyOp can be
 /// eliminated if it is eventually inserted into another tensor (and some other
 /// conditions are met).
 ///
 /// E.g.:
 /// %0 = tensor.empty()
registry
         .insert<bufferization::BufferizationDialect, tensor::TensorDialect>();
   }
 };
 } // namespace
 
 void EmptyTensorElimination::runOnOperation() {
   Operation *op = getOperation();
   OneShotBufferizationOptions options;
+  options.allowReturnAllocs = true;
   OneShotAnalysisState state(op, options);
   if (failed(analyzeOp(op, state))) {
     signalPassFailure();
     return;
   }
 
   IRRewriter rewriter(op->getContext());
   if (failed(bufferization::insertSliceAnchoredEmptyTensorEliminationStep(
           rewriter, op, state)))
     signalPassFailure();
 }
 
 std::unique_ptr<Pass> mlir::bufferization::createEmptyTensorEliminationPass() {
   return std::make_unique<EmptyTensorElimination>();
 }

mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-empty-tensor-elimination.mlir

%r1 = scf.foreach_thread (%iv) in (%c512) shared_outs(%o = %t) -> (tensor<?xf32>) {
     // Self-copy canonicalizes away later.
     scf.foreach_thread.perform_concurrently {
       tensor.parallel_insert_slice %f into %o[42][%sz][1]: tensor<?xf32> into tensor<?xf32>
     }
   }
 
   return %r1: tensor<?xf32>
 }
+
+// -----
+
+// CHECK-LABEL: func @eleminate_multiple_ops(
+//  CHECK-SAME:   %[[FUNC_ARG:[0-9a-zA-Z]*]]: memref<?xf32>
+//  CHECK-SAME:   %[[sz:[0-9a-zA-Z]*]]: index
+func.func @eleminate_multiple_ops(%t: tensor<?xf32> {bufferization.buffer_layout = affine_map<(d0) -> (d0)>}, %sz: index, %c: i1)
+    -> (tensor<?xf32>)
+{
+  %cst1 = arith.constant 0.0: f32
+  %cst2 = arith.constant 1.0: f32
+
+  // CHECK: %[[r:.*]] = scf.if %{{.*}} -> (memref
+  %if = scf.if %c -> tensor<?xf32> {
+    // CHECK: %[[T_SUBVIEW_1:.*]] =  memref.subview %[[FUNC_ARG]][42] [%[[sz]]] [1]
+    %a1 = tensor.empty(%sz) : tensor<?xf32>
+    // CHECK: linalg.fill ins({{.*}} : f32) outs(%[[T_SUBVIEW_1]] : memref<?xf32
+    %f1 = linalg.fill ins(%cst1 : f32) outs(%a1 : tensor<?xf32>) -> tensor<?xf32>
+    // CHECK: scf.yield %[[T_SUBVIEW_1]]
+    scf.yield %f1 : tensor<?xf32>
+  } else {
+      // CHECK: %[[T_SUBVIEW_2:.*]] =  memref.subview %[[FUNC_ARG]][42] [%[[sz]]] [1]
+    %a2 = tensor.empty(%sz) : tensor<?xf32>
+    // CHECK: linalg.fill ins({{.*}} : f32) outs(%[[T_SUBVIEW_2]] : memref<?xf32
+    %f2 = linalg.fill ins(%cst2 : f32) outs(%a2 : tensor<?xf32>) -> tensor<?xf32>
+    // CHECK: scf.yield %[[T_SUBVIEW_2]]
+    scf.yield %f2 : tensor<?xf32>
+  }
+
+  // Self-copy could canonicalize away later.
+  // CHECK: %[[T_SUBVIEW_3:.*]] =  memref.subview %[[FUNC_ARG]][42] [%[[sz]]] [1]
+  // CHECK: memref.copy %[[r]], %[[T_SUBVIEW_3]]
+  %r1 = tensor.insert_slice %if into %t[42][%sz][1]: tensor<?xf32> into tensor<?xf32>
+  return %r1: tensor<?xf32>
+}