[mlir][bufferize] Treat certain aliasing-only uses like memory reads

[springerm]

Summary

This fixes an issue in One-Shot Bufferize that could lead to missing buffer copies in the future. This bug can currently not be triggered because of the order in which ops are analyzed (always bottom-to-top). However, if we consider different traversal orders for the analysis in the future, this bug can cause subtle issues that are difficult to debug.

Example:

%0 = ...
%1 = tensor.insert ... into %0
%2 = tensor.extract_slice %0
tensor.extract %2[...]

In case of a top-to-bottom analysis of the above IR, the tensor.insert is analyzed before the tensor.extract_slice. In that case, the tensor.insert will bufferize in-place because %2 is not yet known to become an alias of %0 (and therefore causing a conflict).

With this change, the tensor.insert will bufferize out-of-place, regardless of the traversal order.

Note: As this bug cannot be triggered with the current traversal order, there is no test case for this change. However, a subsequent change will introduce an improvement to loop buffferization that depends on this change. That change will have a test case that depends on this fix.

[springerm]

rebase

[nicolasvasilache]

let's please add a test or this because this is load bearing, approving conditioned on that.

[springerm]

address comments

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

 namespace bufferization {
 
 struct OneShotBufferizationOptions;
 class BufferizationAliasInfo;
 class OneShotAnalysisState;
 
 /// Options for analysis-enabled bufferization.
 struct OneShotBufferizationOptions : public BufferizationOptions {
+  enum class AnalysisHeuristic { BottomUp, TopDown };
+
   OneShotBufferizationOptions() = default;
 
   /// Specifies whether returning newly allocated memrefs should be allowed.
   /// Otherwise, a pass failure is triggered.
   bool allowReturnAllocs = false;
+
+  /// The heuristic controls the order in which ops are traversed during the
+  /// analysis.
+  AnalysisHeuristic analysisHeuristic = AnalysisHeuristic::BottomUp;
 };
 
 /// The BufferizationAliasInfo class maintains a list of buffer aliases and
 /// equivalence classes to support bufferization.
 class BufferizationAliasInfo {
 public:
   explicit BufferizationAliasInfo(Operation *rootOp);
 

mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td

Option<"allowReturnAllocs", "allow-return-allocs", "bool",
             /*default=*/"false",
            "Allows returning/yielding new allocations from a block.">,
     Option<"allowUnknownOps", "allow-unknown-ops", "bool",
            /*default=*/"false",
            "Allows unknown (not bufferizable) ops in the input IR.">,
     Option<"analysisFuzzerSeed", "analysis-fuzzer-seed", "unsigned",
            /*default=*/"0",
            "Test only: Analyze ops in random order with a given seed (fuzzer)">,
+    Option<"analysisHeuristic", "analysis-heuristic", "std::string",
+           /*default=*/"\"bottom-up\"",
+           "Heuristic that control the IR traversal during analysis">,
     Option<"bufferizeFunctionBoundaries", "bufferize-function-boundaries",
            "bool", /*default=*/"0",
            "Bufferize function boundaries (experimental).">,
     Option<"copyBeforeWrite", "copy-before-write", "bool", /*default=*/"false",
            "Skip the analysis. Make a buffer copy on every write.">,
     Option<"createDeallocs", "create-deallocs", "bool", /*default=*/"true",
            "Specify if buffers should be deallocated. For compatibility with "
            "core bufferization passes.">,

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

if (s == "fully-dynamic-layout-map")
     return BufferizationOptions::LayoutMapOption::FullyDynamicLayoutMap;
   if (s == "identity-layout-map")
     return BufferizationOptions::LayoutMapOption::IdentityLayoutMap;
   if (s == "infer-layout-map")
     return BufferizationOptions::LayoutMapOption::InferLayoutMap;
   llvm_unreachable("invalid layout map option");
 }
 
+static OneShotBufferizationOptions::AnalysisHeuristic
+parseHeuristicOption(const std::string &s) {
+  if (s == "bottom-up")
+    return OneShotBufferizationOptions::AnalysisHeuristic::BottomUp;
+  if (s == "top-down")
+    return OneShotBufferizationOptions::AnalysisHeuristic::TopDown;
+  llvm_unreachable("invalid analysisheuristic option");
+}
+
 struct OneShotBufferizePass
     : public bufferization::impl::OneShotBufferizeBase<OneShotBufferizePass> {
   OneShotBufferizePass() = default;
 
   explicit OneShotBufferizePass(const OneShotBufferizationOptions &options)
       : options(options) {}
 
   void getDependentDialects(DialectRegistry &registry) const override {
     registry
         .insert<bufferization::BufferizationDialect, memref::MemRefDialect>();
     registerAllocationOpInterfaceExternalModels(registry);
   }
 
   void runOnOperation() override {
     OneShotBufferizationOptions opt;
     if (!options) {
       // Make new bufferization options if none were provided when creating the
       // pass.
       opt.allowReturnAllocs = allowReturnAllocs;
       opt.allowUnknownOps = allowUnknownOps;
       opt.analysisFuzzerSeed = analysisFuzzerSeed;
+      opt.analysisHeuristic = parseHeuristicOption(analysisHeuristic);
       opt.copyBeforeWrite = copyBeforeWrite;
       opt.createDeallocs = createDeallocs;
       opt.functionBoundaryTypeConversion =
           parseLayoutMapOption(functionBoundaryTypeConversion);
       if (mustInferMemorySpace)
         opt.defaultMemorySpace = None;
       opt.printConflicts = printConflicts;
       opt.testAnalysisOnly = testAnalysisOnly;

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

static void getAliasingInplaceWrites(DenseSet<OpOperand *> &res, Value root,
   });
 }
 
 // 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))
+    for (auto &use : alias.getUses()) {
+      // Read of a value that aliases root.
+      if (state.bufferizesToMemoryRead(use)) {
+        res.insert(&use);
+        continue;
+      }
+
+      // Read of a dependent value in the SSA use-def chain. E.g.:
+      //
+      // %0 = ...
+      // %1 = tensor.extract_slice %0 {not_analyzed_yet}
+      // "read"(%1)
+      //
+      // In the above example, getAliasingReads(%0) includes the first OpOperand
+      // of the tensor.extract_slice op. The extract_slice itself does not read
+      // but its aliasing result is eventually fed into an op that does.
+      //
+      // Note: This is considered a "read" only if the use does not bufferize to
+      // a memory write. (We already ruled out memory reads. In case of a memory
+      // write, the buffer would be entirely overwritten; in the above example
+      // there would then be no flow of data from the extract_slice operand to
+      // its result's uses.)
+      if (!state.bufferizesToMemoryWrite(use)) {
+        SmallVector<OpResult> opResults = state.getAliasingOpResult(use);
+        if (llvm::any_of(opResults,
+                         [&](OpResult r) { return state.isValueRead(r); }))
           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
static LogicalResult inPlaceAnalysis(SmallVector<Operation *> &ops,
   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);
   }
 
-  // Walk ops in reverse for better interference analysis.
-  for (Operation *op : reverse(ops))
+  // Analyze a single op.
+  auto analyzeOp = [&](Operation *op) {
     for (OpOperand &opOperand : op->getOpOperands())
       if (opOperand.get().getType().isa<TensorType>())
         if (auto bufferizableOp = state.getOptions().dynCastBufferizableOp(op))
           if (failed(bufferizableInPlaceAnalysisImpl(opOperand, aliasInfo,
                                                      state, domInfo)))
             return failure();
+    return success();
+  };
+
+  OneShotBufferizationOptions::AnalysisHeuristic heuristic =
+      static_cast<const OneShotBufferizationOptions &>(state.getOptions())
+          .analysisHeuristic;
+  if (heuristic == OneShotBufferizationOptions::AnalysisHeuristic::BottomUp) {
+    // Default: Walk ops in reverse for better interference analysis.
+    for (Operation *op : reverse(ops))
+      if (failed(analyzeOp(op)))
+        return failure();
+  } else if (heuristic ==
+             OneShotBufferizationOptions::AnalysisHeuristic::TopDown) {
+    for (Operation *op : ops)
+      if (failed(analyzeOp(op)))
+        return failure();
+  } else {
+    llvm_unreachable("unsupported heuristic");
+  }
 
   return success();
 }
 
 /// Return true if the given op has a tensor result or a tensor operand.
 static bool hasTensorSemantics(Operation *op) {
   bool hasTensorResult = any_of(op->getResultTypes(), isaTensor);
   bool hasTensorOperand = any_of(op->getOperandTypes(), isaTensor);

mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize.mlir

 // RUN: mlir-opt %s -one-shot-bufferize="allow-unknown-ops" -split-input-file | FileCheck %s
 
 // Run fuzzer with different seeds.
 // RUN: mlir-opt %s -one-shot-bufferize="test-analysis-only analysis-fuzzer-seed=23" -split-input-file -o /dev/null
 // RUN: mlir-opt %s -one-shot-bufferize="test-analysis-only analysis-fuzzer-seed=59" -split-input-file -o /dev/null
 // RUN: mlir-opt %s -one-shot-bufferize="test-analysis-only analysis-fuzzer-seed=91" -split-input-file -o /dev/null
 
+// Run with top-down analysis.
+// RUN: mlir-opt %s -one-shot-bufferize="allow-unknown-ops analysis-heuristic=top-down" -split-input-file | FileCheck %s --check-prefix=CHECK-TOP-DOWN-ANALYSIS
+
 // Test without analysis: Insert a copy on every buffer write.
 // RUN: mlir-opt %s -allow-unregistered-dialect -one-shot-bufferize="allow-unknown-ops copy-before-write" -split-input-file | FileCheck %s --check-prefix=CHECK-COPY-BEFORE-WRITE
 
 // CHECK-LABEL: func @no_conflict
 //       CHECK:   memref.alloc
 //       CHECK:   memref.store
 //  CHECK-NEXT:   memref.store
 //  CHECK-NEXT:   memref.store

 func.func @alloc_tensor_with_memory_space() -> tensor<5xf32> {
   // CHECK: %[[alloc:.*]] = memref.alloc() {{.*}} : memref<5xf32, 1>
   %0 = bufferization.alloc_tensor() {memory_space = 1 : ui64} : tensor<5xf32>
   // CHECK: %[[r:.*]] = bufferization.to_tensor %[[alloc]]
   // CHECK: memref.dealloc %[[alloc]]
   // CHECK: return %[[r]]
   return %0 : tensor<5xf32>
 }
+
+// -----
+
+// CHECK-LABEL: func @read_of_alias
+// CHECK-TOP-DOWN-ANALYSIS-LABEL: func @read_of_alias
+func.func @read_of_alias(%t: tensor<100xf32>, %pos1: index, %pos2: index,
+                         %pos3: index, %pos4: index, %sz: index, %f: f32)
+  -> (f32, f32)
+{
+  // CHECK: %[[alloc:.*]] = memref.alloc
+  // CHECK: memref.copy
+  // CHECK: memref.store %{{.*}}, %[[alloc]]
+  // CHECK-TOP-DOWN-ANALYSIS: %[[alloc:.*]] = memref.alloc
+  // CHECK-TOP-DOWN-ANALYSIS: memref.copy
+  // CHECK-TOP-DOWN-ANALYSIS: memref.store %{{.*}}, %[[alloc]]
+  %0 = tensor.insert %f into %t[%pos1] : tensor<100xf32>
+  %1 = tensor.extract_slice %t[%pos2][%sz][1] : tensor<100xf32> to tensor<?xf32>
+  %2 = tensor.extract %1[%pos3] : tensor<?xf32>
+  %3 = tensor.extract %0[%pos3] : tensor<100xf32>
+  return %2, %3 : f32, f32
+}