|
7 | 7 |
|
8 | 8 | #include <rmm/cuda_stream.hpp> |
9 | 9 | #include <rmm/detail/error.hpp> |
| 10 | +#include <rmm/mr/cuda_memory_resource.hpp> |
10 | 11 | #include <rmm/mr/pinned_host_memory_resource.hpp> |
11 | 12 | #include <rmm/resource_ref.hpp> |
12 | 13 |
|
| 14 | +#include <cuda/memory_resource> |
13 | 15 | #include <thrust/host_vector.h> |
14 | 16 |
|
15 | 17 | #include <gtest/gtest.h> |
@@ -218,3 +220,126 @@ TEST(ResourceRefConversionAllocator, VectorMove) |
218 | 220 | auto vec2 = std::move(vec1); |
219 | 221 | ASSERT_EQ(vec2[0], 42); |
220 | 222 | } |
| 223 | + |
| 224 | +// -------------------------------------------------------------------------- |
| 225 | +// Tests for forward_property adaptor with RMM resource_ref as upstream. |
| 226 | +// |
| 227 | +// This exercises the fix for https://github.com/rapidsai/rmm/issues/2322: |
| 228 | +// an unconstrained friend get_property in cccl_async_resource_ref caused |
| 229 | +// ambiguity with CCCL's default get_property for dynamic_accessibility_property. |
| 230 | +// -------------------------------------------------------------------------- |
| 231 | + |
| 232 | +// A minimal adaptor that uses cuda::forward_property with an RMM async resource ref as upstream. |
| 233 | +struct forwarding_adaptor |
| 234 | + : cuda::forward_property<forwarding_adaptor, rmm::device_async_resource_ref> { |
| 235 | + explicit forwarding_adaptor(rmm::device_async_resource_ref upstream) : upstream_{upstream} {} |
| 236 | + |
| 237 | + rmm::device_async_resource_ref upstream_resource() const { return upstream_; } |
| 238 | + |
| 239 | + void* allocate(cuda::stream_ref stream, std::size_t bytes, std::size_t alignment) |
| 240 | + { |
| 241 | + return upstream_.allocate(stream, bytes, alignment); |
| 242 | + } |
| 243 | + void deallocate(cuda::stream_ref stream, |
| 244 | + void* ptr, |
| 245 | + std::size_t bytes, |
| 246 | + std::size_t alignment) noexcept |
| 247 | + { |
| 248 | + upstream_.deallocate(stream, ptr, bytes, alignment); |
| 249 | + } |
| 250 | + void* allocate_sync(std::size_t bytes, std::size_t alignment) |
| 251 | + { |
| 252 | + return upstream_.allocate_sync(bytes, alignment); |
| 253 | + } |
| 254 | + void deallocate_sync(void* ptr, std::size_t bytes, std::size_t alignment) noexcept |
| 255 | + { |
| 256 | + upstream_.deallocate_sync(ptr, bytes, alignment); |
| 257 | + } |
| 258 | + |
| 259 | + friend bool operator==(forwarding_adaptor const& lhs, forwarding_adaptor const& rhs) |
| 260 | + { |
| 261 | + return lhs.upstream_ == rhs.upstream_; |
| 262 | + } |
| 263 | + friend bool operator!=(forwarding_adaptor const& lhs, forwarding_adaptor const& rhs) |
| 264 | + { |
| 265 | + return !(lhs == rhs); |
| 266 | + } |
| 267 | + |
| 268 | + private: |
| 269 | + rmm::device_async_resource_ref upstream_; |
| 270 | +}; |
| 271 | + |
| 272 | +// A minimal adaptor using forward_property with an RMM sync resource ref as upstream. |
| 273 | +struct forwarding_sync_adaptor |
| 274 | + : cuda::forward_property<forwarding_sync_adaptor, rmm::device_resource_ref> { |
| 275 | + explicit forwarding_sync_adaptor(rmm::device_resource_ref upstream) : upstream_{upstream} {} |
| 276 | + |
| 277 | + rmm::device_resource_ref upstream_resource() const { return upstream_; } |
| 278 | + |
| 279 | + void* allocate_sync(std::size_t bytes, std::size_t alignment) |
| 280 | + { |
| 281 | + return upstream_.allocate_sync(bytes, alignment); |
| 282 | + } |
| 283 | + void deallocate_sync(void* ptr, std::size_t bytes, std::size_t alignment) noexcept |
| 284 | + { |
| 285 | + upstream_.deallocate_sync(ptr, bytes, alignment); |
| 286 | + } |
| 287 | + |
| 288 | + friend bool operator==(forwarding_sync_adaptor const& lhs, forwarding_sync_adaptor const& rhs) |
| 289 | + { |
| 290 | + return lhs.upstream_ == rhs.upstream_; |
| 291 | + } |
| 292 | + friend bool operator!=(forwarding_sync_adaptor const& lhs, forwarding_sync_adaptor const& rhs) |
| 293 | + { |
| 294 | + return !(lhs == rhs); |
| 295 | + } |
| 296 | + |
| 297 | + private: |
| 298 | + rmm::device_resource_ref upstream_; |
| 299 | +}; |
| 300 | + |
| 301 | +// Compile-time checks: verify that the forwarding adaptor satisfies resource concepts. |
| 302 | +static_assert(cuda::has_property<forwarding_adaptor, cuda::mr::device_accessible>, |
| 303 | + "forwarding_adaptor must have device_accessible property via forward_property"); |
| 304 | + |
| 305 | +static_assert(cuda::has_property<forwarding_sync_adaptor, cuda::mr::device_accessible>, |
| 306 | + "forwarding_sync_adaptor must have device_accessible property via forward_property"); |
| 307 | + |
| 308 | +// Type-erasing a forward_property adaptor into resource_ref triggers the ambiguity |
| 309 | +// from issue #2322. If the constraint on cccl_async_resource_ref::get_property is missing, |
| 310 | +// this will fail to compile when CCCL has dynamic_accessibility_property. |
| 311 | +TEST(ForwardPropertyAdaptor, TypeEraseAsyncAdaptor) |
| 312 | +{ |
| 313 | + rmm::mr::cuda_memory_resource mr{}; |
| 314 | + rmm::device_async_resource_ref upstream{mr}; |
| 315 | + forwarding_adaptor adaptor{upstream}; |
| 316 | + cuda::mr::resource_ref<cuda::mr::device_accessible> erased{adaptor}; |
| 317 | + |
| 318 | + rmm::cuda_stream stream{}; |
| 319 | + void* ptr = erased.allocate(stream, 1024, 256); |
| 320 | + ASSERT_NE(ptr, nullptr); |
| 321 | + erased.deallocate(stream, ptr, 1024, 256); |
| 322 | +} |
| 323 | + |
| 324 | +TEST(ForwardPropertyAdaptor, TypeEraseSyncAdaptor) |
| 325 | +{ |
| 326 | + rmm::mr::cuda_memory_resource mr{}; |
| 327 | + rmm::device_resource_ref upstream{mr}; |
| 328 | + forwarding_sync_adaptor adaptor{upstream}; |
| 329 | + cuda::mr::synchronous_resource_ref<cuda::mr::device_accessible> erased{adaptor}; |
| 330 | + |
| 331 | + void* ptr = erased.allocate_sync(1024); |
| 332 | + ASSERT_NE(ptr, nullptr); |
| 333 | + erased.deallocate_sync(ptr, 1024); |
| 334 | +} |
| 335 | + |
| 336 | +// Verify that get_property still works correctly through the forwarding adaptor. |
| 337 | +TEST(ForwardPropertyAdaptor, GetPropertyDeviceAccessible) |
| 338 | +{ |
| 339 | + rmm::mr::cuda_memory_resource mr{}; |
| 340 | + rmm::device_async_resource_ref upstream{mr}; |
| 341 | + forwarding_adaptor adaptor{upstream}; |
| 342 | + |
| 343 | + // Should compile and not throw - device_accessible is a stateless property |
| 344 | + get_property(adaptor, cuda::mr::device_accessible{}); |
| 345 | +} |
0 commit comments