add ops

torch_mlu_ops-v1.3.2/benchmarks/benchmark_active.py

@@ -0,0 +1,64 @@
+import torch
+import torch_mlu
+import torch_mlu_ops as tmo
+from common import *
+import argparse
+from tabulate import tabulate
+import os
+
+e2e_time_param_dict_list = [{"batch": 1, "seq_len": 5, "hidden_size": 1024,
+                             "act_mode": "gelu", "is_gated": False, "input_dtype": [torch.bfloat16]},
+                            {"batch": 16, "seq_len": 5, "hidden_size": 1024,
+                             "act_mode": "gelu", "is_gated": False, "input_dtype": [torch.bfloat16]},
+                            {"batch": 72, "seq_len": 5, "hidden_size": 1024,
+                             "act_mode": "gelu", "is_gated": False, "input_dtype": [torch.bfloat16]},
+                            {"batch": 1024, "seq_len": 5, "hidden_size": 1024,
+                             "act_mode": "gelu", "is_gated": False, "input_dtype": [torch.bfloat16]},
+                            {"batch": 4096, "seq_len": 5, "hidden_size": 1024,
+                             "act_mode": "gelu", "is_gated": False, "input_dtype": [torch.bfloat16]},
+                            {"batch": 8192, "seq_len": 5, "hidden_size": 1024,
+                             "act_mode": "gelu", "is_gated": False, "input_dtype": [torch.bfloat16]},
+                            {"batch": 32768, "seq_len": 5, "hidden_size": 1024,
+                             "act_mode": "gelu", "is_gated": False, "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 = ["input_shape", "act_mode", "is_gated", "input_dtype", "hardware_time(us)", "e2e_latency(us)", "IO efficiency"]
+  contents = []
+  bd = get_band_width()
+  for params_dict in e2e_time_param_dict_list:
+    batch = params_dict["batch"]
+    seq_len = params_dict["seq_len"]
+    hidden_size = params_dict["hidden_size"]
+    act_mode = params_dict["act_mode"]
+    is_gated = params_dict["is_gated"]
+    input_dtype_list = params_dict["input_dtype"]
+    for dtype in input_dtype_list:
+      if dtype == torch.bfloat16 and not torch_mlu.mlu.is_bf16_supported():
+        dtype = torch.half
+      input = torch.randn(batch, seq_len, hidden_size).to(device).to(dtype)
+      hardware_time, e2e_time = benchmark_forward(tmo.active,
+                                                  input,
+                                                  act_mode,
+                                                  is_gated,
+                                                  repeats=args.repeat_times)
+      io_bytes = input.element_size() * input.nelement() * (2 - 0.5 * is_gated)
+      io_eff = io_bytes / hardware_time / bd
+      content = [f"{batch,seq_len,hidden_size}", f"{act_mode}", f"{is_gated}", 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()