properly handle ELEMENTS_PER_THREAD

Signed-off-by: Andrew Duffy <[email protected]>

Author: a10y

vortex-cuda/kernels/src/sequence.cu

@@ -1,21 +1,26 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+#include "config.cuh"
 #include <stdint.h>
 
+#define MIN(a, b) (((a) < (b)) : (a) : (b))
+
 template<typename ValueT>
 __device__ void sequence(
     ValueT *const output,
     ValueT base,
     ValueT multiplier,
     uint64_t len
 ) {
-    const uint64_t idx = blockIdx.x * blockDim.x + threadIdx.x;
-    if (idx >= len) {
-        return;
-    }
+    const uint64_t worker = blockIdx.x * blockDim.x + threadIdx.x;
 
-    output[idx] = static_cast<ValueT>(idx) * multiplier + base;
+    const uint64_t elemStart = MIN(worker * ELEMENTS_PER_THREAD, len);
+    const uint64_t elemEnd = MIN(elemStart + ELEMENTS_PER_THREAD, len);
+
+    for (uint64_t idx = elemStart; idx < elemEnd; idx++) {
+        output[idx] = static_cast<ValueT>(idx) * multiplier + base;
+    }
 }
 
 #define GENERATE_KERNEL(ValueT, suffix) \