/*
 * Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved.
 *
 * This source code and/or documentation ("Licensed Deliverables") are
 * subject to NVIDIA intellectual property rights under U.S. and
 * international Copyright laws.
 */

#pragma once

#include <cugraph-ops/cuda/stream.hpp>
#include <cugraph-ops/graph/format.hpp>

#include <cuda_runtime_api.h>

#include <cstdint>

namespace cugraph::ops {

/**
 * @brief gathers input embeddings based on the sampled input `indices` array
 *
 * This means that the resulting node indices in the sampled graph based on this
 * `indices` array are assumed to be numbered contiguously from 0 to
 * `n_indices -1`.
 *
 * @param[out] out       output dense array. [on device]
 *                       [dim = `n_indices x n`] [row major]
 * @param[in]  in        input larger array from where to gather vectors.
 *                       [on device] [dim = `M x n`] [row major]
 * @param[in]  indices   input indices array which is used to gather the vectors
 *                       [on device] [len = `n_indices`] [row major]
 * @param[in]  n         vector dimension
 * @param[in]  n_indices length of the `indices` array
 * @param[in]  stream    cuda stream to schedule work on
 *
 * @{
 */
void gather(float* out,
            const float* in,
            const int32_t* indices,
            int32_t n,
            int32_t n_indices,
            const cuda::stream& stream);
void gather(float* out,
            const float* in,
            const int64_t* indices,
            int64_t n,
            int64_t n_indices,
            const cuda::stream& stream);
void gather(int32_t* out,
            const int32_t* in,
            const int32_t* indices,
            int32_t n,
            int32_t n_indices,
            const cuda::stream& stream);
void gather(int32_t* out,
            const int32_t* in,
            const int64_t* indices,
            int64_t n,
            int64_t n_indices,
            const cuda::stream& stream);
/** @} */

}  // namespace cugraph::ops
