import torch
import torch_mlu
import unittest
import torch_mlu_ops as tmo
from common_utils import *
from itertools import product
import random
import os

def generate_token_count(num_expert,
                         total_token_count):
    token_count = torch.randint(low=1, high=1024, size=(num_expert, ), \
                                dtype=torch.int32).to(dtype=torch.float32)
    sum = torch.sum(token_count, dim=-1) * 1.0
    token_count *= total_token_count / sum.item()
    token_count = token_count.to(dtype=torch.int32)
    cusum_token_count = torch.zeros(num_expert + 1, dtype=torch.int32)
    end_expert_token_count = torch.ones(num_expert + 1, dtype=torch.int32) * total_token_count
    cusum_token_count[1:] = torch.cumsum(token_count, dim=-1)
    cusum_token_count = torch.minimum(cusum_token_count, end_expert_token_count)
    cusum_token_count[-1] = total_token_count
    token_count = cusum_token_count[1:] - cusum_token_count[0:-1]
    return token_count, cusum_token_count

def gen_case(num_tokens,
             topk,
             hidden_size,
             multi_scale,
             need_gather,
             num_expert,
             expert_size,
             start_expert_id,
             dtype,
             device):
    input = torch.randn((num_tokens, hidden_size), dtype=dtype, device=device)
    effective_count = num_tokens * topk if need_gather else num_tokens
    if multi_scale:
        token_count, cusum_token_count = generate_token_count(num_expert, effective_count)
        token_count = token_count.to(device=device)
        cusum_token_count = cusum_token_count.to(device=device)
        scale = torch.randn((num_expert, hidden_size), dtype=torch.float32, device=device)
    else:
        token_count = None
        cusum_token_count = None
        scale = torch.randn((hidden_size), dtype=torch.float32, device=device)

    if need_gather:
        gather_ids = torch.randperm(num_tokens * topk, dtype=torch.int32) // topk
        gather_ids = gather_ids.mlu()
    else:
        gather_ids = None

    if expert_size < num_expert and multi_scale:
        token_count = token_count[start_expert_id:]
        cusum_token_count = cusum_token_count[start_expert_id:]
        gather_index_start_position = cusum_token_count[0:1]
        scale = scale[start_expert_id:]
        effective_count = cusum_token_count[-1] - cusum_token_count[0]
    else:
        gather_index_start_position = None

    if not need_gather:
        gather_index_start_position = None

    return input, scale, gather_ids, token_count, cusum_token_count, gather_index_start_position, effective_count

def per_token_smooth_quantize_base(x: torch.Tensor,
                                   smooth: torch.Tensor,
                                   zero: torch.Tensor = None,
                                   token_count: torch.Tensor = None):
    output_shape = x.size()
    output_scale_shape = x.size()[0:-1]
    output, output_scale = QuantByRow(x.flatten(0, -2) * smooth, 8)
    return output.reshape(output_shape), output_scale.reshape(output_scale_shape)

def quantize_base(x: torch.Tensor,
                  scale: torch.Tensor,
                  zero: torch.Tensor = None
                  ) -> torch.Tensor:
   return (x * scale).round().clamp(-128.0, 127.0).to(torch.int8)

