split special perf

FlagOpen · Aug 1, 2024 · 273efc8 · 273efc8
1 parent 145ed76
commit 273efc8
Show file tree

Hide file tree

Showing 3 changed files with 190 additions and 82 deletions.
diff --git a/benchmark/test_distribution_perf.py 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
@@ -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
@@ -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()