Fix nits in benchmark files (#2786)

Author: tengqm

benchmark/test_act_quant.py

@@ -6,7 +6,7 @@
 import flag_gems
 from flag_gems.utils.device_info import get_device_capability
 
-from .performance_utils import GenericBenchmark
+from . import base
 
 M = [1, 40, 164, 512, 3454, 12027, 38594]
 N = [128, 896, 2048, 8192]
@@ -74,11 +74,11 @@ def torch_act_quant(
     return y, s
 
 
-class ActQuantBenchmark(GenericBenchmark):
+class ActQuantBenchmark(base.GenericBenchmark):
     # Only 2D shapes make sense for act_quant
     def set_more_shapes(self):
         self.shapes = SHAPES
-        return None
+        return []
 
 
 @pytest.mark.act_quant
@@ -97,7 +97,7 @@ def input_fn(shape, dtype, device):
         op_name="act_quant",
         torch_op=torch_act_quant,
         input_fn=input_fn,
-        dtypes=[torch.bfloat16],
         gems_op=flag_gems.act_quant_triton,
+        dtypes=[torch.bfloat16],
     )
     bench.run()

benchmark/test_avg_pool3d_perf.py

@@ -3,12 +3,11 @@
 import pytest
 import torch
 
-from benchmark.attri_util import FLOAT_DTYPES, BenchLevel
-from benchmark.performance_utils import Config, GenericBenchmark, generate_tensor_input
+from . import base, consts, utils
 
 
 def avg_pool3d_input_fn(shape, dtype, device):
-    inp = generate_tensor_input(shape, dtype, device)
+    inp = utils.generate_tensor_input(shape, dtype, device)
     # Common case
     yield inp, {
         "kernel_size": 3,
@@ -18,7 +17,7 @@ def avg_pool3d_input_fn(shape, dtype, device):
         "count_include_pad": True,
         "divisor_override": None,
     }
-    if Config.bench_level == BenchLevel.COMPREHENSIVE:
+    if base.Config.bench_level == consts.BenchLevel.COMPREHENSIVE:
         # With count_include_pad=False
         yield inp, {
             "kernel_size": 3,
@@ -49,8 +48,8 @@ def avg_pool3d_input_fn(shape, dtype, device):
             }
 
 
-class AvgPool3dBenchmark(GenericBenchmark):
-    def get_input_iter(self, cur_dtype) -> Generator:
+class AvgPool3dBenchmark(base.GenericBenchmark):
+    def get_input_iter(self, dtype) -> Generator:
         shapes_5d = [
             (4, 3, 16, 56, 56),
             (8, 64, 8, 28, 28),
@@ -59,7 +58,7 @@ def get_input_iter(self, cur_dtype) -> Generator:
         ]
 
         for shape in shapes_5d:
-            yield from self.input_fn(shape, cur_dtype, self.device)
+            yield from self.input_fn(shape, dtype, self.device)
 
 
 @pytest.mark.avg_pool3d
@@ -68,18 +67,18 @@ def test_perf_avg_pool3d():
         input_fn=avg_pool3d_input_fn,
         op_name="avg_pool3d",
         torch_op=torch.ops.aten.avg_pool3d,
-        dtypes=FLOAT_DTYPES,
+        dtypes=consts.FLOAT_DTYPES,
     )
     bench.run()
 
 
-@pytest.mark.avg_pool3d
+@pytest.mark.avg_pool3d_backward
 def test_perf_avg_pool3d_backward():
     bench = AvgPool3dBenchmark(
         input_fn=avg_pool3d_input_fn,
         op_name="avg_pool3d",
         torch_op=torch.ops.aten.avg_pool3d,
-        dtypes=FLOAT_DTYPES,
+        dtypes=consts.FLOAT_DTYPES,
         is_backward=True,
     )
     bench.run()