Merge branch 'master' into competition/conv-transpose2d

Author: goldenfox2025

benchmark/performance_utils.py

@@ -25,6 +25,7 @@
     BenchMode,
     OperationAttribute,
     check_metric_dependencies,
+    model_shapes,
 )
 from .conftest import Config, emit_record_logger
 
@@ -492,9 +493,9 @@ def set_more_shapes(self):
         more_shapes_3d = [(100, 2**i, 100) for i in (0, 8, 16)]
         return more_shapes_1d + more_shapes_2d + more_shapes_3d
 
-    def get_input_iter(self, cur_dtype) -> Generator:
+    def get_input_iter(self, dtype) -> Generator:
         for shape in self.shapes:
-            yield from self.input_fn(shape, cur_dtype, self.device)
+            yield from self.input_fn(shape, dtype, self.device)
 
 
 class GenericBenchmarkFilterShapes(GenericBenchmark):
@@ -641,6 +642,61 @@ def get_tflops(self, op, *args, **kwargs):
         return torch.tensor(inp_shape).prod().item() * 2
 
 
+class BlasBenchmark(Benchmark):
+    """
+    benchmark for blas
+    """
+
+    DEFAULT_METRICS = DEFAULT_METRICS[:] + ["tflops"]
+
+    def __init__(self, *args, input_fn, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.input_fn = input_fn
+
+    def get_input_iter(self, dtype) -> Generator:
+        for b, m, n, k in self.shapes:
+            yield from self.input_fn(b, m, n, k, dtype, self.device, False)
+
+        if Config.bench_level == BenchLevel.COMPREHENSIVE:
+            for b, m, n, k in self.shapes:
+                yield from self.input_fn(b, m, n, k, dtype, self.device, True)
+
+    def set_more_shapes(self):
+        large_k_shapes = [
+            [8, 1848, 1536, 151936],
+            [8, 1848, 1536, 128256],
+            [8, 1848, 1536, 152064],
+            [8, 4096, 1, 152064],
+        ]
+
+        model_shaps = model_shapes()
+        return large_k_shapes + model_shaps
+
+    def get_tflops(self, op, *args, **kwargs):
+        total_flops = 0
+        # shape(m,k)(k,n)
+        # total_flops mxnx2k
+        if self.op_name == "mm":
+            total_flops = args[0].shape[0] * args[0].shape[1] * args[1].shape[1] * 2
+
+        # shape(m,n)(n,p)
+        # total_flops mxpx(2n+1)
+        elif self.op_name == "addmm":
+            total_flops = (
+                args[0].shape[0] * args[1].shape[1] * (args[1].shape[0] * 2 + 1)
+            )
+        # total_flops bxnxpx2m
+        elif self.op_name == "bmm":
+            total_flops = (
+                args[0].shape[0]
+                * args[0].shape[1]
+                * args[1].shape[2]
+                * 2
+                * args[0].shape[2]
+            )
+        return total_flops
+
+
 def generate_tensor_input(shape, dtype, device):
     if dtype in FLOAT_DTYPES:
         return torch.randn(shape, dtype=dtype, device=device)

benchmark/test_addmm.py

@@ -0,0 +1,33 @@
+import pytest
+import torch
+
+import flag_gems
+
+from . import attri_util as attrs
+from . import performance_utils as base
+
+
+def _input_fn(b, m, n, k, dtype, device, b_column_major):
+    inp1 = torch.randn([m, k], dtype=dtype, device=device)
+    bias = torch.randn([m, n], dtype=dtype, device=device)
+    if b_column_major:
+        inp2 = torch.randn([n, k], dtype=dtype, device=device)
+        yield bias, inp1, inp2.t(),
+    else:
+        inp2 = torch.randn([k, n], dtype=dtype, device=device)
+        yield bias, inp1, inp2,
+
+
+@pytest.mark.addmm
+def test_addmm(monkeypatch):
+    if flag_gems.vendor_name == "mthreads":
+        monkeypatch.setenv("MUSA_ENABLE_SQMMA", "1")
+
+    bench = base.BlasBenchmark(
+        op_name="addmm",
+        input_fn=_input_fn,
+        torch_op=torch.addmm,
+        dtypes=attrs.FLOAT_DTYPES,
+    )
+
+    bench.run()

benchmark/test_addmv.py

@@ -0,0 +1,36 @@
+from typing import Generator
+
+import pytest
+import torch
+
+from . import performance_utils as base
+from .attri_util import FLOAT_DTYPES
+
+
+class AddmvBenchmark(base.GenericBenchmark2DOnly):
+    def set_more_shapes(self):
+        return []
+
+    def get_input_iter(self, dtype) -> Generator:
+        for m, n in self.shapes:
+            yield from self.input_fn(m, n, dtype, self.device)
+
+
+def _input_fn(m, n, cur_dtype, device):
+    mat = torch.randn([m, n], dtype=cur_dtype, device=device)
+    vec = torch.randn([n], dtype=cur_dtype, device=device)
+    bias = torch.randn([m], dtype=cur_dtype, device=device)
+
+    # torch.addmv(bias, mat, vec)
+    yield bias, mat, vec
+
+
+@pytest.mark.addmv
+def test_addmv():
+    bench = AddmvBenchmark(
+        op_name="addmv",
+        input_fn=_input_fn,
+        torch_op=torch.addmv,
+        dtypes=FLOAT_DTYPES,
+    )
+    bench.run()

benchmark/test_addr.py

@@ -0,0 +1,36 @@
+from typing import Generator
+
+import pytest
+import torch
+
+from benchmark.attri_util import FLOAT_DTYPES
+
+from . import performance_utils as base
+
+
+class AddrBenchmark(base.BlasBenchmark):
+    def set_more_shapes(self):
+        return []
+
+    def get_input_iter(self, dtype) -> Generator:
+        for shape in self.shapes:
+            m, n = shape[0], shape[1]
+            yield from self.input_fn(m, n, dtype, self.device)
+
+
+def _input_fn(m, n, cur_dtype, device):
+    inp1 = torch.randn([m, n], dtype=cur_dtype, device=device)
+    inp2 = torch.randn([m], dtype=cur_dtype, device=device)
+    inp3 = torch.randn([n], dtype=cur_dtype, device=device)
+    yield inp1, inp2, inp3, {"alpha": 0.5, "beta": 0.5}
+
+
+@pytest.mark.addr
+def test_addr():
+    bench = AddrBenchmark(
+        op_name="addr",
+        input_fn=_input_fn,
+        torch_op=torch.Tensor.addr,
+        dtypes=FLOAT_DTYPES,
+    )
+    bench.run()

benchmark/test_baddbmm.py

@@ -0,0 +1,59 @@
+import pytest
+import torch
+
+from . import attri_util as attrs
+from . import performance_utils as base
+
+
+class BaddbmmBenchmark(base.BlasBenchmark):
+    def set_more_shapes(self):
+        model_shapes_list = attrs.model_shapes()
+
+        skip_shapes = [
+            (4, 8192, 128256, 4096),
+            (4, 8192, 152064, 3584),
+        ]
+
+        filtered = []
+        for shape in model_shapes_list:
+            if shape not in skip_shapes:
+                filtered.append(shape)
+
+        return filtered
+
+    def get_tflops(self, op, *args, **kwargs):
+        # shape(b,m,k)(b,k,n)
+        # total_flops = b * m * n * (2 * k + 1)
+        total_flops = (
+            args[1].shape[0]
+            * args[1].shape[1]
+            * args[2].shape[2]
+            * (args[1].shape[2] * 2 + 1)
+        )
+        return total_flops
+
+
+def _input_fn(b, m, n, k, dtype, device, b_column_major):
+    inp1 = torch.randn([b, m, k], dtype=dtype, device=device, requires_grad=True)
+
+    if b_column_major:
+        inp2 = torch.randn([b, n, k], dtype=dtype, device=device, requires_grad=True)
+        inp2 = inp2.transpose(1, 2).contiguous()
+    else:
+        inp2 = torch.randn([b, k, n], dtype=dtype, device=device, requires_grad=True)
+
+    bias = torch.randn([b, m, n], dtype=dtype, device=device, requires_grad=True)
+
+    yield bias, inp1, inp2
+
+
+@pytest.mark.baddbmm
+def test_baddbmm():
+    bench = BaddbmmBenchmark(
+        op_name="baddbmm",
+        input_fn=_input_fn,
+        torch_op=torch.baddbmm,
+        dtypes=attrs.FLOAT_DTYPES,
+    )
+
+    bench.run()