Revert "fix: errors in existing cuda_kernels (#2877)"

This reverts commit 7687108.
scikit-hep · Dec 11, 2023 · 85538a8 · 85538a8
1 parent 7687108
commit 85538a8
Show file tree

Hide file tree

Showing 6 changed files with 13 additions and 7 deletions.
diff --git a/dev/generate-kernel-signatures.py b/dev/generate-kernel-signatures.py
@@ -15,6 +15,8 @@
     "awkward_ListArray_validity",
     "awkward_BitMaskedArray_to_ByteMaskedArray",
     "awkward_ListArray_compact_offsets",
+    "awkward_new_Identities",
+    "awkward_Identities32_to_Identities64",
     "awkward_ListOffsetArray_flatten_offsets",
     "awkward_IndexedArray_overlay_mask",
     "awkward_IndexedArray_fill_count",
@@ -34,6 +36,7 @@
     "awkward_RegularArray_getitem_next_range_spreadadvanced",
     "awkward_RegularArray_getitem_next_array",
     "awkward_missing_repeat",
+    "awkward_Identities_getitem_carry",
     "awkward_RegularArray_getitem_jagged_expand",
     "awkward_ListArray_getitem_jagged_expand",
     "awkward_ListArray_getitem_next_array",

diff --git a/dev/generate-tests.py b/dev/generate-tests.py
@@ -464,6 +464,8 @@ def gencpukerneltests(specdict):
     "awkward_ListArray_validity",
     "awkward_BitMaskedArray_to_ByteMaskedArray",
     "awkward_ListArray_compact_offsets",
+    "awkward_new_Identities",
+    "awkward_Identities32_to_Identities64",
     "awkward_ListOffsetArray_flatten_offsets",
     "awkward_IndexedArray_overlay_mask",
     "awkward_IndexedArray_fill_count",
@@ -483,6 +485,7 @@ def gencpukerneltests(specdict):
     "awkward_RegularArray_getitem_next_range_spreadadvanced",
     "awkward_RegularArray_getitem_next_array",
     "awkward_missing_repeat",
+    "awkward_Identities_getitem_carry",
     "awkward_RegularArray_getitem_jagged_expand",
     "awkward_ListArray_getitem_jagged_expand",
     "awkward_ListArray_getitem_next_array",

diff --git a/src/awkward/_connect/cuda/cuda_kernels/awkward_ByteMaskedArray_getitem_nextcarry.cu b/src/awkward/_connect/cuda/cuda_kernels/awkward_ByteMaskedArray_getitem_nextcarry.cu
@@ -4,9 +4,9 @@
 // def f(grid, block, args):
 //     (tocarry, mask, length, validwhen, invocation_index, err_code) = args
 //     scan_in_array = cupy.empty(length, dtype=cupy.int64)
-//     cuda_kernel_templates.get_function(fetch_specialization(['awkward_ByteMaskedArray_getitem_nextcarry_a', tocarry.dtype, mask.dtype]))(grid, block, (tocarry, mask, length, validwhen, scan_in_array, invocation_index, err_code))
+//     cuda_kernel_templates.get_function(fetch_specialization(['awkward_ByteMaskedArray_getitem_nextcarry_a', tocarry.dtype, mask.dtype]))(grid, block, (tocarry, mask, validwhen, length, scan_in_array, invocation_index, err_code))
 //     scan_in_array = inclusive_scan(grid, block, (scan_in_array, invocation_index, err_code))
-//     cuda_kernel_templates.get_function(fetch_specialization(['awkward_ByteMaskedArray_getitem_nextcarry_b', tocarry.dtype, mask.dtype]))(grid, block, (tocarry, mask, length, validwhen, scan_in_array, invocation_index, err_code))
+//     cuda_kernel_templates.get_function(fetch_specialization(['awkward_ByteMaskedArray_getitem_nextcarry_b', tocarry.dtype, mask.dtype]))(grid, block, (tocarry, mask, validwhen, length, scan_in_array, invocation_index, err_code))
 // out["awkward_ByteMaskedArray_getitem_nextcarry_a", {dtype_specializations}] = None
 // out["awkward_ByteMaskedArray_getitem_nextcarry_b", {dtype_specializations}] = None
 // END PYTHON

diff --git a/src/awkward/_connect/cuda/cuda_kernels/awkward_ByteMaskedArray_reduce_next_64.cu b/src/awkward/_connect/cuda/cuda_kernels/awkward_ByteMaskedArray_reduce_next_64.cu
@@ -4,9 +4,9 @@
 // def f(grid, block, args):
 //     (nextcarry, nextparents, outindex, mask, parents, length, validwhen, invocation_index, err_code) = args
 //     scan_in_array = cupy.empty(length, dtype=cupy.int64)
