/*
 * 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 of a simple aggregation fused with the concatenation
 *        of additional node features (agg_concat) using node features in a node-to-node
 *        reduction (n2n) operating on a full graph (fg).
 *        The "self" node features are concatenated at the end.
 *
 * @note given dimensions must be non-0.
 * @note `concat_feat` must not overlap with `node_feat` except if
 * `concat_feat == node_feat` and `dim_concat == dim_node`.
 *
 * @param[out] out             the output aggregated embeddings. [on device]
 *                             [dim = `fg.n_nodes x (dim_node + dim_concat)`]
 * @param[out] out_pos         the position of min/max values [on device]
 *                             [dim = `fg.n_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 = `fg.n_nodes x dim_node`]
 * @param[in]  concat_feat     the node features to concatenate after aggregation.
 *                             [on device] [dim = `fg.n_nodes x dim_concat`]
 * @param[in]  dim_node        dimensionality of the node features
 * @param[in]  dim_concat      dimensionality of the additional concatenated features
 * @param[in]  fg              the input CSR full graph (or batch of graphs)
 * @param[in]  op              aggregation operation
 * @param[in]  stream          cuda stream
 * @{
 */
void agg_concat_fg_n2n_fwd(float* out,
                           int32_t* out_pos,
                           const float* node_feat,
                           const float* concat_feat,
                           size_t dim_node,
                           size_t dim_concat,
                           const fg_csr_s32_t& fg,
                           AggOpT op,
                           const cuda::stream& stream);
void agg_concat_fg_n2n_fwd(float* out,
                           int32_t* out_pos,
                           const float* node_feat,
                           const float* concat_feat,
                           size_t dim_node,
                           size_t dim_concat,
                           const fg_csr_batch_s32_t& fg,
                           AggOpT op,
                           const cuda::stream& stream);
void agg_concat_fg_n2n_fwd(float* out,
                           int64_t* out_pos,
                           const float* node_feat,
                           const float* concat_feat,
                           size_t dim_node,
                           size_t dim_concat,
                           const fg_csr_s64_t& fg,
                           AggOpT op,
                           const cuda::stream& stream);
void agg_concat_fg_n2n_fwd(float* out,
                           int64_t* out_pos,
                           const float* node_feat,
                           const float* concat_feat,
                           size_t dim_node,
                           size_t dim_concat,
                           const fg_csr_batch_s64_t& fg,
                           AggOpT op,
                           const cuda::stream& stream);

void agg_concat_fg_n2n_fwd(__half* out,
                           int32_t* out_pos,
                           const __half* node_feat,
                           const __half* concat_feat,
                           size_t dim_node,
                           size_t dim_concat,
                           const fg_csr_s32_t& fg,
                           AggOpT op,
                           const cuda::stream& stream);
void agg_concat_fg_n2n_fwd(__half* out,
                           int32_t* out_pos,
                           const __half* node_feat,
                           const __half* concat_feat,
                           size_t dim_node,
                           size_t dim_concat,
                           const fg_csr_batch_s32_t& fg,
                           AggOpT op,
                           const cuda::stream& stream);
void agg_concat_fg_n2n_fwd(__half* out,
                           int64_t* out_pos,
                           const __half* node_feat,
                           const __half* concat_feat,
                           size_t dim_node,
                           size_t dim_concat,
                           const fg_csr_s64_t& fg,
                           AggOpT op,
                           const cuda::stream& stream);
void agg_concat_fg_n2n_fwd(__half* out,
                           int64_t* out_pos,
                           const __half* node_feat,
                           const __half* concat_feat,
                           size_t dim_node,
                           size_t dim_concat,
                           const fg_csr_batch_s64_t& fg,
                           AggOpT op,
                           const cuda::stream& stream);

void agg_concat_fg_n2n_fwd(__nv_bfloat16* out,
                           int32_t* out_pos,
                           const __nv_bfloat16* node_feat,
                           const __nv_bfloat16* concat_feat,
                           size_t dim_node,
                           size_t dim_concat,
                           const fg_csr_s32_t& fg,
                           AggOpT op,
                           const cuda::stream& stream);
void agg_concat_fg_n2n_fwd(__nv_bfloat16* out,
                           int32_t* out_pos,
                           const __nv_bfloat16* node_feat,
                           const __nv_bfloat16* concat_feat,
                           size_t dim_node,
                           size_t dim_concat,
                           const fg_csr_batch_s32_t& fg,
                           AggOpT op,
                           const cuda::stream& stream);
