# Copyright (c) Microsoft Corporation.
# SPDX-License-Identifier: Apache-2.0

# DeepSpeed Team

import torch
from ..config import DeepSpeedInferenceConfig
from .base import BaseOp


class MLPGemmOp(BaseOp):

    def __init__(self, config: DeepSpeedInferenceConfig):
        super(MLPGemmOp, self).__init__(config)
        if self.config.fp16:
            self.mlp_gemm_func = self.inference_cuda_module.mlp_gemm_fp16
        else:
            self.mlp_gemm_func = self.inference_cuda_module.mlp_gemm_fp32

    def forward(self, input: torch.Tensor, residual: torch.Tensor, input_bias: torch.Tensor,
                weight_interm: torch.Tensor, weight_out: torch.Tensor, bias: torch.Tensor, gamma: torch.Tensor,
                beta: torch.Tensor):
        output, residual_add = self.mlp_gemm_func(input, residual, input_bias, weight_interm, weight_out, bias, gamma,
                                                  beta, self.config.epsilon, self.config.pre_layer_norm,
                                                  self.config.mlp_after_attn, weight_interm.scale, weight_out.scale,
                                                  self.config.q_int8, self.config.mlp_act_func_type)
        return output, residual_add
