[mlir][linalg][bufferize] Rewrite "write into non-writable memory" detection

[springerm]

Summary

Depends On D111378

[springerm]

The purpose of this revision is to make "write into non-writable memory" conflict detection easier to understand.

In particular, I was thinking for a while that the existing code is missing a case.

%r = tensor.extract_slice %t

where:
* %t aliases in-place write
* %r aliases non-writable buffer

This case is not covered by the existing code. The new one does cover it. After some more thinking, I now believe that it is not possible to construct such an example. However, this was not immediately obvious to me. So for the sake of simplicity/understandability, my suggestion would be to go with something similar to this revision.

The main idea here is that there is a conflict in the case of inplace bufferization if:

1. Someone writes to (an alias of) opOperand, opResult or to-be-bufferized op writes itself.
2. And, opOperand or opResult aliases a non-writable buffer.

[nicolasvasilache]

Yes makes sense to isolate the logic and give it a good name.

Could you promote parts of your comment to the commit message ?

Also, if one of the cases here is not expected to be possible to write, how about putting a defensive assertion + TODO on the corresponding branch so that if it shows up in the future we don't silently miss it?

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

public:
   ///      be visible from R).
   ///   2. one cannot find an intermediate clobbering write `C` to W, such that
   ///      C interleaved between W and R (i.e. W -> C -> R where -> denotes
   ///      dominance).
   bool wouldCreateReadAfterWriteInterference(
       Operation *opToBufferize, DenseSet<OpOperand *> &usesRead,
       DenseSet<OpOperand *> &usesWrite, const DominanceInfo &domInfo) const;
 
+  /// Return true if bufferizing `opResult` inplace would create a write to a
+  /// non-writable buffer.
+  bool wouldCreateWriteToNonWritableBuffer(OpResult opResult) const;
+
   /// Assume that result bufferizes in-place with one of the operation's
   /// operands. Return true if it is possible to find an inplace write W (resp.
   /// a read R) among the uses of `aliasInfo[result]`, and a read R (resp. an
   /// inplace write W) among the uses of
   /// `aliasInfo[getAliasingOpOperand(result)]`, such that W and R interfere.
   /// Interference detection is needed to determine which cases may bufferize
   /// inplace without interferences. Such cases comprise:
   ///

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

bool BufferizationAliasInfo::wouldCreateReadAfterWriteInterference(
   usesRead.clear();
   usesWrite.clear();
   getAliasingReads(usesRead, opResult);
   getAliasingInplaceWrites(usesWrite, opOperand);
   return wouldCreateReadAfterWriteInterference(opToBufferize, usesRead,
                                                usesWrite, domInfo);
 }
 
