[mlir][bufferize] Add new debug flag: copy-before-write

[springerm]

Summary

If this flag is set, the analysis is skipped and buffers are copied before every write.

[tpopp]

Awesome! You use copyBeforeWrite and skipAnalysis. I think it would be better to use the same name consistently, like the other options/flags, unless skip-analysis is expected to grow to modifying multiple option values(?)

[springerm]

address comments

[springerm]

In D133288#3769842, @tpopp wrote:

Awesome! You use copyBeforeWrite and skipAnalysis. I think it would be better to use the same name consistently, like the other options/flags, unless skip-analysis is expected to grow to modifying multiple option values(?)

Using copy-before-write everywhere now.

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

struct BufferizationOptions {
   /// Specifies whether dealloc ops should be generated along with alloc ops. If
   /// not, new memory allocations will leak.
   bool createDeallocs = true;
 
   /// Seed for the analysis fuzzer. If set to `0`, the fuzzer is deactivated.
   /// Should be used only with `testAnalysisOnly = true`.
   unsigned analysisFuzzerSeed = 0;
 
+  /// If set to `true`, the analysis is skipped. A buffer is copied before every
+  /// write. This flag cannot be used together with `testAnalysisOnly = true`.
+  bool copyBeforeWrite = false;
+
   /// If set to `true`, does not modify the IR apart from adding attributes (for
   /// checking the results of the analysis) and post analysis steps.
   bool testAnalysisOnly = false;
 
   /// If set to `true`, the IR is annotated with details about RaW conflicts.
   /// For debugging only. Should be used together with `testAnalysisOnly`.
   bool printConflicts = false;
 

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

 class OneShotAnalysisState;
 struct OneShotBufferizationOptions;
 
 /// Analyze `moduleOp` and its nested ops. Bufferization decisions are stored in
 /// `state`.
 LogicalResult analyzeModuleOp(ModuleOp moduleOp, OneShotAnalysisState &state);
 
 /// Bufferize `op` and its nested ops that implement `BufferizableOpInterface`.
-/// Whether buffer copies are needed or not is queried from the given state.
+///
+/// Note: This function does not run One-Shot Analysis. No buffer copies are
+/// inserted unless `options.copyBeforeWrite` is set, in which case buffers are
+/// copied before every write.
 LogicalResult bufferizeModuleOp(ModuleOp moduleOp,
-                                const OneShotAnalysisState &analysisState);
+                                const OneShotBufferizationOptions &options);
 
 /// Run One-Shot Module Bufferization on the given module. Performs a simple
 /// function call analysis to determine which function arguments are
 /// inplaceable. Then analyzes and bufferizes FuncOps one-by-one with One-Shot
 /// Bufferize.
 LogicalResult runOneShotModuleBufferize(
     ModuleOp moduleOp,
     const bufferization::OneShotBufferizationOptions &options);
 
 } // namespace bufferization
 } // namespace mlir
 
 #endif // MLIR_DIALECT_BUFFERIZATION_TRANSFORMS_ONESHOTMODULEBUFFERIZE_H

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

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<"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.">,
     ListOption<"dialectFilter", "dialect-filter", "std::string",
                "Restrict bufferization to ops from these dialects.">,
     Option<"functionBoundaryTypeConversion",
            "function-boundary-type-conversion", "std::string",
            /*default=*/"\"infer-layout-map\"",

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

struct OneShotBufferizePass
   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.copyBeforeWrite = copyBeforeWrite;
       opt.createDeallocs = createDeallocs;
       opt.functionBoundaryTypeConversion =
           parseLayoutMapOption(functionBoundaryTypeConversion);
       if (mustInferMemorySpace)
         opt.defaultMemorySpace = None;
       opt.printConflicts = printConflicts;
       opt.testAnalysisOnly = testAnalysisOnly;
       opt.bufferizeFunctionBoundaries = bufferizeFunctionBoundaries;

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

if (options.testAnalysisOnly)
     annotateOpsWithBufferizationMarkers(op, aliasInfo, state);
 
   return success(!failedAnalysis);
 }
 
 LogicalResult
 bufferization::runOneShotBufferize(Operation *op,
                                    const OneShotBufferizationOptions &options) {
+  assert(!(options.copyBeforeWrite && options.testAnalysisOnly) &&
+         "invalid combination of bufferization flags");
+  if (!options.copyBeforeWrite) {
+    // If a buffer is copied before every write, no analysis is needed.
     OneShotAnalysisState state(op, options);
     if (failed(insertTensorCopies(op, options)))
       return failure();
+  }
   if (options.testAnalysisOnly)
     return success();
-  return bufferizeOp(op, options, /*copyBeforeWrite=*/false);
+  return bufferizeOp(op, options, /*copyBeforeWrite=*/options.copyBeforeWrite);
 }

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