void agg_concat_fg_n2n_fwd(__nv_bfloat16* out,
                           int64_t* out_pos,
                           const __nv_bfloat16* node_feat,
                           const __nv_bfloat16* concat_feat,
                           size_t dim_node,
                           size_t dim_concat,
                           const fg_csr_s64_t& fg,
                           AggOpT op,
                           const cuda::stream& stream);
void agg_concat_fg_n2n_fwd(__nv_bfloat16* out,
                           int64_t* out_pos,
                           const __nv_bfloat16* node_feat,
                           const __nv_bfloat16* concat_feat,
                           size_t dim_node,
                           size_t dim_concat,
                           const fg_csr_batch_s64_t& fg,
                           AggOpT op,
                           const cuda::stream& stream);
/** @} */

/**
 * @brief Computes the backward pass of a simple aggregation fused with the concatenation
 *        of additional node features (agg_concat) using node features in a node-to-node
 *        reduction (n2n) operating on a full graph (fg).
 *        The "self" node features are concatenated at the end.
 *
 * @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 `d_concat_feat` must not overlap with `d_node_feat` except if
 * `d_concat_feat == d_node_feat` and `dim_concat == dim_node`.
 *
 * @param[out] d_node_feat     the output gradients of input node features
 *                             [on device] [dim = `fg.n_nodes x dim_node`].
 *                             Pass `nullptr` if not required.
 * @param[out] d_concat_feat   the output gradients of concatenated node features
 *                             [on device] [dim = `fg.n_nodes x dim_concat`].
 *                             Pass `nullptr` if not required.
 * @param[in]  dout            the input node feature gradients. [on device]
 *                             [dim = `fg.n_nodes x (dim_node + dim_concat)`]
 * @param[in]  out_pos         the position of min/max values [on device]
 *                             [dim = `fg.n_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]  dim_concat      dimensionality of the additional concatenated features
 * @param[in]  fg              the input CSR full graph (or batch of graphs)
 * @param[in]  op              aggregation operation
 * @param[in]  stream          cuda stream
 * @{
 */
void agg_concat_fg_n2n_bwd(float* d_node_feat,
                           float* d_concat_feat,
                           const float* dout,
                           const int32_t* out_pos,
                           size_t dim_node,
                           size_t dim_concat,
                           const fg_csr_s32_t& fg,
                           AggOpT op,
                           const cuda::stream& stream);
void agg_concat_fg_n2n_bwd(float* d_node_feat,
                           float* d_concat_feat,
                           const float* dout,
                           const int32_t* out_pos,
                           size_t dim_node,
                           size_t dim_concat,
                           const fg_csr_batch_s32_t& fg,
                           AggOpT op,
                           const cuda::stream& stream);
void agg_concat_fg_n2n_bwd(float* d_node_feat,
                           float* d_concat_feat,
                           const float* dout,
                           const int64_t* out_pos,
                           size_t dim_node,
                           size_t dim_concat,
                           const fg_csr_s64_t& fg,
                           AggOpT op,
                           const cuda::stream& stream);
void agg_concat_fg_n2n_bwd(float* d_node_feat,
                           float* d_concat_feat,
                           const float* dout,
                           const int64_t* out_pos,
                           size_t dim_node,
                           size_t dim_concat,
                           const fg_csr_batch_s64_t& fg,
                           AggOpT op,
                           const cuda::stream& stream);

void agg_concat_fg_n2n_bwd(__half* d_node_feat,
                           __half* d_concat_feat,
                           const __half* dout,
                           const int32_t* out_pos,
                           size_t dim_node,
                           size_t dim_concat,
                           const fg_csr_s32_t& fg,
                           AggOpT op,
                           const cuda::stream& stream);
void agg_concat_fg_n2n_bwd(__half* d_node_feat,
                           __half* d_concat_feat,
                           const __half* dout,
                           const int32_t* out_pos,
                           size_t dim_node,
                           size_t dim_concat,
                           const fg_csr_batch_s32_t& fg,
                           AggOpT op,
                           const cuda::stream& stream);
void agg_concat_fg_n2n_bwd(__half* d_node_feat,
                           __half* d_concat_feat,
                           const __half* dout,
                           const int64_t* out_pos,
                           size_t dim_node,
                           size_t dim_concat,
                           const fg_csr_s64_t& fg,
                           AggOpT op,
                           const cuda::stream& stream);
void agg_concat_fg_n2n_bwd(__half* d_node_feat,
                           __half* d_concat_feat,
                           const __half* dout,
                           const int64_t* out_pos,
                           size_t dim_node,
                           size_t dim_concat,
                           const fg_csr_batch_s64_t& fg,
                           AggOpT op,
                           const cuda::stream& stream);

