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

/*! \file gemm.h
 *  \brief Functions for matrix multiplication.
 */

#ifndef TRANSFORMER_ENGINE_GEMM_H_
#define TRANSFORMER_ENGINE_GEMM_H_

#include <stdint.h>

#include "transformer_engine.h"

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

/*! \brief Configuration for matrix multiplication. */
typedef void *NVTEMatmulConfig;

/*! \brief Configuration for grouped matrix multiplication. */
typedef void *NVTEGroupedMatmulConfig;

/*! \enum NVTEMatmulConfigAttribute
 * \brief Type of option for matrix multiplication.
 */
enum NVTEMatmulConfigAttribute {
  /*! Bias tensor
   *
   * If provided, the bias tensor is applied in the GEMM epilogue.
   */
  kNVTEMatmulConfigBiasTensor = 0,
  /*! Bias gradient tensor
   *
   * If provided, the bias gradient tensor will be filled in the GEMM epilogue.
   */
  kNVTEMatmulConfigDBiasTensor = 1,
  /*! Whether to compute GELU in GEMM epilogue. */
  kNVTEMatmulConfigWithGELUEpilogue = 2,
  /*! Whether to compute GELU backward in GEMM epilogue. */
  kNVTEMatmulConfigWithDGELUEpilogue = 3,
  /*! Auxilliary tensor for GEMM epilogue.
   *
   * For GELU, this will be filled with the GELU input. For GELU
   * backward, this is expected to already be filled with the GELU
   * input.
   */
  kNVTEMatmulConfigEpilogueAuxTensor = 4,
  /*! Whether to use split accumulator for FP8 GEMM. */
  kNVTEMatmulConfigUseSplitAccumulator = 5,
  /*! Number of streaming multiprocessors to use in GEMM kernel. */
  kNVTEMatmulConfigSMCount = 6,
  kNVTEMatmulConfigNumAttributes
};

/*! \enum NVTEGroupedMatmulConfigAttribute
 * \brief Type of option for grouped matrix multiplication.
 */
enum NVTEGroupedMatmulConfigAttribute {
  /*! Average M dimension hint
   *
   * Optional hint for average M dimension across all matrices in the group.
   * Used by cuBLASLt for algorithm selection heuristics. If not set,
   * computed automatically from D's logical shape.
   */
  kNVTEGroupedMatmulConfigAvgM = 0,
  /*! Average N dimension hint
   *
   * Optional hint for average N dimension across all matrices in the group.
   * Used by cuBLASLt for algorithm selection heuristics. If not set,
   * computed automatically from D's logical shape.
   */
  kNVTEGroupedMatmulConfigAvgN = 1,
  /*! Average K (reduction) dimension hint
   *
   * Optional hint for average K dimension across all matrices in the group.
   * Used by cuBLASLt for algorithm selection heuristics. If not set,
   * computed automatically from A's logical shape.
   */
  kNVTEGroupedMatmulConfigAvgK = 2,
  /*! Number of streaming multiprocessors to use in GEMM kernel. */
  kNVTEGroupedMatmulConfigSMCount = 3,
  /*! Split accumulator mode. Only taken into account on Hopper. Default: true. */
  kNVTEGroupedMatmulConfigUseSplitAccumulator = 4,
  kNVTEGroupedMatmulConfigNumAttributes
};

/*! \brief Create a matrix multiplication configuration. */
NVTEMatmulConfig nvte_create_matmul_config();

/*! \brief Query an option in matrix multiplication configuration.
 *
 *  \param[in]  config        Matrix multiplication configuration.
 *  \param[in]  attr          Option type.
 *  \param[out] buf           Memory address to write option value to.
 *                            Ignored if NULL.
 *  \param[in]  size_in_bytes Size of buf.
 *  \param[out] size_written  Number of bytes that have been written to
 *                            buf. If buf is NULL, then the number of
 *                            bytes that would have been written.
 */
void nvte_get_matmul_config_attribute(NVTEMatmulConfig config, NVTEMatmulConfigAttribute attr,
                                      void *buf, size_t size_in_bytes, size_t *size_written);

