/*
 * 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 <cstdint>
#include <cuda_bf16.h>
#include <cuda_fp16.h>

namespace cugraph::ops {

/**
 * @brief Computes the forward pass for a concetenation of edge features
 * features of the source nodes, and the features of the target nodes
 * to create new edge features (update_efeat_cat) for a static graph
 * in an edge-to-edge fashion (e2e)
 *
 * @note If `dim_edge` is passed as `0`, the edge embeddings are ignored for this
 * operation. In this case, `edge_feat` must also be passed as `nullptr`.
 * @note `node_feat` can't be passed as `nullptr`.
 *
 * @param[out] out             the output embeddings after concatenation.
 *                             [on device]
 *                             [dim = `graph.n_indices x (dim_edge + dim_node + dim_node)`].
 * @param[in]  edge_feat       the input edge embeddings [on device]
 *                             [dim = `graph.n_indices x dim_edge`]
 * @param[in]  node_feat       the input node embeddings
 *                             [on device] [dim = `graph.n_dst_nodes x dim_node`]
 * @param[in]  graph           the input graph (in-graph, CSC).
 * @param[in]  dim_edge        dimension of the input edge embeddings.
 * @param[in]  dim_node        dimension of the input node embeddings.
 * @param[in]  use_src_feat    flag indicating whether to source source embeddings
 * @param[in]  use_dst_feat    flag indicating whether to destination source embeddings
 * @param[in]  stream          cuda stream
 * @{
 */
void update_efeat_e2e_concat_fwd(float* out,
                                 const float* edge_feat,
                                 const float* node_feat,
                                 const csc_s32_t& graph,
                                 size_t dim_edge,
                                 size_t dim_node,
                                 bool use_src_feat,
                                 bool use_dst_feat,
                                 const cuda::stream& stream);
void update_efeat_e2e_concat_fwd(float* out,
                                 const float* edge_feat,
                                 const float* node_feat,
                                 const csc_s64_t& graph,
                                 size_t dim_edge,
                                 size_t dim_node,
                                 bool use_src_feat,
                                 bool use_dst_feat,
                                 const cuda::stream& stream);

void update_efeat_e2e_concat_fwd(__half* out,
                                 const __half* edge_feat,
                                 const __half* node_feat,
                                 const csc_s32_t& graph,
                                 size_t dim_edge,
                                 size_t dim_node,
                                 bool use_src_feat,
                                 bool use_dst_feat,
                                 const cuda::stream& stream);
void update_efeat_e2e_concat_fwd(__half* out,
                                 const __half* edge_feat,
                                 const __half* node_feat,
                                 const csc_s64_t& graph,
                                 size_t dim_edge,
                                 size_t dim_node,
                                 bool use_src_feat,
                                 bool use_dst_feat,
                                 const cuda::stream& stream);

void update_efeat_e2e_concat_fwd(__nv_bfloat16* out,
                                 const __nv_bfloat16* edge_feat,
                                 const __nv_bfloat16* node_feat,
                                 const csc_s32_t& graph,
                                 size_t dim_edge,
                                 size_t dim_node,
                                 bool use_src_feat,
                                 bool use_dst_feat,
                                 const cuda::stream& stream);
void update_efeat_e2e_concat_fwd(__nv_bfloat16* out,
                                 const __nv_bfloat16* edge_feat,
                                 const __nv_bfloat16* node_feat,
                                 const csc_s64_t& graph,
                                 size_t dim_edge,
                                 size_t dim_node,
                                 bool use_src_feat,
                                 bool use_dst_feat,
                                 const cuda::stream& stream);
/** @} */

