This is for compatibility with existing bufferization passes. Also clean up memref type generation a bit.
Differential Revision: https://reviews.llvm.org/D118243
/// 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;
}
/// Return a contiguous MemRefType (i.e. with canonical/empty layout map)
-/// with the same shape as `shapedType` and specified `layout` and
-/// `addressSpace`.
+/// with the same shape as `shapedType` and specified `addressSpace`.
MemRefType getContiguousMemRefType(ShapedType shapedType,
- MemRefLayoutAttrInterface layout = {},
Attribute memorySpace = {});
-/// Return an UnrankedMemRefType with the given element type and memory space.
-UnrankedMemRefType getUnrankedMemRefType(Type elementType,
- Attribute memorySpace = {});
-
/// Return a MemRefType to which the `tensorType` can be bufferized in a
/// composable fashion. The layout must be the most dynamic possible and
/// canonicalize away once bufferization is finished.
-MemRefType getDynamicMemRefType(RankedTensorType tensorType,
- unsigned addressSpace = 0);
+BaseMemRefType getMemRefType(TensorType tensorType,
+ const BufferizationOptions &options,
+ MemRefLayoutAttrInterface layout = {},
+ Attribute memorySpace = {});
/// Creates a memref allocation with the given type and dynamic extents.
FailureOr<Value> createAlloc(OpBuilder &b, Location loc, MemRefType type,
Option<"useLinalgCopy", "use-linalg-copy", "bool",
/*default=*/"false",
"Use a copy operation implemented as a Linalg op.">,
+ Option<"fullyDynamicLayoutMaps", "fully-dynamic-layout-maps", "bool",
+ /*default=*/"true",
+ "Generate MemRef types with dynamic offset+strides by default.">,
Option<"analysisFuzzerSeed", "analysis-fuzzer-seed", "unsigned",
/*default=*/"0",
"Analyze ops in random order with a given seed (fuzzer)">,
#endif
}
-static Value lookupBuffer(RewriterBase &rewriter, Value tensor) {
- assert(tensor.getType().isa<TensorType>() && "unexpected non-tensor type");
+static Value lookupBuffer(RewriterBase &rewriter, Value tensor,
+ const BufferizationOptions &options) {
+ auto tensorType = tensor.getType().dyn_cast<TensorType>();
+ assert(tensorType && "unexpected non-tensor type");
// Replace "%t = to_tensor %m" with %m.
if (auto toTensorOp = tensor.getDefiningOp<bufferization::ToTensorOp>())
// Insert to_memref op.
OpBuilder::InsertionGuard g(rewriter);
setInsertionPointAfter(rewriter, tensor);
- Type memrefType;
- if (auto rankedTensorType = tensor.getType().dyn_cast<RankedTensorType>()) {
- memrefType = getDynamicMemRefType(rankedTensorType);
- } else {
- memrefType = getUnrankedMemRefType(
- tensor.getType().cast<TensorType>().getElementType());
- }
+ Type memrefType = getMemRefType(tensorType, options);
ensureToMemrefOpIsValid(tensor, memrefType);
return rewriter.create<bufferization::ToMemrefOp>(tensor.getLoc(), memrefType,
tensor);
Operation *op = opOperand.getOwner();
Location loc = op->getLoc();
Value operand = opOperand.get();
- Value operandBuffer = lookupBuffer(rewriter, operand);
+ Value operandBuffer = lookupBuffer(rewriter, operand, options);
if (forceInPlace || isInPlace(opOperand))
return operandBuffer;
return isa<FuncOp>(bbArg.getOwner()->getParentOp());
}
-MemRefType
-bufferization::getContiguousMemRefType(ShapedType shapedType,
- MemRefLayoutAttrInterface layout,
- Attribute memorySpace) {
+MemRefType bufferization::getContiguousMemRefType(ShapedType shapedType,
+ Attribute memorySpace) {
+ MemRefLayoutAttrInterface layout = {};
return MemRefType::get(shapedType.getShape(), shapedType.getElementType(),
layout, memorySpace);
}
-UnrankedMemRefType bufferization::getUnrankedMemRefType(Type elementType,
- Attribute memorySpace) {
- return UnrankedMemRefType::get(elementType, memorySpace);
-}
+BaseMemRefType bufferization::getMemRefType(TensorType tensorType,
+ const BufferizationOptions &options,
+ MemRefLayoutAttrInterface layout,
+ Attribute memorySpace) {
+ // Case 1: Unranked memref type.
+ if (auto unrankedTensorType = tensorType.dyn_cast<UnrankedTensorType>()) {
+ assert(!layout && "UnrankedTensorType cannot have a layout map");
+ return UnrankedMemRefType::get(unrankedTensorType.getElementType(),
+ 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.
+ auto rankedTensorType = tensorType.cast<RankedTensorType>();
+ if (layout || !options.fullyDynamicLayoutMaps) {
+ return MemRefType::get(rankedTensorType.getShape(),
+ rankedTensorType.getElementType(), layout,
+ memorySpace);
+ }
-MemRefType bufferization::getDynamicMemRefType(RankedTensorType tensorType,
- unsigned addressSpace) {
+ // Case 3: Ranked memref type with unspecified layout. Choose the most dynamic
+ // one.
// TODO: address space decisions to connect with the actual alloc.
int64_t dynamicOffset = ShapedType::kDynamicStrideOrOffset;
- SmallVector<int64_t> dynamicStrides(tensorType.getRank(),
+ SmallVector<int64_t> dynamicStrides(rankedTensorType.getRank(),
ShapedType::kDynamicStrideOrOffset);
AffineMap stridedLayout = makeStridedLinearLayoutMap(
- dynamicStrides, dynamicOffset, tensorType.getContext());
- return MemRefType::get(tensorType.getShape(), tensorType.getElementType(),
- stridedLayout, addressSpace);
+ dynamicStrides, dynamicOffset, rankedTensorType.getContext());
+ return MemRefType::get(rankedTensorType.getShape(),
+ rankedTensorType.getElementType(), stridedLayout,
+ memorySpace);
}
/// dynamic buffer type supported.
/// A later pass across all CallOps in the module can decide whether to simplify
/// the types of to version according to some cost model.
-static FunctionType getBufferizedFunctionType(MLIRContext *ctx,
- TypeRange argumentTypes,
- TypeRange resultTypes) {
- auto rewrite = [](Type t) -> Type {
+static FunctionType
+getBufferizedFunctionType(MLIRContext *ctx, TypeRange argumentTypes,
+ TypeRange resultTypes,
+ const BufferizationOptions &options) {
+ auto rewrite = [&](Type t) -> Type {
// TODO: non-zero address space.
// TODO: layout information if relevant.
- if (auto rankedTensorType = t.dyn_cast<RankedTensorType>())
- return getDynamicMemRefType(rankedTensorType);
if (auto tensorType = t.dyn_cast<TensorType>())
- return getUnrankedMemRefType(tensorType.getElementType());
+ return getMemRefType(tensorType, options);
return t;
};
auto argTypes = llvm::to_vector<4>(llvm::map_range(argumentTypes, rewrite));
return funcOp->emitError() << "cannot bufferize bodiless function that "
<< "returns a tensor";
FunctionType bufferizedFuncType = getBufferizedFunctionType(
- funcOp.getContext(), funcOp.getType().getInputs(), TypeRange{});
+ funcOp.getContext(), funcOp.getType().getInputs(), TypeRange{},
+ state.getOptions());
funcOp.setType(bufferizedFuncType);
return success();
}
// 2. Rewrite the terminator without the inPlace bufferizable values.
ValueRange retValues{returnValues};
FunctionType bufferizedFuncType = getBufferizedFunctionType(
- funcOp.getContext(), funcOp.getType().getInputs(), retValues.getTypes());
+ funcOp.getContext(), funcOp.getType().getInputs(), retValues.getTypes(),
+ state.getOptions());
OpBuilder b(returnOp);
b.create<ReturnOp>(returnOp.getLoc(), returnValues);
returnOp->erase();
// Get the bufferized FunctionType for funcOp or construct it if not yet
// available.
FunctionType bufferizedFuncType = getBufferizedFunctionType(
- funcOp.getContext(), argumentTypes, resultTypes);
+ funcOp.getContext(), argumentTypes, resultTypes, state.getOptions());
// 3. Rewrite tensor operands as memrefs based on `bufferizedFuncType`.
for (OpOperand &opOperand : callOp->getOpOperands()) {
// Compute new result types.
SmallVector<Type> newResultTypes;
for (Type type : executeRegionOp->getResultTypes()) {
- if (auto rankedTensorType = type.dyn_cast<RankedTensorType>()) {
- newResultTypes.push_back(getDynamicMemRefType(rankedTensorType));
- } else if (auto tensorType = type.dyn_cast<TensorType>()) {
- newResultTypes.push_back(
- getUnrankedMemRefType(tensorType.getElementType()));
+ if (auto tensorType = type.dyn_cast<TensorType>()) {
+ newResultTypes.push_back(getMemRefType(tensorType, state.getOptions()));
} else {
newResultTypes.push_back(type);
}
// Compute new types of the bufferized scf.if op.
SmallVector<Type> newTypes;
for (Type returnType : ifOp->getResultTypes()) {
- if (returnType.isa<TensorType>()) {
- assert(returnType.isa<RankedTensorType>() &&
- "unsupported unranked tensor");
- newTypes.push_back(
- getDynamicMemRefType(returnType.cast<RankedTensorType>()));
+ if (auto tensorType = returnType.dyn_cast<TensorType>()) {
+ newTypes.push_back(getMemRefType(tensorType, state.getOptions()));
} else {
newTypes.push_back(returnType);
}
options->allowUnknownOps = allowUnknownOps;
options->analysisFuzzerSeed = analysisFuzzerSeed;
options->createDeallocs = createDeallocs;
+ options->fullyDynamicLayoutMaps = fullyDynamicLayoutMaps;
options->printConflicts = printConflicts;
options->testAnalysisOnly = testAnalysisOnly;
layout = rankedMemRefType.getLayout();
// Compute the new memref type.
- Type resultMemRefType;
- if (resultTensorType.isa<RankedTensorType>()) {
- resultMemRefType =
- getContiguousMemRefType(resultTensorType, layout, memorySpace);
- } else {
- resultMemRefType =
- getUnrankedMemRefType(resultTensorType.getElementType(), memorySpace);
- }
+ Type resultMemRefType = getMemRefType(resultTensorType, state.getOptions(),
+ layout, memorySpace);
// Replace the op with a memref.cast.
assert(memref::CastOp::areCastCompatible(resultBuffer->getType(),
Location loc = op->getLoc();
auto tensorType = fromElementsOp.getType().cast<RankedTensorType>();
auto shape = tensorType.getShape();
- MemRefType resultType =
- MemRefType::get(tensorType.getShape(), tensorType.getElementType());
+ MemRefType resultType = getContiguousMemRefType(tensorType);
FailureOr<Value> maybeBuffer =
createAlloc(rewriter, loc, resultType, {},
/*deallocMemref=*/state.getOptions().createDeallocs,
// RUN: mlir-opt %s -allow-unregistered-dialect -linalg-comprehensive-module-bufferize="allow-return-memref 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 -linalg-comprehensive-module-bufferize="allow-return-memref allow-unknown-ops fully-dynamic-layout-maps=0" -split-input-file | FileCheck %s --check-prefix=CHECK-NO-LAYOUT-MAP
+
// Run fuzzer with different seeds.
// RUN: mlir-opt %s -allow-unregistered-dialect -linalg-comprehensive-module-bufferize="allow-return-memref test-analysis-only analysis-fuzzer-seed=23" -split-input-file -o /dev/null
// RUN: mlir-opt %s -allow-unregistered-dialect -linalg-comprehensive-module-bufferize="allow-return-memref test-analysis-only analysis-fuzzer-seed=59" -split-input-file -o /dev/null
// RUN: mlir-opt %s -allow-unregistered-dialect -test-comprehensive-function-bufferize="dialect-filter=tensor allow-unknown-ops allow-return-memref" -canonicalize -split-input-file | FileCheck %s --check-prefix=CHECK-TENSOR
// RUN: mlir-opt %s -allow-unregistered-dialect -test-comprehensive-function-bufferize="dialect-filter=scf allow-unknown-ops allow-return-memref" -canonicalize -split-input-file | FileCheck %s --check-prefix=CHECK-SCF
+// CHECK: #[[$MAP:.*]] = affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)>
+
// CHECK-LABEL: func @use_of_unknown_op_1(
-// CHECK-SAME: %[[m1:.*]]: memref<?xf32
+// CHECK-SAME: %[[m1:.*]]: memref<?xf32, #[[$MAP]]>
+// CHECK-NO-LAYOUT-MAP-LABEL: func @use_of_unknown_op_1(
+// CHECK-NO-LAYOUT-MAP-SAME: %[[m1:.*]]: memref<?xf32>)
func @use_of_unknown_op_1(%t1: tensor<?xf32> {linalg.inplaceable = true})
-> vector<5xf32> {
// ToTensorOp is generated because the function is bufferized and has a
// memref block argument.
- // CHECK: %[[m1_tensor:.*]] = bufferization.to_tensor %[[m1]]
+ // CHECK: %[[m1_tensor:.*]] = bufferization.to_tensor %[[m1]] : memref<?xf32, #[[$MAP]]>
// CHECK: %[[dummy:.*]] = "test.dummy_op"(%[[m1_tensor]])
+ // CHECK-NO-LAYOUT-MAP: %[[m1_tensor:.*]] = bufferization.to_tensor %[[m1]] : memref<?xf32>
+ // CHECK-NO-LAYOUT-MAP: %[[dummy:.*]] = "test.dummy_op"(%[[m1_tensor]])
%0 = "test.dummy_op"(%t1) : (tensor<?xf32>) -> tensor<?xf32>
%idx = arith.constant 0 : index
%cst = arith.constant 0.0 : f32
- // CHECK: %[[dummy_memref:.*]] = bufferization.to_memref %[[dummy]]
- // CHECK: vector.transfer_read %[[dummy_memref]]
+ // CHECK: %[[dummy_memref:.*]] = bufferization.to_memref %[[dummy]] : memref<?xf32, #[[$MAP]]>
+ // CHECK: vector.transfer_read %[[dummy_memref]][%{{.*}}], %{{.*}} : memref<?xf32, #[[$MAP]]>
+ // CHECK-NO-LAYOUT-MAP: %[[dummy_memref:.*]] = bufferization.to_memref %[[dummy]] : memref<?xf32>
+ // CHECK-NO-LAYOUT-MAP: vector.transfer_read %[[dummy_memref]][%{{.*}}], %{{.*}} : memref<?xf32>
%1 = vector.transfer_read %0[%idx], %cst : tensor<?xf32>, vector<5xf32>
return %1 : vector<5xf32>
}
// RUN: mlir-opt %s -linalg-comprehensive-module-bufferize="allow-return-memref test-analysis-only analysis-fuzzer-seed=59" -split-input-file -o /dev/null
// RUN: mlir-opt %s -linalg-comprehensive-module-bufferize="allow-return-memref 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 -linalg-comprehensive-module-bufferize="allow-return-memref fully-dynamic-layout-maps=0" -split-input-file | FileCheck %s --check-prefix=CHECK-NO-LAYOUT-MAP
+
// CHECK-LABEL: func @transfer_read(%{{.*}}: memref<?xf32, #map>) -> vector<4xf32> {
+// CHECK-NO-LAYOUT-MAP-LABEL: func @transfer_read(%{{.*}}: memref<?xf32>) -> vector<4xf32>
func @transfer_read(
%A : tensor<?xf32> {linalg.inplaceable = false})
-> (vector<4xf32>)
// CHECK-LABEL: func @fill_inplace(
// CHECK-SAME: %[[A:[a-zA-Z0-9]*]]: memref<?xf32, #[[$map_1d_dyn]]>
+// CHECK-NO-LAYOUT-MAP-LABEL: func @fill_inplace(%{{.*}}: memref<?xf32>) {
func @fill_inplace(
%A : tensor<?xf32> {linalg.inplaceable = true})
-> tensor<?xf32>
/// No linalg.inplaceable flag, must allocate.
// CHECK-LABEL: func @not_inplace(
// CHECK-SAME: %[[A:[a-zA-Z0-9]*]]: memref<?xf32, #[[$map_1d_dyn]]>) -> memref<?xf32> {
+// CHECK-NO-LAYOUT-MAP-LABEL: func @not_inplace(%{{.*}}: memref<?xf32>) -> memref<?xf32>
func @not_inplace(
%A : tensor<?xf32> {linalg.inplaceable = false})
-> tensor<?xf32>
// CHECK-LABEL: func @not_inplace
// CHECK-SAME: %[[A:[a-zA-Z0-9]*]]: memref<?x?xf32, #[[$map_2d_dyn]]>) {
+// CHECK-NO-LAYOUT-MAP-LABEL: func @not_inplace(%{{.*}}: memref<?x?xf32>) {
func @not_inplace(
%A : tensor<?x?xf32> {linalg.inplaceable = true})
-> tensor<?x?xf32>
*this, "dialect-filter",
llvm::cl::desc("Bufferize only ops from the specified dialects"),
llvm::cl::ZeroOrMore, llvm::cl::MiscFlags::CommaSeparated};
+ Option<bool> fullyDynamicLayoutMaps{
+ *this, "fully-dynamic-layout-maps",
+ llvm::cl::desc("Use fully dynamic layout maps on memref types"),
+ llvm::cl::init(true)};
Option<bool> createDeallocs{
*this, "create-deallocs",
llvm::cl::desc("Specify if buffers should be deallocated"),
options->allowUnknownOps = allowUnknownOps;
options->testAnalysisOnly = testAnalysisOnly;
options->analysisFuzzerSeed = analysisFuzzerSeed;
+ options->fullyDynamicLayoutMaps = fullyDynamicLayoutMaps;
options->createDeallocs = createDeallocs;
if (dialectFilter.hasValue()) {