Add micro-benchmark for FastGelu (#10744)

* Add micro-benchmark for FastGelu

* Delete the bert-base case, as it is very similar to the bert-large one.

* Add argument parsing and more user-friendly provider type assertion.

onnxruntime/python/tools/microbench/benchmark.py

@@ -1,9 +1,16 @@
+from argparse import ArgumentParser
 import time
 import numpy
 import onnxruntime as ort
 import torch
 
 
+def add_arguments(parser: ArgumentParser):
+    parser.add_argument("--provider", required=False, type=str, default="rocm", help="Execution provider to use")
+    parser.add_argument("--precision", required=False, type=str, default="fp16", help="Number format to use")
+    parser.add_argument('--profiling', type=bool, default=False, help='If enable profiling')
+
+
 def create_input_output_tensors(inputs, outputs):
     device = "cuda"
     input_tensors = {name: torch.from_numpy(array).to(device) for name, array in inputs.items()}
@@ -29,20 +36,28 @@ def create_io_binding(sess, input_tensors, output_tensors):
     return io_binding
 
 
-def create_session(onnx_file, provider, profiling):
+def create_session(onnx_file, args):
     sess_opt = ort.SessionOptions()
-    sess_opt.enable_profiling = profiling
-    if provider == "rocm":
+    sess_opt.enable_profiling = args.profiling
+    if args.provider == "rocm":
         execution_provider = ["ROCMExecutionProvider"] 
-    else:
+    elif args.provider == "cuda":
         execution_provider = ["CUDAExecutionProvider"] 
-        
+    else:
+        raise ValueError(f"The script doesn't support provider type '{args.provider}' yet.")
+
     sess = ort.InferenceSession(onnx_file, sess_options=sess_opt, providers=execution_provider)
+
+    if args.provider == "rocm":
+        assert 'ROCMExecutionProvider' in sess.get_providers()
+    elif args.provider == "cuda":
+        assert 'CUDAExecutionProvider' in sess.get_providers()
+
     return sess
  
 
-def benchmark(onnx_file, inputs, outputs, provider, profiling=False):
-    sess = create_session(onnx_file, provider, profiling)
+def benchmark(onnx_file, inputs, outputs, args):
+    sess = create_session(onnx_file, args)
     input_tensors, output_tensors = create_input_output_tensors(inputs, outputs)
     io_binding = create_io_binding(sess, input_tensors, output_tensors)
 

onnxruntime/python/tools/microbench/fast_gelu.py

@@ -0,0 +1,60 @@
+import argparse
+import numpy as np
+from benchmark import benchmark, add_arguments
+
+
+def create_inputs_outputs(batch, seq_len, intermediate_dimension, data_type):
+    np.random.seed(0)
+    a = np.random.rand(batch, seq_len, intermediate_dimension).astype(data_type)
+    b = np.random.rand(intermediate_dimension).astype(data_type)
+    c = np.random.rand(batch, seq_len, intermediate_dimension).astype(data_type)
+
+    inputs = {"A": a, "B": b}
+    outputs = {"return_val": c}
+ 
+    return inputs, outputs
+
+
+def add_benchmark_case(benchmark_cases, batch_size, seq_len, intermediate_dimension, data_type, model):
+    benchmark_cases += [
+        (batch_size, seq_len, intermediate_dimension, data_type, model),
+    ]
+
+
+def create_benchmark_cases(precision):
+    benchmark_cases = []
+    if precision == "fp16":
+      model = "models/fast_gelu_fp16.onnx"
+      data_type = np.float16
+    else:
+      model = "models/fast_gelu_fp32.onnx"
+      data_type = np.float32
+
+    # bert-large
+    hidden_size = 1024
+    seq_len = 384
+    batch_size = 1
+    intermediate_dimension = hidden_size * 4
+    add_benchmark_case(benchmark_cases, batch_size, seq_len, intermediate_dimension, data_type, model)
+
+    return benchmark_cases
+
+
+def benchmark_fast_gelu(batch, seq_len, intermediate_dimension, data_type, onnx_file, args):
+    inputs, outputs = create_inputs_outputs(batch, seq_len, intermediate_dimension, data_type)
+    time = benchmark(onnx_file, inputs, outputs, args)
+    return time
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    add_arguments(parser)
+    args = parser.parse_args()
+
+    for (batch, seq_len, intermediate_dimension, data_type, onnx_file) in create_benchmark_cases(args.precision):
+        time = benchmark_fast_gelu(batch, seq_len, intermediate_dimension, data_type, onnx_file, args)
+        print(f"(batch seq_len inter_dim) = ({batch} {seq_len} {intermediate_dimension}), {time:7.4f} ms")
+
+
+if __name__ == "__main__":
+    main()

onnxruntime/python/tools/microbench/matmul.py

@@ -1,5 +1,6 @@
+import argparse
 import numpy as np
-from benchmark import benchmark 
+from benchmark import benchmark, add_arguments
 
 
 def create_inputs_outputs(b1, b2, m, k, n, data_type):
@@ -23,7 +24,7 @@ def add_benchmark_case(benchmark_cases, batch_size, seq_len, hidden_size, interm
     ]
 
 
-def create_benchmark_cases(precision="fp16"):
+def create_benchmark_cases(precision):
     benchmark_cases = []
     if precision == "fp16":
       model = "models/matmul_fp16.onnx"
@@ -51,15 +52,19 @@ def create_benchmark_cases(precision="fp16"):
     return benchmark_cases
 
 
-def benchmark_matmul(b1, b2, m, k, n, data_type, onnx_file):
+def benchmark_matmul(b1, b2, m, k, n, data_type, onnx_file, args):
     inputs, outputs = create_inputs_outputs(b1, b2, m, k, n, data_type)
-    time = benchmark(onnx_file, inputs, outputs, "rocm")
+    time = benchmark(onnx_file, inputs, outputs, args)
     return time
 
 
 def main():
-    for (b1, b2, m, k, n, data_type, onnx_file) in create_benchmark_cases(): 
-        time = benchmark_matmul(b1, b2, m, k, n, data_type, onnx_file)
+    parser = argparse.ArgumentParser()
+    add_arguments(parser)
+    args = parser.parse_args()
+
+    for (b1, b2, m, k, n, data_type, onnx_file) in create_benchmark_cases(args.precision):
+        time = benchmark_matmul(b1, b2, m, k, n, data_type, onnx_file, args)
         tflops = b1 * b2 * m * k * n * 2 / time / 1000000000
         print(f"(b1 b2 m k n) = ({b1} {b2} {m} {k} {n}), {time:7.4f} ms, {tflops:4.2f} tflops")
 

onnxruntime/python/tools/microbench/models/fast_gelu_fp16.onnx

@@ -0,0 +1,23 @@
+p2o:·

++
+
A
+
B
+return_val"FastGelu:
com.microsoft	fast_geluZ,
+
A'
+%
+!
+batch
+	seq_len
+	inter_dimZ
+
B
+
+
+	inter_dimb5
+
+return_val'
+%
+!
+batch
+	seq_len
+	inter_dimB
+
com.microsoft

onnxruntime/python/tools/microbench/models/fast_gelu_fp32.onnx

@@ -0,0 +1,20 @@
+p2o:·

++
+
A
+
B
+return_val"FastGelu:
com.microsoft	fast_geluZ,
+
A'
+%
!
+batch
+	seq_len
+	inter_dimZ
+
B
+

+	inter_dimb5
+
+return_val'
+%
!
+batch
+	seq_len
+	inter_dimB
+
com.microsoft