From 833b7e044c0b28deb8aa57fe90c8a2ebe0101eff Mon Sep 17 00:00:00 2001
From: Lukasz Stafiniak <lukstafi@gmail.com>
Date: Sun, 29 Sep 2024 13:50:13 +0200
Subject: [PATCH] Fix context aliveness in test and README

---
 README.md     | 8 +++++---
 test/saxpy.ml | 4 +++-
 2 files changed, 8 insertions(+), 4 deletions(-)

diff --git a/README.md b/README.md
index 9730645..d47cf92 100644
--- a/README.md
+++ b/README.md
@@ -14,7 +14,7 @@ let prog = Cu.compile_to_ptx ~cu_src:kernel ~name:"saxpy" ~options:[ "--use_fast
 let () =
   Cu.init ();
   let device = Cu.device_get ~ordinal:0 in
-  let _context = Cu.ctx_create ~flags:0 device in
+  let context = Cu.ctx_create ~flags:0 device in
   let module_ = Cu.module_load_data_ex prog [] in
   let kernel = Cu.module_get_function module_ ~name:"saxpy" in
   ...
@@ -34,12 +34,14 @@ let () =
   Cu.memcpy_D_to_H ~dst:hOut ~src:dOut ();
   Cu.mem_free dX;
   ...
-  Cu.module_unload module_
+  Cu.module_unload module_;
+  (* Keep the context alive up till here. *)
+  ignore (Sys.opaque_identity context)
 ```
 
 (the `...` are parts omitted for presentation brevity).
 You can see how a kernel is compiled and launched, how on-device tensors are created, retrieved to host
-(i.e. the CPU), and released. `_context` is (or might be) destroyed when it is finalized by GC.
+(i.e. the CPU), and released.
 
 Note that you don't need to add the include path to the `compile_to_ptx` options.
 
diff --git a/test/saxpy.ml b/test/saxpy.ml
index b75c4a8..0141e35 100644
--- a/test/saxpy.ml
+++ b/test/saxpy.ml
@@ -23,7 +23,7 @@ let%expect_test "SAXPY" =
   Cu.init ();
   if Cu.Device.get_count () > 0 then (
     let device = Cu.Device.get ~ordinal:0 in
-    let _context = Cu.Context.create [] device in
+    let context = Cu.Context.create [] device in
     let module_ = Cu.Module.load_data_ex prog [] in
     let kernel = Cu.Module.get_function module_ ~name:"saxpy" in
     let size = num_threads * num_blocks in
@@ -55,6 +55,8 @@ let%expect_test "SAXPY" =
     Cu.Deviceptr.mem_free dY;
     Cu.Deviceptr.mem_free dOut;
     Cu.Module.unload module_;
+    (* Keep the context around up till here. *)
+    ignore (Sys.opaque_identity context);
     Format.set_margin 110;
     for i = 0 to size - 1 do
       let ( ! ) arr = Host.get arr [| i |] in