-//     cuda_kernel_templates.get_function(fetch_specialization(['awkward_ByteMaskedArray_reduce_next_64_a', nextcarry.dtype, nextparents.dtype, outindex.dtype]))(grid, block, (nextcarry, nextparents, outindex, mask, parents, length, validwhen, scan_in_array, invocation_index, err_code))
+//     cuda_kernel_templates.get_function(fetch_specialization(['awkward_ByteMaskedArray_reduce_next_64_a', nextcarry.dtype, nextparents.dtype, outindex.dtype]))(grid, block, (nextcarry, nextparents, outindex, mask, parents, validwhen, length, scan_in_array, invocation_index, err_code))
 //     scan_in_array = inclusive_scan(grid, block, (scan_in_array, length, invocation_index, err_code))
-//     cuda_kernel_templates.get_function(fetch_specialization(['awkward_ByteMaskedArray_reduce_next_64_b', nextcarry.dtype, nextparents.dtype, outindex.dtype]))(grid, block, (nextcarry, nextparents, outindex, mask, parents, length, validwhen, scan_in_array, invocation_index, err_code))
+//     cuda_kernel_templates.get_function(fetch_specialization(['awkward_ByteMaskedArray_reduce_next_64_b', nextcarry.dtype, nextparents.dtype, outindex.dtype]))(grid, block, (nextcarry, nextparents, outindex, mask, parents, validwhen, length, scan_in_array, invocation_index, err_code))
 // out["awkward_ByteMaskedArray_reduce_next_64_a", {dtype_specializations}] = None
 // out["awkward_ByteMaskedArray_reduce_next_64_b", {dtype_specializations}] = None
 // END PYTHON

diff --git a/.../_connect/cuda/cuda_kernels/awkward_ByteMaskedArray_reduce_next_nonlocal_nextshifts_64.cu b/.../_connect/cuda/cuda_kernels/awkward_ByteMaskedArray_reduce_next_nonlocal_nextshifts_64.cu
@@ -5,10 +5,10 @@
 //     (nextshifts, mask, length, valid_when, invocation_index, err_code) = args
 //     scan_in_array_k = cupy.empty(length, dtype=cupy.int64)
 //     scan_in_array_nullsum = cupy.empty(length, dtype=cupy.int64)
-//     cuda_kernel_templates.get_function(fetch_specialization(["awkward_ByteMaskedArray_reduce_next_nonlocal_nextshifts_64_a", nextshifts.dtype, mask.dtype]))(grid, block, (nextshifts, mask, length, valid_when, scan_in_array_k, scan_in_array_nullsum, invocation_index, err_code))
+//     cuda_kernel_templates.get_function(fetch_specialization(["awkward_ByteMaskedArray_reduce_next_nonlocal_nextshifts_64_a", nextshifts.dtype, mask.dtype]))(grid, block, (nextshifts, mask, valid_when, length, scan_in_array_k, scan_in_array_nullsum, invocation_index, err_code))
 //     scan_in_array_k = inclusive_scan(grid, block, (scan_in_array_k, invocation_index, err_code))
 //     scan_in_array_nullsum = inclusive_scan(grid, block, (scan_in_array_nullsum, invocation_index, err_code))
-//     cuda_kernel_templates.get_function(fetch_specialization(["awkward_ByteMaskedArray_reduce_next_nonlocal_nextshifts_64_a", nextshifts.dtype, mask.dtype]))(grid, block, (nextshifts, mask, length, valid_when, scan_in_array_k, scan_in_array_nullsum, invocation_index, err_code))
+//     cuda_kernel_templates.get_function(fetch_specialization(["awkward_ByteMaskedArray_reduce_next_nonlocal_nextshifts_64_a", nextshifts.dtype, mask.dtype]))(grid, block, (nextshifts, mask, valid_when, length, scan_in_array_k, scan_in_array_nullsum, invocation_index, err_code))
 // out["awkward_ByteMaskedArray_reduce_next_nonlocal_nextshifts_64_a", {dtype_specializations}] = None
 // out["awkward_ByteMaskedArray_reduce_next_nonlocal_nextshifts_64_b", {dtype_specializations}] = None
 // END PYTHON

diff --git a/src/awkward/_connect/cuda/cuda_kernels/awkward_ListArray_getitem_jagged_expand.cu b/src/awkward/_connect/cuda/cuda_kernels/awkward_ListArray_getitem_jagged_expand.cu
@@ -18,7 +18,7 @@ awkward_ListArray_getitem_jagged_expand(T* multistarts,
                                         uint64_t invocation_index,
                                         uint64_t* err_code) {
   if (err_code[0] == NO_ERROR) {
-    int64_t thread_id = (blockIdx.x * blockDim.x + threadIdx.x) / jaggedsize;
+    int64_t thread_id = (blockIdx.x * blockDim.x + threadIdx.x) % length;
     int64_t thready_id = (blockIdx.x * blockDim.x + threadIdx.x) % jaggedsize;
     W start = fromstarts[thread_id];
     X stop = fromstops[thread_id];