Remove Module.unload, attach a finalizer instead

lukstafi · Oct 26, 2024 · 534e6b6 · 534e6b6
1 parent 4f9d96e
commit 534e6b6
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 14 deletions.
diff --git 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
@@ -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
@@ -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}