meltygui 0.1.0__py3-none-any.whl

meltygui/core/graphics/cuda_context_core.py

@@ -0,0 +1,166 @@
+"""Shared CUDA context ownership for rendering and GL interop.
+
+Per-device kernel leases live as long as compiled kernels; the standalone display
+stack entry has a separate lease so popping it cannot invalidate cached kernels.
+Both use CUDA primary contexts, so Torch storage remains directly accessible.
+"""
+import atexit
+from contextlib import contextmanager
+import ctypes
+import sys
+import threading
+
+from meltygui.core.melty import Melty
+from meltygui.core.graphics.gl_state import is_gl_thread
+
+
+class CudaContextRuntime:
+    def __init__(self):
+        self.contexts = {}
+        self.kernels = {}
+        self.legacy_kernels = None
+        self.primary_context = None
+        self.primary_thread = None
+        self.last_logged = None
+        self.driver_library = None
+
+
+def runtime():
+    if Melty.cuda_interop is None:
+        Melty.cuda_interop = CudaContextRuntime()
+    state = Melty.cuda_interop
+    # Transfer the old cache's context leases without reacquiring resources.
+    # Existing interop runtime objects gain the pool on first use after hotswap.
+    if 'contexts' not in state.__dict__:
+        state.contexts = {}
+    state.contexts.update(sys.__dict__.pop('_lsd_cuda_march_contexts', {}))
+    return state
+
+
+def primary_context_for(device):
+    """Retain a render-thread primary context for compiled kernels on a device."""
+    if not is_gl_thread():
+        raise RuntimeError('rendering CUDA contexts belong to the render thread')
+    device = int(device)
+    contexts = runtime().contexts
+    if device not in contexts:
+        import meltygui_pycuda.driver as cuda
+        cuda.init()
+        contexts[device] = cuda.Device(device).retain_primary_context()
+    return contexts[device]
+
+
+@contextmanager
+def using_device(device):
+    """Run on a tensor's GPU and restore PyCUDA and native caller contexts."""
+    with using_context(primary_context_for(device)):
+        yield
+
+
+def ensure_primary_context(device):
+    """Install the standalone render-thread stack entry when none exists."""
+    if not is_gl_thread():
+        return False
+    import meltygui_pycuda.driver as cuda
+    cuda.init()
+    state = runtime()
+    if cuda.Context.get_current() is not None or state.primary_context is not None:
+        return False
+    context = cuda.Device(int(device)).retain_primary_context()
+    try:
+        context.push()
+    except Exception:
+        detach_inactive_primary(context)
+        raise
+    state.primary_context = context
+    state.primary_thread = threading.current_thread()
+    atexit.register(_detach_primary)
+    return True
+
+
+def _driver_library():
+    state = runtime()
+    if state.driver_library is None:
+        library = ctypes.CDLL('libcuda.so.1')
+        query = library.cuGLGetDevices_v2
+        query.argtypes = (ctypes.POINTER(ctypes.c_uint), ctypes.POINTER(ctypes.c_int),
+                          ctypes.c_uint, ctypes.c_uint)
+        query.restype = ctypes.c_int
+        for name in ('cuCtxGetCurrent', 'cuCtxPopCurrent_v2'):
+            function = getattr(library, name)
+            function.argtypes = (ctypes.POINTER(ctypes.c_void_p),)
+            function.restype = ctypes.c_int
+        library.cuCtxSetCurrent.argtypes = (ctypes.c_void_p,)
+        library.cuCtxSetCurrent.restype = ctypes.c_int
+        state.driver_library = library
+    return state.driver_library
+
+
+def _native_context():
+    context = ctypes.c_void_p()
+    result = _driver_library().cuCtxGetCurrent(ctypes.byref(context))
+    if result != 0:
+        raise RuntimeError(f"cuCtxGetCurrent failed: {result}")
+    return context.value
+
+
+def current_device_index():
+    """The CUDA ordinal of the context current on this thread, or None."""
+    import meltygui_pycuda.driver as cuda
+    if cuda.Context.get_current() is None:
+        return None
+    bus = cuda.Context.get_device().pci_bus_id()
+    for index in range(cuda.Device.count()):
+        if cuda.Device(index).pci_bus_id() == bus:
+            return index
+    return None
+
+
+def _detach_primary():
+    state = runtime()
+    if state.primary_context is None:
+        return
+    import meltygui_pycuda.driver as cuda
+    if (threading.current_thread() is not state.primary_thread
+            or cuda.Context.get_current() != state.primary_context):
+        # Popping someone else's stack entry corrupts its owner's lifetime.
+        return
+    cuda.Context.pop()
+    detach_inactive_primary(state.primary_context)
+    state.primary_context = None
+    state.primary_thread = None
+
+
+def detach_inactive_primary(context):
+    """Release an owned primary-context lease AFTER removing its PyCUDA entry.
+
+    PyCUDA's inactive-context detach pushes the native context and assumes that
+    releasing it also pops it. cuDevicePrimaryCtxRelease explicitly does not pop
+    (cuda.h); balance that native push without touching PyCUDA's restored stack.
+    """
+    context.detach()
+    popped = ctypes.c_void_p()
+    result = _driver_library().cuCtxPopCurrent_v2(ctypes.byref(popped))
+    if result != 0:
+        raise RuntimeError(f"cuCtxPopCurrent after primary release failed: {result}")
+
+
+@contextmanager
+def using_context(context):
+    """Temporarily restore a registration's CUDA context, preserving the stack."""
+    import meltygui_pycuda.driver as cuda
+    previous_native = _native_context()
+    pushed = cuda.Context.get_current() != context or previous_native != context.handle
+    if pushed:
+        context.push()
+    try:
+        yield
+    finally:
+        if pushed:
+            cuda.Context.pop()
+        # Torch may also switch the native context inside an already-current
+        # scope, where no PyCUDA push/pop was needed on entry.
+        if _native_context() != previous_native:
+            result = _driver_library().cuCtxSetCurrent(previous_native)
+            if result != 0:
+                raise RuntimeError(f"restoring CUDA context failed: {result}")