/*! \brief Set an option in matrix multiplication configuration.
 *
 *  \param[in,out] config        Matrix multiplication configuration.
 *  \param[in]     attr          Option type.
 *  \param[in]     buf           Memory address to read option value from.
 *  \param[in]     size_in_bytes Size of buf.
 */
void nvte_set_matmul_config_attribute(NVTEMatmulConfig config, NVTEMatmulConfigAttribute attr,
                                      const void *buf, size_t size_in_bytes);

/*! \brief Destroy a matrix multiplication configuration. */
void nvte_destroy_matmul_config(NVTEMatmulConfig config);

/*! \brief Create a grouped matrix multiplication configuration. */
NVTEGroupedMatmulConfig nvte_create_grouped_matmul_config();

/*! \brief Query an option in grouped matrix multiplication configuration.
 *
 *  \param[in] config Grouped matrix multiplication configuration.
 *  \param[in] attr Option type.
 *  \param[out] buf Memory address to write option value. Ignored if
 *                  NULL.
 *  \param[in] size_in_bytes Size of buf.
 *  \param[out] size_written Number of bytes that have been written to
 *                           buf. If buf is NULL, then the number of
 *                           bytes that would have been written.
 */
void nvte_get_grouped_matmul_config_attribute(NVTEGroupedMatmulConfig config,
                                              NVTEGroupedMatmulConfigAttribute attr, void *buf,
                                              size_t size_in_bytes, size_t *size_written);

/*! \brief Set an option in grouped matrix multiplication configuration.
 *
 *  \param[in] config Grouped matrix multiplication configuration.
 *  \param[in] attr Option type.
 *  \param[out] buf Memory address to read option value.
 *  \param[in] size_in_bytes Size of buf.
 */
void nvte_set_grouped_matmul_config_attribute(NVTEGroupedMatmulConfig config,
                                              NVTEGroupedMatmulConfigAttribute attr,
                                              const void *buf, size_t size_in_bytes);

/*! \brief Destroy a grouped matrix multiplication configuration. */
void nvte_destroy_grouped_matmul_config(NVTEGroupedMatmulConfig config);

/*! \brief Compute matrix multiplication of 2 matrices, potentially fused with other operations (deprecated).
 *
 * This has been deprecated in favor of nvte_cublas_gemm_v2.
 *
 * Computes:
 *  - `D = AB` if both `bias` and `pre_gelu_out` are empty tensors
 *  - `D = AB + bias` if `pre_gelu_out` is empty and `bias` is not empty
 *  - `D = GELU(AB + bias)` if both `bias` and `pre_gelu_out` are not empty tensors
 *
 *  \param[in]     A                     The A matrix.
 *  \param[in]     B                     The B matrix.
 *  \param[in,out] D                     Output matrix.
 *  \param[in]     bias                  Bias tensor.
 *  \param[in,out] pre_gelu_out          Output matrix before GELU activation.
 *  \param[in]     transa                Whether A matrix is transposed.
 *  \param[in]     transb                Whether B matrix is transposed.
 *  \param[in]     grad                  Whether this operation is part of the
 *                                       gradient computation.
 *  \param[out]    workspace             Workspace tensor.
 *  \param[in]     accumulate            Whether to accumulate the result into the D matrix.
 *  \param[in]     use_split_accumulator Whether to use split accumulator in the FP8 GEMM.
 *  \param[in]     math_sm_count         Number of GPU SMs to use (default=0: use cuBLAS heuristics)
 *  \param[in]     stream                CUDA stream used for the operation.
 */
void nvte_cublas_gemm(const NVTETensor A, const NVTETensor B, NVTETensor D, const NVTETensor bias,
                      NVTETensor pre_gelu_out, bool transa, bool transb, bool grad,
                      NVTETensor workspace, bool accumulate, bool use_split_accumulator,
                      int math_sm_count, cudaStream_t stream);

