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 <[email protected]>

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 = (

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

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")
+        )

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

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,

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

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

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[