class TestSmoothQuantOp(BtTestCase):
    def run_gen_case(self, dic):
        dump_data = dic.pop('dump_data')
        if dump_data:
            self.launch(*dic.values())
        else:
            x = create_tensor_from_dic(dic['x'])
            smooth = create_tensor_from_dic(dic['smooth'])
            zero = None if dic['zero']['data'] is None else create_tensor_from_dic(dic['zero'])
            token_count = None if dic['token_count']['data'] is None else dic['token_count']['data']
            gather_index = None if dic['gather_index']['data'] is None else dic['gather_index']['data']
            gather_index_start_position = None if dic['gather_index_start_position']['data'] is None else dic['gather_index_start_position']['data']
            output = None if dic['output']['data'] is None else create_tensor_from_dic(dic['output'])
            output_scale = None if dic['output_scale']['data'] is None else create_tensor_from_dic(dic['output_scale'])
            dynamic_quant = dic['dynamic_quant']['data']
            self.launch(x, smooth, zero, token_count, gather_index, gather_index_start_position, output, output_scale, dynamic_quant)

    def launch(self, *args):
        tmo_out = tmo.moe_quantize(*args)
        out_base = None if args[6] is None else args[6].clone()
        scale_base = None if args[7] is None else args[7].clone()
        args = list(args)
        args[6] = out_base
        args[7] = scale_base
        torch_out = self.op_impl_base(*args)

        if args[-1]:
            self.assertTensorsEqual(torch_out[0].cpu().reshape(-1).float(),
                                        tmo_out[0].cpu().reshape(-1).float(),
                                        0.01, use_MSE=True, use_RAE=True)
            self.assertTensorsEqual(torch_out[1].cpu().reshape(-1).float(),
                                        tmo_out[1].cpu().reshape(-1).float(),
                                        0.003, use_MSE=True, use_RAE=True)
        else:
            self.assertTensorsEqual(torch_out.cpu().reshape(-1).float(),
                                        tmo_out.cpu().reshape(-1).float(),
                                        0.01, use_MSE=True, use_RAE=True)


    def op_impl_base(self, *args):
        x, smooth, zero, token_count, gather_index, gather_index_start_position, \
            output, output_scale, dynamic_quant = args
        input = x.to(dtype=torch.float32).cpu()
        input_scale = smooth.cpu()
        cusum_token_count = None
        if token_count is not None:
            token_count = token_count.cpu()
            cusum_token_count = torch.zeros(token_count.shape[0] + 1, dtype=torch.int32)
            cusum_token_count[1:] = torch.cumsum(token_count, dim=-1)

        if gather_index_start_position is not None:
            gather_index_start_position = gather_index_start_position.cpu()

        gather_index_start = 0
        if gather_index_start_position is not None:
            gather_index_start = gather_index_start_position[0]

        if cusum_token_count is not None:
            gather_index_end = cusum_token_count[-1] + gather_index_start
        elif gather_index is not None:
            gather_index_end = gather_index.numel()
        else:
            gather_index_end = input.numel()

        if gather_index is not None:
            gather_index = gather_index.cpu()
            gathered_input = input[gather_index[gather_index_start : gather_index_end]]
        else:
            gathered_input = input[gather_index_start : gather_index_end]

        if cusum_token_count is not None:
            for i in range(token_count.shape[0]):
                gathered_input[cusum_token_count[i] : cusum_token_count[i+1]] *= input_scale[i]
        else:
            gathered_input *= input_scale

        if output is None:
            if not dynamic_quant:
                return gathered_input.round().clamp(-128.0, 127.0).to(torch.int8), None
            else:
                return QuantByRow(gathered_input, 8)
        else:
            if not dynamic_quant:
                output.copy_(gathered_input.round().clamp(-128.0, 127.0).to(torch.int8))
                output_scale = None
            else:
                out, scale = QuantByRow(gathered_input, 8)
                output_fl = output.flatten()
                output_fl[:out.numel()].copy_(out.flatten())
                # output = output_fl
                output_scale_fl = output_scale.flatten()
                output_scale_fl[:scale.numel()].copy_(scale.flatten())
                # output_scale = output_scale_fl
        return (output, output_scale) if dynamic_quant else (output,)

    def test_random_case(self):
        torch.manual_seed(333)
        test_cases = 100
        num_tokens_list = torch.randint(low=1, high=4096, size=(test_cases, ), dtype=torch.int32)
        topk_list = torch.randint(low=1, high=16, size=(test_cases, ), dtype=torch.int32)
        hidden_size_list = torch.randint(low=128, high=8193, size=(test_cases, ), dtype=torch.int32)
        num_expert_list = torch.randint(low=1, high=129, size=(test_cases, ), dtype=torch.int32)
        expert_size_list = torch.randint(low=1, high=129, size=(test_cases, ), dtype=torch.int32)
        expert_size_list = torch.minimum(expert_size_list, num_expert_list)
        start_expert_id_list = torch.randint(low=0, high=129, size=(test_cases, ), dtype=torch.int32)
        start_expert_id_list = torch.minimum(start_expert_id_list, num_expert_list - expert_size_list)
        start_expert_id_list = torch.maximum(start_expert_id_list, torch.zeros(test_cases, dtype=torch.int32))
        multi_scale_list = torch.randint(low=0, high=2, size=(test_cases, ), dtype=torch.int32)
        need_gather_list = torch.randint(low=0, high=2, size=(test_cases, ), dtype=torch.int32)
        input_with_stride_list = torch.randint(low=0, high=2, size=(test_cases, ), dtype=torch.int32)
        dynamic_quant_list = torch.randint(low=0, high=2, size=(test_cases, ), dtype=torch.int32)
        dtype_list = torch.randint(low=0, high=10, size=(test_cases, ), dtype=torch.int32)
        dtypes = [torch.half, torch.half, torch.float32]
        if torch_mlu.mlu.is_bf16_supported():
            dtypes += [torch.bfloat16, torch.bfloat16]
        dtype_list = random.choices(dtypes, k=test_cases)

        device = "mlu"
        for i in range(test_cases):
            num_tokens = num_tokens_list[i].item()
            topk = topk_list[i].item()
            hidden_size = hidden_size_list[i].item()
            num_expert = num_expert_list[i].item()
            expert_size = expert_size_list[i].item()
            start_expert_id = start_expert_id_list[i].item()
            multi_scale = multi_scale_list[i].item() == 1
            need_gather = need_gather_list[i].item() == 1
            input_with_stride = input_with_stride_list[i].item() == 1
            dynamic_quant = dynamic_quant_list[i].item() == 1
            dtype = dtype_list[i]

            if not multi_scale or not torch_mlu.mlu.is_bf16_supported():
                need_gather = False

            inputs = gen_case(num_tokens,
                              topk,
                              hidden_size,
                              multi_scale,
                              need_gather,
                              num_expert,
                              expert_size,
                              start_expert_id,
                              dtype,
                              device)

            input = inputs[0]
            input_scale = inputs[1]
            gather_ids = inputs[2]
            token_count = inputs[3]
            cusum_token_count = inputs[4]
            gather_index_start_position = inputs[5]
            effective_count = inputs[6]

            if input_with_stride:
                hidden_size = hidden_size - 64
                input = input[..., hidden_size : ]
                input_scale = input_scale[..., hidden_size : ].contiguous()

            print("num_tokens={}, topk={}, hidden_size={}, num_expert={}, expert_size={}, "
                  "start_expert_id={}, multi_scale={}, need_gather={}, input_with_stride={}, "
                  "dynamic_quant={}, dtype={}, testing...".format(
                      num_tokens, topk, hidden_size, num_expert, expert_size, start_expert_id, \
                      multi_scale, need_gather, input_with_stride, dynamic_quant, dtype))

            torch_quant, torch_output_scale = self.op_impl_base(input,
                                                                input_scale,
                                                                None,
                                                                token_count,
                                                                gather_ids,
                                                                gather_index_start_position,
                                                                None,
                                                                None,
                                                                dynamic_quant)
            tmo_output_scale = None
            if dynamic_quant:
                tmo_output, tmo_output_scale = \
                    tmo.moe_quantize(input, input_scale, None, token_count,
                                     gather_ids, gather_index_start_position,
                                     None, None, dynamic_quant)
            else:
                tmo_output, = \
                    tmo.moe_quantize(input, input_scale, None, token_count,
                                     gather_ids, gather_index_start_position,
                                     None, None, dynamic_quant)
            tmo_output = tmo_output[:effective_count]
            if tmo_output_scale is not None:
                tmo_output_scale = tmo_output_scale[:effective_count]

            self.assertTensorsEqual(torch_quant.cpu().reshape(-1).float(),
                                    tmo_output.cpu().reshape(-1).float(),
                                    0.01, use_MSE=True, use_RAE=True)

            if dynamic_quant:
                self.assertTensorsEqual(torch_output_scale.cpu().reshape(-1).float(),
                                        tmo_output_scale.cpu().reshape(-1).float(),
                                        0.003, use_MSE=True, use_RAE=True)

    def test_interface(self):
        channel = 16
        dtype_list = [torch.half, torch.float]
        if torch_mlu.mlu.is_bf16_supported():
            dtype_list.append(torch.bfloat16)
        for dtype in dtype_list:
            input = torch.randn((2, 3, 2, channel), dtype=dtype, device="mlu")
            input_scale = torch.randn((channel)).float().mlu()

            print("test tmo.per_token_smooth_quantize...")
            torch_quant, torch_scale = per_token_smooth_quantize_base(input, input_scale)
            tmo_quant, tmo_scale = tmo.per_token_smooth_quantize(input, input_scale)

            self.assertTensorsEqual(torch_quant.cpu().float(), tmo_quant.cpu().float(),
                                    0.01, use_MSE=True, use_RAE=True)
            self.assertTensorsEqual(torch_scale.cpu().float(), tmo_scale.cpu().float(),
                            0.003, use_MSE=True, use_RAE=True)

            print("test tmo.quantize ...")
            input_scale = input_scale * 100.0
            torch_quant = quantize_base(input, input_scale)
            tmo_quant = tmo.quantize(input, input_scale)
            self.assertTensorsEqual(torch_quant.cpu().float(), tmo_quant.cpu().float(),
                                    0.003, use_MSE=True, use_RAE=True)

        print("test tmo.moe_quantize inplace ...")
        num_tokens, topk, hidden_size, multi_scale, need_gather, num_expert, expert_size, start_expert_id, dtype, dynamic_quant = \
        1024, 5, 512, True, True, 32, 4, 4, torch.half, True
        inputs = gen_case(num_tokens, topk, hidden_size, multi_scale, need_gather, num_expert,
                            expert_size, start_expert_id, dtype, 'mlu')
        input = inputs[0]
        input_scale = inputs[1]
        gather_ids = inputs[2]
        token_count = inputs[3]
        cusum_token_count = inputs[4]
        gather_index_start_position = inputs[5]
        effective_count = inputs[6]
        tmo_output = torch.empty(num_tokens*topk, hidden_size, dtype=torch.int8, device='mlu')
        tmo_output_scale = torch.empty(num_tokens*topk, dtype=torch.float, device='mlu')
        if 'MLU370' not in torch_mlu.mlu.get_device_name():
            tmo.moe_quantize(input, input_scale, None, token_count,
                                gather_ids, gather_index_start_position,
                                tmo_output, tmo_output_scale, dynamic_quant)
            tmo_output = tmo_output[:effective_count]
            tmo_output_scale = tmo_output_scale[:effective_count]
            torch_output = torch.empty_like(tmo_output)
            torch_output_scale = torch.empty_like(tmo_output_scale)
            self.op_impl_base(input, input_scale,
                    None, token_count, gather_ids, gather_index_start_position,
                    torch_output, torch_output_scale, dynamic_quant)
            self.assertTensorsEqual(torch_output.cpu().reshape(-1).float(),
                                    tmo_output.cpu().reshape(-1).float(),
                                    0.01, use_MSE=True, use_RAE=True)
            self.assertTensorsEqual(torch_output_scale.cpu().reshape(-1).float(),
                                    tmo_output_scale.cpu().reshape(-1).float(),
                                    0.003, use_MSE=True, use_RAE=True)

        input = input.reshape(2, 4, 8, 16, -1)
        input_scale = input_scale[0]
        tmo_output = torch.empty(input.size(), dtype=torch.int8, device='mlu')
        tmo_output_scale = torch.empty(input.size()[:-1], dtype=torch.float, device='mlu')
        tmo.moe_quantize(input, input_scale, None, None, None, None,
                         tmo_output, tmo_output_scale, dynamic_quant)
        torch_quant = torch.empty_like(tmo_output)
        torch_output_scale = torch.empty_like(tmo_output_scale)
        self.op_impl_base(input, input_scale,
                None, None, None, None, torch_quant, torch_output_scale, dynamic_quant)
        self.assertTensorsEqual(torch_quant.cpu().reshape(-1).float(),
                                tmo_output.cpu().reshape(-1).float(),
                                0.01, use_MSE=True, use_RAE=True)
        self.assertTensorsEqual(torch_output_scale.cpu().reshape(-1).float(),
                                tmo_output_scale.cpu().reshape(-1).float(),
                                0.003, use_MSE=True, use_RAE=True)

    @unittest.skipIf('TMO_MEM_CHECK' in os.environ, "Skipping test_prevent due to ASan issues")
    def test_prevent(self):
        num_tokens, topk, hidden_size, multi_scale, need_gather, num_expert, expert_size, start_expert_id, dtype, dynamic_quant = \
        1024, 5, 512, True, True, 32, 4, 4, torch.half, True
        inputs = gen_case(num_tokens, topk, hidden_size, multi_scale, need_gather, num_expert,
                            expert_size, start_expert_id, dtype, 'mlu')
        input = inputs[0]
        input_scale = inputs[1]
        gather_ids = inputs[2]
        token_count = inputs[3]
        cusum_token_count = inputs[4]
        gather_index_start_position = inputs[5]
        effective_count = inputs[6]
        tmo_output = torch.empty(input.size(), dtype=torch.int8, device='mlu')
        tmo_output_scale = torch.empty(input.size()[:-1], dtype=torch.float, device='mlu')
        func = tmo.moe_quantize

        self.assertException("input must be mlu tensor.",
                             func, input.cpu(), input_scale, None, None, None, None,
                             None, None, dynamic_quant)
        self.assertException(None,
                             func, input, input_scale, None, token_count, gather_ids,
                             gather_index_start_position.cpu(), None, None, dynamic_quant)
        self.assertException("not support output_scale if dynamic_quant = false",
                             func, input, input_scale, None, token_count, gather_ids,
                             gather_index_start_position, None, tmo_output_scale, False)
        self.assertException("input.dim() == 2 if has gather_index or token_count",
                             func, input.reshape(32, 32, -1), input_scale, None, None, gather_ids,
                             gather_index_start_position, None, None, dynamic_quant)
        self.assertException("input.dim() >= 2",
                             func, input.reshape(-1), input_scale, None, None, None, None,
                             None, None, dynamic_quant)
        self.assertException("output.dim() >= 2",
                             func, input, input_scale, None, token_count, gather_ids,
                             gather_index_start_position, tmo_output.reshape(-1), tmo_output_scale, True)
        self.assertException("input and output must have the same shape",
                             func, input.reshape(2, 512, -1), input_scale, None, None, None, None,
                             tmo_output.reshape(32, 32, -1), None, dynamic_quant)
        self.assertException("output_scale_shape must be equal to input_shape[0:-1]",
                             func, input.reshape(2, 512, -1), input_scale[0], None, None, None, None,
                             None, tmo_output_scale.reshape(32, 32), dynamic_quant)
        self.assertException("gather_index must exist if gather_index_start_position has value",
                             func, input, input_scale, None, token_count, None,
                             gather_index_start_position, None, None, True)
        self.assertException("gather_index.dim() == 1",
                             func, input, input_scale, None, token_count, gather_ids.reshape(1, -1),
                             gather_index_start_position, None, None, True)


    def test_perf_case(self):
        num_tokens_list = [1, 72, 512]
        topk = 5
        hidden_size_list = [2048, 4096, 5120, 8192]
        # [num_expert, start_expert_id, expert_size]
        expert_options_list = [[8, 0, 8], [32, 24, 8]]
        multi_scale_list = [True, False]
        need_gather_list = [True, False]
        dynamic_quant_list = [True]
        dtype_list = [torch.half, torch.bfloat16]

        device = 'mlu'
        args = product(num_tokens_list, hidden_size_list, expert_options_list,\
                       multi_scale_list, need_gather_list, dynamic_quant_list, dtype_list)
        for num_tokens, hidden_size, expert_options, multi_scale, need_gather, dynamic_quant, dtype in args:
            num_expert = expert_options[0]
            start_expert_id = expert_options[1]
            expert_size = expert_options[2]

            if not multi_scale or not torch_mlu.mlu.is_bf16_supported():
                need_gather = False

            if not torch_mlu.mlu.is_bf16_supported():
                continue

            torch.manual_seed(444)
            inputs = gen_case(num_tokens,
                              topk,
                              hidden_size,
                              multi_scale,
                              need_gather,
                              num_expert,
                              expert_size,
                              start_expert_id,
                              dtype,
                              device)

            input = inputs[0]
            input_scale = inputs[1]
            gather_ids = inputs[2]
            token_count = inputs[3]
            cusum_token_count = inputs[4]
            gather_index_start_position = inputs[5]
            effective_count = inputs[6]

            print("num_tokens={}, hidden_size={}, num_expert={}, expert_size={}, "
                  "start_expert_id={}, multi_scale={}, need_gather={}, input_with_stride={}, "
                  "dynamic_quant={}, dtype={}, testing...".format(
                      num_tokens, hidden_size, num_expert, expert_size, start_expert_id, \
                      multi_scale, need_gather, False, dynamic_quant, dtype))

            torch_quant, torch_output_scale = self.op_impl_base(input,
                                                                input_scale,
                                                                None,
                                                                token_count,
                                                                gather_ids,
                                                                gather_index_start_position,
                                                                None,
                                                                None,
                                                                dynamic_quant)
            notify_start = torch.mlu.Event(enable_timing=True)
            notify_end = torch.mlu.Event(enable_timing=True)
            notify_start.record()
            loop = 10
            for _ in range(loop):
                if dynamic_quant:
                    tmo_output, tmo_output_scale = \
                        tmo.moe_quantize(input, input_scale, None, token_count,
                                         gather_ids, gather_index_start_position,
                                         None, None, dynamic_quant)
                else:
                    tmo_output, = \
                        tmo.moe_quantize(input, input_scale, None, token_count,
                                         gather_ids, gather_index_start_position,
                                         None, None, dynamic_quant)

            notify_end.record()
            notify_end.synchronize()
            time = notify_start.hardware_time(notify_end) / loop

            tmo_output = tmo_output[:effective_count]
            if tmo_output_scale is not None:
                tmo_output_scale = tmo_output_scale[:effective_count]
            print("time is: {:.1f}us".format(time))

            self.assertTensorsEqual(torch_quant.cpu().reshape(-1).float(),
                                    tmo_output.cpu().reshape(-1).float(),
                                    0.01, use_MSE=True, use_RAE=True)
            if dynamic_quant:
                self.assertTensorsEqual(torch_output_scale.cpu().reshape(-1).float(),
                                        tmo_output_scale.cpu().reshape(-1).float(),
                                        0.003, use_MSE=True, use_RAE=True)

    def test_inductor(self):
        multi_scale_list = [True, False]
        need_gather_list = [True, False] if 'MLU370' not in torch_mlu.mlu.get_device_name() else [False]
        dynamic_quant_list = [True, False]
        num_tokens, hidden_size, num_expert, start_expert_id, expert_size, topk, \
            dtype, device = 1, 2048, 32, 24, 8, 5, torch.float16, 'mlu'
        params = product(multi_scale_list, need_gather_list, dynamic_quant_list)
        # quantize
        print(f"check ops.quantize...")
        input, input_scale, gather_ids, token_count, _, _, _ = gen_case(num_tokens,
                                                                        topk,
                                                                        hidden_size,
                                                                        False,
                                                                        False,
                                                                        num_expert,
                                                                        expert_size,
                                                                        start_expert_id,
                                                                        dtype,
                                                                        device)
        output = torch.empty(input.size(), dtype=torch.int8, device=device)
        args = (input, input_scale, output, torch.Tensor(), None, None, None, None, 'per_token', False)
        self.base_opcheck(torch.ops.torch_mlu_ops.smooth_quant, args)

        # per_token_smooth_quantize
        print(f"check ops.per_token_smooth_quantize...")
        output_scale = torch.empty(input.size()[:-1], dtype=input_scale.dtype, device=device)
        args = (input, input_scale, output, output_scale, None, token_count, None, None, 'per_token', True)
        self.base_opcheck(torch.ops.torch_mlu_ops.smooth_quant, args)

        # moe_quantize
        for multi_scale, need_gather, dynamic_quant in params:
            if not multi_scale and need_gather:
                continue
            print(f"check ops.moe_quantize multi_scale: {multi_scale}, need_gather: {need_gather}, dynamic_quant: {dynamic_quant} ...")
            input, input_scale, gather_ids, token_count, _, \
                gather_index_start_position, _ = gen_case(num_tokens,
                                                            topk,
                                                            hidden_size,
                                                            multi_scale,
                                                            need_gather,
                                                            num_expert,
                                                            expert_size,
                                                            start_expert_id,
                                                            dtype,
                                                            device)
            output_shape = list(input.size())
            output_scale_shape = list(input.size()[:-1])
            if gather_ids is not None:
                output_tokens = gather_ids.size(0)
                output_shape[0] = output_tokens
                output_scale_shape[0] = output_tokens
            output = torch.empty(output_shape, dtype=torch.int8, device=device)
            output_scale = torch.empty(output_scale_shape, dtype=input_scale.dtype, device=device) if dynamic_quant else None
            args = (input, input_scale,
                    output, output_scale, None,
                    token_count, gather_ids, gather_index_start_position,
                    'per_token', dynamic_quant)
            self.base_opcheck(torch.ops.torch_mlu_ops.smooth_quant, args)

if __name__ == '__main__':
    exit(run_unittest(TestSmoothQuantOp))