/*! \brief Compute matrix multiplication of 2 matrices, potentially fused with other operations.
 *
 * Computes:
 *  - `D = alpha * op(A) * op(B) + beta * C`
 *
 *  \param[in]  transa    Whether to transpose A matrix.
 *  \param[in]  transb    Whether to transpose B matrix.
 *  \param[in]  alpha     Scaling factor applied to matmul output.
 *  \param[in]  A         A matrix.
 *  \param[in]  B         B matrix.
 *  \param[in]  beta      Scaling factor applied to C matrix.
 *  \param[in]  C         C matrix.
 *  \param[out] D         Output matrix.
 *  \param[in]  workspace Workspace tensor.
 *  \param[in]  config    Additional configuration.
 *  \param[in]  stream    CUDA stream used for the operation.
 */
void nvte_cublas_gemm_v2(int transa, int transb, const float *alpha, const NVTETensor A,
                         const NVTETensor B, const float *beta, const NVTETensor C, NVTETensor D,
                         NVTETensor workspace, NVTEMatmulConfig config, cudaStream_t stream);

/*! \brief Compute matrix multiplication of 2 matrices, potentially fused with other operations,
 * allowing for using a scaling factor for the GEMM result and the accumulation input (deprecated)
 *
 * This has been deprecated in favor of nvte_cublas_gemm_v2.
 *
 * Computes:
 *  - `D = alpha*AB` if both `bias` and `pre_gelu_out` are empty tensors
 *  - `D = alpha*AB + bias` if `pre_gelu_out` is empty and `bias` is not empty
 *  - `D = GELU(alpha*AB + bias)` if both `bias` and `pre_gelu_out` are not empty tensors
 *
 *  \param[in]     A                     The A matrix.
 *  \param[in]     B                     The B matrix.
 *  \param[in,out] D                     Output matrix.
 *  \param[in]     bias                  Bias tensor.
 *  \param[in,out] pre_gelu_out          Output matrix before GELU activation.
 *  \param[in]     transa                Whether A matrix is transposed.
 *  \param[in]     transb                Whether B matrix is transposed.
 *  \param[in]     grad                  Whether this operation is part of the
 *                                       gradient computation.
 *  \param[out]    workspace             Workspace tensor.
 *  \param[in]     alpha                 Scaling factor applied to the result of the GEMM
 *  \param[in]     beta                  Scaling factor applied to original value of D when
 *                                       accumulating into it. beta=0 means no accumulation.
 *  \param[in]     use_split_accumulator Whether to use split accumulator in the FP8 GEMM.
 *  \param[in]     math_sm_count         Number of GPU SMs to use (default=0: use cuBLAS heuristics)
 *  \param[in]     stream                CUDA stream used for the operation.
 */
void nvte_cublas_gemm_scaled(const NVTETensor A, const NVTETensor B, NVTETensor D,
                             const NVTETensor bias, NVTETensor pre_gelu_out, bool transa,
                             bool transb, bool grad, NVTETensor workspace, float alpha, float beta,
                             bool use_split_accumulator, int math_sm_count, cudaStream_t stream);

/*! \brief Compute matrix multiplication of 2 matrices with chunking and atomic counters.
 *
 * \warning   Cublas atomic gemm uses a beta API and is not tested for all use cases.
 *
 * Computes:
 *  - `D = AB` if both `bias` and `pre_gelu_out` are empty tensors
 *  - `D = AB + bias` if `pre_gelu_out` is empty and `bias` is not empty
 *  - `D = GELU(AB + bias)` if both `bias` and `pre_gelu_out` are not empty tensors
 *
 *  \param[in]     A                     The A matrix.
 *  \param[in]     B                     The B matrix.
 *  \param[in,out] D                     Output matrix.
 *  \param[in]     bias                  Bias tensor.
 *  \param[in,out] pre_gelu_out          Output matrix before GELU activation.
 *  \param[in]     transa                Whether A matrix is transposed.
 *  \param[in]     transb                Whether B matrix is transposed.
 *  \param[in]     grad                  Whether this operation is part of the
 *                                       gradient computation.
 *  \param[out]    workspace             Workspace tensor.
 *  \param[in]     accumulate            Whether to accumulate the result into the D matrix.
 *  \param[in]     use_split_accumulator Whether to use split accumulator in the FP8 GEMM.
 *  \param[in]     math_sm_count         Number of GPU SMs to use (default=0: use cuBLAS heuristics)
 *  \param[in]     m_split               Number of chunks/splits along m-dimension for Atomic GEMM.
 *  \param[in]     n_split               Number of chunks/splits along n-dimension for Atomic GEMM.
 *  \param[in]     gemm_producer         Whether Atomic GEMM is the producer or consumer.
 *  \param[in,out] counter               counter[chunk_i]=0 indicates chunk_i has been produced.
 *  \param[in]     stream                CUDA stream used for the operation.
 */