meltygui/core/graphics/cuda_interop_core.py

@@ -0,0 +1,136 @@
+"""CUDA context selection and registered-buffer operations for OpenGL interop.
+
+The rendering context determines the device. Never register GL resources from an
+unrelated CUDA device, even when Torch's current device happens to be that device.
+Tensor adaptation and texture ownership live in model/cuda_texture_model.py.
+"""
+from contextlib import contextmanager
+import ctypes
+
+from meltygui.core.graphics import cuda_context_core
+from meltygui.core.graphics.gl_state import (
+    ResourceDeletionDeferred, current_context, is_gl_thread,
+)
+
+
+def log_once(message):
+    state = cuda_context_core.runtime()
+    if message != state.last_logged:
+        print(f"[cuda_interop] {message}")
+        state.last_logged = message
+
+
+def gl_devices():
+    """CUDA ordinals associated with the current GL context, or none.
+
+    PyCUDA does not expose cuGLGetDevices; use the CUDA driver's public API.
+    The versioned symbol and enum values are declared in CUDA's cudaGL.h.
+    """
+    if not is_gl_thread() or current_context() is None:
+        return ()
+    import meltygui_pycuda.driver as cuda
+    cuda.init()
+    capacity = cuda.Device.count()
+    devices = (ctypes.c_int * capacity)()
+    count = ctypes.c_uint()
+    # CU_GL_DEVICE_LIST_ALL: all GPUs used by the current GL context.
+    result = cuda_context_core._driver_library().cuGLGetDevices_v2(ctypes.byref(count), devices, capacity, 1)
+    if result != 0:
+        return ()
+    return tuple(devices[:count.value])
+
+
+def cuda_ready():
+    """Whether GL interop is safe in the current thread and context pair."""
+    if not is_gl_thread() or current_context() is None:
+        return False
+    try:
+        import meltygui_pycuda.driver as cuda
+        import meltygui_pycuda.gl
+        context = cuda.Context.get_current()
+        if context is None or context.handle != cuda_context_core._native_context():
+            return False
+        device = cuda_context_core.current_device_index()
+        return device is not None and device in gl_devices()
+    except Exception:
+        return False
+
+
+def ensure_context():
+    """For standalone/test setup, retain the GL device's primary context.
+
+    An existing compatible context is left alone. An incompatible one is
+    rejected, not silently replaced. Only the context retained here is detached.
+    """
+    if not is_gl_thread() or current_context() is None:
+        return False
+    try:
+        import meltygui_pycuda.driver as cuda
+        import meltygui_pycuda.gl
+        devices = gl_devices()
+        if not devices:
+            return False
+        if cuda.Context.get_current() is not None:
+            return cuda_ready()
+        return cuda_context_core.ensure_primary_context(devices[0])
+    except Exception as error:
+        log_once(f"ensure_context failed: {error}")
+        return False
+
+
+def register_buffer(buffer):
+    if not cuda_ready():
+        raise RuntimeError("CUDA registration requires a compatible current GL/CUDA context pair")
+    from meltygui_pycuda import gl as cuda_gl
+    return cuda_gl.RegisteredBuffer(buffer, cuda_gl.graphics_map_flags.WRITE_DISCARD)
+
+
+@contextmanager
+def interop_context():
+    """Activate a GL-compatible context for one adapter operation.
+
+    Torch can change the native context while PyCUDA still records the rendering
+    context. Reactivate that context locally, validate its GL device, and restore
+    the caller's native context afterward. A genuinely incompatible PyCUDA context
+    is rejected. Standalone callers with no PyCUDA context get lazy setup.
+    """
+    if not is_gl_thread() or current_context() is None:
+        yield None
+        return
+    try:
+        import meltygui_pycuda.driver as cuda
+        import meltygui_pycuda.gl
+        context = cuda.Context.get_current()
+        if context is None and ensure_context():
+            context = cuda.Context.get_current()
+    except Exception:
+        context = None
+    if context is None:
+        yield None
+        return
+    with cuda_context_core.using_context(context):
+        yield context if cuda_ready() else None
+
+
+def copy_to_buffer(registered, context, pointer, nbytes):
+    """Copy ready device storage into a mapped GL buffer and unmap on failure."""
+    import meltygui_pycuda.driver as cuda
+    with cuda_context_core.using_context(context):
+        mapping = registered.map()
+        try:
+            destination, capacity = mapping.device_ptr_and_size()
+            if capacity < nbytes:
+                raise ValueError(f"CUDA buffer has {capacity} bytes; upload needs {nbytes}")
+            cuda.memcpy_dtod(destination, pointer, nbytes)
+            cuda.Context.synchronize()
+        finally:
+            mapping.unmap()
+
+
+def unregister_buffer(registered, context):
+    try:
+        with cuda_context_core.using_context(context):
+            registered.unregister()
+    except Exception as error:
+        log_once(f"unregister deferred: {error}")
+        raise ResourceDeletionDeferred from error

