[NFC][mlir] Add support for llvm style casting for mlir types

[Tyker]

Summary

Note:
when operating on a Type hierarchy with LeafType inheriting from MiddleType which inherits from mlir::Type.
calling LeafType::classof(MiddleType) will always return false.
because classof call the static getTypeID from its parent instead of the dynamic Type::getTypeID
so classof in this context will check if the TypeID of LeafType is the same as the TypeID of MiddleType which is always false.
It is bypassed in this commit inside CastInfo<To, From>::isPossible by calling classof with an mlir::Type.
but other unsuspecting users of LeafType::classof(MiddleType) would still get an incorrect result.

[mehdi_amini]

Sweet!

mlir/include/mlir/IR/Types.h

324

Can you add a one-line comment here?

[bzcheeseman]

+1 to the concept, many thanks for this! I think it makes more sense to implement the member versions with the llvm-style versions, it should simplify the member style code because then you won't have to manually implement the classof/isa/cast innards it'd just be cast<T>(*this) or something.

mlir/unittests/IR/TypeTest.cpp

29

_or_null is deprecated, please use _if_present

[Tyker]

In D132009#3728861, @bzcheeseman wrote:

+1 to the concept, many thanks for this! I think it makes more sense to implement the member versions with the llvm-style versions, it should simplify the member style code because then you won't have to manually implement the classof/isa/cast innards it'd just be cast<T>(*this) or something.

I think both casting system will stay available. so I am not sure it matters.
i found out that some Types can "inherit" from other type so it makes sense to cast from a more then just mlir::Type.
so I added support for that. it may be good to rereview

mlir/unittests/IR/TypeTest.cpp

29

it isn't appear to be attributed as such.

[bzcheeseman]

mlir/unittests/IR/TypeTest.cpp

29

We didn't want to add a bunch of build spew for all LLVM users immediately, it will be marked deprecated at some point (hopefully) soon.

[mehdi_amini]

I think both casting system will stay available. so I am not sure it matters.

Actually we should likely deprecated the old-style in the future, so the suggestion above gives a nicer path to get there.

[rriddle]

In D132009#3731678, @mehdi_amini wrote:

I think both casting system will stay available. so I am not sure it matters.

Actually we should likely deprecated the old-style in the future, so the suggestion above gives a nicer path to get there.

+1, we intend to get rid of the old method.

[Tyker]

I updated to to have the member version of cast/isa/dyn_cast depend on the out of line ones.

and I found an interesting bug. when having a type hierarchy like the following

struct LeafType;

struct MiddleType : Type::TypeBase<MiddleType, Type, TypeStorage> {
  using Base::Base;
  static bool classof(Type ty) {
    return ty.getTypeID() == TypeID::get<LeafType>() || Base::classof(ty);
  }
};

struct LeafType : Type::TypeBase<LeafType, MiddleType, TypeStorage> {
  using Base::Base;
};

calling LeafType::classof(MiddleType) will always return false.

here is the current implementation of StorageUserBase::classof

template <typename T>
static bool classof(T val) {
  static_assert(std::is_convertible<ConcreteT, T>::value,
                "casting from a non-convertible type");
  return val.getTypeID() == getTypeID();
}

When T is not mlir::Type, val.getTypeID() will call the static getTypeID from its parent instead of the dynamic Type::getTypeID
so classof in the context will check if the TypeID of LeafType is the same as the TypeID of MiddleType which is always false.

It is "solved" in the current version of the patch inside CastInfo<To, From>::isPossible by calling classof with an mlir::Type.
but other unsuspecting users of LeafType::classof(MiddleType) would still get an incorrect result.

is there any benefit to having Type::getTypeID and StorageUserBase::getTypeID have the same name ?
should the StorageUserBase::getTypeID change name to avoid this issue?
or make classof Take a mlir::Type.

also is StorageUserBase used for attributes ?

[bzcheeseman]

mlir/include/mlir/IR/Types.h

328

Not sure if you saw https://llvm.org/docs/HowToSetUpLLVMStyleRTTI.html#advanced-use-cases but you might be able to use some cast traits since Type is nullable. I think you can use CastIsPossible, NullableValueCastFailed, and DefaultDoCastIfPossible. The default CastIsPossible also checks if the types are the same or related by inheritance and always allows casting from derived to base.

For context, that example nullable type was more or less based on mlir::Type.

[Tyker]

mlir/include/mlir/IR/Types.h

328

thanks for the tip, I added NullableValueCastFailed and DefaultDoCastIfPossible

[bzcheeseman]

Thanks for fixing the or_null to if_present too (where it makes sense). LGTM as long as everything passes CI :)

mlir/include/mlir/IR/Types.h

341

I take it this means you can't use the CastIsPossible<To, From> trait?

[Tyker]

mlir/include/mlir/IR/Types.h

341

I couldn't get it to work easily. I think i could have by insisting a bit more but it was in my opinion more trouble then it was worth.

