Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
39 changes: 39 additions & 0 deletions cuda_core/tests/helpers/__init__.py
Original file line number Diff line number Diff line change
@@ -1,9 +1,13 @@
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

import functools
import os
import pathlib
import sys
from typing import Union

from cuda.core._utils.cuda_utils import handle_return

CUDA_PATH = os.environ.get("CUDA_PATH")
CUDA_INCLUDE_PATH = None
Expand All @@ -24,3 +28,38 @@
# Import shared platform helpers for tests across repos
sys.path.insert(0, str(pathlib.Path(__file__).resolve().parents[3] / "cuda_python_test_helpers"))
from cuda_python_test_helpers import * # noqa: F403


@functools.cache
def supports_ipc_mempool(device_id: Union[int, object]) -> bool:
"""Return True if mempool IPC via POSIX file descriptor is supported.

Uses cuDeviceGetAttribute(CU_DEVICE_ATTRIBUTE_MEMPOOL_SUPPORTED_HANDLE_TYPES)
to check for CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR support. Does not
require an active CUDA context.
"""
if IS_WSL: # noqa: F405
return False

try:
# Lazy import to avoid hard dependency when not running GPU tests
try:
from cuda.bindings import driver # type: ignore
except Exception:
from cuda import cuda as driver # type: ignore

# Initialize CUDA
handle_return(driver.cuInit(0))

# Resolve device id from int or Device-like object
dev_id = int(getattr(device_id, "device_id", device_id))

# Query supported mempool handle types bitmask
attr = driver.CUdevice_attribute.CU_DEVICE_ATTRIBUTE_MEMPOOL_SUPPORTED_HANDLE_TYPES
mask = handle_return(driver.cuDeviceGetAttribute(attr, dev_id))

# Check POSIX FD handle type support via bitmask
posix_fd = driver.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR
return (int(mask) & int(posix_fd)) != 0
except Exception:
return False
3 changes: 1 addition & 2 deletions cuda_core/tests/test_memory.py
Original file line number Diff line number Diff line change
Expand Up @@ -38,15 +38,14 @@
from cuda.core._memory import IPCBufferDescriptor
from cuda.core._utils.cuda_utils import CUDAError, handle_return
from cuda.core.utils import StridedMemoryView
from helpers import IS_WINDOWS
from helpers import IS_WINDOWS, supports_ipc_mempool
from helpers.buffers import DummyUnifiedMemoryResource

from conftest import (
create_managed_memory_resource_or_skip,
skip_if_managed_memory_unsupported,
skip_if_pinned_memory_unsupported,
)
from cuda_python_test_helpers import supports_ipc_mempool

POOL_SIZE = 2097152 # 2MB size

Expand Down
40 changes: 0 additions & 40 deletions cuda_python_test_helpers/cuda_python_test_helpers/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,20 +2,15 @@
# SPDX-License-Identifier: Apache-2.0

import ctypes
import functools
import os
import platform
import sys
from contextlib import suppress
from typing import Union

from cuda.core._utils.cuda_utils import handle_return

__all__ = [
"IS_WINDOWS",
"IS_WSL",
"libc",
"supports_ipc_mempool",
"under_compute_sanitizer",
]

Expand Down Expand Up @@ -68,38 +63,3 @@ def under_compute_sanitizer() -> bool:
# Another common indicator: sanitizer injectors are configured via env vars.
inj = os.environ.get("CUDA_INJECTION64_PATH", "")
return "compute-sanitizer" in inj or "cuda-memcheck" in inj


@functools.cache
def supports_ipc_mempool(device_id: Union[int, object]) -> bool:
"""Return True if mempool IPC via POSIX file descriptor is supported.

Uses cuDeviceGetAttribute(CU_DEVICE_ATTRIBUTE_MEMPOOL_SUPPORTED_HANDLE_TYPES)
to check for CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR support. Does not
require an active CUDA context.
"""
if _detect_wsl():
return False

try:
# Lazy import to avoid hard dependency when not running GPU tests
try:
from cuda.bindings import driver # type: ignore
except Exception:
from cuda import cuda as driver # type: ignore

# Initialize CUDA
handle_return(driver.cuInit(0))

# Resolve device id from int or Device-like object
dev_id = int(getattr(device_id, "device_id", device_id))

# Query supported mempool handle types bitmask
attr = driver.CUdevice_attribute.CU_DEVICE_ATTRIBUTE_MEMPOOL_SUPPORTED_HANDLE_TYPES
mask = handle_return(driver.cuDeviceGetAttribute(attr, dev_id))

# Check POSIX FD handle type support via bitmask
posix_fd = driver.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR
return (int(mask) & int(posix_fd)) != 0
except Exception:
return False
Loading