for (func::FuncOp funcOp : orderedFuncOps) {
     // Mark op as fully analyzed.
     funcState.analyzedFuncOps[funcOp] = FuncOpAnalysisState::Analyzed;
   }
 
   return success();
 }
 
 LogicalResult mlir::bufferization::bufferizeModuleOp(
-    ModuleOp moduleOp, const OneShotAnalysisState &analysisState) {
-  auto const &options = static_cast<const OneShotBufferizationOptions &>(
-      analysisState.getOptions());
+    ModuleOp moduleOp, const OneShotBufferizationOptions &options) {
   assert(options.bufferizeFunctionBoundaries &&
          "expected that function boundary bufferization is activated");
   IRRewriter rewriter(moduleOp.getContext());
 
   // A list of functions in the order in which they are analyzed + bufferized.
   SmallVector<func::FuncOp> orderedFuncOps;
 
   // A mapping of FuncOps to their callers.
   FuncCallerMap callerMap;
 
   if (failed(getFuncOpsOrderedByCalls(moduleOp, orderedFuncOps, callerMap)))
     return failure();
 
   // Bufferize functions.
   for (func::FuncOp funcOp : orderedFuncOps) {
     // Note: It would be good to apply cleanups here but we cannot as aliasInfo
     // would be invalidated.
-    if (failed(bufferizeOp(funcOp, options, /*copyBeforeWrite=*/false)))
+    if (failed(bufferizeOp(funcOp, options, options.copyBeforeWrite)))
       return failure();
     // Change buffer return types to more precise layout maps.
     if (options.functionBoundaryTypeConversion ==
         BufferizationOptions::LayoutMapOption::InferLayoutMap)
       foldMemRefCasts(funcOp);
   }
 
   // Post-pass cleanup of function argument attributes.
   moduleOp.walk([&](func::FuncOp op) {
     for (BlockArgument bbArg : op.getArguments())
       removeBufferizationAttributes(bbArg);
   });
 
   return success();
 }
 
 LogicalResult mlir::bufferization::runOneShotModuleBufferize(
     ModuleOp moduleOp, const OneShotBufferizationOptions &options) {
   assert(options.bufferizeFunctionBoundaries &&
          "expected that function boundary bufferization is activated");
+  assert(!(options.copyBeforeWrite && options.testAnalysisOnly) &&
+         "invalid combination of bufferization flags");
+  if (!options.copyBeforeWrite) {
     OneShotAnalysisState analysisState(moduleOp, options);
     if (failed(insertTensorCopies(moduleOp, options)))
       return failure();
+  }
   if (options.testAnalysisOnly)
     return success();
-  if (failed(bufferizeModuleOp(moduleOp, analysisState)))
+  if (failed(bufferizeModuleOp(moduleOp, options)))
     return failure();
   return success();
 }

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
 
+// 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
+//  CHECK-NEXT:   memref.store
+// CHECK-COPY-BEFORE-WRITE-LABEL: func @no_conflict
+//       CHECK-COPY-BEFORE-WRITE:   memref.alloc
+//       CHECK-COPY-BEFORE-WRITE:   memref.store
+//       CHECK-COPY-BEFORE-WRITE:   memref.store
+//       CHECK-COPY-BEFORE-WRITE:   memref.store
+//       CHECK-COPY-BEFORE-WRITE:   memref.alloc
+//       CHECK-COPY-BEFORE-WRITE:   memref.copy
+//       CHECK-COPY-BEFORE-WRITE:   memref.store
+func.func @no_conflict(%fill: f32, %f: f32, %idx: index) -> tensor<3xf32> {
+  %t = tensor.from_elements %fill, %fill, %fill : tensor<3xf32>
+  %i = tensor.insert %f into %t[%idx] : tensor<3xf32>
+  return %i : tensor<3xf32>
+}
+
+// -----
+
 // CHECK-LABEL: func @use_tensor_func_arg(
 //  CHECK-SAME:     %[[A:.*]]: tensor<?xf32>
 func.func @use_tensor_func_arg(%A : tensor<?xf32>) -> (vector<4xf32>) {
   %c0 = arith.constant 0 : index
   %f0 = arith.constant 0.0 : f32
 
   // CHECK: %[[A_memref:.*]] = bufferization.to_memref %[[A]]
   // CHECK: %[[res:.*]] = vector.transfer_read %[[A_memref]]