void agg_concat_fg_n2n_bwd(__nv_bfloat16* d_node_feat,
                           __nv_bfloat16* d_concat_feat,
                           const __nv_bfloat16* dout,
                           const int32_t* out_pos,
                           size_t dim_node,
                           size_t dim_concat,
                           const fg_csr_s32_t& fg,
                           AggOpT op,
                           const cuda::stream& stream);
void agg_concat_fg_n2n_bwd(__nv_bfloat16* d_node_feat,
                           __nv_bfloat16* d_concat_feat,
                           const __nv_bfloat16* dout,
                           const int32_t* out_pos,
                           size_t dim_node,
                           size_t dim_concat,
                           const fg_csr_batch_s32_t& fg,
                           AggOpT op,
                           const cuda::stream& stream);
void agg_concat_fg_n2n_bwd(__nv_bfloat16* d_node_feat,
                           __nv_bfloat16* d_concat_feat,
                           const __nv_bfloat16* dout,
                           const int64_t* out_pos,
                           size_t dim_node,
                           size_t dim_concat,
                           const fg_csr_s64_t& fg,
                           AggOpT op,
                           const cuda::stream& stream);
void agg_concat_fg_n2n_bwd(__nv_bfloat16* d_node_feat,
                           __nv_bfloat16* d_concat_feat,
                           const __nv_bfloat16* dout,
                           const int64_t* out_pos,
                           size_t dim_node,
                           size_t dim_concat,
                           const fg_csr_batch_s64_t& fg,
                           AggOpT op,
                           const cuda::stream& stream);
/** @} */

/**
 * @brief Computes the forward pass of a simple aggregation fused with the concatenation
 *        of additional node features (agg_concat) using edge features in an edge-to-node
 *        reduction (e2n) operating on a full graph (fg).
 *        The "self" node features are concatenated at the end.
 *
 * @note given dimensions must be non-0.
 * @note If `fg.ef_indices` are set, then we use this as an edge index to obtain
 * the index of an edge feature given the index of an edge.
 * @note With an arbitrary edge feature index, the size of `edge_feat` is
 * arbitrary. The number of rows is `fg.n_indices` for a single graph or
 * `fg_batch.n_edges` for a batch of graphs.
 * Below we'll refer to the number of rows simply as `n_edges`.
 *
 * @param[out] out             the output aggregated embeddings. [on device]
 *                             [dim = `fg.n_nodes x (dim_edge + dim_concat)`]
 * @param[out] out_pos         the position of min/max values [on device]
 *                             [dim = `fg.n_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]  concat_feat     the node features to concatenate after aggregation.
 *                             [on device] [dim = `fg.n_nodes x dim_concat`]
 * @param[in]  dim_edge        dimensionality of the edge features
 * @param[in]  dim_concat      dimensionality of the additional concatenated features
 * @param[in]  fg              the input CSR full graph (or batch of graphs)
 * @param[in]  op              aggregation operation
 * @param[in]  stream          cuda stream
 * @{
 */
void agg_concat_fg_e2n_fwd(float* out,
                           int32_t* out_pos,
                           const float* edge_feat,
                           const float* concat_feat,
                           size_t dim_edge,
                           size_t dim_concat,
                           const fg_csr_s32_t& fg,
                           AggOpT op,
                           const cuda::stream& stream);
void agg_concat_fg_e2n_fwd(float* out,
                           int64_t* out_pos,
                           const float* edge_feat,
                           const float* concat_feat,
                           size_t dim_edge,
                           size_t dim_concat,
                           const fg_csr_s64_t& fg,
                           AggOpT op,
                           const cuda::stream& stream);

void agg_concat_fg_e2n_fwd(float* out,
                           int32_t* out_pos,
                           const float* edge_feat,
                           const float* concat_feat,
                           size_t dim_edge,
                           size_t dim_concat,
                           const fg_csr_batch_s32_t& fg,
                           AggOpT op,
                           const cuda::stream& stream);
void agg_concat_fg_e2n_fwd(float* out,
                           int64_t* out_pos,
                           const float* edge_feat,
                           const float* concat_feat,
                           size_t dim_edge,
                           size_t dim_concat,
                           const fg_csr_batch_s64_t& fg,
                           AggOpT op,
                           const cuda::stream& stream);

