From 273efc8958cc0c0b6bd033ebcf79b5e594ccfaad Mon Sep 17 00:00:00 2001 From: Bowen12992 Date: Thu, 1 Aug 2024 16:50:57 +0800 Subject: [PATCH] split special perf --- benchmark/test_distribution_perf.py | 84 +++++++++++++++++ benchmark/test_special_perf.py | 83 +---------------- benchmark/test_tensor_constructor_perf.py | 105 ++++++++++++++++++++++ 3 files changed, 190 insertions(+), 82 deletions(-) create mode 100644 benchmark/test_distribution_perf.py create mode 100644 benchmark/test_tensor_constructor_perf.py diff --git a/benchmark/test_distribution_perf.py b/benchmark/test_distribution_perf.py new file mode 100644 index 00000000..19f04bfc --- /dev/null +++ b/benchmark/test_distribution_perf.py @@ -0,0 +1,84 @@ +import torch + +from .performance_utils import ( + FLOAT_DTYPES, + POINTWISE_BATCH, + SIZES, + Benchmark, + unary_arg, +) + + +def test_perf_rand(): + def rand_kwargs(dtype, batch, size): + return {"size": (batch, size), "dtype": dtype, "device": "cuda"} + + bench = Benchmark( + op_name="rand", + torch_op=torch.rand, + arg_func=None, + dtypes=FLOAT_DTYPES, + batch=POINTWISE_BATCH, + sizes=SIZES, + kwargs_func=rand_kwargs, + ) + bench.run() + + +def test_perf_randn(): + def randn_kwargs(dtype, batch, size): + return {"size": (batch, size), "dtype": dtype, "device": "cuda"} + + bench = Benchmark( + op_name="randn", + torch_op=torch.randn, + arg_func=None, + dtypes=FLOAT_DTYPES, + batch=POINTWISE_BATCH, + sizes=SIZES, + kwargs_func=randn_kwargs, + ) + bench.run() + + +def test_perf_rand_like(): + bench = Benchmark( + op_name="rand_like", + torch_op=torch.rand_like, + arg_func=unary_arg, + dtypes=FLOAT_DTYPES, + batch=POINTWISE_BATCH, + sizes=SIZES, + ) + bench.run() + + +def test_perf_normal(): + def normal_arg(dtype, batch, size): + loc = torch.full(size=(size, batch), fill_value=3.0, dtype=dtype, device="cuda") + scale = torch.full( + size=(size, batch), fill_value=10.0, dtype=dtype, device="cuda" + ) + return loc, scale + + bench = Benchmark( + op_name="distributions.normal.Normal", + torch_op=torch.distributions.normal.Normal, + arg_func=normal_arg, + dtypes=FLOAT_DTYPES, + batch=POINTWISE_BATCH, + sizes=SIZES, + ) + bench.run() + + +def test_perf_uniform(): + bench = Benchmark( + op_name="uniform_", + torch_op=torch.Tensor.uniform_, + arg_func=unary_arg, + dtypes=FLOAT_DTYPES, + batch=POINTWISE_BATCH, + sizes=SIZES, + ) + bench.run() diff --git a/benchmark/test_special_perf.py b/benchmark/test_special_perf.py index 7d156e8a..73b760aa 100644 --- a/benchmark/test_special_perf.py +++ b/benchmark/test_special_perf.py @@ -1,56 +1,6 @@ import torch -from .performance_utils import ( - FLOAT_DTYPES, - POINTWISE_BATCH, - SIZES, - Benchmark, - unary_arg, -) - - -def test_perf_rand(): - def rand_kwargs(dtype, batch, size): - return {"size": (batch, size), "dtype": dtype, "device": "cuda"} - - bench = Benchmark( - op_name="rand", - torch_op=torch.rand, - arg_func=None, - dtypes=FLOAT_DTYPES, - batch=POINTWISE_BATCH, - sizes=SIZES, - kwargs_func=rand_kwargs, - ) - bench.run() - - -def test_perf_randn(): - def randn_kwargs(dtype, batch, size): - return {"size": (batch, size), "dtype": dtype, "device": "cuda"} - - bench = Benchmark( - op_name="randn", - torch_op=torch.randn, - arg_func=None, - dtypes=FLOAT_DTYPES, - batch=POINTWISE_BATCH, - sizes=SIZES, - kwargs_func=randn_kwargs, - ) - bench.run() - - -def test_perf_rand_like(): - bench = Benchmark( - op_name="rand_like", - torch_op=torch.rand_like, - arg_func=unary_arg, - dtypes=FLOAT_DTYPES, - batch=POINTWISE_BATCH, - sizes=SIZES, - ) - bench.run() +from .performance_utils import POINTWISE_BATCH, SIZES, Benchmark def test_perf_embedding(): @@ -74,37 +24,6 @@ def embedding_kwargs(dtype, batch, size): bench.run() -def test_perf_normal(): - def normal_arg(dtype, batch, size): - loc = torch.full(size=(size, batch), fill_value=3.0, dtype=dtype, device="cuda") - scale = torch.full( - size=(size, batch), fill_value=10.0, dtype=dtype, device="cuda" - ) - return loc, scale - - bench = Benchmark( - op_name="distributions.normal.Normal", - torch_op=torch.distributions.normal.Normal, - arg_func=normal_arg, - dtypes=FLOAT_DTYPES, - batch=POINTWISE_BATCH, - sizes=SIZES, - ) - bench.run() - - -def test_perf_uniform(): - bench = Benchmark( - op_name="uniform_", - torch_op=torch.Tensor.uniform_, - arg_func=unary_arg, - dtypes=FLOAT_DTYPES, - batch=POINTWISE_BATCH, - sizes=SIZES, - ) - bench.run() - - def test_perf_resolve_neg(): def resolve_neg_arg(dtype, batch, size): x = torch.randn(size=(batch, size), dtype=dtype, device="cuda") diff --git a/benchmark/test_tensor_constructor_perf.py b/benchmark/test_tensor_constructor_perf.py new file mode 100644 index 00000000..e033c154 --- /dev/null +++ b/benchmark/test_tensor_constructor_perf.py @@ -0,0 +1,105 @@ +import torch + +from .performance_utils import ( + FLOAT_DTYPES, + POINTWISE_BATCH, + SIZES, + Benchmark, + unary_arg, +) + + +def test_perf_ones(): + def ones_kwargs(dtype, batch, size): + return {"size": (batch, size), "dtype": dtype, "device": "cuda"} + + bench = Benchmark( + op_name="ones", + torch_op=torch.ones, + arg_func=None, + dtypes=FLOAT_DTYPES, + batch=POINTWISE_BATCH, + sizes=SIZES, + kwargs_func=ones_kwargs, + ) + bench.run() + + +def test_perf_zeros(): + def zeros_kwargs(dtype, batch, size): + return {"size": (batch, size), "dtype": dtype, "device": "cuda"} + + bench = Benchmark( + op_name="zeros", + torch_op=torch.zeros, + arg_func=None, + dtypes=FLOAT_DTYPES, + batch=POINTWISE_BATCH, + sizes=SIZES, + kwargs_func=zeros_kwargs, + ) + bench.run() + + +def test_perf_full(): + def full_kwargs(dtype, batch, size): + return { + "size": (batch, size), + "fill_value": 3.1415926, + "dtype": dtype, + "device": "cuda", + } + + bench = Benchmark( + op_name="full", + torch_op=torch.full, + arg_func=None, + dtypes=FLOAT_DTYPES, + batch=POINTWISE_BATCH, + sizes=SIZES, + kwargs_func=full_kwargs, + ) + bench.run() + + +def test_perf_ones_like(): + bench = Benchmark( + op_name="ones_like", + torch_op=torch.ones_like, + arg_func=unary_arg, + dtypes=FLOAT_DTYPES, + batch=POINTWISE_BATCH, + sizes=SIZES, + ) + bench.run() + + +def test_perf_zeros_like(): + bench = Benchmark( + op_name="zeros_like", + torch_op=torch.zeros_like, + arg_func=unary_arg, + dtypes=FLOAT_DTYPES, + batch=POINTWISE_BATCH, + sizes=SIZES, + ) + bench.run() + + +def test_perf_full_like(): + def full_kwargs(dtype, batch, size): + return { + "input": torch.randn([batch, size], dtype=dtype, device="cuda"), + "fill_value": 3.1415926, + } + + bench = Benchmark( + op_name="full_like", + torch_op=torch.full_like, + arg_func=None, + dtypes=FLOAT_DTYPES, + batch=POINTWISE_BATCH, + sizes=SIZES, + kwargs_func=full_kwargs, + ) + bench.run()