From b9d32fcc046c253690db7d3008850139ecc8de1b Mon Sep 17 00:00:00 2001
From: Steve Bronder <sbronder@flatironinstitute.org>
Date: Tue, 3 Nov 2026 16:12:46 -0500
Subject: [PATCH] Share OpenCL kernel-cache registration across translation
 units

Use one inline OpenCL context wrapper so cached kernels registered from
different translation units are all invalidated after device selection.
Exercise kernels before and after a device switch through another TU.
---
 stan/math/opencl/opencl_context.hpp           |  3 ++-
 test/unit/multiple_translation_units1.cpp     | 10 ++++++++
 test/unit/multiple_translation_units2.cpp     | 15 ++++++++++++
 test/unit/multiple_translation_units_test.cpp | 23 +++++++++++++++++++
 4 files changed, 50 insertions(+), 1 deletion(-)

diff --git a/stan/math/opencl/opencl_context.hpp b/stan/math/opencl/opencl_context.hpp
index e2373df126d..3a1c95fd0a5 100644
--- a/stan/math/opencl/opencl_context.hpp
+++ b/stan/math/opencl/opencl_context.hpp
@@ -433,7 +433,8 @@ class opencl_context {
     kernel_caches_.push_back(cache);
   }
 };
-static opencl_context opencl_context;
+// Share the context and its kernel-cache registry across translation units.
+inline opencl_context opencl_context;
 }  // namespace math
 }  // namespace stan
 
diff --git a/test/unit/multiple_translation_units1.cpp b/test/unit/multiple_translation_units1.cpp
index 3e9a69ef42d..1d3d6b44309 100644
--- a/test/unit/multiple_translation_units1.cpp
+++ b/test/unit/multiple_translation_units1.cpp
@@ -1,3 +1,13 @@
 #include <test/unit/multiple_translation_units_includes.hpp>
 
 inline stan::math::var function1() { return 0; }
+
+#ifdef STAN_OPENCL
+const void* opencl_context_address1() { return &stan::math::opencl_context; }
+
+double opencl_cached_kernel1() {
+  stan::math::matrix_cl<double> matrix = stan::math::constant(1.0, 2, 2);
+  stan::math::matrix_cl<double> result = matrix * matrix;
+  return stan::math::from_matrix_cl(result).sum();
+}
+#endif
diff --git a/test/unit/multiple_translation_units2.cpp b/test/unit/multiple_translation_units2.cpp
index 8958b90aaa0..2494e080d3d 100644
--- a/test/unit/multiple_translation_units2.cpp
+++ b/test/unit/multiple_translation_units2.cpp
@@ -1,3 +1,18 @@
 #include <test/unit/multiple_translation_units_includes.hpp>
 
 inline stan::math::var function2() { return 0; }
+
+#ifdef STAN_OPENCL
+const void* opencl_context_address2() { return &stan::math::opencl_context; }
+
+double opencl_cached_kernel2() {
+  stan::math::matrix_cl<double> matrix = stan::math::constant(2.0, 2, 2);
+  stan::math::matrix_cl<double> result = matrix + matrix;
+  return stan::math::from_matrix_cl(result).sum();
+}
+
+void select_opencl_device2() {
+  stan::math::opencl_context.select_device(OPENCL_PLATFORM_ID,
+                                           OPENCL_DEVICE_ID);
+}
+#endif
diff --git a/test/unit/multiple_translation_units_test.cpp b/test/unit/multiple_translation_units_test.cpp
index 8894851aba2..c2759bcf4dc 100644
--- a/test/unit/multiple_translation_units_test.cpp
+++ b/test/unit/multiple_translation_units_test.cpp
@@ -4,3 +4,26 @@ TEST(multiple_translation_units, compile) {
   SUCCEED() << "this test compiling indicates that compiling the math library "
             << "public umbrella headers with multiple translation units is ok.";
 }
+
+#ifdef STAN_OPENCL
+const void* opencl_context_address1();
+const void* opencl_context_address2();
+double opencl_cached_kernel1();
+double opencl_cached_kernel2();
+void select_opencl_device2();
+
+TEST(multiple_translation_units, shares_opencl_context) {
+  EXPECT_EQ(opencl_context_address1(), opencl_context_address2());
+}
+
+TEST(multiple_translation_units, opencl_cached_kernels_after_device_switch) {
+  EXPECT_DOUBLE_EQ(8.0, opencl_cached_kernel1());
+  EXPECT_DOUBLE_EQ(16.0, opencl_cached_kernel2());
+
+  // Switching in another translation unit must invalidate both kernel caches.
+  select_opencl_device2();
+
+  EXPECT_DOUBLE_EQ(8.0, opencl_cached_kernel1());
+  EXPECT_DOUBLE_EQ(16.0, opencl_cached_kernel2());
+}
+#endif