cuda-core: add pixi examples env and fix runnable examples (NVIDIA#1718)

* cuda-core: add pixi examples env and fix example run regressions

Add a dedicated Pixi `examples` environment with the optional runtime dependencies needed by the example suite and CUDA include-path activation. Also fix two runnable examples so they execute correctly under the new env.

Made-with: Cursor

* cuda-core: require pytorch-gpu in examples env

Use the explicit `pytorch-gpu` metapackage for the examples environment so PyTorch resolves to a CUDA-enabled build by design.

Made-with: Cursor

* 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

* cuda-core: apply pre-commit simplification in CPU example

Adopt ruff's SIM117 recommendation by collapsing nested context managers in the CPU strided memory example run path.

Made-with: Cursor

Author: cpcloud

cuda_core/examples/jit_lto_fractal.py

@@ -76,7 +76,7 @@ def __init__(self):
 
         # Setup the launch configuration such that each thread will be generating one pixel, and subdivide
         # the problem into 16x16 chunks.
-        self.grid = (self.width / 16, self.height / 16, 1.0)
+        self.grid = (self.width // 16, self.height // 16, 1)
         self.block = (16, 16, 1)
         self.config = LaunchConfig(grid=self.grid, block=self.block)
 

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
+def _create_cpu_program():
     # Here is a concrete (very naive!) implementation on CPU:
     cpu_code = string.Template(r"""
     extern "C"
@@ -101,31 +98,36 @@ def run():
         source_extension=".cpp",
         extra_compile_args=["-std=c++11"],
     )
-    temp_dir = tempfile.mkdtemp()
+    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
-        del cpu_func, my_func
-        # clean up temp directory
-        shutil.rmtree(temp_dir)
+
+def run():
+    cpu_prog = _create_cpu_program()
+    with tempfile.TemporaryDirectory() as temp_dir, _compiled_cpu_func(cpu_prog, temp_dir) as cpu_func:
+        _run_example(cpu_prog, cpu_func)
 
 
 if __name__ == "__main__":