utils.hpp

/*
 * Copyright (c) 2021-2025, NVIDIA CORPORATION.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#pragma once
 
#include <chrono>
#include <cstring>
#include <future>
#include <optional>
#include <stdexcept>
#include <tuple>
#include <type_traits>
 
#include <kvikio/error.hpp>
#include <kvikio/shim/cuda.hpp>
 
namespace kvikio {
 
// cuFile defines a page size to 4 KiB
inline constexpr std::size_t page_size = 4096;
 
[[nodiscard]] off_t convert_size2off(std::size_t x);
 
[[nodiscard]] ssize_t convert_size2ssize(std::size_t x);
 
[[nodiscard]] CUdeviceptr convert_void2deviceptr(void const* devPtr);
 
template <typename T, std::enable_if_t<std::is_integral_v<T>>* = nullptr>
[[nodiscard]] std::int64_t convert_to_64bit(T value)
{
  if constexpr (std::numeric_limits<T>::max() > std::numeric_limits<std::int64_t>::max()) {
    KVIKIO_EXPECT(value <= std::numeric_limits<std::int64_t>::max(),
                  "convert_to_64bit(x): x too large to fit std::int64_t",
                  std::overflow_error);
  }
  return std::int64_t(value);
}
 
[[nodiscard]] inline std::uint64_t convert_to_64bit(std::uint64_t value) { return value; }
 
template <typename T, std::enable_if_t<std::is_floating_point_v<T>>* = nullptr>
[[nodiscard]] double convert_to_64bit(T value)
{
  return double(value);
}
 
#ifdef KVIKIO_CUDA_FOUND
bool is_host_memory(void const* ptr);
#else
constexpr bool is_host_memory(void const* ptr) { return true; }
#endif
 
[[nodiscard]] int get_device_ordinal_from_pointer(CUdeviceptr dev_ptr);
 
[[nodiscard]] KVIKIO_EXPORT CUcontext get_primary_cuda_context(int ordinal);
 
[[nodiscard]] std::optional<CUcontext> get_context_associated_pointer(CUdeviceptr dev_ptr);
 
[[nodiscard]] bool current_context_can_access_pointer(CUdeviceptr dev_ptr);
 
[[nodiscard]] CUcontext get_context_from_pointer(void const* devPtr);
 
class PushAndPopContext {
 private:
  CUcontext _ctx;
 
 public:
  PushAndPopContext(CUcontext ctx);
  PushAndPopContext(PushAndPopContext const&)            = delete;
  PushAndPopContext& operator=(PushAndPopContext const&) = delete;
  PushAndPopContext(PushAndPopContext&&)                 = delete;
  PushAndPopContext&& operator=(PushAndPopContext&&)     = delete;
  ~PushAndPopContext();
};
 
// Find the base and offset of the memory allocation `devPtr` is in
std::tuple<void*, std::size_t, std::size_t> get_alloc_info(void const* devPtr,
                                                           CUcontext* ctx = nullptr);
 
template <typename T>
std::future<std::decay_t<T>> make_ready_future(T&& t)
{
  std::promise<std::decay_t<T>> p;
  auto fut = p.get_future();
  p.set_value(std::forward<T>(t));
  return fut;
}
 
template <typename T>
bool is_future_done(T const& future)
{
  KVIKIO_EXPECT(future.valid(),
                "The future object does not refer to a valid shared state.",
                std::invalid_argument);
  return future.wait_for(std::chrono::seconds(0)) != std::future_status::timeout;
}
 
}  // namespace kvikio