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

#pragma once

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

#include <cstdint>
#include <cuda_bf16.h>
#include <cuda_fp16.h>

namespace cugraph::ops {

/**
 * @brief Computes the forward pass for a vanilla aggregation (agg_simple)
 *        operating on a graph in a node-to-node reduction (n2n)
 *
 * @note given dimensions must be non-0.
 *
 * @param[out] out             the output aggregated embeddings. [on device]
 *                             [dim = `graph.n_dst_nodes x dim_node`]
 * @param[out] out_pos         the position of min/max values [on device]
 *                             [dim = `graph.n_dst_nodes x dim_node`]
 *                             Only required for AggOpT::kMin and AggOpT::kMax,
 *                             may be `nullptr` otherwise.
 * @param[in]  node_feat       the input node features. [on device]
 *                             [dim = `graph.n_dst_nodes x dim_node`]
 * @param[in]  dim_node        dimensionality of the node features
 * @param[in]  graph           the input graph
 * @param[in]  op              aggregation operation
 * @param[in]  stream          cuda stream
 * @{
 */
void agg_simple_n2n_fwd(float* out,
                        int32_t* out_pos,
                        const float* node_feat,
                        size_t dim_node,
                        const csc_s32_t& graph,
                        AggOpT op,
                        const cuda::stream& stream);
void agg_simple_n2n_fwd(float* out,
                        int64_t* out_pos,
                        const float* node_feat,
                        size_t dim_node,
                        const csc_s64_t& graph,
                        AggOpT op,
                        const cuda::stream& stream);

void agg_simple_n2n_fwd(__half* out,
                        int32_t* out_pos,
                        const __half* node_feat,
                        size_t dim_node,
                        const csc_s32_t& graph,
                        AggOpT op,
                        const cuda::stream& stream);
void agg_simple_n2n_fwd(__half* out,
                        int64_t* out_pos,
                        const __half* node_feat,
                        size_t dim_node,
                        const csc_s64_t& graph,
                        AggOpT op,
                        const cuda::stream& stream);

void agg_simple_n2n_fwd(__nv_bfloat16* out,
                        int32_t* out_pos,
                        const __nv_bfloat16* node_feat,
                        size_t dim_node,
                        const csc_s32_t& graph,
                        AggOpT op,
                        const cuda::stream& stream);
void agg_simple_n2n_fwd(__nv_bfloat16* out,
                        int64_t* out_pos,
                        const __nv_bfloat16* node_feat,
                        size_t dim_node,
                        const csc_s64_t& graph,
                        AggOpT op,
                        const cuda::stream& stream);

void agg_simple_n2n_fwd(float* out,
                        int32_t* out_pos,
                        const float* node_feat,
                        size_t dim_node,
                        const bipartite_csc_s32_t& graph,
                        AggOpT op,
                        const cuda::stream& stream);
void agg_simple_n2n_fwd(float* out,
                        int64_t* out_pos,
                        const float* node_feat,
                        size_t dim_node,
                        const bipartite_csc_s64_t& graph,
                        AggOpT op,
                        const cuda::stream& stream);

void agg_simple_n2n_fwd(__half* out,
                        int32_t* out_pos,
                        const __half* node_feat,
                        size_t dim_node,
                        const bipartite_csc_s32_t& graph,
                        AggOpT op,
                        const cuda::stream& stream);
void agg_simple_n2n_fwd(__half* out,
                        int64_t* out_pos,
                        const __half* node_feat,
                        size_t dim_node,
                        const bipartite_csc_s64_t& graph,
                        AggOpT op,
                        const cuda::stream& stream);

void agg_simple_n2n_fwd(__nv_bfloat16* out,
                        int32_t* out_pos,
                        const __nv_bfloat16* node_feat,
                        size_t dim_node,
                        const bipartite_csc_s32_t& graph,
                        AggOpT op,
                        const cuda::stream& stream);
void agg_simple_n2n_fwd(__nv_bfloat16* out,
                        int64_t* out_pos,
                        const __nv_bfloat16* node_feat,
                        size_t dim_node,
                        const bipartite_csc_s64_t& graph,
                        AggOpT op,
                        const cuda::stream& stream);
/** @} */

/**
 * @brief Computes the backward pass for a vanilla aggregation (agg_simple)
 *        operating on a graph in a node-to-node reduction (n2n)
 *
 * @note Dimensions passed must be the same as in forward, regardless of whether
 *       the gradient for a given feature is actually required or not.
 *
 * @param[out] d_node_feat     the output gradients of input node features
 *                             [on device] [dim = `.n_dst_nodes x dim_node`]
 * @param[in]  dout            the input node feature gradients. [on device]
 *                             [dim = `.n_dst_nodes x dim_node`]
 * @param[in]  out_pos         the position of min/max values [on device]
 *                             [dim = `graph.n_dst_nodes x dim_node`]
 *                             Only required for AggOpT::kMin and AggOpT::kMax,
 *                             may be `nullptr` otherwise.
 * @param[in]  dim_node        dimensionality of the node features
 * @param[in]  graph           the input graph
 * @param[in]  op              aggregation operation
 * @param[in]  stream          cuda stream
 * @{
 */
void agg_simple_n2n_bwd(float* d_node_feat,
                        const float* dout,
                        const int32_t* out_pos,
                        size_t dim_node,
                        const csc_s32_t& graph,
                        AggOpT op,
                        const cuda::stream& stream);
void agg_simple_n2n_bwd(float* d_node_feat,
                        const float* dout,
                        const int64_t* out_pos,
                        size_t dim_node,
                        const csc_s64_t& graph,
                        AggOpT op,
                        const cuda::stream& stream);

void agg_simple_n2n_bwd(__half* d_node_feat,
                        const __half* dout,
                        const int32_t* out_pos,
                        size_t dim_node,
                        const csc_s32_t& graph,
                        AggOpT op,
                        const cuda::stream& stream);
void agg_simple_n2n_bwd(__half* d_node_feat,
                        const __half* dout,
                        const int64_t* out_pos,
                        size_t dim_node,
                        const csc_s64_t& graph,
                        AggOpT op,
                        const cuda::stream& stream);

void agg_simple_n2n_bwd(__nv_bfloat16* d_node_feat,
                        const __nv_bfloat16* dout,
                        const int32_t* out_pos,
                        size_t dim_node,
                        const csc_s32_t& graph,
                        AggOpT op,
                        const cuda::stream& stream);
void agg_simple_n2n_bwd(__nv_bfloat16* d_node_feat,
                        const __nv_bfloat16* dout,
                        const int64_t* out_pos,
                        size_t dim_node,
                        const csc_s64_t& graph,
                        AggOpT op,
                        const cuda::stream& stream);

void agg_simple_n2n_bwd(float* d_node_feat,
                        const float* dout,
                        const int32_t* out_pos,
                        size_t dim_node,
                        const bipartite_csc_s32_t& graph,
                        AggOpT op,
                        const cuda::stream& stream);
void agg_simple_n2n_bwd(float* d_node_feat,
                        const float* dout,
                        const int64_t* out_pos,
                        size_t dim_node,
                        const bipartite_csc_s64_t& graph,
                        AggOpT op,
                        const cuda::stream& stream);

void agg_simple_n2n_bwd(__half* d_node_feat,
                        const __half* dout,
                        const int32_t* out_pos,
                        size_t dim_node,
                        const bipartite_csc_s32_t& graph,
                        AggOpT op,
                        const cuda::stream& stream);
void agg_simple_n2n_bwd(__half* d_node_feat,
                        const __half* dout,
                        const int64_t* out_pos,
                        size_t dim_node,
                        const bipartite_csc_s64_t& graph,
                        AggOpT op,
                        const cuda::stream& stream);

void agg_simple_n2n_bwd(__nv_bfloat16* d_node_feat,
                        const __nv_bfloat16* dout,
                        const int32_t* out_pos,
                        size_t dim_node,
                        const bipartite_csc_s32_t& graph,
                        AggOpT op,
                        const cuda::stream& stream);
void agg_simple_n2n_bwd(__nv_bfloat16* d_node_feat,
                        const __nv_bfloat16* dout,
                        const int64_t* out_pos,
                        size_t dim_node,
                        const bipartite_csc_s64_t& graph,
                        AggOpT op,
                        const cuda::stream& stream);
/** @} */

/**
 * @brief Computes the forward pass for a simple aggregation (agg_simple)
 *        using edge features in an edge-to-node reduction (e2n)
 *        operating on a graph
 *
 * @note given dimensions must be non-0.
 * @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` may
 * arbitrary.
 *
 * @param[out] out             the output aggregated embeddings. [on device]
 *                             [dim = `graph.n_dst_nodes x dim_edge`]
 * @param[out] out_pos         the position of min/max values [on device]
 *                             [dim = `graph.n_dst_nodes x dim_edge`]
 *                             Only required for AggOpT::kMin and AggOpT::kMax,
 *                             may be `nullptr` otherwise.
 * @param[in]  edge_feat       the input edge features. [on device]
 *                             [dim = `.n_edges x dim_edge`]
 * @param[in]  dim_edge        dimensionality of the edge features
 * @param[in]  graph           the input graph
 * @param[in]  op              aggregation operation
 * @param[in]  stream          cuda stream
 * @{
 */
void agg_simple_e2n_fwd(float* out,
                        int32_t* out_pos,
                        const float* edge_feat,
                        size_t dim_edge,
                        const csc_s32_t& graph,
                        AggOpT op,
                        const cuda::stream& stream);
void agg_simple_e2n_fwd(float* out,
                        int64_t* out_pos,
                        const float* edge_feat,
                        size_t dim_edge,
                        const csc_s64_t& graph,
                        AggOpT op,
                        const cuda::stream& stream);

void agg_simple_e2n_fwd(__half* out,
                        int32_t* out_pos,
                        const __half* edge_feat,
                        size_t dim_edge,
                        const csc_s32_t& graph,
                        AggOpT op,
                        const cuda::stream& stream);
void agg_simple_e2n_fwd(__half* out,
                        int64_t* out_pos,
                        const __half* edge_feat,
                        size_t dim_edge,
                        const csc_s64_t& graph,
                        AggOpT op,
                        const cuda::stream& stream);

void agg_simple_e2n_fwd(__nv_bfloat16* out,
                        int32_t* out_pos,
                        const __nv_bfloat16* edge_feat,
                        size_t dim_edge,
                        const csc_s32_t& graph,
                        AggOpT op,
                        const cuda::stream& stream);
void agg_simple_e2n_fwd(__nv_bfloat16* out,
                        int64_t* out_pos,
                        const __nv_bfloat16* edge_feat,
                        size_t dim_edge,
                        const csc_s64_t& graph,
                        AggOpT op,
                        const cuda::stream& stream);

void agg_simple_e2n_fwd(float* out,
                        int32_t* out_pos,
                        const float* edge_feat,
                        size_t dim_edge,
                        const bipartite_csc_s32_t& graph,
                        AggOpT op,
                        const cuda::stream& stream);
void agg_simple_e2n_fwd(float* out,
                        int64_t* out_pos,
                        const float* edge_feat,
                        size_t dim_edge,
                        const bipartite_csc_s64_t& graph,
                        AggOpT op,
                        const cuda::stream& stream);

void agg_simple_e2n_fwd(__half* out,
                        int32_t* out_pos,
                        const __half* edge_feat,
                        size_t dim_edge,
                        const bipartite_csc_s32_t& graph,
                        AggOpT op,
                        const cuda::stream& stream);
void agg_simple_e2n_fwd(__half* out,
                        int64_t* out_pos,
                        const __half* edge_feat,
                        size_t dim_edge,
                        const bipartite_csc_s64_t& graph,
                        AggOpT op,
                        const cuda::stream& stream);

void agg_simple_e2n_fwd(__nv_bfloat16* out,
                        int32_t* out_pos,
                        const __nv_bfloat16* edge_feat,
                        size_t dim_edge,
                        const bipartite_csc_s32_t& graph,
                        AggOpT op,
                        const cuda::stream& stream);
void agg_simple_e2n_fwd(__nv_bfloat16* out,
                        int64_t* out_pos,
                        const __nv_bfloat16* edge_feat,
                        size_t dim_edge,
                        const bipartite_csc_s64_t& graph,
                        AggOpT op,
                        const cuda::stream& stream);
/** @} */

/**
 * @brief Computes the backward pass for a simple aggregation (agg_simple)
 *        using edge features in an edge-to-node reduction (e2n)
 *        operating on a graph
 *
 * @note Dimensions passed must be the same as in forward, regardless of whether
 *       the gradient for a given feature is actually required or not.
 * @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 `d_edge_feat` may
 * arbitrary.
 *
 * @param[out] d_edge_feat     the output gradients of input edge features
 *                             [on device] [dim = `n_edges x dim_edge`]
 * @param[in]  dout            the input node feature gradients. [on device]
 *                             [dim = `graph.n_dst_nodes x dim_edge`]
 * @param[in]  out_pos         the position of min/max values [on device]
 *                             [dim = `graph.n_dst_nodes x dim_edge`]
 *                             Only required for AggOpT::kMin and AggOpT::kMax,
 *                             may be `nullptr` otherwise.
 * @param[in]  dim_edge        dimensionality of the edge features
 * @param[in]  graph           the input graph
 * @param[in]  op              aggregation operation
 * @param[in]  stream          cuda stream
 * @{
 */
void agg_simple_e2n_bwd(float* d_edge_feat,
                        const float* dout,
                        const int32_t* out_pos,
                        size_t dim_edge,
                        const csc_s32_t& graph,
                        AggOpT op,
                        const cuda::stream& stream);
void agg_simple_e2n_bwd(float* d_edge_feat,
                        const float* dout,
                        const int64_t* out_pos,
                        size_t dim_edge,
                        const csc_s64_t& graph,
                        AggOpT op,
                        const cuda::stream& stream);

void agg_simple_e2n_bwd(__half* d_edge_feat,
                        const __half* dout,
                        const int32_t* out_pos,
                        size_t dim_edge,
                        const csc_s32_t& graph,
                        AggOpT op,
                        const cuda::stream& stream);
void agg_simple_e2n_bwd(__half* d_edge_feat,
                        const __half* dout,
                        const int64_t* out_pos,
                        size_t dim_edge,
                        const csc_s64_t& graph,
                        AggOpT op,
                        const cuda::stream& stream);

void agg_simple_e2n_bwd(__nv_bfloat16* d_edge_feat,
                        const __nv_bfloat16* dout,
                        const int32_t* out_pos,
                        size_t dim_edge,
                        const csc_s32_t& graph,
                        AggOpT op,
                        const cuda::stream& stream);
void agg_simple_e2n_bwd(__nv_bfloat16* d_edge_feat,
                        const __nv_bfloat16* dout,
                        const int64_t* out_pos,
                        size_t dim_edge,
                        const csc_s64_t& graph,
                        AggOpT op,
                        const cuda::stream& stream);

void agg_simple_e2n_bwd(float* d_edge_feat,
                        const float* dout,
                        const int32_t* out_pos,
                        size_t dim_edge,
                        const bipartite_csc_s32_t& graph,
                        AggOpT op,
                        const cuda::stream& stream);
void agg_simple_e2n_bwd(float* d_edge_feat,
                        const float* dout,
                        const int64_t* out_pos,
                        size_t dim_edge,
                        const bipartite_csc_s64_t& graph,
                        AggOpT op,
                        const cuda::stream& stream);

void agg_simple_e2n_bwd(__half* d_edge_feat,
                        const __half* dout,
                        const int32_t* out_pos,
                        size_t dim_edge,
                        const bipartite_csc_s32_t& graph,
                        AggOpT op,
                        const cuda::stream& stream);
void agg_simple_e2n_bwd(__half* d_edge_feat,
                        const __half* dout,
                        const int64_t* out_pos,
                        size_t dim_edge,
                        const bipartite_csc_s64_t& graph,
                        AggOpT op,
                        const cuda::stream& stream);

void agg_simple_e2n_bwd(__nv_bfloat16* d_edge_feat,
                        const __nv_bfloat16* dout,
                        const int32_t* out_pos,
                        size_t dim_edge,
                        const bipartite_csc_s32_t& graph,
                        AggOpT op,
                        const cuda::stream& stream);
void agg_simple_e2n_bwd(__nv_bfloat16* d_edge_feat,
                        const __nv_bfloat16* dout,
                        const int64_t* out_pos,
                        size_t dim_edge,
                        const bipartite_csc_s64_t& graph,
                        AggOpT op,
                        const cuda::stream& stream);
/** @} */

/**
 * @brief Computes the forward pass for a simple aggregation (agg_simple)
 *        using edge features and node features in an edge-to-node reduction (e2n)
 *        and a node-to-node reduction (n2n) operating on a graph.
 *        The results of both individual aggregations are simply concatenated.
 *
 * @note given dimensions must be non-0.
 * @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` may
 * arbitrary.
 *
 * @param[out] out             the output aggregated embeddings. [on device]
 *                             [dim = `graph.n_dst_nodes x (dim_node + dim_edge)`]
 * @param[out] out_pos         the position of min/max values [on device]
 *                             [dim = `graph.n_dst_nodes x (dim_node + dim_edge)`]
 *                             Only required for AggOpT::kMin and AggOpT::kMax,
 *                             may be `nullptr` otherwise.
 * @param[in]  node_feat       the input node features. [on device]
 *                             [dim = `graph.n_dst_nodes x dim_node`]
 * @param[in]  edge_feat       the input edge features. [on device]
 *                             [dim = `.n_edges x dim_edge`]
 * @param[in]  dim_node        dimensionality of the node features
 * @param[in]  dim_edge        dimensionality of the edge features
 * @param[in]  graph           the input graph
 * @param[in]  op              aggregation operation
 * @param[in]  stream          cuda stream
 * @{
 */
void agg_simple_n2n_e2n_fwd(float* out,
                            int32_t* out_pos,
                            const float* node_feat,
                            const float* edge_feat,
                            size_t dim_node,
                            size_t dim_edge,
                            const csc_s32_t& graph,
                            AggOpT op,
                            const cuda::stream& stream);
void agg_simple_n2n_e2n_fwd(float* out,
                            int64_t* out_pos,
                            const float* node_feat,
                            const float* edge_feat,
                            size_t dim_node,
                            size_t dim_edge,
                            const csc_s64_t& graph,
                            AggOpT op,
                            const cuda::stream& stream);

void agg_simple_n2n_e2n_fwd(__half* out,
                            int32_t* out_pos,
                            const __half* node_feat,
                            const __half* edge_feat,
                            size_t dim_node,
                            size_t dim_edge,
                            const csc_s32_t& graph,
                            AggOpT op,
                            const cuda::stream& stream);
void agg_simple_n2n_e2n_fwd(__half* out,
                            int64_t* out_pos,
                            const __half* node_feat,
                            const __half* edge_feat,
                            size_t dim_node,
                            size_t dim_edge,
                            const csc_s64_t& graph,
                            AggOpT op,
                            const cuda::stream& stream);

void agg_simple_n2n_e2n_fwd(__nv_bfloat16* out,
                            int32_t* out_pos,
                            const __nv_bfloat16* node_feat,
                            const __nv_bfloat16* edge_feat,
                            size_t dim_node,
                            size_t dim_edge,
                            const csc_s32_t& graph,
                            AggOpT op,
                            const cuda::stream& stream);
void agg_simple_n2n_e2n_fwd(__nv_bfloat16* out,
                            int64_t* out_pos,
                            const __nv_bfloat16* node_feat,
                            const __nv_bfloat16* edge_feat,
                            size_t dim_node,
                            size_t dim_edge,
                            const csc_s64_t& graph,
                            AggOpT op,
                            const cuda::stream& stream);

void agg_simple_n2n_e2n_fwd(float* out,
                            int32_t* out_pos,
                            const float* node_feat,
                            const float* edge_feat,
                            size_t dim_node,
                            size_t dim_edge,
                            const bipartite_csc_s32_t& graph,
                            AggOpT op,
                            const cuda::stream& stream);
void agg_simple_n2n_e2n_fwd(float* out,
                            int64_t* out_pos,
                            const float* node_feat,
                            const float* edge_feat,
                            size_t dim_node,
                            size_t dim_edge,
                            const bipartite_csc_s64_t& graph,
                            AggOpT op,
                            const cuda::stream& stream);

void agg_simple_n2n_e2n_fwd(__half* out,
                            int32_t* out_pos,
                            const __half* node_feat,
                            const __half* edge_feat,
                            size_t dim_node,
                            size_t dim_edge,
                            const bipartite_csc_s32_t& graph,
                            AggOpT op,
                            const cuda::stream& stream);
void agg_simple_n2n_e2n_fwd(__half* out,
                            int64_t* out_pos,
                            const __half* node_feat,
                            const __half* edge_feat,
                            size_t dim_node,
                            size_t dim_edge,
                            const bipartite_csc_s64_t& graph,
                            AggOpT op,
                            const cuda::stream& stream);

void agg_simple_n2n_e2n_fwd(__nv_bfloat16* out,
                            int32_t* out_pos,
                            const __nv_bfloat16* node_feat,
                            const __nv_bfloat16* edge_feat,
                            size_t dim_node,
                            size_t dim_edge,
                            const bipartite_csc_s32_t& graph,
                            AggOpT op,
                            const cuda::stream& stream);
void agg_simple_n2n_e2n_fwd(__nv_bfloat16* out,
                            int64_t* out_pos,
                            const __nv_bfloat16* node_feat,
                            const __nv_bfloat16* edge_feat,
                            size_t dim_node,
                            size_t dim_edge,
                            const bipartite_csc_s64_t& graph,
                            AggOpT op,
                            const cuda::stream& stream);
/** @} */

/**
 * @brief Computes the forward pass for a simple aggregation (agg_simple)
 *        using edge features and node features in an edge-to-node reduction (e2n)
 *        and a node-to-node reduction (n2n) operating on a graph.
 *        The results of both individual aggregations are simply concatenated.
 *
 * @note Dimensions passed must be the same as in forward, regardless of whether
 *       the gradient for a given feature is actually required or not.
 * @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 `d_edge_feat` may
 * arbitrary.
 *
 * @param[out] d_node_feat     the output gradients of input node features
 *                             [on device] [dim = `graph.n_dst_nodes x dim_node`]
 * @param[out] d_edge_feat     the output gradients of input edge features
 *                             [on device] [dim = `n_edges x dim_edge`]
 * @param[in]  dout            the input node feature gradients. [on device]
 *                             [dim = `graph.n_dst_nodes x (dim_node + dim_edge)`]
 * @param[in]  out_pos         the position of min/max values [on device]
 *                             [dim = `graph.n_dst_nodes x (dim_node + dim_edge)`]
 *                             Only required for AggOpT::kMin and AggOpT::kMax,
 *                             may be `nullptr` otherwise.
 * @param[in]  dim_node        dimensionality of the node features
 * @param[in]  dim_edge        dimensionality of the edge features
 * @param[in]  graph           the input graph
 * @param[in]  op              aggregation operation
 * @param[in]  stream          cuda stream
 * @{
 */
void agg_simple_n2n_e2n_bwd(float* d_node_feat,
                            float* d_edge_feat,
                            const float* dout,
                            const int32_t* out_pos,
                            size_t dim_node,
                            size_t dim_edge,
                            const csc_s32_t& graph,
                            AggOpT op,
                            const cuda::stream& stream);
void agg_simple_n2n_e2n_bwd(float* d_node_feat,
                            float* d_edge_feat,
                            const float* dout,
                            const int64_t* out_pos,
                            size_t dim_node,
                            size_t dim_edge,
                            const csc_s64_t& graph,
                            AggOpT op,
                            const cuda::stream& stream);

void agg_simple_n2n_e2n_bwd(__half* d_node_feat,
                            __half* d_edge_feat,
                            const __half* dout,
                            const int32_t* out_pos,
                            size_t dim_node,
                            size_t dim_edge,
                            const csc_s32_t& graph,
                            AggOpT op,
                            const cuda::stream& stream);
void agg_simple_n2n_e2n_bwd(__half* d_node_feat,
                            __half* d_edge_feat,
                            const __half* dout,
                            const int64_t* out_pos,
                            size_t dim_node,
                            size_t dim_edge,
                            const csc_s64_t& graph,
                            AggOpT op,
                            const cuda::stream& stream);

void agg_simple_n2n_e2n_bwd(__nv_bfloat16* d_node_feat,
                            __nv_bfloat16* d_edge_feat,
                            const __nv_bfloat16* dout,
                            const int32_t* out_pos,
                            size_t dim_node,
                            size_t dim_edge,
                            const csc_s32_t& graph,
                            AggOpT op,
                            const cuda::stream& stream);
void agg_simple_n2n_e2n_bwd(__nv_bfloat16* d_node_feat,
                            __nv_bfloat16* d_edge_feat,
                            const __nv_bfloat16* dout,
                            const int64_t* out_pos,
                            size_t dim_node,
                            size_t dim_edge,
                            const csc_s64_t& graph,
                            AggOpT op,
                            const cuda::stream& stream);

void agg_simple_n2n_e2n_bwd(float* d_node_feat,
                            float* d_edge_feat,
                            const float* dout,
                            const int32_t* out_pos,
                            size_t dim_node,
                            size_t dim_edge,
                            const bipartite_csc_s32_t& graph,
                            AggOpT op,
                            const cuda::stream& stream);
void agg_simple_n2n_e2n_bwd(float* d_node_feat,
                            float* d_edge_feat,
                            const float* dout,
                            const int64_t* out_pos,
                            size_t dim_node,
                            size_t dim_edge,
                            const bipartite_csc_s64_t& graph,
                            AggOpT op,
                            const cuda::stream& stream);

void agg_simple_n2n_e2n_bwd(__half* d_node_feat,
                            __half* d_edge_feat,
                            const __half* dout,
                            const int32_t* out_pos,
                            size_t dim_node,
                            size_t dim_edge,
                            const bipartite_csc_s32_t& graph,
                            AggOpT op,
                            const cuda::stream& stream);
void agg_simple_n2n_e2n_bwd(__half* d_node_feat,
                            __half* d_edge_feat,
                            const __half* dout,
                            const int64_t* out_pos,
                            size_t dim_node,
                            size_t dim_edge,
                            const bipartite_csc_s64_t& graph,
                            AggOpT op,
                            const cuda::stream& stream);

void agg_simple_n2n_e2n_bwd(__nv_bfloat16* d_node_feat,
                            __nv_bfloat16* d_edge_feat,
                            const __nv_bfloat16* dout,
                            const int32_t* out_pos,
                            size_t dim_node,
                            size_t dim_edge,
                            const bipartite_csc_s32_t& graph,
                            AggOpT op,
                            const cuda::stream& stream);
void agg_simple_n2n_e2n_bwd(__nv_bfloat16* d_node_feat,
                            __nv_bfloat16* d_edge_feat,
                            const __nv_bfloat16* dout,
                            const int64_t* out_pos,
                            size_t dim_node,
                            size_t dim_edge,
                            const bipartite_csc_s64_t& graph,
                            AggOpT op,
                            const cuda::stream& stream);
/** @} */

}  // namespace cugraph::ops
