feat!: support polymorphic handles and configs

src/cloneable.h

@@ -0,0 +1,21 @@
+#ifndef INFINI_OPS_CLONEABLE_H_
+#define INFINI_OPS_CLONEABLE_H_
+
+#include <memory>
+#include <type_traits>
+
+namespace infini::ops {
+
+template <typename Base, typename Derived>
+class Cloneable : public Base {
+ public:
+  std::unique_ptr<Base> Clone() const override {
+    static_assert(std::is_final_v<Derived>,
+                  "Cloneable requires a final derived class.");
+    return std::make_unique<Derived>(static_cast<const Derived&>(*this));
+  }
+};
+
+}  // namespace infini::ops
+
+#endif

src/config.h

@@ -2,11 +2,20 @@
 #define INFINI_OPS_CONFIG_H_
 
 #include <cstddef>
+#include <memory>
+
+#include "cloneable.h"
 
 namespace infini::ops {
 
 class Config {
  public:
+  virtual ~Config() = default;
+
+  virtual std::unique_ptr<Config> Clone() const {
+    return std::make_unique<Config>(*this);
+  }
+
   std::size_t implementation_index() const { return implementation_index_; }
 
   void set_implementation_index(std::size_t implementation_index) {

src/handle.h

@@ -2,11 +2,20 @@
 #define INFINI_OPS_HANDLE_H_
 
 #include <cstddef>
+#include <memory>
+
+#include "cloneable.h"
 
 namespace infini::ops {
 
 class Handle {
  public:
+  virtual ~Handle() = default;
+
+  virtual std::unique_ptr<Handle> Clone() const {
+    return std::make_unique<Handle>(*this);
+  }
+
   void* stream() const { return stream_; }
 
   void* workspace() const { return workspace_; }

src/operator.h

@@ -163,9 +163,9 @@ class OperatorBase {
 
   virtual std::size_t workspace_size_in_bytes() const { return 0; }
 
-  void set_handle(const Handle& handle) { handle_ = handle; }
+  void set_handle(const Handle& handle) { handle_ = handle.Clone(); }
 
-  void set_config(const Config& config) { config_ = config; }
+  void set_config(const Config& config) { config_ = config.Clone(); }
 
   void set_stream(void* stream) { stream_ = stream; }
 
@@ -176,9 +176,9 @@ class OperatorBase {
   }
 
  protected:
-  Handle handle_;
+  std::unique_ptr<Handle> handle_;
 
-  Config config_;
+  std::unique_ptr<Config> config_;
 
   void* stream_{nullptr};
 

tests/test_cpp_api.py

@@ -84,6 +84,32 @@ def test_cpp_configless_calls_use_first_active_implementation(tmp_path):
     _run([str(binary)])
 
 
+def test_cpp_polymorphic_context_smoke(tmp_path):
+    install_prefix = _install_prefix()
+    include_dir = install_prefix / "include"
+    library_dir = _library_dir(install_prefix)
+    source = tmp_path / "polymorphic_context.cc"
+    binary = tmp_path / "polymorphic_context"
+    source.write_text(_POLYMORPHIC_CONTEXT_SOURCE)
+
+    _run(
+        [
+            _compiler("CXX", "c++"),
+            "-std=c++17",
+            "-Werror",
+            f"-I{include_dir}",
+            str(source),
+            f"-L{library_dir}",
+            "-linfiniops",
+            "-linfinirt",
+            f"-Wl,-rpath,{library_dir}",
+            "-o",
+            str(binary),
+        ]
+    )
+    _run([str(binary)])
+
+
 @pytest.mark.parametrize(
     "header",
     (
@@ -474,3 +500,86 @@ class Operator<ConfiglessSelection, Device::Type::kCpu, 16>
     }
     """
 ).lstrip()
+
+
+_POLYMORPHIC_CONTEXT_SOURCE = textwrap.dedent(
+    r"""
+    #include <operator.h>
+
+    #include <cstddef>
+    #include <type_traits>
+
+    namespace infini::ops {
+
+    class DerivedConfig final : public Cloneable<Config, DerivedConfig> {
+     public:
+      explicit DerivedConfig(int value) : value_{value} {}
+
+      int value() const { return value_; }
+
+     private:
+      int value_;
+    };
+
+    class DerivedHandle final : public Cloneable<Handle, DerivedHandle> {
+     public:
+      explicit DerivedHandle(int value) : value_{value} {}
+
+      int value() const { return value_; }
+
+     private:
+      int value_;
+    };
+
+    class PolymorphicOwner final : public OperatorBase {
+     public:
+      int config_value() const {
+        return static_cast<const DerivedConfig&>(*config_).value();
+      }
+
+      std::size_t implementation_index() const {
+        return config_->implementation_index();
+      }
+
+      int handle_value() const {
+        return static_cast<const DerivedHandle&>(*handle_).value();
+      }
+
+      void* handle_stream() const { return handle_->stream(); }
+    };
+
+    }  // namespace infini::ops
+
+    int main() {
+      using namespace infini::ops;
+
+      static_assert(std::has_virtual_destructor_v<Config>);
+      static_assert(std::has_virtual_destructor_v<Handle>);
+
+      PolymorphicOwner owner;
+
+      {
+        DerivedConfig config{17};
+        config.set_implementation_index(3);
+        owner.set_config(config);
+        config.set_implementation_index(9);
+      }
+
+      if (owner.config_value() != 17) return 1;
+      if (owner.implementation_index() != 3) return 2;
+
+      int stream;
+      {
+        DerivedHandle handle{23};
+        handle.set_stream(&stream);
+        owner.set_handle(handle);
+        handle.set_stream(nullptr);
+      }
+
+      if (owner.handle_value() != 23) return 3;
+      if (owner.handle_stream() != &stream) return 4;
+
+      return 0;
+    }
+    """
+).lstrip()