void agg_concat_fg_e2n_fwd(__half* out,
                           int32_t* out_pos,
                           const __half* edge_feat,
                           const __half* concat_feat,
                           size_t dim_edge,
                           size_t dim_concat,
                           const fg_csr_s32_t& fg,
                           AggOpT op,
                           const cuda::stream& stream);
void agg_concat_fg_e2n_fwd(__half* out,
                           int64_t* out_pos,
                           const __half* edge_feat,
                           const __half* concat_feat,
                           size_t dim_edge,
                           size_t dim_concat,
                           const fg_csr_s64_t& fg,
                           AggOpT op,
                           const cuda::stream& stream);

void agg_concat_fg_e2n_fwd(__half* out,
                           int32_t* out_pos,
                           const __half* edge_feat,
                           const __half* concat_feat,
                           size_t dim_edge,
                           size_t dim_concat,
                           const fg_csr_batch_s32_t& fg,
                           AggOpT op,
                           const cuda::stream& stream);
void agg_concat_fg_e2n_fwd(__half* out,
                           int64_t* out_pos,
                           const __half* edge_feat,
                           const __half* concat_feat,
                           size_t dim_edge,
                           size_t dim_concat,
                           const fg_csr_batch_s64_t& fg,
                           AggOpT op,
                           const cuda::stream& stream);

void agg_concat_fg_e2n_fwd(__nv_bfloat16* out,
                           int32_t* out_pos,
                           const __nv_bfloat16* edge_feat,
                           const __nv_bfloat16* concat_feat,
                           size_t dim_edge,
                           size_t dim_concat,
                           const fg_csr_s32_t& fg,
                           AggOpT op,
                           const cuda::stream& stream);
void agg_concat_fg_e2n_fwd(__nv_bfloat16* out,
                           int64_t* out_pos,
                           const __nv_bfloat16* edge_feat,
                           const __nv_bfloat16* concat_feat,
                           size_t dim_edge,
                           size_t dim_concat,
                           const fg_csr_s64_t& fg,
                           AggOpT op,
                           const cuda::stream& stream);

void agg_concat_fg_e2n_fwd(__nv_bfloat16* out,
                           int32_t* out_pos,
                           const __nv_bfloat16* edge_feat,
                           const __nv_bfloat16* concat_feat,
                           size_t dim_edge,
                           size_t dim_concat,
                           const fg_csr_batch_s32_t& fg,
                           AggOpT op,
                           const cuda::stream& stream);
void agg_concat_fg_e2n_fwd(__nv_bfloat16* out,
                           int64_t* out_pos,
                           const __nv_bfloat16* edge_feat,
                           const __nv_bfloat16* concat_feat,
                           size_t dim_edge,
                           size_t dim_concat,
                           const fg_csr_batch_s64_t& fg,
                           AggOpT op,
                           const cuda::stream& stream);
/** @} */

/**
 * @brief Computes the backward pass of a simple aggregation fused with the concatenation
 *        of additional node features (agg_concat) using edge features in an edge-to-node
 *        reduction (e2n) operating on a full graph (fg).
 *        The "self" node features are concatenated at the end.
 *
 * @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 `fg.ef_indices` are set, then we use this as an edge index to obtain
 * the index of an edge feature given the index of an edge.
 * @note With an arbitrary edge feature index, the size of `d_edge_feat` is
 * arbitrary.
 * @note When specifying shapes of buffers, we'll refer to
 * both `fg.n_indices` for a single graph or `fg_batch.n_edges`
 * for a batch of graphs simply as `n_edges`.
 *
 * @param[out] d_edge_feat     the output gradients of input edge features
 *                             [on device] [dim = `n_edges x dim_edge`].
 *                             Pass `nullptr` if not required.
 * @param[out] d_concat_feat   the output gradients of concatenated node features
 *                             [on device] [dim = `fg.n_nodes x dim_concat`].
 *                             Pass `nullptr` if not required.
 * @param[in]  dout            the input node feature gradients. [on device]
 *                             [dim = `fg.n_nodes x (dim_edge + dim_concat)`]
 * @param[in]  out_pos         the position of min/max values [on device]
 *                             [dim = `fg.n_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]  dim_concat      dimensionality of the additional concatenated features
 * @param[in]  fg              the input CSR full graph (or batch of graphs)
 * @param[in]  op              aggregation operation
 * @param[in]  stream          cuda stream
 * @{
 */
void agg_concat_fg_e2n_bwd(float* d_edge_feat,
                           float* d_concat_feat,
                           const float* dout,
                           const int32_t* out_pos,
                           size_t dim_edge,
                           size_t dim_concat,
                           const fg_csr_s32_t& fg,
                           AggOpT op,
                           const cuda::stream& stream);
