From c2e670262c2ae845b4da38f3e775c680612ad9b9 Mon Sep 17 00:00:00 2001 From: Haoyang Li Date: Mon, 7 Sep 2026 02:10:25 +0000 Subject: [PATCH 1/4] PERF: Lazily initialize the JIT runtime context --- cpp/src/runtime/context.cpp | 67 +++++++++++++++++++++---------------- cpp/src/runtime/context.hpp | 3 ++ 2 files changed, 41 insertions(+), 29 deletions(-) diff --git a/cpp/src/runtime/context.cpp b/cpp/src/runtime/context.cpp index cac9723417b4..4f2e21f0e5bb 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); } @@ -71,45 +67,58 @@ void context::preload_nvcomp() void context::initialize_jit() { - CUDF_FUNC_RANGE(); + std::call_once(_jit_init_flag, [&] { + CUDF_FUNC_RANGE(); - // make sure the required directories exist - std::filesystem::create_directories(_config.rtcx_cache_dir); - std::filesystem::create_directories(_config.jit_bundle_dir); - std::filesystem::create_directories(_config.jit_pch_dir); - std::filesystem::create_directories(_config.jit_tmp_dir); + // make sure the required directories exist + std::filesystem::create_directories(_config.rtcx_cache_dir); + std::filesystem::create_directories(_config.jit_bundle_dir); + std::filesystem::create_directories(_config.jit_pch_dir); + std::filesystem::create_directories(_config.jit_tmp_dir); - _nvrtc_version = rtcx::nvrtc_version(); - _nvjitlink_version = rtcx::nvjitlink_version(); + rtcx::initialize(); + _rtcx_initialized = true; - auto limits = rtcx::cache_limits{.num_mem_blobs = _config.kernel_cache_limit_process, - .num_mem_libraries = _config.kernel_cache_limit_process}; + _nvrtc_version = rtcx::nvrtc_version(); + _nvjitlink_version = rtcx::nvjitlink_version(); - _rtcx_cache = std::make_unique(_config.rtcx_cache_dir, - _config.jit_tmp_dir, - limits, - bool{_config.preload_jit_cache}, - bool{_config.disable_jit_cache}); + auto limits = rtcx::cache_limits{.num_mem_blobs = _config.kernel_cache_limit_process, + .num_mem_libraries = _config.kernel_cache_limit_process}; - if (_config.clear_jit_cache) { - _rtcx_cache->clear_memory_store(); - _rtcx_cache->clear_disk_store(); - } + _rtcx_cache = std::make_unique(_config.rtcx_cache_dir, + _config.jit_tmp_dir, + limits, + bool{_config.preload_jit_cache}, + bool{_config.disable_jit_cache}); - // note that jit_bundle depends on rtcx_cache, so we ensure rtcx_cache is initialized first. - _jit_bundle = std::make_unique(_config.jit_bundle_dir, *_rtcx_cache); + if (_config.clear_jit_cache) { + _rtcx_cache->clear_memory_store(); + _rtcx_cache->clear_disk_store(); + } + + // note that jit_bundle depends on rtcx_cache, so we ensure rtcx_cache is initialized first. + _jit_bundle = std::make_unique(_config.jit_bundle_dir, *_rtcx_cache); + }); } 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() +{ + initialize_jit(); + return *_rtcx_cache; +} -jit_bundle_t& context::jit_bundle() { return *_jit_bundle; } +jit_bundle_t& context::jit_bundle() +{ + initialize_jit(); + return *_jit_bundle; +} bool context::dump_codegen() const { return _config.dump_codegen; } diff --git a/cpp/src/runtime/context.hpp b/cpp/src/runtime/context.hpp index 93f46f8d83ab..38d8076faf44 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; From dec7bfa5faece0cee449989ac97a0c0e3ba8a407 Mon Sep 17 00:00:00 2001 From: Haoyang Li Date: Mon, 7 Sep 2026 02:52:40 +0000 Subject: [PATCH 2/4] PERF: Preserve explicit JIT preloading --- cpp/include/cudf/context.hpp | 4 +++- cpp/src/runtime/context.cpp | 1 + 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/cpp/include/cudf/context.hpp b/cpp/include/cudf/context.hpp index 2363c83b815f..1d50010f9e93 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 program cache + 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 4f2e21f0e5bb..0b85da9a16e3 100644 --- a/cpp/src/runtime/context.cpp +++ b/cpp/src/runtime/context.cpp @@ -140,6 +140,7 @@ std::optional context::nvjitlink_version() const { return _nvjitlink_ve void context::initialize_components(detail::init_flags flags) { CUDF_FUNC_RANGE(); + if (has_flag(flags, detail::init_flags::INITIALIZE_JIT)) { initialize_jit(); } if (has_flag(flags, detail::init_flags::LOAD_NVCOMP)) { preload_nvcomp(); } } From b75b44bbfbeddb433c2f036cb4a70d30e06fc6ab Mon Sep 17 00:00:00 2001 From: Haoyang Li Date: Mon, 7 Sep 2026 03:29:00 +0000 Subject: [PATCH 3/4] PERF: Narrow lazy JIT initialization change --- cpp/include/cudf/context.hpp | 4 +- cpp/src/runtime/context.cpp | 59 +++++++++++---------- cpp/src/runtime/context.hpp | 2 + cpp/tests/utilities_tests/context_tests.cpp | 55 ++++++++++++++----- 4 files changed, 76 insertions(+), 44 deletions(-) diff --git a/cpp/include/cudf/context.hpp b/cpp/include/cudf/context.hpp index 1d50010f9e93..2363c83b815f 100644 --- a/cpp/include/cudf/context.hpp +++ b/cpp/include/cudf/context.hpp @@ -19,12 +19,10 @@ 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 program cache - INITIALIZE_JIT = 1 << 1, /// @brief Default initialization steps DEFAULT = NONE, /// @brief All initialization steps - ALL = LOAD_NVCOMP | INITIALIZE_JIT + ALL = LOAD_NVCOMP }; /// @brief Bitwise OR operator for init_flags diff --git a/cpp/src/runtime/context.cpp b/cpp/src/runtime/context.cpp index 0b85da9a16e3..08e77da45737 100644 --- a/cpp/src/runtime/context.cpp +++ b/cpp/src/runtime/context.cpp @@ -67,37 +67,43 @@ void context::preload_nvcomp() void context::initialize_jit() { - std::call_once(_jit_init_flag, [&] { - CUDF_FUNC_RANGE(); + CUDF_FUNC_RANGE(); - // make sure the required directories exist - std::filesystem::create_directories(_config.rtcx_cache_dir); - std::filesystem::create_directories(_config.jit_bundle_dir); - std::filesystem::create_directories(_config.jit_pch_dir); - std::filesystem::create_directories(_config.jit_tmp_dir); + // make sure the required directories exist + std::filesystem::create_directories(_config.rtcx_cache_dir); + std::filesystem::create_directories(_config.jit_bundle_dir); + std::filesystem::create_directories(_config.jit_pch_dir); + std::filesystem::create_directories(_config.jit_tmp_dir); - rtcx::initialize(); - _rtcx_initialized = true; + _nvrtc_version = rtcx::nvrtc_version(); + _nvjitlink_version = rtcx::nvjitlink_version(); - _nvrtc_version = rtcx::nvrtc_version(); - _nvjitlink_version = rtcx::nvjitlink_version(); + auto limits = rtcx::cache_limits{.num_mem_blobs = _config.kernel_cache_limit_process, + .num_mem_libraries = _config.kernel_cache_limit_process}; - auto limits = rtcx::cache_limits{.num_mem_blobs = _config.kernel_cache_limit_process, - .num_mem_libraries = _config.kernel_cache_limit_process}; + _rtcx_cache = std::make_unique(_config.rtcx_cache_dir, + _config.jit_tmp_dir, + limits, + bool{_config.preload_jit_cache}, + bool{_config.disable_jit_cache}); - _rtcx_cache = std::make_unique(_config.rtcx_cache_dir, - _config.jit_tmp_dir, - limits, - bool{_config.preload_jit_cache}, - bool{_config.disable_jit_cache}); + if (_config.clear_jit_cache) { + _rtcx_cache->clear_memory_store(); + _rtcx_cache->clear_disk_store(); + } - if (_config.clear_jit_cache) { - _rtcx_cache->clear_memory_store(); - _rtcx_cache->clear_disk_store(); - } + // note that jit_bundle depends on rtcx_cache, so we ensure rtcx_cache is initialized first. + _jit_bundle = std::make_unique(_config.jit_bundle_dir, *_rtcx_cache); +} - // note that jit_bundle depends on rtcx_cache, so we ensure rtcx_cache is initialized first. - _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(); }); } @@ -110,13 +116,13 @@ context::~context() rtcx::cache_t& context::rtcx_cache() { - initialize_jit(); + ensure_jit_initialized(); return *_rtcx_cache; } jit_bundle_t& context::jit_bundle() { - initialize_jit(); + ensure_jit_initialized(); return *_jit_bundle; } @@ -140,7 +146,6 @@ std::optional context::nvjitlink_version() const { return _nvjitlink_ve void context::initialize_components(detail::init_flags flags) { CUDF_FUNC_RANGE(); - if (has_flag(flags, detail::init_flags::INITIALIZE_JIT)) { initialize_jit(); } if (has_flag(flags, detail::init_flags::LOAD_NVCOMP)) { preload_nvcomp(); } } diff --git a/cpp/src/runtime/context.hpp b/cpp/src/runtime/context.hpp index 38d8076faf44..6d5f227d0e0f 100644 --- a/cpp/src/runtime/context.hpp +++ b/cpp/src/runtime/context.hpp @@ -65,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 From 6227357f5fe6457569e82e096aa054ca5d0ad9de Mon Sep 17 00:00:00 2001 From: Haoyang Li Date: Mon, 7 Sep 2026 03:58:33 +0000 Subject: [PATCH 4/4] PERF: Preserve eager JIT initialization for ALL --- cpp/include/cudf/context.hpp | 4 +++- cpp/src/runtime/context.cpp | 1 + 2 files changed, 4 insertions(+), 1 deletion(-) 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 08e77da45737..837b86f88dce 100644 --- a/cpp/src/runtime/context.cpp +++ b/cpp/src/runtime/context.cpp @@ -147,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 {