/*************************************************************************
 * Copyright (C) [2023-2024] by Cambricon, Inc.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
 * OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
 * MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
 * IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
 * CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
 * TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
 * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 *************************************************************************/
#ifndef CSRC_KERNELS_MOE_ADD_BIAS_ACTIVATION_MLUH_
#define CSRC_KERNELS_MOE_ADD_BIAS_ACTIVATION_MLUH_
#include "../kernel_utils.h"
#include "cnnl.h"
namespace tmo {
/**
 * @brief Add bias and activate to all tokens. Different expert with different bias.
 *        If in gated mode, add bias to all input. But only activate in the
 *        input of [:, :inner_size]. Then multiply it to the input of [:, inner_size:].
 *        Else, add bias and activate to all input, has no multiply process.
 * @example
 * is_gated = true, num_expert = 4, total_tokens = 6, inner_size = 2
 * input: (6, 4) = [[2, 4, 5, 6], [1, 4, 5, 3],
 *                  [3, 5, 7, 8], [6, 8, 5, 3],
 *                  [2, 3, 4 ,5], [2, 9, 2, 3]]
 * bias: (4, 4) = [[1, 0, 1, 0], [0, 1, 2, 2], [2, 3, 2, 3], [1, 2, 3, 4]]
 * token_count = [2, 2, 1, 1]
 * first step: add bias
 * [[2+1, 4+0, 5+1, 6+0], [1+1, 4+0, 5+1, 3+0],
 *  [3+0, 5+1, 7+2, 8+2], [6+0, 8+1, 5+2, 3+2],
 *  [2+2, 3+3, 4+2, 5+3], [2+1, 9+2, 2+3, 3+4]]
 * second step: act and mul
 * output: (6, 2) = [[act(3)*6, act(4)*6],  [act(2)*6, act(4)*3],
 *                   [act(3)*9, act(6)*10], [act(6)*7, act(9)*5],
 *                   [act(4)*6, act(6)*8],  [act(3)*5, act(11)*7]]
 * @param queue: The queue for mlu.
 * @param output: Output. Pointer to the MLU memory that stores the result.
 *                When is_gated is true, The shape is [total_tokens, input_size / 2].
 *                In this case, the input_size must be even. Otherwise the shape is [total_tokens,
 * input_size]. The memory can be discontinuous in total_tokens dim. The stride is output_stride.
 * @param input: Input. Pointer to the MLU memory that stores the input tokens.
 *               The shape is [total_tokens, input_size].
 *               When is_gated is true, the shape is [total_tokens, 2 * inner_size].
 *               Otherwise the shape is [total_tokens, inner_size].
 * @param bias: Input. Pointer to the MLU memory that stores the bias. The memory must be
 * continuous. When is_gated is true, the shape is [num_expert, 2 * inner_size]. Otherwise the shape
 * is [num_expert, inner_size]. Bias can be nullptr. If bias is nullptr, has no add bias process.
 * @param cusum_token_count: Input. Pointer to the MLU memory that stores the prefix sum of token
 *                                  counts. The shape is [num_expert + 1]. If cusum_token_count
 *                                  is not nullptr, cusum_token_count, start_expert_id and
 *                                  expert_size together determine which tokens to process.
 *                                  If cusum_token_count is nullptr, process all tokens,
 *                                  the number of which is total_tokens. When bias is not nullptr,
 *                                  cusum_token_count must also not be nullptr.
 * @param num_expert: The number of expert.
 * @param total_tokens: The total number of tokens.
 * @param inner_size: The inner size of output.
 * @param output_stride: The stride of output, must be greater than or equal to inner_size.
 * @param dtype: Data type.
 * @param is_gated: Gated or not.
 * @param act_type: The type of activation. Support gelu and swish.
 * @param start_expert_id: The index of the start expert.
 * @param expert_size: The number of experts to process.
 * @param active_coef: The coefficient used in the swish activation.
 */
KernelStatus invokeGroupAddBiasActivationKernel(cnrtQueue_t queue,
                                                void *output,
                                                const void *input,
                                                const void *bias,
                                                const int *cusum_token_count,
                                                int num_expert,
                                                int total_tokens,
                                                int inner_size,
                                                int output_stride,
                                                cnnlDataType_t dtype,
                                                bool is_gated,
                                                cnnlActivationMode_t act_type,
                                                int start_expert_id,
                                                int expert_size,
                                                float active_coef);
}  // namespace tmo

#endif  // CSRC_KERNELS_MOE_ADD_BIAS_ACTIVATION_MLUH_
