maxjhandsome
/
llvm-project
forked from OSchip/llvm-project
1
0
Fork 0
Code Issues Pull Requests Packages Releases Wiki Activity

[mlir][bufferize] Add a way for ops to fail the analysis 

Add `BufferizableOpInterface::verifyAnalysis`. Ops can implement this method to check for expected invariants and limitations.

The purpose of this change is to introduce a modular way of checking assertions such as `assertScfForAliasingProperties`.

Differential Revision: https://reviews.llvm.org/D120189
This commit is contained in:
Matthias Springer 2022-02-20 05:49:33 +09:00
parent 24bfa24355
commit 4ec00fb3ea
8 changed files with 65 additions and 83 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

17
mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.td
View File

@ -290,6 +290,23 @@ def BufferizableOpInterface : OpInterface<"BufferizableOpInterface"> {
/*defaultImplementation=*/[{
return false;
}]
>,
InterfaceMethod<
/*desc=*/[{
Return `failure` if this op does not pass the analysis. This method
is run during One-Shot Bufferize (after all post-analysis steps). If
the op does not pass the analysis, bufferization is aborted.
This method can be used to check expected invariants and limitations
of the current bufferization implementation.
}],
/*retType=*/"LogicalResult",
/*methodName=*/"verifyAnalysis",
/*args=*/(ins "const BufferizationState &":$state),
/*methodBody=*/"",
/*defaultImplementation=*/[{
return success();
}]
>
];

10
mlir/include/mlir/Dialect/SCF/BufferizableOpInterfaceImpl.h
View File

@ -20,16 +20,6 @@ class BufferizationAliasInfo;
} // namespace bufferization
namespace scf {
/// Assert that yielded values of an scf.for op are aliasing their corresponding
/// bbArgs. This is required because the i-th OpResult of an scf.for op is
/// currently assumed to alias with the i-th iter_arg (in the absence of
/// conflicts).
LogicalResult
assertScfForAliasingProperties(Operation *op,
bufferization::BufferizationState &state,
bufferization::BufferizationAliasInfo &aliasInfo,
SmallVector<Operation *> &newOps);
void registerBufferizableOpInterfaceExternalModels(DialectRegistry &registry);
} // namespace scf
} // namespace mlir

13
mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp
View File

@ -778,6 +778,19 @@ LogicalResult bufferization::analyzeOp(Operation *op,
return failure();
}
// Analysis verification: After setting up alias/equivalence sets, each op
// can check for expected invariants/limitations and fail the analysis if
// necessary.
bool passedAnalysis = true;
op->walk([&](Operation *op) {
if (BufferizableOpInterface bufferizableOp =
options.dynCastBufferizableOp(op))
if (failed(bufferizableOp.verifyAnalysis(state)))
passedAnalysis = false;
});
if (!passedAnalysis)
return failure();
// Annotate operations if we only want to report the analysis.
if (options.testAnalysisOnly)
annotateOpsWithBufferizationMarkers(op, aliasInfo, state);

3
mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferizePass.cpp
View File

@ -105,9 +105,6 @@ void LinalgComprehensiveModuleBufferize::runOnOperation() {
opt = *options;
}
// Only certain scf.for ops are supported by the analysis.
opt.addPostAnalysisStep(scf::assertScfForAliasingProperties);
ModuleOp moduleOp = getOperation();
applyEnablingTransformations(moduleOp);

66
mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp
View File

