/*
 * Copyright (c) 2020-2024, 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 simple multi-head attention layer
 *        (mha_simple) operating on a graph in a node-to-node
 *        reduction (n2n).
 *
 * @note If `graph.map_csc_to_coo` 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 We always assume `dim_head = dim / params.num_heads`.
 * @note The activation in `params` is ignored as of now.
 *
 * @param[out] out             the output embeddings after attention.
 *                             [on device] [dim = `graph.n_dst_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 = `SZ x params.num_heads x graph.n_indices`]
 *                             with `SZ=2` if `score_weight` is passed `nullptr`, `SZ=3` otherwise.
 * @param[in]  key_feat        the "key" embeddings of all input nodes. [on device]
 *                             [dim = `graph.n_dst_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 = `graph.n_dst_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 = `graph.n_dst_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]  score_weight    weight used to scale 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]  graph           the input 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_n2n_fwd(float* out,
                        float* sm_scores,
                        const float* key_feat,
                        const float* query_feat,
                        const float* value_feat,
                        const float* score_bias,
                        const float* score_weight,
                        bool norm_by_dim,
                        const csc_s32_t& graph,
                        const mha_params& params,
                        int dim_node,
                        const cuda::stream& stream);
void mha_simple_n2n_fwd(float* out,
                        float* sm_scores,
                        const float* key_feat,
                        const float* query_feat,
                        const float* value_feat,
                        const float* score_bias,
                        const float* score_weight,
                        bool norm_by_dim,
                        const csc_s64_t& graph,
                        const mha_params& params,
                        int dim_node,
                        const cuda::stream& stream);
/** @} */

/**
 * @brief Computes the forward pass for a simple multi-head attention layer
 *        (mha_simple) operatating on a graph in a node-to-node
 *        reduction (n2n) incorporating edge features (efeat).
 *
 * @note If `graph.map_csc_to_coo` 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 We always assume `dim_head = dim / params.num_heads`.
 * @note The activation in `params` is ignored as of now.
 *
 * @param[out] out             the output embeddings after attention.
 *                             [on device] [dim = `graph.n_dst_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 = `SZ x params.num_heads x graph.n_indices`]
 *                             with `SZ=2` if `score_weight` is passed `nullptr`, `SZ=3` otherwise.
 * @param[in]  key_feat        the "key" embeddings of all input nodes. [on device]
 *                             [dim = `graph.n_dst_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 = `graph.n_dst_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 = `graph.n_dst_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]  score_weight    weight used to scale 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]  graph           the input 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_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,
                              const float* score_weight,
                              bool norm_by_dim,
                              const csc_s32_t& graph,
                              const mha_params& params,
                              int dim_node,
                              const cuda::stream& stream);
void mha_simple_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,
                              const float* score_weight,
                              bool norm_by_dim,
                              const csc_s64_t& graph,
                              const mha_params& params,
                              int dim_node,
                              const cuda::stream& stream);
/** @} */

