diff --git a/scripts/generate_wrappers.py b/scripts/generate_wrappers.py index 88d1ba186..c4c0ee1a8 100644 --- a/scripts/generate_wrappers.py +++ b/scripts/generate_wrappers.py @@ -1534,6 +1534,13 @@ def _append_unique(declaration, definition): declarations.append(declaration) definitions.append(definition) + _append_unique( + f"extern template std::size_t " + f"Operator<{op_type}>::DefaultImplementationIndex(Device::Type);", + f"template std::size_t " + f"Operator<{op_type}>::DefaultImplementationIndex(Device::Type);", + ) + for call in operator.calls: template_arguments = _generate_template_arguments(call) params = _generate_parameters(call) diff --git a/src/operator.h b/src/operator.h index c3c98e00f..1c60596a4 100644 --- a/src/operator.h +++ b/src/operator.h @@ -3,6 +3,7 @@ #include #include +#include #include #include #include @@ -206,7 +207,8 @@ class Operator : public OperatorBase { template static std::unique_ptr Make(const Tensor tensor, Args&&... args) { - return Make({}, tensor, std::forward(args)...); + return Make(DefaultConfig(tensor.device().type()), tensor, + std::as_const(args)...); } template @@ -222,7 +224,10 @@ class Operator : public OperatorBase { template static std::unique_ptr Make(const std::vector tensors, Args&&... args) { - return Make({}, tensors, std::forward(args)...); + assert(!tensors.empty() && "operator tensor list input cannot be empty"); + + return Make(DefaultConfig(tensors.front().device().type()), tensors, + std::as_const(args)...); } template @@ -279,7 +284,7 @@ class Operator : public OperatorBase { template static void Call(const Tensor tensor, const Args&... args) { - return Call({}, {}, tensor, args...); + return Call({}, DefaultConfig(tensor.device().type()), tensor, args...); } template < @@ -329,6 +334,39 @@ class Operator : public OperatorBase { static constexpr std::size_t implementation_index_{implementation_index}; private: + template + static constexpr std::size_t FirstActiveImplementationIndex( + List) { + return static_cast(first); + } + + static std::size_t FirstActiveImplementationIndex(List<>) { + assert(false && "operator has no active implementation for this device"); + std::abort(); + } + + static std::size_t DefaultImplementationIndex(Device::Type dev_type) { + std::size_t default_index{0}; + + DispatchFunc>( + dev_type, + [&](auto device_tag) { + constexpr Device::Type kDev = decltype(device_tag)::value; + default_index = FirstActiveImplementationIndex( + typename ActiveImplementations::type{}); + }, + "Operator::DefaultImplementationIndex"); + + return default_index; + } + + static Config DefaultConfig(Device::Type dev_type) { + Config config; + config.set_implementation_index(DefaultImplementationIndex(dev_type)); + + return config; + } + template static auto CallReturning(const TensorLike& tensor, const Args&... args) { auto out = Key::MakeReturnValue(tensor, args...); diff --git a/tests/test_cpp_api.py b/tests/test_cpp_api.py index 1cc60fe34..52a97542f 100644 --- a/tests/test_cpp_api.py +++ b/tests/test_cpp_api.py @@ -58,6 +58,32 @@ def test_cpp_returning_call_smoke(tmp_path): _run([str(binary)]) +def test_cpp_configless_calls_use_first_active_implementation(tmp_path): + install_prefix = _install_prefix() + include_dir = install_prefix / "include" + library_dir = _library_dir(install_prefix) + source = tmp_path / "configless_active_implementation.cc" + binary = tmp_path / "configless_active_implementation" + source.write_text(_CONFIGLESS_ACTIVE_IMPLEMENTATION_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", ( @@ -332,3 +358,119 @@ class OwningTensor { } """ ).lstrip() + + +_CONFIGLESS_ACTIVE_IMPLEMENTATION_SOURCE = textwrap.dedent( + r""" + #include + + #include + #include + #include + + namespace infini::ops { + + class ConfiglessSelection : public Operator { + public: + ConfiglessSelection(const Tensor input, Tensor out) {} + + ConfiglessSelection(const std::vector inputs, Tensor out) {} + + virtual void operator()(const Tensor input, Tensor out) const = 0; + + virtual void operator()(const std::vector inputs, + Tensor out) const = 0; + + template + static auto MakeReturnValue(const TensorLike& input) { + return TensorLike::Empty(input.shape(), input.dtype(), input.device()); + } + }; + + template <> + class Operator + : public ConfiglessSelection { + public: + using ConfiglessSelection::ConfiglessSelection; + + void operator()(const Tensor input, Tensor out) const override { + const auto* input_data = static_cast(input.data()); + auto* out_data = static_cast(out.data()); + out_data[0] = input_data[0] + 16.0f; + } + + void operator()(const std::vector inputs, + Tensor out) const override { + const auto* first_data = static_cast(inputs[0].data()); + const auto* second_data = static_cast(inputs[1].data()); + auto* out_data = static_cast(out.data()); + out_data[0] = first_data[0] + second_data[0] + 16.0f; + } + }; + + } // namespace infini::ops + + int main() { + float input_data = 1.0f; + float second_data = 2.0f; + float out_data = 0.0f; + const infini::ops::Tensor::Shape shape{1}; + const infini::ops::Device device{infini::ops::Device::Type::kCpu}; + const infini::ops::DataType dtype{infini::ops::DataType::kFloat32}; + infini::ops::Tensor input(&input_data, shape, dtype, device); + infini::ops::Tensor second(&second_data, shape, dtype, device); + infini::ops::Tensor out(&out_data, shape, dtype, device); + + auto tensor_op = infini::ops::ConfiglessSelection::Make(input, out); + (*tensor_op)(input, out); + if (std::fabs(out_data - 17.0f) > 1e-6f) { + return 1; + } + + out_data = 0.0f; + std::vector inputs{input, second}; + auto vector_op = infini::ops::ConfiglessSelection::Make(inputs, out); + (*vector_op)(inputs, out); + if (std::fabs(out_data - 19.0f) > 1e-6f) { + return 2; + } + + float cat_out_data[2] = {}; + const infini::ops::Tensor::Shape cat_shape{2}; + infini::ops::Tensor cat_out(cat_out_data, cat_shape, dtype, device); + auto cat_op = + infini::ops::Cat::Make(inputs, std::int64_t{0}, cat_out); + (*cat_op)(inputs, std::int64_t{0}, cat_out); + if (std::fabs(cat_out_data[0] - 1.0f) > 1e-6f || + std::fabs(cat_out_data[1] - 2.0f) > 1e-6f) { + return 3; + } + + float abs_input_data = -4.0f; + float abs_out_data = 0.0f; + infini::ops::Tensor abs_input(&abs_input_data, shape, dtype, device); + infini::ops::Tensor abs_out(&abs_out_data, shape, dtype, device); + auto abs_op = infini::ops::Abs::Make(abs_input, abs_out); + (*abs_op)(abs_input, abs_out); + if (std::fabs(abs_out_data - 4.0f) > 1e-6f) { + return 4; + } + + auto abs_op_from_rvalue = infini::ops::Abs::Make( + abs_input, + infini::ops::Tensor(&abs_out_data, shape, dtype, device)); + (*abs_op_from_rvalue)(abs_input, abs_out); + if (std::fabs(abs_out_data - 4.0f) > 1e-6f) { + return 5; + } + + abs_out_data = 0.0f; + infini::ops::Abs::Call(abs_input, abs_out); + if (std::fabs(abs_out_data - 4.0f) > 1e-6f) { + return 6; + } + + return 0; + } + """ +).lstrip() diff --git a/tests/test_generate_wrappers.py b/tests/test_generate_wrappers.py index f347d2964..998109de6 100644 --- a/tests/test_generate_wrappers.py +++ b/tests/test_generate_wrappers.py @@ -78,6 +78,33 @@ class Clamp { ) in text +def test_operator_call_instantiations_externalize_default_implementation_lookup(): + module = _load_generator_module() + operator = module._Operator( + "abs", + constructors=[], + calls=[ + module._ParsedFunction( + [ + module._ParsedArgument("const Tensor", "input"), + module._ParsedArgument("Tensor", "out"), + ] + ) + ], + ) + + declarations, definitions = module._generate_operator_call_instantiation_entries( + operator + ) + + signature = ( + "std::size_t " + "Operator<::infini::ops::Abs>::DefaultImplementationIndex(Device::Type);" + ) + assert f"extern template {signature}" in declarations + assert f"template {signature}" in definitions + + def test_operator_call_instantiations_keep_scalar_and_optional_tensor_overloads_distinct( monkeypatch, tmp_path ):