@ -385,6 +385,37 @@ struct ForOpInterface
return success();
}
/// Assert that yielded values of an scf.for op are aliasing with their
/// corresponding bbArgs. This is required because the i-th OpResult of an
/// scf.for op is currently assumed to alias with the i-th iter_arg (in the
/// absence of conflicts).
LogicalResult verifyAnalysis(Operation *op,
const BufferizationState &state) const {
auto forOp = cast<scf::ForOp>(op);
auto yieldOp =
cast<scf::YieldOp>(forOp.getLoopBody().front().getTerminator());
for (OpOperand &operand : yieldOp->getOpOperands()) {
auto tensorType = operand.get().getType().dyn_cast<TensorType>();
if (!tensorType)
continue;
OpOperand &forOperand = forOp.getOpOperandForResult(
forOp->getResult(operand.getOperandNumber()));
auto bbArg = forOp.getRegionIterArgForOpOperand(forOperand);
// Note: This is overly strict. We should check for aliasing bufferized
// values. But we don't have a "must-alias" analysis yet.
if (!state.areEquivalentBufferizedValues(operand.get(), bbArg))
// TODO: this could get resolved with copies but it can also turn into
// swaps so we need to be careful about order of copies.
return yieldOp->emitError()
<< "Yield operand #" << operand.getOperandNumber()
<< " does not bufferize to a buffer that is aliasing the "
"matching"
<< " enclosing scf::for operand";
}
return success();
}
};
/// Bufferization of scf.yield. Bufferized as part of their enclosing ops, so
@ -434,41 +465,6 @@ struct YieldOpInterface
} // namespace scf
} // namespace mlir
LogicalResult mlir::scf::assertScfForAliasingProperties(
Operation *op, BufferizationState &state, BufferizationAliasInfo &aliasInfo,
SmallVector<Operation *> &newOps) {
LogicalResult status = success();
op->walk([&](scf::ForOp forOp) {
auto yieldOp =
cast<scf::YieldOp>(forOp.getLoopBody().front().getTerminator());
for (OpOperand &operand : yieldOp->getOpOperands()) {
auto tensorType = operand.get().getType().dyn_cast<TensorType>();
if (!tensorType)
continue;
OpOperand &forOperand = forOp.getOpOperandForResult(
forOp->getResult(operand.getOperandNumber()));
auto bbArg = forOp.getRegionIterArgForOpOperand(forOperand);
// Note: This is overly strict. We should check for aliasing bufferized
// values. But we don't have a "must-alias" analysis yet.
if (!aliasInfo.areEquivalentBufferizedValues(operand.get(), bbArg)) {
// TODO: this could get resolved with copies but it can also turn into
// swaps so we need to be careful about order of copies.
status =
yieldOp->emitError()
<< "Yield operand #" << operand.getOperandNumber()
<< " does not bufferize to a buffer that is aliasing the matching"
<< " enclosing scf::for operand";
return WalkResult::interrupt();
}
}
return WalkResult::advance();
});
return status;
}
void mlir::scf::registerBufferizableOpInterfaceExternalModels(
DialectRegistry &registry) {
registry.addOpInterface<ExecuteRegionOp, ExecuteRegionOpInterface>();

6
mlir/test/Dialect/Linalg/comprehensive-module-bufferize-invalid.mlir
View File

@ -87,7 +87,7 @@ func @scf_for(%A : tensor<?xf32>,
%B : tensor<?xf32> {linalg.inplaceable = true},
%C : tensor<4xf32>,
%lb : index, %ub : index, %step : index)
-> (tensor<?xf32>, tensor<?xf32>)
-> (f32, f32)
{
%r0:2 = scf.for %i = %lb to %ub step %step iter_args(%tA = %A, %tB = %B)
-> (tensor<?xf32>, tensor<?xf32>)
@ -102,7 +102,9 @@ func @scf_for(%A : tensor<?xf32>,
scf.yield %ttB, %ttA : tensor<?xf32>, tensor<?xf32>
}
return %r0#0, %r0#1: tensor<?xf32>, tensor<?xf32>
%f0 = tensor.extract %r0#0[%step] : tensor<?xf32>
%f1 = tensor.extract %r0#1[%step] : tensor<?xf32>
return %f0, %f1: f32, f32
}
// -----

29
mlir/test/Dialect/Linalg/comprehensive-module-bufferize-partial.mlir
View File

@ -179,35 +179,6 @@ func @simple_tensor_test(%t1 : tensor<?xf32>, %f : f32) -> tensor<?xf32> {
// -----
// CHECK-SCF-LABEL: func @simple_scf_for(
// CHECK-SCF-SAME: %[[t1:.*]]: tensor<?xf32>
func @simple_scf_for(
%t1: tensor<?xf32>, %sz: index, %step: index, %f: f32) -> tensor<?xf32> {
%c0 = arith.constant 0 : index
// CHECK-SCF: %[[t1_memref:.*]] = bufferization.to_memref %[[t1]]
// CHECK-SCF: %[[alloc:.*]] = memref.alloc
// CHECK-SCF: %[[casted:.*]] = memref.cast %[[alloc]]
// CHECK-SCF: memref.copy %[[t1_memref]], %[[alloc]]
// CHECK-SCF: %[[scf_for:.*]] = scf.for %[[iv:.*]] = %{{.*}} to %{{.*}} step %{{.*}} iter_args(%[[arg0:.*]] = %[[casted]]) -> ({{.*}}) {
%0 = scf.for %iv = %c0 to %sz step %step iter_args(%arg0 = %t1) -> tensor<?xf32> {
// CHECK-SCF: %[[arg0_tensor:.*]] = bufferization.to_tensor %[[arg0]]
// CHECK-SCF: %[[insert:.*]] = tensor.insert %{{.*}} into %[[arg0_tensor]]
%1 = tensor.insert %f into %arg0[%iv] : tensor<?xf32>
// CHECK-SCF: %[[insert_memref:.*]] = bufferization.to_memref %[[insert]]
// CHECK-SCF: scf.yield %[[insert_memref]]
scf.yield %1 : tensor<?xf32>
}
// CHECK-SCF: }
// CHECK-SCF: %[[scf_for_tensor:.*]] = bufferization.to_tensor %[[scf_for]]
// CHECK-SCF: return %[[scf_for_tensor]]
return %0 : tensor<?xf32>
}
// -----
// CHECK-SCF-LABEL: func @simple_scf_if(
// CHECK-SCF-SAME: %[[t1:.*]]: tensor<?xf32> {linalg.inplaceable = true}, %[[c:.*]]: i1, %[[pos:.*]]: index
func @simple_scf_if(%t1: tensor<?xf32> {linalg.inplaceable = true}, %c: i1, %pos: index, %f: f32)

4
mlir/test/lib/Dialect/Linalg/TestComprehensiveBufferize.cpp
View File

@ -102,10 +102,6 @@ struct TestComprehensiveFunctionBufferize
void TestComprehensiveFunctionBufferize::runOnOperation() {
auto options = std::make_unique<AnalysisBufferizationOptions>();
if (!allowReturnMemref)
options->addPostAnalysisStep(scf::assertScfForAliasingProperties);
options->allowReturnMemref = allowReturnMemref;
options->allowUnknownOps = allowUnknownOps;
options->testAnalysisOnly = testAnalysisOnly;