cuda-core: simplify strided memory CPU example flow

Refactor the cffi CPU example into small helpers so setup/teardown is easier to follow and avoid global state for compiled symbols. Also direct progress prints to stderr to match example logging behavior.

Made-with: Cursor

Author: cpcloud

cuda_core/examples/strided_memory_view_cpu.py

@@ -15,10 +15,10 @@
 # ################################################################################
 
 import importlib
-import shutil
 import string
 import sys
 import tempfile
+from contextlib import contextmanager
 
 try:
     from cffi import FFI
@@ -61,9 +61,7 @@
 # We assume the 0-th argument supports either DLPack or CUDA Array Interface (both
 # of which are supported by StridedMemoryView).
 @args_viewable_as_strided_memory((0,))
-def my_func(arr):
-    global cpu_func
-    global cpu_prog
+def my_func(arr, cpu_prog, cpu_func):
     # Create a memory view over arr (assumed to be a 1D array of int32). The stream
     # ordering is taken care of, so that arr can be safely accessed on our work
     # stream (ordered after a data stream on which arr is potentially prepared).
@@ -79,8 +77,7 @@ def my_func(arr):
     cpu_func(cpu_prog.cast("int*", view.ptr), size)
 
 
-def run():
-    global my_func, cpu_func, cpu_prog
+def _create_cpu_program():
     # Here is a concrete (very naive!) implementation on CPU:
     cpu_code = string.Template(r"""
     extern "C"
@@ -101,34 +98,37 @@ def run():
         source_extension=".cpp",
         extra_compile_args=["-std=c++11"],
     )
-    temp_dir = tempfile.mkdtemp()
-    cpu_func = None
+    return cpu_prog
+
+
+@contextmanager
+def _compiled_cpu_func(cpu_prog, temp_dir):
     saved_sys_path = sys.path.copy()
     try:
         cpu_prog.compile(tmpdir=temp_dir)
-
         sys.path.append(temp_dir)
         cpu_func = getattr(importlib.import_module("_cpu_obj.lib"), func_name)
+        yield cpu_func
+    finally:
+        sys.path = saved_sys_path
+        # Ensure cffi modules are unloadable before removing the temp build dir.
+        sys.modules.pop("_cpu_obj.lib", None)
+        sys.modules.pop("_cpu_obj", None)
 
-        # Create input array on CPU
-        arr_cpu = np.zeros(1024, dtype=np.int32)
-        print(f"before: {arr_cpu[:10]=}")
 
-        # Run the workload
-        my_func(arr_cpu)
+def _run_example(cpu_prog, cpu_func):
+    arr_cpu = np.zeros(1024, dtype=np.int32)
+    print(f"before: {arr_cpu[:10]=}", file=sys.stderr)
+    my_func(arr_cpu, cpu_prog, cpu_func)
+    print(f"after: {arr_cpu[:10]=}", file=sys.stderr)
+    assert np.allclose(arr_cpu, np.arange(1024, dtype=np.int32))
 
-        # Check the result
-        print(f"after: {arr_cpu[:10]=}")
-        assert np.allclose(arr_cpu, np.arange(1024, dtype=np.int32))
-    finally:
-        sys.path = saved_sys_path
-        # to allow FFI module to unload, we delete references to
-        # to cpu_func
-        if cpu_func is not None:
-            del cpu_func
-        del cpu_prog, my_func
-        # clean up temp directory
-        shutil.rmtree(temp_dir)
+
+def run():
+    cpu_prog = _create_cpu_program()
+    with tempfile.TemporaryDirectory() as temp_dir:
+        with _compiled_cpu_func(cpu_prog, temp_dir) as cpu_func:
+            _run_example(cpu_prog, cpu_func)
 
 
 if __name__ == "__main__":