[flang-commits] [flang] [llvm] [flang-rt][cuda] Skip scope-exit cleanup when the context has a sticky error (PR #213184)

[Zhen Wang via flang-commits]

https://github.com/wangzpgi created https://github.com/llvm/llvm-project/pull/213184

A sticky CUDA error (e.g. an illegal memory access in a kernel) leaves the primary context active but unusable, so CUFDeviceIsActive() reports it as fine and the compiler-generated scope-exit frees abort a program that ran to completion: 'cudaFree(p)' failed with 'cudaErrorIllegalAddress'.

Detect this by freeing a null pointer, a no-op that still reports the sticky error. It runs only once the primary context is known active, so it cannot lazily create one.

>From 05619745e20bb0828804cd624d2dea119673fca1 Mon Sep 17 00:00:00 2001
From: Zhen Wang <zhenw at nvidia.com>
Date: Thu, 30 Jul 2026 16:28:57 -0700
Subject: [PATCH] Treat a sticky-error context as torn down in
 CUFDeviceIsActive

---
 flang-rt/lib/cuda/allocator.cpp         | 7 +++++++
 flang/test/Lower/CUDA/cuda-return01.cuf | 6 ++++++
 2 files changed, 13 insertions(+)

diff --git a/flang-rt/lib/cuda/allocator.cpp b/flang-rt/lib/cuda/allocator.cpp
index 4161caa7cb022..d8287d6f696e3 100644
--- a/flang-rt/lib/cuda/allocator.cpp
+++ b/flang-rt/lib/cuda/allocator.cpp
@@ -59,6 +59,13 @@ static bool deviceContextTornDown() {
       int active{0};
       if (getState(device, &flags, &active) == CUDA_SUCCESS) {
         tornDown = active == 0;
+        // A sticky error (e.g. an illegal kernel memory access) leaves the
+        // primary context active but unusable: later calls all fail, so
+        // scope-exit frees would abort an otherwise successful program. A
+        // null free is a no-op that surfaces this without creating a context.
+        if (!tornDown && cudaFree(nullptr) != cudaSuccess) {
+          tornDown = true;
+        }
       }
     }
   } else {
diff --git a/flang/test/Lower/CUDA/cuda-return01.cuf b/flang/test/Lower/CUDA/cuda-return01.cuf
index 6e5d02c3b3962..9ea5ce1538081 100644
--- a/flang/test/Lower/CUDA/cuda-return01.cuf
+++ b/flang/test/Lower/CUDA/cuda-return01.cuf
@@ -31,13 +31,19 @@ end
 
 program main
   integer, allocatable, device :: a(:)
+  integer, device, pointer :: p(:)
   return
 end
 
+! The descriptor of a device pointer lives in managed memory and is freed at
+! scope exit, under the same device-active guard as the other cleanup.
+
 ! CHECK-LABEL: func.func @_QQmain() attributes {fir.bindc_name = "MAIN"}
 ! CHECK: cuf.alloc !fir.box<!fir.heap<!fir.array<?xi32>>> {bindc_name = "a", data_attr = #cuf.cuda<device>, uniq_name = "_QFEa"} -> !fir.ref<!fir.box<!fir.heap<!fir.array<?xi32>>>>
+! CHECK: %[[PTR:.*]]:2 = hlfir.declare %{{.*}} {data_attr = #cuf.cuda<device>, fortran_attrs = #fir.var_attrs<pointer>, uniq_name = "_QFEp"}
 ! CHECK: %[[ACTIVE:.*]] = fir.call @_FortranACUFDeviceIsActive() {{.*}} : () -> i1
 ! CHECK-NEXT: fir.if %[[ACTIVE]] {
+! CHECK-NEXT: cuf.free %[[PTR]]#0 : !fir.ref<!fir.box<!fir.ptr<!fir.array<?xi32>>>>{{.*}}
 ! CHECK: cuf.deallocate
 ! CHECK: cuf.free{{.*}}
 ! CHECK-NEXT: }