test: overload methods

tests-cuda/test_1809-cuda-jit.py

@@ -4,12 +4,13 @@
 import pytest  # noqa: F401
 
 import awkward as ak  # noqa: F401
-import cupy
 
 numba = pytest.importorskip("numba")
 
 from numba import cuda, types  # noqa: F401, E402
 from numba.core.typing.typeof import typeof, typeof_impl  # noqa: F401, E402
+from numba.extending import overload
+from numba.extending import overload_method
 
 from numba import config
 config.CUDA_LOW_OCCUPANCY_WARNINGS = False
@@ -23,10 +24,15 @@
 
 class ArrayViewArgHandler:
     def prepare_args(self, ty, val, **kwargs):
-        print(repr(val), type(val))
-        if isinstance(val, ak.Array):   
-            return ty, val
+        print("ArrayViewArgHandler::prepare_args", repr(val), type(val))
+        if isinstance(val, ak.Array):
+            print("High level ak.Array to_numpy data pointer is", ak.to_numpy(val).ctypes.data)
+            # A pointer to the memory area of the array as a Python integer.
+            # return types.uint64, val.layout._data.ctypes.data
+            #
+            return types.uint64, ak._connect.numba.arrayview.ArrayView.fromarray(val).lookup.arrayptrs.ctypes.data
         elif isinstance(val, ak._connect.numba.arrayview.ArrayView):
+            print("ak.ArrayView")
             return types.uint64, val._numbaview.lookup.arrayptrs
         else:
             return ty, val
@@ -41,9 +47,13 @@ def prepare_args(self, ty, val, **kwargs):
 def swallow(array):
     pass
 
+@cuda.jit(extensions=[array_view_arg_handler])
+def multiply(array, n):
+    for i in range(len(array)):
+        print(n * array[i])
 
 @cuda.jit
-def passthrough(array):
+def passthrough(array, device=True):
     return array
 
 
@@ -53,26 +63,77 @@ def passthrough2(array):
 
 
 @cuda.jit
-def digest(array):
-    return array[0]
+def digest(array, val):
+    val[0] = 10 * array[0]
 
 
 @cuda.jit
-def digest2(array):
-    tmp = array[0]
-    return tmp, tmp, array[0]
+def digest2(array, tmp):
+    tmp[0] = array[0]
+    tmp[1] = array[0]
+    tmp[2] = array[0]
+
 
+@cuda.jit
+def increment_by_one(an_array):
+    # Thread id in a 1D block
+    tx = cuda.threadIdx.x
+    # Block id in a 1D grid
+    ty = cuda.blockIdx.x
+    # Block width, i.e. number of threads per block
+    bw = cuda.blockDim.x
+    # Compute flattened index inside the array
+    pos = tx + ty * bw
+    if pos < an_array.size:  # Check array boundaries
+        an_array[pos] += 1
+
 def test_numpy_array_1d():
-    nparray = np.array([0, 1, 2, 3], dtype=int)
+    nparray = np.array([10, 1, 2, 3], dtype=int)
     swallow[1, 1](nparray)
+    print("Swallow", nparray)
+    multiply[1, 1](nparray, 3)
+    print("Multiply", nparray)
+#    arr = passthrough[1, 1](np.array([10, 11, 12, 13], dtype=int))
+#    arr1, arr2 = passthrough2[1, 1](arr)
+#    print(arr1, arr2)
+    v0 = np.empty(1)
+    digest[1, 1](nparray, v0)
+    print("Digest", v0[0])
+    v1 = np.empty(3)
+    digest2[1, 1](nparray, v1)
+    print(v1)
+    increment_by_one[1, 3](v1)
+    print(v1)
 
 def test_to_numy_array_1d():
     akarray = ak.Array([0, 1, 2, 3])
     swallow[1, 1](ak.to_numpy(akarray))
 
