From 040c0645e2ce600df1284d314976f797fc5f61c0 Mon Sep 17 00:00:00 2001
From: PeixuanZuo <94887879+PeixuanZuo@users.noreply.github.com>
Date: Wed, 16 Mar 2022 10:48:26 +0800
Subject: [PATCH] [ADD] Add micro-benchmark for Cast (#10870)

* [ADD] Add micro-benchmark for Cast

* [UPDATE] related to bert model and fix the format
---
 onnxruntime/python/tools/microbench/cast.py   | 70 +++++++++++++++++++
 .../microbench/models/cast_fp16tofp32.onnx    | 17 +++++
 .../microbench/models/cast_fp32tofp16.onnx    | 18 +++++
 3 files changed, 105 insertions(+)
 create mode 100644 onnxruntime/python/tools/microbench/cast.py
 create mode 100644 onnxruntime/python/tools/microbench/models/cast_fp16tofp32.onnx
 create mode 100644 onnxruntime/python/tools/microbench/models/cast_fp32tofp16.onnx

diff --git a/onnxruntime/python/tools/microbench/cast.py b/onnxruntime/python/tools/microbench/cast.py
new file mode 100644
index 0000000000..9f1ffae9a4
--- /dev/null
+++ b/onnxruntime/python/tools/microbench/cast.py
@@ -0,0 +1,70 @@
+import argparse
+from dataclasses import dataclass
+import numpy as np
+from benchmark import BenchmarkOp, add_arguments
+
+
+@dataclass
+class OpParam:
+    x : int
+    y : int
+    m : int
+    n : int
+    input_data_type : type
+    output_data_type : type
+
+
+@dataclass
+class ModelParam:
+    token_type_ids_dim0 : int
+    input_ids_dim1 : int
+
+
+class BenchmarkCast(BenchmarkOp):
+    def __init__(self, args):
+        BenchmarkOp.__init__(self, args)
+
+    def create_inputs_outputs(cls, op_param):
+        np.random.seed(0)
+        input_data = np.random.rand(op_param.x, op_param.y, op_param.m, op_param.n).astype(op_param.input_data_type)
+        output_data = np.random.rand(op_param.x, op_param.y, op_param.m, op_param.n).astype(op_param.output_data_type)
+        inputs = {"X": input_data}
+        outputs = {"Y": output_data}
+        return inputs, outputs
+
+    def add_model_cases(self, mp, model, input_data_type, output_data_type):
+        self.add_case(OpParam(1, mp.token_type_ids_dim0, mp.input_ids_dim1, 1024, input_data_type, output_data_type), model)
+        self.add_case(OpParam(1, mp.token_type_ids_dim0, mp.input_ids_dim1, 1, input_data_type, output_data_type), model)
+        self.add_case(OpParam(16, mp.token_type_ids_dim0, mp.input_ids_dim1, mp.input_ids_dim1, input_data_type, output_data_type), model)
+
+    def create_cases(self):
+        model = "models/cast_fp16tofp32.onnx" if self.args.precision == "fp16" else "models/cast_fp32tofp16.onnx"
+        input_data_type = np.float16 if self.args.precision == "fp16" else np.float32
+        output_data_type = np.float32 if self.args.precision == "fp16" else np.float16
+        # huggingface bert-large
+        self.add_case(OpParam(1, 1, 1, 1024, input_data_type, output_data_type), model)
+        self.add_case(OpParam(1, 1, 1024, 1024, input_data_type, output_data_type), model)
+        self.add_case(OpParam(1, 1, 1024, 4096, input_data_type, output_data_type), model)
+        self.add_case(OpParam(1, 1, 1024, 30522, input_data_type, output_data_type), model)
+        # huggingface bert-large with default dims
+        model_param = ModelParam(8, 512)
+        self.add_model_cases(model_param, model, input_data_type, output_data_type)
+        # huggingface bert-large with large input dims
+        model_param = ModelParam(32, 1024)
+        self.add_model_cases(model_param, model, input_data_type, output_data_type)
+
+    def case_profile(cls, op_param, time):
+        profile = f"(x y m n input_data_type) = ({op_param.x} {op_param.y} {op_param.m} {op_param.n} {op_param.input_data_type}), {time:7.4f} ms"
+        return profile
+    
+
+def main():
+    parser = argparse.ArgumentParser()
+    add_arguments(parser)
+    args = parser.parse_args()
+    bm = BenchmarkCast(args)
+    bm.benchmark()
+
+
+if __name__ == "__main__":
+    main()
diff --git a/onnxruntime/python/tools/microbench/models/cast_fp16tofp32.onnx b/onnxruntime/python/tools/microbench/models/cast_fp16tofp32.onnx
new file mode 100644
index 0000000000..6792e343b1
--- /dev/null
+++ b/onnxruntime/python/tools/microbench/models/cast_fp16tofp32.onnx
@@ -0,0 +1,17 @@
+cast_fp16tofp32:a
+
+
X
Y"Cast*	
+to
 
castZ
+
X
+
+
+
x
+
y
+
m
+
nb
+
Y
+

+
x
+
y
+
m
+
nB
\ No newline at end of file
diff --git a/onnxruntime/python/tools/microbench/models/cast_fp32tofp16.onnx b/onnxruntime/python/tools/microbench/models/cast_fp32tofp16.onnx
new file mode 100644
index 0000000000..7b7e148e71
--- /dev/null
+++ b/onnxruntime/python/tools/microbench/models/cast_fp32tofp16.onnx
@@ -0,0 +1,18 @@
+cast_fp32tofp16:a
+
+
X
Y"Cast*	
+to
+ 
castZ
+
X
+

+
x
+
y
+
m
+
nb
+
Y
+
+
+
x
+
y
+
m
+
nB
\ No newline at end of file