From 534e6b64ab3e213c39c4c302f452a361fad98cd9 Mon Sep 17 00:00:00 2001
From: Lukasz Stafiniak <lukstafi@gmail.com>
Date: Sat, 26 Oct 2024 14:35:06 +0200
Subject: [PATCH] Remove `Module.unload`, attach a finalizer instead

---
 CHANGES.md  |  4 ++++
 cudajit.ml  | 19 +++++++++++--------
 cudajit.mli | 11 +++++------
 3 files changed, 20 insertions(+), 14 deletions(-)

diff --git a/CHANGES.md b/CHANGES.md
index 9ac4bdf..6d353ec 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -7,6 +7,10 @@
 - `cuda_call_hook` to help in debugging.
 - `is_success` functions.
 
+### Changes
+
+- Removed `Module.unload`, instead `Module.load_data_ex` attaches an unload as a finalizer.
+
 ## [0.5.0] 2024-09-25
 
 ### Added
diff --git a/cudajit.ml b/cudajit.ml
index 79283da..0fd6f89 100644
--- a/cudajit.ml
+++ b/cudajit.ml
@@ -1536,12 +1536,17 @@ module Module = struct
              | POSITION_INDEPENDENT_CODE c -> [ bi2vp c ])
            options
     in
-    check "cu_module_load_data_ex"
-    @@ Cuda.cu_module_load_data_ex cu_mod
-         (coerce (ptr char) (ptr void) ptx.Nvrtc.ptx)
-         n_opts (CArray.start c_options)
-    @@ CArray.start c_opts_args;
-    !@cu_mod
+    let unload cu_mod = check "cu_module_unload" @@ Cuda.cu_module_unload cu_mod in
+    let result =
+      check "cu_module_load_data_ex"
+      @@ Cuda.cu_module_load_data_ex cu_mod
+           (coerce (ptr char) (ptr void) ptx.Nvrtc.ptx)
+           n_opts (CArray.start c_options)
+      @@ CArray.start c_opts_args;
+      !@cu_mod
+    in
+    Gc.finalise unload result;
+    result
 
   let get_function module_ ~name =
     let open Ctypes in
@@ -1555,8 +1560,6 @@ module Module = struct
     let size_in_bytes = allocate size_t Unsigned.Size_t.zero in
     check "cu_module_get_global" @@ Cuda.cu_module_get_global device size_in_bytes module_ name;
     (Deviceptr !@device, !@size_in_bytes)
-
-  let unload cu_mod = check "cu_module_unload" @@ Cuda.cu_module_unload cu_mod
 end
 
 module Stream = struct
diff --git a/cudajit.mli b/cudajit.mli
index ed45a73..d665e0a 100644
--- a/cudajit.mli
+++ b/cudajit.mli
@@ -572,18 +572,17 @@ module Module : sig
   val load_data_ex : Nvrtc.compile_to_ptx_result -> jit_option list -> t
   (** Currently, the image passed via this call is the PTX source. See
       {{:https://docs.nvidia.com/cuda/cuda-driver-api/group__CUDA__MODULE.html#group__CUDA__MODULE_1g9e8047e9dbf725f0cd7cafd18bfd4d12}
-        cuModuleLoadDataEx}. *)
+        cuModuleLoadDataEx}.
+
+      The module is finalized using
+      {{:https://docs.nvidia.com/cuda/cuda-driver-api/group__CUDA__MODULE.html#group__CUDA__MODULE_1g8ea3d716524369de3763104ced4ea57b}
+        cuModuleUnload}. *)
 
   val get_function : t -> name:string -> func
   (** See
       {{:https://docs.nvidia.com/cuda/cuda-driver-api/group__CUDA__MODULE.html#group__CUDA__MODULE_1ga52be009b0d4045811b30c965e1cb2cf}
         cuModuleGetFunction}. *)
 
-  val unload : t -> unit
-  (** See
-      {{:https://docs.nvidia.com/cuda/cuda-driver-api/group__CUDA__MODULE.html#group__CUDA__MODULE_1g8ea3d716524369de3763104ced4ea57b}
-        cuModuleUnload}. *)
-
   val get_global : t -> name:string -> Deviceptr.t * Unsigned.size_t
   (** See
       {{:https://docs.nvidia.com/cuda/cuda-driver-api/group__CUDA__MODULE.html#group__CUDA__MODULE_1gf3e43672e26073b1081476dbf47a86ab}