mirror of
https://github.com/capstone-engine/llvm-capstone.git
synced 2025-01-10 10:01:42 +00:00
[mlir][bufferize] Better user control of layout maps
This changes replaces the `fully-dynamic-layout-maps` options (which was badly named) with two new options: * `unknown-type-conversion` controls the layout maps on buffer types for which no layout map can be inferred. * `function-boundary-type-conversion` controls the layout maps on buffer types inside of function signatures. Differential Revision: https://reviews.llvm.org/D125615
This commit is contained in:
Matthias Springer
parent
6e23cd2bf0
commit
f287da8a15
@ -336,16 +336,29 @@ simpler memref type (e.g., identity layout map), we expect that canonicalization
|
||||
patterns would clean up unnecessarily dynamic layout maps. (Some of these
|
||||
canonicalization patterns may not be implemented yet.)
|
||||
|
||||
Note that One-Shot Bufferize always generates the most specific memref type when
|
||||
the entire IR is bufferizable. In that case, we do not have to rely on
|
||||
canonicalization patterns to clean up the bufferized IR.
|
||||
One-Shot Bufferize tries to infer the most precise memref type when bufferizing
|
||||
an op. If the entire IR is bufferizable, we do not have to resort to
|
||||
conservatively use fully dynamic layout maps. In that case, we also do not have
|
||||
to rely on canonicalization patterns to clean up the bufferized IR.
|
||||
|
||||
One-Shot Bufferize can be configured to always generate memref types with
|
||||
identity layout when the exact target memref type is not known via
|
||||
`fully-dynamic-layout-maps=0`. This can be useful for legacy code that cannot
|
||||
handle memref types with layout maps. Note that this leads to additional buffer
|
||||
copies when folding a `to_tensor`/`to_memref` pair with memref types that are
|
||||
not cast-compatible.
|
||||
Note: There are some bufferizable ops for which a percise layout map cannot be
|
||||
inferred. E.g., a `tensor.cast` from a `tensor<*xf32>` to a `tensor<?x?xf32>`
|
||||
must be bufferized to a `memref.cast` with a memref type that has a fully
|
||||
dynamic layout map.
|
||||
|
||||
One-Shot Bufferize has an option `unknown-type-conversion` to control the
|
||||
generation of layout maps when no precise layout can be inferred:
|
||||
|
||||
* `fully-dynamic-layout-map` uses fully dynamic layout maps and is the default
|
||||
behavior. This composes well when IR is partially bufferized.
|
||||
* `identity-layout-map` uses static identity layout maps. This option can be
|
||||
useful for legacy code that cannot handle memref types with layout maps.
|
||||
Note that this setting can lead to additional buffer copies when folding a
|
||||
`to_tensor`/`to_memref` pair with memref types that are not cast-compatible.
|
||||
|
||||
Note: The `unknown-type-conversion` option does not affect layout maps of
|
||||
function signatures. There is a separate `function-signature-type-conversion`
|
||||
option that controls layout maps of function parameters and function results.
|
||||
|
||||
## Extending One-Shot Bufferize
|
||||
|
||||
|
||||
@ -62,6 +62,12 @@ struct BufferizationOptions {
|
||||
FilterType type;
|
||||
};
|
||||
|
||||
enum class LayoutMapOption : int8_t {
|
||||
InferLayoutMap = 0,
|
||||
IdentityLayoutMap = 1,
|
||||
FullyDynamicLayoutMap = 2
|
||||
};
|
||||
|
||||
BufferizationOptions();
|
||||
|
||||
/// Return `true` if the filter has at least one ALLOW rule.
|
||||
@ -201,6 +207,45 @@ struct BufferizationOptions {
|
||||
/// bufferized or not.
|
||||
bool bufferizeFunctionBoundaries = false;
|
||||
|
||||
/// This flag controls buffer types on function signatures.
|
||||
///
|
||||
/// * InferLayoutMap: All function parameter types have a fully dynamic layout
|
||||
/// map, but function result types are inferred from the body of the
|
||||
/// function.
|
||||
/// * FullyDynamicLayoutMap: All function parameter types and result types
|
||||
/// have a fully dynamic layout map. This option is most efficient because
|
||||
/// any layout map can be casted to a fully dynamic one.
|
||||
/// * IdentityLayoutMap: All function parameter types and result types have a
|
||||
/// static identity layout (i.e., no layout map). This option may introduce
|
||||
/// additional buffer allocs and copies because layout maps cannot be casted
|
||||
/// away.
|
||||
///
|
||||
/// If `bufferizeFunctionBoundaries` is not set, this flag has no effect. If
|
||||
/// `promoteBufferResultsToOutParams` is set, `kInferMostPreciseLayoutMap` is
|
||||
/// is an invalid option.
|
||||
///
|
||||
/// Note: Inferred layout maps may not be desireable when interacting with
|
||||
/// external functions, because the generated function signatures will be less
|
||||
/// predictable.
|
||||
LayoutMapOption functionBoundaryTypeConversion =
|
||||
LayoutMapOption::InferLayoutMap;
|
||||
|
||||
/// This flag controls buffer types on unknown ops (to_memref wrappers) and in
|
||||
/// other cases where a precise memref type cannot be inferred (e.g., the
|
||||
/// bufferization of "tensor.cast").
|
||||
///
|
||||
/// * InferLayoutMap: This option is invalid and cannot be used.
|
||||
/// * FullyDynamicLayoutMap: Assume that unknown ops have results with fully
|
||||
/// dynamic layout maps after bufferization. This option is most efficient
|
||||
/// because any layout map can be casted to a fully dynamic one.
|
||||
/// * IdentityLayoutMap: Assume that unknown ops have results with static
|
||||
/// identity layout (i.e., no layout map) after bufferization. This option
|
||||
/// introduces additional buffer allocs and copies if the unknown op is
|
||||
/// eventually bufferized to an op that returns a buffer with non-identity
|
||||
/// layout.
|
||||
LayoutMapOption unknownTypeConversion =
|
||||
LayoutMapOption::FullyDynamicLayoutMap;
|
||||
|
||||
/// Specifies whether dealloc ops should be generated along with alloc ops. If
|
||||
/// not, new memory allocations will leak.
|
||||
bool createDeallocs = true;
|
||||
@ -209,10 +254,6 @@ struct BufferizationOptions {
|
||||
/// Should be used only with `testAnalysisOnly = true`.
|
||||
unsigned analysisFuzzerSeed = 0;
|
||||
|
||||
/// Specifies whether fully dynamic layout maps should be used on ranked
|
||||
/// MemRef types. If false, MemRef types will have no layout maps.
|
||||
bool fullyDynamicLayoutMaps = true;
|
||||
|
||||
/// 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;
|
||||
@ -554,9 +595,9 @@ OpTy replaceOpWithNewBufferizedOp(RewriterBase &rewriter, Operation *op,
|
||||
/// If possible, op bufferization implementations should not use this function
|
||||
/// and instead infer precise memref types for tensor results by themselves.
|
||||
///
|
||||
/// Unless a layout map was specified, `options` flags determine what kind of
|
||||
/// layout map will be used. For best composability (without copies), the fully
|
||||
/// dynamic layout map is used by default.
|
||||
/// Unless a layout map was specified, `options.unknownTypeConverter` determines
|
||||
/// what kind of layout map will be used. For best composability (without
|
||||
/// copies), the fully dynamic layout map is used by default.
|
||||
///
|
||||
/// Note: Canonicalization patterns could clean up layout maps and infer more
|
||||
/// precise layout maps after bufferization. However, many possible
|
||||
@ -566,6 +607,17 @@ BaseMemRefType getMemRefType(TensorType tensorType,
|
||||
MemRefLayoutAttrInterface layout = {},
|
||||
Attribute memorySpace = {});
|
||||
|
||||
/// Return a MemRef type with fully dynamic layout. If the given tensor type
|
||||
/// is unranked, return an unranked MemRef type.
|
||||
BaseMemRefType getMemRefTypeWithFullyDynamicLayout(TensorType tensorType,
|
||||
Attribute memorySpace = {});
|
||||
|
||||
/// Return a MemRef type with a static identity layout (i.e., no layout map). If
|
||||
/// the given tensor type is unranked, return an unranked MemRef type.
|
||||
BaseMemRefType
|
||||
getMemRefTypeWithStaticIdentityLayout(TensorType tensorType,
|
||||
Attribute memorySpace = {});
|
||||
|
||||
/// Try to hoist all new buffer allocations until the next hoisting barrier.
|
||||
LogicalResult hoistBufferAllocations(Operation *op,
|
||||
const BufferizationOptions &options);
|
||||
|
||||
@ -196,6 +196,19 @@ def OneShotBufferize : Pass<"one-shot-bufferize", "ModuleOp"> {
|
||||
when encountering an op that is not included in the filter (even if it is
|
||||
bufferizable).
|
||||
|
||||
One-Shot Bufferize will by default assume memref types with fully dynamic
|
||||
layout maps when a precise layout cannot be inferred. E.g., this is the case
|
||||
when wrapping a non-bufferizable op in to_memref/to_tensor ops. This
|
||||
behavior can be overridden with `unknown-type-conversion`. Valid values are
|
||||
`fully-dynamic-layout-map` and `identity-layout-map`.
|
||||
|
||||
Layout maps on function signatures can be controlled with a separate
|
||||
`function-boundary-type-conversion` option, which can be set to
|
||||
`infer-layout-map` in addition to the two possible values mentioned above.
|
||||
When layout maps are referred, function return types may be more precise.
|
||||
Function argument types cannot be inferred and have fully dynamic layout
|
||||
maps in that case.
|
||||
|
||||
For testing/debugging purposes, `test-analysis-only=1 print-conflicts=1`
|
||||
prints analysis results and explains why an OpOperand was decided to
|
||||
bufferize out-of-place. This is useful for understanding why One-Shot
|
||||
@ -254,9 +267,10 @@ def OneShotBufferize : Pass<"one-shot-bufferize", "ModuleOp"> {
|
||||
"core bufferization passes.">,
|
||||
ListOption<"dialectFilter", "dialect-filter", "std::string",
|
||||
"Restrict bufferization to ops from these dialects.">,
|
||||
Option<"fullyDynamicLayoutMaps", "fully-dynamic-layout-maps", "bool",
|
||||
/*default=*/"true",
|
||||
"Generate MemRef types with dynamic offset+strides by default.">,
|
||||
Option<"functionBoundaryTypeConversion",
|
||||
"function-boundary-type-conversion", "std::string",
|
||||
/*default=*/"\"infer-layout-map\"",
|
||||
"Controls layout maps when bufferizing function signatures.">,
|
||||
Option<"testAnalysisOnly", "test-analysis-only", "bool",
|
||||
/*default=*/"false",
|
||||
"Test only: Only run inplaceability analysis and annotate IR">,
|
||||
@ -267,6 +281,9 @@ def OneShotBufferize : Pass<"one-shot-bufferize", "ModuleOp"> {
|
||||
Option<"promoteBufferResultsToOutParams",
|
||||
"promote-buffer-results-to-out-params", "bool", /*default=*/"false",
|
||||
"Replace returned buffers (that were not dropped) with out params.">,
|
||||
Option<"unknownTypeConversion", "unknown-type-conversion", "std::string",
|
||||
/*default=*/"\"fully-dynamic-layout-map\"",
|
||||
"Controls layout maps for non-inferrable memref types.">,
|
||||
];
|
||||
let constructor = "mlir::bufferization::createOneShotBufferizePass()";
|
||||
}
|
||||
|
||||
@ -662,19 +662,38 @@ BaseMemRefType bufferization::getMemRefType(TensorType tensorType,
|
||||
memorySpace);
|
||||
}
|
||||
|
||||
// Case 2: Ranked memref type with specified layout. If fully dynamic layout
|
||||
// maps are not requested, generate a type with `layout`, which is empty (no
|
||||
// layout map) by default.
|
||||
// Case 2: Ranked memref type with specified layout.
|
||||
auto rankedTensorType = tensorType.cast<RankedTensorType>();
|
||||
if (layout || !options.fullyDynamicLayoutMaps) {
|
||||
if (layout) {
|
||||
return MemRefType::get(rankedTensorType.getShape(),
|
||||
rankedTensorType.getElementType(), layout,
|
||||
memorySpace);
|
||||
}
|
||||
|
||||
// Case 3: Ranked memref type with unspecified layout. Choose the most dynamic
|
||||
// one.
|
||||
// TODO: address space decisions to connect with the actual alloc.
|
||||
// Case 3: Configured with "fully dynamic layout maps".
|
||||
if (options.unknownTypeConversion ==
|
||||
BufferizationOptions::LayoutMapOption::FullyDynamicLayoutMap)
|
||||
return getMemRefTypeWithFullyDynamicLayout(tensorType, memorySpace);
|
||||
|
||||
// Case 4: Configured with "static identity layout maps".
|
||||
if (options.unknownTypeConversion ==
|
||||
BufferizationOptions::LayoutMapOption::IdentityLayoutMap)
|
||||
return getMemRefTypeWithStaticIdentityLayout(tensorType, memorySpace);
|
||||
|
||||
llvm_unreachable("InferLayoutMap is an invalid option");
|
||||
}
|
||||
|
||||
BaseMemRefType
|
||||
bufferization::getMemRefTypeWithFullyDynamicLayout(TensorType tensorType,
|
||||
Attribute memorySpace) {
|
||||
// Case 1: Unranked memref type.
|
||||
if (auto unrankedTensorType = tensorType.dyn_cast<UnrankedTensorType>()) {
|
||||
return UnrankedMemRefType::get(unrankedTensorType.getElementType(),
|
||||
memorySpace);
|
||||
}
|
||||
|
||||
// Case 2: Ranked memref type.
|
||||
auto rankedTensorType = tensorType.cast<RankedTensorType>();
|
||||
int64_t dynamicOffset = ShapedType::kDynamicStrideOrOffset;
|
||||
SmallVector<int64_t> dynamicStrides(rankedTensorType.getRank(),
|
||||
ShapedType::kDynamicStrideOrOffset);
|
||||
@ -684,3 +703,22 @@ BaseMemRefType bufferization::getMemRefType(TensorType tensorType,
|
||||
rankedTensorType.getElementType(), stridedLayout,
|
||||
memorySpace);
|
||||
}
|
||||
|
||||
/// Return a MemRef type with a static identity layout (i.e., no layout map). If
|
||||
/// the given tensor type is unranked, return an unranked MemRef type.
|
||||
BaseMemRefType
|
||||
bufferization::getMemRefTypeWithStaticIdentityLayout(TensorType tensorType,
|
||||
Attribute memorySpace) {
|
||||
// Case 1: Unranked memref type.
|
||||
if (auto unrankedTensorType = tensorType.dyn_cast<UnrankedTensorType>()) {
|
||||
return UnrankedMemRefType::get(unrankedTensorType.getElementType(),
|
||||
memorySpace);
|
||||
}
|
||||
|
||||
// Case 2: Ranked memref type.
|
||||
auto rankedTensorType = tensorType.cast<RankedTensorType>();
|
||||
MemRefLayoutAttrInterface layout = {};
|
||||
return MemRefType::get(rankedTensorType.getShape(),
|
||||
rankedTensorType.getElementType(), layout,
|
||||
memorySpace);
|
||||
}
|
||||
|
||||
@ -151,6 +151,17 @@ struct FinalizingBufferizePass
|
||||
}
|
||||
};
|
||||
|
||||
static BufferizationOptions::LayoutMapOption
|
||||
parseLayoutMapOption(std::string s) {
|
||||
if (s == "fully-dynamic-layout-map")
|
||||
return BufferizationOptions::LayoutMapOption::FullyDynamicLayoutMap;
|
||||
if (s == "identity-layout-map")
|
||||
return BufferizationOptions::LayoutMapOption::IdentityLayoutMap;
|
||||
if (s == "infer-layout-map")
|
||||
return BufferizationOptions::LayoutMapOption::InferLayoutMap;
|
||||
llvm_unreachable("invalid layout map option");
|
||||
}
|
||||
|
||||
struct OneShotBufferizePass
|
||||
: public OneShotBufferizeBase<OneShotBufferizePass> {
|
||||
OneShotBufferizePass() : OneShotBufferizeBase<OneShotBufferizePass>() {}
|
||||
@ -175,11 +186,13 @@ struct OneShotBufferizePass
|
||||
opt.alwaysAliasingWithDest = alwaysAliasingWithDest;
|
||||
opt.analysisFuzzerSeed = analysisFuzzerSeed;
|
||||
opt.createDeallocs = createDeallocs;
|
||||
opt.fullyDynamicLayoutMaps = fullyDynamicLayoutMaps;
|
||||
opt.functionBoundaryTypeConversion =
|
||||
parseLayoutMapOption(functionBoundaryTypeConversion);
|
||||
opt.printConflicts = printConflicts;
|
||||
opt.testAnalysisOnly = testAnalysisOnly;
|
||||
opt.bufferizeFunctionBoundaries = bufferizeFunctionBoundaries;
|
||||
opt.promoteBufferResultsToOutParams = promoteBufferResultsToOutParams;
|
||||
opt.unknownTypeConversion = parseLayoutMapOption(unknownTypeConversion);
|
||||
|
||||
BufferizationOptions::OpFilterEntry::FilterFn filterFn =
|
||||
[&](Operation *op) {
|
||||
@ -362,6 +375,9 @@ LogicalResult
|
||||
bufferization::bufferizeOp(Operation *op,
|
||||
BufferizationState &bufferizationState) {
|
||||
const auto &options = bufferizationState.getOptions();
|
||||
assert(options.unknownTypeConversion !=
|
||||
BufferizationOptions::LayoutMapOption::InferLayoutMap &&
|
||||
"invalid layout map option");
|
||||
|
||||
// Keep track of to_memref ops.
|
||||
DenseSet<Operation *> toMemrefOps;
|
||||
@ -371,13 +387,9 @@ bufferization::bufferizeOp(Operation *op,
|
||||
//
|
||||
// We should ideally know the exact memref type of all operands when
|
||||
// bufferizing an op. (This is the case when bufferizing top-to-bottom.)
|
||||
// Otherwise, we have to use a memref type with a fully dynamic layout map,
|
||||
// which has to canonicalize away. This is less efficient.
|
||||
//
|
||||
// If "fullyDynamicLayoutMaps = false", we would have to insert buffer copies
|
||||
// to fold ("finalize") to_memref(to_tensor(x)) ops with non-cast-compatible
|
||||
// layout maps when doing a traversal other than top-to-bottom. These would
|
||||
// not easily fold away.
|
||||
// Otherwise, we have to use a memref type with a fully dynamic layout map to
|
||||
// avoid copies. We are currently missing patterns for layout maps to
|
||||
// canonicalize away (or canonicalize to more precise layouts).
|
||||
SmallVector<Operation *> worklist;
|
||||
op->walk<WalkOrder::PreOrder>([&](Operation *op) {
|
||||
if (hasTensorSemantics(op))
|
||||
@ -478,6 +490,7 @@ BufferizationOptions bufferization::getPartialBufferizationOptions() {
|
||||
BufferizationOptions options;
|
||||
options.allowUnknownOps = true;
|
||||
options.createDeallocs = false;
|
||||
options.fullyDynamicLayoutMaps = false;
|
||||
options.unknownTypeConversion =
|
||||
BufferizationOptions::LayoutMapOption::IdentityLayoutMap;
|
||||
return options;
|
||||
}
|
||||
|
||||
@ -58,15 +58,24 @@ static func::ReturnOp getAssumedUniqueReturnOp(FuncOp funcOp) {
|
||||
|
||||
/// Return the index-th bufferized function argument type. This assumes that the
|
||||
/// specified argument is a tensor. If the tensor is ranked, a layout map may be
|
||||
/// specified by the user. If no layout map is specified, a fully dynamic map is
|
||||
/// used.
|
||||
/// specified by the user. If no layout map is specified, the default layout map
|
||||
/// (as per `options.functionBoundaryTypeConversion`) is used.
|
||||
static BaseMemRefType
|
||||
getBufferizedFunctionArgType(FuncOp funcOp, int64_t index,
|
||||
const BufferizationOptions &options) {
|
||||
auto tensorType =
|
||||
funcOp.getFunctionType().getInput(index).dyn_cast<TensorType>();
|
||||
assert(tensorType && "expected TensorType");
|
||||
BaseMemRefType memrefType = getMemRefType(tensorType, options);
|
||||
|
||||
BaseMemRefType memrefType;
|
||||
if (options.functionBoundaryTypeConversion ==
|
||||
BufferizationOptions::LayoutMapOption::IdentityLayoutMap) {
|
||||
memrefType = getMemRefTypeWithStaticIdentityLayout(tensorType);
|
||||
} else {
|
||||
// Note: Layout maps on function parameters cannot be inferred. The best we
|
||||
// can do at the moment is "fully dynamic".
|
||||
memrefType = getMemRefTypeWithFullyDynamicLayout(tensorType);
|
||||
}
|
||||
|
||||
auto layoutAttr = funcOp.getArgAttrOfType<AffineMapAttr>(
|
||||
index, BufferizationDialect::kBufferLayoutAttrName);
|
||||
@ -386,11 +395,10 @@ struct ReturnOpInterface
|
||||
|
||||
struct FuncOpInterface
|
||||
: public BufferizableOpInterface::ExternalModel<FuncOpInterface, FuncOp> {
|
||||
/// Rewrite function bbArgs and return values into buffer form (using the
|
||||
/// canonical memref layout for now). This function bufferizes the function
|
||||
/// signature and the ReturnOp. When the entire function body has been
|
||||
/// bufferized, function return types can be switched to more concise memref
|
||||
/// types as part of `foldMemRefCasts`.
|
||||
/// Rewrite function bbArgs and return values into buffer form. This function
|
||||
/// bufferizes the function signature and the ReturnOp. When the entire
|
||||
/// function body has been bufferized, function return types can be switched
|
||||
/// to more concise memref types as part of `foldMemRefCasts`.
|
||||
///
|
||||
/// When a tensor function argument is known to be equivalent to a tensor
|
||||
/// result, it is dropped from the return values.
|
||||
@ -439,6 +447,7 @@ struct FuncOpInterface
|
||||
// TODO: Support functions with multiple returns.
|
||||
func::ReturnOp returnOp = getAssumedUniqueReturnOp(funcOp);
|
||||
assert(returnOp && "expected func with single return op");
|
||||
Location loc = returnOp.getLoc();
|
||||
|
||||
// 1. Rewrite the bbArgs. Turn every tensor bbArg into a memref bbArg.
|
||||
Block &frontBlock = funcOp.getBody().front();
|
||||
@ -474,9 +483,11 @@ struct FuncOpInterface
|
||||
SmallVector<Value> returnValues;
|
||||
for (OpOperand &returnOperand : returnOp->getOpOperands()) {
|
||||
Value returnVal = returnOperand.get();
|
||||
auto tensorType = returnVal.getType().dyn_cast<TensorType>();
|
||||
rewriter.setInsertionPoint(returnOp);
|
||||
|
||||
// If not a tensor type just forward it.
|
||||
if (!returnVal.getType().isa<RankedTensorType>()) {
|
||||
if (!tensorType) {
|
||||
returnValues.push_back(returnVal);
|
||||
continue;
|
||||
}
|
||||
@ -485,12 +496,10 @@ struct FuncOpInterface
|
||||
if (options.dropEquivalentFuncResults) {
|
||||
if (Optional<int64_t> equivBbArgIdx = getEquivalentFuncArgIdx(
|
||||
funcOp, funcState, returnOperand.getOperandNumber())) {
|
||||
rewriter.setInsertionPoint(returnOp);
|
||||
Location loc = returnOp.getLoc();
|
||||
// TODO: Use memref type with fully dynamic layout map and add folder
|
||||
// for memref.cast + memref.copy.
|
||||
Value toMemrefOp = rewriter.create<bufferization::ToMemrefOp>(
|
||||
loc,
|
||||
getMemRefType(returnVal.getType().cast<TensorType>(), options),
|
||||
returnVal);
|
||||
loc, getMemRefType(tensorType, options), returnVal);
|
||||
BlockArgument equivBbArg = funcOp.getArgument(*equivBbArgIdx);
|
||||
// Note: This copy will fold away. It must be inserted here to ensure
|
||||
// that `returnVal` still has at least one use and does not fold away.
|
||||
@ -501,7 +510,17 @@ struct FuncOpInterface
|
||||
}
|
||||
}
|
||||
|
||||
returnValues.push_back(*state.getBuffer(rewriter, returnOperand));
|
||||
BaseMemRefType resultType;
|
||||
if (options.functionBoundaryTypeConversion ==
|
||||
BufferizationOptions::LayoutMapOption::IdentityLayoutMap) {
|
||||
resultType = getMemRefTypeWithStaticIdentityLayout(tensorType);
|
||||
} else {
|
||||
// Note: If `InferLayoutMap`, cast are later folded away.
|
||||
resultType = getMemRefTypeWithFullyDynamicLayout(tensorType);
|
||||
}
|
||||
Value toMemrefOp = rewriter.create<bufferization::ToMemrefOp>(
|
||||
loc, resultType, returnVal);
|
||||
returnValues.push_back(toMemrefOp);
|
||||
}
|
||||
|
||||
// 3. Rewrite the terminator without the in-place bufferizable values.
|
||||
|
||||
@ -471,7 +471,10 @@ LogicalResult mlir::bufferization::runOneShotModuleBufferize(
|
||||
// would be invalidated.
|
||||
if (failed(bufferizeOp(funcOp, bufferizationState)))
|
||||
return failure();
|
||||
foldMemRefCasts(funcOp);
|
||||
// Change buffer return types to more precise layout maps.
|
||||
if (options.functionBoundaryTypeConversion ==
|
||||
BufferizationOptions::LayoutMapOption::InferLayoutMap)
|
||||
foldMemRefCasts(funcOp);
|
||||
}
|
||||
|
||||
// Check result.
|
||||
|
||||
@ -6,7 +6,7 @@
|
||||
// RUN: mlir-opt %s -one-shot-bufferize="allow-return-allocs test-analysis-only analysis-fuzzer-seed=91 bufferize-function-boundaries" -split-input-file -o /dev/null
|
||||
|
||||
// Test bufferization using memref types that have no layout map.
|
||||
// RUN: mlir-opt %s -one-shot-bufferize="allow-return-allocs fully-dynamic-layout-maps=0 bufferize-function-boundaries" -split-input-file -o /dev/null
|
||||
// RUN: mlir-opt %s -one-shot-bufferize="allow-return-allocs unknown-type-conversion=identity-layout-map function-boundary-type-conversion=identity-layout-map bufferize-function-boundaries" -split-input-file -o /dev/null
|
||||
|
||||
// CHECK-LABEL: func @write_to_select_op_source
|
||||
// CHECK-SAME: %[[t1:.*]]: memref<?xf32, #{{.*}}>, %[[t2:.*]]: memref<?xf32, #{{.*}}>
|
||||
|
||||
@ -1,7 +1,7 @@
|
||||
// RUN: mlir-opt %s -allow-unregistered-dialect -one-shot-bufferize="allow-return-allocs allow-unknown-ops" -split-input-file | FileCheck %s
|
||||
|
||||
// Test bufferization using memref types that have no layout map.
|
||||
// RUN: mlir-opt %s -allow-unregistered-dialect -one-shot-bufferize="allow-return-allocs allow-unknown-ops fully-dynamic-layout-maps=0" -split-input-file | FileCheck %s --check-prefix=CHECK-NO-LAYOUT-MAP
|
||||
// RUN: mlir-opt %s -allow-unregistered-dialect -one-shot-bufferize="allow-return-allocs allow-unknown-ops unknown-type-conversion=identity-layout-map" -split-input-file | FileCheck %s --check-prefix=CHECK-NO-LAYOUT-MAP
|
||||
|
||||
// Run fuzzer with different seeds.
|
||||
// RUN: mlir-opt %s -allow-unregistered-dialect -one-shot-bufferize="allow-return-allocs test-analysis-only analysis-fuzzer-seed=23" -split-input-file -o /dev/null
|
||||
|
||||
@ -7,7 +7,7 @@
|
||||
// RUN: mlir-opt %s -one-shot-bufferize="bufferize-function-boundaries=1 allow-return-allocs test-analysis-only analysis-fuzzer-seed=91" -split-input-file -o /dev/null
|
||||
|
||||
// Test bufferization using memref types that have no layout map.
|
||||
// RUN: mlir-opt %s -one-shot-bufferize="bufferize-function-boundaries=1 allow-return-allocs fully-dynamic-layout-maps=0" -split-input-file -o /dev/null
|
||||
// RUN: mlir-opt %s -one-shot-bufferize="bufferize-function-boundaries=1 allow-return-allocs unknown-type-conversion=identity-layout-map function-boundary-type-conversion=identity-layout-map" -split-input-file -o /dev/null
|
||||
|
||||
// Make sure that the returned buffer is not deallocated.
|
||||
// TODO: Such buffers currently leak. We need buffer hoisting / ref counting for
|
||||
|
||||
@ -1,4 +1,5 @@
|
||||
// RUN: mlir-opt %s -one-shot-bufferize="bufferize-function-boundaries=1" -split-input-file | FileCheck %s
|
||||
// Note: Default is function-boundary-type-conversion=infer-layout-map
|
||||
// RUN: mlir-opt %s -one-shot-bufferize="bufferize-function-boundaries=1 allow-return-allocs" -split-input-file | FileCheck %s
|
||||
|
||||
// Run fuzzer with different seeds.
|
||||
// RUN: mlir-opt %s -one-shot-bufferize="bufferize-function-boundaries=1 allow-return-allocs test-analysis-only analysis-fuzzer-seed=23" -split-input-file -o /dev/null
|
||||
@ -6,14 +7,29 @@
|
||||
// RUN: mlir-opt %s -one-shot-bufferize="bufferize-function-boundaries=1 allow-return-allocs test-analysis-only analysis-fuzzer-seed=91" -split-input-file -o /dev/null
|
||||
|
||||
// Test bufferization using memref types that have no layout map.
|
||||
// RUN: mlir-opt %s -one-shot-bufferize="bufferize-function-boundaries=1 allow-return-allocs fully-dynamic-layout-maps=0" -split-input-file | FileCheck %s --check-prefix=CHECK-NO-LAYOUT-MAP-LABEL
|
||||
// RUN: mlir-opt %s -one-shot-bufferize="bufferize-function-boundaries=1 allow-return-allocs unknown-type-conversion=identity-layout-map function-boundary-type-conversion=identity-layout-map" -split-input-file | FileCheck %s --check-prefix=CHECK-NO-LAYOUT-MAP
|
||||
|
||||
// Test bufferization using memref types that have fully dynamic layout maps.
|
||||
// RUN: mlir-opt %s -one-shot-bufferize="bufferize-function-boundaries=1 allow-return-allocs function-boundary-type-conversion=fully-dynamic-layout-map" -split-input-file | FileCheck %s --check-prefix=CHECK-FULLY-DYNAMIC-LAYOUT-MAP
|
||||
|
||||
|
||||
// Bufferization of bodiless function with no tensor return value.
|
||||
|
||||
// CHECK-LABEL: func private @private_func
|
||||
// CHECK: #[[$map0:.*]] = affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)>
|
||||
// CHECK: #[[$map1:.*]] = affine_map<(d0, d1)[s0, s1, s2] -> (d0 * s1 + s0 + d1 * s2)>
|
||||
// CHECK-LABEL: func private @private_func(memref<?xf32,
|
||||
// CHECK-SAME: #[[$map0]]>)
|
||||
// CHECK-NO-LAYOUT-MAP-LABEL: func private @private_func(memref<?xf32>)
|
||||
func.func private @private_func(tensor<?xf32>) -> ()
|
||||
|
||||
// CHECK-LABEL: func @empty_func()
|
||||
// CHECK-LABEL: func private @private_func_2d(memref<?x?xf32,
|
||||
// CHECK-SAME: #[[$map1]]>)
|
||||
// CHECK-NO-LAYOUT-MAP-LABEL: func private @private_func_2d(memref<?x?xf32>)
|
||||
func.func private @private_func_2d(tensor<?x?xf32>) -> ()
|
||||
|
||||
// CHECK-LABEL: func @empty_func() {
|
||||
// CHECK-NO-LAYOUT-MAP-LABEL: func @empty_func() {
|
||||
// CHECK-FULLY-DYNAMIC-LAYOUT-MAP-LABEL: func @empty_func() {
|
||||
func.func @empty_func() -> () {
|
||||
return
|
||||
}
|
||||
@ -23,10 +39,45 @@ func.func @empty_func() -> () {
|
||||
// A bodiless function that returns something that is not a tensor.
|
||||
|
||||
// CHECK: func private @external_func_with_return_val(memref<4xi32, #{{.*}}>) -> f32
|
||||
// CHECK-FULLY-DYNAMIC-LAYOUT-MAP: #[[$map1:.*]] = affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)>
|
||||
// CHECK-FULLY-DYNAMIC-LAYOUT-MAP-LABEL: func private @external_func_with_return_val(memref<4xi32,
|
||||
// CHECK-FULLY-DYNAMIC-LAYOUT-MAP-SAME: #[[$map1]]>
|
||||
func.func private @external_func_with_return_val(tensor<4xi32>) -> f32
|
||||
|
||||
// -----
|
||||
|
||||
// A function that returns a non-equivalent tensor with layout map.
|
||||
|
||||
// CHECK: #[[$map2:.*]] = affine_map<(d0, d1)[s0] -> (d0 * 10 + s0 + d1)>
|
||||
// CHECK-LABEL: func @return_extract_slice(%{{.*}}) -> memref<2x?xf32,
|
||||
// CHECK-SAME: #[[$map2]]> {
|
||||
// CHECK: %[[alloc:.*]] = memref.alloc() {{.*}} : memref<20x10xf32>
|
||||
// CHECK: %[[subview:.*]] = memref.subview {{.*}} : memref<20x10xf32> to memref<2x?xf32, #[[$map2]]>
|
||||
// CHECK: return %[[subview]]
|
||||
|
||||
// CHECK-NO-LAYOUT-MAP: #[[$map2:.*]] = affine_map<(d0, d1)[s0] -> (d0 * 10 + s0 + d1)>
|
||||
// CHECK-NO-LAYOUT-MAP-LABEL: func @return_extract_slice(%{{.*}}) -> memref<2x?xf32>
|
||||
// CHECK-NO-LAYOUT-MAP: %[[alloc:.*]] = memref.alloc() {{.*}} : memref<20x10xf32>
|
||||
// CHECK-NO-LAYOUT-MAP: %[[subview:.*]] = memref.subview {{.*}} : memref<20x10xf32> to memref<2x?xf32, #[[$map2]]>
|
||||
// CHECK-NO-LAYOUT-MAP: %[[alloc_no_layout:.*]] = memref.alloc(%{{.*}}) : memref<2x?xf32>
|
||||
// CHECK-NO-LAYOUT-MAP: memref.copy %[[subview]], %[[alloc_no_layout]]
|
||||
// CHECK-NO-LAYOUT-MAP: memref.dealloc %[[alloc]]
|
||||
// CHECK-NO-LAYOUT-MAP: return %[[alloc_no_layout]]
|
||||
|
||||
// CHECK-FULLY-DYNAMIC-LAYOUT-MAP: #[[$map2a:.*]] = affine_map<(d0, d1)[s0, s1, s2] -> (d0 * s1 + s0 + d1 * s2)>
|
||||
// CHECK-FULLY-DYNAMIC-LAYOUT-MAP: #[[$map2b:.*]] = affine_map<(d0, d1)[s0] -> (d0 * 10 + s0 + d1)>
|
||||
// CHECK-FULLY-DYNAMIC-LAYOUT-MAP-LABEL: func @return_extract_slice(%{{.*}}) -> memref<2x?xf32,
|
||||
// CHECK-FULLY-DYNAMIC-LAYOUT-MAP-SAME: #[[$map2a]]> {
|
||||
func.func @return_extract_slice(%idx: index, %sz: index) -> (tensor<2x?xf32>)
|
||||
{
|
||||
%t = linalg.init_tensor [20, 10] : tensor<20x10xf32>
|
||||
%0 = tensor.extract_slice %t[%idx, %idx][2, %sz][1, 1]
|
||||
: tensor<20x10xf32> to tensor<2x?xf32>
|
||||
return %0 : tensor<2x?xf32>
|
||||
}
|
||||
|
||||
// -----
|
||||
|
||||
// CHECK-LABEL: func private @private_func
|
||||
func.func private @private_func(tensor<?xf32>) -> (f32)
|
||||
|
||||
|
||||
@ -6,7 +6,7 @@
|
||||
// RUN: mlir-opt %s -one-shot-bufferize="allow-return-allocs test-analysis-only analysis-fuzzer-seed=91 bufferize-function-boundaries" -split-input-file -o /dev/null
|
||||
|
||||
// Test bufferization using memref types that have no layout map.
|
||||
// RUN: mlir-opt %s -one-shot-bufferize="allow-return-allocs fully-dynamic-layout-maps=0 bufferize-function-boundaries" -split-input-file | FileCheck %s --check-prefix=CHECK-NO-LAYOUT-MAP
|
||||
// RUN: mlir-opt %s -one-shot-bufferize="allow-return-allocs unknown-type-conversion=identity-layout-map function-boundary-type-conversion=identity-layout-map bufferize-function-boundaries" -split-input-file | FileCheck %s --check-prefix=CHECK-NO-LAYOUT-MAP
|
||||
|
||||
// TODO: Some test cases from this file should be moved to other dialects.
|
||||
|
||||
|
||||
@ -6,7 +6,7 @@
|
||||
// RUN: mlir-opt %s -allow-unregistered-dialect -one-shot-bufferize="allow-return-allocs test-analysis-only analysis-fuzzer-seed=91 bufferize-function-boundaries" -split-input-file -o /dev/null
|
||||
|
||||
// Test bufferization using memref types that have no layout map.
|
||||
// RUN: mlir-opt %s -allow-unregistered-dialect -one-shot-bufferize="allow-return-allocs fully-dynamic-layout-maps=0 bufferize-function-boundaries" -split-input-file -o /dev/null
|
||||
// RUN: mlir-opt %s -allow-unregistered-dialect -one-shot-bufferize="allow-return-allocs unknown-type-conversion=identity-layout-map function-boundary-type-conversion=identity-layout-map bufferize-function-boundaries" -split-input-file -o /dev/null
|
||||
|
||||
// CHECK-DAG: #[[$map_1d_dyn:.*]] = affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)>
|
||||
|
||||
|
||||
@ -6,7 +6,7 @@
|
||||
// RUN: mlir-opt %s -one-shot-bufferize="allow-return-allocs test-analysis-only analysis-fuzzer-seed=91 bufferize-function-boundaries" -split-input-file -o /dev/null
|
||||
|
||||
// Test bufferization using memref types that have no layout map.
|
||||
// RUN: mlir-opt %s -one-shot-bufferize="allow-return-allocs fully-dynamic-layout-maps=0 bufferize-function-boundaries" -split-input-file -o /dev/null
|
||||
// RUN: mlir-opt %s -one-shot-bufferize="allow-return-allocs unknown-type-conversion=identity-layout-map bufferize-function-boundaries" -split-input-file -o /dev/null
|
||||
|
||||
// CHECK-DAG: #[[$map_1d_dyn:.*]] = affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)>
|
||||
|
||||
|
||||
Loading…
Reference in New Issue
Block a user