/*************************************************************************
 * Copyright (c) 2022-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 *
 * See LICENSE for license information.
 ************************************************************************/

#ifndef TRANSFORMER_ENGINE_FUSED_ROPE_H_
#define TRANSFORMER_ENGINE_FUSED_ROPE_H_

#include "transformer_engine.h"

#ifdef __cplusplus
extern "C" {
#endif

/*! \brief Apply rotary positional embedding to the input tensor.
 *
 *  \param[in]     input           Input tensor for fused rope.
 *  \param[in]     freqs           The freqs tensor.
 *  \param[out]    output          Output tensor.
 *  \param[in]     s               Length of the s dimension of input.
 *  \param[in]     b               Length of the b dimension of input.
 *  \param[in]     h               Length of the h dimension of input.
 *  \param[in]     d               Length of the d dimension of input.
 *  \param[in]     d2              Length of the d dimension of freqs.
 *  \param[in]     stride_s        Stride of the s dimension of input.
 *  \param[in]     stride_b        Stride of the b dimension of input.
 *  \param[in]     stride_h        Stride of the h dimension of input.
 *  \param[in]     stride_d        Stride of the d dimension of input.
 *  \param[in]     o_stride_s      Stride of the s dimension of output.
 *  \param[in]     o_stride_b      Stride of the b dimension of output.
 *  \param[in]     o_stride_h      Stride of the h dimension of output.
 *  \param[in]     o_stride_d      Stride of the d dimension of output.
 *  \param[in]     stream          CUDA stream used for the operation.
 */
void nvte_fused_rope_forward(const NVTETensor input, const NVTETensor freqs, NVTETensor output,
                             const int s, const int b, const int h, const int d, const int d2,
                             const int stride_s, const int stride_b, const int stride_h,
                             const int stride_d, const int o_stride_s, const int o_stride_b,
                             const int o_stride_h, const int o_stride_d, cudaStream_t stream);

/*! \brief Compute the backward of the fused rope.
 *
 *  \param[in]     output_grads    Incoming gradient tensor for backward.
 *  \param[in]     freqs           The freqs tensor.
 *  \param[out]    input_grads     Input gradient tensor to calculate.
 *  \param[in]     s               Length of the s dimension of output_grads.
 *  \param[in]     b               Length of the b dimension of output_grads.
 *  \param[in]     h               Length of the h dimension of output_grads.
 *  \param[in]     d               Length of the d dimension of output_grads.
 *  \param[in]     d2              Length of the d dimension of freqs.
 *  \param[in]     stride_s        Stride of the s dimension of output_grads.
 *  \param[in]     stride_b        Stride of the b dimension of output_grads.
 *  \param[in]     stride_h        Stride of the h dimension of output_grads.
 *  \param[in]     stride_d        Stride of the d dimension of output_grads.
 *  \param[in]     o_stride_s      Stride of the s dimension of input_grads.
 *  \param[in]     o_stride_b      Stride of the b dimension of input_grads.
 *  \param[in]     o_stride_h      Stride of the h dimension of input_grads.
 *  \param[in]     o_stride_d      Stride of the d dimension of input_grads.
 *  \param[in]     stream          CUDA stream used for the operation.
 */
void nvte_fused_rope_backward(const NVTETensor output_grads, const NVTETensor freqs,
                              NVTETensor input_grads, const int s, const int b, const int h,
                              const int d, const int d2, const int stride_s, const int stride_b,
                              const int stride_h, const int stride_d, const int o_stride_s,
                              const int o_stride_b, const int o_stride_h, const int o_stride_d,
                              cudaStream_t stream);

/*! \brief Apply rotary positional embedding to the input tensor in thd format.
 *
 *  \param[in]     input         Input tensor for fused rope.
 *  \param[in]     cu_seqlens    The cumulative sum of sequence lengths tensor.
 *  \param[in]     freqs         The freqs tensor.
 *  \param[out]    output        Output tensor.
 *  \param[in]     cp_size       Context parallel world size.
 *  \param[in]     cp_rank       Context parallel rank.
 *  \param[in]     max_s         Max sequence length.
 *  \param[in]     b             Batch size.
 *  \param[in]     h             Length of the h dimension of input.
 *  \param[in]     d             Length of the d dimension of input.
 *  \param[in]     d2            Length of the d dimension of freqs.
 *  \param[in]     stride_t      Stride of the t dimension of input.
 *  \param[in]     stride_h      Stride of the h dimension of input.
 *  \param[in]     stride_d      Stride of the d dimension of input.
 *  \param[in]     o_stride_t    Stride of the t dimension of output.
 *  \param[in]     o_stride_h    Stride of the h dimension of output.
 *  \param[in]     o_stride_d    Stride of the d dimension of output.
 *  \param[in]     stream        CUDA stream used for the operation.
 */
void nvte_fused_rope_thd_forward(const NVTETensor input, const NVTETensor cu_seqlens,
                                 const NVTETensor freqs, NVTETensor output, const int cp_size,
                                 const int cp_rank, const int max_s, const int b, const int h,
                                 const int d, const int d2, const int stride_t, const int stride_h,
                                 const int stride_d, const int o_stride_t, const int o_stride_h,
                                 const int o_stride_d, cudaStream_t stream);

/*! \brief Compute the backward of the fused rope in thd format.
 *
 *  \param[in]     output_grads  Incoming gradient tensor for backward.
 *  \param[in]     cu_seqlens    The cumulative sum of sequence lengths tensor.
 *  \param[in]     freqs         The freqs tensor.
 *  \param[out]    input_grads   Input gradient to calculate.
 *  \param[in]     cp_size       Context parallel world size.
 *  \param[in]     cp_rank       Context parallel rank.
 *  \param[in]     max_s         Max sequence length.
 *  \param[in]     b             Batch size.
 *  \param[in]     h             Length of the h dimension of output_grads.
 *  \param[in]     d             Length of the d dimension of output_grads.
 *  \param[in]     d2            Length of the d dimension of freqs.
 *  \param[in]     stride_t      Stride of the t dimension of output_grads.
 *  \param[in]     stride_h      Stride of the h dimension of output_grads.
 *  \param[in]     stride_d      Stride of the d dimension of output_grads.
 *  \param[in]     o_stride_t    Stride of the t dimension of input_grads.
 *  \param[in]     o_stride_h    Stride of the h dimension of input_grads.
 *  \param[in]     o_stride_d    Stride of the d dimension of input_grads.
 *  \param[in]     stream        CUDA stream used for the operation.
 */
void nvte_fused_rope_thd_backward(const NVTETensor output_grads, const NVTETensor cu_seqlens,
                                  const NVTETensor freqs, NVTETensor input_grads, const int cp_size,
                                  const int cp_rank, const int max_s, const int b, const int h,
                                  const int d, const int d2, const int stride_t, const int stride_h,
                                  const int stride_d, const int o_stride_t, const int o_stride_h,
                                  const int o_stride_d, cudaStream_t stream);

#ifdef __cplusplus
}  // extern "C"
#endif

#endif  // TRANSFORMER_ENGINE_FUSED_ROPE_H_
