[mlir][bufferization] Add tensor-like and buffer-like interfaces

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

@@ -0,0 +1,20 @@
+//===- BufferizationTypeInterfaces.h - 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_BUFFERIZATION_IR_BUFFERIZATIONTYPEINTERFACES_H_
+#define MLIR_DIALECT_BUFFERIZATION_IR_BUFFERIZATIONTYPEINTERFACES_H_
+
+#include "mlir/IR/BuiltinTypeInterfaces.h" // for ShapedTypeInterface
+
+//===----------------------------------------------------------------------===//
+// Bufferization Type Interfaces
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Bufferization/IR/BufferizationTypeInterfaces.h.inc"
+
+#endif // MLIR_DIALECT_BUFFERIZATION_IR_BUFFERIZATIONTYPEINTERFACES_H_

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

@@ -0,0 +1,49 @@
+//===- BufferizationTypeInterfaces.td - Type Interfaces ----*- tablegen -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This is the definition file for type interfaces used in Bufferization.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef BUFFERIZATION_TYPE_INTERFACES
+#define BUFFERIZATION_TYPE_INTERFACES
+
+include "mlir/IR/OpBase.td"
+include "mlir/IR/BuiltinTypeInterfaces.td"
+
+def Bufferization_TensorLikeTypeInterface
+    : TypeInterface<"TensorLikeType", [ShapedTypeInterface]> {
+  let cppNamespace = "::mlir::bufferization";
+  let description = [{
+    Indicates that the type that attaches this interface can be treated as a
+    tensor type (similarly to a MLIR builtin tensor) during bufferization.
+
+    Implementing this interface means that the type also implements
+    ShapedTypeInterface.
+
+    The interface currently has no methods as it is used by types to opt into
+    being supported by the bufferization procedures.
+  }];
+}
+
+def Bufferization_MemRefLikeTypeInterface
+    : TypeInterface<"MemRefLikeType", [ShapedTypeInterface]> {
+  let cppNamespace = "::mlir::bufferization";
+  let description = [{
+    Indicates that the type that attaches this interface can be treated as a
+    memref type (similarly to a MLIR builtin memref) during bufferization.
+
+    Implementing this interface means that the type also implements
+    ShapedTypeInterface.
+
+    The interface currently has no methods as it is used by types to opt into
+    being supported by the bufferization procedures.
+  }];
+}
+
+#endif // BUFFERIZATION_TYPE_INTERFACES

mlir/include/mlir/Dialect/Bufferization/IR/CMakeLists.txt

@@ -10,3 +10,9 @@ mlir_tablegen(BufferizationEnums.h.inc -gen-enum-decls)
 mlir_tablegen(BufferizationEnums.cpp.inc -gen-enum-defs)
 add_public_tablegen_target(MLIRBufferizationEnumsIncGen)
 add_dependencies(mlir-headers MLIRBufferizationEnumsIncGen)
+
+set(LLVM_TARGET_DEFINITIONS BufferizationTypeInterfaces.td)
+mlir_tablegen(BufferizationTypeInterfaces.h.inc -gen-type-interface-decls)
+mlir_tablegen(BufferizationTypeInterfaces.cpp.inc -gen-type-interface-defs)
+add_public_tablegen_target(MLIRBufferizationTypeInterfacesIncGen)
+add_dependencies(mlir-headers MLIRBufferizationTypeInterfacesIncGen)

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

@@ -9,8 +9,10 @@
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
+#include "mlir/Dialect/Bufferization/IR/BufferizationTypeInterfaces.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/IR/BuiltinTypes.h"
 #include "mlir/Interfaces/FunctionInterfaces.h"
 #include "mlir/Transforms/InliningUtils.h"
 
@@ -51,6 +53,16 @@ struct BufferizationInlinerInterface : public DialectInlinerInterface {
     return true;
   }
 };
+
+template <typename Tensor>
+struct BuiltinTensorExternalModel
+    : TensorLikeType::ExternalModel<BuiltinTensorExternalModel<Tensor>,
+                                    Tensor> {};
+
+template <typename MemRef>
+struct BuiltinMemRefExternalModel
+    : MemRefLikeType::ExternalModel<BuiltinMemRefExternalModel<MemRef>,
+                                    MemRef> {};
 } // namespace
 
 //===----------------------------------------------------------------------===//
