Merge pull request #6828 from gmarkall/grm-issue-6827

Fix regression in CUDA: Set stream in mapped and managed array device_setup

numba/cuda/cudadrv/devicearray.py

@@ -681,6 +681,7 @@ class MappedNDArray(DeviceNDArrayBase, np.ndarray):
 
     def device_setup(self, gpu_data, stream=0):
         self.gpu_data = gpu_data
+        self.stream = stream
 
 
 class ManagedNDArray(DeviceNDArrayBase, np.ndarray):
@@ -690,6 +691,7 @@ class ManagedNDArray(DeviceNDArrayBase, np.ndarray):
 
     def device_setup(self, gpu_data, stream=0):
         self.gpu_data = gpu_data
+        self.stream = stream
 
 
 def from_array_like(ary, stream=0, gpu_data=None):

numba/cuda/tests/cudapy/test_cuda_array_interface.py

@@ -4,7 +4,7 @@
 from numba import cuda
 from numba.cuda.testing import unittest, ContextResettingTestCase, ForeignArray
 from numba.cuda.testing import skip_on_cudasim, skip_if_external_memmgr
-from numba.tests.support import override_config
+from numba.tests.support import linux_only, override_config
 from unittest.mock import call, patch
 
 
@@ -261,12 +261,32 @@ def test_produce_no_stream(self):
         c_arr = cuda.device_array(10)
         self.assertIsNone(c_arr.__cuda_array_interface__['stream'])
 
+        mapped_arr = cuda.mapped_array(10)
+        self.assertIsNone(mapped_arr.__cuda_array_interface__['stream'])
+
+    @linux_only
+    def test_produce_managed_no_stream(self):
+        managed_arr = cuda.managed_array(10)
+        self.assertIsNone(managed_arr.__cuda_array_interface__['stream'])
+
     def test_produce_stream(self):
         s = cuda.stream()
         c_arr = cuda.device_array(10, stream=s)
         cai_stream = c_arr.__cuda_array_interface__['stream']
         self.assertEqual(s.handle.value, cai_stream)
 
+        s = cuda.stream()
+        mapped_arr = cuda.mapped_array(10, stream=s)
+        cai_stream = mapped_arr.__cuda_array_interface__['stream']
+        self.assertEqual(s.handle.value, cai_stream)
+
+    @linux_only
+    def test_produce_managed_stream(self):
+        s = cuda.stream()
+        managed_arr = cuda.managed_array(10, stream=s)
+        cai_stream = managed_arr.__cuda_array_interface__['stream']
+        self.assertEqual(s.handle.value, cai_stream)
+
     def test_consume_no_stream(self):
         # Create a foreign array with no stream
         f_arr = ForeignArray(cuda.device_array(10))