/**
 * @brief Computes the backward pass for a simple multi-head attention layer
 *        (mha_simple) operating on a graph in a node-to-node
 *        reduction (n2n).
 *
 * @note If `graph.map_csc_to_coo` 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 We always assume `dim_head = dim / params.num_heads`.
 * @note `grad_key_feat`, `grad_query_feat` and `grad_value_feat`
 * may be set to `nullptr` if not needed.
 * @note The activation in `params` is ignored as of now.
 *
 * @param[out] grad_key_feat      the output gradient on the "key" embeddings. [on device]
 *                                [dim = `graph.n_dst_nodes x dim_node`]
 *                                For each node, the features associated with
 *                                the different heads are expected to be concatenated.
 * @param[out] grad_query_feat    the output gradient on the "query" embeddings [on device]
 *                                [dim = `graph.n_dst_nodes x dim_node`]
 *                                For each node, the features associated with
 *                                the different heads are expected to be concatenated.
 * @param[out] grad_value_feat    the output gradient on the "value" embeddings [on device]
 *                                [dim = `graph.n_dst_nodes x dim_node`]
 *                                For each node, the features associated with
 *                                the different heads are expected to be concatenated.
 * @param[inout] grad_workspace   workspace for deterministic gradient reduction,
 *                                if passed as none, non-deterministic path is chosen
 *                                [on device]
 *                                [dim = `graph.n_indices * (2 * 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[out]  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[out]  grad_score_weight gradient on weighted applied 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 = `graph.n_dst_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 = `graph.n_dst_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 = `graph.n_dst_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 = `graph.n_dst_nodes x dim_node`]
 *                                For each node, the features associated with
 *                                the different heads are expected to be concatenated.
 * @param[in]  score_weight       weight used to scale the un-normalized softmax scores
 *                                (not applied when passed as nullptr) [on device]
 *                                [dim = `params.num_heads x n_edges`]
 * @param[in]  sm_scores          the output softmax scores from the forward pass.
 *                                [on device] [dim = `SZ x params.num_heads x graph.n_indices`]
 *                                with `SZ=2` if `score_weight` is passed `nullptr`, `SZ=3`
 * otherwise.
 * @param[in]  norm_by_dim        if set, divide result of dot-product by `sqrt(dim_head)`
 * @param[in]  graph              the input 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_n2n_bwd(float* grad_key_feat,
                        float* grad_query_feat,
                        float* grad_value_feat,
                        float* grad_workspace,
                        float* grad_sm_scores,
                        float* grad_score_bias,
                        float* grad_score_weight,
                        const float* grad_out,
                        const float* key_feat,
                        const float* query_feat,
                        const float* value_feat,
                        const float* score_weight,
                        const float* sm_scores,
                        bool norm_by_dim,
                        const csc_s32_t& graph,
                        const mha_params& params,
                        int dim_node,
                        const cuda::stream& stream);
void mha_simple_n2n_bwd(float* grad_key_feat,
                        float* grad_query_feat,
                        float* grad_value_feat,
                        float* grad_workspace,
                        float* grad_sm_scores,
                        float* grad_score_bias,
                        float* grad_score_weight,
                        const float* grad_out,
                        const float* key_feat,
                        const float* query_feat,
                        const float* value_feat,
                        const float* score_weight,
                        const float* sm_scores,
                        bool norm_by_dim,
                        const csc_s64_t& graph,
                        const mha_params& params,
                        int dim_node,
                        const cuda::stream& stream);
/** @} */

/**
 * @brief Computes the backward pass for a simple multi-head attention layer
 *        (mha_simple) operatating on a graph in a node-to-node
 *        reduction (n2n) incorporating edge features.
 *
 * @note If `graph.map_csc_to_coo` 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 We always assume `dim_head = dim / params.num_heads`.
 * @note `grad_key_feat`, `grad_query_feat`, `grad_value_feat` and
 * `grad_edge_feat` may be set to `nullptr` if not needed.
 * @note The activation in `params` is ignored as of now.
 *
 * @param[out] grad_key_feat      the output gradient on the "key" embeddings. [on device]
 *                                [dim = `graph.n_dst_nodes x dim_node`]
 *                                For each node, the features associated with
 *                                the different heads are expected to be concatenated.
 * @param[out] grad_query_feat    the output gradient on the "query" embeddings [on device]
 *                                [dim = `graph.n_dst_nodes x dim_node`]
 *                                For each node, the features associated with
 *                                the different heads are expected to be concatenated.
 * @param[out] grad_value_feat    the output gradient on the "value" embeddings [on device]
 *                                [dim = `graph.n_dst_nodes x dim_node`]
 *                                For each node, the features associated with
 *                                the different heads are expected to be concatenated.
 * @param[out] grad_edge_feat     the output gradient on the edge feature embeddings [on device]
 *                                [dim = `graph.n_dst_nodes x dim_node`]
 *                                For each edge, the features associated with
 *                                the different heads are expected to be concatenated.
 * @param[inout] grad_workspace   workspace for deterministic gradient reduction,
 *                                if passed as none, non-deterministic path is chosen
 *                                [on device]
 *                                [dim = `graph.n_indices * (2 * 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[out]  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[out]  grad_score_weight gradient on weighted applied 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 = `graph.n_dst_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 = `graph.n_dst_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 = `graph.n_dst_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 = `graph.n_dst_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]  score_weight       weight used to scale the un-normalized softmax scores
 *                                (not applied when passed as nullptr) [on device]
 *                                [dim = `params.num_heads x n_edges`]
 * @param[inout]  sm_scores       the output softmax scores from the forward pass.
 *                                [on device] [dim = `SZ x params.num_heads x graph.n_indices`]
 *                                with `SZ=2` if `score_weight` is passed `nullptr`, `SZ=3`
 * otherwise.
 * @param[in]  norm_by_dim        if set, divide result of dot-product by `sqrt(dim_head)`
 * @param[in]  graph              the input 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_n2n_efeat_bwd(float* grad_key_feat,
                              float* grad_query_feat,
                              float* grad_value_feat,
                              float* grad_edge_feat,
                              float* grad_workspace,
                              float* grad_sm_scores,
                              float* grad_score_bias,
                              float* grad_score_weight,
                              const float* grad_out,
                              const float* key_feat,
                              const float* query_feat,
                              const float* value_feat,
                              const float* edge_feat,
                              const float* score_weight,
                              const float* sm_scores,
                              bool norm_by_dim,
                              const csc_s32_t& graph,
                              const mha_params& params,
                              int dim_node,
                              const cuda::stream& stream);
void mha_simple_n2n_efeat_bwd(float* grad_key_feat,
                              float* grad_query_feat,
                              float* grad_value_feat,
                              float* grad_edge_feat,
                              float* grad_workspace,
                              float* grad_sm_scores,
                              float* grad_score_bias,
                              float* grad_score_weight,
                              const float* grad_out,
                              const float* key_feat,
                              const float* query_feat,
                              const float* value_feat,
                              const float* edge_feat,
                              const float* score_weight,
                              const float* sm_scores,
                              const mha_params& params,
                              int dim_node,
                              const cuda::stream& stream);
/** @} */