@@ -63,6 +75,15 @@ void mlir::bufferization::BufferizationDialect::initialize() {
 #include "mlir/Dialect/Bufferization/IR/BufferizationOps.cpp.inc"
       >();
   addInterfaces<BufferizationInlinerInterface>();
+
+  RankedTensorType::attachInterface<
+      BuiltinTensorExternalModel<RankedTensorType>>(*getContext());
+  UnrankedTensorType::attachInterface<
+      BuiltinTensorExternalModel<UnrankedTensorType>>(*getContext());
+  MemRefType::attachInterface<BuiltinMemRefExternalModel<MemRefType>>(
+      *getContext());
+  UnrankedMemRefType::attachInterface<
+      BuiltinMemRefExternalModel<UnrankedMemRefType>>(*getContext());
 }
 
 LogicalResult BufferizationDialect::verifyRegionArgAttribute(

mlir/test/lib/Dialect/Test/CMakeLists.txt

@@ -93,6 +93,7 @@ mlir_target_link_libraries(MLIRTestDialect PUBLIC
   MLIRTransformUtils
   MLIRTransforms
   MLIRValueBoundsOpInterface
+  MLIRBufferizationDialect
 )
 
 add_mlir_translation_library(MLIRTestFromLLVMIRTranslation

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

@@ -19,6 +19,7 @@ include "TestAttrDefs.td"
 include "TestInterfaces.td"
 include "mlir/IR/BuiltinTypes.td"
 include "mlir/Interfaces/DataLayoutInterfaces.td"
+include "mlir/Dialect/Bufferization/IR/BufferizationTypeInterfaces.td"
 
 // All of the types will extend this class.
 class Test_Type<string name, list<Trait> traits = []>
@@ -403,4 +404,47 @@ def TestTypeOpAsmTypeInterface : Test_Type<"TestTypeOpAsmTypeInterface",
   let mnemonic = "op_asm_type_interface";
 }
 
+def TestTensorType : Test_Type<"TestTensor", [Bufferization_TensorLikeTypeInterface]> {
+  let mnemonic = "test_tensor";
+  let parameters = (ins
+    ArrayRefParameter<"int64_t">:$shape,
+    "mlir::Type":$elementType
+  );
+  let assemblyFormat = "`<` `[` $shape `]` `,` $elementType `>`";
+
+  let extraClassDeclaration = [{
+    // ShapedTypeInterface:
+    bool hasRank() const {
+      return true;
+    }
+    test::TestTensorType cloneWith(std::optional<llvm::ArrayRef<int64_t>> shape,
+                                   mlir::Type elementType) const {
+      return test::TestTensorType::get(
+        getContext(), shape.value_or(getShape()), elementType);
+    }
+  }];
+}
+
+def TestMemrefType : Test_Type<"TestMemref", [Bufferization_MemRefLikeTypeInterface]> {
+  let mnemonic = "test_memref";
+  let parameters = (ins
+    ArrayRefParameter<"int64_t">:$shape,
+    "mlir::Type":$elementType,
+    DefaultValuedParameter<"mlir::Attribute", "nullptr">:$memSpace
+  );
+  let assemblyFormat = "`<` `[` $shape `]` `,` $elementType (`,` $memSpace^)? `>`";
+
+  let extraClassDeclaration = [{
+    // ShapedTypeInterface:
+    bool hasRank() const {
+      return true;
+    }
+    test::TestMemrefType cloneWith(std::optional<llvm::ArrayRef<int64_t>> shape,
+                                   mlir::Type elementType) const {
+      return test::TestMemrefType::get(
+        getContext(), shape.value_or(getShape()), elementType, getMemSpace());
+    }
+  }];
+}
+
 #endif // TEST_TYPEDEFS

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

@@ -18,6 +18,7 @@
 #include <tuple>
 
 #include "TestTraits.h"
+#include "mlir/Dialect/Bufferization/IR/BufferizationTypeInterfaces.h"
 #include "mlir/IR/Diagnostics.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/DialectImplementation.h"

mlir/unittests/IR/InterfaceTest.cpp

@@ -6,6 +6,8 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
+#include "mlir/Dialect/Bufferization/IR/BufferizationTypeInterfaces.h"
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/BuiltinDialect.h"
 #include "mlir/IR/BuiltinOps.h"
@@ -84,3 +86,78 @@ TEST(InterfaceTest, TestImplicitConversion) {
   typeA = typeB;
   EXPECT_EQ(typeA, typeB);
 }
+
+TEST(InterfaceTest, TestBuiltinTensorIsTensorLikeType) {
+  MLIRContext context;
+  // Note: attaches external model to builtins
+  context.loadDialect<bufferization::BufferizationDialect>();
+
+  auto builtinRankedTensor = mlir::RankedTensorType::get(
+      {1, 2, 3}, mlir::IntegerType::get(&context, 32));
+  EXPECT_TRUE(mlir::isa<bufferization::TensorLikeType>(builtinRankedTensor));
+  EXPECT_FALSE(mlir::isa<bufferization::MemRefLikeType>(builtinRankedTensor));
+
+  auto builtinUnrankedTensor =
+      mlir::UnrankedTensorType::get(mlir::IntegerType::get(&context, 32));
+  EXPECT_TRUE(mlir::isa<bufferization::TensorLikeType>(builtinUnrankedTensor));
+  EXPECT_FALSE(mlir::isa<bufferization::MemRefLikeType>(builtinUnrankedTensor));
+}
+
+TEST(InterfaceTest, TestCustomTensorIsTensorLikeType) {
+  MLIRContext context;
+  context.loadDialect<test::TestDialect>();
+
+  auto customTensorType = test::TestTensorType::get(
+      &context, {1, 2, 3}, mlir::IntegerType::get(&context, 32));
+  EXPECT_TRUE(mlir::isa<bufferization::TensorLikeType>(customTensorType));
+
+  auto customCloneType = customTensorType.cloneWith(
+      ArrayRef<int64_t>{3, 4, 5}, customTensorType.getElementType());
+  EXPECT_EQ(customTensorType.getElementType(),
+            customCloneType.getElementType());
+  EXPECT_TRUE(mlir::isa<bufferization::TensorLikeType>(customCloneType));
+  EXPECT_TRUE(mlir::isa<test::TestTensorType>(customCloneType));
+
+  // user-specified conversions
+  bufferization::TensorLikeType baseCopy = customTensorType;
+  std::ignore = baseCopy;
+}
+
+TEST(InterfaceTest, TestBuiltinMemrefIsMemRefLikeType) {
+  MLIRContext context;
+  // Note: attaches external model to builtins
+  context.loadDialect<bufferization::BufferizationDialect>();
+
+  auto builtinRankedMemref =
+      mlir::MemRefType::get({1, 2, 3}, mlir::IntegerType::get(&context, 32));
+  EXPECT_TRUE(mlir::isa<bufferization::MemRefLikeType>(builtinRankedMemref));
+  EXPECT_FALSE(mlir::isa<bufferization::TensorLikeType>(builtinRankedMemref));
+
+  auto builtinUnrankedMemref = mlir::UnrankedMemRefType::get(
+      mlir::IntegerType::get(&context, 32), nullptr);
+  EXPECT_TRUE(mlir::isa<bufferization::MemRefLikeType>(builtinUnrankedMemref));
+  EXPECT_FALSE(mlir::isa<bufferization::TensorLikeType>(builtinUnrankedMemref));
+}
+
+TEST(InterfaceTest, TestCustomMemrefIsMemRefLikeType) {
+  MLIRContext context;
+  context.loadDialect<test::TestDialect>();
+
+  auto customMemrefType = test::TestMemrefType::get(
+      &context, {1, 2, 3}, mlir::IntegerType::get(&context, 32),
+      mlir::StringAttr::get(&context, "some_memspace"));
+  EXPECT_TRUE(mlir::isa<bufferization::MemRefLikeType>(customMemrefType));
+
+  auto customCloneType = customMemrefType.cloneWith(
+      ArrayRef<int64_t>{3, 4, 5}, customMemrefType.getElementType());
+  EXPECT_EQ(customMemrefType.getElementType(),
+            customCloneType.getElementType());
+  EXPECT_TRUE(mlir::isa<bufferization::MemRefLikeType>(customCloneType));
+  EXPECT_TRUE(mlir::isa<test::TestMemrefType>(customCloneType));
+  EXPECT_EQ(customMemrefType.getMemSpace(),
+            mlir::cast<test::TestMemrefType>(customCloneType).getMemSpace());
+
+  // user-specified conversions
+  bufferization::MemRefLikeType baseCopy = customMemrefType;
+  std::ignore = baseCopy;
+}