use detection idiom inside of getItem instead of virtual member

Author: makslevental

mlir/lib/Bindings/Python/IRCore.cpp

@@ -2218,13 +2218,6 @@ class PyBlockArgumentList
                   step),
         operation(std::move(operation)), block(block) {}
 
-  pybind11::object getItem(intptr_t index) override {
-    auto item = this->SliceableT::getItem(index);
-    if (item.ptr() != nullptr)
-      return item.attr(MLIR_PYTHON_MAYBE_DOWNCAST_ATTR)();
-    return item;
-  }
-
   static void bindDerived(ClassTy &c) {
     c.def_property_readonly("types", [](PyBlockArgumentList &self) {
       return getValueTypes(self, self.operation->getContext());
@@ -2274,13 +2267,6 @@ class PyOpOperandList : public Sliceable<PyOpOperandList, PyValue> {
                   step),
         operation(operation) {}
 
-  pybind11::object getItem(intptr_t index) override {
-    auto item = this->SliceableT::getItem(index);
-    if (item.ptr() != nullptr)
-      return item.attr(MLIR_PYTHON_MAYBE_DOWNCAST_ATTR)();
-    return item;
-  }
-
   void dunderSetItem(intptr_t index, PyValue value) {
     index = wrapIndex(index);
     mlirOperationSetOperand(operation->get(), index, value.get());
@@ -2337,13 +2323,6 @@ class PyOpResultList : public Sliceable<PyOpResultList, PyOpResult> {
                   step),
         operation(std::move(operation)) {}
 
-  pybind11::object getItem(intptr_t index) override {
-    auto item = this->SliceableT::getItem(index);
-    if (item.ptr() != nullptr)
-      return item.attr(MLIR_PYTHON_MAYBE_DOWNCAST_ATTR)();
-    return item;
-  }
-
   static void bindDerived(ClassTy &c) {
     c.def_property_readonly("types", [](PyOpResultList &self) {
       return getValueTypes(self, self.operation->getContext());

mlir/lib/Bindings/Python/IRModule.h

@@ -1110,6 +1110,7 @@ class PyConcreteAttribute : public BaseTy {
 /// bindings so such operation always exists).
 class PyValue {
 public:
+  virtual ~PyValue() = default;
   PyValue(PyOperationRef parentOperation, MlirValue value)
       : parentOperation(std::move(parentOperation)), value(value) {}
   operator MlirValue() const { return value; }
@@ -1122,7 +1123,7 @@ class PyValue {
   /// Gets a capsule wrapping the void* within the MlirValue.
   pybind11::object getCapsule();
 
-  virtual pybind11::object maybeDownCast();
+  pybind11::object maybeDownCast();
 
   /// Creates a PyValue from the MlirValue wrapped by a capsule. Ownership of
   /// the underlying MlirValue is still tied to the owning operation.

mlir/lib/Bindings/Python/PybindUtils.h

@@ -10,6 +10,7 @@
 #define MLIR_BINDINGS_PYTHON_PYBINDUTILS_H
 
 #include "mlir-c/Support.h"
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/Twine.h"
 #include "llvm/Support/DataTypes.h"
 
@@ -228,19 +229,29 @@ class Sliceable {
     return linearIndex;
   }
 
+  /// Trait to check if T provides a `maybeDownCast` method.
+  /// Note, you need the & to detect inherited members.
+  template <typename T, typename... Args>
+  using has_maybe_downcast = decltype(&T::maybeDownCast);
+
   /// Returns the element at the given slice index. Supports negative indices
   /// by taking elements in inverse order. Returns a nullptr object if out
   /// of bounds.
-  virtual pybind11::object getItem(intptr_t index) {
+  pybind11::object getItem(intptr_t index) {
     // Negative indices mean we count from the end.
     index = wrapIndex(index);
     if (index < 0) {
       PyErr_SetString(PyExc_IndexError, "index out of range");
       return {};
     }
 
-    return pybind11::cast(
-        static_cast<Derived *>(this)->getRawElement(linearizeIndex(index)));
+    if constexpr (llvm::is_detected<has_maybe_downcast, ElementTy>::value)
+      return static_cast<Derived *>(this)
+          ->getRawElement(linearizeIndex(index))
+          .maybeDownCast();
+    else
+      return pybind11::cast(
+          static_cast<Derived *>(this)->getRawElement(linearizeIndex(index)));
   }
 
   /// Returns a new instance of the pseudo-container restricted to the given

mlir/test/python/dialects/arith_dialect.py

@@ -74,6 +74,9 @@ def __str__(self):
 
         with InsertionPoint(module.body):
             a = arith.constant(value=FloatAttr.get(f16_t, 42.42))
+            # CHECK: ArithValue(%cst = arith.constant 4.240
+            print(a)
+
             b = a + a
             # CHECK: ArithValue(%0 = arith.addf %cst, %cst : f16)
             print(b)

mlir/test/python/ir/value.py

@@ -370,7 +370,8 @@ def dont_cast_int_shouldnt_register(v):
         print(e)
 
     @register_value_caster(IntegerType.static_typeid, replace=True)
-    def dont_cast_int(v) -> Value:
+    def dont_cast_int(v) -> OpResult:
+        assert isinstance(v, OpResult)
         print("don't cast", v.result_number, v)
         return v