Add RUNTIME_FUNCTION annotations to cub's device MR.

Author: alliepiper

cub/cub/detail/device_memory_resource.cuh

@@ -13,10 +13,15 @@
 #  pragma system_header
 #endif // no system header
 
+#include <cub/util_device.cuh>
+
 #include <cuda/__runtime/api_wrapper.h>
 #include <cuda/__stream/stream_ref.h>
+#include <cuda/std/__exception/terminate.h>
 #include <cuda/std/cstdint>
 
+#include <nv/target>
+
 CUB_NAMESPACE_BEGIN
 
 namespace detail
@@ -25,34 +30,51 @@ namespace detail
 // this implementation
 struct device_memory_resource
 {
-  void* allocate(size_t bytes, size_t /* alignment */)
+  CUB_RUNTIME_FUNCTION void* allocate(size_t bytes, size_t /* alignment */)
   {
     void* ptr{nullptr};
-    _CCCL_TRY_CUDA_API(::cudaMallocAsync, "allocate failed to allocate with cudaMalloc", &ptr, bytes, NULL);
+    NV_IF_TARGET(
+      NV_IS_HOST,
+      (_CCCL_TRY_CUDA_API(::cudaMallocAsync, "allocate failed to allocate with cudaMallocAsync", &ptr, bytes, NULL);),
+      (_CubLog("%s\n", "cub::detail::device_memory_resource::allocate not supported from device code.");
+       ::cuda::std::terminate();));
     _CCCL_ASSERT(ptr != nullptr, "allocate failed to allocate with cudaMallocAsync");
     return ptr;
   }
 
-  void deallocate(void* ptr, size_t /* bytes */)
+  CUB_RUNTIME_FUNCTION void deallocate(void* ptr, size_t /* bytes */)
   {
-    _CCCL_TRY_CUDA_API(::cudaFree, "deallocate failed", ptr);
+    NV_IF_TARGET( //
+      NV_IS_HOST,
+      (_CCCL_TRY_CUDA_API(::cudaFree, "deallocate failed", ptr);),
+      (_CubLog("%s\n", "cub::detail::device_memory_resource::deallocate not supported from device code.");
+       ::cuda::std::terminate();));
   }
 
-  void* allocate(::cuda::stream_ref stream, size_t bytes, size_t /* alignment */)
+  CUB_RUNTIME_FUNCTION void* allocate(::cuda::stream_ref stream, size_t bytes, size_t /* alignment */)
   {
     return allocate(stream, bytes);
   }
 
-  void* allocate(::cuda::stream_ref stream, size_t bytes)
+  CUB_RUNTIME_FUNCTION void* allocate(::cuda::stream_ref stream, size_t bytes)
   {
     void* ptr{nullptr};
-    _CCCL_TRY_CUDA_API(::cudaMallocAsync, "allocate failed to allocate with cudaMallocAsync", &ptr, bytes, stream.get());
+    NV_IF_TARGET( //
+      NV_IS_HOST,
+      (_CCCL_TRY_CUDA_API(
+         ::cudaMallocAsync, "allocate failed to allocate with cudaMallocAsync", &ptr, bytes, stream.get());),
+      (_CubLog("%s\n", "cub::detail::device_memory_resource::allocate not supported from device code.");
+       ::cuda::std::terminate();));
     return ptr;
   }
 
-  void deallocate(::cuda::stream_ref stream, void* ptr, size_t /* bytes */)
+  CUB_RUNTIME_FUNCTION void deallocate(::cuda::stream_ref stream, void* ptr, size_t /* bytes */)
   {
-    _CCCL_TRY_CUDA_API(::cudaFreeAsync, "deallocate failed", ptr, stream.get());
+    NV_IF_TARGET( //
+      NV_IS_HOST,
+      (_CCCL_TRY_CUDA_API(::cudaFreeAsync, "deallocate failed", ptr, stream.get());),
+      (_CubLog("%s\n", "cub::detail::device_memory_resource::deallocate not supported from device code.");
+       ::cuda::std::terminate();));
   }
 };
 } // namespace detail