/*
 * Copyright (c) 2020-2023, 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 <cugraph-ops/cuda/stream.hpp>
#include <cugraph-ops/graph/format.hpp>
#include <cugraph-ops/operators/common.hpp>

#include <cstdint>

namespace cugraph::ops::gnn {

/**
 * @brief Computes the forward pass for a simple multi-head attention layer
 *        without using cudnn (mha_simple)
 *        operating on full graphs (fg) in a node-to-node
 *        reduction (n2n)
 *
 * @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 / 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[out] sm_scores       the output softmax scores for backward.
 *                             [on device] [dim = `2 x params.num_heads x n_edges`]
 * @param[in]  key_feat        the "key" 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]  query_feat      the "query" 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]  value_feat      the "value" 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]  score_bias      bias added to the un-normalized softmax scores
 *                             (not applied when passed as nullptr) [on device]
 *                             [dim = `params.num_heads x n_edges`]
 * @param[in]  norm_by_dim     if set, divide result of dot-product by `sqrt(dim_head)`
 * @param[in]  fg              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_simple_fg_n2n_fwd(float* out,
                           float* sm_scores,
                           const float* key_feat,
                           const float* query_feat,
                           const float* value_feat,
                           const float* score_bias,
                           bool norm_by_dim,
                           const fg_csr_s32_t& fg,
                           const mha_params& params,
                           int dim_node,
                           const cuda::stream& stream);
void mha_simple_fg_n2n_fwd(float* out,
                           float* sm_scores,
                           const float* key_feat,
                           const float* query_feat,
                           const float* value_feat,
                           const float* score_bias,
                           bool norm_by_dim,
                           const fg_csr_s64_t& fg,
                           const mha_params& params,
                           int dim_node,
                           const cuda::stream& stream);
void mha_simple_fg_n2n_fwd(float* out,
                           float* sm_scores,
                           const float* key_feat,
                           const float* query_feat,
                           const float* value_feat,
                           const float* score_bias,
                           bool norm_by_dim,
                           const fg_csr_batch_s32_t& fg,
                           const mha_params& params,
                           int dim_node,
                           const cuda::stream& stream);
void mha_simple_fg_n2n_fwd(float* out,
                           float* sm_scores,
                           const float* key_feat,
                           const float* query_feat,
                           const float* value_feat,
                           const float* score_bias,
                           bool norm_by_dim,
                           const fg_csr_batch_s64_t& fg,
                           const mha_params& params,
                           int dim_node,
                           const cuda::stream& stream);
/** @} */

