[mlir][bufferize] Improve buffer writability analysis

[springerm]

Summary

Find writability conflicts (writes to buffers that are not allowed to be written to) by checking SSA use-def chains. This is better than the current writability analysis, which is too conservative and finds false positives.

[bkramer]

thanks

mlir/include/mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h

public:
   /// Find all tensors that are yielded/returned from a block and store them in
   /// `yieldedTensors`. Also include all aliasing tensors in the same block.
   void gatherYieldedTensors(Operation *op);
 
   /// Return true if the buffer of the given tensor value is written to. Must
   /// not be called for values inside not yet analyzed functions.
   bool isValueWritten(Value value) const;
 
+  /// Return true if the buffer of the given tensor value is writable.
+  bool isWritable(Value value) const;
+
 private:
   /// `aliasInfo` keeps track of aliasing and equivalent values. Only internal
   /// functions and `runOneShotBufferize` may access this object.
   BufferizationAliasInfo aliasInfo;
 
   /// A set of all tensors (and maybe aliasing tensors) that yielded from a
   /// block.
   DenseSet<Value> yieldedTensors;

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

bool OneShotAnalysisState::isValueWritten(Value value) const {
   aliasInfo.applyOnAliases(value, [&](Value val) {
     for (OpOperand &use : val.getUses())
       if (isInPlace(use) && bufferizesToMemoryWrite(use))
         isWritten = true;
   });
   return isWritten;
 }
 
+bool OneShotAnalysisState::isWritable(Value value) const {
+  // TODO: Out-of-place bufferized value could be considered writable.
+  if (auto bufferizableOp = getOptions().dynCastBufferizableOp(value))
+    return bufferizableOp.isWritable(value, *this);
+
+  // Query BufferizableOpInterface to see if the BlockArgument is writable.
+  if (auto bbArg = value.dyn_cast<BlockArgument>())
+    if (auto bufferizableOp =
+            getOptions().dynCastBufferizableOp(bbArg.getOwner()->getParentOp()))
+      return bufferizableOp.isWritable(bbArg, *this);
+
+  // Not a bufferizable op: The conservative answer is "not writable".
+  return false;
+}
+
 //===----------------------------------------------------------------------===//
 // Bufferization-specific alias analysis.
 //===----------------------------------------------------------------------===//
 
 /// Return true if opOperand has been decided to bufferize in-place.
 static bool isInplaceMemoryWrite(OpOperand &opOperand,
                                  const BufferizationAliasInfo &aliasInfo,
-                                 AnalysisState &state) {
+                                 const AnalysisState &state) {
   // OpOperands that do not bufferize to a memory write do not write in-place.
   if (!state.bufferizesToMemoryWrite(opOperand))
     return false;
   // Check current bufferization decisions.
   return aliasInfo.isInPlace(opOperand);
 }
 
-/// Return true if, under current bufferization decisions, the buffer of `value`
-/// is not writable.
-static bool aliasesNonWritableBuffer(Value value,
-                                     const BufferizationAliasInfo &aliasInfo,
-                                     AnalysisState &state) {
-  bool foundNonWritableBuffer = false;
-  aliasInfo.applyOnAliases(value, [&](Value v) {
-    // Query BufferizableOpInterface to see if the value is writable.
-    // TODO: Out-of-place bufferized value could be considered writable.
-    if (auto bufferizableOp = state.getOptions().dynCastBufferizableOp(v))
-      if (bufferizableOp && bufferizableOp.isWritable(v, state))
-        return;
-
-    // Query BufferizableOpInterface to see if the BlockArgument is writable.
-    if (auto bbArg = v.dyn_cast<BlockArgument>())
-      if (auto bufferizableOp = state.getOptions().dynCastBufferizableOp(
-              bbArg.getOwner()->getParentOp()))
-        if (bufferizableOp.isWritable(bbArg, state))
-          return;
-
-    foundNonWritableBuffer = true;
-  });
-
-  return foundNonWritableBuffer;
-}
-
-/// Return true if the buffer to which `operand` would bufferize is equivalent
-/// to some buffer write.
-static bool aliasesInPlaceWrite(Value value,
-                                const BufferizationAliasInfo &aliasInfo,
-                                AnalysisState &state) {
-  bool foundInplaceWrite = false;
-  aliasInfo.applyOnAliases(value, [&](Value v) {
-    for (auto &use : v.getUses()) {
-      if (isInplaceMemoryWrite(use, aliasInfo, state)) {
-        foundInplaceWrite = true;
-        return;
-      }
-    }
-  });
-  return foundInplaceWrite;
-}
-
 /// Return true if `a` happens before `b`, i.e., `a` or one of its ancestors
 /// properly dominates `b` and `b` is not inside `a`.
 static bool happensBefore(Operation *a, Operation *b,
                           const DominanceInfo &domInfo) {
   do {
     // TODO: Instead of isProperAncestor + properlyDominates, we should use
     // properlyDominatesImpl(a, b, /*enclosingOpOk=*/false)
     if (a->isProperAncestor(b))
for (OpOperand *uConflictingWrite : usesWrite) {
         return true;
       }
     }
   }
 
   return false;
 }
 
