Validate indices

Author: gmarkall

numba_cuda/numba/cuda/cache_hints.py

@@ -61,12 +61,35 @@ def stwt(array, i, value):
 }
 
 
+def _validate_arguments(instruction, array, index):
+    if not isinstance(array, types.Array):
+        msg = f"{instruction} operates on arrays. Got type {array}"
+        raise NumbaTypeError(msg)
+
+    valid_index = False
+
+    if isinstance(index, types.Integer):
+        if array.ndim != 1:
+            msg = f"Expected {array.ndim} indices, got a scalar"
+            raise NumbaTypeError(msg)
+        valid_index = True
+
+    if isinstance(index, types.UniTuple):
+        if index.count != array.ndim:
+            msg = f"Expected {array.ndim} indices, got {index.count}"
+            raise NumbaTypeError(msg)
+
+        if all([isinstance(t, types.Integer) for t in index.dtype]):
+            valid_index = True
+
+    if not valid_index:
+        raise NumbaTypeError(f"{index} is not a valid index")
+
+
 def ld_cache_operator(operator):
     @intrinsic
     def impl(typingctx, array, index):
-        if not isinstance(array, types.Array):
-            msg = f"ldcs operates on arrays. Got type {array}"
-            raise NumbaTypeError(msg)
+        _validate_arguments(f"ld{operator}", array, index)
 
         # Need to validate bitwidth
 
@@ -111,9 +134,7 @@ def codegen(context, builder, sig, args):
 def st_cache_operator(operator):
     @intrinsic
     def impl(typingctx, array, index, value):
-        if not isinstance(array, types.Array):
-            msg = f"ldcs operates on arrays. Got type {array}"
-            raise NumbaTypeError(msg)
+        _validate_arguments(f"st{operator}", array, index)
 
         # Need to validate bitwidth
 

numba_cuda/numba/cuda/tests/cudapy/test_cache_hints.py

@@ -1,8 +1,8 @@
-from numba import cuda, typeof
+from numba import cuda, errors, typeof, types
 from numba.cuda.testing import unittest, CUDATestCase
 import numpy as np
 
-types = (
+tested_types = (
     np.int8, np.int16, np.int32, np.int64,
     np.uint8, np.uint16, np.uint32, np.uint64,
     np.float16, np.float32, np.float64,
@@ -35,7 +35,7 @@ def f(r, x):
                 for i in range(len(r)):
                     r[i] = operator(x, i)
 
-            for ty in types:
+            for ty in tested_types:
                 with self.subTest(operator=operator, ty=ty):
                     x = np.arange(5).astype(ty)
                     r = np.zeros_like(x)
@@ -58,7 +58,7 @@ def f(r, x):
                 for i in range(len(r)):
                     operator(r, i, x[i])
 
-            for ty in types:
+            for ty in tested_types:
                 with self.subTest(operator=operator, ty=ty):
                     x = np.arange(5).astype(ty)
                     r = np.zeros_like(x)
@@ -74,6 +74,39 @@ def f(r, x):
 
                     self.assertIn(f"st.global.{modifier}.b{bitwidth}", ptx)
 
+    def test_bad_indices(self):
+        def float_indices(x):
+            cuda.ldcs(x, 1.0)
+
+        sig_1d = (types.float32[::1],)
+
+        msg = "float64 is not a valid index"
+        with self.assertRaisesRegex(errors.TypingError, msg):
+            cuda.compile_ptx(float_indices, sig_1d)
+
+        def too_long_indices(x):
+            cuda.ldcs(x, (1, 2))
+
+        msg = "Expected 1 indices, got 2"
+        with self.assertRaisesRegex(errors.TypingError, msg):
+            cuda.compile_ptx(too_long_indices, sig_1d)
+
+        def too_short_indices_scalar(x):
+            cuda.ldcs(x, 1)
+
+        def too_short_indices_tuple(x):
+            cuda.ldcs(x, (1,))
+
+        sig_2d = (types.float32[:,::1],)
+
+        msg = "Expected 2 indices, got a scalar"
+        with self.assertRaisesRegex(errors.TypingError, msg):
+            cuda.compile_ptx(too_short_indices_scalar, sig_2d)
+
+        msg = "Expected 2 indices, got 1"
+        with self.assertRaisesRegex(errors.TypingError, msg):
+            cuda.compile_ptx(too_short_indices_tuple, sig_2d)
+
 
 if __name__ == '__main__':
     unittest.main()