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

[springerm]

Summary

This is to avoid unnecessary traversals of the IR.

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

public:
   /// %a = tensor.empty() : tensor<10xf32>
   /// %b = arith.constant ... : tensor<10xf32>
   /// %r = arith.select %cond, %a, %b : tensor<10xf32>
   /// findDefinitions(%r) = {%b}. %a is excluded because it does not define the
   /// contents of the tensor.
   ///
   /// Note: OpResults of unknown ops are handled conservatively and assumed to
   /// be definitions.
-  SetVector<Value> findDefinitions(Value value) const;
+  const SetVector<Value> &findDefinitions(Value value);
 
   /// Return `true` if the given OpResult has been decided to bufferize inplace.
   virtual bool isInPlace(OpOperand &opOperand) const;
 
   /// Return true if `v1` and `v2` bufferize to equivalent buffers.
   virtual bool areEquivalentBufferizedValues(Value v1, Value v2) const;
 
   /// Return true if `v1` and `v2` may bufferize to aliasing buffers.
protected:
   AnalysisState(const BufferizationOptions &options, TypeID type);
 
 private:
   /// A reference to current bufferization options.
   const BufferizationOptions &options;
 
   /// The type of analysis.
   TypeID type;
+
+  /// Cache definitions of tensor values.
+  DenseMap<Value, SetVector<Value>> cachedDefinitons;
 };
 
 /// Create an AllocTensorOp for the given shaped value (memref or tensor).
 /// If `copy` is set, the shaped value is copied. Otherwise, a tensor with
 /// undefined contents is allocated.
 FailureOr<Value>
 allocateTensorForShapedValue(OpBuilder &b, Location loc, Value shapedValue,
                              bool escape, const BufferizationOptions &options,

mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp

while (!workingSet.empty()) {
     for (OpOperand *o : aliases)
       workingSet.insert(o->get());
   }
 
   return result;
 }
 
 // Find the values that define the contents of the given value.
-llvm::SetVector<Value> AnalysisState::findDefinitions(Value value) const {
-  return findValueInReverseUseDefChain(
+const llvm::SetVector<Value> &AnalysisState::findDefinitions(Value value) {
+  if (!cachedDefinitons.count(value)) {
+    cachedDefinitons[value] = findValueInReverseUseDefChain(
         value, [&](Value v) { return this->bufferizesToMemoryWrite(v); },
         /*followEquivalentOnly=*/false, /*alwaysIncludeLeaves=*/false);
   }
+  return cachedDefinitons[value];
+}
 
 AnalysisState::AnalysisState(const BufferizationOptions &options)
     : AnalysisState(options, TypeID::get<AnalysisState>()) {}
 
 AnalysisState::AnalysisState(const BufferizationOptions &options, TypeID type)
     : options(options), type(type) {
   for (const BufferizationOptions::AnalysisStateInitFn &fn :
        options.stateInitializers)

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

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.findDefinitions(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