Repository: numba/numba
Pull request: CUDA: Make `grid()` and `gridsize()` use 64-bit integers
It is possible to launch more than `2**31` threads, so `cuda.grid()` and `cuda.gridsize()` need to use 64 bit integers.

Fixes #9229.

Problem statement:
cuda.grid(1) returns incorrect value due to signed integer overflow
`cuda.grid(1)` returns an incorrect value in the following code:

```py
import numpy as np
from numba import cuda


@cuda.jit
def bug(out):
    i = cuda.threadIdx.x + cuda.blockIdx.x * cuda.blockDim.x
    i2 = cuda.grid(1)
    if i != i2:
        # this shouldn't happen
        print(i, i2)


free_mem, _ = cuda.current_context().get_memory_info()
# max_num_parallel_iter is around 566281830 on my 4090
max_num_parallel_iter = free_mem // 4 - 128 * 1024 * 1024
dev_out = cuda.device_array(max_num_parallel_iter, dtype=np.float32)
bug[22121216, 256](dev_out)
```

One example of a case which is printed is when `i = 2147551327` and `i2 = -214741596`, which likely indicates a signed 32 bit integer overflow. The bug doesn't appear when setting `max_num_parallel_iter` to an integer constant, which must mean it's somehow related to the calculation from the memory info.

Hints from issue discussion:
Thanks for the report! I can reproduce this, and it is due to an integer overflow.

> The bug doesn't appear when setting `max_num_parallel_iter` to an integer constant, which must mean it's somehow related to the calculation from the memory info.

I think the reason you don't see it is that launches are asynchronous, and setting `max_num_parallel_iter` to a constant changed the timing of execution such that your program manages to exit before the values printed by the kernel get flushed into the terminal. I find that (and I think you should be able to reproduce too) adding `cuda.synchronize()` to the end of the program manifests the error even when there's no computation or device array:

```
from numba import cuda

@cuda.jit
def bug():
    i = cuda.threadIdx.x + cuda.blockIdx.x * cuda.blockDim.x
    i2 = cuda.grid(1)
    if i != i2:
        # this shouldn't happen
        print(i, i2)

bug[22121216, 256]()
cuda.synchronize()
```
I have a fix in https://github.com/gmarkall/numba/commit/c40f740c2f9c0831ac1a7d6fada919ce60605bc7 - I will need to add tests to this for a PR, but intend to do this soon.
@nathanfarlow There is a fix in PR #9235. As a workaround for now, can you write your kernels such that they don't require more than 2^31 threads, or use the manual computation?

One reason what seems like an "obvious" problem has slipped by for so long is that I think people generally don't use so many threads,  so this isn't a common pattern in practice. Instead, it's usually more efficient to use a grid-strided loop with fewer threads, as outlined in: https://towardsdatascience.com/cuda-by-numba-examples-1-4-e0d06651612f#0909
Thanks for looking into this and thanks for the advice! :)

Commit message: Skip Issue #9229 test on cudasim

Patch to review:
diff --git a/docs/upcoming_changes/9235.cuda.rst b/docs/upcoming_changes/9235.cuda.rst
new file mode 100644
index 00000000000..4f72afc8ecb
--- /dev/null
+++ b/docs/upcoming_changes/9235.cuda.rst
@@ -0,0 +1,5 @@
+Make grid() and gridsize() use 64-bit integers
+==============================================
+
+`cuda.grid()` and `cuda.gridsize()` now use 64-bit integers, so they no longer
+overflow when the grid contains more than ``2 ** 31`` threads.
diff --git a/numba/cuda/intrinsics.py b/numba/cuda/intrinsics.py
index 746669ed8d7..ae8096744b7 100644
--- a/numba/cuda/intrinsics.py
+++ b/numba/cuda/intrinsics.py
@@ -15,9 +15,9 @@
 def _type_grid_function(ndim):
     val = ndim.literal_value
     if val == 1:
-        restype = types.int32
+        restype = types.int64
     elif val in (2, 3):
-        restype = types.UniTuple(types.int32, val)
+        restype = types.UniTuple(types.int64, val)
     else:
         raise ValueError('argument can only be 1, 2, 3')
 
@@ -48,7 +48,7 @@ def grid(typingctx, ndim):
 
     def codegen(context, builder, sig, args):
         restype = sig.return_type
-        if restype == types.int32:
+        if restype == types.int64:
             return nvvmutils.get_global_id(builder, dim=1)
         elif isinstance(restype, types.UniTuple):
             ids = nvvmutils.get_global_id(builder, dim=restype.count)
@@ -80,15 +80,16 @@ def gridsize(typingctx, ndim):
     sig = _type_grid_function(ndim)
 
     def _nthreads_for_dim(builder, dim):
+        i64 = ir.IntType(64)
         ntid = nvvmutils.call_sreg(builder, f"ntid.{dim}")
         nctaid = nvvmutils.call_sreg(builder, f"nctaid.{dim}")
-        return builder.mul(ntid, nctaid)
+        return builder.mul(builder.sext(ntid, i64), builder.sext(nctaid, i64))
 
     def codegen(context, builder, sig, args):
         restype = sig.return_type
         nx = _nthreads_for_dim(builder, 'x')
 
-        if restype == types.int32:
+        if restype == types.int64:
             return nx
         elif isinstance(restype, types.UniTuple):
             ny = _nthreads_for_dim(builder, 'y')