+/// Return true if bufferizing `opResult` inplace would create a write to a
+/// non-writable buffer.
+bool BufferizationAliasInfo::wouldCreateWriteToNonWritableBuffer(
+    OpResult opResult) const {
+  Optional<OpOperand *> maybeAliasingOperand = getAliasingOpOperand(opResult);
+  if (!maybeAliasingOperand || !*maybeAliasingOperand)
+    return false;
+
+  // Certain buffers are not writeable:
+  //   1. A function bbArg that is not inplaceable or
+  //   2. A constant op.
+  assert(!aliasesNonWritableBuffer(opResult) &&
+         "expected that opResult does not alias non-writable buffer");
+  bool nonWritable = aliasesNonWritableBuffer((*maybeAliasingOperand)->get());
+  if (!nonWritable)
+    return false;
+
+  // This is a problem only if the buffer is written to via some alias.
+  bool hasWrite = aliasesInPlaceWrite(opResult) ||
+                  aliasesInPlaceWrite((*maybeAliasingOperand)->get()) ||
+                  bufferizesToMemoryWrite(**maybeAliasingOperand);
+  if (!hasWrite)
+    return false;
+
+  LDBG("->the corresponding buffer is not writeable\n");
+  return true;
+}
+
 /// Return true if the source of a `insertSliceOp` bufferizes to an
 /// equivalent ExtractSliceOp that bufferizes inplace.
 bool BufferizationAliasInfo::isSourceEquivalentToAMatchingInplaceExtractSliceOp(
     InsertSliceOp insertSliceOp) const {
   LDBG("isSourceEquivalentToAMatchingInplaceExtractSliceOp: " << *insertSliceOp
                                                               << '\n');
   auto leaderIt = equivalentInfo.findLeader(insertSliceOp.source());
   for (auto mit = leaderIt, meit = equivalentInfo.member_end(); mit != meit;

 static LogicalResult
 bufferizableInPlaceAnalysis(ExtractSliceOp extractSliceOp,
                             BufferizationAliasInfo &aliasInfo,
                             const DominanceInfo &domInfo) {
   LDBG('\n');
   LDBG("Inplace analysis for extract_slice: "
        << printOperationInfo(extractSliceOp) << '\n');
 
-  // If `extractSliceOp` were to be bufferized inplace, it cannot end up
-  // aliasing a write into a non-writable buffer.
-  bool wouldCreateAliasingWriteToNonWritableBuffer =
-      aliasInfo.aliasesInPlaceWrite(extractSliceOp.result()) &&
-      aliasInfo.aliasesNonWritableBuffer(extractSliceOp.source());
-
-  if (wouldCreateAliasingWriteToNonWritableBuffer)
-    LDBG("->the corresponding buffer is not writable\n");
-  else
-    LDBG("->bufferizes to writable inplace buffer\n");
-
-  // In any of extractSliceOp.result's aliases, can we find 2 such that we hit
-  // an interfering write?
   OpResult r = extractSliceOp->getResult(0);
   OpOperand &s = extractSliceOp->getOpOperand(0);
   bool foundInterference =
-      wouldCreateAliasingWriteToNonWritableBuffer ||
+      /* If `extractSliceOp` were to be bufferized inplace, it cannot end up
+         aliasing a write into a non-writable buffer.*/
+      aliasInfo.wouldCreateWriteToNonWritableBuffer(r) ||
+      /* In any of extractSliceOp.result's aliases, can we find 2 such that we
+         hit an interfering write? */
       aliasInfo.wouldCreateReadAfterWriteInterference(r, domInfo);
   if (foundInterference)
     aliasInfo.bufferizeOutOfPlace(r);
   else
     aliasInfo.bufferizeInPlace(r, s);
 
   LDBG("Done inplace analysis for extract_slice\n");
 
bufferizableInPlaceAnalysis(OpOperand &operand,
 
   int64_t resultNumber = result.getResultNumber();
   (void)resultNumber;
   LDBG('\n');
   LDBG("Inplace analysis for <- #" << resultNumber << " -> #"
                                    << operand.getOperandNumber() << " in "
                                    << printValueInfo(result) << '\n');
 
-  // `result` must bufferize to a writable buffer to be a candidate.
-  // This means the operand  must not alias either:
-  //   1. a function bbArg that is not inplaceable or
-  //   2. a constant op.
-  // to be considered for inplace bufferization
-  bool wouldCreateAliasingWriteToNonWritableBuffer =
-      aliasInfo.aliasesNonWritableBuffer(operand.get());
-  if (wouldCreateAliasingWriteToNonWritableBuffer)
-    LDBG("->the corresponding buffer is not writable\n");
-  else
-    LDBG("->bufferizes to writable inplace buffer\n");
-
   bool foundInterference =
-      wouldCreateAliasingWriteToNonWritableBuffer ||
+      aliasInfo.wouldCreateWriteToNonWritableBuffer(result) ||
       aliasInfo.wouldCreateReadAfterWriteInterference(result, domInfo);
 
   if (foundInterference)
     aliasInfo.bufferizeOutOfPlace(result);
   else
     // TODO: Atm, all inplace bufferizations yield equivalent tensors. Support
     // more cases on a per-need basis.
     aliasInfo.bufferizeInPlace(result, operand);

mlir/test/Dialect/Linalg/comprehensive-module-bufferize-analysis.mlir

func @read_dependence_through_scf_and_call(
   // CHECK-SAME: {__inplace_results_attr__ = ["true"]}
   %B2 = linalg.fill(%f2, %I2) : f32, tensor<64xf32> -> tensor<64xf32>
 
   call @bar(%A2) : (tensor<64xf32>) -> ()
   call @bar(%B2) : (tensor<64xf32>) -> ()
   return
 }
 
+// -----
+
 //===----------------------------------------------------------------------===//
 // Transitive cases through extract_slice.
 //===----------------------------------------------------------------------===//
 
+// CHECK-LABEL: func @write_into_constant_via_alias
+func @write_into_constant_via_alias(%v : vector<5xi32>,
+                                    %s1 : index, %s2 : index,
+                                    %s3 : index) -> tensor<?xi32> {
+  %A = constant dense<[1, 2, 3, 4]> : tensor<4xi32>
+  //      CHECK: tensor.extract_slice
+  // CHECK-SAME: {__inplace_results_attr__ = ["false"]}
+  %b = tensor.extract_slice %A[%s1][%s2][1] : tensor<4xi32> to tensor<?xi32>
+  //      CHECK: vector.transfer_write
+  // CHECK-SAME: {__inplace_results_attr__ = ["true"]}
+  %r = vector.transfer_write %v, %b[%s3] : vector<5xi32>, tensor<?xi32>
+  return %r : tensor<?xi32>
+}
+
+// -----
+
 builtin.func @matmul_on_tensors(
     %arg0: tensor<518x518xf32> {linalg.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, linalg.inplaceable = false},
     %arg1: tensor<518x518xf32> {linalg.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, linalg.inplaceable = false},
     %arg2: tensor<256x256xf32> {linalg.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, linalg.inplaceable = true})
     -> tensor<256x256xf32>
 {
   %c0 = constant 0 : index
   %cst_0 = constant 0.000000e+00 : f32