[mlir][bufferization][NFC] Cache definitions of read tensors

[springerm]

Summary

This is to avoid unnecessary traversals of the IR.

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

public:
 
   /// 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;
 
+  /// Find the definitions of the given tensor value or retrieve them from the
+  /// cache.
+  const SetVector<Value> &findDefinitionsCached(Value value);
+
+  /// Reset cached data structures.
+  void resetCache();
+
   /// Union the alias sets of `v1` and `v2`.
   void unionAliasSets(Value v1, Value v2);
 
   /// Union the equivalence classes of `v1` and `v2`.
   void unionEquivalenceClasses(Value v1, Value v2);
 
   /// Base class for OneShotAnalysisState extensions that allow
   /// OneShotAnalysisState to contain user-specified information in the state
struct ValueComparator {
     }
   };
 
   using EquivalenceClassRangeType = llvm::iterator_range<
       llvm::EquivalenceClasses<Value, ValueComparator>::member_iterator>;
   /// Check that aliasInfo for `v` exists and return a reference to it.
   EquivalenceClassRangeType getAliases(Value v) const;
 
+  /// Cache definitions of tensor values.
+  DenseMap<Value, SetVector<Value>> cachedDefinitions;
+
   /// Set of all OpResults that were decided to bufferize in-place.
   llvm::DenseSet<OpOperand *> inplaceBufferized;
 
   /// Auxiliary structure to store all the values a given value may alias with.
   /// Alias information is "may be" conservative: In the presence of branches, a
   /// value may alias with one of multiple other values. The concrete aliasing
   /// value may not even be known at compile time. All such values are
   /// considered to be aliases.

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

 
 #include "mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h"
 #include "mlir/IR/Operation.h"
 
 namespace mlir {
 namespace bufferization {
 class AnalysisState;
 struct BufferizationStatistics;
+class OneShotAnalysisState;
 struct OneShotBufferizationOptions;
 
 /// A function that matches anchor OpOperands for tensor::EmptyOp elimination.
 /// If an OpOperand is matched, the function should populate the SmallVector
 /// with all values that are needed during `RewriteFn` to produce the
 /// replacement value.
 using AnchorMatchFn = std::function<bool(OpOperand &, SmallVector<Value> &)>;
 
 /// A function that rewrites matched anchors.
 using RewriteFn = std::function<Value(OpBuilder &, Location, OpOperand &)>;
 
 /// Try to eliminate tensor::EmptyOps inside `op`.
 ///
 /// * `rewriteFunc` generates the replacement for the tensor::EmptyOp.
 /// * Only tensor::EmptyOps that are anchored on a matching OpOperand as per
 ///   `anchorMatchFunc` are considered. "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 eliminateEmptyTensors(RewriterBase &rewriter, Operation *op,
-                                    bufferization::AnalysisState &state,
+                                    OneShotAnalysisState &state,
                                     AnchorMatchFn anchorMatchFunc,
                                     RewriteFn rewriteFunc);
 
 /// Try to eliminate tensor::EmptyOps inside `op` that are anchored on an
 /// InsertSliceOp, i.e., if it is eventually inserted into another tensor
 /// (and some other conditions are met).
 LogicalResult insertSliceAnchoredEmptyTensorEliminationStep(
-    RewriterBase &rewriter, Operation *op, bufferization::AnalysisState &state);
+    RewriterBase &rewriter, Operation *op, OneShotAnalysisState &state);
 
 /// Resolve RaW and other conflicts by inserting bufferization.alloc_tensor ops.
 /// After applying this transform, the IR can be bufferized without inserting
 /// additional buffer allocations.
 LogicalResult insertTensorCopies(Operation *op,
                                  const OneShotBufferizationOptions &options,
                                  BufferizationStatistics *statistics = nullptr);
 

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

 }
 
 /// 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,
