/*************************************************************************
 * 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"
#include "utils.h"

namespace tmo {
namespace torch_api {

void smooth_quant(const at::Tensor &input,
                  const at::Tensor &input_scale,
                  const at::Tensor &output,
                  const at::Tensor &output_scale,
                  const c10::optional<at::Tensor> &input_zero,
                  const c10::optional<at::Tensor> &token_count,
                  const c10::optional<at::Tensor> &gather_index,
                  const c10::optional<at::Tensor> &gather_index_start_position,
                  const std::string &quant_mode,
                  const bool &dynamic_quant) {
  bool is_pertoken = quant_mode == "per_token";
  CHECK_TENSOR_CONTIGUOUS(input_scale)
  if (output_scale.dim() > 0) {
    CHECK_TENSOR_CONTIGUOUS(output_scale)
  }
  CHECK_OPTIONAL_TENSOR_CONTIGUOUS(input_zero)
  CHECK_OPTIONAL_TENSOR_CONTIGUOUS(token_count)
  CHECK_OPTIONAL_TENSOR_CONTIGUOUS(gather_index)
  CHECK_OPTIONAL_TENSOR_CONTIGUOUS(gather_index_start_position)

  TORCH_CHECK(isMlu(input), "input must be mlu tensor.");
  checkTensorSameAttr<TensorAttr::DEVICE>(input, input_scale, input_zero, token_count, gather_index,
                                          gather_index_start_position);

  TORCH_CHECK(quant_mode == "per_token" || quant_mode == "per_tensor",
              "quant_mode must be 'per_token' or per_tensor.")
  if (dynamic_quant) {
    TORCH_CHECK(is_pertoken, "only support per_token if dynamic_quant == true")
    TORCH_CHECK(output_scale.numel() > 0, "invalid output_scale if dynamic_quant = true")
  } else {
    TORCH_CHECK(output_scale.numel() <= 0, "not support output_scale if dynamic_quant = false")
  }
  if (gather_index.has_value() || token_count.has_value()) {
    TORCH_CHECK(input.dim() == 2, "input.dim() == 2 if has gather_index or token_count")
  }
  TORCH_CHECK(input.dim() >= 2, "input.dim() >= 2")
  TORCH_CHECK(output.dim() >= 2, "output.dim() >= 2")
  if (!gather_index.has_value()) {
    TORCH_CHECK(output.sizes() == input.sizes(), "input and output must have the same shape")
  }
  if (output_scale.numel() > 0) {
    if (!gather_index.has_value()) {
      TORCH_CHECK(std::vector<int64_t>(output_scale.sizes().begin(), output_scale.sizes().end()) ==
                      std::vector<int64_t>(input.sizes().begin(), input.sizes().end() - 1),
                  "output_scale_shape must be equal to input_shape[0:-1]")
    }
  }
  if (gather_index_start_position.has_value()) {
    TORCH_CHECK(gather_index.has_value(),
                "gather_index must exist if gather_index_start_position has value")
    TORCH_CHECK(gather_index.value().dim() == 1, "gather_index.dim() == 1")
  }

  const torch_mlu::mlu::MLUGuard device_guard(input.device());
  // flatten input
  at::Tensor input_flat = input.flatten(0, input.dim() - 2);
  TORCH_CHECK(input_flat.data_ptr() == input.data_ptr(), "check the input strides")
  // flatten output
  at::Tensor output_flat = output.flatten(0, output.dim() - 2);
  TORCH_CHECK(output_flat.data_ptr() == output.data_ptr(), "check the output stride")
  // flatten output_scale
  at::Tensor output_scale_flat;
  if (dynamic_quant) {
    output_scale_flat = output_scale.flatten(0, -1);
  }

  // Convert torch tensor to tensor descriptor
  auto descs = createTensorDescs(
      {input_flat, input_scale, input_zero.value_or(at::Tensor()),
       token_count.value_or(at::Tensor()), gather_index.value_or(at::Tensor()),
       gather_index_start_position.value_or(at::Tensor()), output_flat, output_scale_flat});

  auto input_shape = input.sizes();
  int64_t in_channel = input_shape.back();
  int64_t m =
      std::accumulate(input_shape.begin(), input_shape.end() - 1, 1, std::multiplies<int>());
  SmoothQuantTheory obj(m, in_channel, getCnnlDataType(input.scalar_type()), dynamic_quant);
  cnpxPush(obj);
  CNNL_CHECK_FATAL(cnnlSmoothQuantOnline_v4(
      torch_mlu::getCurrentHandle(), nullptr /*smooth_quant_online_desc*/, descs[0].get(),
      getAtTensorPtr(input_flat),                                  /*input*/
      descs[1].get(), getAtTensorPtr(input_scale),                 /*input_scale*/
      descs[2].get(), getAtTensorPtr(input_zero),                  /*input_zero*/
      descs[3].get(), getAtTensorPtr(token_count),                 /*token_count*/
      descs[4].get(), getAtTensorPtr(gather_index),                /*gather_index*/
      descs[5].get(), getAtTensorPtr(gather_index_start_position), /*gather_index_start_position*/
      descs[6].get(), getAtTensorPtr(output_flat),                 /*output*/
      descs[7].get(), getAtTensorPtr(output_scale_flat),           /*output_scale*/
      nullptr, nullptr,                                            /*output_zero*/
      nullptr, 0 /*worksapce*/));
  cnpxPop();
}

}  //  namespace torch_api
}  //  namespace tmo