/**
 * @brief Computes the forward pass for a simple multi-head attention layer
 *        without using cudnn (mha_simple)
 *        operatating on full graphs (fg) in a node-to-node
 *        reduction (n2n) incorporating edge features (efeat)
 *
 * @note The `eidx` version uses 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` may
 * 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 / 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[out] sm_scores       the output softmax scores for backward.
 *                             [on device] [dim = `2 x params.num_heads x n_edges`]
 * @param[in]  key_feat        the "key" 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]  query_feat      the "query" 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]  value_feat      the "value" 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 edge feature embeddings of all input edges [on device]
 *                             [dim = `n_edges x dim_node`]
 *                             For each edge, the features associated with
 *                             the different heads are expected to be concatenated.
 * @param[in]  score_bias      bias added to the un-normalized softmax scores
 *                             (not applied when passed as nullptr) [on device]
 *                             [dim = `params.num_heads x n_edges`]
 * @param[in]  norm_by_dim     if set, divide result of dot-product by `sqrt(dim_head)`.
 * @param[in]  fg              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_simple_fg_n2n_efeat_fwd(float* out,
                                 float* sm_scores,
                                 const float* key_feat,
                                 const float* query_feat,
                                 const float* value_feat,
                                 const float* edge_feat,
                                 const float* score_bias,
                                 bool norm_by_dim,
                                 const fg_csr_s32_t& fg,
                                 const mha_params& params,
                                 int dim_node,
                                 const cuda::stream& stream);
void mha_simple_fg_n2n_efeat_fwd(float* out,
                                 float* sm_scores,
                                 const float* key_feat,
                                 const float* query_feat,
                                 const float* value_feat,
                                 const float* edge_feat,
                                 const float* score_bias,
                                 bool norm_by_dim,
                                 const fg_csr_s64_t& fg,
                                 const mha_params& params,
                                 int dim_node,
                                 const cuda::stream& stream);
void mha_simple_fg_n2n_efeat_fwd(float* out,
                                 float* sm_scores,
                                 const float* key_feat,
                                 const float* query_feat,
                                 const float* value_feat,
                                 const float* edge_feat,
                                 const float* score_bias,
                                 bool norm_by_dim,
                                 const fg_csr_batch_s32_t& fg,
                                 const mha_params& params,
                                 int dim_node,
                                 const cuda::stream& stream);
void mha_simple_fg_n2n_efeat_fwd(float* out,
                                 float* sm_scores,
                                 const float* key_feat,
                                 const float* query_feat,
                                 const float* value_feat,
                                 const float* edge_feat,
                                 const float* score_bias,
                                 bool norm_by_dim,
                                 const fg_csr_batch_s64_t& fg,
                                 const mha_params& params,
                                 int dim_node,
                                 const cuda::stream& stream);
void mha_simple_fg_n2n_efeat_eidx_fwd(float* out,
                                      float* sm_scores,
                                      const float* key_feat,
                                      const float* query_feat,
                                      const float* value_feat,
                                      const float* edge_feat,
                                      const float* score_bias,
                                      bool norm_by_dim,
                                      const fg_csr_s32_t& fg,
                                      const mha_params& params,
                                      int dim_node,
                                      const cuda::stream& stream);
void mha_simple_fg_n2n_efeat_eidx_fwd(float* out,
                                      float* sm_scores,
                                      const float* key_feat,
                                      const float* query_feat,
                                      const float* value_feat,
                                      const float* edge_feat,
                                      const float* score_bias,
                                      bool norm_by_dim,
                                      const fg_csr_s64_t& fg,
                                      const mha_params& params,
                                      int dim_node,
                                      const cuda::stream& stream);
void mha_simple_fg_n2n_efeat_eidx_fwd(float* out,
                                      float* sm_scores,
                                      const float* key_feat,
                                      const float* query_feat,
                                      const float* value_feat,
                                      const float* edge_feat,
                                      const float* score_bias,
                                      bool norm_by_dim,
                                      const fg_csr_batch_s32_t& fg,
                                      const mha_params& params,
                                      int dim_node,
                                      const cuda::stream& stream);
void mha_simple_fg_n2n_efeat_eidx_fwd(float* out,
                                      float* sm_scores,
                                      const float* key_feat,
                                      const float* query_feat,
                                      const float* value_feat,
                                      const float* edge_feat,
                                      const float* score_bias,
                                      bool norm_by_dim,
                                      const fg_csr_batch_s64_t& fg,
                                      const mha_params& params,
                                      int dim_node,
                                      const cuda::stream& stream);
/** @} */

