[Mlir-commits] [mlir] [mlir][bufferization] Introduce reconcileBufferTypeMismatchFn hook (PR #202667)
Andrei Golubev
llvmlistbot at llvm.org
Tue Jun 9 07:12:31 PDT 2026
https://github.com/andrey-golubev updated https://github.com/llvm/llvm-project/pull/202667
>From a99ce557d0e3ea31e12bfe77d1c98a0753a9540a Mon Sep 17 00:00:00 2001
From: "Golubev, Andrey" <andrey.golubev at intel.com>
Date: Tue, 12 May 2026 09:33:07 +0000
Subject: [PATCH] [mlir][bufferization] Introduce reconcileBufferTypeMismatchFn
hook
This is 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>
---
.../IR/BufferizableOpInterface.h | 15 +
.../IR/UnstructuredControlFlow.h | 51 +--
.../BufferizableOpInterfaceImpl.cpp | 21 +-
.../IR/BufferizableOpInterface.cpp | 32 +-
.../BufferizableOpInterfaceImpl.cpp | 76 ++---
...e-shot-bufferize-memory-space-invalid.mlir | 2 +-
.../test-one-shot-module-bufferize.mlir | 292 ++++++++++++++++++
.../one-shot-bufferize-invalid.mlir | 2 +-
.../SCF/one-shot-bufferize-invalid.mlir | 4 +-
.../TestOneShotModuleBufferize.cpp | 33 +-
mlir/test/lib/Dialect/Test/TestOpDefs.cpp | 73 +++++
mlir/test/lib/Dialect/Test/TestOps.td | 104 ++++++-
12 files changed, 610 insertions(+), 95 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h
index bf94f56497160..6bdaaa901ab05 100644
--- a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h
+++ b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h
@@ -273,6 +273,14 @@ struct BufferizationOptions {
using DefaultMemorySpaceFn =
std::function<std::optional<Attribute>(TensorLikeType t)>;
+ /// Resolve a mismatch between buffer types that were independently inferred
+ /// for the same bufferized value. Returns `failure()` to signal bufferization
+ /// failure; returns a buffer-like type when reconciliation suceeded. By
+ /// default, resolves into a memref with a fully dynamic layout.
+ using ReconcileBufferTypeMismatchFn = std::function<FailureOr<BufferLikeType>(
+ Operation *, BufferLikeType x, BufferLikeType y,
+ const BufferizationOptions &)>;
+
BufferizationOptions();
/// Try to cast the given op to BufferizableOpInterface if the op is allow
@@ -364,6 +372,13 @@ struct BufferizationOptions {
DefaultMemorySpaceFn defaultMemorySpaceFn =
[](TensorLikeType t) -> std::optional<Attribute> { return Attribute(); };
+ /// Hook to reconcile two buffer types that were independently inferred for
+ /// the same bufferized value (e.g. init_arg vs. yielded value in `scf.for`,
+ /// branches of `scf.if`, cases of `scf.index_switch`). 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..d97fe264565ca 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;
+ BufferLikeType callerType;
if (auto memrefType =
dyn_cast<BaseMemRefType>(opOperand->get().getType())) {
// The operand was already bufferized. Take its type directly.
- callerType = memrefType;
+ callerType = cast<BufferLikeType>(memrefType);
} 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,28 +81,35 @@ 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");
+ auto tensorType = value.getType();
+ if (isa<TensorType>(tensorType)) {
+ auto bufferMemRefType = llvm::cast<BaseMemRefType>(bufferType);
+ auto callerMemRefType = llvm::cast<BaseMemRefType>(callerType);
+ if (auto rankedTensorType = dyn_cast<RankedTensorType>(tensorType)) {
+ assert(bufferMemRefType.hasRank() && callerMemRefType.hasRank() &&
+ "expected ranked memrefs");
+ assert(llvm::all_equal({bufferMemRefType.getShape(),
+ callerMemRefType.getShape(),
+ rankedTensorType.getShape()}) &&
+ "expected same shape");
+ } else {
+ assert(!bufferMemRefType.hasRank() && !callerMemRefType.hasRank() &&
+ "expected unranked memrefs");
+ }
}
#endif // NDEBUG
- if (bufferType.getMemorySpace() != callerType.getMemorySpace())
- return op->emitOpError("incoming operands of block argument have "
- "inconsistent memory spaces");
+ auto reconciled = options.reconcileBufferTypeMismatchFn(
+ op, cast<BufferLikeType>(bufferType),
+ cast<BufferLikeType>(callerType), options);
+ if (failed(reconciled)) {
+ return failure();
+ }
- bufferType = getMemRefTypeWithFullyDynamicLayout(
- tensorType, bufferType.getMemorySpace());
+ bufferType = *reconciled;
}
if (!bufferType)
diff --git a/mlir/lib/Dialect/Arith/Transforms/BufferizableOpInterfaceImpl.cpp b/mlir/lib/Dialect/Arith/Transforms/BufferizableOpInterfaceImpl.cpp
index 91d03a0356de1..49e29665c60e1 100644
--- a/mlir/lib/Dialect/Arith/Transforms/BufferizableOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Arith/Transforms/BufferizableOpInterfaceImpl.cpp
@@ -163,10 +163,11 @@ struct SelectOpInterface
// The "true" and the "false" operands must have the same type. If the
// buffers have different types, they differ 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));
+ auto targetType = options.reconcileBufferTypeMismatchFn(
+ selectOp, cast<BufferLikeType>(trueBuffer.getType()),
+ cast<BufferLikeType>(falseBuffer.getType()), options);
if (failed(targetType))
return failure();
if (trueBuffer.getType() != *targetType)
@@ -198,16 +199,10 @@ struct SelectOpInterface
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 different types, they differ 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(
+ op, cast<BufferLikeType>(*trueType), cast<BufferLikeType>(*falseType),
+ options);
}
};
diff --git a/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp b/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
index e60736a20dc3b..a3c44c1cd88a3 100644
--- a/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
+++ b/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
@@ -368,12 +368,42 @@ defaultUnknownTypeConverter(TensorLikeType tensorType, Attribute memorySpace,
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(Operation *op, BufferLikeType x,
+ BufferLikeType y,
+ const BufferizationOptions &) {
+ const auto xMemRef = cast<BaseMemRefType>(x);
+ const auto yMemRef = cast<BaseMemRefType>(y);
+
+ if (xMemRef.getMemorySpace() != yMemRef.getMemorySpace())
+ return op->emitError(
+ "inconsistent memory spaces in buffers provided for reconciliation");
+
+ if (isa<UnrankedMemRefType>(xMemRef)) {
+ // unranked memrefs have no layout.
+ return x;
+ }
+
+ const auto xRankedMemref = cast<MemRefType>(xMemRef);
+ int64_t dynamicOffset = ShapedType::kDynamic;
+ SmallVector<int64_t> dynamicStrides(xRankedMemref.getRank(),
+ ShapedType::kDynamic);
+ auto stridedLayout = StridedLayoutAttr::get(xRankedMemref.getContext(),
+ dynamicOffset, dynamicStrides);
+ return cast<BufferLikeType>(
+ MemRefType::get(xRankedMemref.getShape(), xRankedMemref.getElementType(),
+ stridedLayout, xRankedMemref.getMemorySpace()));
+}
+
} // namespace
// 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
diff --git a/mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp b/mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp
index 16eb9aadc06f0..df5b761342d7c 100644
--- a/mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp
@@ -315,20 +315,8 @@ 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 different: Promote to fully dynamic layout
- // map.
- return cast<BufferLikeType>(getMemRefTypeWithFullyDynamicLayout(
- cast<TensorType>(opResult.getType()),
- thenBaseMemRefType.getMemorySpace()));
+ return options.reconcileBufferTypeMismatchFn(op, thenBufferType,
+ elseBufferType, options);
}
};
@@ -407,23 +395,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 +420,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 different: Promote to fully dynamic
- // layout map.
- bufferType = getMemRefTypeWithFullyDynamicLayout(
- cast<TensorType>(value.getType()), bufferType.getMemorySpace());
+ auto reconciled = options.reconcileBufferTypeMismatchFn(
+ op, bufferType, *yieldedBufferType, options);
+ if (failed(reconciled))
+ return failure();
+ bufferType = *reconciled;
}
return cast<BufferLikeType>(bufferType);
@@ -576,27 +557,24 @@ 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<TensorType>(iterArg.getType())) {
+ auto yieldedBufferType = cast<BaseMemRefType>(yieldedValueBufferType);
+ auto initBufferType = cast<BaseMemRefType>(*initArgBufferType);
+ if (auto yieldedRankedBufferType =
+ dyn_cast<MemRefType>(yieldedBufferType)) {
+ assert(llvm::all_equal(
+ {yieldedRankedBufferType.getShape(),
+ cast<MemRefType>(initBufferType).getShape(),
+ cast<RankedTensorType>(iterTensorType).getShape()}) &&
+ "expected same shape");
+ }
}
#endif // NDEBUG
- // TODO: Properly support with options, for now it is hardcoded MemRef type
- // based approach
- return cast<BufferLikeType>(getMemRefTypeWithFullyDynamicLayout(
- iterTensorType, yieldedBufferType.getMemorySpace()));
+
+ return options.reconcileBufferTypeMismatchFn(loopOp, *initArgBufferType,
+ yieldedValueBufferType, options);
}
/// 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..480fde7f87417 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 {{inconsistent memory spaces in buffers provided for reconciliation}}
// 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..4b5c3d963c5e5 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
@@ -1,8 +1,10 @@
// RUN: mlir-opt %s -test-one-shot-module-bufferize -split-input-file | FileCheck %s
+// RUN: mlir-opt %s -test-one-shot-module-bufferize="infer-function-result-layout=0" -split-input-file | FileCheck %s --check-prefix=PRESERVE-SIGNATURE
#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 +274,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 +299,292 @@ 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, #test.memref_layout<"layout_a">>
+ func.func @layout_for_func() -> tensor<10xf32> {
+ // CHECK: %[[memref:.*]] = "test.create_memref_op"
+ // CHECK-SAME: -> memref<10xf32, #test.memref_layout<"layout_a">>
+ // CHECK: return %[[memref]]
+ %memref = "test.tensor_with_future_layout"() {layout = #layout1}
+ : () -> tensor<10xf32>
+ return %memref : tensor<10xf32>
+ }
+
+ // PRESERVE-SIGNATURE: func.func @layout_for_func
+ // PRESERVE-SIGNATURE-SAME: -> memref<10xf32>
+ // PRESERVE-SIGNATURE: %[[memref:.*]] = "test.create_memref_op"
+ // PRESERVE-SIGNATURE-SAME: -> memref<10xf32, #test.memref_layout<"layout_a">>
+ // PRESERVE-SIGNATURE: %[[cast:.*]] = memref.cast %[[memref]]
+ // PRESERVE-SIGNATURE-SAME: to memref<10xf32>
+ // PRESERVE-SIGNATURE: return %[[cast]]
+}
+
+// -----
+
+#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, #test.memref_layout<"layout_a">>
+ func.func @mismatch_in_if(%cond: i1) -> tensor<10xf32> {
+ // CHECK: %[[ret:.*]] = 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: return %[[ret]]
+ return %ret : tensor<10xf32>
+ }
+
+ // PRESERVE-SIGNATURE: func.func @mismatch_in_if
+ // PRESERVE-SIGNATURE-SAME: -> memref<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, #test.memref_layout<"layout_a">>
+ func.func @mismatch_in_switch(%idx: index) -> tensor<10xf32> {
+ // CHECK: %[[ret:.*]] = 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: return %[[ret]]
+ return %ret : tensor<10xf32>
+ }
+
+ // PRESERVE-SIGNATURE: func.func @mismatch_in_switch
+ // PRESERVE-SIGNATURE-SAME: -> memref<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, #test.memref_layout<"layout_a">>
+ 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: return %[[loop]]
+ return %loop : tensor<10xf32>
+ }
+
+ // PRESERVE-SIGNATURE: func.func @mismatch_in_for
+ // PRESERVE-SIGNATURE-SAME: -> memref<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, #test.memref_layout<"layout_a">>
+ 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: return %[[loop]]
+ return %loop : tensor<10xf32>
+ }
+
+ // PRESERVE-SIGNATURE: func.func @mismatch_in_for
+ // PRESERVE-SIGNATURE-SAME: -> memref<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, #test.memref_layout<"layout_a">>
+ func.func @mismatch_in_scf_execute_region(%cond: i1) -> tensor<10xf32> {
+ // CHECK: %[[out:.*]] = 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: return %[[out]]
+ return %out : tensor<10xf32>
+ }
+
+ // PRESERVE-SIGNATURE: func.func @mismatch_in_scf_execute_region
+ // PRESERVE-SIGNATURE-SAME: -> memref<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, #test.memref_layout<"layout_a">>
+ 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: return %[[select]]
+ return %select : tensor<10xf32>
+ }
+
+ // PRESERVE-SIGNATURE: func.func @mismatch_in_select
+ // PRESERVE-SIGNATURE-SAME: -> memref<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..44f2ef67e9893 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{{inconsistent memory spaces in buffers provided for reconciliation}}
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..e2a628505a0bc 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 {{inconsistent memory spaces in buffers provided for reconciliation}}
%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{{inconsistent memory spaces in buffers provided for reconciliation}}
%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..ad0f88483fbd7 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,29 @@ struct TestOneShotModuleBufferizePass
return bufferization::BufferLikeType{};
});
};
+ // A simple yet distinct (from upstream) policy: compare layouts and return
+ // "smaller" one.
+ opt.reconcileBufferTypeMismatchFn =
+ [](Operation *, 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;
+ };
+ opt.inferFunctionResultLayout = this->inferFunctionResultLayout;
bufferization::BufferizationState bufferizationState;
@@ -98,6 +121,12 @@ struct TestOneShotModuleBufferizePass
bufferizationState)))
signalPassFailure();
}
+
+ Option<bool> inferFunctionResultLayout{
+ *this, "infer-function-result-layout",
+ llvm::cl::desc(
+ "Allows to change the function signature. By default, set to true."),
+ llvm::cl::init(true)};
};
} // namespace
diff --git a/mlir/test/lib/Dialect/Test/TestOpDefs.cpp b/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
index 41f5a6f69fa21..6a71fe6d7cd60 100644
--- a/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
+++ b/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
@@ -2019,6 +2019,79 @@ 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 different
+ // 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));
+}
+
+::mlir::LogicalResult test::TestForceNewLayoutOp::bufferize(
+ ::mlir::RewriterBase &rewriter,
+ const ::mlir::bufferization::BufferizationOptions &options,
+ ::mlir::bufferization::BufferizationState &state) {
+ auto buffer =
+ mlir::bufferization::getBuffer(rewriter, getInput(), options, state);
+ if (mlir::failed(buffer))
+ return failure();
+
+ const auto bufferizedOutType =
+ mlir::bufferization::getBufferType(getOutput(), options, state);
+ if (mlir::failed(bufferizedOutType))
+ return failure();
+
+ auto dummyMemrefOp = test::TestDummyMemrefOp::create(
+ rewriter, getLoc(), *bufferizedOutType, *buffer);
+ mlir::bufferization::replaceOpWithBufferizedValues(rewriter, getOperation(),
+ dummyMemrefOp.getResult());
+ return mlir::success();
+}
+
+mlir::FailureOr<mlir::bufferization::BufferLikeType>
+test::TestForceNewLayoutOp::getBufferType(
+ mlir::Value value, const mlir::bufferization::BufferizationOptions &options,
+ 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.
+ 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 d008343e43a39..c2c83ad3b3ddb 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
@@ -4178,6 +4178,104 @@ 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 creation operation that forces a specific `TestMemRefLayoutAttr`
+ layout to be set during bufferization into a memref. In tensor world, this
+ op does nothing encoding/layout wise. Thus, it allows to create memref
+ layout mismatches that can later be resolved by the bufferization framework
+ while keeping tensor IR valid.
+ }];
+
+ 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 forces a specific `TestMemRefLayoutAttr` to be set
+ during bufferization into a memref. In tensor world, this op does nothing
+ encoding/layout wise. Thus, it allows to create memref layout mismatches
+ that can later be resolved by the bufferization framework while keeping
+ tensor IR valid.
+ }];
+
+ let arguments = (ins
+ Arg<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