meltygui/core/graphics/cuda_kernel_core.py

@@ -0,0 +1,91 @@
+"""Render-thread CUDA compilation and process-lifetime kernel ownership.
+
+Feature views supply CUDA source and entry points. Melty's CUDA runtime retains
+modules alongside their primary contexts, including across definition hotswaps.
+"""
+import sys
+
+from meltygui.core.graphics.cuda_context_core import runtime, using_device
+
+
+def available():
+    try:
+        import meltygui_pycuda.driver
+        import meltygui_pycuda as pycuda  # noqa: F401
+        return True
+    except Exception:
+        return False
+
+
+def host_compiler_flags():
+    """Use an installed CUDA-12-compatible C++ compiler when available."""
+    import shutil
+    for version in ("12", "11", "10"):
+        if shutil.which("gcc-" + version) and shutil.which("g++-" + version):
+            return ["-ccbin", "g++-" + version]
+    return ["-allow-unsupported-compiler"]
+
+
+class CudaKernel:
+    def __init__(self, source, module, functions):
+        self.source = source
+        self.module = module
+        self.functions = functions
+
+
+def _cache():
+    state = runtime()
+    if 'kernels' not in state.__dict__:
+        state.kernels = {}
+    # One-time runtime transfer, not import compatibility. Keep existing compiled
+    # functions alive until their feature is next requested under its context.
+    if 'legacy_kernels' not in state.__dict__ or state.legacy_kernels is None:
+        state.legacy_kernels = {
+            'voxels': sys.__dict__.pop('_lsd_cuda_march_kernels', {}),
+            'lines': sys.__dict__.pop('_melty_cuda_line_kernels', {}),
+        }
+    return state.kernels
+
+
+def kernel_functions(feature, device, source, names):
+    """Compile a feature's entry points once per source/device.
+
+    A failed replacement leaves the last good module owned by the runtime.
+    Replacement and destruction happen with the owning device context active.
+    Unchanged source after hotswap retains the same function objects.
+    """
+    device = int(device)
+    names = tuple(names)
+    # This also enforces the render-thread contract on hits and before adopting
+    # old entries; none of the runtime dictionaries can be mutated by workers.
+    with using_device(device):
+        cache = _cache()
+        key = (feature, device, names)
+        entry = cache.get(key)
+        if entry is not None and entry.source == source:
+            return entry.functions
+        legacy = runtime().legacy_kernels.get(feature, {}).get(device)
+        if legacy is not None:
+            if feature == 'voxels' and isinstance(legacy, dict) and legacy.get('__source__') == hash(source):
+                entry = CudaKernel(source, None, {name: legacy[name] for name in names})
+            elif feature == 'lines' and legacy[0] == source:
+                entry = CudaKernel(source, legacy[1], dict(zip(names, legacy[2:])))
+            if entry is not None and entry.source == source:
+                cache[key] = entry
+                runtime().legacy_kernels[feature].pop(device)
+                return entry.functions
+        import meltygui_pycuda.driver as cuda
+        from meltygui_pycuda.compiler import SourceModule
+        capability = cuda.Device(device).compute_capability()
+        module = SourceModule(source, no_extern_c=True, arch='sm_%d%d' % capability,
+                              options=['-O3'] + host_compiler_flags())
+        try:
+            functions = {name: module.get_function(name) for name in names}
+        except Exception:
+            del module
+            raise
+        cache[key] = CudaKernel(source, module, functions)
+        runtime().legacy_kernels.get(feature, {}).pop(device, None)
+        # Drop replaced and transferred references while the context is active.
+        del entry, legacy
+        return functions