/**
 * @brief Computes the forward pass for a summation of edge features
 * features of the source nodes, and the features of the target nodes
 * to create new edge features (update_efeat_sum) for a static graph
 * in an edge-to-edge fashion (e2e)
 *
 *
 * @note If `edge_feat` is passed as `nullptr`, it is not considered for this operation.
 * In any case, `dim_edge` can't be passed as `0` as it also indicates the dimension
 * of the node features and the output.
 * @note `node_feat` can't be passed as `nullptr`.
 *
 * @param[out] out             the output embeddings after summation.
 *                             [on device] [dim = `graph.n_indices x dim_edge`].
 * @param[in]  edge_feat       the input edge embeddings [on device]
 *                             [dim = `graph.n_indices x dim_edge`]
 * @param[in]  node_feat       the input node embeddings
 *                             [on device] [dim = `graph.n_dst_nodes x dim_edge`]
 * @param[in]  graph           the input graph (in-graph, CSC).
 * @param[in]  dim_edge        dimension of the input edge embeddings.
 * @param[in]  use_src_feat    flag indicating whether to source source embeddings
 * @param[in]  use_dst_feat    flag indicating whether to destination source embeddings
 * @param[in]  stream          cuda stream
 * @{
 */
void update_efeat_e2e_sum_fwd(float* out,
                              const float* edge_feat,
                              const float* node_feat,
                              const csc_s32_t& graph,
                              size_t dim_edge,
                              bool use_src_feat,
                              bool use_dst_feat,
                              const cuda::stream& stream);
void update_efeat_e2e_sum_fwd(float* out,
                              const float* edge_feat,
                              const float* node_feat,
                              const csc_s64_t& graph,
                              size_t dim_edge,
                              bool use_src_feat,
                              bool use_dst_feat,
                              const cuda::stream& stream);

void update_efeat_e2e_sum_fwd(__half* out,
                              const __half* edge_feat,
                              const __half* node_feat,
                              const csc_s32_t& graph,
                              size_t dim_edge,
                              bool use_src_feat,
                              bool use_dst_feat,
                              const cuda::stream& stream);
void update_efeat_e2e_sum_fwd(__half* out,
                              const __half* edge_feat,
                              const __half* node_feat,
                              const csc_s64_t& graph,
                              size_t dim_edge,
                              bool use_src_feat,
                              bool use_dst_feat,
                              const cuda::stream& stream);

void update_efeat_e2e_sum_fwd(__nv_bfloat16* out,
                              const __nv_bfloat16* edge_feat,
                              const __nv_bfloat16* node_feat,
                              const csc_s32_t& graph,
                              size_t dim_edge,
                              bool use_src_feat,
                              bool use_dst_feat,
                              const cuda::stream& stream);
void update_efeat_e2e_sum_fwd(__nv_bfloat16* out,
                              const __nv_bfloat16* edge_feat,
                              const __nv_bfloat16* node_feat,
                              const csc_s64_t& graph,
                              size_t dim_edge,
                              bool use_src_feat,
                              bool use_dst_feat,
                              const cuda::stream& stream);
/** @} */

/**
 * @brief Computes the backward pass for the concat-op of edge features
 * features of the source nodes, and the features of the target nodes
 * to create new edge features (update_efeat_cat) for a static graph
 * in an edge-to-edge fashion (e2e)
 *
 * @note If a gradient w.r.t. a specific input embedding, e.g. `grad_edge_feat`,
 * is not needed, the pointer can be passed as `nullptr` and the gradient computation
 * will be skipped. Have in mind that the corresponding dimension still might be necessary,
 * see note above and forward for details on that.
 *
 * @param[out] grad_edge_feat  the output gradient on the edge embeddings [on device]
 *                             [dim = `graph.n_indices x dim_edge`]
 * @param[out] grad_node_feat  the output gradient on the node embeddings
 *                             [on device] [dim = `graph.n_dst_nodes x dim_node`]
 * @param[in]  grad_out        the input gradient ofn the output embeddings after concatenation.
 *                             [on device]
 *                             [dim = `graph.n_indices x (dim_edge + dim_node + dim_node)`].
 * @param[in]  graph           the input graph (in-graph, CSC).
 * @param[in]  dim_edge        dimension of the input edge embeddings.
 * @param[in]  dim_node        dimension of the input node embeddings.
 * @param[in]  use_src_feat    flag indicating whether to source source embeddings
 * @param[in]  use_dst_feat    flag indicating whether to destination source embeddings
 * @param[in]  stream          cuda stream
 * @{
 */
