/*
 * Copyright (c) 2021-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 DataT>
struct __attribute__((aligned(16))) clebsch_gordan_tensor {
  const DataT* __restrict__ cg_values{nullptr};
  const int16_t* __restrict__ cg_indices{nullptr};
  const int32_t* __restrict__ cg_offsets{nullptr};
  int32_t total_output_irreps{0};
};  // struct clebsch_gordan_tensor

void symmetric_tensor_contraction_n2n_fwd(float* feat_b,
                                          const float* feat_a,
                                          const float* weight,
                                          const int32_t* weight_id,
                                          const clebsch_gordan_tensor<float>& cg_tensor,
                                          const int32_t& stride_b,
                                          const int32_t& stride_a,
                                          const int32_t& stride_w,
                                          const int32_t& stride_feat,
                                          const int32_t& num_batches,
                                          const int32_t& correlation,
                                          const cuda::stream& stream);

void symmetric_tensor_contraction_n2n_fwd(__half* feat_b,
                                          const __half* feat_a,
                                          const __half* weight,
                                          const int32_t* weight_id,
                                          const clebsch_gordan_tensor<__half>& cg_tensor,
                                          const int32_t& stride_b,
                                          const int32_t& stride_a,
                                          const int32_t& stride_w,
                                          const int32_t& stride_feat,
                                          const int32_t& num_batches,
                                          const int32_t& correlation,
                                          const cuda::stream& stream);

void symmetric_tensor_contraction_n2n_fwd(__nv_bfloat16* feat_b,
                                          const __nv_bfloat16* feat_a,
                                          const __nv_bfloat16* weight,
                                          const int32_t* weight_id,
                                          const clebsch_gordan_tensor<__nv_bfloat16>& cg_tensor,
                                          const int32_t& stride_b,
                                          const int32_t& stride_a,
                                          const int32_t& stride_w,
                                          const int32_t& stride_feat,
                                          const int32_t& num_batches,
                                          const int32_t& correlation,
                                          const cuda::stream& stream);

void symmetric_tensor_contraction_n2n_bwd(float* d_feat_a,
                                          float* d_weight,
                                          const float* feat_b,
                                          const float* feat_a,
                                          const float* weight,
                                          const int32_t* weight_id,
                                          const clebsch_gordan_tensor<float>& cg_tensor0,
                                          const clebsch_gordan_tensor<float>& cg_tensor1,
                                          const int32_t& stride_b,
                                          const int32_t& stride_a,
                                          const int32_t& stride_w,
                                          const int32_t& stride_feat,
                                          const int32_t& num_batches,
                                          const int32_t& correlation,
                                          const cuda::stream& stream);

void symmetric_tensor_contraction_n2n_bwd(__half* d_feat_a,
                                          float* d_weight,
                                          const __half* feat_b,
                                          const __half* feat_a,
                                          const __half* weight,
                                          const int32_t* weight_id,
                                          const clebsch_gordan_tensor<__half>& cg_tensor0,
                                          const clebsch_gordan_tensor<__half>& cg_tensor1,
                                          const int32_t& stride_b,
                                          const int32_t& stride_a,
                                          const int32_t& stride_w,
                                          const int32_t& stride_feat,
                                          const int32_t& num_batches,
                                          const int32_t& correlation,
                                          const cuda::stream& stream);

void symmetric_tensor_contraction_n2n_bwd(__nv_bfloat16* d_feat_a,
                                          float* d_weight,
                                          const __nv_bfloat16* feat_b,
                                          const __nv_bfloat16* feat_a,
                                          const __nv_bfloat16* weight,
                                          const int32_t* weight_id,
                                          const clebsch_gordan_tensor<__nv_bfloat16>& cg_tensor0,
                                          const clebsch_gordan_tensor<__nv_bfloat16>& cg_tensor1,
                                          const int32_t& stride_b,
                                          const int32_t& stride_a,
                                          const int32_t& stride_w,
                                          const int32_t& stride_feat,
                                          const int32_t& num_batches,
                                          const int32_t& correlation,
                                          const cuda::stream& stream);

}  // namespace cugraph::ops
