Switch from ConversionDialectInterface to TensorLike API

Noteworthy changes:
* bufferization::getMemRefType() accepts a TensorType instead of Value
  to achieve broader applicability
* BufferizationOptions::UnknownTypeConverterFn accepts a TensorType
  instead of Value to allow it being used in the updated getMemRefType()

Author: andrey-golubev

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

@@ -745,8 +745,8 @@ bool defaultHasTensorSemantics(Operation *op);
 FailureOr<BaseMemRefType> castToMemRef(FailureOr<BufferLikeType> bufferType);
 
 /// This function is a free-standing helper that relies on
-/// bufferization::ConversionInterface to verify the types in tensor and buffer
-/// worlds match.
+/// bufferization::TensorLikeTypeInterface to verify the types in tensor and
+/// buffer worlds match.
 bool typesMatchAfterBufferization(Operation &op, Value tensor, Value buffer);
 } // namespace detail
 

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

@@ -13,8 +13,15 @@
 // Bufferization Type Interfaces
 //===----------------------------------------------------------------------===//
 
+#include "mlir/IR/Diagnostics.h"
 #include "mlir/IR/Types.h"
 
+namespace mlir::bufferization {
+struct BufferizationOptions;
+class BufferizationState;
+class BufferLikeType;
+} // namespace mlir::bufferization
+
 #include "mlir/Dialect/Bufferization/IR/BufferizationTypeInterfaces.h.inc"
 
 #endif // MLIR_DIALECT_BUFFERIZATION_IR_BUFFERIZATIONTYPEINTERFACES_H_

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

@@ -21,10 +21,31 @@ def Bufferization_TensorLikeTypeInterface
   let description = [{
     Indicates that this type is a tensor type (similarly to a MLIR builtin
     tensor) for bufferization purposes.
-
-    The interface currently has no methods as it is used by types to opt into
-    being supported by the bufferization procedures.
   }];
+
+  let methods = [
+    InterfaceMethod<[{
+        Returns a BufferLike type for this TensorLike type.
+      }],
+      /*retTy=*/"::mlir::FailureOr<::mlir::bufferization::BufferLikeType>",
+      /*methodName=*/"getBufferType",
+      /*args=*/(ins
+        "const ::mlir::bufferization::BufferizationOptions &":$options,
+        "const ::mlir::bufferization::BufferizationState &":$state,
+        "::llvm::function_ref<::mlir::InFlightDiagnostic()>":$emitError
+      )
+    >,
+    InterfaceMethod<[{
+        Returns whether a BufferLike type is compatible to this TensorLike type.
+        The BufferLike type is assumed to be created by getBufferType().
+      }],
+      /*retTy=*/"::mlir::LogicalResult",
+      /*methodName=*/"verifyCompatibleBufferType",
+      /*args=*/(ins
+        "::mlir::bufferization::BufferLikeType":$bufferType,
+        "::llvm::function_ref<::mlir::InFlightDiagnostic()>":$emitError)
+    >
+  ];
 }
 
 def Bufferization_BufferLikeTypeInterface

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

@@ -8,7 +8,6 @@
 
 #include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
-#include "mlir/Dialect/Bufferization/IR/BufferizationConversionInterface.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
@@ -720,11 +719,9 @@ bufferization::getBufferType(Value value, const BufferizationOptions &options,
   if (bufferizableOp)
     return bufferizableOp.getBufferType(value, options, state, invocationStack);
 
-  // Op is not bufferizable, use conversion interface.
-  bufferization::ConversionInterface iface(value.getContext());
-  return iface.getBufferType(value, options, state, [&](const Twine &message) {
-    return op->emitError(message);
-  });
+  // Op is not bufferizable.
+  return cast<TensorLikeType>(value.getType())
+      .getBufferType(options, state, [&]() { return op->emitError(); });
 }
 
 bool bufferization::hasTensorSemantics(Operation *op) {
@@ -1059,10 +1056,8 @@ bool bufferization::detail::typesMatchAfterBufferization(Operation &op,
   assert(isa<TensorLikeType>(tensor.getType()) && "expected TensorLikeType");
   assert(isa<BufferLikeType>(buffer.getType()) && "expected BufferLikeType");
 
-  // Op is not bufferizable, use conversion interface.
-  bufferization::ConversionInterface iface(op.getContext());
-  return succeeded(iface.typesMatch(
-      cast<TensorLikeType>(tensor.getType()),
-      cast<BufferLikeType>(buffer.getType()),
-      [&](const Twine &message) { return op.emitError(message); }));
+  return mlir::succeeded(
+      cast<TensorLikeType>(tensor.getType())
+          .verifyCompatibleBufferType(cast<BufferLikeType>(buffer.getType()),
+                                      [&]() { return op.emitError(); }));
 }

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

@@ -57,7 +57,40 @@ struct BufferizationInlinerInterface : public DialectInlinerInterface {
 template <typename Tensor>
 struct BuiltinTensorExternalModel
     : TensorLikeType::ExternalModel<BuiltinTensorExternalModel<Tensor>,
-                                    Tensor> {};
+                                    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";
+
+    return cast<BufferLikeType>(
+        getMemRefType(tensorType, options, /*layout=*/{}, *memSpace));
+  }
+
+  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");
+
+    auto tensorType = cast<ShapedType>(tensor);
+    auto memrefType = cast<ShapedType>(bufferType);
+
+    if (tensorType.getShape() != memrefType.getShape())
+      return emitError() << "shapes do not match";
+
+    if (tensorType.getElementType() != memrefType.getElementType())
+      return emitError() << "element types do not match";
+
+    return mlir::success();
+  }
+};
 
 template <typename MemRef>
 struct BuiltinMemRefExternalModel

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

