Repository: numba/numba

Issue:
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.


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()
```

Resolving patch:
diff --git a/numba/cuda/intrinsics.py b/numba/cuda/intrinsics.py
--- 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
--- 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
 


Test patch:
diff --git a/numba/cuda/tests/cudapy/test_intrinsics.py b/numba/cuda/tests/cudapy/test_intrinsics.py
--- a/numba/cuda/tests/cudapy/test_intrinsics.py
+++ b/numba/cuda/tests/cudapy/test_intrinsics.py
@@ -440,6 +440,31 @@ 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 large enough grid for thread IDs to overflow an int32
+        # (22121216 * 256 = 5663031296, which is greater than 2 ** 32)
+        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[:])