/*
 * 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 <cugraph-ops/operators/common.hpp>

#include <cstdint>

namespace cugraph::ops {

/**
 * @brief Computes the forward pass for a "GAT-like" multi-head attention layer
 *        (mha_gat) operating on full graphs (fg) in a node-to-node
 *        reduction (n2n).
 *
 * @note If `fg.ef_indices` are set, then we use this as an edge index to obtain
 * the index of an edge feature given the index of an edge.
 * @note With an arbitrary edge feature index, the size of `edge_feat` is
 * arbitrary.
 * @note When specifying shapes of buffers, we'll refer to
 * both `fg.n_indices` for a single graph or `fg_batch.n_edges`
 * for a batch of graphs simply as `n_edges`.
 * @note We always assume `dim_head = dim_node / params.num_heads`.
 *
 * @param[out] out             the output embeddings after attention.
 *                             [on device] [dim = `fg.n_nodes x dim_out`].
 *                             For each node, the features associated with
 *                             the different heads are concatenated if
 *                             params.concat_heads = true (`dim_out = dim_node`),
 *                             and averaged otherwise (`dim_out = dim_head`).
 * @param[inout] sm_scores     the output softmax scores for the backward pass.
 *                             [on device] [dim = `2 x params.num_heads x n_edges`]
 * @param[in]  node_feat       the embeddings of all input nodes. [on device]
 *                             [dim = `fg.n_nodes x dim_node`]
 *                             For each node, the features associated with
 *                             the different heads are expected to be concatenated.
 * @param[in]  attn_weights    the weights applied to in/out feat [on device]
 *                             [dim = `dim_node + dim_node`]
 * @param[in]  graph           the input CSR graph.
 * @param[in]  params          the multi-head attention parameters
 * @param[in]  dim_node        dimension of the input embeddings.
 *                             It must be a multiple of `params.num_heads`.
 * @param[in]  stream          cuda stream
 * @{
 */
void mha_gat_fg_n2n_fwd(float* out,
                        float* sm_scores,
                        const float* node_feat,
                        const float* attn_weights,
                        const fg_csr_s32_t& graph,
                        const mha_params& params,
                        int dim_node,
                        const cuda::stream& stream);
void mha_gat_fg_n2n_fwd(float* out,
                        float* sm_scores,
                        const float* node_feat,
                        const float* attn_weights,
                        const fg_csr_s64_t& graph,
                        const mha_params& params,
                        int dim_node,
                        const cuda::stream& stream);
void mha_gat_fg_n2n_fwd(float* out,
                        float* sm_scores,
                        const float* node_feat,
                        const float* attn_weights,
                        const fg_csr_batch_s32_t& graph,
                        const mha_params& params,
                        int dim_node,
                        const cuda::stream& stream);
void mha_gat_fg_n2n_fwd(float* out,
                        float* sm_scores,
                        const float* node_feat,
                        const float* attn_weights,
                        const fg_csr_batch_s64_t& graph,
                        const mha_params& params,
                        int dim_node,
                        const cuda::stream& stream);
/** @} */

/**
 * @brief Computes the forward pass for a "GAT-like" multi-head attention layer
 *        (mha_gat) operating on full graphs (fg) in a node-to-node reduction (n2n)
 *        but using edge features, too, for computing the dot product (efeat).
 *
 * @note If `fg.ef_indices` are set, then we use this as an edge index to obtain
 * the index of an edge feature given the index of an edge.
 * @note With an arbitrary edge feature index, the size of `edge_feat` is
 * arbitrary.
 * @note When specifying shapes of buffers, we'll refer to
 * both `fg.n_indices` for a single graph or `fg_batch.n_edges`
 * for a batch of graphs simply as `n_edges`.
 * @note We always assume `dim_head = dim_node / params.num_heads`.
 *
 * @param[out] out             the output embeddings after attention.
 *                             [on device] [dim = `fg.n_nodes x dim_out`].
 *                             For each node, the features associated with
 *                             the different heads are concatenated if
 *                             params.concat_heads = true (`dim_out = dim_node`),
 *                             and averaged otherwise (`dim_out = dim_head`).
 * @param[inout] sm_scores     the output softmax scores for the backward pass.
 *                             [on device] [dim = `2 x params.num_heads x n_edges`]
 * @param[in]  node_feat       the embeddings of all input nodes. [on device]
 *                             [dim = `fg.n_nodes x dim_node`]
 *                             For each node, the features associated with
 *                             the different heads are expected to be concatenated.
 * @param[in]  edge_feat       the embeddings of all edges [on device]
 *                             [dim = `n_edges x dim_edge`]
 *                             For each edge, the features associated with the different
 *                             heads are expected to be concatenated.
 * @param[in]  attn_weights    the weights applied to in/out feat and edge features [on device]
 *                             [dim = `dim_node + dim_node + dim_edge`]
 * @param[in]  graph           the input CSR graph.
 * @param[in]  params          the multi-head attention parameters
 * @param[in]  dim_node        dimension of the input embeddings.
 *                             It must be a multiple of `params.num_heads`.
 * @param[in]  dim_edge        dimension of the edge embeddings.
 *                             It must be a multiple of `params.num_heads`.
 * @param[in]  stream          cuda stream
 * @{
 */
