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 @@ -9,10 +9,14 @@ #ifndef MLIR_DIALECT_BUFFERIZATION_TRANSFORMS_ONESHOTMODULEBUFFERIZE_H #define MLIR_DIALECT_BUFFERIZATION_TRANSFORMS_ONESHOTMODULEBUFFERIZE_H +#include "llvm/ADT/SmallSet.h" namespace mlir { struct LogicalResult; class ModuleOp; +namespace func { +class FuncOp; +} namespace bufferization { struct BufferizationStatistics; @@ -20,18 +24,26 @@ struct OneShotBufferizationOptions; /// Analyze `moduleOp` and its nested ops. Bufferization decisions are stored in -/// `state`. -LogicalResult analyzeModuleOp(ModuleOp moduleOp, OneShotAnalysisState &state, - BufferizationStatistics *statistics = nullptr); +/// `state`. If the analysis of a FuncOp fails, it is inserted to +/// `notAnalyzedFuncOps`. +LogicalResult +analyzeModuleOp(ModuleOp moduleOp, OneShotAnalysisState &state, + BufferizationStatistics *statistics = nullptr, + llvm::SmallSet *notAnalyzedFuncOps = nullptr); /// Bufferize `op` and its nested ops that implement `BufferizableOpInterface`. /// /// 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 OneShotBufferizationOptions &options, - BufferizationStatistics *statistics = nullptr); +/// inserted except two cases: +/// - `options.copyBeforeWrite` is set, in which case buffers are copied before +/// every write. +/// - `options.copyBeforeWrite` is not set and `notAnalyzedFuncOps` contains +/// FuncOps. These FuncOps were not analyzed. Buffer copies will be inserted +/// only to these FuncOps. +LogicalResult +bufferizeModuleOp(ModuleOp moduleOp, const OneShotBufferizationOptions &options, + llvm::SmallSet ¬AnalyzedFuncOps, + BufferizationStatistics *statistics = nullptr); /// Remove bufferization attributes on every FuncOp arguments in the ModuleOp. void removeBufferizationAttributesInModule(ModuleOp moduleOp); diff --git a/mlir/include/mlir/Dialect/Bufferization/Transforms/Transforms.h b/mlir/include/mlir/Dialect/Bufferization/Transforms/Transforms.h --- a/mlir/include/mlir/Dialect/Bufferization/Transforms/Transforms.h +++ b/mlir/include/mlir/Dialect/Bufferization/Transforms/Transforms.h @@ -10,7 +10,9 @@ #define MLIR_DIALECT_BUFFERIZATION_TRANSFORMS_TRANSFORMS_H #include "mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h" +#include "mlir/Dialect/Func/IR/FuncOps.h" #include "mlir/IR/Operation.h" +#include "llvm/ADT/SmallSet.h" namespace mlir { namespace bufferization { @@ -48,10 +50,12 @@ /// Resolve RaW and other conflicts by inserting bufferization.alloc_tensor ops. /// After applying this transform, the IR can be bufferized without inserting -/// additional buffer allocations. -LogicalResult insertTensorCopies(Operation *op, - const OneShotBufferizationOptions &options, - BufferizationStatistics *statistics = nullptr); +/// additional buffer allocations. `notAnalyzedFuncOps` holds the FuncOps that +/// could not be analyzed. +LogicalResult insertTensorCopies( + Operation *op, const OneShotBufferizationOptions &options, + BufferizationStatistics *statistics = nullptr, + llvm::SmallSet *notAnalyzedFuncOps = nullptr); /// Resolve RaW and other conflicts by inserting bufferization.alloc_tensor ops. /// After applying this transform, the IR can be bufferized without inserting 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 @@ -1004,7 +1004,7 @@ equivalenceAnalysis(ops, aliasInfo, state); } -/// Assert that the current bufferization decisions are consistent. +/// Return success if the current bufferization decisions are consistent. static LogicalResult checkAliasInfoConsistency(Operation *op, const DominanceInfo &domInfo, AnalysisState &state, @@ -1016,10 +1016,11 @@ if (!options.isOpAllowed(op.getOperation())) return WalkResult::advance(); - // Input IR may not contain any ToMemrefOps. These are not supported because - // the analysis cannot follow the data flow through memrefs. + // ToMemrefOps are not supported because the analysis cannot follow the data + // flow through memrefs. if (isa(op.getOperation())) { - op->emitError("to_memref ops not supported during One-Shot Analysis"); + if (options.testAnalysisOnly || !options.bufferizeFunctionBoundaries) + op->emitError("to_memref ops not supported during One-Shot Analysis"); return WalkResult::interrupt(); } @@ -1031,7 +1032,8 @@ // This error can happen if certain "mustBufferizeInPlace" interface // methods are implemented incorrectly, such that the IR already has // a RaW conflict before making any bufferization decisions. - op->emitError("input IR has RaW conflict"); + if (options.testAnalysisOnly || !options.bufferizeFunctionBoundaries) + op->emitError("input IR has RaW conflict"); return WalkResult::interrupt(); } } 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 @@ -361,10 +361,10 @@ funcOp.setType(newFuncType); } -LogicalResult -mlir::bufferization::analyzeModuleOp(ModuleOp moduleOp, - OneShotAnalysisState &state, - BufferizationStatistics *statistics) { +LogicalResult mlir::bufferization::analyzeModuleOp( + ModuleOp moduleOp, OneShotAnalysisState &state, + BufferizationStatistics *statistics, + llvm::SmallSet *notAnalyzedFuncOps) { assert(state.getOptions().bufferizeFunctionBoundaries && "expected that function boundary bufferization is activated"); FuncAnalysisState &funcState = getOrCreateFuncAnalysisState(state); @@ -387,9 +387,12 @@ // Gather equivalence info for CallOps. equivalenceAnalysis(funcOp, aliasInfo, state, funcState); - // Analyze funcOp. - if (failed(analyzeOp(funcOp, state, statistics))) - return failure(); + // Analyze funcOp. If failed, mark it as not analyzed. + if (failed(analyzeOp(funcOp, state, statistics))) { + funcState.analyzedFuncOps[funcOp] = FuncOpAnalysisState::NotAnalyzed; + notAnalyzedFuncOps->insert(funcOp); + continue; + } // Run some extra function analyses. if (failed(aliasingFuncOpBBArgsAnalysis(funcOp, state, funcState)) || @@ -400,6 +403,8 @@ funcState.analyzedFuncOps[funcOp] = FuncOpAnalysisState::Analyzed; } + if (notAnalyzedFuncOps && !notAnalyzedFuncOps->empty()) + return failure(); return success(); } @@ -413,6 +418,7 @@ LogicalResult mlir::bufferization::bufferizeModuleOp( ModuleOp moduleOp, const OneShotBufferizationOptions &options, + llvm::SmallSet ¬AnalyzedFuncOps, BufferizationStatistics *statistics) { assert(options.bufferizeFunctionBoundaries && "expected that function boundary bufferization is activated"); @@ -431,7 +437,12 @@ 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, options.copyBeforeWrite, + // if options.copyBeforeWrite is not set, notAnalyzedFuncOps keeps track of + // FuncOps that were not analyzed. Buffer copies will be inserted to these + // FuncOps. + bool copyBeforeWrite = + options.copyBeforeWrite || notAnalyzedFuncOps.contains(funcOp); + if (failed(bufferizeOp(funcOp, options, copyBeforeWrite, /*opFilter=*/nullptr, statistics))) return failure(); // Change buffer return types to more precise layout maps. @@ -453,13 +464,17 @@ "expected that function boundary bufferization is activated"); assert(!(options.copyBeforeWrite && options.testAnalysisOnly) && "invalid combination of bufferization flags"); + llvm::SmallSet notAnalyzedFuncOps; + if (!options.copyBeforeWrite) { - if (failed(insertTensorCopies(moduleOp, options, statistics))) + if (failed(insertTensorCopies(moduleOp, options, statistics, + ¬AnalyzedFuncOps))) return failure(); } if (options.testAnalysisOnly) - return success(); - if (failed(bufferizeModuleOp(moduleOp, options, statistics))) + return success(notAnalyzedFuncOps.empty()); + if (failed( + bufferizeModuleOp(moduleOp, options, notAnalyzedFuncOps, statistics))) return failure(); return success(); } diff --git a/mlir/lib/Dialect/Bufferization/Transforms/TensorCopyInsertion.cpp b/mlir/lib/Dialect/Bufferization/Transforms/TensorCopyInsertion.cpp --- a/mlir/lib/Dialect/Bufferization/Transforms/TensorCopyInsertion.cpp +++ b/mlir/lib/Dialect/Bufferization/Transforms/TensorCopyInsertion.cpp @@ -98,7 +98,8 @@ LogicalResult mlir::bufferization::insertTensorCopies( Operation *op, const OneShotBufferizationOptions &options, - BufferizationStatistics *statistics) { + BufferizationStatistics *statistics, + llvm::SmallSet *notAnalyzedFuncOps) { // Preprocessing: Resolve currently unsupported bufferization cases. resolveUsesInRepetitiveRegions(op, options); @@ -107,8 +108,21 @@ // analysis depending on whether function boundary bufferization is enabled or // not. if (options.bufferizeFunctionBoundaries) { - if (failed(analyzeModuleOp(cast(op), state, statistics))) - return failure(); + ModuleOp module = cast(op); + if (failed( + analyzeModuleOp(module, state, statistics, notAnalyzedFuncOps))) { + if (options.testAnalysisOnly) + return failure(); + // If some of the FuncOps were not analyzed, run insertTensorCopies() only + // on the analyzed ones. Otherwise, it will run on the entire module. + if (!notAnalyzedFuncOps->empty()) { + for (auto funcOp : module.getOps()) + if (!notAnalyzedFuncOps->contains(funcOp) && + failed(insertTensorCopies(funcOp, state))) + return failure(); + return success(); + } + } } else { if (failed(analyzeOp(op, state, statistics))) return failure(); diff --git a/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-invalid.mlir b/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-invalid.mlir new file mode 100644 --- /dev/null +++ b/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-invalid.mlir @@ -0,0 +1,20 @@ +// RUN: mlir-opt %s -allow-unregistered-dialect -one-shot-bufferize -split-input-file -verify-diagnostics + +func.func @to_memref_op_is_writing( + %t1: tensor {bufferization.writable = true}, %idx1: index, + %idx2: index, %idx3: index, %v1: vector<5xf32>) -> (vector<5xf32>, vector<5xf32>) { + // This is a RaW conflict because to_memref is an inplace write and %t1 is + // read further down. This will likely have to change with partial + // bufferization. + + // expected-error @+1 {{to_memref ops not supported during One-Shot Analysis}} + %0 = bufferization.to_memref %t1 : memref + + // Read from both. + %cst = arith.constant 0.0 : f32 + %r1 = vector.transfer_read %t1[%idx3], %cst : tensor, vector<5xf32> + %r2 = vector.transfer_read %0[%idx3], %cst : memref, vector<5xf32> + + return %r1, %r2 : vector<5xf32>, vector<5xf32> +} + diff --git a/mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize-analysis-invalid.mlir b/mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize-analysis-invalid.mlir new file mode 100644 --- /dev/null +++ b/mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize-analysis-invalid.mlir @@ -0,0 +1,19 @@ +// RUN: mlir-opt %s -allow-unregistered-dialect -one-shot-bufferize="bufferize-function-boundaries=1 test-analysis-only=1" -split-input-file -verify-diagnostics + +func.func @to_memref_op_is_writing( + %t1: tensor {bufferization.writable = true}, %idx1: index, + %idx2: index, %idx3: index, %v1: vector<5xf32>) -> (vector<5xf32>, vector<5xf32>) { + // This is a RaW conflict because to_memref is an inplace write and %t1 is + // read further down. This will likely have to change with partial + // bufferization. + + // expected-error @+1 {{to_memref ops not supported during One-Shot Analysis}} + %0 = bufferization.to_memref %t1 : memref + + // Read from both. + %cst = arith.constant 0.0 : f32 + %r1 = vector.transfer_read %t1[%idx3], %cst : tensor, vector<5xf32> + %r2 = vector.transfer_read %0[%idx3], %cst : memref, vector<5xf32> + + return %r1, %r2 : vector<5xf32>, vector<5xf32> +} diff --git a/mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize-force-copy-before-write.mlir b/mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize-force-copy-before-write.mlir new file mode 100644 --- /dev/null +++ b/mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize-force-copy-before-write.mlir @@ -0,0 +1,39 @@ +// RUN: mlir-opt %s -one-shot-bufferize="bufferize-function-boundaries=1" -drop-equivalent-buffer-results --split-input-file | FileCheck %s + +// ToMemref ops do not pass analysis step. CopyBeforeWrite will be true for the +// FuncOp containing ToMemref, but not for the other FuncOps. + +module { + // CHECK-LABEL: func.func @foo( + // CHECK-SAME: %[[arg0:.*]]: memref>) { + func.func @foo(%arg0: tensor) -> tensor { + // CHECK-NEXT: %[[c0:.*]] = arith.constant 0 : index + %cst = arith.constant 1.000000e+00 : f32 + + // CHECK-NEXT: %[[c1:.*]] = arith.constant 1.000000e+00 : f32 + %c0 = arith.constant 0 : index + + // CHECK-NEXT: memref.store %[[c1]], %[[arg0]]{{\[}}%[[c0]]] : memref> + %inserted = tensor.insert %cst into %arg0[%c0] : tensor + + return %inserted : tensor + } + + // CHECK-LABEL: func.func @contains_to_memref_op( + // CHECK-SAME: %[[arg0:.*]]: memref>, + // CHECK-SAME: %[[arg1:.*]]: index) -> vector<5xf32> { + func.func @contains_to_memref_op(%arg0: tensor {bufferization.writable = true}, %arg1: index) -> vector<5xf32> { + + %0 = bufferization.to_memref %arg0 : memref + + // CHECK: %[[c0:.*]] = arith.constant 0 : index + %cst = arith.constant 0.000000e+00 : f32 + + // CHECK: %[[dim:.*]] = memref.dim %[[arg0]], %[[c0]] : memref> + // CHECK: %[[alloc:.*]] = memref.alloc(%[[dim]]) : memref + // CHECK: memref.copy %[[arg0]], %[[alloc]] : memref> to memref + // CHECK: vector.transfer_read + %1 = vector.transfer_read %0[%arg1], %cst : memref, vector<5xf32> + return %1 : vector<5xf32> + } +} diff --git a/mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize-invalid.mlir b/mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize-invalid.mlir --- a/mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize-invalid.mlir +++ b/mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize-invalid.mlir @@ -211,7 +211,7 @@ func.func @unknown_op(%A : tensor<4xf32>) -> tensor<4xf32> { - // expected-error: @+1 {{op was not bufferized}} + // expected-error @+1 {{op was not bufferized}} %r = "marklar"(%A) : (tensor<4xf32>) -> (tensor<4xf32>) // expected-error @+1 {{operand #0 may return/yield a new buffer allocation}} return %r: tensor<4xf32> @@ -242,26 +242,6 @@ // ----- -func.func @to_memref_op_is_writing( - %t1: tensor {bufferization.writable = true}, %idx1: index, - %idx2: index, %idx3: index, %v1: vector<5xf32>) -> (vector<5xf32>, vector<5xf32>) { - // This is a RaW conflict because to_memref is an inplace write and %t1 is - // read further down. This will likely have to change with partial - // bufferization. - - // expected-error @+1 {{to_memref ops not supported during One-Shot Analysis}} - %0 = bufferization.to_memref %t1 : memref - - // Read from both. - %cst = arith.constant 0.0 : f32 - %r1 = vector.transfer_read %t1[%idx3], %cst : tensor, vector<5xf32> - %r2 = vector.transfer_read %0[%idx3], %cst : memref, vector<5xf32> - - return %r1, %r2 : vector<5xf32>, vector<5xf32> -} - -// ----- - // expected-error @+2 {{failed to bufferize op}} // expected-error @+1 {{cannot bufferize bodiless function that returns a tensor}} func.func private @foo(%t : tensor) -> (f32, tensor, f32)