[Mlir-commits] [mlir] [mlir][bufferization] Introduce reconcileBufferTypeMismatchFn hook (PR #202667)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Tue Jun 9 07:09:34 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Andrei Golubev (andrey-golubev)
<details>
<summary>Changes</summary>
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).
---
Patch is 44.30 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/202667.diff
12 Files Affected:
- (modified) mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h (+15)
- (modified) mlir/include/mlir/Dialect/Bufferization/IR/UnstructuredControlFlow.h (+29-23)
- (modified) mlir/lib/Dialect/Arith/Transforms/BufferizableOpInterfaceImpl.cpp (+8-13)
- (modified) mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp (+31-1)
- (modified) mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp (+28-50)
- (modified) mlir/test/Dialect/Arith/one-shot-bufferize-memory-space-invalid.mlir (+1-1)
- (modified) mlir/test/Dialect/Bufferization/Transforms/test-one-shot-module-bufferize.mlir (+292)
- (modified) mlir/test/Dialect/ControlFlow/one-shot-bufferize-invalid.mlir (+1-1)
- (modified) mlir/test/Dialect/SCF/one-shot-bufferize-invalid.mlir (+2-2)
- (modified) mlir/test/lib/Dialect/Bufferization/TestOneShotModuleBufferize.cpp (+31-2)
- (modified) mlir/test/lib/Dialect/Test/TestOpDefs.cpp (+73)
- (modified) mlir/test/lib/Dialect/Test/TestOps.td (+101-3)
``````````diff
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..d5f11a2968c16 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,36 @@ 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..d6bd113059b47 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);
@@ -575,28 +556,25 @@ static FailureOr<BufferLikeType> computeLoopRegionIterArgBufferType(
if (*initArgBufferType == yieldedValueBufferType)
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");
+ // If there is a mismatch between the yielded buffer type and the init_arg
+ // 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">>...
[truncated]
``````````
</details>
https://github.com/llvm/llvm-project/pull/202667
More information about the Mlir-commits
mailing list