micro optimization + make numpy as required

Author: leofang

cuda_core/cuda/core/_kernel_arg_handler.pyx

@@ -13,11 +13,7 @@ from libcpp cimport vector
 
 import ctypes
 
-# this might be an unnecessary assumption that NumPy does not exist...
-try:
-    import numpy
-except ImportError:
-    numpy = None
+import numpy
 
 from cuda.core._memory import Buffer
 
@@ -43,7 +39,32 @@ ctypedef fused supported_type:
     cpp_double_complex
 
 
-# TODO: cache ctypes/numpy type objects to avoid attribute access
+# cache ctypes/numpy type objects to avoid attribute access
+cdef object ctypes_bool = ctypes.c_bool
+cdef object ctypes_int8 = ctypes.c_int8
+cdef object ctypes_int16 = ctypes.c_int16
+cdef object ctypes_int32 = ctypes.c_int32
+cdef object ctypes_int64 = ctypes.c_int64
+cdef object ctypes_uint8 = ctypes.c_uint8
+cdef object ctypes_uint16 = ctypes.c_uint16
+cdef object ctypes_uint32 = ctypes.c_uint32
+cdef object ctypes_uint64 = ctypes.c_uint64
+cdef object ctypes_float = ctypes.c_float
+cdef object ctypes_double = ctypes.c_double
+cdef object numpy_bool = numpy.bool_
+cdef object numpy_int8 = numpy.int8
+cdef object numpy_int16 = numpy.int16
+cdef object numpy_int32 = numpy.int32
+cdef object numpy_int64 = numpy.int64
+cdef object numpy_uint8 = numpy.uint8
+cdef object numpy_uint16 = numpy.uint16
+cdef object numpy_uint32 = numpy.uint32
+cdef object numpy_uint64 = numpy.uint64
+cdef object numpy_float16 = numpy.float16
+cdef object numpy_float32 = numpy.float32
+cdef object numpy_float64 = numpy.float64
+cdef object numpy_complex64 = numpy.complex64
+cdef object numpy_complex128 = numpy.complex128
 
 
 # limitation due to cython/cython#534
@@ -76,27 +97,27 @@ cdef inline int prepare_ctypes_arg(
         vector.vector[void*]& data_addresses,
         arg,
         const size_t idx) except -1:
-    if isinstance(arg, ctypes.c_bool):
+    if isinstance(arg, ctypes_bool):
         return prepare_arg[cpp_bool](data, data_addresses, arg.value, idx)
-    elif isinstance(arg, ctypes.c_int8):
+    elif isinstance(arg, ctypes_int8):
         return prepare_arg[int8_t](data, data_addresses, arg.value, idx)
-    elif isinstance(arg, ctypes.c_int16):
+    elif isinstance(arg, ctypes_int16):
         return prepare_arg[int16_t](data, data_addresses, arg.value, idx)
-    elif isinstance(arg, ctypes.c_int32):
+    elif isinstance(arg, ctypes_int32):
         return prepare_arg[int32_t](data, data_addresses, arg.value, idx)
-    elif isinstance(arg, ctypes.c_int64):
+    elif isinstance(arg, ctypes_int64):
         return prepare_arg[int64_t](data, data_addresses, arg.value, idx)
-    elif isinstance(arg, ctypes.c_uint8):
+    elif isinstance(arg, ctypes_uint8):
         return prepare_arg[uint8_t](data, data_addresses, arg.value, idx)
-    elif isinstance(arg, ctypes.c_uint16):
+    elif isinstance(arg, ctypes_uint16):
         return prepare_arg[uint16_t](data, data_addresses, arg.value, idx)
-    elif isinstance(arg, ctypes.c_uint32):
+    elif isinstance(arg, ctypes_uint32):
         return prepare_arg[uint32_t](data, data_addresses, arg.value, idx)
-    elif isinstance(arg, ctypes.c_uint64):
+    elif isinstance(arg, ctypes_uint64):
         return prepare_arg[uint64_t](data, data_addresses, arg.value, idx)
