add ops

torch_mlu_ops-v1.3.2/benchmarks/benchmark_quantize.py

@@ -0,0 +1,61 @@
+import torch
+import torch_mlu
+import torch_mlu_ops as tmo
+from common import *
+import argparse
+from tabulate import tabulate
+import os
+import random
+
+params_dict = {"dynamic": [True],
+               "token_num": [1, 72, 490, 512, 525, 1024, 4096, 8192, 32768],
+               "hidden_size": [8192, 1024],
+               "input_dtype": [torch.bfloat16]}
+
+def main():
+  parser = argparse.ArgumentParser()
+  parser.add_argument('--repeat_times', type=int, default=10, help='repeat times for testing')
+  parser.add_argument('--csv', action='store_true', help='write the report data to csv')
+  parser.add_argument('-o', type=str, help='specify the output folder name under --csv mode')
+
+  args = parser.parse_args()
+  device = 'mlu'
+  titles = ["dynamic", "token_num", "hidden_size", "input_dytpe", "hardware_time(us)", "e2e_latency(us)", "IO efficiency"]
+  contents = []
+  params_list = product(params_dict["dynamic"], params_dict["token_num"], params_dict["hidden_size"], params_dict["input_dtype"])
+  bd = get_band_width()
+  for param in params_list:
+    dynamic, token_num, hidden_size, dtype = param[0], param[1], param[2], param[3]
+    input_shape = (token_num, hidden_size)
+    if dtype == torch.bfloat16 and not torch_mlu.mlu.is_bf16_supported():
+      dtype = torch.half
+    input = torch.randn(input_shape).to(device).to(dtype)
+    scale = torch.randn(input_shape[-1]).to(device).to(torch.float32)
+    zero = None
+    if dynamic:
+      hardware_time, e2e_time = benchmark_forward(tmo.per_token_smooth_quantize,
+                                                input,
+                                                scale,
+                                                zero,
+                                                None,
+                                                repeats=args.repeat_times)
+    else:
+      hardware_time, e2e_time = benchmark_forward(tmo.quantize,
+                                                input,
+                                                scale,
+                                                zero,
+                                                repeats=args.repeat_times)
+    io_bytes = (input.element_size() + 1) * input.nelement() + scale.element_size() * scale.nelement()
+    io_eff = io_bytes / hardware_time / bd
+    content = [f"{dynamic}", f"{token_num}", f"{hidden_size}", f"{dtype}", f"{hardware_time}", f"{e2e_time}", f"{io_eff}"]
+    contents.append(content)
+  table = [titles] + contents
+  print(tabulate(table, headers="firstrow", tablefmt="grid"))
+
+  if args.csv:
+    current_file_path = __file__
+    _, file_name = os.path.split(current_file_path)
+    save_to_csv(table, args.o, file_name)
+
+if __name__=="__main__":
+  main()