Replaced PyArray_SIZE(ary) with PyArray_NBYTES(ary).

pycuda.driver.register_host_memory() was passing the size of the array in elements, rather than the length of the array in bytes when creating a register_host_memory class. This resulted in "Invalid Argument" errors when calling asynchronous memory copies on the entire array.
inducer · Aug 4, 2015 · 0f28293 · 0f28293
1 parent ccb694d
commit 0f28293
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 3 deletions.
diff --git a/src/wrapper/wrap_cudadrv.cpp b/src/wrapper/wrap_cudadrv.cpp
@@ -484,7 +484,7 @@ namespace
 
     std::auto_ptr<registered_host_memory> regmem(
         new registered_host_memory(
-          PyArray_DATA(ary.ptr()), PyArray_SIZE(ary.ptr()), flags, ary));
+          PyArray_DATA(ary.ptr()), PyArray_NBYTES(ary.ptr()), flags, ary));
 
     PyObject *new_array_ptr = PyArray_FromInterface(ary.ptr());
     if (new_array_ptr == Py_NotImplemented)

diff --git a/test/test_driver.py b/test/test_driver.py
@@ -573,8 +573,16 @@ def test_register_host_memory(self):
             from py.test import skip
             skip("register_host_memory is not supported on OS X")
 
-        a = drv.aligned_empty((2**20,), np.float64, alignment=4096)
-        drv.register_host_memory(a)
+        import resource
+
+        a = drv.aligned_empty((2**20,), np.float64,
+            alignment=resource.getpagesize())
+        a_pin = drv.register_host_memory(a)
+
+        gpu_ary = drv.mem_alloc_like(a)
+        stream = drv.Stream()
+        drv.memcpy_htod_async(gpu_ary, a_pin, stream)
+        drv.Context.synchronize()
 
     @pytest.mark.xfail
     @mark_cuda_test