[Mlir-commits] [mlir] 633344d - [mlir][bufferization] Introduce castBufferTypeFn hook and API (#205080)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Jul 1 03:54:59 PDT 2026
Author: Andrei Golubev
Date: 2026-07-01T12:54:54+02:00
New Revision: 633344de1ba723503ae3489933d1f318c0bb4473
URL: https://github.com/llvm/llvm-project/commit/633344de1ba723503ae3489933d1f318c0bb4473
DIFF: https://github.com/llvm/llvm-project/commit/633344de1ba723503ae3489933d1f318c0bb4473.diff
LOG: [mlir][bufferization] Introduce castBufferTypeFn hook and API (#205080)
This PR is the second part of the work that aims to allow customizations
in resolving mismatching buffer types.
(see 0b1fe74f079fd330e4d4d85ce33b7f5702da1be7 for part one)
Add a new bufferization hook that lets downstream bufferization
implementations materialize an operation to resolve a pre-reconciled
buffer type mismatch.
The hook + API is used as a fallback mechanism in several upstream
operations, as well as in the general bufferization infrastructure.
Semantically, this is a "cast" creation akin to `createAlloc()` and
`createMemCpy()` but for the cases where a memref::CastOp-like operation
is acceptable.
A valid value is returned whenever a cast-like operation is created,
failure signifies a failed attempt to materialize a cast (in theory,
there could be a fall back to alloc + copy either externally, but this
right now is not considered in most contexts). By default, a
`memref.cast` operation is created that casts a bufferized value to a
reconciled type.
Added:
Modified:
mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h
mlir/lib/Dialect/Arith/Transforms/BufferizableOpInterfaceImpl.cpp
mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp
mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp
mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp
mlir/test/Dialect/Bufferization/Transforms/test-one-shot-module-bufferize.mlir
mlir/test/lib/Dialect/Bufferization/TestOneShotModuleBufferize.cpp
mlir/test/lib/Dialect/Test/TestOpDefs.cpp
mlir/test/lib/Dialect/Test/TestOps.td
mlir/test/lib/Dialect/Test/TestTypeDefs.td
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h
index 9f7a6a638f307..9f49a31c1583b 100644
--- a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h
+++ b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h
@@ -258,6 +258,12 @@ struct BufferizationOptions {
/// Memcpy function: Generate a memcpy between two buffers.
using MemCpyFn =
std::function<LogicalResult(OpBuilder &, Location, Value, Value)>;
+ /// Cast function: Convert a buffer value to a new value with the specified
+ /// type. This method is typically used when a simple cast-like operation is
+ /// sufficient to convert the buffer value, for example, when layout maps
+ /// between buffer value and resulting type do not match.
+ using CastFn =
+ std::function<FailureOr<Value>(OpBuilder &, Location, Type, Value)>;
/// Initializer function for analysis state.
using AnalysisStateInitFn = std::function<void(AnalysisState &)>;
/// Tensor-like -> Buffer-like type conversion.
@@ -300,6 +306,7 @@ struct BufferizationOptions {
/// Helper functions for allocation and memory copying.
std::optional<AllocationFn> allocationFn;
std::optional<MemCpyFn> memCpyFn;
+ std::optional<CastFn> castFn;
/// Create a memref allocation with the given type and dynamic extents.
FailureOr<Value> createAlloc(OpBuilder &b, Location loc, MemRefType type,
@@ -309,6 +316,10 @@ struct BufferizationOptions {
LogicalResult createMemCpy(OpBuilder &b, Location loc, Value from,
Value to) const;
+ /// Creates a cast function from a buffer value to a new type.
+ FailureOr<Value> createCast(OpBuilder &b, Location loc, Type dest,
+ Value value) const;
+
/// Specifies whether not bufferizable ops are allowed in the input. If so,
/// bufferization.to_buffer and bufferization.to_tensor ops are inserted at
/// the boundaries.
diff --git a/mlir/lib/Dialect/Arith/Transforms/BufferizableOpInterfaceImpl.cpp b/mlir/lib/Dialect/Arith/Transforms/BufferizableOpInterfaceImpl.cpp
index 35e472470560a..f810000dd082f 100644
--- a/mlir/lib/Dialect/Arith/Transforms/BufferizableOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Arith/Transforms/BufferizableOpInterfaceImpl.cpp
@@ -172,12 +172,14 @@ struct SelectOpInterface
return selectOp->emitError(
"incompatible buffer types on true/false operands");
}
- if (trueBuffer.getType() != *targetType)
+ if (trueBuffer.getType() != *targetType) {
trueBuffer =
- memref::CastOp::create(rewriter, loc, *targetType, trueBuffer);
- if (falseBuffer.getType() != *targetType)
+ *options.createCast(rewriter, loc, *targetType, trueBuffer);
+ }
+ if (falseBuffer.getType() != *targetType) {
falseBuffer =
- memref::CastOp::create(rewriter, loc, *targetType, falseBuffer);
+ *options.createCast(rewriter, loc, *targetType, falseBuffer);
+ }
}
replaceOpWithNewBufferizedOp<arith::SelectOp>(
diff --git a/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp b/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
index e0a35e1932f3d..7f656164f3a15 100644
--- a/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
+++ b/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
@@ -840,6 +840,19 @@ LogicalResult BufferizationOptions::createMemCpy(OpBuilder &b, Location loc,
return success();
}
+FailureOr<Value> BufferizationOptions::createCast(OpBuilder &b, Location loc,
+ Type dest,
+ Value value) const {
+ if (castFn)
+ return (*castFn)(b, loc, dest, value);
+
+ assert(isa<BaseMemRefType>(dest) && "expected BaseMemRefType");
+ assert(isa<BaseMemRefType>(value.getType()) && "expected BaseMemRefType");
+ assert(memref::CastOp::areCastCompatible(value.getType(), dest) &&
+ "cast incompatible");
+ return memref::CastOp::create(b, loc, dest, value).getResult();
+}
+
//===----------------------------------------------------------------------===//
// Bufferization-specific IRMapping support with debugging.
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp b/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp
index 2a1dcc89a83ed..320f426985750 100644
--- a/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp
+++ b/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp
@@ -60,7 +60,7 @@ FailureOr<Value> mlir::bufferization::castOrReallocMemRefValue(
// a fix extra conditions in `isGuaranteedCastCompatible`.
if (memref::CastOp::areCastCompatible(srcType, destType) &&
isGuaranteedCastCompatible(srcType, destType)) {
- Value casted = memref::CastOp::create(b, value.getLoc(), destType, value);
+ Value casted = *options.createCast(b, value.getLoc(), destType, value);
return casted;
}
@@ -100,6 +100,20 @@ LogicalResult mlir::bufferization::foldToBufferToTensorPair(
return success();
}
+ if (!llvm::isa<BaseMemRefType>(srcType) ||
+ !llvm::isa<BaseMemRefType>(destType)) {
+ // Non-builtin case: the best is to try the user-provided cast.
+ assert(options.castFn.has_value() &&
+ "user-provided cast is required for non-builtin types");
+ auto replacement =
+ options.createCast(rewriter, bufferToTensor.getBuffer().getLoc(),
+ destType, bufferToTensor.getBuffer());
+ if (failed(replacement))
+ return failure();
+ rewriter.replaceOp(toBuffer, *replacement);
+ return success();
+ }
+
auto rankedSrcType = llvm::dyn_cast<MemRefType>(srcType);
auto rankedDestType = llvm::dyn_cast<MemRefType>(destType);
auto unrankedSrcType = llvm::dyn_cast<UnrankedMemRefType>(srcType);
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp b/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp
index 8bdee09d0d6e7..e0262abf313b8 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp
@@ -476,9 +476,10 @@ bufferization::bufferizeBlockSignature(Block *block, RewriterBase &rewriter,
rewriter, operand.getLoc(), *operandBufferType, operand);
// A cast is needed if the operand and the block argument have
diff erent
// bufferized types.
- if (type != *operandBufferType)
- bufferizedOperand = memref::CastOp::create(rewriter, operand.getLoc(),
- type, bufferizedOperand);
+ if (type != *operandBufferType) {
+ bufferizedOperand = *options.createCast(rewriter, operand.getLoc(),
+ type, bufferizedOperand);
+ }
newOperands.push_back(bufferizedOperand);
}
operands.getMutableForwardedOperands().assign(newOperands);
diff --git a/mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp b/mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp
index b9662751aa7f4..401be7440f5a2 100644
--- a/mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp
@@ -31,21 +31,13 @@ namespace {
/// Helper function for loop bufferization. Cast the given buffer to the given
/// memref type.
-static Value castBuffer(OpBuilder &b, Value buffer, Type type) {
+static Value castBuffer(OpBuilder &b, Value buffer, Type type,
+ const BufferizationOptions &options) {
// If the buffer already has the correct type, no cast is needed.
if (buffer.getType() == type)
return buffer;
- // TODO: Properly support with options, for now it is hardcoded MemRef type
- // based approach
- assert(isa<BaseMemRefType>(type) && "expected BaseMemRefType");
- assert(isa<BaseMemRefType>(buffer.getType()) && "expected BaseMemRefType");
- // TODO: In case `type` has a layout map that is not the fully dynamic
- // one, we may not be able to cast the buffer. In that case, the loop
- // iter_arg's layout map must be changed (see uses of `castBuffer`).
- assert(memref::CastOp::areCastCompatible(buffer.getType(), type) &&
- "scf.while op bufferization: cast incompatible");
- return memref::CastOp::create(b, buffer.getLoc(), type, buffer).getResult();
+ return *options.createCast(b, buffer.getLoc(), type, buffer);
}
/// Helper function for loop bufferization. Return "true" if the given value
@@ -114,7 +106,7 @@ struct ConditionOpInterface
whileOp.getAfterArguments()[it.index()], options, state);
if (failed(resultType))
return failure();
- Value buffer = castBuffer(rewriter, *maybeBuffer, *resultType);
+ Value buffer = castBuffer(rewriter, *maybeBuffer, *resultType, options);
newArgs.push_back(buffer);
} else {
newArgs.push_back(value);
@@ -206,7 +198,7 @@ struct ExecuteRegionOpInterface
rewriter.setInsertionPointAfter(newOp);
SmallVector<Value> newResults;
for (const auto &it : llvm::enumerate(executeRegionOp->getResultTypes())) {
- if (isa<TensorType>(it.value())) {
+ if (isa<TensorLikeType>(it.value())) {
newResults.push_back(bufferization::ToTensorOp::create(
rewriter, executeRegionOp.getLoc(), it.value(),
newOp->getResult(it.index())));
@@ -361,7 +353,7 @@ struct IndexSwitchOpInterface
// Compute bufferized result types.
SmallVector<Type> newTypes;
for (Value result : switchOp.getResults()) {
- if (!isa<TensorType>(result.getType())) {
+ if (!isa<TensorLikeType>(result.getType())) {
newTypes.push_back(result.getType());
continue;
}
@@ -762,7 +754,8 @@ struct ForOpInterface
auto targetType = bufferization::getBufferType(result, options, state);
if (failed(targetType))
return failure();
- castedInitArgs.push_back(castBuffer(rewriter, initArg, *targetType));
+ castedInitArgs.push_back(
+ castBuffer(rewriter, initArg, *targetType, options));
}
// Construct a new scf.for op with memref instead of tensor values.
@@ -986,7 +979,8 @@ struct WhileOpInterface
auto targetType = bufferization::getBufferType(beforeArg, options, state);
if (failed(targetType))
return failure();
- castedInitArgs.push_back(castBuffer(rewriter, initArg, *targetType));
+ castedInitArgs.push_back(
+ castBuffer(rewriter, initArg, *targetType, options));
}
// The result types of a WhileOp are the same as the "after" bbArg types.
@@ -1186,14 +1180,14 @@ struct YieldOpInterface
yieldOp->getParentOp()->getResult(it.index()), options, state);
if (failed(resultType))
return failure();
- buffer = castBuffer(rewriter, buffer, *resultType);
+ buffer = castBuffer(rewriter, buffer, *resultType, options);
} else if (auto whileOp =
dyn_cast<scf::WhileOp>(yieldOp->getParentOp())) {
FailureOr<BufferLikeType> resultType = bufferization::getBufferType(
whileOp.getBeforeArguments()[it.index()], options, state);
if (failed(resultType))
return failure();
- buffer = castBuffer(rewriter, buffer, *resultType);
+ buffer = castBuffer(rewriter, buffer, *resultType, options);
}
newResults.push_back(buffer);
} else {
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 b2674e9497927..ccc60d42ca23a 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
@@ -312,7 +312,7 @@ module @BufferizeLayoutForFunction {
%memref = "test.tensor_with_future_layout"() {layout = #layout1}
: () -> tensor<10xf32>
- // CHECK: %[[out:.*]] = memref.cast %[[memref]]
+ // CHECK: %[[out:.*]] = test.dummy_cast_op %[[memref]]
// CHECK-SAME: to memref<10xf32>
// CHECK: return %[[out]]
return %memref : tensor<10xf32>
@@ -341,7 +341,7 @@ module @BufferizeLayoutMismatchInsideScfIf {
} else {
// CHECK: %[[another:.*]] = "test.create_memref_op"
// CHECK-SAME: -> memref<10xf32, #test.memref_layout<"layout_b">>
- // CHECK: %[[cast:.*]] = memref.cast %[[another]]
+ // CHECK: %[[cast:.*]] = test.dummy_cast_op %[[another]]
// CHECK-SAME: to memref<10xf32, #test.memref_layout<"layout_a">>
%another = "test.tensor_with_future_layout"() {layout = #layout2}
: () -> tensor<10xf32>
@@ -349,7 +349,7 @@ module @BufferizeLayoutMismatchInsideScfIf {
scf.yield %another : tensor<10xf32>
}
- // CHECK: %[[out:.*]] = memref.cast %[[if]]
+ // CHECK: %[[out:.*]] = test.dummy_cast_op %[[if]]
// CHECK-SAME: to memref<10xf32>
// CHECK: return %[[out]]
return %ret : tensor<10xf32>
@@ -381,7 +381,7 @@ module @BufferizeLayoutMismatchInsideScfSwitch {
case 1 {
// CHECK: %[[another:.*]] = "test.create_memref_op"
// CHECK-SAME: -> memref<10xf32, #test.memref_layout<"layout_b">>
- // CHECK: %[[cast:.*]] = memref.cast %[[another]]
+ // CHECK: %[[cast:.*]] = test.dummy_cast_op %[[another]]
// CHECK-SAME: to memref<10xf32, #test.memref_layout<"layout_a">>
%another = "test.tensor_with_future_layout"() {layout = #layout2}
: () -> tensor<10xf32>
@@ -391,7 +391,7 @@ module @BufferizeLayoutMismatchInsideScfSwitch {
default {
// CHECK: %[[yet_another:.*]] = "test.create_memref_op"
// CHECK-SAME: -> memref<10xf32, #test.memref_layout<"layout_c">>
- // CHECK: %[[cast:.*]] = memref.cast %[[yet_another]]
+ // CHECK: %[[cast:.*]] = test.dummy_cast_op %[[yet_another]]
// CHECK-SAME: to memref<10xf32, #test.memref_layout<"layout_a">>
%yet_another = "test.tensor_with_future_layout"() {layout = #layout3}
: () -> tensor<10xf32>
@@ -399,7 +399,7 @@ module @BufferizeLayoutMismatchInsideScfSwitch {
scf.yield %yet_another : tensor<10xf32>
}
- // CHECK: %[[out:.*]] = memref.cast %[[switch]]
+ // CHECK: %[[out:.*]] = test.dummy_cast_op %[[switch]]
// CHECK-SAME: to memref<10xf32>
// CHECK: return %[[out]]
return %ret : tensor<10xf32>
@@ -430,7 +430,7 @@ module @BufferizeLayoutMismatchInsideScfFor {
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: %[[cast:.*]] = test.dummy_cast_op %[[conflict]]
// CHECK-SAME: to memref<10xf32, #test.memref_layout<"layout_a">>
%conflict = "test.force_new_layout"(%iter) {layout = #layout2}
: (tensor<10xf32>) -> tensor<10xf32>
@@ -438,7 +438,7 @@ module @BufferizeLayoutMismatchInsideScfFor {
scf.yield %conflict : tensor<10xf32>
}
- // CHECK: %[[out:.*]] = memref.cast %[[loop]]
+ // CHECK: %[[out:.*]] = test.dummy_cast_op %[[loop]]
// CHECK-SAME: to memref<10xf32>
// CHECK: return %[[out]]
return %loop : tensor<10xf32>
@@ -450,6 +450,45 @@ module @BufferizeLayoutMismatchInsideScfFor {
#layout1 = #test.memref_layout<"layout_a">
#layout2 = #test.memref_layout<"layout_b">
+// CHECK-LABEL: @BufferizeLayoutMismatchInsideScfForWithUserType
+module @BufferizeLayoutMismatchInsideScfForWithUserType {
+ // CHECK: func.func @mismatch_in_for
+ // CHECK-SAME: -> !test.test_memref<[10], f32>
+ func.func @mismatch_in_for(
+ %lb: index, %ub: index, %step: index)
+ -> !test.test_tensor<[10], f32> {
+ // CHECK: %[[init:.*]] = "test.create_memref_op"
+ // CHECK-SAME: -> !test.test_memref<[10], f32, #test.memref_layout<"layout_a">>
+ %init = "test.tensor_with_future_layout"() {layout = #layout1}
+ : () -> !test.test_tensor<[10], f32>
+
+ // CHECK: %[[loop:.+]] = scf.for
+ // CHECK-SAME: iter_args(%[[iter:.*]] = %[[init]])
+ // CHECK-SAME: -> (!test.test_memref<[10], f32, #test.memref_layout<"layout_a">>)
+ %loop = scf.for %i = %lb to %ub step %step
+ iter_args(%iter = %init) -> (!test.test_tensor<[10], f32>) {
+ // CHECK: %[[conflict:.*]] = "test.dummy_memref_op"(%[[iter]])
+ // CHECK-SAME: -> !test.test_memref<[10], f32, #test.memref_layout<"layout_b">>
+ // CHECK: %[[cast:.*]] = test.dummy_cast_op %[[conflict]]
+ // CHECK-SAME: to !test.test_memref<[10], f32, #test.memref_layout<"layout_a">>
+ %conflict = "test.force_new_layout"(%iter) {layout = #layout2}
+ : (!test.test_tensor<[10], f32>) -> !test.test_tensor<[10], f32>
+ // CHECK: scf.yield %[[cast]]
+ scf.yield %conflict : !test.test_tensor<[10], f32 >
+ }
+
+ // CHECK: %[[out:.*]] = test.dummy_cast_op %[[loop]]
+ // CHECK-SAME: to !test.test_memref<[10], f32>
+ // CHECK: return %[[out]]
+ return %loop : !test.test_tensor<[10], f32>
+ }
+}
+
+// -----
+
+#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
@@ -495,7 +534,7 @@ module @BufferizeLayoutMismatchInsideScfForWithSubviews {
scf.yield %out : tensor<10xf32>
}
- // CHECK: %[[out:.*]] = memref.cast %[[loop]]
+ // CHECK: %[[out:.*]] = test.dummy_cast_op %[[loop]]
// CHECK-SAME: to memref<10xf32>
// CHECK: return %[[out]]
return %loop : tensor<10xf32>
@@ -528,7 +567,7 @@ module @BufferizeLayoutMismatchInScfExecuteRegion {
^bb2:
// CHECK: %[[another:.*]] = "test.create_memref_op"
// CHECK-SAME: -> memref<10xf32, #test.memref_layout<"layout_b">>
- // CHECK: %[[cast:.*]] = memref.cast %[[another]]
+ // CHECK: %[[cast:.*]] = test.dummy_cast_op %[[another]]
// CHECK-SAME: to memref<10xf32, #test.memref_layout<"layout_a">>
%another = "test.tensor_with_future_layout"() {layout = #layout2}
: () -> tensor<10xf32>
@@ -540,13 +579,59 @@ module @BufferizeLayoutMismatchInScfExecuteRegion {
scf.yield %res : tensor<10xf32>
}
- // CHECK: %[[out:.*]] = memref.cast %[[region]]
+ // CHECK: %[[out:.*]] = test.dummy_cast_op %[[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: @BufferizeLayoutMismatchInScfExecuteRegionWithUserType
+module @BufferizeLayoutMismatchInScfExecuteRegionWithUserType {
+ // CHECK: func.func @mismatch_in_scf_execute_region
+ // CHECK-SAME: -> !test.test_memref<[10], f32>
+ func.func @mismatch_in_scf_execute_region(%cond: i1) -> !test.test_tensor<[10], f32> {
+ // CHECK: %[[region:.*]] = scf.execute_region
+ // CHECK-SAME: -> !test.test_memref<[10], f32, #test.memref_layout<"layout_a">>
+ %out = scf.execute_region -> !test.test_tensor<[10], f32> {
+ cf.cond_br %cond, ^bb1, ^bb2
+
+ ^bb1:
+ // CHECK: %[[one:.*]] = "test.create_memref_op"
+ // CHECK-SAME: -> !test.test_memref<[10], f32, #test.memref_layout<"layout_a">>
+ %one = "test.tensor_with_future_layout"() {layout = #layout1}
+ : () -> !test.test_tensor<[10], f32>
+ // CHECK: cf.br ^bb3(%[[one]]
+ cf.br ^bb3(%one : !test.test_tensor<[10], f32>)
+
+ ^bb2:
+ // CHECK: %[[another:.*]] = "test.create_memref_op"
+ // CHECK-SAME: -> !test.test_memref<[10], f32, #test.memref_layout<"layout_b">>
+ // CHECK: %[[cast:.*]] = test.dummy_cast_op %[[another]]
+ // CHECK-SAME: to !test.test_memref<[10], f32, #test.memref_layout<"layout_a">>
+ %another = "test.tensor_with_future_layout"() {layout = #layout2}
+ : () -> !test.test_tensor<[10], f32>
+ // CHECK: cf.br ^bb3(%[[cast]]
+ cf.br ^bb3(%another : !test.test_tensor<[10], f32>)
+
+ ^bb3(%res: !test.test_tensor<[10], f32>):
+ // CHECK: scf.yield {{%.*}} : !test.test_memref<[10], f32, #test.memref_layout<"layout_a">>
+ scf.yield %res : !test.test_tensor<[10], f32>
+ }
+
+ // CHECK: %[[out:.*]] = test.dummy_cast_op %[[region]]
+ // CHECK-SAME: to !test.test_memref<[10], f32>
+ // CHECK: return %[[out]]
+ return %out : !test.test_tensor<[10], f32>
+ }
+}
+
+
// -----
#layout1 = #test.memref_layout<"layout_a">
@@ -566,13 +651,13 @@ module @BufferizeLayoutMismatchInArithSelect {
%tensor2 = "test.tensor_with_future_layout"() {layout = #layout2}
: () -> tensor<10xf32>
- // CHECK: %[[cast:.*]] = memref.cast %[[memref2]]
+ // CHECK: %[[cast:.*]] = test.dummy_cast_op %[[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: %[[out:.*]] = test.dummy_cast_op %[[select]]
// CHECK-SAME: to memref<10xf32>
// CHECK: return %[[out]]
return %select : tensor<10xf32>
diff --git a/mlir/test/lib/Dialect/Bufferization/TestOneShotModuleBufferize.cpp b/mlir/test/lib/Dialect/Bufferization/TestOneShotModuleBufferize.cpp
index 00eb2dc92e14f..036e05bfc2e02 100644
--- a/mlir/test/lib/Dialect/Bufferization/TestOneShotModuleBufferize.cpp
+++ b/mlir/test/lib/Dialect/Bufferization/TestOneShotModuleBufferize.cpp
@@ -16,6 +16,7 @@
#include "TestAttributes.h" // TestTensorEncodingAttr, TestMemRefLayoutAttr
#include "TestDialect.h"
+#include "TestOps.h"
using namespace mlir;
@@ -30,6 +31,21 @@ getMemRefLayoutForTensorEncoding(RankedTensorType tensorType) {
return {};
}
+test::TestMemRefLayoutAttr
+getLayoutFromBuffer(bufferization::BufferLikeType buffer) {
+ auto layout =
+ llvm::TypeSwitch<bufferization::BufferLikeType,
+ MemRefLayoutAttrInterface>(buffer)
+ .Case([&](MemRefType memref) { return memref.getLayout(); })
+ .Case([&](test::TestMemrefType testMemref) {
+ return cast<MemRefLayoutAttrInterface>(testMemref.getLayout());
+ })
+ .Default([](bufferization::BufferLikeType) {
+ return MemRefLayoutAttrInterface();
+ });
+ return dyn_cast_or_null<test::TestMemRefLayoutAttr>(layout);
+}
+
struct TestOneShotModuleBufferizePass
: public PassWrapper<TestOneShotModuleBufferizePass, OperationPass<>> {
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestOneShotModuleBufferizePass)
@@ -82,9 +98,14 @@ struct TestOneShotModuleBufferizePass
})
.Case([&](test::TestTensorType testTensorType)
-> bufferization::BufferLikeType {
+ assert(memSpace == nullptr &&
+ "memory space is not supported for test types");
+ // Note: unknown type conversion cannot create layout because
+ // test tensor does not have an encoding, this is probably
+ // enough for a test.
return test::TestMemrefType::get(
testTensorType.getContext(), testTensorType.getShape(),
- testTensorType.getElementType(), memSpace);
+ testTensorType.getElementType(), /*layout=*/nullptr);
})
.Default([&](bufferization::TensorLikeType tensor) {
llvm_unreachable("unexpected tensor type");
@@ -97,13 +118,8 @@ struct TestOneShotModuleBufferizePass
[](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);
+ auto lhsLayout = getLayoutFromBuffer(x);
+ auto rhsLayout = getLayoutFromBuffer(y);
if (lhsLayout && rhsLayout) {
return lhsLayout.getDummy().getValue() <=
rhsLayout.getDummy().getValue()
@@ -112,6 +128,9 @@ struct TestOneShotModuleBufferizePass
}
return rhsLayout ? y : x;
};
+ opt.castFn = [&](OpBuilder &b, Location loc, Type dest, Value value) {
+ return test::TestDummyCastOp::create(b, loc, dest, value).getResult();
+ };
// Function signature update only works with memref.cast. Disable it to
// align behaviour for upstream and user casts.
opt.inferFunctionResultLayout = false;
diff --git a/mlir/test/lib/Dialect/Test/TestOpDefs.cpp b/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
index 20d72a2fd09ca..75dae36fe9ce8 100644
--- a/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
+++ b/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
@@ -2068,19 +2068,27 @@ 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));
+ return llvm::TypeSwitch<mlir::Type,
+ mlir::FailureOr<mlir::bufferization::BufferLikeType>>(
+ value.getType())
+ .Case([&](RankedTensorType tensorType) {
+ // 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));
+ })
+ .Case([&](TestTensorType tensorType) {
+ auto layout = cast<MemRefLayoutAttrInterface>(getLayout());
+ return cast<bufferization::BufferLikeType>(
+ TestMemrefType::get(tensorType.getContext(), tensorType.getShape(),
+ tensorType.getElementType(), layout));
+ })
+ .Default([&](Type) { return emitError("unknown type"); });
}
LogicalResult test::TestForceNewLayoutOp::bufferize(
@@ -2107,15 +2115,23 @@ 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));
+ return llvm::TypeSwitch<mlir::Type,
+ mlir::FailureOr<mlir::bufferization::BufferLikeType>>(
+ value.getType())
+ .Case([&](RankedTensorType tensorType) {
+ // 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));
+ })
+ .Case([&](TestTensorType tensorType) {
+ auto layout = cast<MemRefLayoutAttrInterface>(getLayout());
+ return cast<bufferization::BufferLikeType>(
+ TestMemrefType::get(tensorType.getContext(), tensorType.getShape(),
+ tensorType.getElementType(), layout));
+ })
+ .Default([&](Type) { return emitError("unknown type"); });
}
// Define a custom builder for ManyRegionsOp declared in TestOps.td.
diff --git a/mlir/test/lib/Dialect/Test/TestOps.td b/mlir/test/lib/Dialect/Test/TestOps.td
index 29f456d1ec7eb..c12e2abb3ab75 100644
--- a/mlir/test/lib/Dialect/Test/TestOps.td
+++ b/mlir/test/lib/Dialect/Test/TestOps.td
@@ -4239,7 +4239,7 @@ def TestTensorWithFutureLayoutOp
}];
let arguments = (ins TestMemRefLayoutAttr:$layout);
- let results = (outs AnyRankedTensor:$output);
+ let results = (outs Bufferization_TensorLikeTypeInterface:$output);
let extraClassDefinition = [{
bool test::TestTensorWithFutureLayoutOp::bufferizesToMemoryRead(
::mlir::OpOperand&,
@@ -4285,10 +4285,10 @@ def TestForceNewLayoutOp
}];
let arguments = (ins
- AnyRankedTensor:$input,
+ Bufferization_TensorLikeTypeInterface:$input,
TestMemRefLayoutAttr:$layout
);
- let results = (outs AnyRankedTensor:$output);
+ let results = (outs Bufferization_TensorLikeTypeInterface:$output);
let extraClassDefinition = [{
bool test::TestForceNewLayoutOp::bufferizesToMemoryRead(::mlir::OpOperand&,
const ::mlir::bufferization::AnalysisState&) {
@@ -4322,6 +4322,16 @@ def TestForceNewLayoutOp
}];
}
+def TestDummyCastOp : TEST_Op<"dummy_cast_op", []> {
+ let arguments = (ins
+ Arg<Bufferization_BufferLikeTypeInterface>:$source
+ );
+ let results = (outs
+ Arg<Bufferization_BufferLikeTypeInterface>:$dest
+ );
+ let assemblyFormat = "$source attr-dict `:` type($source) `to` type($dest)";
+}
+
//===----------------------------------------------------------------------===//
// Test assembly format references
//===----------------------------------------------------------------------===//
diff --git a/mlir/test/lib/Dialect/Test/TestTypeDefs.td b/mlir/test/lib/Dialect/Test/TestTypeDefs.td
index f405940358ca0..58c340a935bdb 100644
--- a/mlir/test/lib/Dialect/Test/TestTypeDefs.td
+++ b/mlir/test/lib/Dialect/Test/TestTypeDefs.td
@@ -477,9 +477,9 @@ def TestMemrefType : Test_Type<"TestMemref",
let parameters = (ins
ArrayRefParameter<"int64_t">:$shape,
"mlir::Type":$elementType,
- DefaultValuedParameter<"mlir::Attribute", "nullptr">:$memSpace
+ DefaultValuedParameter<"mlir::MemRefLayoutAttrInterface", "nullptr">:$layout
);
- let assemblyFormat = "`<` `[` $shape `]` `,` $elementType (`,` $memSpace^)? `>`";
+ let assemblyFormat = "`<` `[` $shape `]` `,` $elementType (`,` $layout^)? `>`";
let extraClassDeclaration = [{
// ShapedTypeInterface:
@@ -489,7 +489,7 @@ def TestMemrefType : Test_Type<"TestMemref",
test::TestMemrefType cloneWith(std::optional<llvm::ArrayRef<int64_t>> shape,
mlir::Type elementType) const {
return test::TestMemrefType::get(
- getContext(), shape.value_or(getShape()), elementType, getMemSpace());
+ getContext(), shape.value_or(getShape()), elementType, getLayout());
}
}];
}
More information about the Mlir-commits
mailing list