Add Stream::memsetAsync methods + tests

include/cudawrappers/cu.hpp

@@ -693,8 +693,20 @@ class Stream : public Wrapper<CUstream> {
     checkCudaCall(cuMemPrefetchAsync(devPtr, size, dstDevice, _obj));
   }
 
+  void memsetAsync(DeviceMemory &devPtr, unsigned char value, size_t size) {
+    checkCudaCall(cuMemsetD8Async(devPtr, value, size, _obj));
+  }
+
+  void memsetAsync(DeviceMemory &devPtr, unsigned short value, size_t size) {
+    checkCudaCall(cuMemsetD16Async(devPtr, value, size, _obj));
+  }
+
+  void memsetAsync(DeviceMemory &devPtr, unsigned int value, size_t size) {
+    checkCudaCall(cuMemsetD32Async(devPtr, value, size, _obj));
+  }
+
   void zero(DeviceMemory &devPtr, size_t size) {
-    checkCudaCall(cuMemsetD8Async(devPtr, 0, size, _obj));
+    memsetAsync(devPtr, static_cast<unsigned char>(0), size);
   }
 
   void launchKernel(Function &function, unsigned gridX, unsigned gridY,

tests/test_cu.cpp

@@ -195,7 +195,32 @@ TEMPLATE_LIST_TEST_CASE("Test memset", "[memset]", TestTypes) {
   cu::Device device(0);
   cu::Context context(CU_CTX_SCHED_BLOCKING_SYNC, device);
 
-  SECTION("Test memset cu::DeviceMemory synchronously") {
+  SECTION("Test memset cu::DeviceMemory asynchronously") {
+    const size_t N = 3;
+    const size_t size = N * sizeof(TestType);
+    cu::HostMemory a(size);
+    cu::HostMemory b(size);
+    TestType value = 0xAA;
+
+    // Populate the memory with values
+    TestType* const a_ptr = static_cast<TestType*>(a);
+    TestType* const b_ptr = static_cast<TestType*>(b);
+    for (int i = 0; i < N; i++) {
+      a_ptr[i] = 0;
+      b_ptr[i] = value;
+    }
+    cu::DeviceMemory mem(size);
+
+    cu::Stream stream;
+    stream.memcpyHtoDAsync(mem, a, size);
+    stream.memsetAsync(mem, value, N);
+    stream.memcpyDtoHAsync(b, mem, size);
+    stream.synchronize();
+
+    CHECK(static_cast<bool>(memcmp(a, b, size)));
+  }
+
+  SECTION("Test zeroing cu::DeviceMemory synchronously") {
     const size_t N = 3;
     const size_t size = N * sizeof(TestType);
     cu::HostMemory a(size);