Skip to content

Avoid aggregate CUmemLocation initialization #2673

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
rwgk wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Avoid aggregate CUmemLocation initialization #2673

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
7 changes: 3 additions & 4 deletions cuda_core/cuda/core/_memory/_device_memory_resource.pyx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -321,10 +321,9 @@ cpdef str DMR_mempool_get_access(DeviceMemoryResource dmr, int device_id):

cdef int dev_id = Device(device_id).device_id
cdef cydriver.CUmemAccess_flags flags
cdef cydriver.CUmemLocation location = cydriver.CUmemLocation(
type=cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE,
id=dev_id,
)
cdef cydriver.CUmemLocation location
location.type = cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE
location.id = dev_id

with nogil:
HANDLE_RETURN(cydriver.cuMemPoolGetAccess(&flags, as_cu(dmr._h_pool), &location))
Expand Down
22 changes: 10 additions & 12 deletions cuda_core/cuda/core/_memory/_location.pxd
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,24 +15,22 @@ from cuda.bindings cimport cydriver

IF CUDA_CORE_BUILD_MAJOR >= 13:
cdef inline cydriver.CUmemLocation to_cumemlocation(str kind, int loc_id):
cdef cydriver.CUmemLocation cu_loc
if kind == "device":
return cydriver.CUmemLocation(
type=cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE,
id=loc_id)
cu_loc.type = cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE
cu_loc.id = loc_id
elif kind == "host":
return cydriver.CUmemLocation(
type=cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST,
id=0)
cu_loc.type = cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST
cu_loc.id = 0
elif kind == "host_numa":
return cydriver.CUmemLocation(
type=cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST_NUMA,
id=loc_id)
cu_loc.type = cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST_NUMA
cu_loc.id = loc_id
elif kind == "host_numa_current":
return cydriver.CUmemLocation(
type=cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST_NUMA_CURRENT,
id=0)
cu_loc.type = cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST_NUMA_CURRENT
cu_loc.id = 0
else:
raise ValueError(f"unknown location kind: {kind!r}")
return cu_loc
ELSE:
cdef inline cydriver.CUmemLocation to_cumemlocation(str kind, int loc_id):
raise NotImplementedError(
Expand Down
5 changes: 2 additions & 3 deletions cuda_core/cuda/core/_memory/_managed_memory_ops.pyx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -193,9 +193,8 @@ cdef void _do_single_advise(Buffer buf, object advice_value, object loc, bint al
# Driver ignores location for read_mostly / unset_preferred_location
# advice values but still validates the CUmemLocation; pass a
# host placeholder.
cu_loc = cydriver.CUmemLocation(
type=cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST,
id=0)
cu_loc.type = cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST
cu_loc.id = 0
else:
cu_loc = to_cumemlocation(loc.kind, loc.id)
with nogil:
Expand Down
5 changes: 3 additions & 2 deletions cuda_core/cuda/core/_memory/_memory_pool.pyx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -279,8 +279,9 @@ cdef int MP_init_current_pool(
"""
IF CUDA_CORE_BUILD_MAJOR >= 13:
cdef cydriver.CUmemoryPool pool
cdef cydriver.CUmemLocation loc = cydriver.CUmemLocation(
type=loc_type, id=loc_id)
cdef cydriver.CUmemLocation loc
loc.type = loc_type
loc.id = loc_id
with nogil:
HANDLE_RETURN(cydriver.cuMemGetMemPool(&pool, &loc, alloc_type))
self._h_pool = create_mempool_handle_ref(pool)
Expand Down
7 changes: 3 additions & 4 deletions cuda_core/cuda/core/_memory/_peer_access_utils.pyx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -113,10 +113,9 @@ cdef inline tuple _query_peer_access_ids(DeviceMemoryResource mr):
cdef inline bint _peer_access_includes(DeviceMemoryResource mr, int dev_id):
"""Return True if peer access from ``dev_id`` is currently granted."""
cdef cydriver.CUmemAccess_flags flags
cdef cydriver.CUmemLocation location = cydriver.CUmemLocation(
type=cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE,
id=dev_id,
)
cdef cydriver.CUmemLocation location
location.type = cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE
location.id = dev_id
with nogil:
HANDLE_RETURN(cydriver.cuMemPoolGetAccess(&flags, as_cu(mr._h_pool), &location))
return flags == cydriver.CUmemAccess_flags.CU_MEM_ACCESS_FLAGS_PROT_READWRITE
Expand Down
12 changes: 5 additions & 7 deletions cuda_core/cuda/core/graph/_graph_node.pyx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -827,20 +827,18 @@ cdef inline AllocNode GN_alloc(GraphNode self, size_t size, object device,
num_deps = 1

cdef vector[cydriver.CUmemAccessDesc] access_descs
cdef cydriver.CUmemAccessDesc access_desc
cdef int peer_id
cdef list peer_ids = []

if peer_access is not None:
for peer_dev in peer_access:
peer_id = getattr(peer_dev, 'device_id', peer_dev)
peer_ids.append(peer_id)
access_descs.push_back(cydriver.CUmemAccessDesc_st(
cydriver.CUmemLocation_st(
cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE,
peer_id
),
cydriver.CUmemAccess_flags.CU_MEM_ACCESS_FLAGS_PROT_READWRITE
))
access_desc.location.type = cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE
access_desc.location.id = peer_id
access_desc.flags = cydriver.CUmemAccess_flags.CU_MEM_ACCESS_FLAGS_PROT_READWRITE
access_descs.push_back(access_desc)

cdef str memory_type_str = "device" if memory_type is None else str(memory_type)

Expand Down
Loading