void mha_gat_fg_n2n_efeat_fwd(float* out,
                              float* sm_scores,
                              const float* node_feat,
                              const float* edge_feat,
                              const float* attn_weights,
                              const fg_csr_s32_t& graph,
                              const mha_params& params,
                              int dim_node,
                              int dim_edge,
                              const cuda::stream& stream);
void mha_gat_fg_n2n_efeat_fwd(float* out,
                              float* sm_scores,
                              const float* node_feat,
                              const float* edge_feat,
                              const float* attn_weights,
                              const fg_csr_s64_t& graph,
                              const mha_params& params,
                              int dim_node,
                              int dim_edge,
                              const cuda::stream& stream);
void mha_gat_fg_n2n_efeat_fwd(float* out,
                              float* sm_scores,
                              const float* node_feat,
                              const float* edge_feat,
                              const float* attn_weights,
                              const fg_csr_batch_s32_t& graph,
                              const mha_params& params,
                              int dim_node,
                              int dim_edge,
                              const cuda::stream& stream);
void mha_gat_fg_n2n_efeat_fwd(float* out,
                              float* sm_scores,
                              const float* node_feat,
                              const float* edge_feat,
                              const float* attn_weights,
                              const fg_csr_batch_s64_t& graph,
                              const mha_params& params,
                              int dim_node,
                              int dim_edge,
                              const cuda::stream& stream);
/** @} */

/**
 * @brief Computes the backward pass for a "GAT-like" multi-head attention layer
 *        (mha_gat) operating on full graphs (fg) in a node-to-node reduction (n2n).
 *
 * @note If `fg.ef_indices` are set, then we use this as an edge index to obtain
 * the index of an edge feature given the index of an edge.
 * @note With an arbitrary edge feature index, the size of `edge_feat` is
 * arbitrary.
 * @note When specifying shapes of buffers, we'll refer to
 * both `fg.n_indices` for a single graph or `fg_batch.n_edges`
 * for a batch of graphs simply as `n_edges`.
 * @note We always assume `dim_head = dim_node / params.num_heads`.
 * @note `grad_node_feat` and `grad_weights` may be set to `nullptr` if not needed.
 *
 * @param[out] grad_node_feat   the gradients on input embeddings of all nodes
 *                              [on device] [dim = `fg.n_nodes x dim_node`].
 * @param[out] grad_weights     the gradients on all attention weights [on device]
 *                              [dim = `dim_node + dim_node`]
 * @param[inout] grad_sm_scores buffer for the gradient on softmax scores.
 *                              [on device] [dim = `2 x params.num_heads x n_edges`]
 * @param[in]  grad_out         the gradient on the output of the forward pass
 *                              [on device] [dim = `fg.n_nodes x dim_out`]
 *                              For each node, the gradients associated with
 *                              the different heads are assumed to be
 *                              separate if params.concat_heads = true (`dim_out = dim_node`),
 *                              and shared otherwise (`dim_out = dim_head`).
 * @param[in]  node_feat        the embeddings of all input nodes. [on device]
 *                              [dim = `fg.n_nodes x dim_node`]
 *                              For each node, the features associated with
 *                              the different heads are expected to be concatenated.
 * @param[in]  attn_weights     the weights applied to in/out feat [on device]
 *                              [dim = `dim_node + dim_node`]
 * @param[inout]  in_sm_scores  the output softmax scores from the forward pass.
 *                              [on device] [dim = `2 x params.num_heads x n_edges`]
 * @param[in]  graph            the input CSR graph.
 * @param[in]  params           the multi-head attention parameters
 * @param[in]  dim_node         dimension of the input embeddings.
 *                              It must be a multiple of `params.num_heads`.
 * @param[in]  stream           cuda stream
 * @{
 */
void mha_gat_fg_n2n_bwd(float* grad_node_feat,
                        float* grad_weights,
                        float* grad_sm_scores,
                        const float* grad_out,
                        const float* node_feat,
                        const float* attn_weights,
                        const float* in_sm_scores,
                        const fg_csr_s32_t& graph,
                        const mha_params& params,
                        int dim_node,
                        const cuda::stream& stream);
void mha_gat_fg_n2n_bwd(float* grad_node_feat,
                        float* grad_weights,
                        float* grad_sm_scores,
                        const float* grad_out,
                        const float* node_feat,
                        const float* attn_weights,
                        const float* in_sm_scores,
                        const fg_csr_s64_t& graph,
                        const mha_params& params,
                        int dim_node,
                        const cuda::stream& stream);
void mha_gat_fg_n2n_bwd(float* grad_node_feat,
                        float* grad_weights,
                        float* grad_sm_scores,
                        const float* grad_out,
                        const float* node_feat,
                        const float* attn_weights,
                        const float* in_sm_scores,
                        const fg_csr_batch_s32_t& graph,
                        const mha_params& params,
                        int dim_node,
                        const cuda::stream& stream);
