PERF: Lazily initialize the JIT runtime context

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

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<jit_bundle_t>(_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 {

cpp/src/runtime/context.hpp

@@ -9,6 +9,7 @@
 #include <cudf/utilities/export.hpp>
 
 #include <memory>
+#include <mutex>
 #include <optional>
 
 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_t> _rtcx_cache;
   std::unique_ptr<jit_bundle_t> _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);

cpp/tests/utilities_tests/context_tests.cpp

@@ -14,8 +14,28 @@
 
 #include <gtest/gtest.h>
 
+#include <barrier>
+#include <future>
+#include <vector>
+
 struct ContextTest : public cudf::test::BaseFixture {};
 
+namespace {
+
+cudf::size_type compute_column_jit()
+{
+  auto c_0        = cudf::test::fixed_width_column_wrapper<cudf::size_type>{3, 20, 1, 50};
+  auto c_1        = cudf::test::fixed_width_column_wrapper<cudf::size_type>{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<cudf::size_type>{3, 20, 1, 50};
-    auto c_1        = cudf::test::fixed_width_column_wrapper<cudf::size_type>{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<std::future<cudf::size_type>> 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 <typename Lambda>