/**
 * @brief Computes the forward pass for a simple multi-head attention layer
 *        (mha_simple) operating on bipartite graphs in a node-to-node
 *        reduction (n2n).
 *
 * @note If `graph.map_csc_to_coo` 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 We always assume `dim_head = dim / params.num_heads`.
 * @note The activation in `params` is ignored as of now.
 *
 * @param[out] out             the output embeddings after attention.
 *                             [on device] [dim = `graph.n_dst_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 = `SZ x params.num_heads x graph.n_indices`]
 *                             with `SZ=2` if `score_weight` is passed `nullptr`, `SZ=3` otherwise.
 * @param[in]  key_feat        the "key" embeddings of all source nodes. [on device]
 *                             [dim = `graph.n_src_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 destination nodes. [on device]
 *                             [dim = `graph.n_dst_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 source nodes. [on device]
 *                             [dim = `graph.n_src_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 graph.n_indices`]
 * @param[in]  score_weight    weight used to scale 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]  graph           the input graph (`bipartite_csc`).
 * @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_n2n_fwd(float* out,
                        float* sm_scores,
                        const float* key_feat,
                        const float* query_feat,
                        const float* value_feat,
                        const float* score_bias,
                        const float* score_weight,
                        bool norm_by_dim,
                        const bipartite_csc_s32_t& graph,
                        const mha_params& params,
                        int dim_node,
                        const cuda::stream& stream);
void mha_simple_n2n_fwd(float* out,
                        float* sm_scores,
                        const float* key_feat,
                        const float* query_feat,
                        const float* value_feat,
                        const float* score_bias,
                        const float* score_weight,
                        bool norm_by_dim,
                        const bipartite_csc_s64_t& graph,
                        const mha_params& params,
                        int dim_node,
                        const cuda::stream& stream);

/**
 * @brief Computes the forward pass for a simple multi-head attention layer
 *        (mha_simple) operatating on bipartite graphs in a node-to-node
 *        reduction (n2n) incorporating edge features (efeat).
 *
 * @note If `graph.map_csc_to_coo` 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 We always assume `dim_head = dim / params.num_heads`.
 * @note The activation in `params` is ignored as of now.
 *
 * @param[out] out             the output embeddings after attention.
 *                             [on device] [dim = `graph.n_dst_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 = `SZ x params.num_heads x graph.n_indices`]
 *                             with `SZ=2` if `score_weight` is passed `nullptr`, `SZ=3` otherwise.
 * @param[in]  key_feat        the "key" embeddings of all source nodes. [on device]
 *                             [dim = `graph.n_src_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 destination nodes. [on device]
 *                             [dim = `graph.n_dst_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 source nodes. [on device]
 *                             [dim = `graph.n_src_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 = `graph.n_indices 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 graph.n_indices`]
 * @param[in]  score_weight    weight used to scale 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]  graph           the input graph (`bipartite_csc`).
 * @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_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,
                              const float* score_weight,
                              bool norm_by_dim,
                              const bipartite_csc_s32_t& graph,
                              const mha_params& params,
                              int dim_node,
                              const cuda::stream& stream);
void mha_simple_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,
                              const float* score_weight,
                              bool norm_by_dim,
                              const bipartite_csc_s64_t& graph,
                              const mha_params& params,
                              int dim_node,
                              const cuda::stream& stream);
/** @} */