void update_efeat_e2e_concat_bwd(float* grad_edge_feat,
                                 float* grad_node_feat,
                                 const float* grad_out,
                                 const csc_s32_t& graph,
                                 size_t dim_edge,
                                 size_t dim_node,
                                 bool use_src_feat,
                                 bool use_dst_feat,
                                 const cuda::stream& stream);
void update_efeat_e2e_concat_bwd(float* grad_edge_feat,
                                 float* grad_node_feat,
                                 const float* grad_out,
                                 const csc_s64_t& graph,
                                 size_t dim_edge,
                                 size_t dim_node,
                                 bool use_src_feat,
                                 bool use_dst_feat,
                                 const cuda::stream& stream);

void update_efeat_e2e_concat_bwd(__half* grad_edge_feat,
                                 __half* grad_node_feat,
                                 const __half* grad_out,
                                 const csc_s32_t& graph,
                                 size_t dim_edge,
                                 size_t dim_node,
                                 bool use_src_feat,
                                 bool use_dst_feat,
                                 const cuda::stream& stream);
void update_efeat_e2e_concat_bwd(__half* grad_edge_feat,
                                 __half* grad_node_feat,
                                 const __half* grad_out,
                                 const csc_s64_t& graph,
                                 size_t dim_edge,
                                 size_t dim_node,
                                 bool use_src_feat,
                                 bool use_dst_feat,
                                 const cuda::stream& stream);

void update_efeat_e2e_concat_bwd(__nv_bfloat16* grad_edge_feat,
                                 __nv_bfloat16* grad_node_feat,
                                 const __nv_bfloat16* grad_out,
                                 const csc_s32_t& graph,
                                 size_t dim_edge,
                                 size_t dim_node,
                                 bool use_src_feat,
                                 bool use_dst_feat,
                                 const cuda::stream& stream);
void update_efeat_e2e_concat_bwd(__nv_bfloat16* grad_edge_feat,
                                 __nv_bfloat16* grad_node_feat,
                                 const __nv_bfloat16* grad_out,
                                 const csc_s64_t& graph,
                                 size_t dim_edge,
                                 size_t dim_node,
                                 bool use_src_feat,
                                 bool use_dst_feat,
                                 const cuda::stream& stream);
/** @} */

/**
 * @brief Computes the backward pass for a sum-op of edge features
 * features of the source nodes, and the features of the target nodes
 * to create new edge features (update_efeat_sum) for a static graph
 * in an edge-to-edge fashion (e2e)
 *
 *
 * @note If a gradient w.r.t. a specific input embedding, e.g. `grad_edge_feat`,
 * is not needed, the pointer can be passed as `nullptr` and the gradient computation
 * will be skipped.
 * @note `dim_edge` is expected to be the same as in the forward pass.
 *
 * @param[out] grad_edge_feat  the output gradient on the edge embeddings
 *                             [on device][dim = `graph.n_indices x dim_edge`]
 * @param[out] grad_node_feat  the output gradient on the node embeddings
 *                             [on device] [dim = `graph.n_dst_nodes x dim_edge`]
 * @param[in]  grad_out        the input gradient on the output embeddings after summation.
 *                             [on device] [dim = `graph.n_indices x dim_edge`].
 * @param[in]  graph           the input graph (in-graph, CSC).
 * @param[in]  dim_edge        dimension of the input embeddings.
 * @param[in]  use_src_feat    flag indicating whether to source source embeddings
 * @param[in]  use_dst_feat    flag indicating whether to destination source embeddings
 * @param[in]  stream          cuda stream
 * @{
 */
void update_efeat_e2e_sum_bwd(float* grad_edge_feat,
                              float* grad_node_feat,
                              const float* grad_out,
                              const csc_s32_t& graph,
                              size_t dim_edge,
                              bool use_src_feat,
                              bool use_dst_feat,
                              const cuda::stream& stream);
