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>
 
#ifdef KVIKIO_CUDA_FOUND
#include <nvtx3/nvtx3.hpp>
#endif
 
#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(const void* 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()) {
    if (value > std::numeric_limits<std::int64_t>::max()) {
      throw std::overflow_error("convert_to_64bit(x): x too large to fit std::int64_t");
    }
  }
  return std::int64_t(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(const void* ptr);
#else
constexpr bool is_host_memory(const void* 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(const void* devPtr);
 
class PushAndPopContext {
 private:
  CUcontext _ctx;
 
 public:
  PushAndPopContext(CUcontext ctx);
  PushAndPopContext(const PushAndPopContext&)            = 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(const void* devPtr,
                                                           CUcontext* ctx = nullptr);
 
template <typename T>
bool is_future_done(const T& future)
{
  return future.wait_for(std::chrono::seconds(0)) != std::future_status::timeout;
}
 
#ifdef KVIKIO_CUDA_FOUND
struct libkvikio_domain {
  static constexpr char const* name{"libkvikio"};
};
 
// Macro to concatenate two tokens x and y.
#define KVIKIO_CONCAT_HELPER(x, y) x##y
#define KVIKIO_CONCAT(x, y)        KVIKIO_CONCAT_HELPER(x, y)
 
// Macro to create a static, registered string that will not have a name conflict with any
// registered string defined in the same scope.
#define KVIKIO_REGISTER_STRING(msg)                                        \
  [](const char* a_msg) -> auto& {                                         \
    static nvtx3::registered_string_in<libkvikio_domain> a_reg_str{a_msg}; \
    return a_reg_str;                                                      \
  }(msg)
 
// Macro overloads of KVIKIO_NVTX_FUNC_RANGE
#define KVIKIO_NVTX_FUNC_RANGE_IMPL() NVTX3_FUNC_RANGE_IN(libkvikio_domain)
 
#define KVIKIO_NVTX_SCOPED_RANGE_IMPL(msg, val)                                        \
  nvtx3::scoped_range_in<libkvikio_domain> KVIKIO_CONCAT(_kvikio_nvtx_range, __LINE__) \
  {                                                                                    \
    nvtx3::event_attributes                                                            \
    {                                                                                  \
      KVIKIO_REGISTER_STRING(msg), nvtx3::payload { convert_to_64bit(val) }            \
    }                                                                                  \
  }
 
#define KVIKIO_NVTX_MARKER_IMPL(msg, val) \
  nvtx3::mark_in<libkvikio_domain>(       \
    nvtx3::event_attributes{KVIKIO_REGISTER_STRING(msg), nvtx3::payload{convert_to_64bit(val)}})
 
#endif
 
#ifdef KVIKIO_CUDA_FOUND
#define KVIKIO_NVTX_FUNC_RANGE() KVIKIO_NVTX_FUNC_RANGE_IMPL()
#else
#define KVIKIO_NVTX_FUNC_RANGE(...) \
  do {                              \
  } while (0)
#endif
 
#ifdef KVIKIO_CUDA_FOUND
#define KVIKIO_NVTX_SCOPED_RANGE(msg, val) KVIKIO_NVTX_SCOPED_RANGE_IMPL(msg, val)
#else
#define KVIKIO_NVTX_SCOPED_RANGE(msg, val) \
  do {                                     \
  } while (0)
#endif
 
#ifdef KVIKIO_CUDA_FOUND
#define KVIKIO_NVTX_MARKER(message, payload) KVIKIO_NVTX_MARKER_IMPL(message, payload)
#else
#define KVIKIO_NVTX_MARKER(message, payload) \
  do {                                       \
  } while (0)
#endif
 
}  // namespace kvikio