diff --git a/src/cloneable.h b/src/cloneable.h new file mode 100644 index 000000000..71e75ee36 --- /dev/null +++ b/src/cloneable.h @@ -0,0 +1,21 @@ +#ifndef INFINI_OPS_CLONEABLE_H_ +#define INFINI_OPS_CLONEABLE_H_ + +#include +#include + +namespace infini::ops { + +template +class Cloneable : public Base { + public: + std::unique_ptr Clone() const override { + static_assert(std::is_final_v, + "Cloneable requires a final derived class."); + return std::make_unique(static_cast(*this)); + } +}; + +} // namespace infini::ops + +#endif diff --git a/src/config.h b/src/config.h index a8b59a4fd..e156497bd 100644 --- a/src/config.h +++ b/src/config.h @@ -2,11 +2,20 @@ #define INFINI_OPS_CONFIG_H_ #include +#include + +#include "cloneable.h" namespace infini::ops { class Config { public: + virtual ~Config() = default; + + virtual std::unique_ptr Clone() const { + return std::make_unique(*this); + } + std::size_t implementation_index() const { return implementation_index_; } void set_implementation_index(std::size_t implementation_index) { diff --git a/src/handle.h b/src/handle.h index 4deeb83c9..27211f0ec 100644 --- a/src/handle.h +++ b/src/handle.h @@ -2,11 +2,20 @@ #define INFINI_OPS_HANDLE_H_ #include +#include + +#include "cloneable.h" namespace infini::ops { class Handle { public: + virtual ~Handle() = default; + + virtual std::unique_ptr Clone() const { + return std::make_unique(*this); + } + void* stream() const { return stream_; } void* workspace() const { return workspace_; } diff --git a/src/operator.h b/src/operator.h index 1c60596a4..73ee1f39d 100644 --- a/src/operator.h +++ b/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_; - Config config_; + std::unique_ptr config_; void* stream_{nullptr}; diff --git a/tests/test_cpp_api.py b/tests/test_cpp_api.py index 52a97542f..221f8a851 100644 --- a/tests/test_cpp_api.py +++ b/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 } """ ).lstrip() + + +_POLYMORPHIC_CONTEXT_SOURCE = textwrap.dedent( + r""" + #include + + #include + #include + + namespace infini::ops { + + class DerivedConfig final : public Cloneable { + public: + explicit DerivedConfig(int value) : value_{value} {} + + int value() const { return value_; } + + private: + int value_; + }; + + class DerivedHandle final : public Cloneable { + 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(*config_).value(); + } + + std::size_t implementation_index() const { + return config_->implementation_index(); + } + + int handle_value() const { + return static_cast(*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); + static_assert(std::has_virtual_destructor_v); + + 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()