/*************************************************************************
 * 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.
 *************************************************************************/

#include "torch_ops_api.h"

namespace tmo {
namespace torch_api {

void fused_layernorm(const at::Tensor &input,
                     const at::Tensor &output,
                     const c10::optional<at::Tensor> &residual,
                     const c10::optional<at::Tensor> &gamma,
                     const c10::optional<at::Tensor> &beta,
                     const c10::optional<at::Tensor> &bias,
                     const c10::optional<at::Tensor> &quant_scale,
                     const c10::optional<at::Tensor> &residual_out,
                     const c10::optional<at::Tensor> &smooth_quant_scale,
                     const std::string &norm_mode,
                     double eps,
                     bool store_output_before_norm,
                     bool dynamic_quant) {
  // check device and dtype
  checkTensorSameAttr<TensorAttr::ALL>(input, residual, gamma, beta, bias);
  checkTensorSameAttr<TensorAttr::DEVICE>(input, output, residual_out, smooth_quant_scale);
  cnnlQuantizeScheme_t output_quant_scheme = CNNL_QUANTIZE_NONE;
  if (dynamic_quant) {
    TORCH_CHECK(quant_scale.has_value(), "dynamic_quant output, must have quant_scale");
    output_quant_scheme = CNNL_QUANTIZE_PER_TOKEN;
  } else if (quant_scale.has_value()) {
    output_quant_scheme = CNNL_QUANTIZE_PER_CHANNEL;
  }
  // check params
  bool has_residual = residual.has_value();
  bool quant_out = quant_scale.has_value();
  int hidden_size = input.size(-1);
  TORCH_CHECK(input.dim() >= 2, "input.dim() >= 2.");
  TORCH_CHECK(input.stride(-1) == 1, "input last dim must be contiguous.");
  TORCH_CHECK(input.sizes() == output.sizes(), "input and output must have the same shape");
  if (has_residual) {
    TORCH_CHECK(input.sizes() == residual.value().sizes(),
                "input and residual must have the same shape");
  }
  TORCH_CHECK(norm_mode == "layernorm" || norm_mode == "rmsnorm",
              "norm_mode must be 'layernorm' or 'rmsnorm'.");
  cnnlTransformerNormType_t mode =
      norm_mode == "layernorm" ? CNNL_TRANSFORMER_LAYERNORM : CNNL_TRANSFORMER_RMSNORM;
  if (norm_mode == "layernorm") {
    TORCH_CHECK(gamma.has_value() && beta.has_value(), "layernorm mode need gamma and beta.");
    TORCH_CHECK(gamma.value().sizes() == beta.value().sizes(),
                "gamma and beta must have the same shape")
    TORCH_CHECK(gamma.value().dim() == 1 && gamma.value().size(0) == hidden_size,
                "layernorm mode, gamma and beta size must be hidden_size.");
  } else {
    TORCH_CHECK(gamma.has_value(), "rmsnorm mode need gamma.");
    TORCH_CHECK(gamma.value().dim() == 1 && gamma.value().size(0) == hidden_size,
                "rmsnorm mode, gamma size must be hidden_size.");
  }
  if (quant_out) {
    TORCH_CHECK(quant_scale.value().dim() == 1 && quant_scale.value().size(0) == hidden_size,
                "quant_scale shape must be [hidden_size]");
  }
  const torch_mlu::mlu::MLUGuard device_guard(input.device());
  at::Tensor smooth_quant_scale_flat =
      dynamic_quant ? smooth_quant_scale.value().flatten() : at::Tensor();
  at::Tensor input_flat;
  at::Tensor output_flat;
  at::Tensor residual_flat;
  at::Tensor residual_out_flat;
  if (quant_out) {
    // input must be 2-dim, output dim must be same as input
    TORCH_CHECK(input.is_contiguous(), "quant_out is not support when input has stride.");
    input_flat = input.flatten(0, -2);
    output_flat = output.flatten(0, -2);
    residual_flat = has_residual ? residual.value().flatten(0, -2) : residual_flat;
    residual_out_flat =
        store_output_before_norm ? residual_out.value().flatten(0, -2) : residual_out_flat;
  } else if (input.dim() > 3) {
    input_flat = input.flatten(0, -3);
    output_flat = output.flatten(0, -3);
    residual_flat = has_residual ? residual.value().flatten(0, -3) : residual_flat;
    residual_out_flat =
        store_output_before_norm ? residual_out.value().flatten(0, -3) : residual_out_flat;
  } else {  // 2-dim or 3-dim
    input_flat = input;
    output_flat = output;
    residual_flat = has_residual ? residual.value() : residual_flat;
    residual_out_flat = store_output_before_norm ? residual_out.value() : residual_out_flat;
  }
  TORCH_CHECK(input.data_ptr() == input_flat.data_ptr(), "check the strides of input.");
  TORCH_CHECK(output_flat.data_ptr() == output.data_ptr(), "check the strides of output.");
  if (has_residual)
    TORCH_CHECK(residual.value().data_ptr() == residual_flat.data_ptr(),
                "check the strides of residual.");
  if (store_output_before_norm)
    TORCH_CHECK(residual_out.value().data_ptr() == residual_out_flat.data_ptr(),
                "check the strides of residual_out.");

  // create tensor desc
  auto descs = createTensorDescs({input_flat, gamma.value_or(at::Tensor()),
                                  beta.value_or(at::Tensor()), bias.value_or(at::Tensor()),
                                  residual_flat, quant_scale.value_or(at::Tensor()),
                                  residual_out_flat, output_flat, smooth_quant_scale_flat});
  auto compute_dtype = getCnnlDataType(input_flat.scalar_type());
  // forward
  auto handle = torch_mlu::getCurrentHandle();
  FusedNormTheory obj(input_flat.size(0), input_flat.size(-1), has_residual, bias.has_value(),
                      quant_out, dynamic_quant, store_output_before_norm, compute_dtype, norm_mode);
  cnpxPush(obj);
  CNNL_CHECK_FATAL(cnnlFuseNorm_v3(handle, descs[0].get(), getAtTensorPtr(input_flat),  // input
                                   descs[5].get(), getAtTensorPtr(quant_scale),    // input_scale
                                   descs[1].get(), getAtTensorPtr(gamma),          // norm_scale
                                   descs[2].get(), getAtTensorPtr(beta),           // norm_bias
                                   descs[4].get(), getAtTensorPtr(residual_flat),  // residual
                                   descs[3].get(), getAtTensorPtr(bias),           // bias
                                   eps, output_quant_scheme, store_output_before_norm, mode,
                                   compute_dtype, nullptr, 0,                    // set workspace
                                   descs[7].get(), getAtTensorPtr(output_flat),  // output
                                   descs[6].get(),
                                   getAtTensorPtr(residual_out_flat),  // residual_out
                                   descs[8].get(), getAtTensorPtr(smooth_quant_scale_flat)));
  cnpxPop();
}

}  // namespace torch_api
}  // namespace tmo
