From 4053f2b96988b69d68c8cc12913f51c149f1ac46 Mon Sep 17 00:00:00 2001
From: Matsu <47756807+khaki3@users.noreply.github.com>
Date: Mon, 27 Jul 2026 16:02:30 -0700
Subject: [PATCH] [flang-rt][cuda] Keep CUFDeviceIsActive working under
 runtime/driver skew (#212320)

Resolve cuDevicePrimaryCtxGetState via the installed driver version, and
prefer running scope-exit cleanup when the probe cannot prove teardown,
so local device allocations are not leaked when the unversioned lookup
fails.
---
 flang-rt/CMakeLists.txt                       |  6 ++++
 flang-rt/README.md                            |  2 +-
 flang-rt/lib/cuda/allocator.cpp               | 36 ++++++++++++-------
 .../unittests/Runtime/CUDA/AllocatorCUF.cpp   |  8 ++---
 4 files changed, 32 insertions(+), 20 deletions(-)

diff --git a/flang-rt/CMakeLists.txt b/flang-rt/CMakeLists.txt
index 5e62ac669e63..4f480151eb56 100644
--- a/flang-rt/CMakeLists.txt
+++ b/flang-rt/CMakeLists.txt
@@ -163,6 +163,12 @@ endif ()
 option(FLANG_RT_INCLUDE_CUF "Build the CUDA Fortran runtime (libflang_rt.cuda.a)" OFF)
 if (FLANG_RT_INCLUDE_CUF)
   find_package(CUDAToolkit REQUIRED)
+  # CUF runtime uses cudaGetDriverEntryPointByVersion (CUDA 12.5+) and assumes
+  # the toolchain floor used elsewhere for CUDA Fortran (12.8+).
+  if (CUDAToolkit_VERSION VERSION_LESS "12.8")
+    message(FATAL_ERROR
+      "FLANG_RT_INCLUDE_CUF requires CUDA Toolkit 12.8 or later (found ${CUDAToolkit_VERSION})")
+  endif ()
 endif()
 
 # A recently added check in clang emits warnings for feclearexcept and fesetround:
diff --git a/flang-rt/README.md b/flang-rt/README.md
index 68bc9c9f6057..7d8e8da71062 100644
--- a/flang-rt/README.md
+++ b/flang-rt/README.md
@@ -157,7 +157,7 @@ CMake itself provide.
    independent of `FLANG_RT_EXPERIMENTAL_OFFLOAD_SUPPORT=CUDA` and only
    requires a
    [CUDA Toolkit installation](https://cmake.org/cmake/help/latest/module/FindCUDAToolkit.html)
-   (no `CMAKE_CUDA_COMPILER`).
+   12.8 or later (no `CMAKE_CUDA_COMPILER`).
 
 
 ### Experimental CUDA Support
diff --git a/flang-rt/lib/cuda/allocator.cpp b/flang-rt/lib/cuda/allocator.cpp
index 5e3ee590afb5..4161caa7cb02 100644
--- a/flang-rt/lib/cuda/allocator.cpp
+++ b/flang-rt/lib/cuda/allocator.cpp
@@ -25,26 +25,34 @@
 namespace Fortran::runtime::cuda {
 
 static bool deviceContextTornDown() {
-  // Must be transparent to the last-error seen by user code: the one-time
-  // cudaGetDriverEntryPoint() lookup below can fail (e.g. runtime newer than
-  // driver) and leave a sticky error that a later user cudaGetLastError() would
-  // misattribute. Consume an error only if the slot started clean, so we never
-  // discard a pre-existing one.
+  // Keep cudaGetLastError transparent: consume probe-only sticky errors when
+  // the slot started clean, never discarding a pre-existing user error.
   cudaError_t priorErr{cudaPeekAtLastError()};
-  bool tornDown{true};
+  // Prefer cleanup when state cannot be proven torn down (avoids leaks).
+  bool tornDown{false};
   int device{0};
   if (cudaGetDevice(&device) == cudaSuccess) {
-    // Driver API reports primary-context state WITHOUT lazily creating one
-    // (unlike the runtime API); resolve it via cudart to avoid a libcuda link.
-    // Only the current device is checked (single-device assumption).
+    // Driver API reports primary-context state without lazily creating one;
+    // resolve via cudart to avoid a libcuda link (current device only).
     using GetStateFn = CUresult(CUDAAPI *)(CUdevice, unsigned *, int *);
     static GetStateFn getState{[]() -> GetStateFn {
       void *fn{nullptr};
-      if (cudaGetDriverEntryPoint("cuDevicePrimaryCtxGetState", &fn,
-              cudaEnableDefault, nullptr) != cudaSuccess) {
-        return nullptr;
+      // Prefer ByVersion(driver): unversioned lookup uses the runtime version
+      // and fails when the runtime is newer than the driver.
+      int driverVersion{0};
+      if (cudaDriverGetVersion(&driverVersion) == cudaSuccess &&
+          cudaGetDriverEntryPointByVersion("cuDevicePrimaryCtxGetState", &fn,
+              static_cast<unsigned>(driverVersion), cudaEnableDefault,
+              nullptr) == cudaSuccess &&
+          fn) {
+        return reinterpret_cast<GetStateFn>(fn);
       }
-      return reinterpret_cast<GetStateFn>(fn);
+      if (cudaGetDriverEntryPoint("cuDevicePrimaryCtxGetState", &fn,
+              cudaEnableDefault, nullptr) == cudaSuccess &&
+          fn) {
+        return reinterpret_cast<GetStateFn>(fn);
+      }
+      return nullptr;
     }()};
     if (getState) {
       unsigned flags{0};
@@ -53,6 +61,8 @@ static bool deviceContextTornDown() {
         tornDown = active == 0;
       }
     }
+  } else {
+    tornDown = true;
   }
   if (priorErr == cudaSuccess && cudaPeekAtLastError() != cudaSuccess) {
     (void)cudaGetLastError();
diff --git a/flang-rt/unittests/Runtime/CUDA/AllocatorCUF.cpp b/flang-rt/unittests/Runtime/CUDA/AllocatorCUF.cpp
index 580afe430ca0..cf0de6f25b4d 100644
--- a/flang-rt/unittests/Runtime/CUDA/AllocatorCUF.cpp
+++ b/flang-rt/unittests/Runtime/CUDA/AllocatorCUF.cpp
@@ -74,12 +74,8 @@ TEST(AllocatableCUFTest, DescriptorAllocationTest) {
 }
 
 TEST(AllocatableCUFTest, DeviceIsActiveKeepsLastErrorClean) {
-  // CUFDeviceIsActive() is emitted around scope-exit device cleanup. Its
-  // internal context probing (e.g. the one-time driver-entry-point lookup)
-  // must not leave a sticky error behind: user code performing its own
-  // cudaGetLastError() after a subsequent kernel launch would otherwise
-  // misattribute that stale error to the launch. The boolean result depends
-  // on whether a primary context is active in this process and is not checked.
+  // CUFDeviceIsActive() probes primary-context state (including a version-
+  // skew fallback). It must not leave a sticky cudaGetLastError behind.
   (void)cudaGetLastError(); // start from a clean error state
   (void)RTNAME(CUFDeviceIsActive)();
   EXPECT_EQ(cudaGetLastError(), cudaSuccess);