From 2ef498e615849aed2590fcb74c16eca9a2caf2a8 Mon Sep 17 00:00:00 2001 From: Angus Hollands Date: Thu, 13 Apr 2023 18:19:49 +0100 Subject: [PATCH] fix: don't touch for `ascontiguousarray` (#2397) * fix: don't touch for `ascontiguousarray` This is because touching should really only occur if the operation result _depends_ upon the array value, which it does not in this instance. * refactor: use `array.size` * fix: `ascontiguousarray` should use same key and report * fix: use Numpy.instance() --------- Co-authored-by: Jim Pivarski --- src/awkward/_kernels.py | 28 ++++++++++++++++++++++------ src/awkward/_nplikes/array_module.py | 9 +++++---- src/awkward/_nplikes/cupy.py | 4 +++- src/awkward/_nplikes/jax.py | 11 ++--------- src/awkward/_nplikes/numpy.py | 4 +++- src/awkward/_nplikes/numpylike.py | 4 +--- src/awkward/_nplikes/typetracer.py | 9 ++++----- src/awkward/contents/numpyarray.py | 8 ++------ 8 files changed, 42 insertions(+), 35 deletions(-) diff --git a/src/awkward/_kernels.py b/src/awkward/_kernels.py index 571188b92f..890ed7afff 100644 --- a/src/awkward/_kernels.py +++ b/src/awkward/_kernels.py @@ -6,6 +6,7 @@ from typing import Any, Callable import awkward as ak +from awkward._nplikes.cupy import Cupy from awkward._nplikes.jax import Jax from awkward._nplikes.numpy import Numpy from awkward._nplikes.numpylike import NumpyMetadata @@ -14,6 +15,7 @@ KernelKeyType: TypeAlias = tuple # Tuple[str, Unpack[Tuple[metadata.dtype, ...]]] +numpy = Numpy.instance() metadata = NumpyMetadata.instance() @@ -68,7 +70,8 @@ class NumpyKernel(BaseKernel): def _cast(cls, x, t): if issubclass(t, ctypes._Pointer): # Do we have a NumPy-owned array? - if Numpy.is_own_array(x): + if numpy.is_own_array(x): + assert numpy.is_c_contiguous(x), "kernel expects contiguous array" if x.ndim > 0: return ctypes.cast(x.ctypes.data, t) else: @@ -100,14 +103,16 @@ def __call__(self, *args) -> None: class CupyKernel(BaseKernel): - def max_length(self, args): - import awkward._connect.cuda as ak_cuda + def __init__(self, impl: Callable[..., Any], key: KernelKeyType): + super().__init__(impl, key) - cupy = ak_cuda.import_cupy("Awkward Arrays with CUDA") + self._cupy = Cupy.instance() + + def max_length(self, args): max_length = metadata.iinfo(metadata.int64).min # TODO should kernels strip nplike wrapper? Probably for array in args: - if isinstance(array, cupy.ndarray): + if self._cupy.is_own_array(array): max_length = max(max_length, len(array)) return max_length @@ -121,6 +126,15 @@ def calc_blocks(self, length): return 1024, 1, 1 return length, 1, 1 + def _cast(self, x, t): + if issubclass(t, ctypes._Pointer): + # Do we have a NumPy-owned array? + if self._cupy.is_own_array(x): + assert self._cupy.is_c_contiguous(x) + return x + else: + return x + def __call__(self, *args) -> None: import awkward._connect.cuda as ak_cuda @@ -134,8 +148,10 @@ def __call__(self, *args) -> None: cupy.array(ak_cuda.NO_ERROR), [], ) + assert len(args) == len(self._impl.argtypes) + + args = [self._cast(x, t) for x, t in zip(args, self._impl.argtypes)] - assert len(args) == len(self._impl.dir) # The first arg is the invocation index which raises itself by 8 in the kernel if there was no error before. # The second arg is the error_code. args = ( diff --git a/src/awkward/_nplikes/array_module.py b/src/awkward/_nplikes/array_module.py index ed45942a43..89526f1918 100644 --- a/src/awkward/_nplikes/array_module.py +++ b/src/awkward/_nplikes/array_module.py @@ -36,10 +36,11 @@ def asarray( else: return self._module.asarray(obj, dtype=dtype) - def ascontiguousarray( - self, x: ArrayLike, *, dtype: np.dtype | None = None - ) -> ArrayLike: - return self._module.ascontiguousarray(x, dtype=dtype) + def ascontiguousarray(self, x: ArrayLike) -> ArrayLike: + # Allow buffers to _pretend_ to be contiguous already + if x.dtype.metadata is not None and x.dtype.metadata.get("pretend_contiguous"): + return x + return self._module.ascontiguousarray(x) def frombuffer( self, buffer, *, dtype: np.dtype | None = None, count: int = -1 diff --git a/src/awkward/_nplikes/cupy.py b/src/awkward/_nplikes/cupy.py index fb4393cca4..7a86ea4e50 100644 --- a/src/awkward/_nplikes/cupy.py +++ b/src/awkward/_nplikes/cupy.py @@ -154,4 +154,6 @@ def is_own_array_type(cls, type_: type) -> bool: return module == "cupy" def is_c_contiguous(self, x: ArrayLike) -> bool: - return x.flags["C_CONTIGUOUS"] + return x.flags["C_CONTIGUOUS"] or ( + x.dtype.metadata is not None and x.dtype.metadata.get("pretend_contiguous") + ) diff --git a/src/awkward/_nplikes/jax.py b/src/awkward/_nplikes/jax.py index fcd214dd42..3e2c865057 100644 --- a/src/awkward/_nplikes/jax.py +++ b/src/awkward/_nplikes/jax.py @@ -1,8 +1,6 @@ # BSD 3-Clause License; see https://github.com/scikit-hep/awkward-1.0/blob/main/LICENSE from __future__ import annotations -import numpy - import awkward as ak from awkward._nplikes.array_module import ArrayModuleNumpyLike from awkward._nplikes.dispatch import register_nplike @@ -74,10 +72,5 @@ def is_tracer_type(cls, type_: type) -> bool: def is_c_contiguous(self, x: ArrayLike) -> bool: return True - def ascontiguousarray( - self, x: ArrayLike, *, dtype: numpy.dtype | None = None - ) -> ArrayLike: - if dtype is not None: - return x.astype(dtype) - else: - return x + def ascontiguousarray(self, x: ArrayLike) -> ArrayLike: + return x diff --git a/src/awkward/_nplikes/numpy.py b/src/awkward/_nplikes/numpy.py index cf2bd605af..4515c66e5b 100644 --- a/src/awkward/_nplikes/numpy.py +++ b/src/awkward/_nplikes/numpy.py @@ -43,7 +43,9 @@ def is_own_array_type(cls, type_) -> bool: return issubclass(type_, numpy.ndarray) def is_c_contiguous(self, x: ArrayLike) -> bool: - return x.flags["C_CONTIGUOUS"] + return x.flags["C_CONTIGUOUS"] or ( + x.dtype.metadata is not None and x.dtype.metadata.get("pretend_contiguous") + ) def packbits( self, diff --git a/src/awkward/_nplikes/numpylike.py b/src/awkward/_nplikes/numpylike.py index c9c46ea064..6a2c715d8d 100644 --- a/src/awkward/_nplikes/numpylike.py +++ b/src/awkward/_nplikes/numpylike.py @@ -248,9 +248,7 @@ def asarray( ... @abstractmethod - def ascontiguousarray( - self, x: ArrayLike, *, dtype: numpy.dtype | None = None - ) -> ArrayLike: + def ascontiguousarray(self, x: ArrayLike) -> ArrayLike: ... @abstractmethod diff --git a/src/awkward/_nplikes/typetracer.py b/src/awkward/_nplikes/typetracer.py index 06e1f34e31..bcb586ba04 100644 --- a/src/awkward/_nplikes/typetracer.py +++ b/src/awkward/_nplikes/typetracer.py @@ -695,11 +695,10 @@ def populate_shape_and_items(node, dim): else: raise TypeError - def ascontiguousarray( - self, x: ArrayLike, *, dtype: numpy.dtype | None = None - ) -> TypeTracerArray: - try_touch_data(x) - return TypeTracerArray._new(dtype or x.dtype, shape=x.shape) + def ascontiguousarray(self, x: ArrayLike) -> TypeTracerArray: + return TypeTracerArray._new( + x.dtype, shape=x.shape, form_key=x.form_key, report=x.report + ) def frombuffer( self, buffer, *, dtype: np.dtype | None = None, count: int = -1 diff --git a/src/awkward/contents/numpyarray.py b/src/awkward/contents/numpyarray.py index b0972810ad..005558bbe3 100644 --- a/src/awkward/contents/numpyarray.py +++ b/src/awkward/contents/numpyarray.py @@ -703,19 +703,15 @@ def _unique(self, negaxis, starts, parents, outlength): if negaxis is None: contiguous_self = self.to_contiguous() - # Python 3.8 could use math.prod - flattened_shape = 1 - for s in contiguous_self.shape: - flattened_shape = flattened_shape * s offsets = ak.index.Index64.zeros(2, self._backend.index_nplike) - offsets[1] = flattened_shape + offsets[1] = self._data.size dtype = ( np.dtype(np.int64) if self._data.dtype.kind.upper() == "M" else self._data.dtype ) - out = self._backend.nplike.empty(offsets[1], dtype=dtype) + out = self._backend.nplike.empty(self._data.size, dtype=dtype) assert offsets.nplike is self._backend.index_nplike self._handle_error( self._backend[