prefetch.hpp

转到此文件的文档页面。

/*
 * 版权所有 (c) 2024, NVIDIA CORPORATION.
 *
 * 根据 Apache 许可证 2.0 版（“许可证”）获得许可；
 * 除非遵守许可证，否则您不得使用此文件。
 * 您可以在以下位置获取许可证的副本：
 *
 * https://apache.ac.cn/licenses/LICENSE-2.0
 *
 * 除非适用法律要求或书面同意，否则软件
 * 根据许可证分发，并按“原样”提供，
 * 不提供任何明示或暗示的保证。
 * 有关管理权限和
 * 限制的特定语言，请参阅许可证。
 */

#pragma once

#include <rmm/cuda_device.hpp>
#include <rmm/cuda_stream_view.hpp>
#include <rmm/detail/export.hpp>
#include <rmm/error.hpp>

#include <cuda/std/span>

namespace RMM_NAMESPACE {

void prefetch(void const* ptr,
 std::size_t size,
 rmm::cuda_device_id device,
 rmm::cuda_stream_view stream)
{
 auto result = cudaMemPrefetchAsync(ptr, size, device.value(), stream.value());
 // 当非托管内存传递给 cudaMemPrefetchAsync 时会引发 InvalidValue 错误
 // 我们应该将其视为一个无操作
 if (result != cudaErrorInvalidValue && result != cudaSuccess) { RMM_CUDA_TRY(result); }
}

template <typename T>
void prefetch(cuda::std::span<T const> data,
 rmm::cuda_device_id device,
 rmm::cuda_stream_view stream)
{
 prefetch(data.data(), data.size_bytes(), device, stream);
 // end of group


} // namespace RMM_NAMESPACE