+    RewriterBase &rewriter, Operation *op, OneShotAnalysisState &state,
     AnchorMatchFn anchorMatchFunc, RewriteFn rewriteFunc) {
   OpBuilder::InsertionGuard g(rewriter);
 
   WalkResult status = op->walk([&](Operation *op) {
     for (OpOperand &operand : op->getOpOperands()) {
       // Skip operands that do not bufferize inplace.
       if (!state.isInPlace(operand))
         continue;
for (OpOperand &operand : op->getOpOperands()) {
       // Create a replacement for the tensor::EmptyOp.
       rewriter.setInsertionPoint(insertionPoint);
       Value replacement = rewriteFunc(rewriter, emptyTensor.getLoc(), operand);
       if (!replacement)
         continue;
 
       // Replace the tensor::EmptyOp.
       rewriter.replaceOp(emptyTensor.getDefiningOp(), replacement);
+      state.resetCache();
     }
 
     // Advance to the next operation.
     return WalkResult::advance();
   });
 
   return failure(status.wasInterrupted());
 }

 /// source's reverse use-def chain is eliminated if:
 /// * On the reverse use-def chain path from the InsertSliceOp to the
 ///   tensor::EmptyOp, all ops were decided to bufferize inplace and the buffer
 ///   relation is "equivalent" (TODO: can be relaxed if needed).
 /// * The reverse use-def chain has exactly one end, which is the
 ///   tensor::EmptyOp.
 template <typename OpTy>
 static LogicalResult insertSliceLikeAnchoredEmptyTensorEliminationStep(
-    RewriterBase &rewriter, Operation *op, AnalysisState &state) {
+    RewriterBase &rewriter, Operation *op, OneShotAnalysisState &state) {
   return eliminateEmptyTensors(
       rewriter, op, state,
       /*anchorMatchFunc=*/
       [&](OpOperand &operand, SmallVector<Value> &neededValues) {
         auto insertSliceOp = dyn_cast<OpTy>(operand.getOwner());
         if (!insertSliceOp)
           return false;
         if (&operand != &insertSliceOp->getOpOperand(0) /*source*/)
return eliminateEmptyTensors(
             insertOp.getMixedOffsets(), insertOp.getMixedSizes(),
             insertOp.getMixedStrides());
         return extractOp.getResult();
       });
 }
 
 LogicalResult
 mlir::bufferization::insertSliceAnchoredEmptyTensorEliminationStep(
-    RewriterBase &rewriter, Operation *op, AnalysisState &state) {
+    RewriterBase &rewriter, Operation *op, OneShotAnalysisState &state) {
   if (failed(insertSliceLikeAnchoredEmptyTensorEliminationStep<
              tensor::InsertSliceOp>(rewriter, op, state)))
     return failure();
   if (failed(insertSliceLikeAnchoredEmptyTensorEliminationStep<
              tensor::ParallelInsertSliceOp>(rewriter, op, state)))
     return failure();
   return success();
 }

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

op->walk([&](Operation *op) {
 
     // Check all tensor OpResults.
     for (OpResult opResult : op->getOpResults()) {
       if (!opResult.getType().isa<TensorType>())
         continue;
 
       // If there is no preceding definition, the tensor contents are
       // undefined.
-      if (findDefinitions(opResult).empty())
+      if (findDefinitionsCached(opResult).empty())
         for (OpOperand &use : opResult.getUses())
           undefinedTensorUses.insert(&use);
     }
 
     return WalkResult::advance();
   });
 }
 
for (OpOperand *uRead : usesRead) {
     //
     // %0 = "writing_op"(%t) : tensor<?x32> -> tensor<?xf32>
     // %1 = "aliasing_op"(%0) : tensor<?x32> -> tensor<?xf32>
     // %2 = "reading_op"(%1) : : tensor<?x32> -> not_a_tensor_type
     //
     // In the above example, if uRead is the OpOperand of reading_op, the
     // definition is %0. Note that operations that create an alias but do not
     // bufferize to a memory write (such as ExtractSliceOp) are skipped.
-    SetVector<Value> definitions = state.findDefinitions(uRead->get());
+    const SetVector<Value> &definitions =
+        state.findDefinitionsCached(uRead->get());
     if (definitions.empty()) {
       // Fast path: No conflict if there are no definitions.
       LLVM_DEBUG(llvm::dbgs()
                  << "  no conflict: read value has no definitions\n");
       continue;
     }
 
     // Look for conflicting memory writes. Potential conflicts are writes to an
wouldCreateWriteToNonWritableBuffer(OpOperand &operand,
 
   return false;
 }
 
 //===----------------------------------------------------------------------===//
 // Bufferization analyses.
 //===----------------------------------------------------------------------===//
 
+// Find the values that define the contents of the given value.
+const llvm::SetVector<Value> &
+OneShotAnalysisState::findDefinitionsCached(Value value) {
+  if (!cachedDefinitions.count(value)) {
+    cachedDefinitions[value] = findValueInReverseUseDefChain(
+        value, [&](Value v) { return this->bufferizesToMemoryWrite(v); },
+        /*followEquivalentOnly=*/false, /*alwaysIncludeLeaves=*/false);
+  }
+  return cachedDefinitions[value];
+}
+
+void OneShotAnalysisState::resetCache() { cachedDefinitions.clear(); }
+
 /// Determine if `operand` can be bufferized in-place.
 static LogicalResult
 bufferizableInPlaceAnalysisImpl(OpOperand &operand, OneShotAnalysisState &state,
                                 const DominanceInfo &domInfo) {
   LLVM_DEBUG(
       llvm::dbgs() << "//===-------------------------------------------===//\n"
                    << "Analyzing operand #" << operand.getOperandNumber()
                    << " of " << *operand.getOwner() << "\n");