[mlir][bufferization] Refine tensor-buffer compatibility checks

Generally, to_tensor and to_buffer already perform sufficient
verification. However, there are some unnecessary strict constraints:
* builtin tensor requires its buffer counterpart to always be memref
* to_buffer on ranked tensor requires to always return memref

These checks are assertions (i.e. preconditions), however, they actually
prevent an apparently useful bufferization where builtin tensors could
become custom buffers. Lift these assertions, maintaining the verification
procedure unchanged, to allow builtin -> custom bufferizations at
operation boundary level.

Author: andrey-golubev

mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp

@@ -680,16 +680,6 @@ bool AnalysisState::hasUndefinedContents(OpOperand *opOperand) const {
   return false;
 }
 
-// bufferization.to_buffer is not allowed to change the rank.
-static void ensureToBufferOpIsValid(Value tensor, Type memrefType) {
-#ifndef NDEBUG
-  auto rankedTensorType = llvm::dyn_cast<RankedTensorType>(tensor.getType());
-  assert((!rankedTensorType || llvm::cast<MemRefType>(memrefType).getRank() ==
-                                   rankedTensorType.getRank()) &&
-         "to_buffer would be invalid: mismatching ranks");
-#endif
-}
-
 FailureOr<Value> bufferization::getBuffer(RewriterBase &rewriter, Value value,
                                           const BufferizationOptions &options,
                                           const BufferizationState &state) {
@@ -708,7 +698,7 @@ FailureOr<Value> bufferization::getBuffer(RewriterBase &rewriter, Value value,
   FailureOr<BufferLikeType> bufferType = getBufferType(value, options, state);
   if (failed(bufferType))
     return failure();
-  ensureToBufferOpIsValid(value, *bufferType);
+
   return bufferization::ToBufferOp::create(rewriter, value.getLoc(),
                                            *bufferType, value)
       .getResult();

mlir/lib/Dialect/Bufferization/IR/BufferizationDialect.cpp

@@ -54,9 +54,6 @@ struct BuiltinTensorExternalModel
   mlir::LogicalResult verifyCompatibleBufferType(
       mlir::Type tensor, BufferLikeType bufferType,
       llvm::function_ref<mlir::InFlightDiagnostic()> emitError) const {
-    assert(isa<TensorType>(tensor) && "expected tensor type");
-    assert(isa<BaseMemRefType>(bufferType) && "expected memref type");
-
     auto tensorType = cast<ShapedType>(tensor);
     auto memrefType = cast<ShapedType>(bufferType);
 

mlir/test/Dialect/Bufferization/Transforms/tensorlike-bufferlike.mlir

@@ -1,4 +1,4 @@
-// RUN: mlir-opt %s -test-tensorlike-bufferlike -split-input-file | FileCheck %s
+// RUN: mlir-opt %s -test-tensorlike-bufferlike -verify-diagnostics -split-input-file | FileCheck %s
 
 // CHECK: func.func @builtin_unranked
 // CHECK-SAME: {found = {operand_0 = "is_tensor_like", result_0 = "is_buffer_like"}}
@@ -35,3 +35,29 @@ func.func @custom_memref(%t: !test.test_memref<[42], f32>) -> ()
 {
   return
 }
+
+// -----
+
+// CHECK: func.func @builtin_custom_builtin_roundtrip
+// CHECK-SAME: {found = {operand_0 = "is_tensor_like", result_0 = "is_tensor_like"}}
+func.func @builtin_custom_builtin_roundtrip(%t: tensor<42xf32>)
+    -> tensor<42xf32> {
+  %buffer = bufferization.to_buffer %t
+    : tensor<42xf32> to !test.test_memref<[42], f32>
+  %tensor = bufferization.to_tensor %buffer
+    : !test.test_memref<[42], f32> to tensor<42xf32>
+  return %tensor : tensor<42xf32>
+}
+
+// -----
+
+// CHECK: func.func @custom_builtin_custom_roundtrip
+// CHECK-SAME: {found = {operand_0 = "is_tensor_like", result_0 = "is_tensor_like"}}
+func.func @custom_builtin_custom_roundtrip(%t: !test.test_tensor<[42], f32>)
+    -> !test.test_tensor<[42], f32> {
+  %buffer = bufferization.to_buffer %t
+    : !test.test_tensor<[42], f32> to memref<42xf32>
+  %tensor = bufferization.to_tensor %buffer
+    : memref<42xf32> to !test.test_tensor<[42], f32>
+  return %tensor : !test.test_tensor<[42], f32>
+}

mlir/test/lib/Dialect/Test/TestTypes.cpp

@@ -569,11 +569,17 @@ TestTensorType::getBufferType(
 ::mlir::LogicalResult TestTensorType::verifyCompatibleBufferType(
     ::mlir::bufferization::BufferLikeType bufferType,
     ::llvm::function_ref<::mlir::InFlightDiagnostic()> emitError) {
-  auto testMemref = dyn_cast<TestMemrefType>(bufferType);
-  if (!testMemref)
-    return emitError() << "expected TestMemrefType";
+  if (auto testMemref = dyn_cast<TestMemrefType>(bufferType)) {
+    const bool valid = getShape() == testMemref.getShape() &&
+                       getElementType() == testMemref.getElementType();
+    return mlir::success(valid);
+  }
+
+  if (auto builtinMemref = dyn_cast<MemRefType>(bufferType)) {
+    const bool valid = getShape() == builtinMemref.getShape() &&
+                       getElementType() == builtinMemref.getElementType();
+    return mlir::success(valid);
+  }
 
-  const bool valid = getShape() == testMemref.getShape() &&
-                     getElementType() == testMemref.getElementType();
-  return mlir::success(valid);
+  return emitError() << "expected TestMemrefType";
 }