+// Helper function to iterate on aliases of `root` and capture the writes.
+static void getAliasingInplaceWrites(DenseSet<OpOperand *> &res, Value root,
+                                     const BufferizationAliasInfo &aliasInfo,
+                                     const AnalysisState &state) {
+  aliasInfo.applyOnAliases(root, [&](Value alias) {
+    for (auto &use : alias.getUses())
+      // Inplace write to a value that aliases root.
+      if (isInplaceMemoryWrite(use, aliasInfo, state))
+        res.insert(&use);
+  });
+}
+
+// Helper function to iterate on aliases of `root` and capture the reads.
+static void getAliasingReads(DenseSet<OpOperand *> &res, Value root,
+                             const BufferizationAliasInfo &aliasInfo,
+                             const AnalysisState &state) {
+  aliasInfo.applyOnAliases(root, [&](Value alias) {
+    for (auto &use : alias.getUses())
+      // Read to a value that aliases root.
+      if (state.bufferizesToMemoryRead(use))
+        res.insert(&use);
+  });
+}
+
 /// Return true if bufferizing `operand` inplace would create a conflict. A read
 /// R and a write W of the same alias set is a conflict if inplace bufferization
 /// of W changes the value read by R to a value different from the one that
 /// would be expected by tracing back R's origin through SSA use-def chains.
 /// A conflict can only be introduced by a new alias and/or an inplace
 /// bufferization decision.
 ///
 /// Example:

 ///
 /// Note: If `checkConsistencyOnly`, this function may be called with a null
 /// OpResult. In that case, only the consistency of bufferization decisions
 /// involving aliases of the given OpOperand are checked.
 static bool wouldCreateReadAfterWriteInterference(
     OpOperand &operand, const DominanceInfo &domInfo, AnalysisState &state,
     const BufferizationAliasInfo &aliasInfo,
     bool checkConsistencyOnly = false) {
-  // Helper function to iterate on aliases of `root` and capture the reads.
-  auto getAliasingReads = [&](DenseSet<OpOperand *> &res, Value root) {
-    aliasInfo.applyOnAliases(root, [&](Value alias) {
-      for (auto &use : alias.getUses())
-        // Read to a value that aliases root.
-        if (state.bufferizesToMemoryRead(use))
-          res.insert(&use);
-    });
-  };
-
-  // Helper function to iterate on aliases of `root` and capture the writes.
-  auto getAliasingInplaceWrites = [&](DenseSet<OpOperand *> &res, Value root) {
-    aliasInfo.applyOnAliases(root, [&](Value alias) {
-      for (auto &use : alias.getUses())
-        // Inplace write to a value that aliases root.
-        if (isInplaceMemoryWrite(use, aliasInfo, state))
-          res.insert(&use);
-    });
-  };
-
   // Collect reads and writes of all aliases of OpOperand and OpResult.
   DenseSet<OpOperand *> usesRead, usesWrite;
-  getAliasingReads(usesRead, operand.get());
-  getAliasingInplaceWrites(usesWrite, operand.get());
+  getAliasingReads(usesRead, operand.get(), aliasInfo, state);
+  getAliasingInplaceWrites(usesWrite, operand.get(), aliasInfo, state);
   for (OpResult result : state.getAliasingOpResult(operand)) {
-    getAliasingReads(usesRead, result);
-    getAliasingInplaceWrites(usesWrite, result);
+    getAliasingReads(usesRead, result, aliasInfo, state);
+    getAliasingInplaceWrites(usesWrite, result, aliasInfo, state);
   }
   if (!checkConsistencyOnly && state.bufferizesToMemoryWrite(operand))
     usesWrite.insert(&operand);
 
   return hasReadAfterWriteInterference(usesRead, usesWrite, domInfo, state,
                                        aliasInfo);
 }
 