void agg_concat_fg_e2n_bwd(float* d_edge_feat,
                           float* d_concat_feat,
                           const float* dout,
                           const int64_t* out_pos,
                           size_t dim_edge,
                           size_t dim_concat,
                           const fg_csr_s64_t& fg,
                           AggOpT op,
                           const cuda::stream& stream);

void agg_concat_fg_e2n_bwd(float* d_edge_feat,
                           float* d_concat_feat,
                           const float* dout,
                           const int32_t* out_pos,
                           size_t dim_edge,
                           size_t dim_concat,
                           const fg_csr_batch_s32_t& fg,
                           AggOpT op,
                           const cuda::stream& stream);
void agg_concat_fg_e2n_bwd(float* d_edge_feat,
                           float* d_concat_feat,
                           const float* dout,
                           const int64_t* out_pos,
                           size_t dim_edge,
                           size_t dim_concat,
                           const fg_csr_batch_s64_t& fg,
                           AggOpT op,
                           const cuda::stream& stream);

void agg_concat_fg_e2n_bwd(__half* d_edge_feat,
                           __half* d_concat_feat,
                           const __half* dout,
                           const int32_t* out_pos,
                           size_t dim_edge,
                           size_t dim_concat,
                           const fg_csr_s32_t& fg,
                           AggOpT op,
                           const cuda::stream& stream);
void agg_concat_fg_e2n_bwd(__half* d_edge_feat,
                           __half* d_concat_feat,
                           const __half* dout,
                           const int64_t* out_pos,
                           size_t dim_edge,
                           size_t dim_concat,
                           const fg_csr_s64_t& fg,
                           AggOpT op,
                           const cuda::stream& stream);

void agg_concat_fg_e2n_bwd(__half* d_edge_feat,
                           __half* d_concat_feat,
                           const __half* dout,
                           const int32_t* out_pos,
                           size_t dim_edge,
                           size_t dim_concat,
                           const fg_csr_batch_s32_t& fg,
                           AggOpT op,
                           const cuda::stream& stream);
void agg_concat_fg_e2n_bwd(__half* d_edge_feat,
                           __half* d_concat_feat,
                           const __half* dout,
                           const int64_t* out_pos,
                           size_t dim_edge,
                           size_t dim_concat,
                           const fg_csr_batch_s64_t& fg,
                           AggOpT op,
                           const cuda::stream& stream);

void agg_concat_fg_e2n_bwd(__nv_bfloat16* d_edge_feat,
                           __nv_bfloat16* d_concat_feat,
                           const __nv_bfloat16* dout,
                           const int32_t* out_pos,
                           size_t dim_edge,
                           size_t dim_concat,
                           const fg_csr_s32_t& fg,
                           AggOpT op,
                           const cuda::stream& stream);
void agg_concat_fg_e2n_bwd(__nv_bfloat16* d_edge_feat,
                           __nv_bfloat16* d_concat_feat,
                           const __nv_bfloat16* dout,
                           const int64_t* out_pos,
                           size_t dim_edge,
                           size_t dim_concat,
                           const fg_csr_s64_t& fg,
                           AggOpT op,
                           const cuda::stream& stream);

void agg_concat_fg_e2n_bwd(__nv_bfloat16* d_edge_feat,
                           __nv_bfloat16* d_concat_feat,
                           const __nv_bfloat16* dout,
                           const int32_t* out_pos,
                           size_t dim_edge,
                           size_t dim_concat,
                           const fg_csr_batch_s32_t& fg,
                           AggOpT op,
                           const cuda::stream& stream);
void agg_concat_fg_e2n_bwd(__nv_bfloat16* d_edge_feat,
                           __nv_bfloat16* d_concat_feat,
                           const __nv_bfloat16* dout,
                           const int64_t* out_pos,
                           size_t dim_edge,
                           size_t dim_concat,
                           const fg_csr_batch_s64_t& fg,
                           AggOpT op,
                           const cuda::stream& stream);
/** @} */