void update_efeat_e2e_sum_bwd(float* grad_edge_feat,
                              float* grad_node_feat,
                              const float* grad_out,
                              const csc_s64_t& graph,
                              size_t dim_edge,
                              bool use_src_feat,
                              bool use_dst_feat,
                              const cuda::stream& stream);

void update_efeat_e2e_sum_bwd(__half* grad_edge_feat,
                              __half* grad_node_feat,
                              const __half* grad_out,
                              const csc_s32_t& graph,
                              size_t dim_edge,
                              bool use_src_feat,
                              bool use_dst_feat,
                              const cuda::stream& stream);
void update_efeat_e2e_sum_bwd(__half* grad_edge_feat,
                              __half* grad_node_feat,
                              const __half* grad_out,
                              const csc_s64_t& graph,
                              size_t dim_edge,
                              bool use_src_feat,
                              bool use_dst_feat,
                              const cuda::stream& stream);

void update_efeat_e2e_sum_bwd(__nv_bfloat16* grad_edge_feat,
                              __nv_bfloat16* grad_node_feat,
                              const __nv_bfloat16* grad_out,
                              const csc_s32_t& graph,
                              size_t dim_edge,
                              bool use_src_feat,
                              bool use_dst_feat,
                              const cuda::stream& stream);
void update_efeat_e2e_sum_bwd(__nv_bfloat16* grad_edge_feat,
                              __nv_bfloat16* grad_node_feat,
                              const __nv_bfloat16* grad_out,
                              const csc_s64_t& graph,
                              size_t dim_edge,
                              bool use_src_feat,
                              bool use_dst_feat,
                              const cuda::stream& stream);
/** @} */

/**
 * @brief Computes the forward pass for a concetenation of edge features
 * features of the source nodes, and the features of the target nodes
 * to create new edge features (update_efeat_cat) for a bipartite graph
 * in an edge-to-edge fashion (e2e)
 *
 * @note If any of `dim_edge`, `dim_src`, or `dim_dst` is set to `0`, the corresponding
 * device array is not considered for this operation. In those cases, the corresponding
 * feature pointer, e.g. `edge_feat`, must also be passed as `nullptr`.
 *
 * @param[out] out             the output embeddings after concatenation.
 *                             [on device]
 *                             [dim = `graph.n_indices x (dim_edge + dim_src + dim_dst)`].
 * @param[in]  edge_feat       the input edge embeddings [on device]
 *                             [dim = `graph.n_indices x dim_edge`]
 * @param[in]  src_feat        the input node embeddings
 *                             [on device] [dim = `graph.n_src_nodes x dim_src`]
 * @param[in]  dst_feat        the input destination-node embeddings
 *                             [on device] [dim = `graph.n_dst_nodes x dim_dst`]
 * @param[in]  graph           the input graph (in-graph, CSC).
 * @param[in]  dim_edge        dimension of the input edge embeddings.
 * @param[in]  dim_src         dimension of the input source node embeddings.
 * @param[in]  dim_dst         dimension of the input destinaton node embeddings.
 * @param[in]  stream          cuda stream
 * @{
 */
void update_efeat_e2e_concat_fwd(float* out,
                                 const float* edge_feat,
                                 const float* src_feat,
                                 const float* dst_feat,
                                 const bipartite_csc_s32_t& graph,
                                 size_t dim_edge,
                                 size_t dim_src,
                                 size_t dim_dst,
                                 const cuda::stream& stream);
void update_efeat_e2e_concat_fwd(float* out,
                                 const float* edge_feat,
                                 const float* src_feat,
                                 const float* dst_feat,
                                 const bipartite_csc_s64_t& graph,
                                 size_t dim_edge,
                                 size_t dim_src,
                                 size_t dim_dst,
                                 const cuda::stream& stream);