CastIsPossible by default passes pointers to classof. so it doesn't work for us.
and adding our own specialization would mean we are ambiguous with the CastIsPossible that detect the up cast.

we could always add an and not up-cast in our SFINAE but this is in my opinion more trouble that it is worth.

[Tyker]

I would still like a comment for what to do about this issue (which i bypassed):

In D132009#3733887, @Tyker wrote:

and I found an interesting bug. when having a type hierarchy like the following

struct LeafType;

struct MiddleType : Type::TypeBase<MiddleType, Type, TypeStorage> {
  using Base::Base;
  static bool classof(Type ty) {
    return ty.getTypeID() == TypeID::get<LeafType>() || Base::classof(ty);
  }
};

struct LeafType : Type::TypeBase<LeafType, MiddleType, TypeStorage> {
  using Base::Base;
};

calling LeafType::classof(MiddleType) will always return false.

here is the current implementation of StorageUserBase::classof

template <typename T>
static bool classof(T val) {
  static_assert(std::is_convertible<ConcreteT, T>::value,
                "casting from a non-convertible type");
  return val.getTypeID() == getTypeID();
}

When T is not mlir::Type, val.getTypeID() will call the static getTypeID from its parent instead of the dynamic Type::getTypeID
so classof in the context will check if the TypeID of LeafType is the same as the TypeID of MiddleType which is always false.

It is "solved" in the current version of the patch inside CastInfo<To, From>::isPossible by calling classof with an mlir::Type.
but other unsuspecting users of LeafType::classof(MiddleType) would still get an incorrect result.

is there any benefit to having Type::getTypeID and StorageUserBase::getTypeID have the same name ?
should the StorageUserBase::getTypeID change name to avoid this issue?
or make classof Take a mlir::Type.

also is StorageUserBase used for attributes ?

[bondhugula]

Can you please add a commit summary along the lines captured in the comments?

[Tyker]

has been committed as bb63d249f8612f87e819071663d81f516a2bec74

mlir/include/mlir/IR/Types.h

public:
   Type &operator=(const Type &other) = default;
 
   bool operator==(Type other) const { return impl == other.impl; }
   bool operator!=(Type other) const { return !(*this == other); }
   explicit operator bool() const { return impl; }
 
   bool operator!() const { return impl == nullptr; }
 
-  template <typename U>
-  bool isa() const;
-  template <typename First, typename Second, typename... Rest>
+  template <typename... Tys>
   bool isa() const;
-  template <typename First, typename... Rest>
+  template <typename... Tys>
   bool isa_and_nonnull() const;
   template <typename U>
   U dyn_cast() const;
   template <typename U>
   U dyn_cast_or_null() const;
   template <typename U>
   U cast() const;
 
public:
   template <template <typename T> class Trait>
   bool hasTrait() {
     return getAbstractType().hasTrait<Trait>();
   }
 
   /// Return the abstract type descriptor for this type.
   const AbstractTy &getAbstractType() { return impl->getAbstractType(); }
 
+  /// Return the Type implementation.
+  ImplType *getImpl() const { return impl; }
+
 protected:
   ImplType *impl{nullptr};
 };
 
 inline raw_ostream &operator<<(raw_ostream &os, Type type) {
   type.print(os);
   return os;
 }

 // Type Utils
 //===----------------------------------------------------------------------===//
 
 // Make Type hashable.
 inline ::llvm::hash_code hash_value(Type arg) {
   return DenseMapInfo<const Type::ImplType *>::getHashValue(arg.impl);
 }
 
-template <typename U>
-bool Type::isa() const {
-  assert(impl && "isa<> used on a null type.");
-  return U::classof(*this);
-}
-
-template <typename First, typename Second, typename... Rest>
+template <typename... Tys>
 bool Type::isa() const {
-  return isa<First>() || isa<Second, Rest...>();
+  return llvm::isa<Tys...>(*this);
 }
 
-template <typename First, typename... Rest>
+template <typename... Tys>
 bool Type::isa_and_nonnull() const {
-  return impl && isa<First, Rest...>();
+  return llvm::isa_and_present<Tys...>(*this);
 }
 
 template <typename U>
 U Type::dyn_cast() const {
-  return isa<U>() ? U(impl) : U(nullptr);
+  return llvm::dyn_cast<U>(*this);
 }
+
 template <typename U>
 U Type::dyn_cast_or_null() const {
-  return (impl && isa<U>()) ? U(impl) : U(nullptr);
+  return llvm::dyn_cast_or_null<U>(*this);
 }
+
 template <typename U>
 U Type::cast() const {
-  assert(isa<U>());
-  return U(impl);
+  return llvm::cast<U>(*this);
 }
 
 } // namespace mlir
 
 namespace llvm {
 
 // Type hash just like pointers.
 template <>
static inline void *getAsVoidPointer(mlir::Type I) {
     return const_cast<void *>(I.getAsOpaquePointer());
   }
   static inline mlir::Type getFromVoidPointer(void *P) {
     return mlir::Type::getFromOpaquePointer(P);
   }
   static constexpr int NumLowBitsAvailable = 3;
 };
 
