Address code review feedback

* Remove BufferizationState from TensorLikeType::getBufferType()
* Rename castToMemRef to asMemRefType (+ add extra docs)
* Improve ToTensorOp's docs
* Apply minor suggestions

Author: andrey-golubev

mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h

@@ -742,7 +742,12 @@ AliasingValueList unknownGetAliasingValues(OpOperand &opOperand);
 bool defaultHasTensorSemantics(Operation *op);
 
 /// This is a helper function used when buffer type is guaranteed to be memref.
-FailureOr<BaseMemRefType> castToMemRef(FailureOr<BufferLikeType> bufferType);
+/// It performs two actions: failure state checking and an explicit llvm::cast<>
+/// from the buffer-like type interface to a BaseMemRefType. This allows easier
+/// management of differences in C++ types at the API boundaries. Valid buffer
+/// type is casted to the memref type. Otherwise, the failure state is
+/// propagated i.e. asMemRefType(mlir::failure()) returns mlir::failure().
+FailureOr<BaseMemRefType> asMemRefType(FailureOr<BufferLikeType> bufferType);
 
 /// This function is a free-standing helper that relies on
 /// bufferization::TensorLikeTypeInterface to verify the types in tensor and

mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td

@@ -404,7 +404,10 @@ def Bufferization_ToTensorOp : Bufferization_Op<"to_tensor", [
   let summary = "create a buffer-like type from a tensor-like type";
   let description = [{
     An operation that creates a tensor from a buffer. The result value is a
-    tensor-like type whose shape and element type match the buffer-like operand.
+    tensor-like type that must match the corresponding buffer-like operand as
+    per TensorLikeType::verifyCompatibleBufferType(). For builtins (TensorType
+    and BaseMemRefType), this means that shapes and element types match between
+    the tensor and the buffer.
 
     The opposite of this op is `to_buffer`. Together, these two ops are
     useful for source/target materializations when doing type conversions

mlir/include/mlir/Dialect/Bufferization/IR/BufferizationTypeInterfaces.td

@@ -31,7 +31,6 @@ def Bufferization_TensorLikeTypeInterface
       /*methodName=*/"getBufferType",
       /*args=*/(ins
         "const ::mlir::bufferization::BufferizationOptions &":$options,
-        "const ::mlir::bufferization::BufferizationState &":$state,
         "::llvm::function_ref<::mlir::InFlightDiagnostic()>":$emitError
       )
     >,

mlir/lib/Dialect/Arith/Transforms/BufferizableOpInterfaceImpl.cpp

@@ -164,7 +164,7 @@ struct SelectOpInterface
     // buffers have different types, they differ only in their layout map. Cast
     // both of them to the most dynamic MemRef type.
     if (trueBuffer.getType() != falseBuffer.getType()) {
-      auto targetType = bufferization::detail::castToMemRef(
+      auto targetType = bufferization::detail::asMemRefType(
           bufferization::getBufferType(selectOp.getResult(), options, state));
       if (failed(targetType))
         return failure();
@@ -188,10 +188,10 @@ struct SelectOpInterface
     auto selectOp = cast<arith::SelectOp>(op);
     assert(value == selectOp.getResult() && "invalid value");
     auto trueType =
-        bufferization::detail::castToMemRef(bufferization::getBufferType(
+        bufferization::detail::asMemRefType(bufferization::getBufferType(
             selectOp.getTrueValue(), options, state, invocationStack));
     auto falseType =
-        bufferization::detail::castToMemRef(bufferization::getBufferType(
+        bufferization::detail::asMemRefType(bufferization::getBufferType(
             selectOp.getFalseValue(), options, state, invocationStack));
     if (failed(trueType) || failed(falseType))
       return failure();

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

@@ -214,7 +214,7 @@ FailureOr<Value> bufferization::allocateTensorForShapedValue(
   if (copy)
     return allocTensorOp.getResult();
   auto copyBufferType =
-      detail::castToMemRef(getBufferType(tensor, options, state));
+      detail::asMemRefType(getBufferType(tensor, options, state));
   if (failed(copyBufferType))
     return failure();
   std::optional<Attribute> memorySpace = copyBufferType->getMemorySpace();
@@ -720,8 +720,9 @@ bufferization::getBufferType(Value value, const BufferizationOptions &options,
     return bufferizableOp.getBufferType(value, options, state, invocationStack);
 
   // Op is not bufferizable.
-  return cast<TensorLikeType>(value.getType())
-      .getBufferType(options, state, [&]() { return op->emitError(); });
+  return cast<TensorLikeType>(value.getType()).getBufferType(options, [&]() {
+    return op->emitError();
+  });
 }
 
 bool bufferization::hasTensorSemantics(Operation *op) {
@@ -965,7 +966,7 @@ FailureOr<BaseMemRefType> bufferization::detail::defaultGetBufferType(
     // If the OpResult has an equivalent OpOperand, both OpResult and
     // OpOperand bufferize to the exact same buffer type.
     Value equivalentOperand = aliases.getAliases().front().opOperand->get();
-    return castToMemRef(getBufferType(equivalentOperand, options,
+    return asMemRefType(getBufferType(equivalentOperand, options,
                                       bufferizationState, invocationStack));
   }
 
@@ -1043,19 +1044,15 @@ bool bufferization::detail::defaultHasTensorSemantics(Operation *op) {
 }
 
 FailureOr<BaseMemRefType>
-bufferization::detail::castToMemRef(FailureOr<BufferLikeType> bufferType) {
+bufferization::detail::asMemRefType(FailureOr<BufferLikeType> bufferType) {
   if (failed(bufferType))
     return failure();
-  assert(isa<BaseMemRefType>(*bufferType) && "expected memref type");
   return cast<BaseMemRefType>(*bufferType);
 }
 
 bool bufferization::detail::typesMatchAfterBufferization(Operation &op,
                                                          Value tensor,
                                                          Value buffer) {
-  assert(isa<TensorLikeType>(tensor.getType()) && "expected TensorLikeType");
-  assert(isa<BufferLikeType>(buffer.getType()) && "expected BufferLikeType");
-
   return mlir::succeeded(
       cast<TensorLikeType>(tensor.getType())
           .verifyCompatibleBufferType(cast<BufferLikeType>(buffer.getType()),

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

@@ -60,10 +60,8 @@ struct BuiltinTensorExternalModel
                                     Tensor> {
   llvm::FailureOr<BufferLikeType> getBufferType(
       mlir::Type tensor, const BufferizationOptions &options,
-      const BufferizationState &state,
       llvm::function_ref<mlir::InFlightDiagnostic()> emitError) const {
     auto tensorType = cast<TensorType>(tensor);
-    // Fall back to tensor -> memref conversion.
     auto memSpace = options.defaultMemorySpaceFn(tensorType);
     if (!memSpace.has_value())
       return emitError() << "could not infer memory space";
@@ -75,7 +73,6 @@ struct BuiltinTensorExternalModel
   mlir::LogicalResult verifyCompatibleBufferType(
       mlir::Type tensor, BufferLikeType bufferType,
       llvm::function_ref<mlir::InFlightDiagnostic()> emitError) const {
-    // Fall back to tensor, memref checking.
     assert(isa<TensorType>(tensor) && "expected tensor type");
     assert(isa<BaseMemRefType>(bufferType) && "expected memref type");
 

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

@@ -234,7 +234,7 @@ AllocTensorOp::getBufferType(Value value, const BufferizationOptions &options,
     memorySpace = *getMemorySpace();
   } else if (getCopy()) {
     auto copyBufferType =
-        bufferization::detail::castToMemRef(bufferization::getBufferType(
+        bufferization::detail::asMemRefType(bufferization::getBufferType(
             getCopy(), options, state, invocationStack));
     if (failed(copyBufferType))
       return failure();

mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp

@@ -293,7 +293,7 @@ struct IfOpInterface
       thenBufferType = cast<BaseMemRefType>(thenValue.getType());
     } else {
       auto maybeBufferType =
-          bufferization::detail::castToMemRef(bufferization::getBufferType(
+          bufferization::detail::asMemRefType(bufferization::getBufferType(
               thenValue, options, state, invocationStack));
       if (failed(maybeBufferType))
         return failure();
@@ -304,7 +304,7 @@ struct IfOpInterface
       elseBufferType = cast<BaseMemRefType>(elseValue.getType());
     } else {
       auto maybeBufferType =
-          bufferization::detail::castToMemRef(bufferization::getBufferType(
+          bufferization::detail::asMemRefType(bufferization::getBufferType(
               elseValue, options, state, invocationStack));
       if (failed(maybeBufferType))
         return failure();
@@ -408,7 +408,7 @@ struct IndexSwitchOpInterface
         return bufferType;
       auto maybeBufferType = bufferization::getBufferType(
           yieldedValue, options, state, invocationStack);
-      return bufferization::detail::castToMemRef(maybeBufferType);
+      return bufferization::detail::asMemRefType(maybeBufferType);
     };
 
     // Compute buffer type of the default case.
@@ -527,7 +527,7 @@ static FailureOr<BaseMemRefType> computeLoopRegionIterArgBufferType(
     const BufferizationOptions &options, const BufferizationState &state,
     SmallVector<Value> &invocationStack) {
   // Determine the buffer type of the init_arg.
-  auto initArgBufferType = bufferization::detail::castToMemRef(
+  auto initArgBufferType = bufferization::detail::asMemRefType(
       bufferization::getBufferType(initArg, options, state, invocationStack));
   if (failed(initArgBufferType))
     return failure();
@@ -555,7 +555,7 @@ static FailureOr<BaseMemRefType> computeLoopRegionIterArgBufferType(
     // Note: This typically triggers a recursive call for the buffer type of
     // the iter_arg.
     auto maybeBufferType =
-        bufferization::detail::castToMemRef(bufferization::getBufferType(
+        bufferization::detail::asMemRefType(bufferization::getBufferType(
             yieldedValue, options, state, invocationStack));
     if (failed(maybeBufferType))
       return failure();
@@ -1083,7 +1083,7 @@ struct WhileOpInterface
       // scf.condition was already bufferized.
       return cast<BaseMemRefType>(conditionYieldedVal.getType());
     }
-    return bufferization::detail::castToMemRef(bufferization::getBufferType(
+    return bufferization::detail::asMemRefType(bufferization::getBufferType(
         conditionYieldedVal, options, state, invocationStack));
   }
 
@@ -1312,13 +1312,13 @@ struct ForallOpInterface
     if (auto bbArg = dyn_cast<BlockArgument>(value))
       // A tensor block argument has the same bufferized type as the
       // corresponding output operand.
-      return bufferization::detail::castToMemRef(
+      return bufferization::detail::asMemRefType(
           bufferization::getBufferType(forallOp.getTiedOpOperand(bbArg)->get(),
                                        options, state, invocationStack));
 
     // The bufferized result type is the same as the bufferized type of the
     // corresponding output operand.
-    return bufferization::detail::castToMemRef(bufferization::getBufferType(
+    return bufferization::detail::asMemRefType(bufferization::getBufferType(
         forallOp.getOutputs()[cast<OpResult>(value).getResultNumber()], options,
         state, invocationStack));
   }

mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp

@@ -55,7 +55,7 @@ struct CastOpInterface
                 SmallVector<Value> &invocationStack) const {
     auto castOp = cast<tensor::CastOp>(op);
     auto maybeSrcBufferType =
-        bufferization::detail::castToMemRef(bufferization::getBufferType(
+        bufferization::detail::asMemRefType(bufferization::getBufferType(
             castOp.getSource(), options, state, invocationStack));
     if (failed(maybeSrcBufferType))
       return failure();
@@ -501,7 +501,7 @@ struct FromElementsOpInterface
         /*copy=*/false);
     if (failed(tensorAlloc))
       return failure();
-    FailureOr<BaseMemRefType> memrefType = bufferization::detail::castToMemRef(
+    FailureOr<BaseMemRefType> memrefType = bufferization::detail::asMemRefType(
         bufferization::getBufferType(*tensorAlloc, options, state));
     if (failed(memrefType))
       return failure();
@@ -760,7 +760,7 @@ struct PadOpInterface
     // Infer memory space from the source tensor.
     auto padOp = cast<tensor::PadOp>(op);
     auto maybeSrcBufferType =
-        bufferization::detail::castToMemRef(bufferization::getBufferType(
+        bufferization::detail::asMemRefType(bufferization::getBufferType(
             padOp.getSource(), options, state, invocationStack));
     if (failed(maybeSrcBufferType))
       return failure();

mlir/test/lib/Dialect/Test/TestTypeDefs.td

@@ -432,7 +432,6 @@ def TestTensorType : Test_Type<"TestTensor",
     // TensorLikeTypeInterface:
     ::mlir::FailureOr<::mlir::bufferization::BufferLikeType>
     getBufferType(const ::mlir::bufferization::BufferizationOptions& options,
-                  const ::mlir::bufferization::BufferizationState& state,
                   ::llvm::function_ref<::mlir::InFlightDiagnostic()> emitError);
 
     ::mlir::LogicalResult verifyCompatibleBufferType(