void update_efeat_e2e_concat_fwd(__half* out,
                                 const __half* edge_feat,
                                 const __half* src_feat,
                                 const __half* dst_feat,
                                 const bipartite_csc_s32_t& graph,
                                 size_t dim_edge,
                                 size_t dim_src,
                                 size_t dim_dst,
                                 const cuda::stream& stream);
void update_efeat_e2e_concat_fwd(__half* out,
                                 const __half* edge_feat,
                                 const __half* src_feat,
                                 const __half* dst_feat,
                                 const bipartite_csc_s64_t& graph,
                                 size_t dim_edge,
                                 size_t dim_src,
                                 size_t dim_dst,
                                 const cuda::stream& stream);

void update_efeat_e2e_concat_fwd(__nv_bfloat16* out,
                                 const __nv_bfloat16* edge_feat,
                                 const __nv_bfloat16* src_feat,
                                 const __nv_bfloat16* dst_feat,
                                 const bipartite_csc_s32_t& graph,
                                 size_t dim_edge,
                                 size_t dim_src,
                                 size_t dim_dst,
                                 const cuda::stream& stream);
void update_efeat_e2e_concat_fwd(__nv_bfloat16* out,
                                 const __nv_bfloat16* edge_feat,
                                 const __nv_bfloat16* src_feat,
                                 const __nv_bfloat16* dst_feat,
                                 const bipartite_csc_s64_t& graph,
                                 size_t dim_edge,
                                 size_t dim_src,
                                 size_t dim_dst,
                                 const cuda::stream& stream);
/** @} */

/**
 * @brief Computes the forward pass for a summation of edge features
 * features of the source nodes, and the features of the target nodes
 * to create new edge features (update_efeat_sum) for a bipartite graph
 * in an edge-to-edge fashion (e2e)
 *
 *
 * @note If any pointers `edge_feat`, `src_feat`, or `dst_feat` is passed as `nullptr`,
 * the corresponding features are not considered for the summmation operation.
 *
 * @param[out] out             the output embeddings after summation.
 *                             [on device] [dim = `graph.n_indices x dim_edge`].
 * @param[in]  edge_feat       the input edge embeddings [on device]
 *                             [dim = `graph.n_indices x dim_edge`]
 * @param[in]  src_feat        the input source-node embeddings
 *                             [on device] [dim = `graph.n_src_nodes x dim_edge`]
 * @param[in]  dst_feat        the input destination-node embeddings
 *                             [on device] [dim = `graph.n_dst_nodes x dim_edge`]
 * @param[in]  graph           the input graph (in-graph, CSC).
 * @param[in]  dim_edge        dimension of the input edge embeddings.
 * @param[in]  stream          cuda stream
 * @{
 */
void update_efeat_e2e_sum_fwd(float* out,
                              const float* edge_feat,
                              const float* src_feat,
                              const float* dst_feat,
                              const bipartite_csc_s32_t& graph,
                              size_t dim_edge,
                              const cuda::stream& stream);
void update_efeat_e2e_sum_fwd(float* out,
                              const float* edge_feat,
                              const float* src_feat,
                              const float* dst_feat,
                              const bipartite_csc_s64_t& graph,
                              size_t dim_edge,
                              const cuda::stream& stream);

void update_efeat_e2e_sum_fwd(__half* out,
                              const __half* edge_feat,
                              const __half* src_feat,
                              const __half* dst_feat,
                              const bipartite_csc_s32_t& graph,
                              size_t dim_edge,
                              const cuda::stream& stream);
void update_efeat_e2e_sum_fwd(__half* out,
                              const __half* edge_feat,
                              const __half* src_feat,
                              const __half* dst_feat,
                              const bipartite_csc_s64_t& graph,
                              size_t dim_edge,
                              const cuda::stream& stream);

void update_efeat_e2e_sum_fwd(__nv_bfloat16* out,
                              const __nv_bfloat16* edge_feat,
                              const __nv_bfloat16* src_feat,
                              const __nv_bfloat16* dst_feat,
                              const bipartite_csc_s32_t& graph,
                              size_t dim_edge,
                              const cuda::stream& stream);