void nvte_cublas_atomic_gemm(const NVTETensor A, const NVTETensor B, NVTETensor D,
                             const NVTETensor bias, NVTETensor pre_gelu_out, bool transa,
                             bool transb, bool grad, NVTETensor workspace, bool accumulate,
                             bool use_split_accumulator, int math_sm_count, int m_split,
                             int n_split, bool gemm_producer, const NVTETensor counter,
                             cudaStream_t stream);

/*! \brief Compute multiple pairs of matrix multiplication, potentially fused with other operations,
 * on multiple streams.
 *
 * Computes:
 *  - `D = AB` if both `bias` and `pre_gelu_out` are empty tensors
 *  - `D = AB + bias` if `pre_gelu_out` is empty and `bias` is not empty
 *  - `D = GELU(AB + bias)` if both `bias` and `pre_gelu_out` are not empty tensors
 *
 *  \param[in]     A                     The list of A matrices.
 *  \param[in]     B                     The list of B matrices.
 *  \param[in,out] D                     List of output matrices.
 *  \param[in]     bias                  List of bias tensors.
 *  \param[in,out] pre_gelu_out          List of output matrix before GELU activation.
 *  \param[in]     num_gemms             Number of GEMMs to compute.
 *  \param[in]     transa                Whether A matrix is transposed.
 *  \param[in]     transb                Whether B matrix is transposed.
 *  \param[in]     grad                  Whether this operation is part of the
 *                                       gradient computation.
 *  \param[out]    workspace             List of workspace tensors.
 *  \param[in]     accumulate            Whether to accumulate the result into the D matrix.
 *  \param[in]     use_split_accumulator Whether to use split accumulator in the FP8 GEMM.
 *  \param[in]     math_sm_count         Number of GPU SMs to use (default=0: use cuBLAS heuristics)
 *  \param[in]     stream                CUDA stream to wait on.
 */
void nvte_multi_tensor_gemm(const NVTETensor *A, const NVTETensor *B, NVTETensor *D,
                            const NVTETensor *bias, NVTETensor *pre_gelu_out, const int num_gemms,
                            bool transa, bool transb, bool grad, NVTETensor *workspace,
                            bool accumulate, bool use_split_accumulator, int math_sm_count,
                            cudaStream_t stream);

/*! \brief Return the required size in bytes for the setup workspace of grouped GEMM.
 *
 * The setup workspace stores pointer arrays and per-matrix dimension arrays used
 * by the grouped GEMM kernel. Its size depends only on the number of tensors (GEMMs)
 * in the group and is independent of matrix dimensions.
 *
 * Pass the result as the size of the workspace_setup tensor in nvte_grouped_gemm.
 *
 *  \param[in] num_tensors  Number of tensors (GEMMs) in the group.
 *  \return Required size in bytes for workspace_setup.
 */
size_t nvte_get_grouped_gemm_setup_workspace_size(size_t num_tensors);

/*! \brief Convert a device array of int32 values to int64 values.
 *
 *  Useful for preparing group_sizes for nvte_grouped_gemm when the caller
 *  holds int32 sizes and needs int64 values on the device.
 *
 *  \param[in]  src     Device pointer to source int32 array.
 *  \param[out] dst     Device pointer to destination int64 array.
 *  \param[in]  n       Number of elements.
 *  \param[in]  stream  CUDA stream.
 */
