[Mlir-commits] [mlir] [MLIR][ExecutionEngine] Tolerate CUDA_ERROR_DEINITIALIZED in mgpuModuleUnload (PR #190563)

[Jared Hoberock]

https://github.com/jaredhoberock updated https://github.com/llvm/llvm-project/pull/190563

>From 703a90c6f1ab5c74f6d344bc26fefe6296ceb28d Mon Sep 17 00:00:00 2001
From: Jared Hoberock <jaredhoberock at gmail.com>
Date: Sun, 5 Apr 2026 19:25:47 -0500
Subject: [PATCH] [MLIR][ExecutionEngine] Tolerate CUDA_ERROR_DEINITIALIZED in
 mgpuModuleUnload

mgpuModuleUnload may be called from a global destructor (registered by
SelectObjectAttr's appendToGlobalDtors) after the CUDA primary context
has already been destroyed during program shutdown. In this case,
cuModuleUnload returns CUDA_ERROR_DEINITIALIZED, which is benign since
the module's resources are already freed with the context.

This matches the approach used by Clang's CUDA support, TVM, and GCC's
libgomp NVPTX plugin for the same shutdown ordering issue.
---
 mlir/lib/ExecutionEngine/CudaRuntimeWrappers.cpp | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/ExecutionEngine/CudaRuntimeWrappers.cpp b/mlir/lib/ExecutionEngine/CudaRuntimeWrappers.cpp
index 7bf6804902479..24c88b9fa587b 100644
--- a/mlir/lib/ExecutionEngine/CudaRuntimeWrappers.cpp
+++ b/mlir/lib/ExecutionEngine/CudaRuntimeWrappers.cpp
@@ -127,7 +127,12 @@ extern "C" MLIR_CUDA_WRAPPERS_EXPORT CUmodule mgpuModuleLoad(void *data) {
 }
 
 extern "C" MLIR_CUDA_WRAPPERS_EXPORT void mgpuModuleUnload(CUmodule module) {
-  CUDA_REPORT_IF_ERROR(cuModuleUnload(module));
+  // At program exit, the CUDA primary context may already be destroyed.
+  // CUDA_ERROR_DEINITIALIZED is benign — the module's resources are already
+  // freed with the context.
+  CUresult result = cuModuleUnload(module);
+  if (result != CUDA_SUCCESS && result != CUDA_ERROR_DEINITIALIZED)
+    CUDA_REPORT_IF_ERROR(result);
 }
 
 extern "C" MLIR_CUDA_WRAPPERS_EXPORT CUfunction