diff --git a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h --- a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h +++ b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h @@ -287,6 +287,10 @@ /// 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; diff --git a/mlir/include/mlir/Dialect/Bufferization/Transforms/OneShotModuleBufferize.h b/mlir/include/mlir/Dialect/Bufferization/Transforms/OneShotModuleBufferize.h --- a/mlir/include/mlir/Dialect/Bufferization/Transforms/OneShotModuleBufferize.h +++ b/mlir/include/mlir/Dialect/Bufferization/Transforms/OneShotModuleBufferize.h @@ -23,9 +23,12 @@ 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 diff --git a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td --- a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td +++ b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td @@ -276,6 +276,8 @@ 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.">, diff --git a/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp b/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp --- a/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp +++ b/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp @@ -193,6 +193,7 @@ opt.allowReturnAllocs = allowReturnAllocs; opt.allowUnknownOps = allowUnknownOps; opt.analysisFuzzerSeed = analysisFuzzerSeed; + opt.copyBeforeWrite = copyBeforeWrite; opt.createDeallocs = createDeallocs; opt.functionBoundaryTypeConversion = parseLayoutMapOption(functionBoundaryTypeConversion); diff --git a/mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp b/mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp --- a/mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp +++ b/mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp @@ -1025,10 +1025,15 @@ LogicalResult bufferization::runOneShotBufferize(Operation *op, const OneShotBufferizationOptions &options) { - OneShotAnalysisState state(op, options); - if (failed(insertTensorCopies(op, options))) - return failure(); + 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); } diff --git a/mlir/lib/Dialect/Bufferization/Transforms/OneShotModuleBufferize.cpp b/mlir/lib/Dialect/Bufferization/Transforms/OneShotModuleBufferize.cpp --- a/mlir/lib/Dialect/Bufferization/Transforms/OneShotModuleBufferize.cpp +++ b/mlir/lib/Dialect/Bufferization/Transforms/OneShotModuleBufferize.cpp @@ -401,9 +401,7 @@ } LogicalResult mlir::bufferization::bufferizeModuleOp( - ModuleOp moduleOp, const OneShotAnalysisState &analysisState) { - auto const &options = static_cast( - analysisState.getOptions()); + ModuleOp moduleOp, const OneShotBufferizationOptions &options) { assert(options.bufferizeFunctionBoundaries && "expected that function boundary bufferization is activated"); IRRewriter rewriter(moduleOp.getContext()); @@ -421,7 +419,7 @@ 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 == @@ -442,12 +440,16 @@ ModuleOp moduleOp, const OneShotBufferizationOptions &options) { assert(options.bufferizeFunctionBoundaries && "expected that function boundary bufferization is activated"); - OneShotAnalysisState analysisState(moduleOp, options); - if (failed(insertTensorCopies(moduleOp, options))) - return failure(); + 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(); } diff --git a/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize.mlir b/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize.mlir --- a/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize.mlir +++ b/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize.mlir @@ -5,6 +5,31 @@ // 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 func.func @use_tensor_func_arg(%A : tensor) -> (vector<4xf32>) {