-    elif isinstance(arg, ctypes.c_float):
+    elif isinstance(arg, ctypes_float):
         return prepare_arg[float](data, data_addresses, arg.value, idx)
-    elif isinstance(arg, ctypes.c_double):
+    elif isinstance(arg, ctypes_double):
         return prepare_arg[double](data, data_addresses, arg.value, idx)
     else:
         return 1
@@ -107,37 +128,34 @@ cdef inline int prepare_numpy_arg(
         vector.vector[void*]& data_addresses,
         arg,
         const size_t idx) except -1:
-    if not numpy:
-        return 1
-
-    if isinstance(arg, numpy.bool_):
+    if isinstance(arg, numpy_bool):
         return prepare_arg[cpp_bool](data, data_addresses, arg, idx)
-    elif isinstance(arg, numpy.int8):
+    elif isinstance(arg, numpy_int8):
         return prepare_arg[int8_t](data, data_addresses, arg, idx)
-    elif isinstance(arg, numpy.int16):
+    elif isinstance(arg, numpy_int16):
         return prepare_arg[int16_t](data, data_addresses, arg, idx)
-    elif isinstance(arg, numpy.int32):
+    elif isinstance(arg, numpy_int32):
         return prepare_arg[int32_t](data, data_addresses, arg, idx)
-    elif isinstance(arg, numpy.int64):
+    elif isinstance(arg, numpy_int64):
         return prepare_arg[int64_t](data, data_addresses, arg, idx)
-    elif isinstance(arg, numpy.uint8):
+    elif isinstance(arg, numpy_uint8):
         return prepare_arg[uint8_t](data, data_addresses, arg, idx)
-    elif isinstance(arg, numpy.uint16):
+    elif isinstance(arg, numpy_uint16):
         return prepare_arg[uint16_t](data, data_addresses, arg, idx)
-    elif isinstance(arg, numpy.uint32):
+    elif isinstance(arg, numpy_uint32):
         return prepare_arg[uint32_t](data, data_addresses, arg, idx)
-    elif isinstance(arg, numpy.uint64):
+    elif isinstance(arg, numpy_uint64):
         return prepare_arg[uint64_t](data, data_addresses, arg, idx)
-    elif isinstance(arg, numpy.float16):
+    elif isinstance(arg, numpy_float16):
         # use int16 as a proxy
         return prepare_arg[int16_t](data, data_addresses, arg, idx)
-    elif isinstance(arg, numpy.float32):
+    elif isinstance(arg, numpy_float32):
         return prepare_arg[float](data, data_addresses, arg, idx)
-    elif isinstance(arg, numpy.float64):
+    elif isinstance(arg, numpy_float64):
         return prepare_arg[double](data, data_addresses, arg, idx)
-    elif isinstance(arg, numpy.complex64):
+    elif isinstance(arg, numpy_complex64):
         return prepare_arg[cpp_single_complex](data, data_addresses, arg, idx)
-    elif isinstance(arg, numpy.complex128):
+    elif isinstance(arg, numpy_complex128):
         return prepare_arg[cpp_double_complex](data, data_addresses, arg, idx)
     else:
         return 1
@@ -185,9 +203,9 @@ cdef class ParamHolder:
                 continue
 
             not_prepared = prepare_numpy_arg(self.data, self.data_addresses, arg, i)
-            if not_prepared != 0:
+            if not_prepared:
                 not_prepared = prepare_ctypes_arg(self.data, self.data_addresses, arg, i)
-            if not_prepared != 0:
+            if not_prepared:
                 # TODO: support ctypes/numpy struct
                 raise TypeError
 

cuda_core/cuda/core/utils.py

@@ -1 +1,5 @@
+# Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: LicenseRef-NVIDIA-SOFTWARE-LICENSE
+
 from cuda.core._memoryview import GPUMemoryView, viewable

cuda_core/pyproject.toml

@@ -41,6 +41,9 @@ classifiers = [
     "Environment :: GPU :: NVIDIA CUDA :: 11",
     "Environment :: GPU :: NVIDIA CUDA :: 12",
 ]
+dependencies = [
+    "numpy",
+]
 
 
 [tool.setuptools]