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

#pragma once

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

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

namespace cugraph::ops {

/**
 * @brief a wrapper struct containing informations
 * about tensor-product paths
 */
template <typename MathT>
struct __attribute__((aligned(16))) tp_info {
  /** offsets into `path_offsets_and_dims` for each "target" */
  const int32_t* __restrict__ path_csr_offsets{nullptr};
  /** "sources" of all paths and their offsets and dimensions */
  const int32_t* __restrict__ path_offsets_and_dims{nullptr};
  /** clebsch-gordan values for each path */
  const MathT* __restrict__ path_cg_values{nullptr};
  /** number of "target" segments */
  int32_t num_target_segments{0};
  /** number of path (i.e. all paths between segments) */
  int32_t num_paths{0};
};  // struct tp_info

struct __attribute__((aligned(16))) tp_data_sizes {
  int64_t batch_size;
  bool shared_a;
  bool shared_b;
  bool shared_w;
  bool shared_bias;
  int32_t stride_a;
  int32_t stride_b;
  int32_t stride_w;
  int32_t stride_o;
};  // struct tp_data_sizes

enum class ConnectionModeT : uint8_t {
  kUVW = 0,
  // UVW with U spherical harmonic
  k1VW,  // NOLINT
  // UVW with V spherical harmonic
  kU1W,  // NOLINT
  kUVU,
  kUVV,
  kUUW,
  kUUU,
  // FullTP, no weight
  kUVUV,
  // FullTP, U spherical harmonic
  k1V1V,
  // FullTP, V spherical harmonic
  kU1U1,
  // Linear
  kUUVV,
};

void segmented_transpose(float* tensor,
                         const float* tensor_transpose,
                         const int32_t* segment_info,
                         int32_t num_segments,
                         int64_t batch_size,
                         int64_t stride,
                         bool is_transposed,
                         const cuda::stream& stream);
void segmented_transpose(__half* tensor,
                         const __half* tensor_transpose,
                         const int32_t* segment_info,
                         int32_t num_segments,
                         int64_t batch_size,
                         int64_t stride,
                         bool is_transposed,
                         const cuda::stream& stream);
void segmented_transpose(__nv_bfloat16* tensor,
                         const __nv_bfloat16* tensor_transpose,
                         const int32_t* segment_info,
                         int32_t num_segments,
                         int64_t batch_size,
                         int64_t stride,
                         bool is_transposed,
                         const cuda::stream& stream);

void fused_tp_n2n_fwd(float* out,
                      const float* in_a,
                      const float* in_b,
                      const float* weight,
                      const float* bias,
                      ConnectionModeT mode,
                      const tp_info<float>& info,
                      const tp_data_sizes& sizes,
                      const cuda::stream& stream);
void fused_tp_n2n_fwd(__half* out,
                      const __half* in_a,
                      const __half* in_b,
                      const __half* weight,
                      const __half* bias,
                      ConnectionModeT mode,
                      const tp_info<float>& info,
                      const tp_data_sizes& sizes,
                      const cuda::stream& stream);
void fused_tp_n2n_fwd(__nv_bfloat16* out,
                      const __nv_bfloat16* in_a,
                      const __nv_bfloat16* in_b,
                      const __nv_bfloat16* weight,
                      const __nv_bfloat16* bias,
                      ConnectionModeT mode,
                      const tp_info<float>& info,
                      const tp_data_sizes& sizes,
                      const cuda::stream& stream);
void fused_tp_n2n_fwd(__half* out,
                      const float* in_a,
                      const __half* in_b,
                      const __half* weight,
                      const __half* bias,
                      ConnectionModeT mode,
                      const tp_info<float>& info,
                      const tp_data_sizes& sizes,
                      const cuda::stream& stream);
void fused_tp_n2n_fwd(__nv_bfloat16* out,
                      const float* in_a,
                      const __nv_bfloat16* in_b,
                      const __nv_bfloat16* weight,
                      const __nv_bfloat16* bias,
                      ConnectionModeT mode,
                      const tp_info<float>& info,
                      const tp_data_sizes& sizes,
                      const cuda::stream& stream);
void fused_tp_n2n_fwd(__half* out,
                      const __half* in_a,
                      const float* in_b,
                      const __half* weight,
                      const __half* bias,
                      ConnectionModeT mode,
                      const tp_info<float>& info,
                      const tp_data_sizes& sizes,
                      const cuda::stream& stream);
void fused_tp_n2n_fwd(__nv_bfloat16* out,
                      const __nv_bfloat16* in_a,
                      const float* in_b,
                      const __nv_bfloat16* weight,
                      const __nv_bfloat16* bias,
                      ConnectionModeT mode,
                      const tp_info<float>& info,
                      const tp_data_sizes& sizes,
                      const cuda::stream& stream);

void fused_tp_n2n_bwd(float* grad_in_a,
                      float* grad_in_b,
                      float* grad_weight,
                      float* grad_bias,
                      const float* grad_out,
                      const float* in_a,
                      const float* in_b,
                      const float* weight,
                      const float* bias,
                      ConnectionModeT mode,
                      const tp_info<float>& info_bwd_dgrad_a,
                      const tp_info<float>& info_bwd_dgrad_b,
                      const tp_info<float>& info_bwd_dgrad_w,
                      const tp_data_sizes& sizes,
                      const cuda::stream& stream);
void fused_tp_n2n_bwd(__half* grad_in_a,
                      __half* grad_in_b,
                      __half* grad_weight,
                      __half* grad_bias,
                      const __half* grad_out,
                      const __half* in_a,
                      const __half* in_b,
                      const __half* weight,
                      const __half* bias,
                      ConnectionModeT mode,
                      const tp_info<float>& info_bwd_dgrad_a,
                      const tp_info<float>& info_bwd_dgrad_b,
                      const tp_info<float>& info_bwd_dgrad_w,
                      const tp_data_sizes& sizes,
                      const cuda::stream& stream);
void fused_tp_n2n_bwd(__nv_bfloat16* grad_in_a,
                      __nv_bfloat16* grad_in_b,
                      __nv_bfloat16* grad_weight,
                      __nv_bfloat16* grad_bias,
                      const __nv_bfloat16* grad_out,
                      const __nv_bfloat16* in_a,
                      const __nv_bfloat16* in_b,
                      const __nv_bfloat16* weight,
                      const __nv_bfloat16* bias,
                      ConnectionModeT mode,
                      const tp_info<float>& info_bwd_dgrad_a,
                      const tp_info<float>& info_bwd_dgrad_b,
                      const tp_info<float>& info_bwd_dgrad_w,
                      const tp_data_sizes& sizes,
                      const cuda::stream& stream);
void fused_tp_n2n_bwd(float* grad_in_a,
                      __half* grad_in_b,
                      __half* grad_weight,
                      __half* grad_bias,
                      const __half* grad_out,
                      const float* in_a,
                      const __half* in_b,
                      const __half* weight,
                      const __half* bias,
                      ConnectionModeT mode,
                      const tp_info<float>& info_bwd_dgrad_a,
                      const tp_info<float>& info_bwd_dgrad_b,
                      const tp_info<float>& info_bwd_dgrad_w,
                      const tp_data_sizes& sizes,
                      const cuda::stream& stream);
void fused_tp_n2n_bwd(float* grad_in_a,
                      __nv_bfloat16* grad_in_b,
                      __nv_bfloat16* grad_weight,
                      __nv_bfloat16* grad_bias,
                      const __nv_bfloat16* grad_out,
                      const float* in_a,
                      const __nv_bfloat16* in_b,
                      const __nv_bfloat16* weight,
                      const __nv_bfloat16* bias,
                      ConnectionModeT mode,
                      const tp_info<float>& info_bwd_dgrad_a,
                      const tp_info<float>& info_bwd_dgrad_b,
                      const tp_info<float>& info_bwd_dgrad_w,
                      const tp_data_sizes& sizes,
                      const cuda::stream& stream);
void fused_tp_n2n_bwd(__half* grad_in_a,
                      float* grad_in_b,
                      __half* grad_weight,
                      __half* grad_bias,
                      const __half* grad_out,
                      const __half* in_a,
                      const float* in_b,
                      const __half* weight,
                      const __half* bias,
                      ConnectionModeT mode,
                      const tp_info<float>& info_bwd_dgrad_a,
                      const tp_info<float>& info_bwd_dgrad_b,
                      const tp_info<float>& info_bwd_dgrad_w,
                      const tp_data_sizes& sizes,
                      const cuda::stream& stream);
void fused_tp_n2n_bwd(__nv_bfloat16* grad_in_a,
                      float* grad_in_b,
                      __nv_bfloat16* grad_weight,
                      __nv_bfloat16* grad_bias,
                      const __nv_bfloat16* grad_out,
                      const __nv_bfloat16* in_a,
                      const float* in_b,
                      const __nv_bfloat16* weight,
                      const __nv_bfloat16* bias,
                      ConnectionModeT mode,
                      const tp_info<float>& info_bwd_dgrad_a,
                      const tp_info<float>& info_bwd_dgrad_b,
                      const tp_info<float>& info_bwd_dgrad_w,
                      const tp_data_sizes& sizes,
                      const cuda::stream& stream);

void fused_tp_n2n_bwd_bwd(float* grad_in_a,
                          float* grad_in_b,
                          float* grad_weight,
                          float* grad_grad_out,
                          const float* grad_grad_in_a,
                          const float* grad_grad_in_b,
                          const float* grad_grad_weight,
                          const float* grad_out,
                          const float* in_a,
                          const float* in_b,
                          const float* weight,
                          ConnectionModeT mode,
                          const tp_info<float>& info_fwd,
                          const tp_info<float>& info_bwd_dgrad_a,
                          const tp_info<float>& info_bwd_dgrad_b,
                          const tp_info<float>& info_bwd_dgrad_w,
                          const tp_data_sizes& sizes,
                          const cuda::stream& stream);
void fused_tp_n2n_bwd_bwd(__half* grad_in_a,
                          __half* grad_in_b,
                          __half* grad_weight,
                          __half* grad_grad_out,
                          const __half* grad_grad_in_a,
                          const __half* grad_grad_in_b,
                          const __half* grad_grad_weight,
                          const __half* grad_out,
                          const __half* in_a,
                          const __half* in_b,
                          const __half* weight,
                          ConnectionModeT mode,
                          const tp_info<float>& info_fwd,
                          const tp_info<float>& info_bwd_dgrad_a,
                          const tp_info<float>& info_bwd_dgrad_b,
                          const tp_info<float>& info_bwd_dgrad_w,
                          const tp_data_sizes& sizes,
                          const cuda::stream& stream);
void fused_tp_n2n_bwd_bwd(__nv_bfloat16* grad_in_a,
                          __nv_bfloat16* grad_in_b,
                          __nv_bfloat16* grad_weight,
                          __nv_bfloat16* grad_grad_out,
                          const __nv_bfloat16* grad_grad_in_a,
                          const __nv_bfloat16* grad_grad_in_b,
                          const __nv_bfloat16* grad_grad_weight,
                          const __nv_bfloat16* grad_out,
                          const __nv_bfloat16* in_a,
                          const __nv_bfloat16* in_b,
                          const __nv_bfloat16* weight,
                          ConnectionModeT mode,
                          const tp_info<float>& info_fwd,
                          const tp_info<float>& info_bwd_dgrad_a,
                          const tp_info<float>& info_bwd_dgrad_b,
                          const tp_info<float>& info_bwd_dgrad_w,
                          const tp_data_sizes& sizes,
                          const cuda::stream& stream);
void fused_tp_n2n_bwd_bwd(float* grad_in_a,
                          __half* grad_in_b,
                          __half* grad_weight,
                          __half* grad_grad_out,
                          const float* grad_grad_in_a,
                          const __half* grad_grad_in_b,
                          const __half* grad_grad_weight,
                          const __half* grad_out,
                          const float* in_a,
                          const __half* in_b,
                          const __half* weight,
                          ConnectionModeT mode,
                          const tp_info<float>& info_fwd,
                          const tp_info<float>& info_bwd_dgrad_a,
                          const tp_info<float>& info_bwd_dgrad_b,
                          const tp_info<float>& info_bwd_dgrad_w,
                          const tp_data_sizes& sizes,
                          const cuda::stream& stream);
void fused_tp_n2n_bwd_bwd(float* grad_in_a,
                          __nv_bfloat16* grad_in_b,
                          __nv_bfloat16* grad_weight,
                          __nv_bfloat16* grad_grad_out,
                          const float* grad_grad_in_a,
                          const __nv_bfloat16* grad_grad_in_b,
                          const __nv_bfloat16* grad_grad_weight,
                          const __nv_bfloat16* grad_out,
                          const float* in_a,
                          const __nv_bfloat16* in_b,
                          const __nv_bfloat16* weight,
                          ConnectionModeT mode,
                          const tp_info<float>& info_fwd,
                          const tp_info<float>& info_bwd_dgrad_a,
                          const tp_info<float>& info_bwd_dgrad_b,
                          const tp_info<float>& info_bwd_dgrad_w,
                          const tp_data_sizes& sizes,
                          const cuda::stream& stream);
void fused_tp_n2n_bwd_bwd(__half* grad_in_a,
                          float* grad_in_b,
                          __half* grad_weight,
                          __half* grad_grad_out,
                          const __half* grad_grad_in_a,
                          const float* grad_grad_in_b,
                          const __half* grad_grad_weight,
                          const __half* grad_out,
                          const __half* in_a,
                          const float* in_b,
                          const __half* weight,
                          ConnectionModeT mode,
                          const tp_info<float>& info_fwd,
                          const tp_info<float>& info_bwd_dgrad_a,
                          const tp_info<float>& info_bwd_dgrad_b,
                          const tp_info<float>& info_bwd_dgrad_w,
                          const tp_data_sizes& sizes,
                          const cuda::stream& stream);
void fused_tp_n2n_bwd_bwd(__nv_bfloat16* grad_in_a,
                          float* grad_in_b,
                          __nv_bfloat16* grad_weight,
                          __nv_bfloat16* grad_grad_out,
                          const __nv_bfloat16* grad_grad_in_a,
                          const float* grad_grad_in_b,
                          const __nv_bfloat16* grad_grad_weight,
                          const __nv_bfloat16* grad_out,
                          const __nv_bfloat16* in_a,
                          const float* in_b,
                          const __nv_bfloat16* weight,
                          ConnectionModeT mode,
                          const tp_info<float>& info_fwd,
                          const tp_info<float>& info_bwd_dgrad_a,
                          const tp_info<float>& info_bwd_dgrad_b,
                          const tp_info<float>& info_bwd_dgrad_w,
                          const tp_data_sizes& sizes,
                          const cuda::stream& stream);

}  // namespace cugraph::ops