void nvte_convert_int32_to_int64(const int32_t *src, int64_t *dst, size_t n, cudaStream_t stream);

/*! \brief Convert int32 array to int64 while scaling each element by a multiplier.
 *
 *  Computes dst[i] = (int64_t)src[i] * multiplier for each i in [0, n).
 *  CUDA-graph safe (no host-device synchronization).
 *
 *  \param[in]  src         Device pointer to source int32 array.
 *  \param[out] dst         Device pointer to destination int64 array.
 *  \param[in]  n           Number of elements.
 *  \param[in]  multiplier  Scale factor applied to each element.
 *  \param[in]  stream      CUDA stream.
 */
void nvte_convert_int32_to_int64_with_multiplier(const int32_t *src, int64_t *dst, size_t n,
                                                 int64_t multiplier, cudaStream_t stream);

/*! \brief Compute exclusive prefix-sum offsets from per-group first-dimension sizes.
 *
 *  Writes n_groups+1 values to offsets: offsets[0]=0,
 *  offsets[i] = sum(first_dims[0..i-1] * last_dim) for i in [1, n_groups].
 *  This is CUDA-graph safe (no host-device synchronization).
 *
 *  \param[in]  first_dims  Device pointer to int64 array of length n_groups.
 *  \param[out] offsets     Device pointer to int64 array of length n_groups+1.
 *  \param[in]  n_groups    Number of groups.
 *  \param[in]  last_dim    Common last dimension (number of columns).
 *  \param[in]  stream      CUDA stream.
 */
void nvte_compute_grouped_tensor_offsets(const int64_t *first_dims, int64_t *offsets,
                                         size_t n_groups, int64_t last_dim, cudaStream_t stream);

/* EXPERIMENTAL FEATURE AND SUBJECT TO CHANGE. */
/*! \brief Grouped matrix multiplication: D = alpha * op(A) @ op(B) + beta * C
 *
 * \note Requires cuBLAS 13.2+ (CUDA 13.1+) and Blackwell (SM100) or newer GPU architecture.
 *       Will error at runtime if compiled with an older cuBLAS version or run on
 *       a pre-Blackwell GPU.
 *
 * Performs batched GEMM on a collection of matrices with potentially different shapes.
 * All tensors in the group must have compatible dimensions for matrix multiplication.
 * Uses NVTEGroupedTensor to efficiently handle collections of tensors with contiguous
 * memory layout and shape metadata.
 *
 *  \param[in]  A                Input grouped tensor A.
 *  \param[in]  transa           Whether to transpose A matrices.
 *  \param[in]  B                Input grouped tensor B.
 *  \param[in]  transb           Whether to transpose B matrices.
 *  \param[in]  C                Input grouped tensor C (can be NULL for beta=0).
 *  \param[out] D                Output grouped tensor D.
 *  \param[in]  alpha            Scale multipliers for A @ B (NVTETensor with num_tensors elements).
 *  \param[in]  beta             Scale multipliers for C (NVTETensor with num_tensors elements).
 *  \param[in]  workspace_setup  Workspace tensor for pointer array setup.
 *  \param[in]  workspace_cublas Workspace tensor for cuBLAS operations.
 *  \param[in]  config           Additional configuration (can be NULL for defaults).
 *  \param[in]  stream           CUDA stream for the operation.
 *
 * Requirements:
 * - cuBLAS 13.2+ (CUDA 13.1+)
 * - Blackwell (SM100) or newer GPU architecture
 * - A, B, C (if provided), D must have the same num_tensors
 * - For each i: D[i] = alpha[i] * op(A[i]) @ op(B[i]) + beta[i] * C[i]
 * - Shape compatibility: if transa=false, transb=false:
 *   - A[i]: (M[i], K[i]), B[i]: (K[i], N[i]), D[i]: (M[i], N[i])
 */
void nvte_grouped_gemm(const NVTEGroupedTensor A, int transa, const NVTEGroupedTensor B, int transb,
                       const NVTEGroupedTensor C, NVTEGroupedTensor D, const NVTETensor alpha,
                       const NVTETensor beta, NVTETensor workspace_setup,
                       NVTETensor workspace_cublas, NVTEGroupedMatmulConfig config,
                       cudaStream_t stream);