/**
 * @brief Computes the forward pass of a simple aggregation fused with the concatenation
 *        of additional node features (agg_concat)
 *        using edge features and node features in an edge-to-node reduction (e2n)
 *        and a node-to-node reduction (n2n) operating on a full graph (fg).
 *        The results of both individual aggregations are simply concatenated. The "self"
 *        node features are concatenated at the end.
 *
 * @note given dimensions must be non-0.
 * @note If `fg.ef_indices` are set, then we use this as an edge index to obtain
 * the index of an edge feature given the index of an edge.
 * @note With an arbitrary edge feature index, the size of `edge_feat` may be
 * arbitrary.
 * @note When specifying shapes of buffers, we'll refer to
 * both `fg.n_indices` for a single graph or `fg_batch.n_edges`
 * for a batch of graphs simply as `n_edges`.
 * @note `concat_feat` must not overlap with `node_feat` except if
 * `concat_feat == node_feat` and `dim_concat == dim_node`.
 *
 * @param[out] out             the output aggregated embeddings. [on device]
 *                             [dim = `fg.n_nodes x (dim_node + dim_edge + dim_concat)`]
 * @param[out] out_pos         the position of min/max values [on device]
 *                             [dim = `fg.n_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 = `fg.n_nodes x dim_node`]
 * @param[in]  edge_feat       the input edge features. [on device]
 *                             [dim = `n_edges x dim_edge`]
 * @param[in]  concat_feat     the node features to concatenate after aggregation.
 *                             [on device] [dim = `fg.n_nodes x dim_concat`]
 * @param[in]  dim_node        dimensionality of the node features
 * @param[in]  dim_edge        dimensionality of the edge features
 * @param[in]  dim_concat      dimensionality of the additional concatenated features
 * @param[in]  fg              the input CSR full graph (or batch of graphs)
 * @param[in]  op              aggregation operation
 * @param[in]  stream          cuda stream
 * @{
 */
void agg_concat_fg_n2n_e2n_fwd(float* out,
                               int32_t* out_pos,
                               const float* node_feat,
                               const float* edge_feat,
                               const float* concat_feat,
                               size_t dim_node,
                               size_t dim_edge,
                               size_t dim_concat,
                               const fg_csr_s32_t& fg,
                               AggOpT op,
                               const cuda::stream& stream);
void agg_concat_fg_n2n_e2n_fwd(float* out,
                               int32_t* out_pos,
                               const float* node_feat,
                               const float* edge_feat,
                               const float* concat_feat,
                               size_t dim_node,
                               size_t dim_edge,
                               size_t dim_concat,
                               const fg_csr_batch_s32_t& fg,
                               AggOpT op,
                               const cuda::stream& stream);
void agg_concat_fg_n2n_e2n_fwd(float* out,
                               int64_t* out_pos,
                               const float* node_feat,
                               const float* edge_feat,
                               const float* concat_feat,
                               size_t dim_node,
                               size_t dim_edge,
                               size_t dim_concat,
                               const fg_csr_s64_t& fg,
                               AggOpT op,
                               const cuda::stream& stream);
void agg_concat_fg_n2n_e2n_fwd(float* out,
                               int64_t* out_pos,
                               const float* node_feat,
                               const float* edge_feat,
                               const float* concat_feat,
                               size_t dim_node,
                               size_t dim_edge,
                               size_t dim_concat,
                               const fg_csr_batch_s64_t& fg,
                               AggOpT op,
                               const cuda::stream& stream);

void agg_concat_fg_n2n_e2n_fwd(__half* out,
                               int32_t* out_pos,
                               const __half* node_feat,
                               const __half* edge_feat,
                               const __half* concat_feat,
                               size_t dim_node,
                               size_t dim_edge,
                               size_t dim_concat,
                               const fg_csr_s32_t& fg,
                               AggOpT op,
                               const cuda::stream& stream);
void agg_concat_fg_n2n_e2n_fwd(__half* out,
                               int32_t* out_pos,
                               const __half* node_feat,
                               const __half* edge_feat,
                               const __half* concat_feat,
                               size_t dim_node,
                               size_t dim_edge,
                               size_t dim_concat,
                               const fg_csr_batch_s32_t& fg,
                               AggOpT op,
                               const cuda::stream& stream);
void agg_concat_fg_n2n_e2n_fwd(__half* out,
                               int64_t* out_pos,
                               const __half* node_feat,
                               const __half* edge_feat,
                               const __half* concat_feat,
                               size_t dim_node,
                               size_t dim_edge,
                               size_t dim_concat,
                               const fg_csr_s64_t& fg,
                               AggOpT op,
                               const cuda::stream& stream);
void agg_concat_fg_n2n_e2n_fwd(__half* out,
                               int64_t* out_pos,
                               const __half* node_feat,
                               const __half* edge_feat,
                               const __half* concat_feat,
                               size_t dim_node,
                               size_t dim_edge,
                               size_t dim_concat,
                               const fg_csr_batch_s64_t& fg,
                               AggOpT op,
                               const cuda::stream& stream);

