diff --git a/numba_cuda/numba/cuda/dispatcher.py b/numba_cuda/numba/cuda/dispatcher.py index f8cf84d56..68f32b971 100644 --- a/numba_cuda/numba/cuda/dispatcher.py +++ b/numba_cuda/numba/cuda/dispatcher.py @@ -1626,21 +1626,7 @@ def _compile_for_args(self, *args, **kws): def typeof_pyval(self, val): # Based on _DispatcherBase.typeof_pyval, but differs from it to support # the CUDA Array Interface. - try: - return typeof(val, Purpose.argument) - except ValueError: - if ( - interface := getattr(val, "__cuda_array_interface__") - ) is not None: - # When typing, we don't need to synchronize on the array's - # stream - this is done when the kernel is launched. - - return typeof( - cuda.from_cuda_array_interface(interface, sync=False), - Purpose.argument, - ) - else: - raise + return typeof(val, Purpose.argument) def specialize(self, *args): """ diff --git a/numba_cuda/numba/cuda/np/numpy_support.py b/numba_cuda/numba/cuda/np/numpy_support.py index 14d24f90b..b2102fc72 100644 --- a/numba_cuda/numba/cuda/np/numpy_support.py +++ b/numba_cuda/numba/cuda/np/numpy_support.py @@ -4,6 +4,7 @@ import collections import ctypes import itertools +import functools import operator import re @@ -21,11 +22,12 @@ numpy_version = tuple(map(int, np.__version__.split(".")[:2])) +@functools.lru_cache def strides_from_shape( shape: tuple[int, ...], itemsize: int, *, order: str ) -> tuple[int, ...]: """Compute strides for a contiguous array with given shape and order.""" - if len(shape) == 0: + if not shape: # 0-D arrays have empty strides return () limits = slice(1, None) if order == "C" else slice(None, -1) diff --git a/numba_cuda/numba/cuda/typing/typeof.py b/numba_cuda/numba/cuda/typing/typeof.py index 2d3e1e826..4e095afbb 100644 --- a/numba_cuda/numba/cuda/typing/typeof.py +++ b/numba_cuda/numba/cuda/typing/typeof.py @@ -5,6 +5,7 @@ from functools import singledispatch import ctypes import enum +import operator import numpy as np from numpy.random.bit_generator import BitGenerator @@ -318,17 +319,13 @@ def _typeof_cuda_array_interface(val, c): Array Interface. These are typed as regular Array types, with lowering handled in numba.cuda.np.arrayobj. """ - # Only handle constants, not arguments (arguments use regular array typing) - if c.purpose == Purpose.argument: - return None - dtype = numpy_support.from_dtype(np.dtype(val["typestr"])) shape = val["shape"] ndim = len(shape) strides = val.get("strides") # Determine layout - if ndim == 0: + if not ndim: layout = "C" elif strides is None: layout = "C" @@ -340,18 +337,14 @@ def _typeof_cuda_array_interface(val, c): c_strides = numpy_support.strides_from_shape( shape, itemsize, order="C" ) - layout = ( - "C" if all(x == y for x, y in zip(strides, c_strides)) else "A" - ) + layout = "C" if all(map(operator.eq, strides, c_strides)) else "A" elif strides[0] == itemsize: f_strides = numpy_support.strides_from_shape( shape, itemsize, order="F" ) - layout = ( - "F" if all(x == y for x, y in zip(strides, f_strides)) else "A" - ) + layout = "F" if all(map(operator.eq, strides, f_strides)) else "A" else: layout = "A" - readonly = val["data"][1] + _, readonly = val["data"] return types.Array(dtype, ndim, layout, readonly=readonly)