@@ -0,0 +1,21 @@
+//===- BufferizationTypeInterfaces.cpp - Type Interfaces --------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Bufferization/IR/BufferizationTypeInterfaces.h"
+
+//===----------------------------------------------------------------------===//
+// Bufferization Type Interfaces
+//===----------------------------------------------------------------------===//
+
+namespace mlir {
+namespace bufferization {
+
+#include "mlir/Dialect/Bufferization/IR/BufferizationTypeInterfaces.cpp.inc"
+
+} // namespace bufferization
+} // namespace mlir

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

@@ -6,7 +6,7 @@ add_mlir_dialect_library(MLIRBufferizationDialect
   BufferizationDialect.cpp
   BufferViewFlowOpInterface.cpp
   UnstructuredControlFlow.cpp
-  BufferizationConversionInterface.cpp
+  BufferizationTypeInterfaces.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Bufferization

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

@@ -11,7 +11,6 @@
 #include "TestTypes.h"
 #include "mlir/Bytecode/BytecodeImplementation.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
-#include "mlir/Dialect/Bufferization/IR/BufferizationConversionInterface.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/AsmState.h"
@@ -285,44 +284,6 @@ getDynamicCustomParserPrinterOp(TestDialect *dialect) {
                                   verifier, regionVerifier, parser, printer);
 }
 
-namespace {
-
-struct TestConverter : bufferization::ConversionDialectInterface {
-  TestConverter(Dialect *dialect)
-      : bufferization::ConversionDialectInterface(dialect) {}
-
-  FailureOr<bufferization::BufferLikeType>
-  getBufferType(Value value, const bufferization::BufferizationOptions &options,
-                const bufferization::BufferizationState &state,
-                function_ref<InFlightDiagnostic(const Twine &)> emitError)
-      const override {
-    auto testTensor = dyn_cast<TestTensorType>(value.getType());
-    if (!testTensor)
-      return emitError("expected TestTensorType");
-
-    return cast<bufferization::BufferLikeType>(
-        TestMemrefType::get(value.getContext(), testTensor.getShape(),
-                            testTensor.getElementType(), nullptr));
-  }
-
-  LogicalResult typesMatch(bufferization::TensorLikeType tensor,
-                           bufferization::BufferLikeType buffer,
-                           function_ref<InFlightDiagnostic(const Twine &)>
-                               emitError) const override {
-    auto testTensor = dyn_cast<TestTensorType>(tensor);
-    auto testMemref = dyn_cast<TestMemrefType>(buffer);
-    if (!testTensor || !testMemref)
-      return emitError("expected TestTensorType and TestMemrefType");
-
-    const bool valid =
-        testTensor.getShape() == testMemref.getShape() &&
-        testTensor.getElementType() == testMemref.getElementType();
-    return success(valid);
-  }
-};
-
-} // namespace
-
 //===----------------------------------------------------------------------===//
 // TestDialect
 //===----------------------------------------------------------------------===//
@@ -372,7 +333,6 @@ void TestDialect::initialize() {
   registerDynamicOp(getDynamicCustomParserPrinterOp(this));
   registerInterfaces();
   allowUnknownOperations();
-  addInterface<TestConverter>();
 
   // Instantiate our fallback op interface that we'll use on specific
   // unregistered op.