Coverage for cuda/core/_kernel_arg_handler.pyx: 77.23%
224 statements
« prev ^ index » next coverage.py v7.17.2, created at 2026-11-05 01:11 +0000
« prev ^ index » next coverage.py v7.17.2, created at 2026-11-05 01:11 +0000
1# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2#
3# SPDX-License-Identifier: Apache-2.0
5from cpython.mem cimport PyMem_Malloc, PyMem_Free
6from libc.stdint cimport (intptr_t,
7 int8_t, int16_t, int32_t, int64_t,
8 uint8_t, uint16_t, uint32_t, uint64_t,)
9from libcpp cimport bool as cpp_bool
10from libcpp.complex cimport complex as cpp_complex
11from libcpp cimport nullptr
12from libcpp cimport vector
14import ctypes
15from typing import Sequence, Any
17import numpy
19from cuda.core._memory import Buffer
20from cuda.core._memory._buffer cimport Buffer as cyBuffer, Buffer_check_open
21from cuda.core._tensor_map import TensorMapDescriptor as _TensorMapDescriptor_py
22from cuda.core._tensor_map cimport TensorMapDescriptor
23from cuda.core.graph._graph_definition cimport GraphCondition
24from cuda.core._utils.cuda_utils import driver
25from cuda.bindings cimport cydriver
28ctypedef cpp_complex.complex[float] cpp_single_complex
29ctypedef cpp_complex.complex[double] cpp_double_complex
32# We need an identifier for fp16 for copying scalars on the host. This is a minimal
33# implementation borrowed from cuda_fp16.h.
34cdef extern from *:
35 """
36 #if __cplusplus >= 201103L
37 #define __CUDA_ALIGN__(n) alignas(n) /* C++11 kindly gives us a keyword for this */
38 #else
39 #if defined(__GNUC__)
40 #define __CUDA_ALIGN__(n) __attribute__ ((aligned(n)))
41 #elif defined(_MSC_VER)
42 #define __CUDA_ALIGN__(n) __declspec(align(n))
43 #else
44 #define __CUDA_ALIGN__(n)
45 #endif /* defined(__GNUC__) */
46 #endif /* __cplusplus >= 201103L */
48 typedef struct __CUDA_ALIGN__(2) {
49 /**
50 * Storage field contains bits representation of the \p half floating-point number.
51 */
52 unsigned short x;
53 } __half_raw;
54 """
55 ctypedef struct __half_raw:
56 unsigned short x
59ctypedef fused supported_type:
60 cpp_bool
61 int8_t
62 int16_t
63 int32_t
64 int64_t
65 uint8_t
66 uint16_t
67 uint32_t
68 uint64_t
69 __half_raw
70 float
71 double
72 intptr_t
73 cpp_single_complex
74 cpp_double_complex
77# cache ctypes/numpy type objects to avoid attribute access
78cdef object ctypes_bool = ctypes.c_bool
79cdef object ctypes_int8 = ctypes.c_int8
80cdef object ctypes_int16 = ctypes.c_int16
81cdef object ctypes_int32 = ctypes.c_int32
82cdef object ctypes_int64 = ctypes.c_int64
83cdef object ctypes_uint8 = ctypes.c_uint8
84cdef object ctypes_uint16 = ctypes.c_uint16
85cdef object ctypes_uint32 = ctypes.c_uint32
86cdef object ctypes_uint64 = ctypes.c_uint64
87cdef object ctypes_float = ctypes.c_float
88cdef object ctypes_double = ctypes.c_double
89cdef object numpy_bool = numpy.bool_
90cdef object numpy_int8 = numpy.int8
91cdef object numpy_int16 = numpy.int16
92cdef object numpy_int32 = numpy.int32
93cdef object numpy_int64 = numpy.int64
94cdef object numpy_uint8 = numpy.uint8
95cdef object numpy_uint16 = numpy.uint16
96cdef object numpy_uint32 = numpy.uint32
97cdef object numpy_uint64 = numpy.uint64
98cdef object numpy_float16 = numpy.float16
99cdef object numpy_float32 = numpy.float32
100cdef object numpy_float64 = numpy.float64
101cdef object numpy_complex64 = numpy.complex64
102cdef object numpy_complex128 = numpy.complex128
105cdef object tensor_map_descriptor_type = _TensorMapDescriptor_py
108# limitation due to cython/cython#534
109ctypedef void* voidptr
112# Cython can't infer the overload without at least one input argument with fused type
113cdef inline int prepare_arg(
114 vector.vector[void*]& data,
115 vector.vector[void*]& data_addresses,
116 arg, # important: keep it a Python object and don't cast
117 const size_t idx,
118 const supported_type* __unused=NULL) except -1:
119 cdef void* ptr = PyMem_Malloc(sizeof(supported_type)) 2zbw x AbBbCby z d e m n o p A B f g q r s t C D E F G H h i j k DbEbFbGbHbIbJbKbkblbmbnbobpbqbrbsbtbubLbMbNbwbObPbQbRbSbTbUbVbxba WbXbYbZb0b1b2bl 3b4b5b6b7b~ abbb8bL ) b u v I c hb$ # ' % ; : / . - , + * ybib! 9 W X 6 7 4 5 8 3 0 1 Y Z 2 ? = jb^ _ ] vbgbfbebdbcb} | { V J M U T S R Q P O N ( 9b!b#b$b%b@ [ 'b`
120 # note: this should also work once ctypes has complex support:
121 # python/cpython#121248
122 if supported_type is cpp_single_complex:
123 (<supported_type*>ptr)[0] = cpp_complex.complex[float](arg.real, arg.imag) 1$?
124 elif supported_type is cpp_double_complex:
125 (<supported_type*>ptr)[0] = cpp_complex.complex[double](arg.real, arg.imag) 2b # = jb
126 elif supported_type is __half_raw:
127 (<supported_type*>ptr).x = <int16_t>(arg.view(numpy_int16)) 1'^
128 else:
129 (<supported_type*>ptr)[0] = <supported_type>(arg) 2zbw x AbBbCby z d e m n o p A B f g q r s t C D E F G H h i j k DbEbFbGbHbIbJbKbkblbmbnbobpbqbrbsbtbubLbMbNbwbObPbQbRbSbTbUbVbxba WbXbYbZb0b1b2bl 3b4b5b6b7b~ abbb8bL ) b u v I c hb$ # ' % ; : / . - , + * ybib! 9 W X 6 7 4 5 8 3 0 1 Y Z 2 ? = jb^ _ ] vbgbfbebdbcb} | { V J M U T S R Q P O N ( 9b!b#b$b%b@ [ 'b`
130 data_addresses[idx] = ptr # take the address to the scalar 2zbw x AbBbCby z d e m n o p A B f g q r s t C D E F G H h i j k DbEbFbGbHbIbJbKbkblbmbnbobpbqbrbsbtbubLbMbNbwbObPbQbRbSbTbUbVbxba WbXbYbZb0b1b2bl 3b4b5b6b7b~ abbb8bL ) b u v I c hb$ # ' % ; : / . - , + * ybib! 9 W X 6 7 4 5 8 3 0 1 Y Z 2 ? = jb^ _ ] vbgbfbebdbcb} | { V J M U T S R Q P O N ( 9b!b#b$b%b@ [ 'b`
131 data[idx] = ptr # for later dealloc 2zbw x AbBbCby z d e m n o p A B f g q r s t C D E F G H h i j k DbEbFbGbHbIbJbKbkblbmbnbobpbqbrbsbtbubLbMbNbwbObPbQbRbSbTbUbVbxba WbXbYbZb0b1b2bl 3b4b5b6b7b~ abbb8bL ) b u v I c hb$ # ' % ; : / . - , + * ybib! 9 W X 6 7 4 5 8 3 0 1 Y Z 2 ? = jb^ _ ] vbgbfbebdbcb} | { V J M U T S R Q P O N ( 9b!b#b$b%b@ [ 'b`
132 return 0 2zbw x AbBbCby z d e m n o p A B f g q r s t C D E F G H h i j k DbEbFbGbHbIbJbKbkblbmbnbobpbqbrbsbtbubLbMbNbwbObPbQbRbSbTbUbVbxba WbXbYbZb0b1b2bl 3b4b5b6b7b~ abbb8bL ) b u v I c hb$ # ' % ; : / . - , + * ybib! 9 W X 6 7 4 5 8 3 0 1 Y Z 2 ? = jb^ _ ] vbgbfbebdbcb} | { V J M U T S R Q P O N ( 9b!b#b$b%b@ [ 'b`
135cdef inline int prepare_tensor_map_arg(
136 vector.vector[void*]& data,
137 vector.vector[void*]& data_addresses,
138 TensorMapDescriptor arg,
139 const size_t idx) except -1:
140 # cuLaunchKernel copies argument bytes during launch, so a TensorMap
141 # descriptor can point directly at its internal CUtensorMap storage.
142 data_addresses[idx] = arg._get_data_ptr()
143 return 0
146cdef inline int prepare_ctypes_arg(
147 vector.vector[void*]& data,
148 vector.vector[void*]& data_addresses,
149 arg,
150 const size_t idx) except -1:
151 cdef object arg_type = type(arg) 1wxyzdemnopABfgqrstCDEFGHhijkalLbKuvIc!9WX67458301YZ2VJMUTSRQPON
152 if arg_type is ctypes_bool: 1wxyzdemnopABfgqrstCDEFGHhijkalLbKuvIc!9WX67458301YZ2VJMUTSRQPON
153 return prepare_arg[cpp_bool](data, data_addresses, arg.value, idx) 1mnqr!
154 elif arg_type is ctypes_int8: 1wxyzdemnopABfgqrstCDEFGHhijkalLbKuvIc9WX67458301YZ2VJMUTSRQPON
155 return prepare_arg[int8_t](data, data_addresses, arg.value, idx) 19
156 elif arg_type is ctypes_int16: 1wxyzdemnopABfgqrstCDEFGHhijkalLbKuvIcWX67458301YZ2VJMUTSRQPON
157 return prepare_arg[int16_t](data, data_addresses, arg.value, idx) 18
158 elif arg_type is ctypes_int32: 1wxyzdemnopABfgqrstCDEFGHhijkalLbKuvIcWX6745301YZ2VJMUTSRQPON
159 return prepare_arg[int32_t](data, data_addresses, arg.value, idx) 167
160 elif arg_type is ctypes_int64: 1wxyzdemnopABfgqrstCDEFGHhijkalLbKuvIcWX45301YZ2VJMUTSRQPON
161 return prepare_arg[int64_t](data, data_addresses, arg.value, idx) 145
162 elif arg_type is ctypes_uint8: 1wxyzdemnopABfgqrstCDEFGHhijkalLbKuvIcWX301YZ2VJMUTSRQPON
163 return prepare_arg[uint8_t](data, data_addresses, arg.value, idx) 13
164 elif arg_type is ctypes_uint16: 1wxyzdemnopABfgqrstCDEFGHhijkalLbKuvIcWX01YZ2VJMUTSRQPON
165 return prepare_arg[uint16_t](data, data_addresses, arg.value, idx) 12
166 elif arg_type is ctypes_uint32: 1wxyzdemnopABfgqrstCDEFGHhijkalLbKuvIcWX01YZVJMUTSRQPON
167 return prepare_arg[uint32_t](data, data_addresses, arg.value, idx) 101
168 elif arg_type is ctypes_uint64: 1wxyzdemnopABfgqrstCDEFGHhijkalLbKuvIcWXYZVJMUTSRQPON
169 return prepare_arg[uint64_t](data, data_addresses, arg.value, idx) 1YZ
170 elif arg_type is ctypes_float: 1wxyzdemnopABfgqrstCDEFGHhijkalLbKuvIcWXVJMUTSRQPON
171 return prepare_arg[float](data, data_addresses, arg.value, idx) 1X
172 elif arg_type is ctypes_double: 1wxyzdemnopABfgqrstCDEFGHhijkalLbKuvIcWVJMUTSRQPON
173 return prepare_arg[double](data, data_addresses, arg.value, idx) 1W
174 else:
175 # If no exact types are found, fallback to slower `isinstance` check
176 if isinstance(arg, ctypes_bool): 1wxyzdemnopABfgqrstCDEFGHhijkalLbKuvIcVJMUTSRQPON
177 return prepare_arg[cpp_bool](data, data_addresses, arg.value, idx) 1LV
178 elif isinstance(arg, ctypes_int8): 1wxyzdemnopABfgqrstCDEFGHhijkalLbKuvIcJMUTSRQPON
179 return prepare_arg[int8_t](data, data_addresses, arg.value, idx) 1U
180 elif isinstance(arg, ctypes_int16): 1wxyzdemnopABfgqrstCDEFGHhijkalLbKuvIcJMTSRQPON
181 return prepare_arg[int16_t](data, data_addresses, arg.value, idx) 1T
182 elif isinstance(arg, ctypes_int32): 1wxyzdemnopABfgqrstCDEFGHhijkalLbKuvIcJMSRQPON
183 return prepare_arg[int32_t](data, data_addresses, arg.value, idx) 1LS
184 elif isinstance(arg, ctypes_int64): 1wxyzdemnopABfgqrstCDEFGHhijkalLbKuvIcJMRQPON
185 return prepare_arg[int64_t](data, data_addresses, arg.value, idx) 1R
186 elif isinstance(arg, ctypes_uint8): 1wxyzdemnopABfgqrstCDEFGHhijkalLbKuvIcJMQPON
187 return prepare_arg[uint8_t](data, data_addresses, arg.value, idx) 1Q
188 elif isinstance(arg, ctypes_uint16): 1wxyzdemnopABfgqrstCDEFGHhijkalLbKuvIcJMPON
189 return prepare_arg[uint16_t](data, data_addresses, arg.value, idx) 1P
190 elif isinstance(arg, ctypes_uint32): 1wxyzdemnopABfgqrstCDEFGHhijkalLbKuvIcJMON
191 return prepare_arg[uint32_t](data, data_addresses, arg.value, idx) 1O
192 elif isinstance(arg, ctypes_uint64): 1wxyzdemnopABfgqrstCDEFGHhijkalLbKuvIcJMN
193 return prepare_arg[uint64_t](data, data_addresses, arg.value, idx) 1N
194 elif isinstance(arg, ctypes_float): 1wxyzdemnopABfgqrstCDEFGHhijkalLbKuvIcJM
195 return prepare_arg[float](data, data_addresses, arg.value, idx) 1LM
196 elif isinstance(arg, ctypes_double): 1wxyzdemnopABfgqrstCDEFGHhijkalbKuvIcJ
197 return prepare_arg[double](data, data_addresses, arg.value, idx) 1J
198 else:
199 return 1 1wxyzdemnopABfgqrstCDEFGHhijkalbKuvIc
202cdef inline int prepare_numpy_arg(
203 vector.vector[void*]& data,
204 vector.vector[void*]& data_addresses,
205 arg,
206 const size_t idx) except -1:
207 cdef object arg_type = type(arg) 2w x y z d e m n o p A B f g q r s t C D E F G H h i j k a l ~ abbbL ) b K u v I c hb$ # ' % ; : / . - , + * ib! 9 W X 6 7 4 5 8 3 0 1 Y Z 2 ? = ^ _ ] gbfbebdbcb} | { V J M U T S R Q P O N ( @ [ `
208 if arg_type is numpy_bool: 2w x y z d e m n o p A B f g q r s t C D E F G H h i j k a l ~ abbbL ) b K u v I c hb$ # ' % ; : / . - , + * ib! 9 W X 6 7 4 5 8 3 0 1 Y Z 2 ? = ^ _ ] gbfbebdbcb} | { V J M U T S R Q P O N ( @ [ `
209 return prepare_arg[cpp_bool](data, data_addresses, arg, idx) 2o p s t hbib
210 elif arg_type is numpy_int8: 2w x y z d e m n o p A B f g q r s t C D E F G H h i j k a l ~ abbbL ) b K u v I c $ # ' % ; : / . - , + * ! 9 W X 6 7 4 5 8 3 0 1 Y Z 2 ? = ^ _ ] gbfbebdbcb} | { V J M U T S R Q P O N ( @ [ `
211 return prepare_arg[int8_t](data, data_addresses, arg, idx) 2gb
212 elif arg_type is numpy_int16: 2w x y z d e m n o p A B f g q r s t C D E F G H h i j k a l ~ abbbL ) b K u v I c $ # ' % ; : / . - , + * ! 9 W X 6 7 4 5 8 3 0 1 Y Z 2 ? = ^ _ ] fbebdbcb} | { V J M U T S R Q P O N ( @ [ `
213 return prepare_arg[int16_t](data, data_addresses, arg, idx) 2fb
214 elif arg_type is numpy_int32: 2w x y z d e m n o p A B f g q r s t C D E F G H h i j k a l ~ abbbL ) b K u v I c $ # ' % ; : / . - , + * ! 9 W X 6 7 4 5 8 3 0 1 Y Z 2 ? = ^ _ ] ebdbcb} | { V J M U T S R Q P O N ( @ [ `
215 return prepare_arg[int32_t](data, data_addresses, arg, idx) 2~ abbbu v eb
216 elif arg_type is numpy_int64: 2w x y z d e m n o p A B f g q r s t C D E F G H h i j k a l L ) b K u v I c $ # ' % ; : / . - , + * ! 9 W X 6 7 4 5 8 3 0 1 Y Z 2 ? = ^ _ ] dbcb} | { V J M U T S R Q P O N ( @ [ `
217 return prepare_arg[int64_t](data, data_addresses, arg, idx) 2db
218 elif arg_type is numpy_uint8: 2w x y z d e m n o p A B f g q r s t C D E F G H h i j k a l L ) b K u v I c $ # ' % ; : / . - , + * ! 9 W X 6 7 4 5 8 3 0 1 Y Z 2 ? = ^ _ ] cb} | { V J M U T S R Q P O N ( @ [ `
219 return prepare_arg[uint8_t](data, data_addresses, arg, idx) 2cb
220 elif arg_type is numpy_uint16: 1wxyzdemnopABfgqrstCDEFGHhijkalL)bKuvIc$#'%;:/.-,+*!9WX67458301YZ2?=^_]}|{VJMUTSRQPON(@[`
221 return prepare_arg[uint16_t](data, data_addresses, arg, idx) 1}
222 elif arg_type is numpy_uint32: 1wxyzdemnopABfgqrstCDEFGHhijkalL)bKuvIc$#'%;:/.-,+*!9WX67458301YZ2?=^_]|{VJMUTSRQPON(@[`
223 return prepare_arg[uint32_t](data, data_addresses, arg, idx) 1|
224 elif arg_type is numpy_uint64: 1wxyzdemnopABfgqrstCDEFGHhijkalL)bKuvIc$#'%;:/.-,+*!9WX67458301YZ2?=^_]{VJMUTSRQPON(@[`
225 return prepare_arg[uint64_t](data, data_addresses, arg, idx) 1{`
226 elif arg_type is numpy_float16: 1wxyzdemnopABfgqrstCDEFGHhijkalL)bKuvIc$#'%;:/.-,+*!9WX67458301YZ2?=^_]VJMUTSRQPON(@[
227 return prepare_arg[__half_raw](data, data_addresses, arg, idx) 1^
228 elif arg_type is numpy_float32: 1wxyzdemnopABfgqrstCDEFGHhijkalL)bKuvIc$#'%;:/.-,+*!9WX67458301YZ2?=_]VJMUTSRQPON(@[
229 return prepare_arg[float](data, data_addresses, arg, idx) 1_@[
230 elif arg_type is numpy_float64: 1wxyzdemnopABfgqrstCDEFGHhijkalL)bKuvIc$#'%;:/.-,+*!9WX67458301YZ2?=]VJMUTSRQPON(
231 return prepare_arg[double](data, data_addresses, arg, idx) 1]
232 elif arg_type is numpy_complex64: 1wxyzdemnopABfgqrstCDEFGHhijkalL)bKuvIc$#'%;:/.-,+*!9WX67458301YZ2?=VJMUTSRQPON(
233 return prepare_arg[cpp_single_complex](data, data_addresses, arg, idx) 1?
234 elif arg_type is numpy_complex128: 1wxyzdemnopABfgqrstCDEFGHhijkalL)bKuvIc$#'%;:/.-,+*!9WX67458301YZ2=VJMUTSRQPON(
235 return prepare_arg[cpp_double_complex](data, data_addresses, arg, idx) 1=
236 else:
237 # If no exact types are found, fallback to slower `isinstance` check
238 if isinstance(arg, numpy_bool): 1wxyzdemnopABfgqrstCDEFGHhijkalL)bKuvIc$#'%;:/.-,+*!9WX67458301YZ2VJMUTSRQPON(
239 return prepare_arg[cpp_bool](data, data_addresses, arg, idx)
240 elif isinstance(arg, numpy_int8): 1wxyzdemnopABfgqrstCDEFGHhijkalL)bKuvIc$#'%;:/.-,+*!9WX67458301YZ2VJMUTSRQPON(
241 return prepare_arg[int8_t](data, data_addresses, arg, idx) 1;
242 elif isinstance(arg, numpy_int16): 1wxyzdemnopABfgqrstCDEFGHhijkalL)bKuvIc$#'%:/.-,+*!9WX67458301YZ2VJMUTSRQPON(
243 return prepare_arg[int16_t](data, data_addresses, arg, idx) 1:
244 elif isinstance(arg, numpy_int32): 1wxyzdemnopABfgqrstCDEFGHhijkalL)bKuvIc$#'%/.-,+*!9WX67458301YZ2VJMUTSRQPON(
245 return prepare_arg[int32_t](data, data_addresses, arg, idx) 1)/
246 elif isinstance(arg, numpy_int64): 1wxyzdemnopABfgqrstCDEFGHhijkalL)bKuvIc$#'%.-,+*!9WX67458301YZ2VJMUTSRQPON(
247 return prepare_arg[int64_t](data, data_addresses, arg, idx) 1.
248 elif isinstance(arg, numpy_uint8): 1wxyzdemnopABfgqrstCDEFGHhijkalL)bKuvIc$#'%-,+*!9WX67458301YZ2VJMUTSRQPON(
249 return prepare_arg[uint8_t](data, data_addresses, arg, idx) 1-
250 elif isinstance(arg, numpy_uint16): 1wxyzdemnopABfgqrstCDEFGHhijkalL)bKuvIc$#'%,+*!9WX67458301YZ2VJMUTSRQPON(
251 return prepare_arg[uint16_t](data, data_addresses, arg, idx) 1,
252 elif isinstance(arg, numpy_uint32): 1wxyzdemnopABfgqrstCDEFGHhijkalL)bKuvIc$#'%+*!9WX67458301YZ2VJMUTSRQPON(
253 return prepare_arg[uint32_t](data, data_addresses, arg, idx) 1+
254 elif isinstance(arg, numpy_uint64): 1wxyzdemnopABfgqrstCDEFGHhijkalL)bKuvIc$#'%*!9WX67458301YZ2VJMUTSRQPON(
255 return prepare_arg[uint64_t](data, data_addresses, arg, idx) 1*
256 elif isinstance(arg, numpy_float16): 1wxyzdemnopABfgqrstCDEFGHhijkalL)bKuvIc$#'%!9WX67458301YZ2VJMUTSRQPON(
257 return prepare_arg[__half_raw](data, data_addresses, arg, idx) 1'
258 elif isinstance(arg, numpy_float32): 1wxyzdemnopABfgqrstCDEFGHhijkalL)bKuvIc$#%!9WX67458301YZ2VJMUTSRQPON(
259 return prepare_arg[float](data, data_addresses, arg, idx) 1)(
260 elif isinstance(arg, numpy_float64): 1wxyzdemnopABfgqrstCDEFGHhijkalLbKuvIc$#%!9WX67458301YZ2VJMUTSRQPON
261 return prepare_arg[double](data, data_addresses, arg, idx) 1%
262 elif isinstance(arg, numpy_complex64): 1wxyzdemnopABfgqrstCDEFGHhijkalLbKuvIc$#!9WX67458301YZ2VJMUTSRQPON
263 return prepare_arg[cpp_single_complex](data, data_addresses, arg, idx) 1$
264 elif isinstance(arg, numpy_complex128): 1wxyzdemnopABfgqrstCDEFGHhijkalLbKuvIc#!9WX67458301YZ2VJMUTSRQPON
265 return prepare_arg[cpp_double_complex](data, data_addresses, arg, idx) 1#
266 else:
267 return 1 1wxyzdemnopABfgqrstCDEFGHhijkalLbKuvIc!9WX67458301YZ2VJMUTSRQPON
270cdef class ParamHolder:
272 def __init__(self, kernel_args: Sequence[Any]) -> None:
273 if len(kernel_args) == 0: 2(bzb)b*bw x +b,b-bAb.b/b:b;b=bBb?b@b[bCby z d e m n o p A B f g q r s t C D E F G H h i j k ]b^b_b`b{b|b}b~bacbcccdcecfcgchcicjckclcmcncocpcqcrcscDbEbtcucvcwcxcFbGbHbIbJbKbkblbmbnbobpbqbrbsbtbubLbycMbNbwbObPbQbRbSbTbzcAcBcCcDcEcFcGcUbVbxbHcIca WbXbYbJcKcZb0b1b2bLcMcNcOcPcQcl 3b4b5b6b7bRc~ abbb8bScL ) b K u v I c Tchb$ # ' % ; : / . - , + * ybib! 9 W X 6 7 4 5 8 3 0 1 Y Z 2 ? = jb^ _ ] vbgbfbebdbcb} | { V J M U T S R Q P O N ( 9b!b#b$bUcVcWc%b@ [ XcYcZc0c1c2c3c4c5c6c7c8c9c!c#c$c%c'c(c)c*c+c,c'b`
274 self.ptr = 0 2(b)b*b+b,b-b.b/b:b;b=b?b@b[bh i j k ]b^b_b`b{b|b}b~bacbcccdcecfcgchcicjckclcmcncocpcqcrcsctcucvcwcxcycwbzcAcBcCcDcEcFcGcxbHcIcJcKcLcMcNcOcPcQcl RcScTcUcVcWcXcYcZc0c1c2c3c4c5c6c7c8c9c!c#c$c%c'c(c)c*c+c,c
275 return 2(b)b*b+b,b-b.b/b:b;b=b?b@b[bh i j k ]b^b_b`b{b|b}b~bacbcccdcecfcgchcicjckclcmcncocpcqcrcsctcucvcwcxcycwbzcAcBcCcDcEcFcGcxbHcIcJcKcLcMcNcOcPcQcl RcScTcUcVcWcXcYcZc0c1c2c3c4c5c6c7c8c9c!c#c$c%c'c(c)c*c+c,c
277 cdef size_t n_args = len(kernel_args) 2zbw x AbBbCby z d e m n o p A B f g q r s t C D E F G H h i j k DbEbFbGbHbIbJbKbkblbmbnbobpbqbrbsbtbubLbMbNbwbObPbQbRbSbTbUbVbxba WbXbYbZb0b1b2bl 3b4b5b6b7b~ abbb8bL ) b K u v I c hb$ # ' % ; : / . - , + * ybib! 9 W X 6 7 4 5 8 3 0 1 Y Z 2 ? = jb^ _ ] vbgbfbebdbcb} | { V J M U T S R Q P O N ( 9b!b#b$b%b@ [ 'b`
278 cdef size_t i
279 cdef int not_prepared
280 cdef object arg_type
281 self.data = vector.vector[voidptr](n_args, nullptr) 2zbw x AbBbCby z d e m n o p A B f g q r s t C D E F G H h i j k DbEbFbGbHbIbJbKbkblbmbnbobpbqbrbsbtbubLbMbNbwbObPbQbRbSbTbUbVbxba WbXbYbZb0b1b2bl 3b4b5b6b7b~ abbb8bL ) b K u v I c hb$ # ' % ; : / . - , + * ybib! 9 W X 6 7 4 5 8 3 0 1 Y Z 2 ? = jb^ _ ] vbgbfbebdbcb} | { V J M U T S R Q P O N ( 9b!b#b$b%b@ [ 'b`
282 self.data_addresses = vector.vector[voidptr](n_args) 2zbw x AbBbCby z d e m n o p A B f g q r s t C D E F G H h i j k DbEbFbGbHbIbJbKbkblbmbnbobpbqbrbsbtbubLbMbNbwbObPbQbRbSbTbUbVbxba WbXbYbZb0b1b2bl 3b4b5b6b7b~ abbb8bL ) b K u v I c hb$ # ' % ; : / . - , + * ybib! 9 W X 6 7 4 5 8 3 0 1 Y Z 2 ? = jb^ _ ] vbgbfbebdbcb} | { V J M U T S R Q P O N ( 9b!b#b$b%b@ [ 'b`
283 for i, arg in enumerate(kernel_args): 2zbw x AbBbCby z d e m n o p A B f g q r s t C D E F G H h i j k DbEbFbGbHbIbJbKbkblbmbnbobpbqbrbsbtbubLbMbNbwbObPbQbRbSbTbUbVbxba WbXbYbZb0b1b2bl 3b4b5b6b7b~ abbb8bL ) b K u v I c hb$ # ' % ; : / . - , + * ybib! 9 W X 6 7 4 5 8 3 0 1 Y Z 2 ? = jb^ _ ] vbgbfbebdbcb} | { V J M U T S R Q P O N ( 9b!b#b$b%b@ [ 'b`
284 arg_type = type(arg) 2zbw x AbBbCby z d e m n o p A B f g q r s t C D E F G H h i j k DbEbFbGbHbIbJbKbkblbmbnbobpbqbrbsbtbubLbMbNbwbObPbQbRbSbTbUbVbxba WbXbYbZb0b1b2bl 3b4b5b6b7b~ abbb8bL ) b K u v I c hb$ # ' % ; : / . - , + * ybib! 9 W X 6 7 4 5 8 3 0 1 Y Z 2 ? = jb^ _ ] vbgbfbebdbcb} | { V J M U T S R Q P O N ( 9b!b#b$b%b@ [ 'b`
285 if arg_type is Buffer: 2zbw x AbBbCby z d e m n o p A B f g q r s t C D E F G H h i j k DbEbFbGbHbIbJbKbkblbmbnbobpbqbrbsbtbubLbMbNbwbObPbQbRbSbTbUbVbxba WbXbYbZb0b1b2bl 3b4b5b6b7b~ abbb8bL ) b K u v I c hb$ # ' % ; : / . - , + * ybib! 9 W X 6 7 4 5 8 3 0 1 Y Z 2 ? = jb^ _ ] vbgbfbebdbcb} | { V J M U T S R Q P O N ( 9b!b#b$b%b@ [ 'b`
286 Buffer_check_open(<cyBuffer>arg) 2DbEbkblbmbnbobpbqbrbsbtbubLbMbNbwbObPbQbRbSbTbUbVbxba WbXbYbZb~ abbb9b!b#b$b%b@ [ 'b`
287 # we need the address of where the actual buffer address is stored
288 if type(arg.handle) is int: 2DbEbkblbmbnbobpbqbrbsbtbubLbMbNbwbObPbQbRbSbTbUbVbxba WbXbYbZb~ abbb9b!b#b$b%b@ [ 'b`
289 # see note below on handling int arguments
290 prepare_arg[intptr_t](self.data, self.data_addresses, arg.handle, i) 2DbEbkblbmbnbobpbqbrbsbtbubLbMbNbwbObPbQbRbSbTbUbVbxba WbXbYbZb~ abbb9b!b#b$b%b@ [ 'b`
291 continue 2DbEbkblbmbnbobpbqbrbsbtbubLbMbNbwbObPbQbRbSbTbUbVbxba WbXbYbZb~ abbb9b!b#b$b%b@ [ 'b`
292 else:
293 # it's a CUdeviceptr:
294 self.data_addresses[i] = <void*><intptr_t>(arg.handle.getPtr())
295 continue
296 elif arg_type is bool: 2zbw x AbBbCby z d e m n o p A B f g q r s t C D E F G H h i j k FbGbHbIbJbKbkblbmbnbobpbqbrbsbtbuba 0b1b2bl 3b4b5b6b7b~ abbb8bL ) b K u v I c hb$ # ' % ; : / . - , + * ybib! 9 W X 6 7 4 5 8 3 0 1 Y Z 2 ? = jb^ _ ] vbgbfbebdbcb} | { V J M U T S R Q P O N ( @ [ `
297 prepare_arg[cpp_bool](self.data, self.data_addresses, arg, i) 2d e f g l yb
298 continue 2d e f g l yb
299 elif arg_type is int: 2zbw x AbBbCby z d e m n o p A B f g q r s t C D E F G H h i j k FbGbHbIbJbKbkblbmbnbobpbqbrbsbtbuba 0b1b2bl 3b4b5b6b7b~ abbb8bL ) b K u v I c hb$ # ' % ; : / . - , + * ybib! 9 W X 6 7 4 5 8 3 0 1 Y Z 2 ? = jb^ _ ] vbgbfbebdbcb} | { V J M U T S R Q P O N ( @ [ `
300 # Here's the dilemma: We want to have a fast path to pass in Python
301 # integers as pointer addresses, but one could also (mistakenly) pass
302 # it with the intention of passing a scalar integer. It's a mistake
303 # bacause a Python int is ambiguous (arbitrary width). Our judgement
304 # call here is to treat it as a pointer address, without any warning!
305 prepare_arg[intptr_t](self.data, self.data_addresses, arg, i) 2zbw x AbBbCby z d e m n o p A B f g q r s t C D E F G H h i j k FbGbHbIbJbKbkblbmbnbobpbqbrbsbtbub0b1b2b3b4b5b6b7b8bI c hb$ # ' % ; : / . - , + * ybib! 9 W X 6 7 4 5 8 3 0 1 Y Z 2 ? = jb^ _ ] vbgbfbebdbcb} | { V J M U T S R Q P O N (
306 continue 2zbw x AbBbCby z d e m n o p A B f g q r s t C D E F G H h i j k FbGbHbIbJbKbkblbmbnbobpbqbrbsbtbub0b1b2b3b4b5b6b7b8bI c hb$ # ' % ; : / . - , + * ybib! 9 W X 6 7 4 5 8 3 0 1 Y Z 2 ? = jb^ _ ] vbgbfbebdbcb} | { V J M U T S R Q P O N (
307 elif arg_type is float: 2w x y z d e m n o p A B f g q r s t C D E F G H h i j k a l ~ abbbL ) b K u v I c hb$ # ' % ; : / . - , + * ib! 9 W X 6 7 4 5 8 3 0 1 Y Z 2 ? = jb^ _ ] vbgbfbebdbcb} | { V J M U T S R Q P O N ( @ [ `
308 prepare_arg[double](self.data, self.data_addresses, arg, i) 2vb
309 continue 2vb
310 elif arg_type is complex: 2w x y z d e m n o p A B f g q r s t C D E F G H h i j k a l ~ abbbL ) b K u v I c hb$ # ' % ; : / . - , + * ib! 9 W X 6 7 4 5 8 3 0 1 Y Z 2 ? = jb^ _ ] gbfbebdbcb} | { V J M U T S R Q P O N ( @ [ `
311 prepare_arg[cpp_double_complex](self.data, self.data_addresses, arg, i) 2jb
312 continue 2jb
313 elif arg_type is tensor_map_descriptor_type: 2w x y z d e m n o p A B f g q r s t C D E F G H h i j k a l ~ abbbL ) b K u v I c hb$ # ' % ; : / . - , + * ib! 9 W X 6 7 4 5 8 3 0 1 Y Z 2 ? = ^ _ ] gbfbebdbcb} | { V J M U T S R Q P O N ( @ [ `
314 prepare_tensor_map_arg(self.data, self.data_addresses, <TensorMapDescriptor>arg, i)
315 continue
317 not_prepared = prepare_numpy_arg(self.data, self.data_addresses, arg, i) 2w x y z d e m n o p A B f g q r s t C D E F G H h i j k a l ~ abbbL ) b K u v I c hb$ # ' % ; : / . - , + * ib! 9 W X 6 7 4 5 8 3 0 1 Y Z 2 ? = ^ _ ] gbfbebdbcb} | { V J M U T S R Q P O N ( @ [ `
318 if not_prepared: 2w x y z d e m n o p A B f g q r s t C D E F G H h i j k a l ~ abbbL ) b K u v I c hb$ # ' % ; : / . - , + * ib! 9 W X 6 7 4 5 8 3 0 1 Y Z 2 ? = ^ _ ] gbfbebdbcb} | { V J M U T S R Q P O N ( @ [ `
319 not_prepared = prepare_ctypes_arg(self.data, self.data_addresses, arg, i) 1wxyzdemnopABfgqrstCDEFGHhijkalLbKuvIc!9WX67458301YZ2VJMUTSRQPON
320 if not_prepared: 2w x y z d e m n o p A B f g q r s t C D E F G H h i j k a l ~ abbbL ) b K u v I c hb$ # ' % ; : / . - , + * ib! 9 W X 6 7 4 5 8 3 0 1 Y Z 2 ? = ^ _ ] gbfbebdbcb} | { V J M U T S R Q P O N ( @ [ `
321 # TODO: revisit this treatment if we decide to cythonize cuda.core
322 if arg_type is driver.CUgraphConditionalHandle: 1wxyzdemnopABfgqrstCDEFGHhijkalbKuvIc
323 prepare_arg[cydriver.CUgraphConditionalHandle](self.data, self.data_addresses, <intptr_t>int(arg), i) 1I
324 continue 1I
325 elif arg_type is GraphCondition: 1wxyzdemnopABfgqrstCDEFGHhijkalbKuvc
326 prepare_arg[cydriver.CUgraphConditionalHandle]( 1wxyzdemnopABfgqrstCDEFGHhijkl
327 self.data, self.data_addresses,
328 <intptr_t><unsigned long long>(<GraphCondition>arg)._c_handle, i) 1wxyzdemnopABfgqrstCDEFGHhijkl
329 continue 1wxyzdemnopABfgqrstCDEFGHhijkl
330 # If no exact types are found, fallback to slower `isinstance` check
331 elif isinstance(arg, Buffer): 1abKuvc
332 Buffer_check_open(<cyBuffer>arg) 1uv
333 if isinstance(arg.handle, int): 1uv
334 prepare_arg[intptr_t](self.data, self.data_addresses, arg.handle, i) 1uv
335 continue 1uv
336 else:
337 self.data_addresses[i] = <void*><intptr_t>(arg.handle.getPtr())
338 continue
339 elif isinstance(arg, bool): 1abKc
340 prepare_arg[cpp_bool](self.data, self.data_addresses, arg, i)
341 continue
342 elif isinstance(arg, int): 1abKc
343 prepare_arg[intptr_t](self.data, self.data_addresses, arg, i) 1b
344 continue 1b
345 elif isinstance(arg, float): 1abKc
346 prepare_arg[double](self.data, self.data_addresses, arg, i) 1b
347 continue 1b
348 elif isinstance(arg, complex): 1abKc
349 prepare_arg[cpp_double_complex](self.data, self.data_addresses, arg, i) 1b
350 continue 1b
351 elif isinstance(arg, driver.CUgraphConditionalHandle): 1aKc
352 prepare_arg[cydriver.CUgraphConditionalHandle](self.data, self.data_addresses, arg, i) 1c
353 continue 1c
354 elif isinstance(arg, GraphCondition): 1aK
355 prepare_arg[cydriver.CUgraphConditionalHandle](
356 self.data, self.data_addresses,
357 <intptr_t><unsigned long long>(<GraphCondition>arg)._c_handle, i)
358 continue
359 # TODO: support ctypes/numpy struct
360 raise TypeError("the argument is of unsupported type: " + str(type(arg))) 1aK
362 self.kernel_args = kernel_args 2zbw x AbBbCby z d e m n o p A B f g q r s t C D E F G H h i j k DbEbFbGbHbIbJbKbkblbmbnbobpbqbrbsbtbubLbMbNbwbObPbQbRbSbTbUbVbxba WbXbYbZb0b1b2bl 3b4b5b6b7b~ abbb8bL ) b u v I c hb$ # ' % ; : / . - , + * ybib! 9 W X 6 7 4 5 8 3 0 1 Y Z 2 ? = jb^ _ ] vbgbfbebdbcb} | { V J M U T S R Q P O N ( 9b!b#b$b%b@ [ 'b`
363 self.ptr = <intptr_t>self.data_addresses.data() 2zbw x AbBbCby z d e m n o p A B f g q r s t C D E F G H h i j k DbEbFbGbHbIbJbKbkblbmbnbobpbqbrbsbtbubLbMbNbwbObPbQbRbSbTbUbVbxba WbXbYbZb0b1b2bl 3b4b5b6b7b~ abbb8bL ) b u v I c hb$ # ' % ; : / . - , + * ybib! 9 W X 6 7 4 5 8 3 0 1 Y Z 2 ? = jb^ _ ] vbgbfbebdbcb} | { V J M U T S R Q P O N ( 9b!b#b$b%b@ [ 'b`
365 def __dealloc__(self) -> None:
366 for data in self.data: 2(bzb)b*bw x +b,b-bAb.b/b:b;b=bBb?b@b[bCby z d e m n o p A B f g q r s t C D E F G H h i j k ]b^b_b`b{b|b}b~bacbcccdcecfcgchcicjckclcmcncocpcqcrcscDbEbtcucvcwcxcFbGbHbIbJbKbkblbmbnbobpbqbrbsbtbubLbycMbNbwbObPbQbRbSbTbzcAcBcCcDcEcFcGcUbVbxbHcIca WbXbYbJcKcZb0b1b2bLcMcNcOcPcQcl 3b4b5b6b7bRc~ abbb8bScL ) b K u v I c Tchb$ # ' % ; : / . - , + * ybib! 9 W X 6 7 4 5 8 3 0 1 Y Z 2 ? = jb^ _ ] vbgbfbebdbcb} | { V J M U T S R Q P O N ( 9b!b#b$bUcVcWc%b@ [ XcYcZc0c1c2c3c4c5c6c7c8c9c!c#c$c%c'c(c)c*c+c,c'b`
367 if data: 2zbw x AbBbCby z d e m n o p A B f g q r s t C D E F G H h i j k DbEbFbGbHbIbJbKbkblbmbnbobpbqbrbsbtbubLbMbNbwbObPbQbRbSbTbUbVbxba WbXbYbZb0b1b2bl 3b4b5b6b7b~ abbb8bL ) b K u v I c hb$ # ' % ; : / . - , + * ybib! 9 W X 6 7 4 5 8 3 0 1 Y Z 2 ? = jb^ _ ] vbgbfbebdbcb} | { V J M U T S R Q P O N ( 9b!b#b$b%b@ [ 'b`
368 PyMem_Free(data) 2zbw x AbBbCby z d e m n o p A B f g q r s t C D E F G H h i j k DbEbFbGbHbIbJbKbkblbmbnbobpbqbrbsbtbubLbMbNbwbObPbQbRbSbTbUbVbxba WbXbYbZb0b1b2bl 3b4b5b6b7b~ abbb8bL ) b u v I c hb$ # ' % ; : / . - , + * ybib! 9 W X 6 7 4 5 8 3 0 1 Y Z 2 ? = jb^ _ ] vbgbfbebdbcb} | { V J M U T S R Q P O N ( 9b!b#b$b%b@ [ 'b`