void mha_gat_fg_n2n_bwd(float* grad_node_feat,
                        float* grad_weights,
                        float* grad_sm_scores,
                        const float* grad_out,
                        const float* node_feat,
                        const float* attn_weights,
                        const float* in_sm_scores,
                        const fg_csr_batch_s64_t& graph,
                        const mha_params& params,
                        int dim_node,
                        const cuda::stream& stream);
/** @} */

/**
 * @brief Computes the backward pass for a "GAT-like" multi-head attention layer
 *        (mha_gat) operating on full graphs (fg) in a node-to-node reduction (n2n)
 *        but using edge features, too, for computing the dot product (efeat).
 *
 * @note If `fg.ef_indices` are set, then we use this as an edge index to obtain
 * the index of an edge feature given the index of an edge.
 * @note With an arbitrary edge feature index, the size of `edge_feat` is
 * arbitrary.
 * @note When specifying shapes of buffers, we'll refer to
 * both `fg.n_indices` for a single graph or `fg_batch.n_edges`
 * for a batch of graphs simply as `n_edges`.
 * @note We always assume `dim_head = dim_node / params.num_heads`.
 * @note `grad_node_feat`, `grad_edge_feat` and `grad_weights` may be set to
 * `nullptr` if not needed.
 *
 * @param[out] grad_node_feat   the gradients on input embeddings of all nodes
 *                              [on device] [dim = `fg.n_nodes x dim_node`].
 * @param[out] grad_edge_feat   the gradient on embeddings of all edges [on device]
 *                              [dim = `n_edges x dim_edge`]
 * @param[out] grad_weights     the gradients on all attention weights [on device]
 *                              [dim = `dim_node + dim_node + dim_edge`]
 * @param[inout] grad_sm_scores the gradient on softmax scores.
 *                              [on device] [dim = `2 x params.num_heads x n_edges`]
 * @param[in]  grad_out         the gradient on the output of the forward pass
 *                              [on device] [dim = `fg.n_nodes x dim_node`]
 * @param[in]  node_feat        the embeddings of all input nodes.
 *                              [on device] [dim = `fg.n_nodes x dim_out`]
 *                              For each node, the gradients associated with
 *                              the different heads are assumed to be
 *                              separate if params.concat_heads = true (`dim_out = dim_node`),
 *                              and shared otherwise (`dim_out = dim_head`).
 * @param[in]  edge_feat        the embeddings of all edges [on device]
 *                              [dim = `n_edges x dim_edge`]
 *                              For each edge, the features associated with the different
 *                              heads are expected to be concatenated.
 * @param[in]  attn_weights     the weights applied to in/out feat [on device]
 *                              [dim = `dim_node + dim_node + dim_edge`]
 * @param[inout]  in_sm_scores  the output softmax scores from the forward pass.
 *                              [on device] [dim = `2 x params.num_heads x n_edges`]
 * @param[in]  graph            the input CSR graph.
 * @param[in]  params           the multi-head attention parameters
 * @param[in]  dim_node         dimension of the input embeddings.
 *                              It must be a multiple of `params.num_heads`.
 * @param[in]  dim_edge         dimension of the edge embeddings
 *                              It must be a multiple of `params.num_heads`.
 * @param[in]  stream           cuda stream
 * @{
 */
void mha_gat_fg_n2n_efeat_bwd(float* grad_node_feat,
                              float* grad_edge_feat,
                              float* grad_weights,
                              float* grad_sm_scores,
                              const float* grad_out,
                              const float* node_feat,
                              const float* edge_feat,
                              const float* attn_weights,
                              const float* in_sm_scores,
                              const fg_csr_s32_t& graph,
                              const mha_params& params,
                              int dim_node,
                              int dim_edge,
                              const cuda::stream& stream);
void mha_gat_fg_n2n_efeat_bwd(float* grad_node_feat,
                              float* grad_edge_feat,
                              float* grad_weights,
                              float* grad_sm_scores,
                              const float* grad_out,
                              const float* node_feat,
                              const float* edge_feat,
                              const float* attn_weights,
                              const float* in_sm_scores,
                              const fg_csr_s64_t& graph,
                              const mha_params& params,
                              int dim_node,
                              int dim_edge,
                              const cuda::stream& stream);
void mha_gat_fg_n2n_efeat_bwd(float* grad_node_feat,
                              float* grad_edge_feat,
                              float* grad_weights,
                              float* grad_sm_scores,
                              const float* grad_out,
                              const float* node_feat,
                              const float* edge_feat,
                              const float* attn_weights,
                              const float* in_sm_scores,
                              const fg_csr_batch_s32_t& graph,
                              const mha_params& params,
                              int dim_node,
                              int dim_edge,
                              const cuda::stream& stream);
void mha_gat_fg_n2n_efeat_bwd(float* grad_node_feat,
                              float* grad_edge_feat,
                              float* grad_weights,
                              float* grad_sm_scores,
                              const float* grad_out,
                              const float* node_feat,
                              const float* edge_feat,
                              const float* attn_weights,
                              const float* in_sm_scores,
                              const fg_csr_batch_s64_t& graph,
                              const mha_params& params,
                              int dim_node,
                              int dim_edge,
                              const cuda::stream& stream);
/** @} */

}  // namespace cugraph::ops
