Skip to content

Commit c52107f

Browse files
juenglinjuenglin
committed
chore: refactor pyx files for readability
1 parent 94e3a76 commit c52107f

18 files changed

Lines changed: 102 additions & 126 deletions

‎cuda_bindings/tests/cython/test_ccudart.pyx‎

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -59,11 +59,8 @@ cdef extern from *:
5959

6060
def test_ccudart_interoperable():
6161
# struct
62-
cdef dim3 oldDim, newDim
63-
oldDim.x = 1
64-
oldDim.y = 2
65-
oldDim.z = 3
66-
newDim = copy_and_append_dim3(oldDim)
62+
cdef dim3 oldDim = [1, 2, 3]
63+
cdef dim3 newDim = copy_and_append_dim3(oldDim)
6764
assert oldDim.x + 1 == newDim.x
6865
assert oldDim.y + 1 == newDim.y
6966
assert oldDim.z + 1 == newDim.z

‎cuda_core/cuda/core/_device.pyx‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1316,7 +1316,6 @@ class Device:
13161316
cdef object res
13171317
cdef SMResource sm_res
13181318
cdef WorkqueueResource wq_res
1319-
cdef GreenCtxHandle h_green
13201319

13211320
if options is None:
13221321
raise ValueError(
@@ -1348,7 +1347,7 @@ class Device:
13481347
else:
13491348
raise TypeError(f"Unsupported context resource type: {type(res)}")
13501349

1351-
h_green = create_green_ctx_handle(
1350+
cdef GreenCtxHandle h_green = create_green_ctx_handle(
13521351
c_resources.data(),
13531352
<unsigned int>(c_resources.size()),
13541353
<cydriver.CUdevice>(self._device_id),

‎cuda_core/cuda/core/_device_resources.pyx‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -250,7 +250,6 @@ cdef object _resolve_split_by_count_request(SMResourceOptions options):
250250
cdef list counts = _broadcast_field(options.count, n_groups)
251251
cdef object first = counts[0]
252252
cdef object value
253-
cdef unsigned int min_count
254253

255254
if options.coscheduled_sm_count is not None:
256255
raise RuntimeError(
@@ -270,7 +269,7 @@ cdef object _resolve_split_by_count_request(SMResourceOptions options):
270269
"use CUDA 13.1 or newer for per-group counts"
271270
)
272271

273-
min_count = _to_sm_count(first)
272+
cdef unsigned int min_count = _to_sm_count(first)
274273
return n_groups, min_count
275274

276275

‎cuda_core/cuda/core/_dlpack.pyx‎

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -115,10 +115,8 @@ cdef inline int setup_dl_tensor_device(DLTensor* dl_tensor, object buf) except -
115115

116116

117117
cdef inline int setup_dl_tensor_dtype(DLTensor* dl_tensor) except -1 nogil:
118-
cdef DLDataType* dtype = &dl_tensor.dtype
119-
dtype.code = <uint8_t>kDLInt
120-
dtype.lanes = <uint16_t>1
121-
dtype.bits = <uint8_t>8
118+
dl_tensor.dtype = DLDataType(
119+
code=<uint8_t>kDLInt, bits=<uint8_t>8, lanes=<uint16_t>1)
122120
return 0
123121

124122

‎cuda_core/cuda/core/_graphics.pyx‎

Lines changed: 8 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -209,10 +209,9 @@ cdef class GraphicsResource:
209209
return self
210210

211211
def _get_mapped_buffer(self) -> object:
212-
cdef Buffer buf
213212
if self._mapped_buffer is None:
214213
return None
215-
buf = <Buffer>self._mapped_buffer
214+
cdef Buffer buf = <Buffer>self._mapped_buffer
216215
if not buf._h_ptr:
217216
self._mapped_buffer = None
218217
return None
@@ -250,28 +249,24 @@ cdef class GraphicsResource:
250249
CUDAError
251250
If the mapping fails.
252251
"""
253-
cdef Stream s_obj
254-
cdef cydriver.CUgraphicsResource raw
255-
cdef cydriver.CUstream cy_stream
256252
cdef cydriver.CUdeviceptr dev_ptr = 0
257253
cdef size_t size = 0
258-
cdef Buffer buf
259254
if not self._handle:
260255
raise RuntimeError("GraphicsResource has been closed")
261256
if self._get_mapped_buffer() is not None:
262257
raise RuntimeError("GraphicsResource is already mapped")
263258

264-
s_obj = Stream_accept(stream)
265-
raw = as_cu(self._handle)
266-
cy_stream = as_cu(s_obj._h_stream)
259+
cdef Stream s_obj = Stream_accept(stream)
260+
cdef cydriver.CUgraphicsResource raw = as_cu(self._handle)
261+
cdef cydriver.CUstream cy_stream = as_cu(s_obj._h_stream)
267262
with nogil:
268263
HANDLE_RETURN(
269264
cydriver.cuGraphicsMapResources(1, &raw, cy_stream)
270265
)
271266
HANDLE_RETURN(
272267
cydriver.cuGraphicsResourceGetMappedPointer(&dev_ptr, &size, raw)
273268
)
274-
buf = Buffer_from_deviceptr_handle(
269+
cdef Buffer buf = Buffer_from_deviceptr_handle(
275270
deviceptr_create_mapped_graphics(dev_ptr, self._handle, s_obj._h_stream),
276271
size,
277272
None,
@@ -299,14 +294,12 @@ cdef class GraphicsResource:
299294
CUDAError
300295
If the unmapping fails.
301296
"""
302-
cdef object buf_obj
303-
cdef Buffer buf
304297
if not self._handle:
305298
raise RuntimeError("GraphicsResource has been closed")
306-
buf_obj = self._get_mapped_buffer()
299+
cdef object buf_obj = self._get_mapped_buffer()
307300
if buf_obj is None:
308301
raise RuntimeError("GraphicsResource is not mapped")
309-
buf = <Buffer>buf_obj
302+
cdef Buffer buf = <Buffer>buf_obj
310303
buf.close(stream=stream)
311304
self._mapped_buffer = None
312305

@@ -332,11 +325,10 @@ cdef class GraphicsResource:
332325
Optional override for the stream used to close the currently
333326
mapped buffer, if one exists.
334327
"""
335-
cdef object buf_obj
336328
cdef Buffer buf
337329
if not self._handle:
338330
return
339-
buf_obj = self._get_mapped_buffer()
331+
cdef object buf_obj = self._get_mapped_buffer()
340332
if buf_obj is not None:
341333
buf = <Buffer>buf_obj
342334
buf.close(stream=stream)

‎cuda_core/cuda/core/_linker.pyx‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -544,11 +544,10 @@ cdef inline void Linker_add_code_object(Linker self, object object_code) except
544544
cdef cydriver.CUjitInputType c_drv_input_type
545545
cdef const char* c_data_ptr
546546
cdef size_t c_data_size
547-
cdef const char* c_name_ptr
548547
cdef const char* c_file_ptr
549548

550549
name_bytes = f"{object_code.name}".encode()
551-
c_name_ptr = <const char*>name_bytes
550+
cdef const char* c_name_ptr = <const char*>name_bytes
552551

553552
input_types = _nvjitlink_input_types if self._use_nvjitlink else _driver_input_types
554553
py_input_type = input_types.get(object_code.code_type)

‎cuda_core/cuda/core/_memory/_buffer.pyx‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -422,8 +422,7 @@ cdef class Buffer:
422422
if not isinstance(max_version, tuple) or len(max_version) != 2:
423423
raise BufferError(f"Expected max_version tuple[int, int], got {max_version}")
424424
versioned = max_version >= (1, 0)
425-
capsule = make_py_capsule(self, versioned)
426-
return capsule
425+
return make_py_capsule(self, versioned)
427426

428427
def __dlpack_device__(self) -> tuple[int, int]:
429428
return classify_dl_device(self)

‎cuda_core/cuda/core/_memory/_device_memory_resource.pyx‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -321,10 +321,10 @@ cpdef str DMR_mempool_get_access(DeviceMemoryResource dmr, int device_id):
321321

322322
cdef int dev_id = Device(device_id).device_id
323323
cdef cydriver.CUmemAccess_flags flags
324-
cdef cydriver.CUmemLocation location
325-
326-
location.type = cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE
327-
location.id = dev_id
324+
cdef cydriver.CUmemLocation location = cydriver.CUmemLocation(
325+
type=cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE,
326+
id=dev_id,
327+
)
328328

329329
with nogil:
330330
HANDLE_RETURN(cydriver.cuMemPoolGetAccess(&flags, as_cu(dmr._h_pool), &location))

‎cuda_core/cuda/core/_memory/_managed_memory_ops.pyx‎

Lines changed: 16 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -80,36 +80,37 @@ cdef tuple _coerce_batch_buffers(object buffers, str what):
8080

8181

8282
cdef tuple _broadcast_locations(object location, Py_ssize_t n, bint allow_none, str what):
83-
cdef object coerced
8483
if isinstance(location, Sequence):
8584
if len(location) != n:
8685
raise ValueError(
8786
f"{what}: location length {len(location)} does not match "
8887
f"targets length {n}"
8988
)
9089
return tuple(_coerce_location(loc, allow_none=allow_none) for loc in location)
91-
coerced = _coerce_location(location, allow_none=allow_none)
90+
cdef object coerced = _coerce_location(location, allow_none=allow_none)
9291
return tuple([coerced] * n)
9392

9493

9594
IF CUDA_CORE_BUILD_MAJOR >= 13:
9695
# Convert a _LocSpec dataclass to a cydriver.CUmemLocation struct.
9796
cdef inline cydriver.CUmemLocation _to_cumemlocation(object loc):
98-
cdef cydriver.CUmemLocation out
9997
cdef str kind = loc.kind
10098
if kind == "device":
101-
out.type = cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE
102-
out.id = <int>loc.id
99+
return cydriver.CUmemLocation(
100+
type=cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE,
101+
id=<int>loc.id)
103102
elif kind == "host":
104-
out.type = cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST
105-
out.id = 0
103+
return cydriver.CUmemLocation(
104+
type=cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST,
105+
id=0)
106106
elif kind == "host_numa":
107-
out.type = cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST_NUMA
108-
out.id = <int>loc.id
107+
return cydriver.CUmemLocation(
108+
type=cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST_NUMA,
109+
id=<int>loc.id)
109110
else: # host_numa_current
110-
out.type = cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST_NUMA_CURRENT
111-
out.id = 0
112-
return out
111+
return cydriver.CUmemLocation(
112+
type=cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST_NUMA_CURRENT,
113+
id=0)
113114
ELSE:
114115
# CUDA 12 cuMemPrefetchAsync takes a device ordinal (-1 = host).
115116
cdef inline int _to_legacy_device(object loc) except? -2:
@@ -223,8 +224,9 @@ cdef void _do_single_advise(Buffer buf, object advice_value, object loc, bint al
223224
# Driver ignores location for read_mostly / unset_preferred_location
224225
# advice values but still validates the CUmemLocation; pass a
225226
# host placeholder.
226-
cu_loc.type = cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST
227-
cu_loc.id = 0
227+
cu_loc = cydriver.CUmemLocation(
228+
type=cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST,
229+
id=0)
228230
else:
229231
cu_loc = _to_cumemlocation(loc)
230232
with nogil:

‎cuda_core/cuda/core/_memory/_memory_pool.pyx‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -275,10 +275,9 @@ cdef int MP_init_current_pool(
275275
Requires CUDA 13+.
276276
"""
277277
IF CUDA_CORE_BUILD_MAJOR >= 13:
278-
cdef cydriver.CUmemLocation loc
279278
cdef cydriver.CUmemoryPool pool
280-
loc.id = loc_id
281-
loc.type = loc_type
279+
cdef cydriver.CUmemLocation loc = cydriver.CUmemLocation(
280+
type=loc_type, id=loc_id)
282281
with nogil:
283282
HANDLE_RETURN(cydriver.cuMemGetMemPool(&pool, &loc, alloc_type))
284283
self._h_pool = create_mempool_handle_ref(pool)

0 commit comments

Comments
 (0)