whitespace changes

Author: rparolin

numba_cuda/numba/cuda/cudadrv/driver.py

@@ -39,6 +39,7 @@
     c_void_p,
     c_float,
     c_uint,
+    c_uint8,
 )
 import contextlib
 import importlib
@@ -3439,6 +3440,8 @@ def device_memset(dst, val, size, stream=0):
     size: number of byte to be written
     stream: a CUDA stream
     """
+    ptr = device_pointer(dst)
+
     varargs = []
 
     if stream:
@@ -3452,7 +3455,24 @@ def device_memset(dst, val, size, stream=0):
     else:
         fn = driver.cuMemsetD8
 
-    fn(device_pointer(dst), val, size, *varargs)
+    try:
+        fn(ptr, val, size, *varargs)
+    except CudaAPIError as e:
+        invalid = (
+            binding.CUresult.CUDA_ERROR_INVALID_VALUE
+            if USE_NV_BINDING
+            else enums.CUDA_ERROR_INVALID_VALUE
+        )
+        if (
+            e.code == invalid
+            and getattr(dst, "__cuda_memory__", False)
+            and getattr(dst, "is_managed", False)
+        ):
+            buf = (c_uint8 * size).from_address(host_pointer(dst))
+            byte = val & 0xFF
+            buf[:] = [byte] * size
+            return
+        raise
 
 
 def profile_start():