void agg_concat_fg_n2n_e2n_fwd(__nv_bfloat16* out,
                               int32_t* out_pos,
                               const __nv_bfloat16* node_feat,
                               const __nv_bfloat16* edge_feat,
                               const __nv_bfloat16* concat_feat,
                               size_t dim_node,
                               size_t dim_edge,
                               size_t dim_concat,
                               const fg_csr_s32_t& fg,
                               AggOpT op,
                               const cuda::stream& stream);
void agg_concat_fg_n2n_e2n_fwd(__nv_bfloat16* out,
                               int32_t* out_pos,
                               const __nv_bfloat16* node_feat,
                               const __nv_bfloat16* edge_feat,
                               const __nv_bfloat16* concat_feat,
                               size_t dim_node,
                               size_t dim_edge,
                               size_t dim_concat,
                               const fg_csr_batch_s32_t& fg,
                               AggOpT op,
                               const cuda::stream& stream);
void agg_concat_fg_n2n_e2n_fwd(__nv_bfloat16* out,
                               int64_t* out_pos,
                               const __nv_bfloat16* node_feat,
                               const __nv_bfloat16* edge_feat,
                               const __nv_bfloat16* concat_feat,
                               size_t dim_node,
                               size_t dim_edge,
                               size_t dim_concat,
                               const fg_csr_s64_t& fg,
                               AggOpT op,
                               const cuda::stream& stream);
void agg_concat_fg_n2n_e2n_fwd(__nv_bfloat16* out,
                               int64_t* out_pos,
                               const __nv_bfloat16* node_feat,
                               const __nv_bfloat16* edge_feat,
                               const __nv_bfloat16* concat_feat,
                               size_t dim_node,
                               size_t dim_edge,
                               size_t dim_concat,
                               const fg_csr_batch_s64_t& fg,
                               AggOpT op,
                               const cuda::stream& stream);
/** @} */

/**
 * @brief Computes the backward pass of a simple aggregation fused with the concatenation
 *        of the "self" node features (agg_concat)
 *        using edge features and node features in an edge-to-node reduction (e2n)
 *        and a node-to-node reduction (n2n) operating on a full graph (fg).
 *        The results of both individual aggregations are simply concatenated. The "self"
 *        node features are concatenated at the end.
 *
 * @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 `fg.ef_indices` are set, then we use this as an edge index to obtain
 * the index of an edge feature given the index of an edge.
 * @note With an arbitrary edge feature index, the size of `d_edge_feat` may be
 * arbitrary.
 * @note When specifying shapes of buffers, we'll refer to
 * both `fg.n_indices` for a single graph or `fg_batch.n_edges`
 * for a batch of graphs simply as `n_edges`.
 * @note `d_concat_feat` must not overlap with `d_node_feat` except if
 * `d_concat_feat == d_node_feat` and `dim_concat == dim_node`.
 *
 * @param[out] d_node_feat     the output gradients of input node features
 *                             [on device] [dim = `fg.n_nodes x dim_node`].
 *                             Pass `nullptr` if not required.
 * @param[out] d_edge_feat     the output gradients of input edge features
 *                             [on device] [dim = `n_edges x dim_edge`].
 *                             Pass `nullptr` if not required.
 * @param[out] d_concat_feat   the output gradients of concatenated node features
 *                             [on device] [dim = `fg.n_nodes x dim_concat`].
 *                             Pass `nullptr` if not required.
 * @param[in]  dout            the input node feature gradients. [on device]
 *                             [dim = `fg.n_nodes x (dim_node + dim_edge + dim_concat)`]
 * @param[in]  out_pos         the position of min/max values [on device]
 *                             [dim = `fg.n_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]  dim_concat      dimensionality of the additional concatenated features
 * @param[in]  fg              the input CSR full graph (or batch of graphs)
 * @param[in]  op              aggregation operation
 * @param[in]  stream          cuda stream
 * @{
 */
void agg_concat_fg_n2n_e2n_bwd(float* d_node_feat,
                               float* d_edge_feat,
                               float* d_concat_feat,
                               const float* dout,
                               const int32_t* out_pos,
                               size_t dim_node,
                               size_t dim_edge,
                               size_t dim_concat,
                               const fg_csr_s32_t& fg,
                               AggOpT op,
                               const cuda::stream& stream);
void agg_concat_fg_n2n_e2n_bwd(float* d_node_feat,
                               float* d_edge_feat,
                               float* d_concat_feat,
                               const float* dout,
                               const int32_t* out_pos,
                               size_t dim_node,
                               size_t dim_edge,
                               size_t dim_concat,
                               const fg_csr_batch_s32_t& fg,
                               AggOpT op,
                               const cuda::stream& stream);
