[Mlir-commits] [mlir] f7dd9a3 - [mlir][bufferize] Add new debug flag: copy-before-write
Matthias Springer
llvmlistbot at llvm.org
Mon Sep 5 05:41:31 PDT 2022
Author: Matthias Springer
Date: 2022-09-05T14:41:19+02:00
New Revision: f7dd9a32060ab5b024424d24f6fbbd54b301a491
URL: https://github.com/llvm/llvm-project/commit/f7dd9a32060ab5b024424d24f6fbbd54b301a491
DIFF: https://github.com/llvm/llvm-project/commit/f7dd9a32060ab5b024424d24f6fbbd54b301a491.diff
LOG: [mlir][bufferize] Add new debug flag: copy-before-write
If this flag is set, the analysis is skipped and buffers are copied before every write.
Differential Revision: https://reviews.llvm.org/D133288
Added:
Modified:
mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h
mlir/include/mlir/Dialect/Bufferization/Transforms/OneShotModuleBufferize.h
mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp
mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp
mlir/lib/Dialect/Bufferization/Transforms/OneShotModuleBufferize.cpp
mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize.mlir
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h
index dc2b12f433a2..cd2be2925c68 100644
--- a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h
+++ b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h
@@ -287,6 +287,10 @@ struct BufferizationOptions {
/// 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
index 782b2c4aeeda..aafa002f4206 100644
--- a/mlir/include/mlir/Dialect/Bufferization/Transforms/OneShotModuleBufferize.h
+++ b/mlir/include/mlir/Dialect/Bufferization/Transforms/OneShotModuleBufferize.h
@@ -23,9 +23,12 @@ struct OneShotBufferizationOptions;
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
index 97684fb4a931..f2d699985bfe 100644
--- a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
@@ -276,6 +276,8 @@ def OneShotBufferize : Pass<"one-shot-bufferize", "ModuleOp"> {
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
index 1a328d871235..74c275c73c57 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp
@@ -193,6 +193,7 @@ struct OneShotBufferizePass
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
index 46420f172f7b..100fc7767994 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp
@@ -1025,10 +1025,15 @@ LogicalResult bufferization::analyzeOp(Operation *op,
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
index 9e6692999ea8..2442337add00 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/OneShotModuleBufferize.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/OneShotModuleBufferize.cpp
@@ -401,9 +401,7 @@ mlir::bufferization::analyzeModuleOp(ModuleOp moduleOp,
}
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());
@@ -421,7 +419,7 @@ LogicalResult mlir::bufferization::bufferizeModuleOp(
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 @@ LogicalResult mlir::bufferization::runOneShotModuleBufferize(
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
index 3bc380f59eef..363011ce0760 100644
--- 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<?xf32>
func.func @use_tensor_func_arg(%A : tensor<?xf32>) -> (vector<4xf32>) {
More information about the Mlir-commits
mailing list