Continue renaming benchmark modules (flagos-ai#2719)

Author: tengqm

benchmark/test_safe_softmax.py

@@ -3,14 +3,13 @@
 import pytest
 import torch
 
-from . import attri_util as attr_utils
-from . import performance_utils as utils
+from . import base, consts, utils
 
 
-class SafeSoftmaxBenchmark(utils.Benchmark):
-    def get_input_iter(self, cur_dtype) -> Generator:
+class SafeSoftmaxBenchmark(base.Benchmark):
+    def get_input_iter(self, dtype) -> Generator:
         for shape in self.shapes:
-            inp = utils.generate_tensor_input(shape, cur_dtype, self.device)
+            inp = utils.generate_tensor_input(shape, dtype, self.device)
             yield inp, -1, None
 
 
@@ -19,7 +18,7 @@ def test_safe_softmax():
     bench = SafeSoftmaxBenchmark(
         op_name="_safe_softmax",
         torch_op=torch.ops.aten._safe_softmax,
-        dtypes=attr_utils.FLOAT_DTYPES,
+        dtypes=consts.FLOAT_DTYPES,
     )
 
     bench.run()

benchmark/test_scaled_dot_product_attention.py

@@ -3,24 +3,20 @@
 
 import flag_gems
 
-from .performance_utils import GenericBenchmark
+from . import base
 
 
-class AttentionBenchmark(GenericBenchmark):
-    """
-    benchmark for attention
-    """
-
+class AttentionBenchmark(base.GenericBenchmark):
     def set_more_shapes(self):
         # self.shapes is a list of tuples, each containing three elements:
         # (batch, num_heads, seq_len, head_size).
-        return None
+        return []
 
 
 @pytest.mark.scaled_dot_product_attention
 @pytest.mark.parametrize("dropout_p", [0.0])
 @pytest.mark.parametrize("is_causal", [True, False])
-def test_perf_scaled_dot_product_attention(monkeypatch, dropout_p, is_causal):
+def test_scaled_dot_product_attention(monkeypatch, dropout_p, is_causal):
     if flag_gems.vendor_name == "hygon":
         monkeypatch.setenv("TRITON_HIP_USE_NEW_STREAM_PIPELINE", "0")
 

benchmark/test_scaled_softmax.py

@@ -5,8 +5,7 @@
 
 import flag_gems
 
-from . import performance_utils as base
-from . import utils
+from . import base, utils
 
 try:
     from transformer_engine.pytorch import cpp_extensions as tex

benchmark/test_scatter_add.py

@@ -4,11 +4,10 @@
 import flag_gems
 from flag_gems.utils import shape_utils
 
-from . import attri_util as attr_utils
-from . import performance_utils as utils
+from . import base, consts
 
 
-class TensorSelectBenchmark(utils.GenericBenchmark2DOnly):
+class TensorSelectBenchmark(base.GenericBenchmark2DOnly):
     def set_more_metrics(self):
         return ["gbps"]
 
@@ -62,6 +61,6 @@ def scatter_input_fn(shape, dtype, device):
         torch_op=torch.Tensor.scatter_add_,
         input_fn=scatter_input_fn,
         get_gbps=_get_gbps,
-        dtypes=attr_utils.FLOAT_DTYPES,
+        dtypes=consts.FLOAT_DTYPES,
     )
     bench.run()

benchmark/test_scatter_reduce.py

@@ -4,10 +4,10 @@
 import flag_gems
 from flag_gems.utils import shape_utils
 
-from . import performance_utils as utils
+from . import base
 
 
-class TensorSelectBenchmark(utils.GenericBenchmark2DOnly):
+class TensorSelectBenchmark(base.GenericBenchmark2DOnly):
     def set_more_metrics(self):
         return ["gbps"]
 

benchmark/test_scatter_src.py

@@ -4,11 +4,10 @@
 import flag_gems
 from flag_gems.utils import shape_utils
 
-from . import attri_util as attr_utils
-from . import performance_utils as utils
+from . import base, consts
 
 
-class TensorSelectBenchmark(utils.GenericBenchmark2DOnly):
+class TensorSelectBenchmark(base.GenericBenchmark2DOnly):
     def set_more_metrics(self):
         return ["gbps"]
 
