[flang][cuda] Fix device-variable leaks and reset-safe cleanup (#209865)

Example:
```fortran
program main
  use cudafor
  integer, allocatable, device :: b(:)
  integer :: istat

  allocate(b(1024))
  b = 2
  istat = cudaDeviceReset()
end program
```

Managed descriptors and main-program allocatable data could leak. After
`cudaDeviceReset()`, compiler-generated scope-exit cleanup could also
access
stale descriptors or recreate a CUDA context.

Fix: free device data and descriptors in the correct order, and avoid
unnecessary
managed temporary descriptors. Each routine’s CUDA scope-exit cleanup is
grouped
under one `CUFDeviceIsActive()` check using
`cuDevicePrimaryCtxGetState`.
Routines without CUDA cleanup—and individual runtime frees—incur no
query.

flang-rt/lib/cuda/allocator.cpp

@@ -19,8 +19,39 @@
 #include "flang/Runtime/CUDA/common.h"
 #include "flang/Support/Fortran.h"
 
+#include "cuda.h"
+#include "cuda_runtime.h"
+
 namespace Fortran::runtime::cuda {
 
+static bool deviceContextTornDown() {
+  int device{0};
+  if (cudaGetDevice(&device) != cudaSuccess) {
+    return true;
+  }
+  // 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).
+  using GetStateFn = CUresult(CUDAAPI *)(CUdevice, unsigned *, int *);
+  static GetStateFn getState{[]() -> GetStateFn {
+    void *fn{nullptr};
+    if (cudaGetDriverEntryPoint("cuDevicePrimaryCtxGetState", &fn,
+            cudaEnableDefault, nullptr) != cudaSuccess) {
+      return nullptr;
+    }
+    return reinterpret_cast<GetStateFn>(fn);
+  }()};
+  if (!getState) {
+    return true;
+  }
+  unsigned flags{0};
+  int active{0};
+  if (getState(device, &flags, &active) != CUDA_SUCCESS) {
+    return true;
+  }
+  return active == 0;
+}
+
 struct DeviceAllocation {
   void *ptr;
   std::size_t size;
@@ -141,6 +172,8 @@ void RTDEF(CUFRegisterAllocator)() {
       kUnifiedAllocatorPos, {&CUFAllocUnified, CUFFreeUnified});
 }
 
+bool RTDEF(CUFDeviceIsActive)() { return !deviceContextTornDown(); }
+
 cudaStream_t RTDECL(CUFGetAssociatedStream)(void *p) {
   int pos = findAllocation(p);
   if (pos >= 0) {
@@ -192,6 +225,8 @@ void *CUFAllocDevice(std::size_t sizeInBytes,
   return p;
 }
 
+// Scope-exit cleanup is guarded in lowering; explicit deallocation after a
+// reset is unsupported, keeping cudaFreeAsync free of context-query overhead.
 void CUFFreeDevice(void *p) {
   CriticalSection critical{lock};
   int pos = findAllocation(p);

flang/include/flang/Lower/AbstractConverter.h

@@ -358,6 +358,7 @@ public:
   virtual const fir::KindMapping &getKindMap() = 0;
 
   virtual Fortran::lower::StatementContext &getFctCtx() = 0;
+  virtual Fortran::lower::StatementContext &getCudaCleanupCtx() = 0;
 
   /// Generate STAT and ERRMSG from a list of StatOrErrmsg
   virtual std::pair<mlir::Value, mlir::Value>

flang/include/flang/Lower/Bridge.h

@@ -127,6 +127,10 @@ public:
 
   Fortran::lower::StatementContext &fctCtx() { return functionContext; }
 
+  Fortran::lower::StatementContext &cudaCleanupCtx() {
+    return cudaCleanupContext;
+  }
+
   Fortran::lower::StatementContext &openAccCtx() { return openAccContext; }
 
   bool validModule() { return getModule(); }
@@ -163,6 +167,7 @@ private:
 
   Fortran::semantics::SemanticsContext &semanticsContext;
   Fortran::lower::StatementContext functionContext;
+  Fortran::lower::StatementContext cudaCleanupContext;
   Fortran::lower::StatementContext openAccContext;
   const Fortran::common::IntrinsicTypeDefaultKinds &defaultKinds;
   const Fortran::evaluate::IntrinsicProcTable &intrinsics;

flang/include/flang/Optimizer/Builder/Runtime/CUDA/Descriptor.h

@@ -31,6 +31,10 @@ void genSyncGlobalDescriptor(fir::FirOpBuilder &builder, mlir::Location loc,
 void genDescriptorCheckSection(fir::FirOpBuilder &builder, mlir::Location loc,
                                mlir::Value desc);
 
+/// Generate a call returning (as i1) whether the device's primary context is
+/// alive, to guard scope-exit frees against a user cudaDeviceReset().
+mlir::Value genDeviceIsActive(fir::FirOpBuilder &builder, mlir::Location loc);
+
 } // namespace fir::runtime::cuda
 
 #endif // FORTRAN_OPTIMIZER_BUILDER_RUNTIME_CUDA_DESCRIPTOR_H_

flang/include/flang/Runtime/CUDA/descriptor.h

@@ -41,6 +41,10 @@ void RTDECL(CUFSyncGlobalDescriptor)(
 void RTDECL(CUFDescriptorCheckSection)(
     const Descriptor *, const char *sourceFile = nullptr, int sourceLine = 0);
 
+/// True unless the device's primary context has been torn down (e.g. by a user
+/// cudaDeviceReset()); used to skip scope-exit device frees safely.
+bool RTDECL(CUFDeviceIsActive)();
+
 } // extern "C"
 
 } // namespace Fortran::runtime::cuda

flang/lib/Lower/Bridge.cpp

@@ -39,6 +39,7 @@
 #include "flang/Optimizer/Builder/Character.h"
 #include "flang/Optimizer/Builder/FIRBuilder.h"
 #include "flang/Optimizer/Builder/Runtime/Assign.h"
+#include "flang/Optimizer/Builder/Runtime/CUDA/Descriptor.h"
 #include "flang/Optimizer/Builder/Runtime/Character.h"
 #include "flang/Optimizer/Builder/Runtime/Derived.h"
 #include "flang/Optimizer/Builder/Runtime/EnvironmentDefaults.h"
@@ -1246,6 +1247,13 @@ public:
     return bridge.fctCtx();
   }
 
+  Fortran::lower::StatementContext &getCudaCleanupCtx() override final {
+    if (!activeConstructStack.empty() &&
+        activeConstructStack.back().eval.isA<Fortran::parser::BlockConstruct>())
+      return activeConstructStack.back().stmtCtx;
+    return bridge.cudaCleanupCtx();
+  }
+
   /// Initializes values for STAT and ERRMSG
   std::pair<mlir::Value, mlir::Value>
   genStatAndErrmsg(mlir::Location loc,
@@ -1972,11 +1980,20 @@ private:
   void genExitRoutine(bool earlyReturn, mlir::ValueRange retval = {}) {
     if (blockIsUnterminated()) {
       bridge.openAccCtx().finalizeAndKeep();
+      if (bridge.cudaCleanupCtx().hasCode()) {
+        mlir::Location loc = toLocation();
+        mlir::Value active =
+            fir::runtime::cuda::genDeviceIsActive(*builder, loc);
+        builder->genIfThen(loc, active)
+            .genThen([&]() { bridge.cudaCleanupCtx().finalizeAndKeep(); })
+            .end();
+      }
       bridge.fctCtx().finalizeAndKeep();
       mlir::func::ReturnOp::create(*builder, toLocation(), retval);
     }
     if (!earlyReturn) {
       bridge.openAccCtx().pop();
+      bridge.cudaCleanupCtx().pop();
       bridge.fctCtx().pop();
     }
   }
@@ -6322,6 +6339,7 @@ private:
   void startNewFunction(Fortran::lower::pft::FunctionLikeUnit &funit) {
     assert(!builder && "expected nullptr");
     bridge.fctCtx().pushScope();
+    bridge.cudaCleanupCtx().pushScope();
     bridge.openAccCtx().pushScope();
     const Fortran::semantics::Scope &scope = funit.getScope();
     LLVM_DEBUG(llvm::dbgs() << "\n[bridge - startNewFunction]";

flang/lib/Lower/ConvertVariable.cpp

@@ -1063,12 +1063,15 @@ enum class VariableCleanUp { Finalize, Deallocate };
 /// 7.5.6.3 point 3 or if it is an allocatable that must be deallocated. Note
 /// that deallocation will trigger finalization if the type has any.
 static std::optional<VariableCleanUp>
-needDeallocationOrFinalization(const Fortran::lower::pft::Variable &var) {
+needDeallocationOrFinalization(const Fortran::lower::pft::Variable &var,
+                               bool deallocateMainProgramVariable = false) {
   if (!var.hasSymbol())
     return std::nullopt;
   const Fortran::semantics::Symbol &sym = var.getSymbol();
   const Fortran::semantics::Scope &owner = sym.owner();
   if (owner.kind() == Fortran::semantics::Scope::Kind::MainProgram) {
+    if (deallocateMainProgramVariable)
+      return VariableCleanUp::Deallocate;
     // The standard does not require finalizing main program variables.
     return std::nullopt;
   }
@@ -1267,8 +1270,10 @@ static void instantiateLocal(Fortran::lower::AbstractConverter &converter,
     Fortran::lower::defaultInitializeAtRuntime(converter, var.getSymbol(),
                                                symMap);
   auto *builder = &converter.getFirOpBuilder();
-  if (needCUDAAlloc(var.getSymbol()) &&
-      !cuf::isCUDADeviceContext(builder->getRegion())) {
+  bool needsHostCudaCleanup = needCUDAAlloc(var.getSymbol()) &&
+                              !cuf::isCUDADeviceContext(builder->getRegion());
+  bool deallocateMainProgramVariable = false;
+  if (needsHostCudaCleanup) {
     cuf::DataAttributeAttr dataAttr =
         Fortran::lower::translateSymbolCUFDataAttribute(builder->getContext(),
                                                         var.getSymbol());
@@ -1276,10 +1281,15 @@ static void instantiateLocal(Fortran::lower::AbstractConverter &converter,
     fir::ExtendedValue exv =
         converter.getSymbolExtendedValue(var.getSymbol(), &symMap);
     auto *sym = &var.getSymbol();
-    const Fortran::semantics::Scope &owner = sym->owner();
-    if (owner.kind() != Fortran::semantics::Scope::Kind::MainProgram &&
-        dataAttr.getValue() != cuf::DataAttribute::Shared) {
-      converter.getFctCtx().attachCleanup([builder, loc, exv, sym]() {
+    // Free the CUF storage, including in the main program: this is storage
+    // cleanup, not finalization, so it must not be skipped there.
+    if (dataAttr.getValue() != cuf::DataAttribute::Shared) {
+      deallocateMainProgramVariable =
+          sym->owner().kind() == Fortran::semantics::Scope::Kind::MainProgram &&
+          Fortran::semantics::IsAllocatable(*sym);
+      Fortran::lower::StatementContext &cudaCleanupCtx =
+          converter.getCudaCleanupCtx();
+      cudaCleanupCtx.attachCleanup([builder, loc, exv, sym]() {
         cuf::DataAttributeAttr dataAttr =
             Fortran::lower::translateSymbolCUFDataAttribute(
                 builder->getContext(), *sym);
@@ -1289,14 +1299,17 @@ static void instantiateLocal(Fortran::lower::AbstractConverter &converter,
   }
 
   if (std::optional<VariableCleanUp> cleanup =
-          needDeallocationOrFinalization(var)) {
+          needDeallocationOrFinalization(var, deallocateMainProgramVariable)) {
     auto *builder = &converter.getFirOpBuilder();
     mlir::Location loc = converter.getCurrentLocation();
     fir::ExtendedValue exv =
         converter.getSymbolExtendedValue(var.getSymbol(), &symMap);
+    Fortran::lower::StatementContext &cleanupCtx =
+        needsHostCudaCleanup ? converter.getCudaCleanupCtx()
+                             : converter.getFctCtx();
     switch (*cleanup) {
     case VariableCleanUp::Finalize:
-      converter.getFctCtx().attachCleanup([builder, loc, exv]() {
+      cleanupCtx.attachCleanup([builder, loc, exv]() {
         mlir::Value box = builder->createBox(loc, exv);
         fir::runtime::genDerivedTypeDestroy(*builder, loc, box);
       });
@@ -1304,7 +1317,7 @@ static void instantiateLocal(Fortran::lower::AbstractConverter &converter,
     case VariableCleanUp::Deallocate:
       auto *converterPtr = &converter;
       auto *sym = &var.getSymbol();
-      converter.getFctCtx().attachCleanup([converterPtr, loc, exv, sym]() {
+      cleanupCtx.attachCleanup([converterPtr, loc, exv, sym]() {
         const fir::MutableBoxValue *mutableBox =
             exv.getBoxOf<fir::MutableBoxValue>();
         assert(mutableBox &&

flang/lib/Optimizer/Builder/Runtime/CUDA/Descriptor.cpp

@@ -47,3 +47,11 @@ void fir::runtime::cuda::genDescriptorCheckSection(fir::FirOpBuilder &builder,
       builder, loc, fTy, desc, sourceFile, sourceLine)};
   fir::CallOp::create(builder, loc, func, args);
 }
+
+mlir::Value fir::runtime::cuda::genDeviceIsActive(fir::FirOpBuilder &builder,
+                                                  mlir::Location loc) {
+  mlir::func::FuncOp func =
+      fir::runtime::getRuntimeFunc<mkRTKey(CUFDeviceIsActive)>(loc, builder);
+  auto call = fir::CallOp::create(builder, loc, func, mlir::ValueRange{});
+  return builder.createConvert(loc, builder.getI1Type(), call.getResult(0));
+}

flang/lib/Optimizer/CodeGen/CodeGen.cpp

@@ -3838,11 +3838,10 @@ struct LoadOpConversion : public fir::FIROpConversion<fir::LoadOp> {
       if (auto callOp = mlir::dyn_cast_or_null<mlir::LLVM::CallOp>(
               inputBoxStorage.getDefiningOp())) {
         if (callOp.getCallee() &&
-            ((*callOp.getCallee())
-                 .starts_with(RTNAME_STRING(CUFAllocDescriptor)) ||
-             (*callOp.getCallee()).starts_with("__tgt_acc_get_deviceptr"))) {
-          // CUDA Fortran local descriptor are allocated in managed memory. So
-          // new storage must be allocated the same way.
+            (*callOp.getCallee()).starts_with("__tgt_acc_get_deviceptr")) {
+          // The device pointer descriptor is allocated in managed memory, so
+          // new storage must be allocated the same way. A CUFAllocDescriptor
+          // source is handled below if the load is used by a GPU launch.
           auto mod = load->getParentOfType<mlir::ModuleOp>();
           newBoxStorage =
               genCUFAllocDescriptor(loc, rewriter, mod, boxTy, lowerTy());

flang/test/Fir/CUDA/cuda-code-gen.mlir

@@ -18,7 +18,7 @@ module attributes {dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<f80, dense<128> :
   }
 
   // CHECK-LABEL: llvm.func @_QQmain()
-  // CHECK-COUNT-2: llvm.call @_FortranACUFAllocDescriptor 
+  // CHECK-COUNT-1: llvm.call @_FortranACUFAllocDescriptor
 
   fir.global linkonce @_QQclX3C737464696E3E00 constant : !fir.char<1,8> {
     %0 = fir.string_lit "<stdin>\00"(8) : !fir.char<1,8>
@@ -125,7 +125,7 @@ module attributes {dlti.dl_spec = #dlti.dl_spec<f80 = dense<128> : vector<2xi64>
 }
 
 // CHECK-LABEL: llvm.func @_QQmain()
-// CHECK-COUNT-4: llvm.call @_FortranACUFAllocDescriptor
+// CHECK-COUNT-3: llvm.call @_FortranACUFAllocDescriptor
 
 // -----
 
@@ -169,7 +169,7 @@ module attributes {dlti.dl_spec = #dlti.dl_spec<!llvm.ptr<270> = dense<32> : vec
 }
 
 // CHECK-LABEL: llvm.func @_QQmain()
-// CHECK-COUNT-3: llvm.call @_FortranACUFAllocDescriptor
+// CHECK-COUNT-2: llvm.call @_FortranACUFAllocDescriptor
 
 // -----
 

flang/test/Lower/CUDA/cuda-allocatable.cuf

@@ -118,12 +118,14 @@ end subroutine
 ! CHECK: %[[LOAD_B:.*]] = fir.load %[[BOX_B_DECL]]#0 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
 ! CHECK: fir.call @_FortranAAllocatableSetBounds
 ! CHECK: %{{.*}} = cuf.allocate %[[BOX_A_DECL]]#0 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>> source(%[[LOAD_B]] : !fir.box<!fir.heap<!fir.array<?xf32>>>) {data_attr = #cuf.cuda<device>} -> i32
+! CHECK: fir.call @_FortranACUFDeviceIsActive()
 ! CHECK: fir.if
-! CHECK: fir.freemem
 ! CHECK: fir.if %{{.*}} {
 ! CHECK:   %{{.*}} = cuf.deallocate %[[BOX_A_DECL]]#0 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>> {data_attr = #cuf.cuda<device>} -> i32
 ! CHECK: }
 ! CHECK: cuf.free %[[BOX_A_DECL]]#0 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>> {data_attr = #cuf.cuda<device>}
+! CHECK: fir.if
+! CHECK: fir.freemem
 
 subroutine sub6()
   real, allocatable, device :: a(:)

flang/test/Lower/CUDA/cuda-derived.cuf

@@ -15,7 +15,11 @@ contains
 ! CHECK-LABEL: func.func @_QMm1Psub1()
 ! CHECK: %[[ALLOC:.*]] = cuf.alloc !fir.type<_QMm1Tty_device{x:!fir.box<!fir.heap<!fir.array<?xi32>>>}> {bindc_name = "a", data_attr = #cuf.cuda<managed>, uniq_name = "_QMm1Fsub1Ea"} -> !fir.ref<!fir.type<_QMm1Tty_device{x:!fir.box<!fir.heap<!fir.array<?xi32>>>}>>
 ! CHECK: %[[DECL:.*]]:2 = hlfir.declare %[[ALLOC]] {data_attr = #cuf.cuda<managed>, uniq_name = "_QMm1Fsub1Ea"} : (!fir.ref<!fir.type<_QMm1Tty_device{x:!fir.box<!fir.heap<!fir.array<?xi32>>>}>>) -> (!fir.ref<!fir.type<_QMm1Tty_device{x:!fir.box<!fir.heap<!fir.array<?xi32>>>}>>, !fir.ref<!fir.type<_QMm1Tty_device{x:!fir.box<!fir.heap<!fir.array<?xi32>>>}>>)
-! CHECK: cuf.free %[[DECL]]#0 : !fir.ref<!fir.type<_QMm1Tty_device{x:!fir.box<!fir.heap<!fir.array<?xi32>>>}>> {data_attr = #cuf.cuda<managed>}
+! CHECK: %[[ACTIVE:.*]] = fir.call @_FortranACUFDeviceIsActive() {{.*}} : () -> i1
+! CHECK-NEXT: fir.if %[[ACTIVE]] {
+! CHECK: fir.call @_FortranADestroy
+! CHECK-NEXT: cuf.free %[[DECL]]#0 : !fir.ref<!fir.type<_QMm1Tty_device{x:!fir.box<!fir.heap<!fir.array<?xi32>>>}>> {data_attr = #cuf.cuda<managed>}
+! CHECK-NEXT: }
 
 end module
 

flang/test/Lower/CUDA/cuda-return01.cuf

@@ -9,7 +9,10 @@ end
 
 ! CHECK: func.func @_QPsub1()
 ! CHECK: %[[DECL:.*]]:2 = hlfir.declare
-! CHECK: cuf.free %[[DECL]]#0 : !fir.ref<!fir.array<10xi32>>
+! CHECK: %[[SUB1_ACTIVE:.*]] = fir.call @_FortranACUFDeviceIsActive() {{.*}} : () -> i1
+! CHECK-NEXT: fir.if %[[SUB1_ACTIVE]] {
+! CHECK: cuf.free %[[DECL]]#0 : !fir.ref<!fir.array<10xi32>>{{.*}}
+! CHECK-NEXT: }
 ! CHECK: return
 ! CHECK: }
 
@@ -19,9 +22,12 @@ subroutine sub2
 end
 
 ! CHECK-LABEL: func.func @_QPsub2()
+! CHECK: %[[SUB2_ACTIVE:.*]] = fir.call @_FortranACUFDeviceIsActive() {{.*}} : () -> i1
+! CHECK-NEXT: fir.if %[[SUB2_ACTIVE]] {
 ! CHECK: fir.if
-! CHECK: cuf.deallocate 
-! CHECK: cuf.free
+! CHECK: cuf.deallocate
+! CHECK: cuf.free{{.*}}
+! CHECK-NEXT: }
 
 program main
   integer, allocatable, device :: a(:)
@@ -30,4 +36,9 @@ end
 
 ! 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-NOT: cuf.free
+! CHECK: %[[ACTIVE:.*]] = fir.call @_FortranACUFDeviceIsActive() {{.*}} : () -> i1
+! CHECK-NEXT: fir.if %[[ACTIVE]] {
+! CHECK: cuf.deallocate
+! CHECK: cuf.free{{.*}}
+! CHECK-NEXT: }
+! CHECK-NEXT: return

flang/test/Lower/CUDA/cuda-return02.cuf

@@ -17,8 +17,16 @@ end
 ! CHECK: %[[DECL:.*]]:2 = hlfir.declare
 ! CHECK: cf.cond_br %{{.*}}, ^bb1, ^bb2
 ! CHECK-NEXT: ^bb1:
+! CHECK-NEXT: %[[ACTIVE1:.*]] = fir.call @_FortranACUFDeviceIsActive()
+! CHECK-NEXT: fir.if %[[ACTIVE1]] {
+! CHECK-NEXT: cuf.free %[[DECL]]#0 : !fir.ref<!fir.array<10xi32>>{{.*}}
+! CHECK-NEXT: }
 ! CHECK-NEXT: return
 ! CHECK-NEXT: ^bb2:
+! CHECK-NEXT: %[[ACTIVE2:.*]] = fir.call @_FortranACUFDeviceIsActive()
+! CHECK-NEXT: fir.if %[[ACTIVE2]] {
+! CHECK-NEXT: cuf.free %[[DECL]]#0 : !fir.ref<!fir.array<10xi32>>{{.*}}
+! CHECK-NEXT: }
 ! CHECK-NEXT: return
 ! CHECK-NEXT: }
 
@@ -42,5 +50,9 @@ end
 ! CHECK: ^bb2:
 ! CHECK: cf.br ^bb3
 ! CHECK: ^bb3:
-! CHECK: cuf.free %[[DECL]]#0 : !fir.ref<!fir.array<10xi32>>
+! CHECK-NEXT: %[[SUB_ACTIVE:.*]] = fir.call @_FortranACUFDeviceIsActive()
+! CHECK-NEXT: fir.if %[[SUB_ACTIVE]] {
+! CHECK-NEXT: cuf.free %[[DECL]]#0 : !fir.ref<!fir.array<10xi32>>{{.*}}
+! CHECK-NEXT: }
+! CHECK-NEXT: return
 ! CHECK: }