/* EXPERIMENTAL FEATURE AND SUBJECT TO CHANGE. */
/*! \brief Grouped matrix multiplication with discrete A input tensors.
 *
 * Identical to nvte_grouped_gemm, but A is provided as a list of tensors
 * instead of NVTEGroupedTensor. This enables discrete per-expert weights as inputA
 * for Grouped GEMM.
 *
 *  \param[in]  A_list           List of A tensors (length = num_a_tensors).
 *  \param[in]  num_a_tensors    Number of tensors in A_list.
 *  \param[in]  transa           Whether to transpose A matrices.
 *  \param[in]  B                Input grouped tensor B.
 *  \param[in]  transb           Whether to transpose B matrices.
 *  \param[in]  C                Input grouped tensor C (can be NULL for beta=0).
 *  \param[out] D                Output grouped tensor D.
 *  \param[in]  alpha            Scale multipliers for A @ B (NVTETensor with num_tensors elements).
 *  \param[in]  beta             Scale multipliers for C (NVTETensor with num_tensors elements).
 *  \param[in]  workspace_setup  Workspace tensor for pointer array setup.
 *  \param[in]  workspace_cublas Workspace tensor for cuBLAS operations.
 *  \param[in]  config           Additional configuration (can be NULL for defaults).
 *  \param[in]  stream           CUDA stream for the operation.
 */
void nvte_grouped_gemm_with_discrete_inputA(const NVTETensor *A_list, size_t num_a_tensors,
                                            int transa, const NVTEGroupedTensor B, int transb,
                                            const NVTEGroupedTensor C, NVTEGroupedTensor D,
                                            const NVTETensor alpha, const NVTETensor beta,
                                            NVTETensor workspace_setup, NVTETensor workspace_cublas,
                                            NVTEGroupedMatmulConfig config, cudaStream_t stream);

/* EXPERIMENTAL FEATURE AND SUBJECT TO CHANGE. */
/*! \brief Grouped matrix multiplication with discrete output tensors.
*
* Identical to nvte_grouped_gemm, but C and D are provided as lists of tensors
* instead of NVTEGroupedTensor. This enables accumulation into non-contiguous
* per-expert buffers (for wgrads).
*
*  \param[in]  A                Input grouped tensor A.
*  \param[in]  transa           Whether to transpose A matrices.
*  \param[in]  B                Input grouped tensor B.
*  \param[in]  transb           Whether to transpose B matrices.
*  \param[in]  C_list           Optional list of C tensors (length = num_c_tensors).
*  \param[in]  num_c_tensors    Number of tensors in C_list (Can be 0 if C is not provided).
*  \param[out] D_list           List of D tensors (length = num_d_tensors).
*  \param[in]  num_d_tensors    Number of tensors in D_list.
*  \param[in]  alpha            Scale multipliers for A @ B (NVTETensor with num_tensors elements).
*  \param[in]  beta             Scale multipliers for C (NVTETensor with num_tensors elements).
*  \param[in]  workspace_setup  Workspace tensor for pointer array setup.
*  \param[in]  workspace_cublas Workspace tensor for cuBLAS operations.
*  \param[in]  config           Additional configuration (can be NULL for defaults).
*  \param[in]  stream           CUDA stream for the operation.
*  \note  All tensors in C_list and D_list must share the same dtype.
*/
void nvte_grouped_gemm_with_discrete_out(const NVTEGroupedTensor A, int transa,
                                         const NVTEGroupedTensor B, int transb,
                                         const NVTETensor *C_list, size_t num_c_tensors,
                                         NVTETensor *D_list, size_t num_d_tensors,
                                         const NVTETensor alpha, const NVTETensor beta,
                                         NVTETensor workspace_setup, NVTETensor workspace_cublas,
                                         NVTEGroupedMatmulConfig config, cudaStream_t stream);

