add CastInfo to support cast Interface to Op

mlir/include/mlir/IR/OpDefinition.h

@@ -2144,6 +2144,9 @@ struct DenseMapInfo<T,
   static bool isEqual(T lhs, T rhs) { return lhs == rhs; }
 };
 
+/// Add support for llvm style casts.
+/// We provide a cast between To and From if To and From is mlir::OpInterface or
+/// derives from it.
 template <typename To, typename From>
 struct CastInfo<
     To, From,
@@ -2157,7 +2160,7 @@ struct CastInfo<
         void>> : NullableValueCastFailed<To>,
                  DefaultDoCastIfPossible<To, From, CastInfo<To, From>> {
 
-  static bool isPossible(From &val) {
+  static inline bool isPossible(From &val) {
     if constexpr (std::is_same_v<To, From>)
       return true;
     else
@@ -2166,7 +2169,53 @@ struct CastInfo<
               const_cast<std::remove_const_t<From> &>(val).getOperation());
   }
 
-  static To doCast(From &val) {
+  static inline To doCast(From &val) {
+    return To(const_cast<std::remove_const_t<From> &>(val).getOperation());
+  }
+};
+
+template <typename OpT, typename = void>
+struct is_concrete_op_type : public std::false_type {};
+
+template <typename OpT, template <typename T> typename... Traits>
+constexpr auto concrete_op_base_type_impl(std::tuple<Traits<OpT>...>) {
+  return mlir::Op<OpT, Traits...>(nullptr);
+}
+
+template <typename OpT>
+using concrete_op_base_type =
+    decltype(concrete_op_base_type_impl<OpT>(typename OpT::traits()));
+
+template <typename OpT>
+struct is_concrete_op_type<
+    OpT, std::enable_if_t<std::is_base_of_v<concrete_op_base_type<OpT>, OpT>>>
+    : public std::true_type {};
+
+/// Add support for llvm style casts.
+/// We provide a cast between To and From if To is mlir::Op<ConcreteType,
+/// Trait0, Trait1, ...> or derives from it and From is mlir::OpInterface or
+/// derives from it.
+template <typename To, typename From>
+struct CastInfo<
+    To, From,
+    std::enable_if_t<
+        is_concrete_op_type<To>() &&
+        std::is_base_of_v<mlir::OpInterface<std::remove_const_t<From>,
+                                            typename std::remove_const_t<
+                                                From>::InterfaceTraits>,
+                          std::remove_const_t<From>>>>
+    : NullableValueCastFailed<To>,
+      DefaultDoCastIfPossible<To, From, CastInfo<To, From>> {
+
+  static inline bool isPossible(From &val) {
+    if constexpr (std::is_same_v<To, From>)
+      return true;
+    else
+      return isa<To>(
+          const_cast<std::remove_const_t<From> &>(val).getOperation());
+  }
+
+  static inline To doCast(From &val) {
     return To(const_cast<std::remove_const_t<From> &>(val).getOperation());
   }
 };

mlir/include/mlir/TableGen/Class.h

@@ -521,6 +521,8 @@ class ParentClass {
   /// Write the parent class declaration.
   void writeTo(raw_indented_ostream &os) const;
 
+  friend class OpClass;
+
 private:
   /// The fully resolved C++ name of the parent class.
   std::string name;

mlir/tools/mlir-tblgen/OpClass.cpp

@@ -36,7 +36,16 @@ OpClass::OpClass(StringRef name, std::string extraClassDeclaration,
 }
 
 void OpClass::finalize() {
+  std::string traitList;
+  llvm::raw_string_ostream os(traitList);
+  iterator_range parentTemplateParams(std::begin(parent.templateParams) + 1,
+                                      std::end(parent.templateParams));
+  llvm::interleaveComma(parentTemplateParams, os, [&](auto &trait) {
+    os << trait << "<" << getClassName().str() << ">";
+  });
+  declare<UsingDeclaration>("traits", "std::tuple<" + traitList + ">");
   Class::finalize();
+
   declare<VisibilityDeclaration>(Visibility::Public);
   declare<ExtraClassDeclaration>(extraClassDeclaration, extraClassDefinition);
 }

mlir/unittests/IR/InterfaceTest.cpp

@@ -88,7 +88,7 @@ TEST(InterfaceTest, TestImplicitConversion) {
   EXPECT_EQ(typeA, typeB);
 }
 
-TEST(OperationInterfaceTest, CastOpToInterface) {
+TEST(OperationInterfaceTest, CastInterfaceToOpOrInterface) {
   DialectRegistry registry;
   MLIRContext ctx;
 
@@ -105,13 +105,20 @@ TEST(OperationInterfaceTest, CastOpToInterface) {
   OwningOpRef<ModuleOp> module = parseSourceString<ModuleOp>(ir, &ctx);
   Operation &op = cast<func::FuncOp>(module->front()).getBody().front().front();
 
+  static_assert(std::is_base_of_v<llvm::concrete_op_base_type<arith::AddIOp>,
+                                  arith::AddIOp>,
+                "");
+  static_assert(llvm::is_concrete_op_type<arith::AddIOp>(), "");
+  static_assert(!llvm::is_concrete_op_type<OpAsmOpInterface>(), "");
+
   OpAsmOpInterface interface = llvm::cast<OpAsmOpInterface>(op);
 
-  bool constantOp =
-      llvm::TypeSwitch<OpAsmOpInterface, bool>(interface)
-          .Case<VectorUnrollOpInterface, arith::ConstantOp>([&](auto op) {
-            return std::is_same_v<decltype(op), arith::ConstantOp>;
-          });
+  bool constantOp = llvm::TypeSwitch<OpAsmOpInterface, bool>(interface)
+                        .Case<arith::AddIOp, arith::ConstantOp>([&](auto op) {
+                          bool is_same =
+                              std::is_same_v<decltype(op), arith::ConstantOp>;
+                          return is_same;
+                        });
 
   EXPECT_TRUE(constantOp);