/**
 * @brief Computes the backward pass for a simple multi-head attention layer
 *        without using cudnn (mha_simple)
 *        operating on full graphs (fg) in a node-to-node
 *        reduction (n2n)
 *
 * @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 / params.num_heads`.
 *
 * @param[in]  grad_key_feat   the output gradient on the "key" embeddings. [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]  grad_query_feat the output gradient on the "query" embeddings [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] grad_value_feat  the output gradient on the "value" embeddings [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[out] grad_sm_scores  the gradient on softmax scores for backward.
 *                             [on device] [dim = `2 x params.num_heads x n_edges`]
 * @param[in]  grad_score_bias gradient on bias added to the un-normalized softmax scores
 *                             (not computed when passed as nullptr) [on device]
 *                             [dim = `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]  key_feat        the "key" 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]  query_feat      the "query" 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] value_feat       the "value" 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]  in_sm_scores    the output softmax scores for backward.
 *                             [on device] [dim = `2 x params.num_heads x n_edges`]
 * @param[in]  norm_by_dim     if set, divide result of dot-product by `sqrt(dim_head)`
 * @param[in]  fg              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_simple_fg_n2n_bwd(float* grad_key_feat,
                           float* grad_query_feat,
                           float* grad_value_feat,
                           float* grad_sm_scores,
                           float* grad_score_bias,
                           const float* grad_out,
                           const float* key_feat,
                           const float* query_feat,
                           const float* value_feat,
                           const float* in_sm_scores,
                           bool norm_by_dim,
                           const fg_csr_s32_t& fg,
                           const mha_params& params,
                           int dim_node,
                           const cuda::stream& stream);
void mha_simple_fg_n2n_bwd(float* grad_key_feat,
                           float* grad_query_feat,
                           float* grad_value_feat,
                           float* grad_sm_scores,
                           float* grad_score_bias,
                           const float* grad_out,
                           const float* key_feat,
                           const float* query_feat,
                           const float* value_feat,
                           const float* in_sm_scores,
                           bool norm_by_dim,
                           const fg_csr_s64_t& fg,
                           const mha_params& params,
                           int dim_node,
                           const cuda::stream& stream);
void mha_simple_fg_n2n_bwd(float* grad_key_feat,
                           float* grad_query_feat,
                           float* grad_value_feat,
                           float* grad_sm_scores,
                           float* grad_score_bias,
                           const float* grad_out,
                           const float* key_feat,
                           const float* query_feat,
                           const float* value_feat,
                           const float* in_sm_scores,
                           bool norm_by_dim,
                           const fg_csr_batch_s32_t& fg,
                           const mha_params& params,
                           int dim_node,
                           const cuda::stream& stream);
void mha_simple_fg_n2n_bwd(float* grad_key_feat,
                           float* grad_query_feat,
                           float* grad_value_feat,
                           float* grad_sm_scores,
                           float* grad_score_bias,
                           const float* grad_out,
                           const float* key_feat,
                           const float* query_feat,
                           const float* value_feat,
                           const float* in_sm_scores,
                           bool norm_by_dim,
                           const fg_csr_batch_s64_t& fg,
                           const mha_params& params,
                           int dim_node,
                           const cuda::stream& stream);
/** @} */