/*! \brief Grouped Bias add for grouped GEMM outputs.
*
* output[row,col] += bias[col].
* Requires uniform last-dimension across all output tensors and bias tensors.
*/
void nvte_grouped_bias_add(const NVTEGroupedTensor output, const NVTEGroupedTensor bias,
                           cudaStream_t stream);

/*! \brief Grouped Scaled Bias add for grouped GEMM outputs.
*
* output[row,col] += bias[col] * scale[row], where biases are per-group
* and scales are per-row across all groups.
* Requires uniform last-dimension across all output tensors and bias tensors.
*/
void nvte_grouped_scaled_bias_add(const NVTEGroupedTensor output, const NVTEGroupedTensor bias,
                                  const NVTETensor scale, cudaStream_t stream);

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

#ifdef __cplusplus

/*! \namespace transformer_engine
 */
namespace transformer_engine {

/*! \brief TE/JAX cudaGraph requires the cuBLAS initialization to happen outside of the capturing
 * region. This function is a helper to call cublasCreate() which allocate memory for the handle.
 * The function will be called in the initialize phase of the related XLA custom calls.
 */

void nvte_cublas_handle_init();

/*! \struct MatmulConfigWrapper
 *  \brief C++ wrapper for NVTEMatmulConfig.
 */
class MatmulConfigWrapper {
 public:
  MatmulConfigWrapper() : config_{nvte_create_matmul_config()} {}

  MatmulConfigWrapper(const MatmulConfigWrapper &) = delete;
  MatmulConfigWrapper &operator=(const MatmulConfigWrapper &) = delete;

  /*! \brief Move constructor. */
  MatmulConfigWrapper(MatmulConfigWrapper &&other) : config_{other.config_} {
    other.config_ = nullptr;
  }
  /*! \brief Move-assignment operator. */
  MatmulConfigWrapper &operator=(MatmulConfigWrapper &&other) {
    if (config_ != nullptr) {
      nvte_destroy_matmul_config(config_);
    }
    config_ = other.config_;
    other.config_ = nullptr;
    return *this;
  }

  ~MatmulConfigWrapper() {
    if (config_ != nullptr) {
      nvte_destroy_matmul_config(config_);
      config_ = nullptr;
    }
  }

  /*! \brief Get the underlying NVTEMatmulConfig.
   *
   *  \return NVTEMatmulConfig held by this MatmulConfigWrapper.
   */
  operator NVTEMatmulConfig() const noexcept { return config_; }

  /*! \brief Set bias tensor. */
  void set_bias_tensor(NVTETensor bias_tensor) {
    nvte_set_matmul_config_attribute(config_, kNVTEMatmulConfigBiasTensor, &bias_tensor,
                                     sizeof(NVTETensor));
  }

  /*! \brief Set bias gradient tensor. */
  void set_dbias_tensor(NVTETensor dbias_tensor) {
    nvte_set_matmul_config_attribute(config_, kNVTEMatmulConfigDBiasTensor, &dbias_tensor,
                                     sizeof(NVTETensor));
  }

  /*! \brief Set whether to compute GELU in GEMM epilogue. */
  void set_with_gelu_epilogue(bool with_gelu_epilogue) {
    const auto val = static_cast<uint8_t>(with_gelu_epilogue);
    nvte_set_matmul_config_attribute(config_, kNVTEMatmulConfigWithGELUEpilogue, &val, sizeof(val));
  }

  /*! \brief Set whether to compute GELU backward in GEMM epilogue. */
  void set_with_dgelu_epilogue(bool with_dgelu_epilogue) {
    const auto val = static_cast<uint8_t>(with_dgelu_epilogue);
    nvte_set_matmul_config_attribute(config_, kNVTEMatmulConfigWithDGELUEpilogue, &val,
                                     sizeof(val));
  }

  /*! \brief Set auxilliary tensor for GEMM epilogue. */
  void set_epilogue_aux_tensor(NVTETensor epilogue_aux_tensor) {
    nvte_set_matmul_config_attribute(config_, kNVTEMatmulConfigEpilogueAuxTensor,
                                     &epilogue_aux_tensor, sizeof(NVTETensor));
  }