void update_efeat_e2e_sum_fwd(__nv_bfloat16* out,
                              const __nv_bfloat16* edge_feat,
                              const __nv_bfloat16* src_feat,
                              const __nv_bfloat16* dst_feat,
                              const bipartite_csc_s64_t& graph,
                              size_t dim_edge,
                              const cuda::stream& stream);
/** @} */

/**
 * @brief Computes the backward pass for the concat-op of edge features
 * features of the source nodes, and the features of the target nodes
 * to create new edge features (update_efeat_cat) for a bipartite graph
 * in an edge-to-edge fashion (e2e)
 *
 * @note If any of `dim_edge`, `dim_src`, or `dim_dst` is set to `0` in the forward pass,
 * the corresponding device array is also not considered in the backward pass. In any case,
 * the dimensions are expected to be the same as in the forward pass.
 * @note If a gradient w.r.t. a specific input embedding, e.g. `grad_edge_feat`,
 * is not needed, the pointer can be passed as `nullptr` and the gradient computation
 * will be skipped. Have in mind that the corresponding dimension still might be necessary,
 * see note above and forward for details on that.
 *
 * @param[out] grad_edge_feat  the output gradient on the edge embeddings [on device]
 *                             [dim = `graph.n_indices x dim_edge`]
 * @param[out] grad_src_feat   the output gradient on the source-node embeddings
 *                             [on device] [dim = `graph.n_src_nodes x dim_src`]
 * @param[out] grad_dst_feat   the output gradient on the destionation-node embeddings
 *                             [on device] [dim = `graph.n_dst_nodes x dim_dst`]
 * @param[in]  grad_out        the input gradient ofn the output embeddings after concatenation.
 *                             [on device]
 *                             [dim = `graph.n_indices x (dim_edge + dim_src + dim_dst)`].
 * @param[in]  graph           the input graph (in-graph, CSC).
 * @param[in]  dim_edge        dimension of the input edge embeddings.
 * @param[in]  dim_src         dimension of the input source node embeddings.
 * @param[in]  dim_dst         dimension of the input destinaton node embeddings.
 * @param[in]  stream          cuda stream
 * @{
 */
void update_efeat_e2e_concat_bwd(float* grad_edge_feat,
                                 float* grad_src_feat,
                                 float* grad_dst_feat,
                                 const float* grad_out,
                                 const bipartite_csc_s32_t& graph,
                                 size_t dim_edge,
                                 size_t dim_src,
                                 size_t dim_dst,
                                 const cuda::stream& stream);
void update_efeat_e2e_concat_bwd(float* grad_edge_feat,
                                 float* grad_src_feat,
                                 float* grad_dst_feat,
                                 const float* grad_out,
                                 const bipartite_csc_s64_t& graph,
                                 size_t dim_edge,
                                 size_t dim_src,
                                 size_t dim_dst,
                                 const cuda::stream& stream);

void update_efeat_e2e_concat_bwd(__half* grad_edge_feat,
                                 __half* grad_src_feat,
                                 __half* grad_dst_feat,
                                 const __half* grad_out,
                                 const bipartite_csc_s32_t& graph,
                                 size_t dim_edge,
                                 size_t dim_src,
                                 size_t dim_dst,
                                 const cuda::stream& stream);
void update_efeat_e2e_concat_bwd(__half* grad_edge_feat,
                                 __half* grad_src_feat,
                                 __half* grad_dst_feat,
                                 const __half* grad_out,
                                 const bipartite_csc_s64_t& graph,
                                 size_t dim_edge,
                                 size_t dim_src,
                                 size_t dim_dst,
                                 const cuda::stream& stream);

void update_efeat_e2e_concat_bwd(__nv_bfloat16* grad_edge_feat,
                                 __nv_bfloat16* grad_src_feat,
                                 __nv_bfloat16* grad_dst_feat,
                                 const __nv_bfloat16* grad_out,
                                 const bipartite_csc_s32_t& graph,
                                 size_t dim_edge,
                                 size_t dim_src,
                                 size_t dim_dst,
                                 const cuda::stream& stream);