-/// Return true if bufferizing `opOperand` inplace would create a write to a
-/// non-writable buffer.
+/// Check the reverse SSA use-def chain (following aliasing OpOperands) for
+/// non-writable tensor values. Stop searching when an out-of-place bufferized
+/// OpOperand was found (or when the OpOperand was not bufferized yet).
+/// `currentOpOperand` is assumed to be in-place, even if that decision was not
+/// materialized in `aliasInfo` yet.
 static bool
-wouldCreateWriteToNonWritableBuffer(OpOperand &opOperand,
+hasPrecedingAliasingNonWritableTensor(Value value, OpOperand *currentOpOperand,
                                       const BufferizationAliasInfo &aliasInfo,
-                                    AnalysisState &state) {
-  // Certain buffers are not writeable:
-  //   1. A function bbArg that is not inplaceable or
-  //   2. A constant op.
-  bool nonWritable =
-      aliasesNonWritableBuffer(opOperand.get(), aliasInfo, state);
-  if (!nonWritable)
-    return false;
+                                      const OneShotAnalysisState &state) {
+  SmallVector<Value> worklist;
+  worklist.push_back(value);
+  while (!worklist.empty()) {
+    Value nextVal = worklist.pop_back_val();
+    if (!state.isWritable(nextVal))
+      return true;
+
+    // If `nextVal` is not a BlockArgument: End of use-def chain reached.
+    auto opResult = nextVal.dyn_cast<OpResult>();
+    if (!opResult)
+      continue;
+
+    // Follow reverse SSA use-def chain.
+    SmallVector<OpOperand *> aliasingOpOperands =
+        state.getAliasingOpOperand(opResult);
+    for (OpOperand *opOperand : aliasingOpOperands)
+      if (aliasInfo.isInPlace(*opOperand) || currentOpOperand == opOperand)
+        worklist.push_back(opOperand->get());
+  }
+}
 
-  // This is a problem only if the buffer is written to via some alias.
-  bool hasWrite = aliasesInPlaceWrite(opOperand.get(), aliasInfo, state) ||
-                  state.bufferizesToMemoryWrite(opOperand);
+/// Return true if bufferizing `operand` inplace would create a write to a
+/// non-writable buffer.
+static bool wouldCreateWriteToNonWritableBuffer(
+    OpOperand &operand, const BufferizationAliasInfo &aliasInfo,
+    OneShotAnalysisState &state, bool checkConsistencyOnly = false) {
+  // Collect writes of all aliases of OpOperand and OpResult.
+  DenseSet<OpOperand *> usesWrite;
+  getAliasingInplaceWrites(usesWrite, operand.get(), aliasInfo, state);
+  for (OpResult result : state.getAliasingOpResult(operand)) {
+    getAliasingInplaceWrites(usesWrite, result, aliasInfo, state);
+  }
+  if (!checkConsistencyOnly && state.bufferizesToMemoryWrite(operand))
+    usesWrite.insert(&operand);
 
-  for (OpResult opResult : state.getAliasingOpResult(opOperand))
-    hasWrite |= aliasesInPlaceWrite(opResult, aliasInfo, state);
+  // Assuming that `operand` bufferizes in-place: For each write (to each
+  // alias), check if there is a non-writable tensor in the reverse SSA use-def
+  // chain.
+  for (OpOperand *uWrite : usesWrite)
+    if (hasPrecedingAliasingNonWritableTensor(uWrite->get(), &operand,
+                                              aliasInfo, state))
+      return true;
 
-  return hasWrite;
+  return false;
 }
 
 //===----------------------------------------------------------------------===//
 // Bufferization analyses.
 //===----------------------------------------------------------------------===//
 
 /// Determine if `operand` can be bufferized in-place.
 static LogicalResult bufferizableInPlaceAnalysisImpl(
-    OpOperand &operand, BufferizationAliasInfo &aliasInfo, AnalysisState &state,
-    const DominanceInfo &domInfo) {
+    OpOperand &operand, BufferizationAliasInfo &aliasInfo,
+    OneShotAnalysisState &state, const DominanceInfo &domInfo) {
   bool foundInterference =
       wouldCreateWriteToNonWritableBuffer(operand, aliasInfo, state) ||
       wouldCreateReadAfterWriteInterference(operand, domInfo, state, aliasInfo);
 
   if (foundInterference)
     aliasInfo.bufferizeOutOfPlace(operand);
   else
     aliasInfo.bufferizeInPlace(operand, state);

 /// When bufferized inplace, an ExtractSliceOp does not by itself create any
 /// read or write from memory. Instead, it has the effect of merging the alias
 /// sets of the source and the result buffers.
 ///
 /// An analysis is required to ensure inplace bufferization would not result in
 /// RaW dependence violations.
 static LogicalResult inPlaceAnalysis(SmallVector<Operation *> &ops,
                                      BufferizationAliasInfo &aliasInfo,
-                                     AnalysisState &state,
+                                     OneShotAnalysisState &state,
                                      const DominanceInfo &domInfo,
                                      unsigned analysisFuzzerSeed = 0) {
   if (analysisFuzzerSeed) {
     // This is a fuzzer. For testing purposes only. Randomize the order in which
     // operations are analyzed. The bufferization quality is likely worse, but
     // we want to make sure that no assertions are triggered anywhere.
     std::mt19937 g(analysisFuzzerSeed);
     llvm::shuffle(ops.begin(), ops.end(), g);
static bool hasTensorSemantics(Operation *op) {
   bool hasTensorResult = any_of(op->getResultTypes(), isaTensor);
   bool hasTensorOperand = any_of(op->getOperandTypes(), isaTensor);
   return hasTensorResult || hasTensorOperand;
 }
 
 /// Analyze all ops that are contained in `op`.
 static LogicalResult inPlaceAnalysis(Operation *op,
                                      BufferizationAliasInfo &aliasInfo,
-                                     AnalysisState &state,
+                                     OneShotAnalysisState &state,
                                      const DominanceInfo &domInfo,
                                      unsigned analysisFuzzerSeed = 0) {
   // Collect ops so we can build our own reverse traversal.
   SmallVector<Operation *> ops;
   op->walk([&](Operation *op) {
     // No tensors => no buffers.
     if (!hasTensorSemantics(op))
       return;

mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-allow-return-allocs.mlir

// CHECK: } else {
     scf.yield %t : tensor<?xf32>
   }
   // CHECK: }
   // CHECK: %[[r_tensor:.*]] = bufferization.to_tensor %[[r]]
   // CHECK: memref.dealloc %[[r]]
   // CHECK: return %[[r_tensor]]
   return %r : tensor<?xf32>
 }
+
+// -----
+
+// CHECK-LABEL: func @write_to_alloc_tensor_or_readonly_tensor(
+//  CHECK-SAME:     %[[arg0:.*]]: tensor<i32>
+func.func @write_to_alloc_tensor_or_readonly_tensor(%arg0: tensor<i32>,
+                                                    %cond: i1, %val: i32)
+  -> tensor<i32>
+{
+  // CHECK: %[[r:.*]] = scf.if {{.*}} {
+  // CHECK:   %[[arg0_m:.*]] = bufferization.to_memref %[[arg0]]
+  // CHECK:   %[[clone:.*]] = bufferization.clone %[[arg0_m]]
+  // CHECK:   scf.yield %[[clone]]
+  // CHECK: } else {
+  // CHECK:   %[[alloc:.*]] = memref.alloc
+  // CHECK:   memref.store %{{.*}}, %[[alloc]]
+  // CHECK:   %[[casted:.*]] = memref.cast %[[alloc]]
+  // CHECK:   scf.yield %[[casted]]
+  // CHECK: }
+  // CHECK: %[[r_t:.*]] = bufferization.to_tensor %[[r]]
+  // CHECK: memref.dealloc %[[r]]
+  // CHECK: return %[[r_t]]
+  %3 = scf.if %cond -> (tensor<i32>) {
+    scf.yield %arg0 : tensor<i32>
+  } else {
+    %7 = bufferization.alloc_tensor() : tensor<i32>
+    %8 = tensor.insert %val into %7[] : tensor<i32>
+    scf.yield %8 : tensor<i32>
+  }
+  return %3 : tensor<i32>
+}