Split benchmark for binary pointwise operators (#2691)

Author: tengqm

benchmark/performance_utils.py

@@ -697,6 +697,53 @@ def get_tflops(self, op, *args, **kwargs):
         return total_flops
 
 
+class BinaryPointwiseBenchmark(Benchmark):
+    """
+    Base class for benchmarking binary pointwise operations.
+    """
+
+    DEFAULT_METRICS = DEFAULT_METRICS[:] + ["tflops"]
+
+    def set_more_shapes(self):
+        special_shapes_2d = [[1024, 2**i] for i in range(0, 20, 4)]
+        shapes_3d = [[64, 64, 2**i] for i in range(0, 20, 4)]
+        return special_shapes_2d + shapes_3d
+
+    def get_input_iter(self, dtype) -> Generator:
+        for shape in self.shapes:
+            inp1 = generate_tensor_input(shape, dtype, self.device)
+            inp2 = generate_tensor_input(shape, dtype, self.device)
+            yield inp1, inp2
+
+    def get_tflops(self, op, *args, **kwargs):
+        shape1 = list(args[0].shape)
+        shape2 = list(args[0].shape)
+        return torch.tensor(shape1).prod().item() + torch.tensor(shape2).prod().item()
+
+
+class ScalarBinaryPointwiseBenchmark(Benchmark):
+    """
+    Base class for benchmarking binary pointwise operations with scalar input.
+    """
+
+    DEFAULT_METRICS = DEFAULT_METRICS[:] + ["tflops"]
+
+    def set_more_shapes(self):
+        special_shapes_2d = [[1024, 2**i] for i in range(0, 20, 4)]
+        shapes_3d = [[64, 64, 2**i] for i in range(0, 20, 4)]
+        return special_shapes_2d + shapes_3d
+
+    def get_input_iter(self, cur_dtype) -> Generator:
+        for shape in self.shapes:
+            inp1 = 0.001  # Scalar input
+            inp2 = generate_tensor_input(shape, cur_dtype, self.device)
+            yield inp1, inp2
+
+    def get_tflops(self, op, *args, **kwargs):
+        shape = list(args[1].shape)  # Second argument is the tensor
+        return torch.tensor(shape).prod().item()
+
+
 def generate_tensor_input(shape, dtype, device):
     if dtype in FLOAT_DTYPES:
         return torch.randn(shape, dtype=dtype, device=device)

benchmark/test_add.py

@@ -0,0 +1,26 @@
+import pytest
+import torch
+
+from . import attri_util as attrs
+from . import performance_utils as base
+
+
+@pytest.mark.add
+def test_add():
+    bench = base.BinaryPointwiseBenchmark(
+        op_name="add",
+        torch_op=torch.add,
+        dtypes=attrs.FLOAT_DTYPES + attrs.COMPLEX_DTYPES,
+    )
+    bench.run()
+
+
+@pytest.mark.add_
+def test_add_inplace():
+    bench = base.BinaryPointwiseBenchmark(
+        op_name="add_",
+        torch_op=lambda a, b: a.add_(b),
+        dtypes=attrs.FLOAT_DTYPES,
+        is_inplace=True,
+    )
+    bench.run()

benchmark/test_allclose.py

@@ -0,0 +1,15 @@
+import pytest
+import torch
+
+from . import attri_util as attrs
+from . import performance_utils as base
+
+
+@pytest.mark.allclose
+def test_allclose():
+    bench = base.BinaryPointwiseBenchmark(
+        op_name="allclose",
+        torch_op=torch.allclose,
+        dtypes=attrs.FLOAT_DTYPES + attrs.INT_DTYPES,
+    )
+    bench.run()

benchmark/test_binary_pointwise_perf.py

@@ -0,0 +1,26 @@
+import pytest
+import torch
+
+from . import attri_util as attrs
+from . import performance_utils as base
+
+
+@pytest.mark.bitwise_and
+def test_bitwise_and():
+    bench = base.BinaryPointwiseBenchmark(
+        op_name="bitwise_and",
+        torch_op=torch.bitwise_and,
+        dtypes=attrs.INT_DTYPES + attrs.BOOL_DTYPES,
+    )
+    bench.run()
+
+
+@pytest.mark.bitwise_and_
+def test_bitwise_and_inplace():
+    bench = base.BinaryPointwiseBenchmark(
+        op_name="bitwise_and_",
+        torch_op=lambda a, b: a.bitwise_and_(b),
+        dtypes=attrs.INT_DTYPES + attrs.BOOL_DTYPES,
+        is_inplace=True,
+    )
+    bench.run()

benchmark/test_bitwise_and.py

@@ -0,0 +1,26 @@
+import pytest
+import torch
+
+from . import attri_util as attrs
+from . import performance_utils as base
+
+
+@pytest.mark.bitwise_or
+def test_bitwise_or():
+    bench = base.BinaryPointwiseBenchmark(
+        op_name="bitwise_or",
+        torch_op=torch.bitwise_or,
+        dtypes=attrs.INT_DTYPES + attrs.BOOL_DTYPES,
+    )
+    bench.run()
+
+
+@pytest.mark.bitwise_or_
+def test_bitwise_or_inplace():
+    bench = base.BinaryPointwiseBenchmark(
+        op_name="bitwise_or_",
+        torch_op=lambda a, b: a.bitwise_or_(b),
+        dtypes=attrs.INT_DTYPES + attrs.BOOL_DTYPES,
+        is_inplace=True,
+    )
+    bench.run()

benchmark/test_bitwise_or.py

@@ -0,0 +1,27 @@
+import pytest
+import torch
+
+from . import attri_util as attrs
+from . import performance_utils as base
+
+
+# TODO(0x45f): Fix OOM when dtypes includes COMPLEX_DTYPES is included (Issue #2693).
+@pytest.mark.div
+def test_div():
+    bench = base.BinaryPointwiseBenchmark(
+        op_name="div",
+        torch_op=torch.div,
+        dtypes=attrs.FLOAT_DTYPES,
+    )
+    bench.run()
+
+
+@pytest.mark.div_
+def test_div_inplace():
+    bench = base.BinaryPointwiseBenchmark(
+        op_name="div_",
+        torch_op=lambda a, b: a.div_(b),
+        dtypes=attrs.FLOAT_DTYPES,
+        is_inplace=True,
+    )
+    bench.run()

benchmark/test_div.py

@@ -0,0 +1,15 @@
+import pytest
+
+from . import attri_util as attrs
+from . import performance_utils as base
+
+
+@pytest.mark.dunder_ior
+def test_dunder_ior_inplace():
+    bench = base.BinaryPointwiseBenchmark(
+        op_name="dunder_ior",
+        torch_op=lambda a, b: a.__ior__(b),
+        dtypes=attrs.INT_DTYPES + attrs.BOOL_DTYPES,
+        is_inplace=True,
+    )
+    bench.run()

benchmark/test_dunder_ior.py

@@ -0,0 +1,14 @@
+import pytest
+
+from . import attri_util as attrs
+from . import performance_utils as base
+
+
+@pytest.mark.dunder_or
+def test_dunder_or():
+    bench = base.BinaryPointwiseBenchmark(
+        op_name="dunder_or",
+        torch_op=lambda a, b: a | b,
+        dtypes=attrs.INT_DTYPES + attrs.BOOL_DTYPES,
+    )
+    bench.run()

benchmark/test_dunder_or.py

@@ -0,0 +1,15 @@
+import pytest
+import torch
+
+from . import attri_util as attrs
+from . import performance_utils as base
+
+
+@pytest.mark.eq
+def test_eq():
+    bench = base.BinaryPointwiseBenchmark(
+        op_name="eq",
+        torch_op=torch.eq,
+        dtypes=attrs.FLOAT_DTYPES,
+    )
+    bench.run()