import torch
import torch_atb
import torch_npu


def quant_per_tensor(in_tensor: torch.Tensor, input_scale: torch.Tensor):
    return torch_npu.npu_quantize(in_tensor, input_scale, None, torch.qint8, -1, True)


def apply_atb(x, weight, weight_scale):
    linear_param = torch_atb.LinearParam()
    linear_param.has_bias = False
    linear_param.out_data_type = torch_atb.AclDataType.ACL_BF16
    linear = torch_atb.Operation(linear_param)
    return linear.forward([x, weight, weight_scale])


def op_prof(x, weight, weight_scale):
    prof = torch_npu.profiler.profile(
        activities=[
            torch_npu.profiler.ProfilerActivity.CPU,
            torch_npu.profiler.ProfilerActivity.NPU
        ],
        schedule=torch_npu.profiler.schedule(
            wait=1, warmup=1, active=1, repeat=1, skip_first=1),
        on_trace_ready=torch_npu.profiler.tensorboard_trace_handler(
            "./result"),
        with_stack=True,
        record_shapes=True,
        profile_memory=False,
        with_modules=False,
        with_flops=False,
    )
    prof.start()
    for step in range(5):
        apply_atb(x, weight, weight_scale)
        prof.step()
    prof.stop()


def main():
    M, K, N = 3, 7, 9
    print(f"M: {M}, K: {K}, N: {N}")
    x = torch.randn(M, K, dtype=torch.bfloat16).npu()
    weight = torch.randint(0, 10, (N, K), dtype=torch.int8).npu()
    weight_scale = torch.randn(N, dtype=torch.float32).npu()
    input_scale = torch.randn(K, dtype=torch.float32).npu()
    x = quant_per_tensor(x, input_scale)
    print("------------------- input -------------------------")
    print("x", x)
    print("weight", weight)
    print("weight_scale", weight_scale)

    atb = apply_atb(x, weight, weight_scale)

    print("\n------------------- output -------------------------")
    print("atb: ", atb[0])

    op_prof(x, weight, weight_scale)

main()

python demo.py
M: 3, K: 7, N: 9
------------------- input -------------------------
x tensor([[ 3,  0,  2,  0,  1,  1,  3],
        [ 2,  0,  0, -1, -1,  0, -6],
        [-7,  2,  3,  1,  0, -1, -3]], device='npu:0', dtype=torch.int8)
weight tensor([[0, 8, 3, 5, 8, 0, 3],
        [5, 6, 1, 8, 8, 1, 2],
        [0, 2, 9, 0, 6, 6, 2],
        [9, 6, 9, 6, 9, 1, 6],
        [2, 8, 5, 8, 3, 2, 7],
        [9, 7, 6, 8, 8, 0, 5],
        [1, 5, 7, 8, 5, 6, 1],
        [7, 5, 3, 5, 3, 5, 6],
        [9, 6, 9, 0, 4, 1, 2]], device='npu:0', dtype=torch.int8)
[W compiler_depend.ts:137] Warning: Warning: Device do not support double dtype now, dtype cast repalce with float. (function operator())
weight_scale tensor([ 4.2073e-01, -1.6913e+00, -4.5058e-01, -4.4355e-01,  8.4547e-01,
        -1.0513e-01, -8.8801e-01, -9.9923e-01,  5.2528e-01], device='npu:0')

------------------- output -------------------------
Segmentation fault (core dumped)
(KunLun) root@Atb_Test2:/home/sbw/matmul/atb/atb_case# [ERROR] TBE Subprocess[task_distribute] raise error[], main process disappeared!
[ERROR] TBE Subprocess[task_distribute] raise error[], main process disappeared!
[ERROR] TBE Subprocess[task_distribute] raise error[], main process disappeared!
[ERROR] TBE Subprocess[task_distribute] raise error[], main process disappeared!
[ERROR] TBE Subprocess[task_distribute] raise error[], main process disappeared!
[ERROR] TBE Subprocess[task_distribute] raise error[], main process disappeared!
[ERROR] TBE Subprocess[task_distribute] raise error[], main process disappeared!
[ERROR] TBE Subprocess[task_distribute] raise error[], main process disappeared!
/usr/local/miniconda3/envs/KunLun/lib/python3.10/multiprocessing/resource_tracker.py:224: UserWarning: resource_tracker: There appear to be 30 leaked semaphore objects to clean up at shutdown
  warnings.warn('resource_tracker: There appear to be %d '