/**
 * @brief Computes the backward pass for a simple multi-head attention layer
 *        without using cudnn (mha_simple)
 *        operatating on full graphs (fg) in a node-to-node
 *        reduction (n2n) incorporating edge features
 *
 * @note The `eidx` version uses 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` may
 * 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 / params.num_heads`.
 *
 * @param[in]  grad_key_feat   the output gradient on the "key" embeddings. [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]  grad_query_feat the output gradient on the "query" embeddings [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] grad_value_feat  the output gradient on the "value" embeddings [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] grad_edge_feat   the output gradient on the edge feature embeddings [on device]
 *                             [dim = `fg.n_nodes x dim_node`]
 *                             For each edge, the features associated with
 *                             the different heads are expected to be concatenated.
 * @param[out] grad_sm_scores  the gradient on softmax scores for backward.
 *                             [on device] [dim = `2 x params.num_heads x n_edges`]
 * @param[in]  grad_score_bias gradient on bias added to the un-normalized softmax scores
 *                             (not computed when passed as nullptr) [on device]
 *                             [dim = `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]  key_feat        the "key" 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]  query_feat      the "query" 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] value_feat       the "value" 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 output gradient on the edge feature embeddings [on device]
 *                             [dim = `n_edges x dim_node`]
 *                             For each edge, the features associated with
 *                             the different heads are expected to be concatenated.
 * @param[in]  in_sm_scores    the output softmax scores for backward.
 *                             [on device] [dim = `2 x params.num_heads x n_edges`]
 * @param[in]  norm_by_dim     if set, divide result of dot-product by `sqrt(dim_head)`
 * @param[in]  fg              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_simple_fg_n2n_efeat_bwd(float* grad_key_feat,
                                 float* grad_query_feat,
                                 float* grad_value_feat,
                                 float* grad_edge_feat,
                                 float* grad_sm_scores,
                                 float* grad_score_bias,
                                 const float* grad_out,
                                 const float* key_feat,
                                 const float* query_feat,
                                 const float* value_feat,
                                 const float* edge_feat,
                                 const float* in_sm_scores,
                                 bool norm_by_dim,
                                 const fg_csr_s32_t& fg,
                                 const mha_params& params,
                                 int dim_node,
                                 const cuda::stream& stream);
void mha_simple_fg_n2n_efeat_bwd(float* grad_key_feat,
                                 float* grad_query_feat,
                                 float* grad_value_feat,
                                 float* grad_edge_feat,
                                 float* grad_sm_scores,
                                 float* grad_score_bias,
                                 const float* grad_out,
                                 const float* key_feat,
                                 const float* query_feat,
                                 const float* value_feat,
                                 const float* edge_feat,
                                 const float* in_sm_scores,
                                 bool norm_by_dim,
                                 const fg_csr_s64_t& fg,
                                 const mha_params& params,
                                 int dim_node,
                                 const cuda::stream& stream);
void mha_simple_fg_n2n_efeat_bwd(float* grad_key_feat,
                                 float* grad_query_feat,
                                 float* grad_value_feat,
                                 float* grad_edge_feat,
                                 float* grad_sm_scores,
                                 float* grad_score_bias,
                                 const float* grad_out,
                                 const float* key_feat,
                                 const float* query_feat,
                                 const float* value_feat,
                                 const float* edge_feat,
                                 const float* in_sm_scores,
                                 bool norm_by_dim,
                                 const fg_csr_batch_s32_t& fg,
                                 const mha_params& params,
                                 int dim_node,
                                 const cuda::stream& stream);
void mha_simple_fg_n2n_efeat_bwd(float* grad_key_feat,
                                 float* grad_query_feat,
                                 float* grad_value_feat,
                                 float* grad_edge_feat,
                                 float* grad_sm_scores,
                                 float* grad_score_bias,
                                 const float* grad_out,
                                 const float* key_feat,
                                 const float* query_feat,
                                 const float* value_feat,
                                 const float* edge_feat,
                                 const float* in_sm_scores,
                                 bool norm_by_dim,
                                 const fg_csr_batch_s64_t& fg,
                                 const mha_params& params,
                                 int dim_node,
                                 const cuda::stream& stream);
void mha_simple_fg_n2n_efeat_eidx_bwd(float* grad_key_feat,
                                      float* grad_query_feat,
                                      float* grad_value_feat,
                                      float* grad_edge_feat,
                                      float* grad_sm_scores,
                                      float* grad_score_bias,
                                      const float* grad_out,
                                      const float* key_feat,
                                      const float* query_feat,
                                      const float* value_feat,
                                      const float* edge_feat,
                                      const float* in_sm_scores,
                                      bool norm_by_dim,
                                      const fg_csr_s32_t& fg,
                                      const mha_params& params,
                                      int dim_node,
                                      const cuda::stream& stream);
void mha_simple_fg_n2n_efeat_eidx_bwd(float* grad_key_feat,
                                      float* grad_query_feat,
                                      float* grad_value_feat,
                                      float* grad_edge_feat,
                                      float* grad_sm_scores,
                                      float* grad_score_bias,
                                      const float* grad_out,
                                      const float* key_feat,
                                      const float* query_feat,
                                      const float* value_feat,
                                      const float* edge_feat,
                                      const float* in_sm_scores,
                                      bool norm_by_dim,
                                      const fg_csr_s64_t& fg,
                                      const mha_params& params,
                                      int dim_node,
                                      const cuda::stream& stream);
void mha_simple_fg_n2n_efeat_eidx_bwd(float* grad_key_feat,
                                      float* grad_query_feat,
                                      float* grad_value_feat,
                                      float* grad_edge_feat,
                                      float* grad_sm_scores,
                                      float* grad_score_bias,
                                      const float* grad_out,
                                      const float* key_feat,
                                      const float* query_feat,
                                      const float* value_feat,
                                      const float* edge_feat,
                                      const float* in_sm_scores,
                                      bool norm_by_dim,
                                      const fg_csr_batch_s32_t& fg,
                                      const mha_params& params,
                                      int dim_node,
                                      const cuda::stream& stream);
void mha_simple_fg_n2n_efeat_eidx_bwd(float* grad_key_feat,
                                      float* grad_query_feat,
                                      float* grad_value_feat,
                                      float* grad_edge_feat,
                                      float* grad_sm_scores,
                                      float* grad_score_bias,
                                      const float* grad_out,
                                      const float* key_feat,
                                      const float* query_feat,
                                      const float* value_feat,
                                      const float* edge_feat,
                                      const float* in_sm_scores,
                                      bool norm_by_dim,
                                      const fg_csr_batch_s64_t& fg,
                                      const mha_params& params,
                                      int dim_node,
                                      const cuda::stream& stream);
/** @} */

}  // namespace cugraph::ops::gnn
