Skip to content

Commit a1f4b2e

Browse files
FindHaofacebook-github-bot
authored andcommitted
Add multiple ops support for --op argument (#2490)
Summary: Allow users benchmark multiple ops in a single run. The ops can be split by commas, `--op fp8_gemm,addmm` Example output: ``` % python run_benchmark.py triton --op fp8_gemm,addmm --num-inputs 1 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:03<00:00, 3.12s/it] x_val torch_fp8_gemm-gbps torch_fp8_gemm-gbps torch_fp8_gemm-latency torch_fp8_gemm-tflops triton_fp8_gemm-gbps triton_fp8_gemm-gbps triton_fp8_gemm-latency triton_fp8_gemm-tflops ------------------ --------------------- --------------------- ------------------------ ----------------------- ---------------------- ---------------------- ------------------------- ------------------------ (1024, 1024, 1024) 462.202 462.202 0.00907462 236.647 630.43 630.43 0.00665309 322.78 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:05<00:00, 5.90s/it] (M, N, K) aten_addmm-best_config aten_addmm-gbps aten_addmm-tflops triton_addmm-best_config triton_addmm-gbps triton_addmm-tflops pt2_triton_matmul-best_config pt2_triton_matmul-gbps pt2_triton_matmul-tflops ------------------ ------------------------ ----------------- ------------------- ------------------------------------------------------------------------------------------------------------- ------------------- --------------------- ------------------------------- ------------------------ -------------------------- (20120, 512, 1536) 818.112 247.544 {'BLOCK_M': 128, 'BLOCK_N': 256, 'BLOCK_K': 64, 'GROUP_M': 8, 'num_warps': 8, 'num_ctas': 1, 'num_stages': 3} 911.569 275.823 889.125 269.031 ``` Pull Request resolved: #2490 Reviewed By: xuzhao9 Differential Revision: D63862548 Pulled By: FindHao fbshipit-source-id: 9d4afa6051d4191bc2e3288f59e2820627647b91
1 parent 12820bc commit a1f4b2e

File tree

1 file changed

+13
-2
lines changed

1 file changed

+13
-2
lines changed

userbenchmark/triton/run.py

Lines changed: 13 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,12 @@
2929

3030
def get_parser(args=None):
3131
parser = argparse.ArgumentParser(allow_abbrev=False)
32-
parser.add_argument("--op", type=str, required=False, help="Operator to benchmark.")
32+
parser.add_argument(
33+
"--op",
34+
type=str,
35+
required=False,
36+
help="Operators to benchmark. Split with comma if multiple.",
37+
)
3338
parser.add_argument(
3439
"--mode",
3540
choices=["fwd", "bwd", "fwd_bwd", "fwd_no_grad"],
@@ -188,5 +193,11 @@ def run(args: List[str] = []):
188193
run_ci()
189194
return
190195

196+
if args.op:
197+
ops = args.op.split(",")
198+
else:
199+
ops = []
191200
with gpu_lockdown(args.gpu_lockdown):
192-
_run(args, extra_args)
201+
for op in ops:
202+
args.op = op
203+
_run(args, extra_args)

0 commit comments

Comments
 (0)