/**
 * @brief Computes the backward pass for a simple multi-head attention layer
 *        (mha_simple) operating on bipartite graphs in a node-to-node
 *        reduction (n2n).
 *
 * @note If `graph.map_csc_to_coo` 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 We always assume `dim_head = dim / params.num_heads`.
 * @note `grad_key_feat`, `grad_query_feat` and `grad_value_feat`
 * may be set to `nullptr` if not needed.
 * @note The activation in `params` is ignored as of now.
 *
 * @param[out] grad_key_feat      the output gradient on the "key" embeddings. [on device]
 *                                [dim = `graph.n_src_nodes x dim_node`]
 *                                For each node, the features associated with
 *                                the different heads are expected to be concatenated.
 * @param[out] grad_query_feat    the output gradient on the "query" embeddings [on device]
 *                                [dim = `graph.n_dst_nodes x dim_node`]
 *                                For each node, the features associated with
 *                                the different heads are expected to be concatenated.
 * @param[out] grad_value_feat    the output gradient on the "value" embeddings [on device]
 *                                [dim = `graph.n_src_nodes x dim_node`]
 *                                For each node, the features associated with
 *                                the different heads are expected to be concatenated.
 * @param[inout] grad_workspace   workspace for deterministic gradient reduction,
 *                                if passed as none, non-deterministic path is chosen
 *                                [on device]
 *                                [dim = `graph.n_indices * (2 * dim_node)]
 * @param[inout] grad_sm_scores   buffer for the gradient on softmax scores.
 *                                [on device] [dim = `2 x params.num_heads x graph.n_indices`]
 * @param[out]  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 graph.n_indices`]
 * @param[out]  grad_score_weight gradient on weighted applied 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 = `graph.n_dst_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 source nodes. [on device]
 *                                [dim = `graph.n_src_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 destination nodes. [on device]
 *                                [dim = `graph.n_dst_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 source nodes. [on device]
 *                                [dim = `graph.n_src_nodes x dim_node`]
 *                                For each node, the features associated with
 *                                the different heads are expected to be concatenated.
 * @param[in]  score_weight       weight used to scale the un-normalized softmax scores
 *                                (not applied when passed as `nullptr`) [on device]
 *                                [dim = `params.num_heads x n_edges`]
 * @param[in]  sm_scores          the output softmax scores from the forward pass.
 *                                [on device] [dim = `SZ x params.num_heads x graph.n_indices`]
 *                                with `SZ=2` if `score_weight` is passed `nullptr`, `SZ=3`
 * otherwise.
 * @param[in]  norm_by_dim        if set, divide result of dot-product by `sqrt(dim_head)`.
 * @param[in]  graph              the input graph (`bipartite_csc`).
 * @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_n2n_bwd(float* grad_key_feat,
                        float* grad_query_feat,
                        float* grad_value_feat,
                        float* grad_workspace,
                        float* grad_sm_scores,
                        float* grad_score_bias,
                        float* grad_score_weight,
                        const float* grad_out,
                        const float* key_feat,
                        const float* query_feat,
                        const float* value_feat,
                        const float* score_weight,
                        const float* sm_scores,
                        bool norm_by_dim,
                        const bipartite_csc_s32_t& graph,
                        const mha_params& params,
                        int dim_node,
                        const cuda::stream& stream);
void mha_simple_n2n_bwd(float* grad_key_feat,
                        float* grad_query_feat,
                        float* grad_value_feat,
                        float* grad_workspace,
                        float* grad_sm_scores,
                        float* grad_score_bias,
                        float* grad_score_weight,
                        const float* grad_out,
                        const float* key_feat,
                        const float* query_feat,
                        const float* value_feat,
                        const float* score_weight,
                        const float* sm_scores,
                        bool norm_by_dim,
                        const bipartite_csc_s64_t& graph,
                        const mha_params& params,
                        int dim_node,
                        const cuda::stream& stream);
/** @} */

