[MLIR][Operation] Simplify Operation casting, NFC

[bzcheeseman]

Summary

We can simplify the code needed to implement dyn_cast/cast/isa support for MLIR operations with documented interfaces via the CastInfo structures. This will also provide an example of how to use CastInfo.

Depends on D123901

[rriddle]

mlir/include/mlir/IR/Operation.h

827–833

I wonder if we should have a trait for forwarding the reference variant to the pointer variant of a ToValueCast. Though maybe this is rare enough where it doesn't matter.

[bzcheeseman]

mlir/include/mlir/IR/Operation.h

827–833

Yeah I don't think it happens that often, not many systems have fallible reference-to-reference casting. But luckily, if it does, we can add a trait for it :)

[rriddle]

mlir/include/mlir/IR/Operation.h

827–833

Well this is reference-to-value, which for anything that crosses the reference-value boundary would be common.

[bzcheeseman]

mlir/include/mlir/IR/Operation.h

827–833

Fair enough - though we're just adding a convenience specialization on top so that we don't have to use dyn_cast<T>(&operation). If more MLIR structures want this it should be easy to add a forwarding trait, I just haven't seen the demand yet.

mlir/include/mlir/IR/Operation.h

public:
   /// the instruction is unlinked.
   Region *getParentRegion() { return block ? block->getParent() : nullptr; }
 
   /// Returns the closest surrounding operation that contains this operation
   /// or nullptr if this is a top-level operation.
   Operation *getParentOp() { return block ? block->getParentOp() : nullptr; }
 
   /// Return the closest surrounding parent operation that is of type 'OpTy'.
-  template <typename OpTy> OpTy getParentOfType() {
+  template <typename OpTy>
+  OpTy getParentOfType() {
     auto *op = this;
     while ((op = op->getParentOp()))
       if (auto parentOp = dyn_cast<OpTy>(op))
         return parentOp;
     return OpTy();
   }
 
   /// Returns the closest surrounding parent operation with trait `Trait`.
public:
   /// - the elements in "operands" will correspond directly to the operands of
   /// the operation, but may be null if non-constant. If folding is successful,
   /// this fills in the `results` vector. If not, `results` is unspecified.
   LogicalResult fold(ArrayRef<Attribute> operands,
                      SmallVectorImpl<OpFoldResult> &results);
 
   /// Returns true if the operation was registered with a particular trait, e.g.
   /// hasTrait<OperandsAreSignlessIntegerLike>().
-  template <template <typename T> class Trait> bool hasTrait() {
+  template <template <typename T> class Trait>
+  bool hasTrait() {
     return name.hasTrait<Trait>();
   }
 
   /// Returns true if the operation *might* have the provided trait. This
   /// means that either the operation is unregistered, or it was registered with
   /// the provide trait.
-  template <template <typename T> class Trait> bool mightHaveTrait() {
+  template <template <typename T> class Trait>
+  bool mightHaveTrait() {
     return name.mightHaveTrait<Trait>();
   }
 
   //===--------------------------------------------------------------------===//
   // Operation Walkers
   //===--------------------------------------------------------------------===//
 
   /// Walk the operation by calling the callback for each nested operation

 inline raw_ostream &operator<<(raw_ostream &os, const Operation &op) {
   const_cast<Operation &>(op).print(os, OpPrintingFlags().useLocalScope());
   return os;
 }
 
 } // namespace mlir
 
 namespace llvm {
-/// Provide isa functionality for operation casts.
-template <typename T> struct isa_impl<T, ::mlir::Operation> {
-  static inline bool doit(const ::mlir::Operation &op) {
-    return T::classof(const_cast<::mlir::Operation *>(&op));
+/// Cast from an (const) Operation * to a derived operation type.
+template <typename T>
+struct CastInfo<T, ::mlir::Operation *>
+    : public ValueFromPointerCast<T, ::mlir::Operation,
+                                  CastInfo<T, ::mlir::Operation *>> {
+  static bool isPossible(::mlir::Operation *op) {
+    return std::is_same<T, ::mlir::Operation *>::value ? true : T::classof(op);
   }
 };
-
-/// Allow isa<Operation *> on operations.
-template <> struct isa_impl<::mlir::Operation *, ::mlir::Operation> {
-  static inline bool doit(const ::mlir::Operation &op) { return true; }
-};
-
-/// Provide specializations for operation casts as the resulting T is value
-/// typed.
-template <typename T> struct cast_retty_impl<T, ::mlir::Operation *> {
-  using ret_type = T;
-};
-template <typename T> struct cast_retty_impl<T, ::mlir::Operation> {
-  using ret_type = T;
-};
-template <class T>
-struct cast_convert_val<T, ::mlir::Operation, ::mlir::Operation> {
-  static T doit(::mlir::Operation &val) { return T(&val); }
-};
-template <class T>
-struct cast_convert_val<T, ::mlir::Operation *, ::mlir::Operation *> {
-  static T doit(::mlir::Operation *val) { return T(val); }
+template <typename T>
+struct CastInfo<T, const ::mlir::Operation *>
+    : public ConstStrippingForwardingCast<T, const ::mlir::Operation *,
+                                          CastInfo<T, ::mlir::Operation *>> {};
+
+/// Cast from an (const) Operation & to a derived operation type.
+template <typename T>
+struct CastInfo<T, ::mlir::Operation>
+    : public NullableValueCastFailed<T>,
+      public DefaultDoCastIfPossible<T, ::mlir::Operation,
+                                     CastInfo<T, ::mlir::Operation>> {
+  using Self = CastInfo<T, ::mlir::Operation>;
+  // Provide isPossible here because here we have the const-stripping from
+  // ConstStrippingCast.
+  static bool isPossible(::mlir::Operation &val) {
+    return CastInfo<T, ::mlir::Operation *>::isPossible(&val);
+  }
+  static T doCast(::mlir::Operation &val) { return T(&val); }
 };
+template <typename T>
+struct CastInfo<T, const ::mlir::Operation>
+    : public ConstStrippingForwardingCast<T, const ::mlir::Operation,
+                                          CastInfo<T, ::mlir::Operation>> {};
 } // namespace llvm
 
 #endif // MLIR_IR_OPERATION_H