Skip to content

Commit 0c8b5d6

Browse files
committed
Use rmm prefetc API so we don't have to handle CUDA 12/13 API differences
1 parent 97773bd commit 0c8b5d6

File tree

1 file changed

+7
-6
lines changed

1 file changed

+7
-6
lines changed

cpp/libcugraph_etl/include/hash/concurrent_unordered_map.cuh

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright (c) 2017-2024, NVIDIA CORPORATION. All rights reserved.
2+
* Copyright (c) 2017-2025, NVIDIA CORPORATION. All rights reserved.
33
*
44
* Licensed under the Apache License, Version 2.0 (the "License");
55
* you may not use this file except in compliance with the License.
@@ -26,6 +26,7 @@
2626

2727
#include <rmm/cuda_stream_view.hpp>
2828
#include <rmm/mr/device/polymorphic_allocator.hpp>
29+
#include <rmm/prefetch.hpp>
2930

3031
#include <cuda/atomic>
3132
#include <thrust/pair.h>
@@ -472,10 +473,10 @@ class concurrent_unordered_map {
472473
cudaError_t status = cudaPointerGetAttributes(&hashtbl_values_ptr_attributes, m_hashtbl_values);
473474

474475
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);
477478
}
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);
479480
}
480481

481482
/**
@@ -545,8 +546,8 @@ class concurrent_unordered_map {
545546
if (cudaSuccess == status && isPtrManaged(hashtbl_values_ptr_attributes)) {
546547
int dev_id = 0;
547548
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);
550551
}
551552
}
552553

0 commit comments

Comments
 (0)