/**
 * @brief Computes the backward pass for a simple multi-head attention layer
 *        (mha_simple) operatating on bipartite graphs in a node-to-node
 *        reduction (n2n) incorporating edge features.
 *
 * @note If `graph.map_csc_to_coo` 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 We always assume `dim_head = dim / params.num_heads`.
 * @note `grad_key_feat`, `grad_query_feat`, `grad_value_feat` and
 * `grad_edge_feat` may be set to `nullptr` if not needed.
 * @note The activation in `params` is ignored as of now.
 *
 * @param[out] grad_key_feat      the output gradient on the "key" embeddings. [on device]
 *                                [dim = `graph.n_src_nodes x dim_node`]
 *                                For each node, the features associated with
 *                                the different heads are expected to be concatenated.
 * @param[out] grad_query_feat    the output gradient on the "query" embeddings [on device]
 *                                [dim = `graph.n_dst_nodes x dim_node`]
 *                                For each node, the features associated with
 *                                the different heads are expected to be concatenated.
 * @param[out] grad_value_feat    the output gradient on the "value" embeddings [on device]
 *                                [dim = `graph.n_src_nodes x dim_node`]
 *                                For each node, the features associated with
 *                                the different heads are expected to be concatenated.
 * @param[out] grad_edge_feat     the output gradient on the edge feature embeddings [on device]
 *                                [dim = `graph.n_indices x dim_node`]
 *                                For each edge, the features associated with
 *                                the different heads are expected to be concatenated.
 * @param[inout] grad_workspace   workspace for deterministic gradient reduction,
 *                                if passed as none, non-deterministic path is chosen
 *                                [on device]
 *                                [dim = `graph.n_indices * (2 * dim_node)]
 * @param[inout] grad_sm_scores   buffer for the gradient on softmax scores.
 *                                [on device] [dim = `2 x params.num_heads x graph.n_indices`]
 * @param[out]  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 graph.n_indices`]
 * @param[out]  grad_score_weight gradient on weighted applied 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 = `graph.n_dst_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 source nodes. [on device]
 *                                [dim = `graph.n_src_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 destination nodes. [on device]
 *                                [dim = `graph.n_dst_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 source nodes. [on device]
 *                                [dim = `graph.n_src_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 = `graph.n_indices x dim_node`]
 *                                For each edge, the features associated with
 *                                the different heads are expected to be concatenated.
 * @param[in]  score_weight       weight used to scale the un-normalized softmax scores
 *                                (not applied when passed as nullptr) [on device]
 *                                [dim = `params.num_heads x n_edges`]
 * @param[in]  sm_scores          the output softmax scores from the forward pass.
 *                                [on device] [dim = `SZ x params.num_heads x graph.n_indices`]
 *                                with `SZ=2` if `score_weight` is passed `nullptr`, `SZ=3`
 * otherwise.
 * @param[in]  norm_by_dim        if set, divide result of dot-product by `sqrt(dim_head)`
 * @param[in]  graph              the input graph (`bipartite_csc`).
 * @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_n2n_efeat_bwd(float* grad_key_feat,
                              float* grad_query_feat,
                              float* grad_value_feat,
                              float* grad_edge_feat,
                              float* grad_workspace,
                              float* grad_sm_scores,
                              float* grad_score_bias,
                              float* grad_score_weight,
                              const float* grad_out,
                              const float* key_feat,
                              const float* query_feat,
                              const float* value_feat,
                              const float* edge_feat,
                              const float* score_weight,
                              const float* sm_scores,
                              bool norm_by_dim,
                              const bipartite_csc_s32_t& graph,
                              const mha_params& params,
                              int dim_node,
                              const cuda::stream& stream);
void mha_simple_n2n_efeat_bwd(float* grad_key_feat,
                              float* grad_query_feat,
                              float* grad_value_feat,
                              float* grad_edge_feat,
                              float* grad_workspace,
                              float* grad_sm_scores,
                              float* grad_score_bias,
                              float* grad_score_weight,
                              const float* grad_out,
                              const float* key_feat,
                              const float* query_feat,
                              const float* value_feat,
                              const float* edge_feat,
                              const float* score_weight,
                              const float* sm_scores,
                              bool norm_by_dim,
                              const bipartite_csc_s64_t& graph,
                              const mha_params& params,
                              int dim_node,
                              const cuda::stream& stream);
/** @} */

}  // namespace cugraph::ops
