[Mlir-commits] [mlir] [mlir][bufferization] Introduce castBufferTypeFn hook and API (PR #205080)

llvmlistbot at llvm.org llvmlistbot at llvm.org
Mon Jun 22 04:15:17 PDT 2026


llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-scf

Author: Andrei Golubev (andrey-golubev)

<details>
<summary>Changes</summary>

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.

---

Patch is 31.93 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/205080.diff


11 Files Affected:

- (modified) mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h (+11) 
- (modified) mlir/lib/Dialect/Arith/Transforms/BufferizableOpInterfaceImpl.cpp (+3-6) 
- (modified) mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp (+17) 
- (modified) mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp (+15-1) 
- (modified) mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp (+4-3) 
- (modified) mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp (+12-26) 
- (modified) mlir/test/Dialect/Bufferization/Transforms/test-one-shot-module-bufferize.mlir (+98-13) 
- (modified) mlir/test/lib/Dialect/Bufferization/TestOneShotModuleBufferize.cpp (+27-8) 
- (modified) mlir/test/lib/Dialect/Test/TestOpDefs.cpp (+38-22) 
- (modified) mlir/test/lib/Dialect/Test/TestOps.td (+13-3) 
- (modified) mlir/test/lib/Dialect/Test/TestTypeDefs.td (+3-3) 


``````````diff
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..7ef115b7ef003 100644
--- a/mlir/lib/Dialect/Arith/Transforms/BufferizableOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Arith/Transforms/BufferizableOpInterfaceImpl.cpp
@@ -172,12 +172,9 @@ struct SelectOpInterface
         return selectOp->emitError(
             "incompatible buffer types on true/false operands");
       }
-      if (trueBuffer.getType() != *targetType)
-        trueBuffer =
-            memref::CastOp::create(rewriter, loc, *targetType, trueBuffer);
-      if (falseBuffer.getType() != *targetType)
-        falseBuffer =
-            memref::CastOp::create(rewriter, loc, *targetType, falseBuffer);
+      trueBuffer = *options.createCast(rewriter, loc, *targetType, trueBuffer);
+      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..79ab923864871 100644
--- a/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
+++ b/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
@@ -840,6 +840,23 @@ LogicalResult BufferizationOptions::createMemCpy(OpBuilder &b, Location loc,
   return success();
 }
 
+FailureOr<Value> BufferizationOptions::createCast(OpBuilder &b, Location loc,
+                                                  Type dest,
+                                                  Value value) const {
+  // If the value already has the correct type, no cast is needed.
+  if (value.getType() == dest)
+    return value;
+
+  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 different
       // 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..db26bf2f82b6a 100644
--- a/mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp
@@ -29,25 +29,6 @@ namespace mlir {
 namespace scf {
 namespace {
 
-/// Helper function for loop bufferization. Cast the given buffer to the given
-/// memref type.
-static Value castBuffer(OpBuilder &b, Value buffer, Type type) {
-  // 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();
-}
-
 /// Helper function for loop bufferization. Return "true" if the given value
 /// is guaranteed to not alias with an external tensor apart from values in
 /// `exceptions`. A value is external if it is defined outside of the given
@@ -114,7 +95,8 @@ struct ConditionOpInterface
             whileOp.getAfterArguments()[it.index()], options, state);
         if (failed(resultType))
           return failure();
-        Value buffer = castBuffer(rewriter, *maybeBuffer, *resultType);
+        Value buffer = *options.createCast(rewriter, maybeBuffer->getLoc(),
+                                           *resultType, *maybeBuffer);
         newArgs.push_back(buffer);
       } else {
         newArgs.push_back(value);
@@ -206,7 +188,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 +343,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 +744,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(*options.createCast(rewriter, initArg.getLoc(),
+                                                   *targetType, initArg));
     }
 
     // Construct a new scf.for op with memref instead of tensor values.
@@ -986,7 +969,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(*options.createCast(rewriter, initArg.getLoc(),
+                                                   *targetType, initArg));
     }
 
     // The result types of a WhileOp are the same as the "after" bbArg types.
@@ -1186,14 +1170,16 @@ struct YieldOpInterface
               yieldOp->getParentOp()->getResult(it.index()), options, state);
           if (failed(resultType))
             return failure();
-          buffer = castBuffer(rewriter, buffer, *resultType);
+          buffer = *options.createCast(rewriter, buffer.getLoc(), *resultType,
+                                       buffer);
         } 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 = *options.createCast(rewriter, buffer.getLoc(), *resultType,
+                                       buffer);
         }
         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">...
[truncated]

``````````

</details>


https://github.com/llvm/llvm-project/pull/205080


More information about the Mlir-commits mailing list