diff --git a/cpp/include/cudf/context.hpp b/cpp/include/cudf/context.hpp index 2363c83b815f..165e200e1293 100644 --- a/cpp/include/cudf/context.hpp +++ b/cpp/include/cudf/context.hpp @@ -19,10 +19,12 @@ enum class init_flags : std::uint32_t { NONE = 0, /// @brief Load the nvCOMP library during initialization LOAD_NVCOMP = 1 << 0, + /// @brief Initialize the JIT runtime and caches during initialization + INITIALIZE_JIT = 1 << 1, /// @brief Default initialization steps DEFAULT = NONE, /// @brief All initialization steps - ALL = LOAD_NVCOMP + ALL = LOAD_NVCOMP | INITIALIZE_JIT }; /// @brief Bitwise OR operator for init_flags diff --git a/cpp/src/runtime/context.cpp b/cpp/src/runtime/context.cpp index cac9723417b4..837b86f88dce 100644 --- a/cpp/src/runtime/context.cpp +++ b/cpp/src/runtime/context.cpp @@ -53,12 +53,8 @@ int32_t get_current_device_compute_capability() context::context(context_config cfg, detail::init_flags flags) : _config{std::move(cfg)}, _device_properties{ - get_driver_version(), get_runtime_version(), get_current_device_compute_capability()}, - _nvrtc_version{0}, - _nvjitlink_version{0} + get_driver_version(), get_runtime_version(), get_current_device_compute_capability()} { - rtcx::initialize(); - initialize_jit(); initialize_components(flags); } @@ -100,16 +96,35 @@ void context::initialize_jit() _jit_bundle = std::make_unique(_config.jit_bundle_dir, *_rtcx_cache); } +void context::ensure_jit_initialized() +{ + std::call_once(_jit_init_flag, [this] { + if (!_rtcx_initialized) { + rtcx::initialize(); + _rtcx_initialized = true; + } + initialize_jit(); + }); +} + context::~context() { _jit_bundle.reset(); _rtcx_cache.reset(); - rtcx::teardown(); + if (_rtcx_initialized) { rtcx::teardown(); } } -rtcx::cache_t& context::rtcx_cache() { return *_rtcx_cache; } +rtcx::cache_t& context::rtcx_cache() +{ + ensure_jit_initialized(); + return *_rtcx_cache; +} -jit_bundle_t& context::jit_bundle() { return *_jit_bundle; } +jit_bundle_t& context::jit_bundle() +{ + ensure_jit_initialized(); + return *_jit_bundle; +} bool context::dump_codegen() const { return _config.dump_codegen; } @@ -132,6 +147,7 @@ void context::initialize_components(detail::init_flags flags) { CUDF_FUNC_RANGE(); if (has_flag(flags, detail::init_flags::LOAD_NVCOMP)) { preload_nvcomp(); } + if (has_flag(flags, detail::init_flags::INITIALIZE_JIT)) { ensure_jit_initialized(); } } namespace { diff --git a/cpp/src/runtime/context.hpp b/cpp/src/runtime/context.hpp index 93f46f8d83ab..6d5f227d0e0f 100644 --- a/cpp/src/runtime/context.hpp +++ b/cpp/src/runtime/context.hpp @@ -9,6 +9,7 @@ #include #include +#include #include namespace rtcx { @@ -53,6 +54,8 @@ class context { private: context_config _config; + std::once_flag _jit_init_flag; + bool _rtcx_initialized = false; std::unique_ptr _rtcx_cache; std::unique_ptr _jit_bundle; device_properties _device_properties; @@ -62,6 +65,8 @@ class context { private: void preload_nvcomp(); + void ensure_jit_initialized(); + void initialize_jit(); void initialize_components(detail::init_flags flags); diff --git a/cpp/tests/utilities_tests/context_tests.cpp b/cpp/tests/utilities_tests/context_tests.cpp index 35d2d9f229e4..0d72fd922264 100644 --- a/cpp/tests/utilities_tests/context_tests.cpp +++ b/cpp/tests/utilities_tests/context_tests.cpp @@ -14,8 +14,28 @@ #include +#include +#include +#include + struct ContextTest : public cudf::test::BaseFixture {}; +namespace { + +cudf::size_type compute_column_jit() +{ + auto c_0 = cudf::test::fixed_width_column_wrapper{3, 20, 1, 50}; + auto c_1 = cudf::test::fixed_width_column_wrapper{10, 7, 20, 0}; + auto table = cudf::table_view{{c_0, c_1}}; + auto col_ref_0 = cudf::ast::column_reference(0); + auto col_ref_1 = cudf::ast::column_reference(1); + auto expression = cudf::ast::operation(cudf::ast::ast_operator::ADD, col_ref_0, col_ref_1); + + return cudf::compute_column_jit(table, expression)->size(); +} + +} // namespace + TEST_F(ContextTest, MultipleInitializeCalls) { cudf::detail::initialize(cudf::detail::init_flags::DEFAULT); @@ -24,25 +44,32 @@ TEST_F(ContextTest, MultipleInitializeCalls) EXPECT_NO_THROW(cudf::detail::initialize(cudf::detail::init_flags::ALL)); } -TEST_F(ContextTest, JitCacheUse) +TEST_F(ContextTest, ConcurrentFirstJitCacheUse) { - auto compute_column = [] { - auto c_0 = cudf::test::fixed_width_column_wrapper{3, 20, 1, 50}; - auto c_1 = cudf::test::fixed_width_column_wrapper{10, 7, 20, 0}; - auto table = cudf::table_view{{c_0, c_1}}; - auto col_ref_0 = cudf::ast::column_reference(0); - auto col_ref_1 = cudf::ast::column_reference(1); - auto expression = cudf::ast::operation(cudf::ast::ast_operator::ADD, col_ref_0, col_ref_1); - - auto result = cudf::compute_column_jit(table, expression); - EXPECT_EQ(result->size(), cudf::size_type{4}); - }; + constexpr auto num_threads = 4; + std::barrier start{num_threads}; + std::vector> results; + results.reserve(num_threads); + + for (auto i = 0; i < num_threads; ++i) { + results.push_back(std::async(std::launch::async, [&] { + start.arrive_and_wait(); + return compute_column_jit(); + })); + } + + for (auto& result : results) { + EXPECT_EQ(result.get(), cudf::size_type{4}); + } +} +TEST_F(ContextTest, JitCacheUse) +{ cudf::detail::initialize(cudf::detail::init_flags::DEFAULT); - ASSERT_NO_THROW(compute_column()); + EXPECT_EQ(compute_column_jit(), cudf::size_type{4}); cudf::detail::initialize(cudf::detail::init_flags::DEFAULT); - ASSERT_NO_THROW(compute_column()); + EXPECT_EQ(compute_column_jit(), cudf::size_type{4}); } template