Split benchmarks for distribution operators

benchmark/test_bernoulli.py

@@ -0,0 +1,22 @@
+import pytest
+import torch
+
+from . import attri_util as attr_utils
+from . import performance_utils as utils
+
+
+def input_fn(shape, cur_dtype, device):
+    self = torch.randn(shape, dtype=cur_dtype, device=device)
+    p = 0.5
+    yield self, p
+
+
+@pytest.mark.bernoulli_
+def test_bernoulli_inplace():
+    bench = utils.GenericBenchmark(
+        op_name="bernoulli_",
+        input_fn=input_fn,
+        torch_op=torch.Tensor.bernoulli_,
+        dtypes=attr_utils.FLOAT_DTYPES,
+    )
+    bench.run()

benchmark/test_exponential.py

@@ -0,0 +1,16 @@
+import pytest
+import torch
+
+from . import attri_util as attr_utils
+from . import performance_utils as utils
+
+
+@pytest.mark.exponential_
+def test_exponential_inplace():
+    bench = utils.GenericBenchmark(
+        op_name="exponential_",
+        input_fn=utils.unary_input_fn,
+        torch_op=torch.Tensor.exponential_,
+        dtypes=attr_utils.FLOAT_DTYPES,
+    )
+    bench.run()

benchmark/test_normal.py

@@ -0,0 +1,40 @@
+import pytest
+import torch
+
+from . import attri_util as attr_utils
+from . import performance_utils as utils
+
+
+def normal_input_fn(shape, cur_dtype, device):
+    loc = torch.full(shape, fill_value=3.0, dtype=cur_dtype, device=device)
+    scale = torch.full(shape, fill_value=10.0, dtype=cur_dtype, device=device)
+    yield loc, scale
+
+
+@pytest.mark.normal
+def test_normal():
+    bench = utils.GenericBenchmark(
+        input_fn=normal_input_fn,
+        op_name="normal",
+        torch_op=torch.normal,
+        dtypes=attr_utils.FLOAT_DTYPES,
+    )
+    bench.run()
+
+
+def normal_inplace_input_fn(shape, cur_dtype, device):
+    self = torch.randn(shape, dtype=cur_dtype, device=device)
+    loc = 3.0
+    scale = 10.0
+    yield self, loc, scale
+
+
+@pytest.mark.normal_
+def test_normal_inplace():
+    bench = utils.GenericBenchmark(
+        input_fn=normal_inplace_input_fn,
+        op_name="normal_",
+        torch_op=torch.Tensor.normal_,
+        dtypes=attr_utils.FLOAT_DTYPES,
+    )
+    bench.run()

benchmark/test_uniform.py

@@ -0,0 +1,16 @@
+import pytest
+import torch
+
+from . import attri_util as attr_utils
+from . import performance_utils as utils
+
+
+@pytest.mark.uniform_
+def test_uniform_inplace():
+    bench = utils.GenericBenchmark(
+        input_fn=utils.unary_input_fn,
+        op_name="uniform_",
+        torch_op=torch.Tensor.uniform_,
+        dtypes=attr_utils.FLOAT_DTYPES,
+    )
+    bench.run()

tools/test-op.sh

@@ -90,7 +90,8 @@ done
 # Run benchmark test if necessary
 for item in "${PERF_TEST_CASES[@]}"; do
   echo "Running benchmark tests for ${item}"
-  pytest -s --level core --record log ${item}
+  echo "pytest -s ${item} --level core --record log"
+  pytest -s ${item} --level core --record log
 done
 
 # Process coverage data only when full-range testing