void update_efeat_e2e_concat_bwd(__nv_bfloat16* grad_edge_feat,
                                 __nv_bfloat16* grad_src_feat,
                                 __nv_bfloat16* grad_dst_feat,
                                 const __nv_bfloat16* grad_out,
                                 const bipartite_csc_s64_t& graph,
                                 size_t dim_edge,
                                 size_t dim_src,
                                 size_t dim_dst,
                                 const cuda::stream& stream);
/** @} */

/**
 * @brief Computes the backward pass for a sum-op of edge features
 * features of the source nodes, and the features of the target nodes
 * to create new edge features (update_efeat_sum) for a bipartite graph
 * in an edge-to-edge fashion (e2e)
 *
 *
 * @note If a gradient w.r.t. a specific input embedding, e.g. `grad_edge_feat`,
 * is not needed, the pointer can be passed as `nullptr` and the gradient computation
 * will be skipped.
 * @note `dim_edge` is expected to be the same as in the forward pass.
 *
 * @param[out] grad_edge_feat  the output gradient on the edge embeddings
 *                             [on device][dim = `graph.n_indices x dim_edge`]
 * @param[out] grad_src_feat   the output gradient on the source-node embeddings
 *                             [on device] [dim = `graph.n_src_nodes x dim_edge`]
 * @param[out] grad_dst_feat   the output gradient on the destination-node embeddings
 *                             [on device] [dim = `graph.n_dst_nodes x dim_edge`]
 * @param[in]  grad_out        the input gradient on the output embeddings after summation.
 *                             [on device] [dim = `graph.n_indices x dim_edge`].
 * @param[in]  graph           the input graph (in-graph, CSC).
 * @param[in]  dim_edge        dimension of the input embeddings.
 * @param[in]  stream          cuda stream
 * @{
 */
void update_efeat_e2e_sum_bwd(float* grad_edge_feat,
                              float* grad_src_feat,
                              float* grad_dst_feat,
                              const float* grad_out,
                              const bipartite_csc_s32_t& graph,
                              size_t dim_edge,
                              const cuda::stream& stream);
void update_efeat_e2e_sum_bwd(float* grad_edge_feat,
                              float* grad_src_feat,
                              float* grad_dst_feat,
                              const float* grad_out,
                              const bipartite_csc_s64_t& graph,
                              size_t dim_edge,
                              const cuda::stream& stream);

void update_efeat_e2e_sum_bwd(__half* grad_edge_feat,
                              __half* grad_src_feat,
                              __half* grad_dst_feat,
                              const __half* grad_out,
                              const bipartite_csc_s32_t& graph,
                              size_t dim_edge,
                              const cuda::stream& stream);
void update_efeat_e2e_sum_bwd(__half* grad_edge_feat,
                              __half* grad_src_feat,
                              __half* grad_dst_feat,
                              const __half* grad_out,
                              const bipartite_csc_s64_t& graph,
                              size_t dim_edge,
                              const cuda::stream& stream);

void update_efeat_e2e_sum_bwd(__nv_bfloat16* grad_edge_feat,
                              __nv_bfloat16* grad_src_feat,
                              __nv_bfloat16* grad_dst_feat,
                              const __nv_bfloat16* grad_out,
                              const bipartite_csc_s32_t& graph,
                              size_t dim_edge,
                              const cuda::stream& stream);
void update_efeat_e2e_sum_bwd(__nv_bfloat16* grad_edge_feat,
                              __nv_bfloat16* grad_src_feat,
                              __nv_bfloat16* grad_dst_feat,
                              const __nv_bfloat16* grad_out,
                              const bipartite_csc_s64_t& graph,
                              size_t dim_edge,
                              const cuda::stream& stream);
/** @} */

}  // namespace cugraph::ops
