Make equality and ordering consistent

Co-authored-by: Cursor <cursoragent@cursor.com>

Author: andralex

cudax/include/cuda/experimental/__stf/places/data_place_extension.cuh

@@ -58,7 +58,7 @@ class exec_place;
  *   int get_device_ordinal() const override { return my_device_id; }
  *   ::std::string to_string() const override { return "my_custom_place"; }
  *   size_t hash() const override { return std::hash<int>{}(my_device_id); }
- *   bool equals(const data_place_extension& other) const override { ... }
+ *   bool operator==(const data_place_extension& other) const override { ... }
  * };
  * @endcode
  */
@@ -104,15 +104,15 @@ public:
    * @param other The other extension to compare with
    * @return true if the extensions represent the same place
    */
-  virtual bool equals(const data_place_extension& other) const = 0;
+  virtual bool operator==(const data_place_extension& other) const = 0;
 
   /**
    * @brief Compare ordering with another extension
    *
    * @param other The other extension to compare with
    * @return true if this extension is less than the other
    */
-  virtual bool less_than(const data_place_extension& other) const = 0;
+  virtual bool operator<(const data_place_extension& other) const = 0;
 
   /**
    * @brief Create a physical memory allocation for this place (VMM API)

cudax/include/cuda/experimental/__stf/places/exec/cuda_stream.cuh

@@ -71,13 +71,13 @@ public:
 
     bool operator==(const exec_place::impl& rhs) const override
     {
-      auto other = dynamic_cast<const impl*>(&rhs);
-      if (!other)
+      if (typeid(*this) != typeid(rhs))
       {
         return false;
       }
+      const auto& other = static_cast<const impl&>(rhs);
       // Compare by stream handle
-      return dstream.stream == other->dstream.stream;
+      return dstream.stream == other.dstream.stream;
     }
 
     size_t hash() const override

cudax/include/cuda/experimental/__stf/places/exec/green_context.cuh

@@ -91,17 +91,17 @@ public:
       return hash_all(view_.g_ctx, view_.pool, view_.devid);
     }
 
-    bool equals(const data_place_extension& other) const override
+    bool operator==(const data_place_extension& other) const override
     {
-      const auto* other_gc = dynamic_cast<const extension*>(&other);
-      if (!other_gc)
+      if (typeid(*this) != typeid(other))
       {
         return false;
       }
-      return view_ == other_gc->view_;
+      const auto& other_gc = static_cast<const extension&>(other);
+      return view_ == other_gc.view_;
     }
 
-    bool less_than(const data_place_extension& other) const override
+    bool operator<(const data_place_extension& other) const override
     {
       if (typeid(*this) != typeid(other))
       {
@@ -424,14 +424,13 @@ public:
 
     bool operator==(const exec_place::impl& rhs) const override
     {
-      // First, check if rhs is also a green context impl
-      auto other = dynamic_cast<const impl*>(&rhs);
-      if (!other)
+      if (typeid(*this) != typeid(rhs))
       {
         return false;
       }
+      const auto& other = static_cast<const impl&>(rhs);
       // Compare green context handles
-      return g_ctx == other->g_ctx;
+      return g_ctx == other.g_ctx;
     }
 
     size_t hash() const override

cudax/include/cuda/experimental/__stf/places/places.cuh

@@ -174,7 +174,7 @@ public:
     // If both are extensions, delegate to the extension
     if (is_extension() && rhs.is_extension())
     {
-      return extension->less_than(*rhs.extension);
+      return *extension < *rhs.extension;
     }
 
     // Extensions sort after non-extensions
@@ -1881,7 +1881,7 @@ inline bool data_place::operator==(const data_place& rhs) const
   if (is_extension())
   {
     _CCCL_ASSERT(devid == extension_devid, "");
-    return (rhs.devid == extension_devid && extension->equals(*rhs.extension));
+    return (rhs.devid == extension_devid && *extension == *rhs.extension);
   }
 
   return (get_grid() == rhs.get_grid() && (get_partitioner() == rhs.get_partitioner()));