[mlir][SCF][bufferize] Implement `resolveConflicts` for SCF ops

[springerm]

Summary

scf::ForOp and scf::WhileOp must insert buffer copies not only for out-of-place bufferizations, but also to enforce additional invariants wrt. to buffer aliasing behavior. This is currently happening in the respective bufferize methods. With this change, the tensor copy insertion pass will also enforce these invariants by inserting copies. The bufferize methods can then be simplified and made independent of the AnalysisState data structure in a subsequent change.

Depends On D126821

[springerm]

update

[aartbik]

mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h

234

in variants -> invariants

mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h

struct BufferizationOptions {
   /// bufferization.to_memref and bufferization.to_tensor ops are inserted at
   /// the boundaries.
   bool allowUnknownOps = false;
 
   /// Specifies whether function boundaries (ops in the func dialect) should be
   /// bufferized or not.
   bool bufferizeFunctionBoundaries = false;
 
+  /// Certain ops have aliasing OpOperand/OpResult invariants (e.g., scf.for).
+  /// If this flag is set to `false`, those invariants are no longer enforced
+  /// with buffer copies.
+  ///
+  /// Note: Deactivating this flag can lead to incorrect bufferization results
+  /// when used incorrectly. This flag is useful with
+  /// `AlwaysCopyBufferizationState` which bufferizes all writing tensor
+  /// OpOperands out-of-place.
+  bool enforceAliasingInvariants = true;
+
   /// This flag controls buffer types on function signatures.
   ///
   /// * InferLayoutMap: All function parameter types have a fully dynamic layout
   ///   map, but function result types are inferred from the body of the
   ///   function.
   /// * FullyDynamicLayoutMap: All function parameter types and result types
   ///   have a fully dynamic layout map. This option is most efficient because
   ///   any layout map can be casted to a fully dynamic one.

mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp

bool isWritable(Operation *op, Value value,
     // inplace from the perspective of ops nested under:
     //   1. Either the matching iter operand is not bufferized inplace and an
     //      alloc + optional copy makes the bbArg itself inplaceable.
     //   2. Or the matching iter operand is bufferized inplace and bbArg just
     //      bufferizes to that too.
     return true;
   }
 
+  LogicalResult resolveConflicts(Operation *op, RewriterBase &rewriter,
+                                 const AnalysisState &state) const {
+    auto bufferizableOp = cast<BufferizableOpInterface>(op);
+    if (failed(bufferizableOp.resolveTensorOpOperandConflicts(rewriter, state)))
+      return failure();
+
+    if (!state.getOptions().enforceAliasingInvariants)
+      return success();
+
+    // According to the `getAliasing...` implementations, a bufferized OpResult
+    // may alias only with the corresponding bufferized init_arg and with no
+    // other buffers. I.e., the i-th OpResult may alias with the i-th init_arg;
+    // but not with any other OpOperand. If a corresponding OpResult/init_arg
+    // pair bufferizes to equivalent buffers, this aliasing requirement is
+    // satisfied. Otherwise, we cannot be sure and must yield a new buffer copy.
+    // (New buffer copies do not alias with any buffer.)
+    auto forOp = cast<scf::ForOp>(op);
+    auto yieldOp =
+        cast<scf::YieldOp>(forOp.getLoopBody().front().getTerminator());
+    OpBuilder::InsertionGuard g(rewriter);
+    rewriter.setInsertionPoint(yieldOp);
+
+    // Indices of all iter_args that have tensor type. These are the ones that
+    // are bufferized.
+    DenseSet<int64_t> indices = getTensorIndices(forOp.getInitArgs());
+    // For every yielded value, is the value equivalent to its corresponding
+    // bbArg?
+    DenseSet<int64_t> equivalentYields = getEquivalentBuffers(
+        forOp.getRegionIterArgs(), yieldOp.getResults(), state);
+    SmallVector<Value> yieldValues;
+    for (int64_t idx = 0;
+         idx < static_cast<int64_t>(yieldOp.getResults().size()); ++idx) {
+      Value value = yieldOp.getResults()[idx];
+      if (!indices.contains(idx) || equivalentYields.contains(idx)) {
+        yieldValues.push_back(value);
+        continue;
+      }
+      Value alloc = rewriter.create<bufferization::AllocTensorOp>(
+          yieldOp.getLoc(), value.getType().cast<RankedTensorType>(),
+          /*dynamicSizes=*/ValueRange(), value, /*escape=*/true);
+      yieldValues.push_back(alloc);
+    }
+
+    rewriter.updateRootInPlace(
+        yieldOp, [&]() { yieldOp.getResultsMutable().assign(yieldValues); });
+    return success();
+  }
+
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           BufferizationState &state) const {
     auto forOp = cast<scf::ForOp>(op);
     auto oldYieldOp =
         cast<scf::YieldOp>(forOp.getLoopBody().front().getTerminator());
     Block *oldLoopBody = &forOp.getLoopBody().front();
 
     // Indices of all iter_args that have tensor type. These are the ones that
bool isWritable(Operation *op, Value value,
     // inplace from the perspective of ops nested under:
     //   1. Either the matching iter operand is not bufferized inplace and an
     //      alloc + optional copy makes the bbArg itself inplaceable.
     //   2. Or the matching iter operand is bufferized inplace and bbArg just
     //      bufferizes to that too.
     return true;
   }
 
+  LogicalResult resolveConflicts(Operation *op, RewriterBase &rewriter,
+                                 const AnalysisState &state) const {
+    auto bufferizableOp = cast<BufferizableOpInterface>(op);
+    if (failed(bufferizableOp.resolveTensorOpOperandConflicts(rewriter, state)))
+      return failure();
+
+    if (!state.getOptions().enforceAliasingInvariants)
+      return success();
+
+    // According to the `getAliasing...` implementations, a bufferized OpResult
+    // may alias only with the corresponding bufferized init_arg and with no
+    // other buffers. I.e., the i-th OpResult may alias with the i-th init_arg;
+    // but not with any other OpOperand. If a corresponding OpResult/init_arg
+    // pair bufferizes to equivalent buffers, this aliasing requirement is
+    // satisfied. Otherwise, we cannot be sure and must yield a new buffer copy.
+    // (New buffer copies do not alias with any buffer.)
+    OpBuilder::InsertionGuard g(rewriter);
+    auto whileOp = cast<scf::WhileOp>(op);
+    auto conditionOp = whileOp.getConditionOp();
+    auto yieldOp = whileOp.getYieldOp();
+
+    // Indices of all bbArgs that have tensor type. These are the ones that
+    // are bufferized. The "before" and "after" regions may have different args.
+    DenseSet<int64_t> indicesBefore = getTensorIndices(whileOp.getInits());
+    DenseSet<int64_t> indicesAfter =
+        getTensorIndices(whileOp.getAfterArguments());
+
+    // For every yielded value, is the value equivalent to its corresponding
+    // bbArg?
+    DenseSet<int64_t> equivalentYieldsBefore = getEquivalentBuffers(
+        whileOp.getBeforeArguments(), conditionOp.getArgs(), state);
+    DenseSet<int64_t> equivalentYieldsAfter = getEquivalentBuffers(
+        whileOp.getAfterArguments(), whileOp.getYieldOp().getResults(), state);
+
+    // Update "before" region.
+    rewriter.setInsertionPoint(conditionOp);
+    SmallVector<Value> beforeYieldValues;
+    for (int64_t idx = 0;
+         idx < static_cast<int64_t>(conditionOp.getArgs().size()); ++idx) {
+      Value value = conditionOp.getArgs()[idx];
+      if (!indicesBefore.contains(idx) ||
+          equivalentYieldsBefore.contains(idx)) {
+        beforeYieldValues.push_back(value);
+        continue;
+      }
+      Value alloc = rewriter.create<bufferization::AllocTensorOp>(
+          conditionOp.getLoc(), value.getType().cast<RankedTensorType>(),
+          /*dynamicSizes=*/ValueRange(), value, /*escape=*/true);
+      beforeYieldValues.push_back(alloc);
+    }
+    rewriter.updateRootInPlace(conditionOp, [&]() {
+      conditionOp.getArgsMutable().assign(beforeYieldValues);
+    });
+
+    // Update "after" region.
+    rewriter.setInsertionPoint(yieldOp);
+    SmallVector<Value> afterYieldValues;
+    for (int64_t idx = 0;
+         idx < static_cast<int64_t>(yieldOp.getResults().size()); ++idx) {
+      Value value = yieldOp.getResults()[idx];
+      if (!indicesAfter.contains(idx) || equivalentYieldsAfter.contains(idx)) {
+        afterYieldValues.push_back(value);
+        continue;
+      }
+      Value alloc = rewriter.create<bufferization::AllocTensorOp>(
+          yieldOp.getLoc(), value.getType().cast<RankedTensorType>(),
+          /*dynamicSizes=*/ValueRange(), value, /*escape=*/true);
+      afterYieldValues.push_back(alloc);
+    }
+    rewriter.updateRootInPlace(yieldOp, [&]() {
+      yieldOp.getResultsMutable().assign(afterYieldValues);
+    });
+
+    return success();
+  }
+
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           BufferizationState &state) const {
     auto whileOp = cast<scf::WhileOp>(op);
 
     assert(whileOp.getBefore().getBlocks().size() == 1 &&
            "regions with multiple blocks not supported");
     Block *beforeBody = &whileOp.getBefore().front();
     assert(whileOp.getAfter().getBlocks().size() == 1 &&
bool isMemoryWrite(Operation *op, OpResult opResult,
     return true;
   }
 
   BufferRelation bufferRelation(Operation *op, OpResult opResult,
                                 const AnalysisState &state) const {
     return BufferRelation::Equivalent;
   }
 
+  LogicalResult resolveConflicts(Operation *op, RewriterBase &rewriter,
+                                 const AnalysisState &state) const {
+    auto bufferizableOp = cast<BufferizableOpInterface>(op);
+    if (failed(bufferizableOp.resolveTensorOpOperandConflicts(rewriter, state)))
+      return failure();
+
+    OpBuilder::InsertionGuard g(rewriter);
+    auto foreachThreadOp = cast<ForeachThreadOp>(op);
+    for (OpResult opResult : foreachThreadOp->getOpResults()) {
+      SmallVector<OpOperand *> destOperands =
+          state.getAliasingOpOperand(opResult);
+      assert(destOperands.size() == 1 &&
+             "expected exactly one aliasing OpOperand");
+      assert(isa<ParallelInsertSliceOp>(destOperands.front()->getOwner()) &&
+             "expected ParallelInsertSliceOp");
+
+      // Nothing to do if there is no conflict.
+      if (state.isInPlace(*destOperands.front()))
+        continue;
+
+      // Create AllocTensorOp.
+      bool isYielded = state.isTensorYielded(opResult);
+      auto resultType = opResult.getType().cast<RankedTensorType>();
+      Value alloc = rewriter.create<bufferization::AllocTensorOp>(
+          op->getLoc(), resultType, /*dynamicDims=*/ValueRange(),
+          /*copy=*/destOperands.front()->get(),
+          /*escape=*/isYielded);
+
+      // Update terminator operand.
+      rewriter.updateRootInPlace(destOperands.front()->getOwner(),
+                                 [&]() { destOperands.front()->set(alloc); });
+    }
+
+    return success();
+  }
+
   LogicalResult bufferize(Operation *op, RewriterBase &b,
                           BufferizationState &state) const {
     OpBuilder::InsertionGuard g(b);
     auto foreachThreadOp = cast<ForeachThreadOp>(op);
 
     // Gather new results of the ForeachThreadOp.
     SmallVector<Value> newResults;
     for (OpResult opResult : foreachThreadOp->getOpResults()) {
bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
     return &opOperand == &op->getOpOperand(1) /*dest*/;
   }
 
   BufferRelation bufferRelation(Operation *op, OpResult opResult,
                                 const AnalysisState &state) const {
     return BufferRelation::Equivalent;
   }
 
+  LogicalResult resolveConflicts(Operation *op, RewriterBase &rewriter,
+                                 const AnalysisState &state) const {
+    return success();
+  }
+
   LogicalResult bufferize(Operation *op, RewriterBase &b,
                           BufferizationState &state) const {
     // Will be bufferized as part of ForeachThreadOp.
     return failure();
   }
 
   // TODO: This is copied from TensorInterfaceImpl.cpp. Find a way to share
   // the code.

mlir/test/Dialect/SCF/one-shot-bufferize-tensor-copy-insertion.mlir

+// RUN: mlir-opt %s -tensor-copy-insertion="allow-return-allocs" -allow-unregistered-dialect -split-input-file | FileCheck %s
+// RUN: mlir-opt %s -tensor-copy-insertion="bufferize-function-boundaries allow-return-allocs" -split-input-file | FileCheck %s --check-prefix=CHECK-FUNC
+
+// CHECK-LABEL: func @scf_for(
+//  CHECK-SAME:     %[[A:.*]]: tensor<?xf32>, %[[B:.*]]: tensor<?xf32>
+func.func @scf_for(%A : tensor<?xf32>, %B : tensor<?xf32>,
+                   %lb : index, %ub : index, %step : index)
+  -> (tensor<?xf32>, tensor<?xf32>)
+{
+  // CHECK: %[[A_copy:.*]] = bufferization.alloc_tensor() copy(%[[A]]) {escape = false} : tensor<?xf32>
+  // CHECK: %[[B_copy:.*]] = bufferization.alloc_tensor() copy(%[[B]]) {escape = false} : tensor<?xf32>
+  // CHECK:   %[[for:.*]]:2 = scf.for {{.*}} iter_args(%[[iter1:.*]] = %[[A_copy]], %[[iter2:.*]] = %[[B_copy]])
+  %r0:2 = scf.for %i = %lb to %ub step %step iter_args(%tA = %A, %tB = %B)
+      -> (tensor<?xf32>, tensor<?xf32>)
+  {
+    // CHECK: scf.yield %[[iter1]], %[[iter2]]
+    scf.yield %tA, %tB : tensor<?xf32>, tensor<?xf32>
+  }
+
+  return %r0#0, %r0#1 : tensor<?xf32>, tensor<?xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @scf_for_swapping_yields(
+//  CHECK-SAME:     %[[A:.*]]: tensor<?xf32>, %[[B:.*]]: tensor<?xf32>
+func.func @scf_for_swapping_yields(%A : tensor<?xf32>, %B : tensor<?xf32>,
+                                   %lb : index, %ub : index, %step : index)
+  -> (tensor<?xf32>, tensor<?xf32>)
+{
+  // CHECK: %[[A_copy:.*]] = bufferization.alloc_tensor() copy(%[[A]]) {escape = false} : tensor<?xf32>
+  // CHECK: %[[B_copy:.*]] = bufferization.alloc_tensor() copy(%[[B]]) {escape = false} : tensor<?xf32>
+  // CHECK:   %[[for:.*]]:2 = scf.for {{.*}} iter_args(%[[iter1:.*]] = %[[A_copy]], %[[iter2:.*]] = %[[B_copy]])
+  %r0:2 = scf.for %i = %lb to %ub step %step iter_args(%tA = %A, %tB = %B)
+      -> (tensor<?xf32>, tensor<?xf32>)
+  {
+    // Yield tensors in different order.
+    // CHECK-DAG: %[[yield1:.*]] = bufferization.alloc_tensor() copy(%[[iter2]]) {escape = true} : tensor<?xf32>
+    // CHECK-DAG: %[[yield2:.*]] = bufferization.alloc_tensor() copy(%[[iter1]]) {escape = true} : tensor<?xf32>
+    // CHECK: scf.yield %[[yield1]], %[[yield2]]
+    scf.yield %tB, %tA : tensor<?xf32>, tensor<?xf32>
+  }
+
+  return %r0#0, %r0#1 : tensor<?xf32>, tensor<?xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @scf_while(
+//  CHECK-SAME:     %[[A:.*]]: tensor<5xi1>, %[[B:.*]]: tensor<5xi1>
+func.func @scf_while(%A: tensor<5xi1>, %B: tensor<5xi1>, %idx: index)
+  -> (tensor<5xi1>, tensor<5xi1>)
+{
+  // CHECK: %[[A_copy:.*]] = bufferization.alloc_tensor() copy(%[[A]]) {escape = false} : tensor<5xi1>
+  // CHECK: %[[B_copy:.*]] = bufferization.alloc_tensor() copy(%[[B]]) {escape = false} : tensor<5xi1>
+  // CHECK: %[[loop:.*]]:2 = scf.while (%[[w0:.*]] = %[[A_copy]], %[[w1:.*]] = %[[B_copy]]) {{.*}} {
+  %r0, %r1 = scf.while (%w0 = %A, %w1 = %B)
+      : (tensor<5xi1>, tensor<5xi1>) -> (tensor<5xi1>, tensor<5xi1>) {
+    // CHECK: %[[condition:.*]] = tensor.extract %[[w0]]
+    %condition = tensor.extract %w0[%idx] : tensor<5xi1>
+    // Yield tensors in different order.
+    // CHECK: scf.condition(%[[condition]]) %[[w0]], %[[w1]]
+    scf.condition(%condition) %w0, %w1 : tensor<5xi1>, tensor<5xi1>
+  } do {
+  ^bb0(%b0: tensor<5xi1>, %b1: tensor<5xi1>):
+    // CHECK: } do {
+    // CHECK: ^bb0(%[[b0:.*]]: tensor<5xi1>, %[[b1:.*]]: tensor<5xi1>):
+    // CHECK: scf.yield %[[b0]], %[[b1]]
+    // CHECK: }
+    scf.yield %b0, %b1 : tensor<5xi1>, tensor<5xi1>
+  }
+
+  return %r0, %r1 : tensor<5xi1>, tensor<5xi1>
+}
+
+// -----
+
+// CHECK-LABEL: func @scf_while_non_equiv_condition_and_body(
+//  CHECK-SAME:     %[[A:.*]]: tensor<5xi1>, %[[B:.*]]: tensor<5xi1>
+func.func @scf_while_non_equiv_condition_and_body(%A: tensor<5xi1>,
+                                                  %B: tensor<5xi1>,
+                                                  %idx: index)
+  -> (tensor<5xi1>, tensor<5xi1>)
+{
+  // CHECK: %[[A_copy:.*]] = bufferization.alloc_tensor() copy(%[[A]]) {escape = false} : tensor<5xi1>
+  // CHECK: %[[B_copy:.*]] = bufferization.alloc_tensor() copy(%[[B]]) {escape = false} : tensor<5xi1>
+  // CHECK: %[[loop:.*]]:2 = scf.while (%[[w0:.*]] = %[[A_copy]], %[[w1:.*]] = %[[B_copy]]) {{.*}} {
+  %r0, %r1 = scf.while (%w0 = %A, %w1 = %B)
+      : (tensor<5xi1>, tensor<5xi1>) -> (tensor<5xi1>, tensor<5xi1>) {
+    // CHECK: %[[condition:.*]] = tensor.extract %[[w0]]
+    %condition = tensor.extract %w0[%idx] : tensor<5xi1>
+    // Yield tensors in different order.
+    // CHECK-DAG: %[[yield0:.*]] = bufferization.alloc_tensor() copy(%[[w1]]) {escape = true} : tensor<5xi1>
+    // CHECK-DAG: %[[yield1:.*]] = bufferization.alloc_tensor() copy(%[[w0]]) {escape = true} : tensor<5xi1>
+    // CHECK: scf.condition(%[[condition]]) %[[yield0]], %[[yield1]]
+    scf.condition(%condition) %w1, %w0 : tensor<5xi1>, tensor<5xi1>
+  } do {
+  ^bb0(%b0: tensor<5xi1>, %b1: tensor<5xi1>):
+    // CHECK: } do {
+    // CHECK: ^bb0(%[[b0:.*]]: tensor<5xi1>, %[[b1:.*]]: tensor<5xi1>):
+    // CHECK-DAG: %[[yield2:.*]] = bufferization.alloc_tensor() copy(%[[b1]]) {escape = true} : tensor<5xi1>
+    // CHECK-DAG: %[[yield3:.*]] = bufferization.alloc_tensor() copy(%[[b0]]) {escape = true} : tensor<5xi1>
+    // CHECK: scf.yield %[[yield2]], %[[yield3]]
+    // CHECK: }
+    scf.yield %b1, %b0 : tensor<5xi1>, tensor<5xi1>
+  }
+
+  return %r0, %r1 : tensor<5xi1>, tensor<5xi1>
+}
+
+// -----
+
+// CHECK-LABEL: func @scf_foreach_thread_out_of_place(
+//  CHECK-SAME:     %[[arg0:.*]]: tensor<100xf32>, %[[arg1:.*]]: tensor<100xf32>
+// CHECK-FUNC-LABEL: func @scf_foreach_thread_out_of_place(
+func.func @scf_foreach_thread_out_of_place(%in: tensor<100xf32>,
+                                           %out: tensor<100xf32>) {
+  %c1 = arith.constant 1 : index
+  %num_threads = arith.constant 100 : index
+
+  // CHECK-FUNC-NOT: alloc_tensor
+  // CHECK: %[[alloc:.*]] = bufferization.alloc_tensor() copy(%[[arg1]]) {escape = false} : tensor<100xf32>
+  // CHECK: scf.foreach_thread
+  %result = scf.foreach_thread (%thread_idx) in (%num_threads) -> tensor<100xf32> {
+      // CHECK: tensor.extract_slice
+      // CHECK: scf.foreach_thread.perform_concurrently
+      // CHECK: scf.foreach_thread.parallel_insert_slice %{{.*}} into %[[alloc]]
+      %1 = tensor.extract_slice %in[%thread_idx][1][1] : tensor<100xf32> to tensor<1xf32>
+      scf.foreach_thread.perform_concurrently {
+        scf.foreach_thread.parallel_insert_slice %1 into %out[%thread_idx][1][1] :
+          tensor<1xf32> into tensor<100xf32>
+      }
+  }
+  return
+}