[Mlir-commits] [mlir] 0b1fe74 - [mlir][bufferization] Introduce reconcileBufferTypeMismatchFn hook (#202667)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Jun 22 02:20:26 PDT 2026
Author: Andrei Golubev
Date: 2026-06-22T11:20:19+02:00
New Revision: 0b1fe74f079fd330e4d4d85ce33b7f5702da1be7
URL: https://github.com/llvm/llvm-project/commit/0b1fe74f079fd330e4d4d85ce33b7f5702da1be7
DIFF: https://github.com/llvm/llvm-project/commit/0b1fe74f079fd330e4d4d85ce33b7f5702da1be7.diff
LOG: [mlir][bufferization] Introduce reconcileBufferTypeMismatchFn hook (#202667)
This PR is the first part of the work that aims to allow customizations
in resolving mismatching buffer types.
Add a new bufferization hook that lets downstream bufferization
implementations define how to handle buffer mismatches that appear
during type inference in various upstream scenarios.
The hook is used as a fallback mechanism in several upstream operations.
For example, when bufferizing block signatures (scf.execute_region), and
resolving "branch" conflicts (scf.if, scfl.index_switch, scf.for,
arith.select).
The hook returns a valid buffer type when reconciliation succeeded;
failure indicates reconciliation failure and should be treated as
bufferization failure. The caller of the hook is expected to use the
returned buffer type. By default, a memref with fully-dynamic layout map
is returned (for unranked case, buffers are assumed to match).
---------
Co-authored-by: Dmitrii Makarenko <dmitrii.makarenko at intel.com>
Added:
Modified:
mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h
mlir/include/mlir/Dialect/Bufferization/IR/UnstructuredControlFlow.h
mlir/lib/Dialect/Arith/Transforms/BufferizableOpInterfaceImpl.cpp
mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp
mlir/test/Dialect/Arith/one-shot-bufferize-memory-space-invalid.mlir
mlir/test/Dialect/Bufferization/Transforms/test-one-shot-module-bufferize.mlir
mlir/test/Dialect/ControlFlow/one-shot-bufferize-invalid.mlir
mlir/test/Dialect/SCF/one-shot-bufferize-invalid.mlir
mlir/test/lib/Dialect/Bufferization/TestOneShotModuleBufferize.cpp
mlir/test/lib/Dialect/Test/TestOpDefs.cpp
mlir/test/lib/Dialect/Test/TestOps.td
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h
index bf94f56497160..9f7a6a638f307 100644
--- a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h
+++ b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h
@@ -273,6 +273,13 @@ struct BufferizationOptions {
using DefaultMemorySpaceFn =
std::function<std::optional<Attribute>(TensorLikeType t)>;
+ /// Resolve a mismatch between buffer types that were independently inferred,
+ /// which results in a conflict at the "merge" point. Returns `failure()` to
+ /// signal bufferization failure; returns a buffer-like type when
+ /// reconciliation suceeded.
+ using ReconcileBufferTypeMismatchFn = std::function<FailureOr<BufferLikeType>(
+ BufferLikeType, BufferLikeType, const BufferizationOptions &)>;
+
BufferizationOptions();
/// Try to cast the given op to BufferizableOpInterface if the op is allow
@@ -364,6 +371,16 @@ struct BufferizationOptions {
DefaultMemorySpaceFn defaultMemorySpaceFn =
[](TensorLikeType t) -> std::optional<Attribute> { return Attribute(); };
+ /// Hook to resolve a mismatch between conflicting buffer types that were
+ /// independently inferred and have to now "converge" to a common buffer type
+ /// (e.g. due to
diff erences in iterations of a loop or branches of
+ /// if-statements). Depending on the situation and the types involved, this
+ /// may produce a "joined" type (e.g. a type combining properties of both), or
+ /// either one of the two types, etc. The default keeps the framework
+ /// behavior: promote to fully-dynamic layout on layout mismatch, fail on
+ /// memory-space mismatch.
+ ReconcileBufferTypeMismatchFn reconcileBufferTypeMismatchFn = nullptr;
+
/// 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;
diff --git a/mlir/include/mlir/Dialect/Bufferization/IR/UnstructuredControlFlow.h b/mlir/include/mlir/Dialect/Bufferization/IR/UnstructuredControlFlow.h
index e8a81c74bd77a..344d7a8f9408e 100644
--- a/mlir/include/mlir/Dialect/Bufferization/IR/UnstructuredControlFlow.h
+++ b/mlir/include/mlir/Dialect/Bufferization/IR/UnstructuredControlFlow.h
@@ -47,8 +47,7 @@ struct OpWithUnstructuredControlFlowBufferizableOpInterfaceExternalModel
// operand types of all forwarded values. If these are all the same type,
// take that type. Otherwise, take only the memory space and fall back to a
// buffer type with a fully dynamic layout map.
- BaseMemRefType bufferType;
- auto tensorType = cast<TensorType>(value.getType());
+ BufferLikeType bufferType;
for (OpOperand *opOperand :
detail::getCallerOpOperands(cast<BlockArgument>(value))) {
@@ -59,19 +58,18 @@ struct OpWithUnstructuredControlFlowBufferizableOpInterfaceExternalModel
continue;
// Compute the bufferized type of the forwarded operand.
- BaseMemRefType callerType;
- if (auto memrefType =
- dyn_cast<BaseMemRefType>(opOperand->get().getType())) {
+ BufferLikeType callerType;
+ if (auto bufferType =
+ dyn_cast<BufferLikeType>(opOperand->get().getType())) {
// The operand was already bufferized. Take its type directly.
- callerType = memrefType;
+ callerType = bufferType;
} else {
FailureOr<BufferLikeType> maybeCallerType =
bufferization::getBufferType(opOperand->get(), options, state,
invocationStack);
if (failed(maybeCallerType))
return failure();
- assert(isa<BaseMemRefType>(*maybeCallerType) && "expected memref type");
- callerType = cast<BaseMemRefType>(*maybeCallerType);
+ callerType = *maybeCallerType;
}
if (!bufferType) {
@@ -83,34 +81,34 @@ struct OpWithUnstructuredControlFlowBufferizableOpInterfaceExternalModel
if (bufferType == callerType)
continue;
- // If the computed buffer type does not match the computed buffer type
- // of the earlier forwarded operands, fall back to a buffer type with a
- // fully dynamic layout map.
+ // If the computed buffer type does not match the computed buffer type of
+ // the earlier forwarded operands, fall back to a reconciled buffer type.
#ifndef NDEBUG
- if (auto rankedTensorType = dyn_cast<RankedTensorType>(tensorType)) {
- assert(bufferType.hasRank() && callerType.hasRank() &&
- "expected ranked memrefs");
- assert(llvm::all_equal({bufferType.getShape(), callerType.getShape(),
- rankedTensorType.getShape()}) &&
- "expected same shape");
- } else {
- assert(!bufferType.hasRank() && !callerType.hasRank() &&
- "expected unranked memrefs");
+ if (auto tensorType = dyn_cast<TensorLikeType>(value.getType())) {
+ const auto emitOpError = [&]() { return op->emitOpError(); };
+ assert(succeeded(tensorType.verifyCompatibleBufferType(bufferType,
+ emitOpError)) &&
+ "incompatible buffer type");
+ assert(succeeded(tensorType.verifyCompatibleBufferType(callerType,
+ emitOpError)) &&
+ "incompatible caller type");
}
#endif // NDEBUG
- if (bufferType.getMemorySpace() != callerType.getMemorySpace())
- return op->emitOpError("incoming operands of block argument have "
- "inconsistent memory spaces");
+ auto reconciled = options.reconcileBufferTypeMismatchFn(
+ bufferType, callerType, options);
+ if (failed(reconciled)) {
+ return op->emitError("incoming operands of block argument have "
+ "incompatible buffer types");
+ }
- bufferType = getMemRefTypeWithFullyDynamicLayout(
- tensorType, bufferType.getMemorySpace());
+ bufferType = *reconciled;
}
if (!bufferType)
return op->emitOpError("could not infer buffer type of block argument");
- return cast<BufferLikeType>(bufferType);
+ return bufferType;
}
protected:
diff --git a/mlir/lib/Dialect/Arith/Transforms/BufferizableOpInterfaceImpl.cpp b/mlir/lib/Dialect/Arith/Transforms/BufferizableOpInterfaceImpl.cpp
index 91d03a0356de1..35e472470560a 100644
--- a/mlir/lib/Dialect/Arith/Transforms/BufferizableOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Arith/Transforms/BufferizableOpInterfaceImpl.cpp
@@ -163,12 +163,15 @@ struct SelectOpInterface
// The "true" and the "false" operands must have the same type. If the
// buffers have
diff erent types, they
diff er only in their layout map. Cast
- // both of them to the most dynamic MemRef type.
+ // both of them to the reconciled type.
if (trueBuffer.getType() != falseBuffer.getType()) {
- auto targetType = bufferization::detail::asMemRefType(
- bufferization::getBufferType(selectOp.getResult(), options, state));
- if (failed(targetType))
- return failure();
+ auto targetType = options.reconcileBufferTypeMismatchFn(
+ cast<BufferLikeType>(trueBuffer.getType()),
+ cast<BufferLikeType>(falseBuffer.getType()), options);
+ if (failed(targetType)) {
+ return selectOp->emitError(
+ "incompatible buffer types on true/false operands");
+ }
if (trueBuffer.getType() != *targetType)
trueBuffer =
memref::CastOp::create(rewriter, loc, *targetType, trueBuffer);
@@ -188,26 +191,17 @@ struct SelectOpInterface
SmallVector<Value> &invocationStack) const {
auto selectOp = cast<arith::SelectOp>(op);
assert(value == selectOp.getResult() && "invalid value");
- auto trueType =
- bufferization::detail::asMemRefType(bufferization::getBufferType(
- selectOp.getTrueValue(), options, state, invocationStack));
- auto falseType =
- bufferization::detail::asMemRefType(bufferization::getBufferType(
- selectOp.getFalseValue(), options, state, invocationStack));
+ auto trueType = bufferization::getBufferType(
+ selectOp.getTrueValue(), options, state, invocationStack);
+ auto falseType = bufferization::getBufferType(
+ selectOp.getFalseValue(), options, state, invocationStack);
if (failed(trueType) || failed(falseType))
return failure();
if (*trueType == *falseType)
return cast<BufferLikeType>(*trueType);
- if (trueType->getMemorySpace() != falseType->getMemorySpace())
- return op->emitError("inconsistent memory space on true/false operands");
-
- // If the buffers have
diff erent types, they
diff er only in their layout
- // map.
- auto memrefType = llvm::cast<MemRefType>(*trueType);
- return cast<BufferLikeType>(getMemRefTypeWithFullyDynamicLayout(
- RankedTensorType::get(memrefType.getShape(),
- memrefType.getElementType()),
- memrefType.getMemorySpace()));
+
+ return options.reconcileBufferTypeMismatchFn(*trueType, *falseType,
+ options);
}
};
diff --git a/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp b/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
index e60736a20dc3b..e0a35e1932f3d 100644
--- a/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
+++ b/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
@@ -347,6 +347,17 @@ bool OpFilter::isOpAllowed(Operation *op) const {
namespace {
+// A helper overload for bufferization::getMemRefTypeWithFullyDynamicLayout().
+BaseMemRefType getMemRefTypeWithFullyDynamicLayout(ArrayRef<int64_t> shape,
+ mlir::Type elementType,
+ Attribute memorySpace) {
+ int64_t dynamicOffset = ShapedType::kDynamic;
+ SmallVector<int64_t> dynamicStrides(shape.size(), ShapedType::kDynamic);
+ auto stridedLayout = StridedLayoutAttr::get(elementType.getContext(),
+ dynamicOffset, dynamicStrides);
+ return MemRefType::get(shape, elementType, stridedLayout, memorySpace);
+}
+
/// Default function arg type converter: Use a fully dynamic layout map.
BufferLikeType
defaultFunctionArgTypeConverter(TensorLikeType type, Attribute memorySpace,
@@ -354,7 +365,8 @@ defaultFunctionArgTypeConverter(TensorLikeType type, Attribute memorySpace,
const BufferizationOptions &options) {
if (auto tensorType = mlir::dyn_cast<TensorType>(type)) {
return cast<BufferLikeType>(
- getMemRefTypeWithFullyDynamicLayout(tensorType, memorySpace));
+ bufferization::getMemRefTypeWithFullyDynamicLayout(tensorType,
+ memorySpace));
}
// If not builtin, fallback to unknown type conversion.
@@ -364,8 +376,29 @@ defaultFunctionArgTypeConverter(TensorLikeType type, Attribute memorySpace,
BufferLikeType
defaultUnknownTypeConverter(TensorLikeType tensorType, Attribute memorySpace,
const BufferizationOptions &options) {
- return cast<BufferLikeType>(getMemRefTypeWithFullyDynamicLayout(
- cast<TensorType>(tensorType), memorySpace));
+ return cast<BufferLikeType>(
+ bufferization::getMemRefTypeWithFullyDynamicLayout(
+ cast<TensorType>(tensorType), memorySpace));
+}
+
+/// Default reconcile hook: memory space mismatch is an error, layout mismatch
+/// is resolved by promoting to fully dynamic.
+FailureOr<BufferLikeType>
+defaultReconcileBufferTypeMismatch(BufferLikeType x, BufferLikeType y,
+ const BufferizationOptions &) {
+ const auto xMemRef = cast<BaseMemRefType>(x);
+ const auto yMemRef = cast<BaseMemRefType>(y);
+
+ if (xMemRef.getMemorySpace() != yMemRef.getMemorySpace())
+ return failure();
+
+ if (isa<UnrankedMemRefType>(xMemRef)) {
+ // unranked memrefs have no layout.
+ return x;
+ }
+
+ return cast<BufferLikeType>(::getMemRefTypeWithFullyDynamicLayout(
+ xMemRef.getShape(), xMemRef.getElementType(), xMemRef.getMemorySpace()));
}
} // namespace
@@ -373,7 +406,8 @@ defaultUnknownTypeConverter(TensorLikeType tensorType, Attribute memorySpace,
// Default constructor for BufferizationOptions.
BufferizationOptions::BufferizationOptions()
: functionArgTypeConverterFn(defaultFunctionArgTypeConverter),
- unknownTypeConverterFn(defaultUnknownTypeConverter) {}
+ unknownTypeConverterFn(defaultUnknownTypeConverter),
+ reconcileBufferTypeMismatchFn(defaultReconcileBufferTypeMismatch) {}
bool BufferizationOptions::isOpAllowed(Operation *op) const {
// Special case: If function boundary bufferization is deactivated, do not
@@ -821,15 +855,8 @@ bufferization::getMemRefTypeWithFullyDynamicLayout(TensorType tensorType,
}
// Case 2: Ranked memref type.
- auto rankedTensorType = llvm::cast<RankedTensorType>(tensorType);
- int64_t dynamicOffset = ShapedType::kDynamic;
- SmallVector<int64_t> dynamicStrides(rankedTensorType.getRank(),
- ShapedType::kDynamic);
- auto stridedLayout = StridedLayoutAttr::get(tensorType.getContext(),
- dynamicOffset, dynamicStrides);
- return MemRefType::get(rankedTensorType.getShape(),
- rankedTensorType.getElementType(), stridedLayout,
- memorySpace);
+ return ::getMemRefTypeWithFullyDynamicLayout(
+ tensorType.getShape(), tensorType.getElementType(), memorySpace);
}
/// Return a MemRef type with a static identity layout (i.e., no layout map). If
diff --git a/mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp b/mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp
index 16eb9aadc06f0..b9662751aa7f4 100644
--- a/mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp
@@ -315,20 +315,12 @@ struct IfOpInterface
if (thenBufferType == elseBufferType)
return cast<BufferLikeType>(thenBufferType);
- // Memory space mismatch.
- auto thenBaseMemRefType = dyn_cast<BaseMemRefType>(thenBufferType);
- auto elseBaseMemRefType = dyn_cast<BaseMemRefType>(elseBufferType);
- if (thenBaseMemRefType && elseBaseMemRefType &&
- thenBaseMemRefType.getMemorySpace() !=
- elseBaseMemRefType.getMemorySpace())
- return op->emitError("inconsistent memory space on then/else branches");
-
- // TODO: Properly support with options, for now it is hardcoded MemRef type
- // based approach Layout maps are
diff erent: Promote to fully dynamic layout
- // map.
- return cast<BufferLikeType>(getMemRefTypeWithFullyDynamicLayout(
- cast<TensorType>(opResult.getType()),
- thenBaseMemRefType.getMemorySpace()));
+ auto reconciled = options.reconcileBufferTypeMismatchFn(
+ thenBufferType, elseBufferType, options);
+ if (failed(reconciled))
+ return op->emitError("incompatible buffer types on then/else branches");
+
+ return *reconciled;
}
};
@@ -407,23 +399,20 @@ struct IndexSwitchOpInterface
assert(value.getDefiningOp() == op && "invalid value");
int64_t resultNum = cast<OpResult>(value).getResultNumber();
- // TODO: Properly support with options, for now it is hardcoded MemRef type
- // based approach Helper function to get buffer type of a case.
- auto getYieldedBufferType = [&](Block &b) -> FailureOr<BaseMemRefType> {
+ auto getYieldedBufferType = [&](Block &b) -> FailureOr<BufferLikeType> {
auto yieldOp = cast<scf::YieldOp>(b.getTerminator());
Value yieldedValue = yieldOp->getOperand(resultNum);
- if (auto bufferType = dyn_cast<BaseMemRefType>(yieldedValue.getType()))
+ if (auto bufferType = dyn_cast<BufferLikeType>(yieldedValue.getType()))
return bufferType;
- auto maybeBufferType = bufferization::getBufferType(
- yieldedValue, options, state, invocationStack);
- return bufferization::detail::asMemRefType(maybeBufferType);
+ return bufferization::getBufferType(yieldedValue, options, state,
+ invocationStack);
};
// Compute buffer type of the default case.
auto maybeBufferType = getYieldedBufferType(switchOp.getDefaultBlock());
if (failed(maybeBufferType))
return failure();
- BaseMemRefType bufferType = *maybeBufferType;
+ BufferLikeType bufferType = *maybeBufferType;
// Compute buffer types of all other cases.
for (int64_t i = 0, numCases = switchOp.getNumCases(); i < numCases; ++i) {
@@ -435,15 +424,11 @@ struct IndexSwitchOpInterface
if (bufferType == *yieldedBufferType)
continue;
- // Memory space mismatch.
- if (bufferType.getMemorySpace() != yieldedBufferType->getMemorySpace())
- return op->emitError("inconsistent memory space on switch cases");
-
- // TODO: Properly support with options, for now it is hardcoded MemRef
- // type based approach Layout maps are
diff erent: Promote to fully dynamic
- // layout map.
- bufferType = getMemRefTypeWithFullyDynamicLayout(
- cast<TensorType>(value.getType()), bufferType.getMemorySpace());
+ auto reconciled = options.reconcileBufferTypeMismatchFn(
+ bufferType, *yieldedBufferType, options);
+ if (failed(reconciled))
+ return op->emitError("incompatible buffer types on switch cases");
+ bufferType = *reconciled;
}
return cast<BufferLikeType>(bufferType);
@@ -558,9 +543,9 @@ static FailureOr<BufferLikeType> computeLoopRegionIterArgBufferType(
// Compute the buffer type of the yielded value.
BufferLikeType yieldedValueBufferType;
- if (isa<BufferLikeType>(yieldedValue.getType())) {
+ if (auto bufferType = dyn_cast<BufferLikeType>(yieldedValue.getType())) {
// scf.yield was already bufferized.
- yieldedValueBufferType = cast<BufferLikeType>(yieldedValue.getType());
+ yieldedValueBufferType = bufferType;
} else {
// Note: This typically triggers a recursive call for the buffer type of
// the iter_arg.
@@ -576,27 +561,27 @@ static FailureOr<BufferLikeType> computeLoopRegionIterArgBufferType(
return yieldedValueBufferType;
// If there is a mismatch between the yielded buffer type and the init_arg
- // buffer type, the buffer type must be promoted to a fully dynamic layout
- // map.
- auto yieldedBufferType = cast<BaseMemRefType>(yieldedValueBufferType);
- auto iterTensorType = cast<TensorType>(iterArg.getType());
- auto initBufferType = llvm::cast<BaseMemRefType>(*initArgBufferType);
- if (initBufferType.getMemorySpace() != yieldedBufferType.getMemorySpace())
- return loopOp->emitOpError(
- "init_arg and yielded value bufferize to inconsistent memory spaces");
+ // buffer type, the buffer type must be reconciled.
#ifndef NDEBUG
- if (auto yieldedRankedBufferType = dyn_cast<MemRefType>(yieldedBufferType)) {
- assert(
- llvm::all_equal({yieldedRankedBufferType.getShape(),
- cast<MemRefType>(initBufferType).getShape(),
- cast<RankedTensorType>(iterTensorType).getShape()}) &&
- "expected same shape");
+ if (auto iterTensorType = dyn_cast<TensorLikeType>(iterArg.getType())) {
+ const auto emitOpError = [&]() { return loopOp->emitOpError(); };
+ assert(succeeded(iterTensorType.verifyCompatibleBufferType(
+ yieldedValueBufferType, emitOpError)) &&
+ "incompatible yielded type");
+ assert(succeeded(iterTensorType.verifyCompatibleBufferType(
+ *initArgBufferType, emitOpError)) &&
+ "incompatible init_arg type");
}
#endif // NDEBUG
- // TODO: Properly support with options, for now it is hardcoded MemRef type
- // based approach
- return cast<BufferLikeType>(getMemRefTypeWithFullyDynamicLayout(
- iterTensorType, yieldedBufferType.getMemorySpace()));
+
+ auto reconciled = options.reconcileBufferTypeMismatchFn(
+ *initArgBufferType, yieldedValueBufferType, options);
+ if (failed(reconciled)) {
+ return loopOp->emitError(
+ "init_arg and yielded value bufferize to incompatible buffer types");
+ }
+
+ return *reconciled;
}
/// Return `true` if the given loop may have 0 iterations.
diff --git a/mlir/test/Dialect/Arith/one-shot-bufferize-memory-space-invalid.mlir b/mlir/test/Dialect/Arith/one-shot-bufferize-memory-space-invalid.mlir
index deda8bb74b323..3a6f3ce050e51 100644
--- a/mlir/test/Dialect/Arith/one-shot-bufferize-memory-space-invalid.mlir
+++ b/mlir/test/Dialect/Arith/one-shot-bufferize-memory-space-invalid.mlir
@@ -5,7 +5,7 @@ func.func @inconsistent_memory_space_arith_select(%c: i1) -> tensor<10xf32> {
// bufferized.
%0 = bufferization.alloc_tensor() {memory_space = 0 : ui64} : tensor<10xf32>
%1 = bufferization.alloc_tensor() {memory_space = 1 : ui64} : tensor<10xf32>
- // expected-error @+2 {{inconsistent memory space on true/false operands}}
+ // expected-error @+2 {{incompatible buffer types on true/false operands}}
// expected-error @+1 {{failed to bufferize op}}
%r = arith.select %c, %0, %1 : tensor<10xf32>
func.return %r : tensor<10xf32>
diff --git a/mlir/test/Dialect/Bufferization/Transforms/test-one-shot-module-bufferize.mlir b/mlir/test/Dialect/Bufferization/Transforms/test-one-shot-module-bufferize.mlir
index 6ea68bec24b53..b2674e9497927 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/test-one-shot-module-bufferize.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/test-one-shot-module-bufferize.mlir
@@ -3,6 +3,7 @@
#enc1 = #test.tensor_encoding<"hello">
#enc2 = #test.tensor_encoding<"not hello">
+// CHECK-LABEL: @BufferizeEncodingThroughFunctionBoundaryAndCustomOps
module @BufferizeEncodingThroughFunctionBoundaryAndCustomOps {
// CHECK: func @inner_func(
// CHECK-SAME: %[[arg0:.*]]: memref<?xf32, #test.memref_layout<"hello">>)
@@ -272,6 +273,7 @@ func.func @scf_while_inplace(
#enc1 = #test.tensor_encoding<"custom">
+// CHECK-LABEL: @BufferizeEncodingForCustomOpsInsideScf
module @BufferizeEncodingForCustomOpsInsideScf {
// CHECK: func.func @custom_encoding_inside_scf(
// CHECK-SAME: %[[arg:.*]]: memref<42xf64, #test.memref_layout<"custom">>,
@@ -296,3 +298,283 @@ module @BufferizeEncodingForCustomOpsInsideScf {
return %loop : tensor<42xf64, #enc1>
}
}
+
+// -----
+
+#layout1 = #test.memref_layout<"layout_a">
+
+module @BufferizeLayoutForFunction {
+ // CHECK: func.func @layout_for_func
+ // CHECK-SAME: -> memref<10xf32>
+ func.func @layout_for_func() -> tensor<10xf32> {
+ // CHECK: %[[memref:.*]] = "test.create_memref_op"
+ // CHECK-SAME: -> memref<10xf32, #test.memref_layout<"layout_a">>
+ %memref = "test.tensor_with_future_layout"() {layout = #layout1}
+ : () -> tensor<10xf32>
+
+ // CHECK: %[[out:.*]] = memref.cast %[[memref]]
+ // CHECK-SAME: to memref<10xf32>
+ // CHECK: return %[[out]]
+ return %memref : tensor<10xf32>
+ }
+}
+
+// -----
+
+#layout1 = #test.memref_layout<"layout_a">
+#layout2 = #test.memref_layout<"layout_b">
+
+// CHECK-LABEL: @BufferizeLayoutMismatchInsideScfIf
+module @BufferizeLayoutMismatchInsideScfIf {
+ // CHECK: func.func @mismatch_in_if
+ // CHECK-SAME: -> memref<10xf32>
+ func.func @mismatch_in_if(%cond: i1) -> tensor<10xf32> {
+ // CHECK: %[[if:.*]] = scf.if
+ // CHECK-SAME: -> (memref<10xf32, #test.memref_layout<"layout_a">>)
+ %ret = scf.if %cond -> tensor<10xf32> {
+ // CHECK: %[[one:.*]] = "test.create_memref_op"
+ // CHECK-SAME: -> memref<10xf32, #test.memref_layout<"layout_a">>
+ %one = "test.tensor_with_future_layout"() {layout = #layout1}
+ : () -> tensor<10xf32>
+ // CHECK: scf.yield %[[one]]
+ scf.yield %one : tensor<10xf32>
+ } else {
+ // CHECK: %[[another:.*]] = "test.create_memref_op"
+ // CHECK-SAME: -> memref<10xf32, #test.memref_layout<"layout_b">>
+ // CHECK: %[[cast:.*]] = memref.cast %[[another]]
+ // CHECK-SAME: to memref<10xf32, #test.memref_layout<"layout_a">>
+ %another = "test.tensor_with_future_layout"() {layout = #layout2}
+ : () -> tensor<10xf32>
+ // CHECK: scf.yield %[[cast]]
+ scf.yield %another : tensor<10xf32>
+ }
+
+ // CHECK: %[[out:.*]] = memref.cast %[[if]]
+ // CHECK-SAME: to memref<10xf32>
+ // CHECK: return %[[out]]
+ return %ret : tensor<10xf32>
+ }
+}
+
+// -----
+
+#layout1 = #test.memref_layout<"layout_a">
+#layout2 = #test.memref_layout<"layout_b">
+#layout3 = #test.memref_layout<"layout_c">
+
+// CHECK-LABEL: @BufferizeLayoutMismatchInsideScfSwitch
+module @BufferizeLayoutMismatchInsideScfSwitch {
+ // CHECK: func.func @mismatch_in_switch
+ // CHECK-SAME: -> memref<10xf32>
+ func.func @mismatch_in_switch(%idx: index) -> tensor<10xf32> {
+ // CHECK: %[[switch:.*]] = scf.index_switch
+ // CHECK-SAME: -> memref<10xf32, #test.memref_layout<"layout_a">>
+ %ret = scf.index_switch %idx -> tensor<10xf32>
+ case 0 {
+ // CHECK: %[[one:.*]] = "test.create_memref_op"
+ // CHECK-SAME: -> memref<10xf32, #test.memref_layout<"layout_a">>
+ %one = "test.tensor_with_future_layout"() {layout = #layout1}
+ : () -> tensor<10xf32>
+ // CHECK: scf.yield %[[one]]
+ scf.yield %one : tensor<10xf32>
+ }
+ case 1 {
+ // CHECK: %[[another:.*]] = "test.create_memref_op"
+ // CHECK-SAME: -> memref<10xf32, #test.memref_layout<"layout_b">>
+ // CHECK: %[[cast:.*]] = memref.cast %[[another]]
+ // CHECK-SAME: to memref<10xf32, #test.memref_layout<"layout_a">>
+ %another = "test.tensor_with_future_layout"() {layout = #layout2}
+ : () -> tensor<10xf32>
+ // CHECK: scf.yield %[[cast]]
+ scf.yield %another : tensor<10xf32>
+ }
+ default {
+ // CHECK: %[[yet_another:.*]] = "test.create_memref_op"
+ // CHECK-SAME: -> memref<10xf32, #test.memref_layout<"layout_c">>
+ // CHECK: %[[cast:.*]] = memref.cast %[[yet_another]]
+ // CHECK-SAME: to memref<10xf32, #test.memref_layout<"layout_a">>
+ %yet_another = "test.tensor_with_future_layout"() {layout = #layout3}
+ : () -> tensor<10xf32>
+ // CHECK: scf.yield %[[cast]]
+ scf.yield %yet_another : tensor<10xf32>
+ }
+
+ // CHECK: %[[out:.*]] = memref.cast %[[switch]]
+ // CHECK-SAME: to memref<10xf32>
+ // CHECK: return %[[out]]
+ return %ret : tensor<10xf32>
+ }
+}
+
+// -----
+
+#layout1 = #test.memref_layout<"layout_a">
+#layout2 = #test.memref_layout<"layout_b">
+
+// CHECK-LABEL: @BufferizeLayoutMismatchInsideScfFor
+module @BufferizeLayoutMismatchInsideScfFor {
+ // CHECK: func.func @mismatch_in_for
+ // CHECK-SAME: -> memref<10xf32>
+ func.func @mismatch_in_for(
+ %lb: index, %ub: index, %step: index)
+ -> tensor<10xf32> {
+ // CHECK: %[[init:.*]] = "test.create_memref_op"
+ // CHECK-SAME: -> memref<10xf32, #test.memref_layout<"layout_a">>
+ %init = "test.tensor_with_future_layout"() {layout = #layout1}
+ : () -> tensor<10xf32>
+
+ // CHECK: %[[loop:.+]] = scf.for
+ // CHECK-SAME: iter_args(%[[iter:.*]] = %[[init]])
+ // CHECK-SAME: -> (memref<10xf32, #test.memref_layout<"layout_a">>)
+ %loop = scf.for %i = %lb to %ub step %step
+ iter_args(%iter = %init) -> (tensor<10xf32>) {
+ // CHECK: %[[conflict:.*]] = "test.dummy_memref_op"(%[[iter]])
+ // CHECK-SAME: -> memref<10xf32, #test.memref_layout<"layout_b">>
+ // CHECK: %[[cast:.*]] = memref.cast %[[conflict]]
+ // CHECK-SAME: to memref<10xf32, #test.memref_layout<"layout_a">>
+ %conflict = "test.force_new_layout"(%iter) {layout = #layout2}
+ : (tensor<10xf32>) -> tensor<10xf32>
+ // CHECK: scf.yield %[[cast]]
+ scf.yield %conflict : tensor<10xf32>
+ }
+
+ // CHECK: %[[out:.*]] = memref.cast %[[loop]]
+ // CHECK-SAME: to memref<10xf32>
+ // CHECK: return %[[out]]
+ return %loop : tensor<10xf32>
+ }
+}
+
+// -----
+
+#layout1 = #test.memref_layout<"layout_a">
+#layout2 = #test.memref_layout<"layout_b">
+
+// Test that custom layout can co-exist in principle within an "end-to-end"
+// SCF example (`for { extract slice -> custom op -> insert slice }`) without
+// bufferization failing completely due to a layout mismatch. The fact that the
+// produced IR itself is rather dumb (e.g. memref.subview drops user-specified
+// layout) is out of scope for now.
+
+// CHECK-LABEL: @BufferizeLayoutMismatchInsideScfForWithSubviews
+module @BufferizeLayoutMismatchInsideScfForWithSubviews {
+ // CHECK: func.func @mismatch_in_for
+ // CHECK-SAME: -> memref<10xf32>
+ func.func @mismatch_in_for(
+ %lb: index, %ub: index, %step: index)
+ -> tensor<10xf32> {
+ %c0 = arith.constant 0 : index
+
+ // CHECK: %[[init:.*]] = "test.create_memref_op"
+ // CHECK-SAME: -> memref<10xf32, #test.memref_layout<"layout_a">>
+ %init = "test.tensor_with_future_layout"() {layout = #layout1}
+ : () -> tensor<10xf32>
+
+ // CHECK: %[[loop:.+]] = scf.for
+ // CHECK-SAME: iter_args(%[[inout:.*]] = %[[init]])
+ // CHECK-SAME: -> (memref<10xf32, #test.memref_layout<"layout_a">>)
+ %loop = scf.for %i = %lb to %ub step %step
+ iter_args(%inout = %init) -> (tensor<10xf32>) {
+ // CHECK: %[[in:.*]] = memref.subview %[[inout]]
+ // CHECK-SAME: to memref<5xf32, strided<[1], offset: ?>>
+ %in = tensor.extract_slice %inout[%c0] [5] [1]
+ : tensor<10xf32> to tensor<5xf32>
+
+ // CHECK: %[[conflict:.*]] = "test.dummy_memref_op"(%[[in]])
+ // CHECK-SAME: -> memref<5xf32, #test.memref_layout<"layout_b">>
+ %conflict = "test.force_new_layout"(%in) {layout = #layout2}
+ : (tensor<5xf32>) -> tensor<5xf32>
+
+ // CHECK: %[[out:.*]] = memref.subview %[[inout]]
+ // CHECK-SAME: to memref<5xf32, strided<[1], offset: ?>>
+ // CHECK: memref.copy %[[conflict]], %[[out]]
+ %out = tensor.insert_slice %conflict into %inout[%c0] [5] [1]
+ : tensor<5xf32> into tensor<10xf32>
+
+ // CHECK: scf.yield %[[inout]]
+ scf.yield %out : tensor<10xf32>
+ }
+
+ // CHECK: %[[out:.*]] = memref.cast %[[loop]]
+ // CHECK-SAME: to memref<10xf32>
+ // CHECK: return %[[out]]
+ return %loop : tensor<10xf32>
+ }
+}
+
+// -----
+
+#layout1 = #test.memref_layout<"layout_a">
+#layout2 = #test.memref_layout<"layout_b">
+
+// CHECK-LABEL: @BufferizeLayoutMismatchInScfExecuteRegion
+module @BufferizeLayoutMismatchInScfExecuteRegion {
+ // CHECK: func.func @mismatch_in_scf_execute_region
+ // CHECK-SAME: -> memref<10xf32>
+ func.func @mismatch_in_scf_execute_region(%cond: i1) -> tensor<10xf32> {
+ // CHECK: %[[region:.*]] = scf.execute_region
+ // CHECK-SAME: -> memref<10xf32, #test.memref_layout<"layout_a">>
+ %out = scf.execute_region -> tensor<10xf32> {
+ cf.cond_br %cond, ^bb1, ^bb2
+
+ ^bb1:
+ // CHECK: %[[one:.*]] = "test.create_memref_op"
+ // CHECK-SAME: -> memref<10xf32, #test.memref_layout<"layout_a">>
+ %one = "test.tensor_with_future_layout"() {layout = #layout1}
+ : () -> tensor<10xf32>
+ // CHECK: cf.br ^bb3(%[[one]]
+ cf.br ^bb3(%one : tensor<10xf32>)
+
+ ^bb2:
+ // CHECK: %[[another:.*]] = "test.create_memref_op"
+ // CHECK-SAME: -> memref<10xf32, #test.memref_layout<"layout_b">>
+ // CHECK: %[[cast:.*]] = memref.cast %[[another]]
+ // CHECK-SAME: to memref<10xf32, #test.memref_layout<"layout_a">>
+ %another = "test.tensor_with_future_layout"() {layout = #layout2}
+ : () -> tensor<10xf32>
+ // CHECK: cf.br ^bb3(%[[cast]]
+ cf.br ^bb3(%another : tensor<10xf32>)
+
+ ^bb3(%res: tensor<10xf32>):
+ // CHECK: scf.yield {{%.*}} : memref<10xf32, #test.memref_layout<"layout_a">>
+ scf.yield %res : tensor<10xf32>
+ }
+
+ // CHECK: %[[out:.*]] = memref.cast %[[region]]
+ // CHECK-SAME: to memref<10xf32>
+ // CHECK: return %[[out]]
+ return %out : tensor<10xf32>
+ }
+}
+
+// -----
+
+#layout1 = #test.memref_layout<"layout_a">
+#layout2 = #test.memref_layout<"layout_b">
+
+// CHECK-LABEL: @BufferizeLayoutMismatchInArithSelect
+module @BufferizeLayoutMismatchInArithSelect {
+ // CHECK: func.func @mismatch_in_select(%[[cond:.*]]: i1)
+ // CHECK-SAME: -> memref<10xf32>
+ func.func @mismatch_in_select(%cond: i1) -> tensor<10xf32> {
+ // CHECK: %[[memref1:.*]] = "test.create_memref_op"
+ // CHECK-SAME: -> memref<10xf32, #test.memref_layout<"layout_a">>
+ %tensor1 = "test.tensor_with_future_layout"() {layout = #layout1}
+ : () -> tensor<10xf32>
+ // CHECK: %[[memref2:.*]] = "test.create_memref_op"
+ // CHECK-SAME: -> memref<10xf32, #test.memref_layout<"layout_b">>
+ %tensor2 = "test.tensor_with_future_layout"() {layout = #layout2}
+ : () -> tensor<10xf32>
+
+ // CHECK: %[[cast:.*]] = memref.cast %[[memref2]]
+ // CHECK-SAME: to memref<10xf32, #test.memref_layout<"layout_a">>
+
+ // CHECK: %[[select:.*]] = arith.select %[[cond]], %[[memref1]], %[[cast]]
+ %select = arith.select %cond, %tensor1, %tensor2 : i1, tensor<10xf32>
+
+ // CHECK: %[[out:.*]] = memref.cast %[[select]]
+ // CHECK-SAME: to memref<10xf32>
+ // CHECK: return %[[out]]
+ return %select : tensor<10xf32>
+ }
+}
diff --git a/mlir/test/Dialect/ControlFlow/one-shot-bufferize-invalid.mlir b/mlir/test/Dialect/ControlFlow/one-shot-bufferize-invalid.mlir
index 479497ccb6f44..e13ce690d30e9 100644
--- a/mlir/test/Dialect/ControlFlow/one-shot-bufferize-invalid.mlir
+++ b/mlir/test/Dialect/ControlFlow/one-shot-bufferize-invalid.mlir
@@ -1,7 +1,7 @@
// RUN: mlir-opt -one-shot-bufferize="bufferize-function-boundaries" -split-input-file %s -verify-diagnostics
// expected-error @below{{failed to bufferize op}}
-// expected-error @below{{incoming operands of block argument have inconsistent memory spaces}}
+// expected-error @below{{incoming operands of block argument have incompatible buffer types}}
func.func @inconsistent_memory_space() -> tensor<5xf32> {
%0 = bufferization.alloc_tensor() {memory_space = 0 : ui64} : tensor<5xf32>
cf.br ^bb1(%0: tensor<5xf32>)
diff --git a/mlir/test/Dialect/SCF/one-shot-bufferize-invalid.mlir b/mlir/test/Dialect/SCF/one-shot-bufferize-invalid.mlir
index 9d206b1895f65..855ef00f6a614 100644
--- a/mlir/test/Dialect/SCF/one-shot-bufferize-invalid.mlir
+++ b/mlir/test/Dialect/SCF/one-shot-bufferize-invalid.mlir
@@ -5,7 +5,7 @@ func.func @inconsistent_memory_space_scf_if(%c: i1) -> tensor<10xf32> {
// bufferized.
%0 = bufferization.alloc_tensor() {memory_space = 0 : ui64} : tensor<10xf32>
%1 = bufferization.alloc_tensor() {memory_space = 1 : ui64} : tensor<10xf32>
- // expected-error @+1 {{inconsistent memory space on then/else branches}}
+ // expected-error @+1 {{incompatible buffer types on then/else branches}}
%r = scf.if %c -> tensor<10xf32> {
// expected-error @+1 {{failed to bufferize op}}
scf.yield %0 : tensor<10xf32>
@@ -46,7 +46,7 @@ func.func @execute_region_no_yield(%t: tensor<5xf32>) -> tensor<5xf32> {
func.func @inconsistent_memory_space_scf_for(%lb: index, %ub: index, %step: index) -> tensor<10xf32> {
%0 = bufferization.alloc_tensor() {memory_space = 0 : ui64} : tensor<10xf32>
%1 = bufferization.alloc_tensor() {memory_space = 1 : ui64} : tensor<10xf32>
- // expected-error @below{{init_arg and yielded value bufferize to inconsistent memory spaces}}
+ // expected-error @below{{init_arg and yielded value bufferize to incompatible buffer types}}
%2 = scf.for %iv = %lb to %ub step %step iter_args(%arg = %0) -> tensor<10xf32> {
// expected-error @below {{failed to bufferize op}}
scf.yield %1 : tensor<10xf32>
diff --git a/mlir/test/lib/Dialect/Bufferization/TestOneShotModuleBufferize.cpp b/mlir/test/lib/Dialect/Bufferization/TestOneShotModuleBufferize.cpp
index 87f95e36e9183..00eb2dc92e14f 100644
--- a/mlir/test/lib/Dialect/Bufferization/TestOneShotModuleBufferize.cpp
+++ b/mlir/test/lib/Dialect/Bufferization/TestOneShotModuleBufferize.cpp
@@ -35,8 +35,8 @@ struct TestOneShotModuleBufferizePass
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestOneShotModuleBufferizePass)
TestOneShotModuleBufferizePass() = default;
- TestOneShotModuleBufferizePass(const TestOneShotModuleBufferizePass &pass) =
- default;
+ TestOneShotModuleBufferizePass(const TestOneShotModuleBufferizePass &pass)
+ : PassWrapper(pass) {}
void getDependentDialects(DialectRegistry ®istry) const override {
registry.insert<test::TestDialect>();
@@ -91,6 +91,30 @@ struct TestOneShotModuleBufferizePass
return bufferization::BufferLikeType{};
});
};
+ // A simple yet distinct (from upstream) policy: compare layouts and return
+ // "smaller" one.
+ opt.reconcileBufferTypeMismatchFn =
+ [](bufferization::BufferLikeType x, bufferization::BufferLikeType y,
+ const bufferization::BufferizationOptions &)
+ -> FailureOr<bufferization::BufferLikeType> {
+ auto getLayout = [](bufferization::BufferLikeType t) {
+ auto m = dyn_cast<MemRefType>(t);
+ return m ? dyn_cast<test::TestMemRefLayoutAttr>(m.getLayout())
+ : test::TestMemRefLayoutAttr();
+ };
+ auto lhsLayout = getLayout(x);
+ auto rhsLayout = getLayout(y);
+ if (lhsLayout && rhsLayout) {
+ return lhsLayout.getDummy().getValue() <=
+ rhsLayout.getDummy().getValue()
+ ? x
+ : y;
+ }
+ return rhsLayout ? y : x;
+ };
+ // Function signature update only works with memref.cast. Disable it to
+ // align behaviour for upstream and user casts.
+ opt.inferFunctionResultLayout = false;
bufferization::BufferizationState bufferizationState;
diff --git a/mlir/test/lib/Dialect/Test/TestOpDefs.cpp b/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
index 4a274090798c1..20d72a2fd09ca 100644
--- a/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
+++ b/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
@@ -2047,6 +2047,77 @@ test::TestCreateTensorOp::getBufferType(
options);
}
+::mlir::LogicalResult test::TestTensorWithFutureLayoutOp::bufferize(
+ ::mlir::RewriterBase &rewriter,
+ const ::mlir::bufferization::BufferizationOptions &options,
+ ::mlir::bufferization::BufferizationState &state) {
+ const auto bufferizedOutType =
+ mlir::bufferization::getBufferType(getOutput(), options, state);
+ if (mlir::failed(bufferizedOutType))
+ return failure();
+
+ auto createMemrefOp =
+ test::TestCreateMemrefOp::create(rewriter, getLoc(), *bufferizedOutType);
+ mlir::bufferization::replaceOpWithBufferizedValues(
+ rewriter, getOperation(), createMemrefOp.getResult());
+ return mlir::success();
+}
+
+mlir::FailureOr<mlir::bufferization::BufferLikeType>
+test::TestTensorWithFutureLayoutOp::getBufferType(
+ mlir::Value value, const mlir::bufferization::BufferizationOptions &,
+ const mlir::bufferization::BufferizationState &,
+ llvm::SmallVector<::mlir::Value> &) {
+ auto tensorType = dyn_cast<RankedTensorType>(value.getType());
+ if (!tensorType)
+ return failure();
+
+ // Set the memref layout to the op's 'layout' attribute, ignoring
+ // any pre-existing tensor encoding. This is what lets two
+ // `test.tensor_with_layout` ops produce *bufferized* memrefs with
diff erent
+ // layouts while keeping their *tensor* result types identical -- which is
+ // required to construct SCF iter_arg/branch mismatches that the verifier
+ // still accepts.
+ auto layout = cast<MemRefLayoutAttrInterface>(getLayout());
+ return cast<bufferization::BufferLikeType>(MemRefType::get(
+ tensorType.getShape(), tensorType.getElementType(), layout));
+}
+
+LogicalResult test::TestForceNewLayoutOp::bufferize(
+ RewriterBase &rewriter, const bufferization::BufferizationOptions &options,
+ bufferization::BufferizationState &state) {
+ auto buffer = bufferization::getBuffer(rewriter, getInput(), options, state);
+ if (failed(buffer))
+ return failure();
+
+ const auto bufferizedOutType =
+ bufferization::getBufferType(getOutput(), options, state);
+ if (failed(bufferizedOutType))
+ return failure();
+
+ auto dummyMemrefOp = test::TestDummyMemrefOp::create(
+ rewriter, getLoc(), *bufferizedOutType, *buffer);
+ bufferization::replaceOpWithBufferizedValues(rewriter, getOperation(),
+ dummyMemrefOp.getResult());
+ return success();
+}
+
+FailureOr<bufferization::BufferLikeType>
+test::TestForceNewLayoutOp::getBufferType(
+ mlir::Value value, const bufferization::BufferizationOptions &options,
+ const bufferization::BufferizationState &,
+ llvm::SmallVector<::mlir::Value> &) {
+ auto tensorType = dyn_cast<RankedTensorType>(value.getType());
+ if (!tensorType)
+ return failure();
+
+ // Set the memref layout to the op's 'layout' attribute, ignoring any
+ // pre-existing tensor encoding.
+ auto layout = cast<MemRefLayoutAttrInterface>(getLayout());
+ return cast<bufferization::BufferLikeType>(MemRefType::get(
+ tensorType.getShape(), tensorType.getElementType(), layout));
+}
+
// Define a custom builder for ManyRegionsOp declared in TestOps.td.
// OpBuilder<(ins "::std::unique_ptr<::mlir::Region>":$firstRegion,
// "::std::unique_ptr<::mlir::Region>":$secondRegion)>
diff --git a/mlir/test/lib/Dialect/Test/TestOps.td b/mlir/test/lib/Dialect/Test/TestOps.td
index 31002d1f17a75..df321bd5feabc 100644
--- a/mlir/test/lib/Dialect/Test/TestOps.td
+++ b/mlir/test/lib/Dialect/Test/TestOps.td
@@ -1009,7 +1009,7 @@ def OpWithResultShapeInterfaceOp : TEST_Op<"op_with_result_shape_interface",
def ReifyShapedTypeUsingReifyResultShapesOp :
TEST_Op<"reify_shaped_type_using_reify_result_shapes",
- [DeclareOpInterfaceMethods<ReifyRankedShapedTypeOpInterface,
+ [DeclareOpInterfaceMethods<ReifyRankedShapedTypeOpInterface,
["reifyResultShapes"]>]> {
let description = [{
Test that when resolving a single dimension of a result for an operation
@@ -1024,7 +1024,7 @@ def ReifyShapedTypeUsingReifyResultShapesOp :
def ReifyShapedTypeUsingReifyShapeOfResultOp :
TEST_Op<"reify_shaped_type_using_reify_shape_of_result",
- [DeclareOpInterfaceMethods<ReifyRankedShapedTypeOpInterface,
+ [DeclareOpInterfaceMethods<ReifyRankedShapedTypeOpInterface,
["reifyResultShapes", "reifyShapeOfResult"]>]> {
let description = [{
Test that when resolving a single dimension of a result for an operation
@@ -1040,7 +1040,7 @@ def ReifyShapedTypeUsingReifyShapeOfResultOp :
def ReifyShapedTypeUsingReifyDimOfResultOp :
TEST_Op<"reify_shaped_type_using_reify_dim_of_result",
- [DeclareOpInterfaceMethods<ReifyRankedShapedTypeOpInterface,
+ [DeclareOpInterfaceMethods<ReifyRankedShapedTypeOpInterface,
["reifyResultShapes", "reifyShapeOfResult", "reifyDimOfResult"]>]> {
let description = [{
Test that when resolving a single dimension of a result for an operation
@@ -4198,6 +4198,103 @@ def TestCreateMemrefOp : TEST_Op<"create_memref_op"> {
let results = (outs Arg<Bufferization_BufferLikeTypeInterface>:$output);
}
+def TestTensorWithFutureLayoutOp
+ : TEST_Op<"tensor_with_future_layout",
+ [DeclareOpInterfaceMethods<
+ BufferizableOpInterface,
+ ["bufferize", "getBufferType", "bufferizesToMemoryRead",
+ "bufferizesToMemoryWrite", "getAliasingValues",
+ "getAliasingOpOperands", "bufferizesToAllocation"]>]> {
+ let description = [{
+ A tensor operation that creates a new value and then forces a specific
+ `TestMemRefLayoutAttr` layout during bufferization into a memref. In tensor
+ world, this op does nothing encoding/layout wise.
+ }];
+
+ let arguments = (ins TestMemRefLayoutAttr:$layout);
+ let results = (outs AnyRankedTensor:$output);
+ let extraClassDefinition = [{
+ bool test::TestTensorWithFutureLayoutOp::bufferizesToMemoryRead(
+ ::mlir::OpOperand&,
+ const ::mlir::bufferization::AnalysisState&) {
+ return false;
+ }
+ bool test::TestTensorWithFutureLayoutOp::bufferizesToMemoryWrite(
+ ::mlir::OpOperand&,
+ const ::mlir::bufferization::AnalysisState&) {
+ return false;
+ }
+ bool test::TestTensorWithFutureLayoutOp::bufferizesToAllocation(
+ mlir::Value) {
+ return true;
+ }
+ ::mlir::bufferization::AliasingValueList
+ test::TestTensorWithFutureLayoutOp::getAliasingValues(::mlir::OpOperand&,
+ const ::mlir::bufferization::AnalysisState&) {
+ return {};
+ }
+ ::mlir::bufferization::AliasingOpOperandList
+ test::TestTensorWithFutureLayoutOp::getAliasingOpOperands(::mlir::Value,
+ const ::mlir::bufferization::AnalysisState&) {
+ return {};
+ }
+ }];
+}
+
+def TestForceNewLayoutOp
+ : TEST_Op<"force_new_layout",
+ [DeclareOpInterfaceMethods<
+ BufferizableOpInterface,
+ ["bufferize", "getBufferType", "bufferizesToMemoryRead",
+ "bufferizesToMemoryWrite", "getAliasingValues",
+ "getAliasingOpOperands"]>]> {
+ let description = [{
+ A tensor operation that modifies a pre-existing value and then forces a
+ specific `TestMemRefLayoutAttr` to be set during bufferization into a
+ memref. In tensor world, this op does nothing encoding/layout wise. It is
+ similar to TestTensorWithFutureLayoutOp, but it keeps the IR values
+ "connected" which can be particularly useful in e.g. SCF loop's body (one
+ can modify an initArg's layout and yield the new-layout value).
+ }];
+
+ let arguments = (ins
+ AnyRankedTensor:$input,
+ TestMemRefLayoutAttr:$layout
+ );
+ let results = (outs AnyRankedTensor:$output);
+ let extraClassDefinition = [{
+ bool test::TestForceNewLayoutOp::bufferizesToMemoryRead(::mlir::OpOperand&,
+ const ::mlir::bufferization::AnalysisState&) {
+ return true;
+ }
+ bool test::TestForceNewLayoutOp::bufferizesToMemoryWrite(::mlir::OpOperand&,
+ const ::mlir::bufferization::AnalysisState&) {
+ return true;
+ }
+ ::mlir::bufferization::AliasingValueList
+ test::TestForceNewLayoutOp::getAliasingValues(::mlir::OpOperand&,
+ const ::mlir::bufferization::AnalysisState&) {
+ auto relation = getInput().getType() == getOutput().getType()
+ ? ::mlir::bufferization::BufferRelation::Equivalent
+ : ::mlir::bufferization::BufferRelation::Unknown;
+ return {{getOutput(), relation, /*isDefinite=*/true}};
+ }
+
+ ::mlir::bufferization::AliasingOpOperandList
+ test::TestForceNewLayoutOp::getAliasingOpOperands(::mlir::Value value,
+ const ::mlir::bufferization::AnalysisState&) {
+ if (value != getOutput())
+ return {};
+
+ auto relation = getInput().getType() == getOutput().getType()
+ ? ::mlir::bufferization::BufferRelation::Equivalent
+ : ::mlir::bufferization::BufferRelation::Unknown;
+ return {{&getOperation()->getOpOperand(0), relation,
+ /*isDefinite=*/true}};
+ }
+ }];
+}
+
//===----------------------------------------------------------------------===//
// Test assembly format references
//===----------------------------------------------------------------------===//
More information about the Mlir-commits
mailing list