+/// Add support for llvm style casts.
+/// We provide a cast between To and From if From is mlir::Type or derives from
+/// it
+template <typename To, typename From>
+struct CastInfo<To, From,
+                typename std::enable_if<
+                    std::is_same_v<mlir::Type, std::remove_const_t<From>> ||
+                    std::is_base_of_v<mlir::Type, From>>::type>
+    : NullableValueCastFailed<To>,
+      DefaultDoCastIfPossible<To, From, CastInfo<To, From>> {
+  /// Arguments are taken as mlir::Type here and not as From.
+  /// Because when casting from an intermediate type of the hierarchy to one of
+  /// its children, the val.getTypeID() inside T::classof will use the static
+  /// getTypeID of the parent instead of the non-static Type::getTypeID return
+  /// the dynamic ID. so T::classof would end up comparing the static TypeID of
+  /// The children to the static TypeID of its parent making it impossible to
+  /// downcast from the parent to the child
+  static inline bool isPossible(mlir::Type ty) {
+    /// Return a constant true instead of a dynamic true when casting to self or
+    /// up the hierarchy
+    return std::is_same_v<To, std::remove_const_t<From>> ||
+           std::is_base_of_v<To, From> || To::classof(ty);
+  }
+  static inline To doCast(mlir::Type ty) { return To(ty.getImpl()); }
+};
+
 } // namespace llvm
 
 #endif // MLIR_IR_TYPES_H

mlir/unittests/IR/CMakeLists.txt

 add_mlir_unittest(MLIRIRTests
   AttributeTest.cpp
   DialectTest.cpp
   InterfaceTest.cpp
   InterfaceAttachmentTest.cpp
   OperationSupportTest.cpp
   PatternMatchTest.cpp
   ShapedTypeTest.cpp
   SubElementInterfaceTest.cpp
+  TypeTest.cpp
 
   DEPENDS
   MLIRTestInterfaceIncGen
 )
 target_include_directories(MLIRIRTests PRIVATE "${MLIR_BINARY_DIR}/test/lib/Dialect/Test")
 target_link_libraries(MLIRIRTests
   PRIVATE
   MLIRIR
   MLIRTestDialect)

mlir/unittests/IR/TypeTest.cpp

+//===- TypeTest.cpp - Type API unit tests ---------------------------------===//
+//
+// 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/IR/Dialect.h"
+#include "mlir/IR/Types.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/Value.h"
+#include "gtest/gtest.h"
+
+using namespace mlir;
+
+/// Mock implementations of a Type hierarchy
+struct LeafType;
+
+struct MiddleType : Type::TypeBase<MiddleType, Type, TypeStorage> {
+  using Base::Base;
+  static bool classof(Type ty) {
+    return ty.getTypeID() == TypeID::get<LeafType>() || Base::classof(ty);
+  }
+};
+
+struct LeafType : Type::TypeBase<LeafType, MiddleType, TypeStorage> {
+  using Base::Base;
+};
+
+struct FakeDialect : Dialect {
+  FakeDialect(MLIRContext *context)
+      : Dialect(getDialectNamespace(), context, TypeID::get<FakeDialect>()) {
+    addTypes<MiddleType, LeafType>();
+  }
+  static constexpr ::llvm::StringLiteral getDialectNamespace() {
+    return ::llvm::StringLiteral("fake");
+  }
+};
+
+TEST(Type, Casting) {
+  MLIRContext ctx;
+  ctx.loadDialect<FakeDialect>();
+
+  Type intTy = IntegerType::get(&ctx, 8);
+  Type nullTy;
+  MiddleType middleTy = MiddleType::get(&ctx);
+  MiddleType leafTy = LeafType::get(&ctx);
+  Type leaf2Ty = LeafType::get(&ctx);
+
+  EXPECT_TRUE(isa<IntegerType>(intTy));
+  EXPECT_FALSE(isa<FunctionType>(intTy));
+  EXPECT_FALSE(llvm::isa_and_present<IntegerType>(nullTy));
+  EXPECT_TRUE(isa<MiddleType>(middleTy));
+  EXPECT_FALSE(isa<LeafType>(middleTy));
+  EXPECT_TRUE(isa<MiddleType>(leafTy));
+  EXPECT_TRUE(isa<LeafType>(leaf2Ty));
+  EXPECT_TRUE(isa<LeafType>(leafTy));
+
+  EXPECT_TRUE(static_cast<bool>(dyn_cast<IntegerType>(intTy)));
+  EXPECT_FALSE(static_cast<bool>(dyn_cast<FunctionType>(intTy)));
+  EXPECT_FALSE(static_cast<bool>(llvm::cast_if_present<FunctionType>(nullTy)));
+  EXPECT_FALSE(
+      static_cast<bool>(llvm::dyn_cast_if_present<IntegerType>(nullTy)));
+
+  EXPECT_EQ(8u, cast<IntegerType>(intTy).getWidth());
+}