refactor: implement Backend registry (2 of 2) (#2390)

* refactor: move `_backends` to module

* refactor: add new backend lookup

* refactor: add new backend lookup

* wip: imports

* refactor: split backends into modules

* fix: simplify return

src/awkward/__init__.py

@@ -5,6 +5,7 @@
 
 # NumPy-like alternatives
 import awkward._backends
+import awkward._nplikes
 
 # layout classes; functionality that used to be in C++ (in Awkward 1.x)
 import awkward.index

src/awkward/_backends/backend.py

@@ -0,0 +1,47 @@
+# BSD 3-Clause License; see https://github.com/scikit-hep/awkward-1.0/blob/main/LICENSE
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+
+import awkward as ak
+from awkward._nplikes.numpy import Numpy
+from awkward._nplikes.numpylike import NumpyLike, NumpyMetadata
+from awkward._singleton import Singleton
+from awkward._typing import Callable, Tuple, TypeAlias, TypeVar, Unpack
+
+np = NumpyMetadata.instance()
+numpy = Numpy.instance()
+
+
+T = TypeVar("T", covariant=True)
+KernelKeyType: TypeAlias = Tuple[str, Unpack[Tuple[np.dtype, ...]]]
+KernelType: TypeAlias = Callable[..., None]
+
+
+class Backend(Singleton, ABC):
+    name: str
+
+    @property
+    @abstractmethod
+    def nplike(self) -> NumpyLike:
+        raise NotImplementedError
+
+    @property
+    @abstractmethod
+    def index_nplike(self) -> NumpyLike:
+        raise NotImplementedError
+
+    def __getitem__(self, key: KernelKeyType) -> KernelType:
+        raise NotImplementedError
+
+    def apply_reducer(
+        self,
+        reducer: ak._reducers.Reducer,
+        layout: ak.contents.NumpyArray,
+        parents: ak.index.Index,
+        outlength: int,
+    ) -> ak.contents.NumpyArray:
+        return reducer.apply(layout, parents, outlength)
+
+    def apply_ufunc(self, ufunc, method, args, kwargs):
+        return getattr(ufunc, method)(*args, **kwargs)

src/awkward/_backends/cupy.py

@@ -0,0 +1,42 @@
+# BSD 3-Clause License; see https://github.com/scikit-hep/awkward-1.0/blob/main/LICENSE
+from __future__ import annotations
+
+from awkward._backends.backend import Backend, KernelKeyType
+from awkward._backends.dispatch import register_backend
+from awkward._kernels import CupyKernel, NumpyKernel
+from awkward._nplikes.cupy import Cupy
+from awkward._nplikes.numpy import Numpy
+from awkward._nplikes.numpylike import NumpyMetadata
+from awkward._typing import Final
+
+np = NumpyMetadata.instance()
+numpy = Numpy.instance()
+
+
+@register_backend(Cupy)
+class CupyBackend(Backend):
+    name: Final[str] = "cuda"
+
+    _cupy: Cupy
+
+    @property
+    def nplike(self) -> Cupy:
+        return self._cupy
+
+    @property
+    def index_nplike(self) -> Cupy:
+        return self._cupy
+
+    def __init__(self):
+        self._cupy = Cupy.instance()
+
+    def __getitem__(self, index: KernelKeyType) -> CupyKernel | NumpyKernel:
+        from awkward._connect import cuda
+
+        cupy = cuda.import_cupy("Awkward Arrays with CUDA")
+        _cuda_kernels = cuda.initialize_cuda_kernels(cupy)
+        func = _cuda_kernels[index]
+        if func is not None:
+            return CupyKernel(func, index)
+        else:
+            raise AssertionError(f"CuPyKernel not found: {index!r}")