  /*! \brief Set whether to use split accumulator for FP8 GEMM. */
  void set_use_split_accumulator(bool use_split_accumulator) {
    const auto val = static_cast<uint8_t>(use_split_accumulator);
    nvte_set_matmul_config_attribute(config_, kNVTEMatmulConfigUseSplitAccumulator, &val,
                                     sizeof(val));
  }

  /*! \brief Set number of streaming multiprocessors to use in GEMM kernel. */
  void set_sm_count(int sm_count) {
    const auto val = static_cast<int32_t>(sm_count);
    nvte_set_matmul_config_attribute(config_, kNVTEMatmulConfigSMCount, &val, sizeof(val));
  }

 private:
  /*! \brief Wrapped NVTEMatmulConfig. */
  NVTEMatmulConfig config_ = nullptr;
};

/*! \struct GroupedMatmulConfigWrapper
 *  \brief C++ wrapper for NVTEGroupedMatmulConfig.
 */
class GroupedMatmulConfigWrapper {
 public:
  GroupedMatmulConfigWrapper() : config_{nvte_create_grouped_matmul_config()} {}

  GroupedMatmulConfigWrapper(const GroupedMatmulConfigWrapper &) = delete;
  GroupedMatmulConfigWrapper &operator=(const GroupedMatmulConfigWrapper &) = delete;

  GroupedMatmulConfigWrapper(GroupedMatmulConfigWrapper &&other) : config_{other.config_} {
    other.config_ = nullptr;
  }
  GroupedMatmulConfigWrapper &operator=(GroupedMatmulConfigWrapper &&other) {
    if (config_ != nullptr) {
      nvte_destroy_grouped_matmul_config(config_);
    }
    config_ = other.config_;
    other.config_ = nullptr;
    return *this;
  }

  ~GroupedMatmulConfigWrapper() {
    if (config_ != nullptr) {
      nvte_destroy_grouped_matmul_config(config_);
      config_ = nullptr;
    }
  }

  /*! \brief Get the underlying NVTEGroupedMatmulConfig.
   *
   *  \return NVTEGroupedMatmulConfig held by this GroupedMatmulConfigWrapper.
   */
  operator NVTEGroupedMatmulConfig() const noexcept { return config_; }

  /*! \brief Set average M dimension hint for algorithm selection. */
  void set_avg_m(int64_t avg_m) {
    nvte_set_grouped_matmul_config_attribute(config_, kNVTEGroupedMatmulConfigAvgM, &avg_m,
                                             sizeof(int64_t));
  }

  /*! \brief Set average N dimension hint for algorithm selection. */
  void set_avg_n(int64_t avg_n) {
    nvte_set_grouped_matmul_config_attribute(config_, kNVTEGroupedMatmulConfigAvgN, &avg_n,
                                             sizeof(int64_t));
  }

  /*! \brief Set average K dimension hint for algorithm selection. */
  void set_avg_k(int64_t avg_k) {
    nvte_set_grouped_matmul_config_attribute(config_, kNVTEGroupedMatmulConfigAvgK, &avg_k,
                                             sizeof(int64_t));
  }

  /*! \brief Set number of streaming multiprocessors to use. */
  void set_sm_count(int sm_count) {
    nvte_set_grouped_matmul_config_attribute(config_, kNVTEGroupedMatmulConfigSMCount, &sm_count,
                                             sizeof(int));
  }

  /*! \brief Set split accumulator mode. Only taken into account on Hopper. */
  void set_use_split_accumulator(bool use_split_accumulator) {
    const auto val = static_cast<uint8_t>(use_split_accumulator);
    nvte_set_grouped_matmul_config_attribute(config_, kNVTEGroupedMatmulConfigUseSplitAccumulator,
                                             &val, sizeof(val));
  }

 private:
  /*! \brief Wrapped NVTEGroupedMatmulConfig. */
  NVTEGroupedMatmulConfig config_ = nullptr;
};

}  // namespace transformer_engine

#endif  // __cplusplus

#endif  // TRANSFORMER_ENGINE_GEMM_H_
