[Mlir-commits] [mlir] [mlir][bufferization] Introduce reconcileBufferTypeMismatchFn hook (PR #202667)

[Andrei Golubev]

================
@@ -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),
----------------
andrey-golubev wrote:

good catch, thanks!

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