-#def test_array_1d():
-#    akarray = ak.Array([0, 1, 2, 3])
-#    swallow[1, 1](akarray))
+def test_mem_management():
+    # copy host->device a numpy array:
+    ary = np.arange(10)
+    d_ary = cuda.to_device(ary)
+
+
+def test_mem_management1():
+    ary = np.arange(10)
+    # enqueue the transfer to a stream:
+    stream = cuda.stream()
+    d_ary = cuda.to_device(ary, stream=stream)
+    # copy device->host:
+    hary = d_ary.copy_to_host()
+
+def test_mem_management2():
+    ary = np.arange(10)
+    d_ary = cuda.to_device(ary)
+    # copy device->host to an existing array:
+    ary = np.empty(shape=d_ary.shape, dtype=d_ary.dtype)
+    d_ary.copy_to_host(ary)
+
+def test_array_1d():    
+    akarray = ak.Array([0, 1, 2, 3])
+    swallow[1, 1](akarray)
+
 
 def test_array_njit():
     @numba.njit
@@ -110,3 +171,48 @@ def something(array):
     something[1, 1](ak.to_numpy(akarray))
 
 
+@overload(len, target='cuda')
+def grid_group_len(seq):
+    if isinstance(seq, cuda.types.GridGroup):
+        def len_impl(seq):
+            n = cuda.gridsize(1)
+            return n
+        return len_impl
+
+
+@cuda.jit
+def fun_len():
+    if cuda.grid(1) == 0:
+        print("Grid size is", len(cuda.cg.this_grid()))
+
+def test_fun_overload():
+    fun_len[1, 1]()
+    fun_len[1, 2]()
+    fun_len[1, 3]()
+    cuda.synchronize()
+
+
+@overload_method(types.Array, 'sum', target='cuda')
+def array_sum(arr):
+    if arr.ndim != 1:
+        # Only implement 1D for this quick example
+        return None
+
+    def sum_impl(arr):
+        res = 0 
+        for i in range(len(arr)):
+            res += arr[i]
+        return res 
+    return sum_impl
+
+@cuda.jit
+def fun_sum(arr):
+    print("Sum is", arr.sum())
+
+import numpy as np
+
+def test_method_overload():
+    fun_sum[1, 1](np.arange(5))
+    fun_sum[1, 1](np.arange(10))
+    cuda.synchronize()
+

tests-cuda/test_1809-cuda-overload.py

@@ -0,0 +1,38 @@
+# BSD 3-Clause License; see https://github.com/scikit-hep/awkward-1.0/blob/main/LICENSE
+
+import numpy as np  # noqa: F401
+import pytest  # noqa: F401
+
+import awkward as ak  # noqa: F401
+
+numba = pytest.importorskip("numba")
+
+from numba.extending import overload_method
+from numba import cuda, types
+
+from numba import config
+config.CUDA_LOW_OCCUPANCY_WARNINGS = False
+config.CUDA_WARN_ON_IMPLICIT_COPY = False
+
+@overload_method(types.Array, 'sum', target='cuda')
+def array_sum(arr):
+    if arr.ndim != 1:
+        # Only implement 1D for this quick example
+        return None
+
+    def sum_impl(arr):
+        res = 0 
+        for i in range(len(arr)):
+            res += arr[i]
+        return res 
+    return sum_impl
+
+@cuda.jit
+def f(arr):
+    print("Sum is", arr.sum())
+
+
+def test_overload_method():
+    f[1, 1](np.arange(5))
+    f[1, 1](np.arange(10))
+    cuda.synchronize()

tests/test_0572-numba-array-ndim.py

@@ -17,6 +17,7 @@ def f1(array):
 
     assert f1(ak.highlevel.Array([[1, 2, 3], [], [4, 5]])) == 2
     assert f1(ak.highlevel.Array([[[1], [2, 3]], [], [[4, 5], []]])) == 3
+    assert f1(ak.highlevel.Array({"x": [1, 2, 3], "y": [None, None, 4]})) == 1
 
     with pytest.raises(numba.core.errors.TypingError):
         f1(ak.highlevel.Record({"x": [1, 2, 3], "y": [4]}))