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

#pragma once

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

#include <cstdint>

namespace cugraph::ops {

/**
 * @brief a wrapper struct containing informations
 * about tensor-product paths
 */
template <typename DataT>
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 DataT* __restrict__ path_cg_values{nullptr};
  /** number of "targets" in all paths */
  int32_t num_paths{0};
};  // struct tp_info

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

void fused_tp_n2n_fwd(float* out,
                      const float* in_a,
                      const float* in_b,
                      const float* weight,
                      ConnectionModeT mode,
                      const tp_info<float>& info,
                      int64_t batch_size,
                      int32_t stride_a,
                      int32_t stride_b,
                      int32_t stride_w,
                      int32_t stride_o,
                      const cuda::stream& stream);

void fused_tp_n2n_bwd(float* grad_in_a,
                      float* grad_in_b,
                      float* grad_weight,
                      const float* grad_out,
                      const float* in_a,
                      const float* in_b,
                      const float* weight,
                      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,
                      int64_t batch_size,
                      int32_t stride_a,
                      int32_t stride_b,
                      int32_t stride_w,
                      int32_t stride_o,
                      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,
                          int64_t batch_size,
                          int32_t stride_a,
                          int32_t stride_b,
                          int32_t stride_w,
                          int32_t stride_o,
                          const cuda::stream& stream);

}  // namespace cugraph::ops