@@ -85,7 +84,7 @@ def test_scatter_src():
         torch_op=torch.scatter,
         input_fn=scatter_input_fn_factory(),
         get_gbps=gather_scatter_gbps,
-        dtypes=attr_utils.FLOAT_DTYPES,
+        dtypes=consts.FLOAT_DTYPES,
     )
     bench.run()
 
@@ -97,7 +96,7 @@ def test_scatter_src_inplace():
         torch_op=torch.Tensor.scatter_,
         input_fn=scatter_inplace_input_fn_factory(),
         get_gbps=gather_scatter_gbps,
-        dtypes=attr_utils.FLOAT_DTYPES,
+        dtypes=consts.FLOAT_DTYPES,
         is_inplace=True,
     )
 

benchmark/test_select_scatter.py

@@ -6,11 +6,10 @@
 import flag_gems
 from flag_gems.utils import shape_utils
 
-from . import attri_util as attr_utils
-from . import performance_utils as utils
+from . import base, consts
 
 
-class TensorSelectBenchmark(utils.GenericBenchmark2DOnly):
+class TensorSelectBenchmark(base.GenericBenchmark2DOnly):
     def set_more_metrics(self):
         return ["gbps"]
 
@@ -55,7 +54,7 @@ def test_select_scatter():
         op_name="select_scatter",
         torch_op=torch.select_scatter,
         input_fn=_input_fn,
-        dtypes=attr_utils.FLOAT_DTYPES,
+        dtypes=consts.FLOAT_DTYPES,
         get_gbps=_get_gbps,
     )
 

benchmark/test_selu.py

@@ -1,14 +1,13 @@
 import pytest
 import torch
 
-from . import attri_util as attrs
-from . import performance_utils as base
+from . import base, consts
 
 
 @pytest.mark.selu
 def test_selu():
     bench = base.UnaryPointwiseBenchmark(
-        op_name="selu", torch_op=torch.nn.functional.selu, dtypes=attrs.FLOAT_DTYPES
+        op_name="selu", torch_op=torch.nn.functional.selu, dtypes=consts.FLOAT_DTYPES
     )
     bench.run()
 
@@ -18,7 +17,7 @@ def test_selu_inplace():
     bench = base.UnaryPointwiseBenchmark(
         op_name="selu_",
         torch_op=torch.ops.aten.selu_,
-        dtypes=attrs.FLOAT_DTYPES,
+        dtypes=consts.FLOAT_DTYPES,
         is_inplace=True,
     )
     bench.run()

benchmark/test_sgn.py

@@ -1,15 +1,14 @@
 import pytest
 
-from . import attri_util as attrs
-from . import performance_utils as base
+from . import base, consts
 
 
 @pytest.mark.sgn_
 def test_sgn_inplace():
     bench = base.UnaryPointwiseBenchmark(
         op_name="atan_",
         torch_op=lambda a: a.sgn_(),
-        dtypes=attrs.FLOAT_DTYPES,
+        dtypes=consts.FLOAT_DTYPES,
         is_inplace=True,
     )
     bench.run()

benchmark/test_sigmoid.py

@@ -1,14 +1,13 @@
 import pytest
 import torch
 
-from . import attri_util as attrs
-from . import performance_utils as base
+from . import base, consts
 
 
 @pytest.mark.sigmoid
 def test_sigmoid():
     bench = base.UnaryPointwiseBenchmark(
-        op_name="sigmoid", torch_op=torch.sigmoid, dtypes=attrs.FLOAT_DTYPES
+        op_name="sigmoid", torch_op=torch.sigmoid, dtypes=consts.FLOAT_DTYPES
     )
     bench.run()
 
@@ -18,7 +17,7 @@ def test_sigmoid_inplace():
     bench = base.UnaryPointwiseBenchmark(
         op_name="sigmoid_",
         torch_op=torch.sigmoid_,
-        dtypes=attrs.FLOAT_DTYPES,
+        dtypes=consts.FLOAT_DTYPES,
         is_inplace=True,
     )
     bench.run()