diff --git a/numba/cuda/nvvmutils.py b/numba/cuda/nvvmutils.py
index adf8e619e04..9a7dcde02b4 100644
--- a/numba/cuda/nvvmutils.py
+++ b/numba/cuda/nvvmutils.py
@@ -217,9 +217,10 @@ def nctaid(self, xyz):
         return call_sreg(self.builder, 'nctaid.%s' % xyz)
 
     def getdim(self, xyz):
-        tid = self.tid(xyz)
-        ntid = self.ntid(xyz)
-        nctaid = self.ctaid(xyz)
+        i64 = ir.IntType(64)
+        tid = self.builder.sext(self.tid(xyz), i64)
+        ntid = self.builder.sext(self.ntid(xyz), i64)
+        nctaid = self.builder.sext(self.ctaid(xyz), i64)
         res = self.builder.add(self.builder.mul(ntid, nctaid), tid)
         return res
 
diff --git a/numba/cuda/tests/cudapy/test_intrinsics.py b/numba/cuda/tests/cudapy/test_intrinsics.py
index 6d03c322da7..b97d098a043 100644
--- a/numba/cuda/tests/cudapy/test_intrinsics.py
+++ b/numba/cuda/tests/cudapy/test_intrinsics.py
@@ -440,6 +440,30 @@ def test_simple_gridsize1d(self):
         compiled[nctaid, ntid](ary)
         self.assertEqual(ary[0], nctaid * ntid)
 
+    @skip_on_cudasim('Requires too many threads')
+    def test_issue_9229(self):
+        # Ensure that grid and grid size are correct - #9229 showed that they
+        # overflowed an int32.
+        @cuda.jit
+        def f(grid_error, gridsize_error):
+            i1 = cuda.grid(1)
+            i2 = cuda.blockIdx.x * cuda.blockDim.x + cuda.threadIdx.x
+            gs1 = cuda.gridsize(1)
+            gs2 = cuda.blockDim.x * cuda.gridDim.x
+            if i1 != i2:
+                grid_error[0] = 1
+            if gs1 != gs2:
+                gridsize_error[0] = 1
+
+        grid_error = np.zeros(1, dtype=np.uint64)
+        gridsize_error = np.zeros(1, dtype=np.uint64)
+
+        # A lerge enough grid for thread IDs to overflow an int32
+        f[22121216, 256](grid_error, gridsize_error)
+
+        self.assertEqual(grid_error[0], 0)
+        self.assertEqual(gridsize_error[0], 0)
+
     @skip_on_cudasim('Tests PTX emission')
     def test_selp(self):
         sig = (int64[:], int64, int64[:])


Code review comments:
diff_hunk: @@ -440,6 +440,30 @@ def test_simple_gridsize1d(self):
         compiled[nctaid, ntid](ary)
         self.assertEqual(ary[0], nctaid * ntid)
 
+    @skip_on_cudasim('Requires too many threads')
+    def test_issue_9229(self):
+        # Ensure that grid and grid size are correct - #9229 showed that they
+        # overflowed an int32.
+        @cuda.jit
+        def f(grid_error, gridsize_error):
+            i1 = cuda.grid(1)
+            i2 = cuda.blockIdx.x * cuda.blockDim.x + cuda.threadIdx.x
+            gs1 = cuda.gridsize(1)
+            gs2 = cuda.blockDim.x * cuda.gridDim.x
+            if i1 != i2:
+                grid_error[0] = 1
+            if gs1 != gs2:
+                gridsize_error[0] = 1
+
+        grid_error = np.zeros(1, dtype=np.uint64)
+        gridsize_error = np.zeros(1, dtype=np.uint64)
+
+        # A lerge enough grid for thread IDs to overflow an int32
+        f[22121216, 256](grid_error, gridsize_error)
path: numba/cuda/tests/cudapy/test_intrinsics.py
text: @user1:
Should the blocks per grid be more easily related to the size of an int32 given TPB is 256?

@author:
I felt the way I'd written it (taken in conjunction with the comment on the previous line) was sufficient - if you'd like a specific change could you please suggest what you have in mind?

@user1:
I'd guess at `(2 ** 32 // 256) + 1` but a comment suggesting the source of the current number would be ok too (I couldn't work out why that specific number was there, but realise now it's from the reproducer in the ticket). Thanks!

diff_hunk: @@ -440,6 +440,30 @@ def test_simple_gridsize1d(self):
         compiled[nctaid, ntid](ary)
         self.assertEqual(ary[0], nctaid * ntid)
 
+    @skip_on_cudasim('Requires too many threads')
+    def test_issue_9229(self):
+        # Ensure that grid and grid size are correct - #9229 showed that they
+        # overflowed an int32.
+        @cuda.jit
+        def f(grid_error, gridsize_error):
+            i1 = cuda.grid(1)
+            i2 = cuda.blockIdx.x * cuda.blockDim.x + cuda.threadIdx.x
+            gs1 = cuda.gridsize(1)
+            gs2 = cuda.blockDim.x * cuda.gridDim.x
+            if i1 != i2:
+                grid_error[0] = 1
+            if gs1 != gs2:
+                gridsize_error[0] = 1
+
+        grid_error = np.zeros(1, dtype=np.uint64)
+        gridsize_error = np.zeros(1, dtype=np.uint64)
+
+        # A lerge enough grid for thread IDs to overflow an int32
path: numba/cuda/tests/cudapy/test_intrinsics.py
text: @user1:
```suggestion
        # A large enough grid for thread IDs to overflow an int32
```