void agg_concat_fg_n2n_e2n_bwd(float* d_node_feat,
                               float* d_edge_feat,
                               float* d_concat_feat,
                               const float* dout,
                               const int64_t* out_pos,
                               size_t dim_node,
                               size_t dim_edge,
                               size_t dim_concat,
                               const fg_csr_s64_t& fg,
                               AggOpT op,
                               const cuda::stream& stream);
void agg_concat_fg_n2n_e2n_bwd(float* d_node_feat,
                               float* d_edge_feat,
                               float* d_concat_feat,
                               const float* dout,
                               const int64_t* out_pos,
                               size_t dim_node,
                               size_t dim_edge,
                               size_t dim_concat,
                               const fg_csr_batch_s64_t& fg,
                               AggOpT op,
                               const cuda::stream& stream);

void agg_concat_fg_n2n_e2n_bwd(__half* d_node_feat,
                               __half* d_edge_feat,
                               __half* d_concat_feat,
                               const __half* dout,
                               const int32_t* out_pos,
                               size_t dim_node,
                               size_t dim_edge,
                               size_t dim_concat,
                               const fg_csr_s32_t& fg,
                               AggOpT op,
                               const cuda::stream& stream);
void agg_concat_fg_n2n_e2n_bwd(__half* d_node_feat,
                               __half* d_edge_feat,
                               __half* d_concat_feat,
                               const __half* dout,
                               const int32_t* out_pos,
                               size_t dim_node,
                               size_t dim_edge,
                               size_t dim_concat,
                               const fg_csr_batch_s32_t& fg,
                               AggOpT op,
                               const cuda::stream& stream);
void agg_concat_fg_n2n_e2n_bwd(__half* d_node_feat,
                               __half* d_edge_feat,
                               __half* d_concat_feat,
                               const __half* dout,
                               const int64_t* out_pos,
                               size_t dim_node,
                               size_t dim_edge,
                               size_t dim_concat,
                               const fg_csr_s64_t& fg,
                               AggOpT op,
                               const cuda::stream& stream);
void agg_concat_fg_n2n_e2n_bwd(__half* d_node_feat,
                               __half* d_edge_feat,
                               __half* d_concat_feat,
                               const __half* dout,
                               const int64_t* out_pos,
                               size_t dim_node,
                               size_t dim_edge,
                               size_t dim_concat,
                               const fg_csr_batch_s64_t& fg,
                               AggOpT op,
                               const cuda::stream& stream);

void agg_concat_fg_n2n_e2n_bwd(__nv_bfloat16* d_node_feat,
                               __nv_bfloat16* d_edge_feat,
                               __nv_bfloat16* d_concat_feat,
                               const __nv_bfloat16* dout,
                               const int32_t* out_pos,
                               size_t dim_node,
                               size_t dim_edge,
                               size_t dim_concat,
                               const fg_csr_s32_t& fg,
                               AggOpT op,
                               const cuda::stream& stream);
void agg_concat_fg_n2n_e2n_bwd(__nv_bfloat16* d_node_feat,
                               __nv_bfloat16* d_edge_feat,
                               __nv_bfloat16* d_concat_feat,
                               const __nv_bfloat16* dout,
                               const int32_t* out_pos,
                               size_t dim_node,
                               size_t dim_edge,
                               size_t dim_concat,
                               const fg_csr_batch_s32_t& fg,
                               AggOpT op,
                               const cuda::stream& stream);
void agg_concat_fg_n2n_e2n_bwd(__nv_bfloat16* d_node_feat,
                               __nv_bfloat16* d_edge_feat,
                               __nv_bfloat16* d_concat_feat,
                               const __nv_bfloat16* dout,
                               const int64_t* out_pos,
                               size_t dim_node,
                               size_t dim_edge,
                               size_t dim_concat,
                               const fg_csr_s64_t& fg,
                               AggOpT op,
                               const cuda::stream& stream);
void agg_concat_fg_n2n_e2n_bwd(__nv_bfloat16* d_node_feat,
                               __nv_bfloat16* d_edge_feat,
                               __nv_bfloat16* d_concat_feat,
                               const __nv_bfloat16* dout,
                               const int64_t* out_pos,
                               size_t dim_node,
                               size_t dim_edge,
                               size_t dim_concat,
                               const fg_csr_batch_s64_t& fg,
                               AggOpT op,
                               const cuda::stream& stream);
/** @} */

}  // namespace cugraph::ops
