|
1 | 1 | /*
|
2 |
| - * Copyright (c) 2017-2024, NVIDIA CORPORATION. All rights reserved. |
| 2 | + * Copyright (c) 2017-2025, NVIDIA CORPORATION. All rights reserved. |
3 | 3 | *
|
4 | 4 | * Licensed under the Apache License, Version 2.0 (the "License");
|
5 | 5 | * you may not use this file except in compliance with the License.
|
|
26 | 26 |
|
27 | 27 | #include <rmm/cuda_stream_view.hpp>
|
28 | 28 | #include <rmm/mr/device/polymorphic_allocator.hpp>
|
| 29 | +#include <rmm/prefetch.hpp> |
29 | 30 |
|
30 | 31 | #include <cuda/atomic>
|
31 | 32 | #include <thrust/pair.h>
|
@@ -472,10 +473,10 @@ class concurrent_unordered_map {
|
472 | 473 | cudaError_t status = cudaPointerGetAttributes(&hashtbl_values_ptr_attributes, m_hashtbl_values);
|
473 | 474 |
|
474 | 475 | if (cudaSuccess == status && isPtrManaged(hashtbl_values_ptr_attributes)) {
|
475 |
| - RAFT_CUDA_TRY(cudaMemPrefetchAsync( |
476 |
| - m_hashtbl_values, m_capacity * sizeof(value_type), dev_id, stream.value())); |
| 476 | + rmm::prefetch( |
| 477 | + m_hashtbl_values, m_capacity * sizeof(value_type), rmm::cuda_device_id{dev_id}, stream); |
477 | 478 | }
|
478 |
| - RAFT_CUDA_TRY(cudaMemPrefetchAsync(this, sizeof(*this), dev_id, stream.value())); |
| 479 | + rmm::prefetch(this, sizeof(*this), rmm::cuda_device_id{dev_id}, stream); |
479 | 480 | }
|
480 | 481 |
|
481 | 482 | /**
|
@@ -545,8 +546,8 @@ class concurrent_unordered_map {
|
545 | 546 | if (cudaSuccess == status && isPtrManaged(hashtbl_values_ptr_attributes)) {
|
546 | 547 | int dev_id = 0;
|
547 | 548 | RAFT_CUDA_TRY(cudaGetDevice(&dev_id));
|
548 |
| - RAFT_CUDA_TRY(cudaMemPrefetchAsync( |
549 |
| - m_hashtbl_values, m_capacity * sizeof(value_type), dev_id, stream.value())); |
| 549 | + rmm::prefetch( |
| 550 | + m_hashtbl_values, m_capacity * sizeof(value_type), rmm::cuda_device_id{dev_id}, stream); |
550 | 551 | }
|
551 | 552 | }
|
552 | 553 |
|
|
0 commit comments