Split benchmark for special operators (#2661)

Author: tengqm

benchmark/test_assert_async.py

@@ -0,0 +1,55 @@
+import pytest
+import torch
+
+import flag_gems
+
+from . import performance_utils as utils
+
+
+def _input_fn(shape, dtype, device):
+    if dtype == torch.bool:
+        tensor = torch.ones(shape, dtype=dtype, device=device)
+    else:
+        tensor = torch.ones(shape, dtype=dtype, device=device)
+
+    msg = "Benchmark assert_async"
+
+    yield (
+        tensor,
+        msg,
+    )
+
+
+class AssertAsyncBenchmark(utils.GenericBenchmark):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def set_shapes(self, shape_file_path=None):
+        self.shapes = [
+            (),
+            (1,),
+            (1, 1),
+            (1, 1, 1),
+        ]
+
+    def set_more_shapes(self):
+        return None
+
+
+@pytest.mark.assert_async
+def test_assert_async():
+    bench = AssertAsyncBenchmark(
+        op_name="assert_async",
+        input_fn=_input_fn,
+        torch_op=torch._assert_async,
+        dtypes=[
+            torch.bool,
+            torch.int32,
+            torch.float32,
+            torch.float16,
+            torch.bfloat16,
+        ],
+    )
+
+    bench.set_gems(flag_gems._assert_async)
+    bench.run()

benchmark/test_conj_physical.py

@@ -0,0 +1,54 @@
+import pytest
+import torch
+
+import flag_gems
+
+from . import attri_util as attr_utils
+from . import performance_utils as utils
+
+
+def _input_fn(shape, dtype, device):
+    if dtype.is_complex:
+        float_dtype = torch.float32 if dtype == torch.complex64 else torch.float64
+        real = torch.randn(shape, dtype=float_dtype, device=device)
+        imag = torch.randn(shape, dtype=float_dtype, device=device)
+        input_tensor = torch.complex(real, imag).to(dtype)
+    elif dtype.is_floating_point:
+        input_tensor = torch.randn(shape, dtype=dtype, device=device)
+    else:
+        input_tensor = torch.randn(shape, device=device).to(dtype)
+    yield (input_tensor,)
+
+
+class Conj_physicalBenchmark(utils.GenericBenchmarkExcluse3D):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def set_shapes(self, shape_file_path=None):
+        conj_physical_shapes = [
+            (256,),
+            (2048, 2048),
+            (128, 512, 256),
+            (32, 64),
+            (512, 1024),
+            (2, 3, 4),
+        ]
+        self.shapes = conj_physical_shapes
+
+    def set_more_shapes(self):
+        return None
+
+
+@pytest.mark.conj_physical
+def test_conj_physical():
+    dtypes = attr_utils.FLOAT_DTYPES + attr_utils.INT_DTYPES + attr_utils.COMPLEX_DTYPES
+
+    bench = Conj_physicalBenchmark(
+        input_fn=_input_fn,
+        op_name="conj_physical",
+        torch_op=torch.conj_physical,
+        dtypes=dtypes,
+    )
+
+    bench.set_gems(flag_gems.conj_physical)
+    bench.run()

benchmark/test_contiguous.py

@@ -0,0 +1,29 @@
+import pytest
+import torch
+
+from . import attri_util as attr_utils
+from . import performance_utils as utils
+
+
+def _input_fn(shape, dtype, device):
+    if dtype in attr_utils.FLOAT_DTYPES:
+        inp = torch.randn(shape, dtype=dtype, device=device)
+    else:
+        inp = torch.randint(
+            low=-10000, high=10000, size=shape, dtype=dtype, device="cpu"
+        ).to(device)
+    inp = inp[::2]
+
+    yield inp,
+
+
+@pytest.mark.contiguous
+def test_contiguous():
+    bench = utils.GenericBenchmark(
+        op_name="contiguous",
+        input_fn=_input_fn,
+        torch_op=torch.Tensor.contiguous,
+        dtypes=attr_utils.FLOAT_DTYPES + attr_utils.INT_DTYPES,
+    )
+
+    bench.run()

benchmark/test_diag.py

@@ -0,0 +1,27 @@
+import random
+
+import pytest
+import torch
+
+from . import attri_util as attr_utils
+from . import performance_utils as utils
+
+
+def _input_fn(shape, dtype, device):
+    input = utils.generate_tensor_input(shape, dtype, device)
+    diagonal = random.randint(-4, 4)
+    yield input, {
+        "diagonal": diagonal,
+    },
+
+
+@pytest.mark.diag
+def test_diag():
+    bench = utils.GenericBenchmarkExcluse3D(
+        op_name="diag",
+        input_fn=_input_fn,
+        torch_op=torch.diag,
+        dtypes=attr_utils.FLOAT_DTYPES + attr_utils.INT_DTYPES + attr_utils.BOOL_DTYPES,
+    )
+
+    bench.run()

benchmark/test_diag_embed.py

@@ -0,0 +1,30 @@
+import pytest
+import torch
+
+from . import attri_util as attr_utils
+from . import performance_utils as utils
+
+
+class EmbeddingBenchmark(utils.GenericBenchmark2DOnly):
+    def set_more_shapes(self):
+        return None
+
+
+def _input_fn(shape, dtype, device):
+    inp = utils.generate_tensor_input(shape, dtype, device)
+    yield {"input": inp},
+
+    if utils.Config.bench_level == utils.BenchLevel.COMPREHENSIVE:
+        yield {"input": inp, "offset": 1, "dim1": 0, "dim2": -1},
+
+
+@pytest.mark.diag_embed
+def test_diag_embed():
+    bench = EmbeddingBenchmark(
+        op_name="diag_embed",
+        input_fn=_input_fn,
+        torch_op=torch.diag_embed,
+        dtypes=attr_utils.FLOAT_DTYPES + attr_utils.INT_DTYPES + attr_utils.BOOL_DTYPES,
+    )
+
+    bench.run()

benchmark/test_diagonal.py

@@ -0,0 +1,26 @@
+import pytest
+import torch
+
+from . import attri_util as attr_utils
+from . import performance_utils as utils
+
+
+def _input_fn(shape, dtype, device):
+    inp = utils.generate_tensor_input(shape, dtype, device)
+    yield inp,
+
+    if utils.Config.bench_level == utils.BenchLevel.COMPREHENSIVE:
+        yield inp, {"offset": 1, "dim1": 0, "dim2": -1},
+
+
+@pytest.mark.diagonal_backward
+def test_diagonal_backward():
+    bench = utils.GenericBenchmarkExcluse1D(
+        op_name="diagonal_backward",
+        input_fn=_input_fn,
+        torch_op=torch.diagonal,
+        dtypes=attr_utils.FLOAT_DTYPES,
+        is_backward=True,
+    )
+
+    bench.run()

benchmark/test_embedding.py

@@ -0,0 +1,70 @@
+import pytest
+import torch
+
+from . import performance_utils as utils
+
+
+class EmbeddingBenchmark(utils.GenericBenchmark2DOnly):
+    def set_more_shapes(self):
+        # TODO: add more shapes
+        return None
+
+
+def embedding_input_fn(shape, dtype, device):
+    num_embeddings, embedding_dim = shape
+    indices = torch.randint(0, num_embeddings, (num_embeddings,), device=device)
+    weight = torch.randn((num_embeddings, embedding_dim), device=device, dtype=dtype)
+    yield {"input": indices, "weight": weight},
+
+    if utils.Config.bench_level == utils.BenchLevel.COMPREHENSIVE:
+        indices_2d = torch.randint(
+            0,
+            num_embeddings,
+            (num_embeddings, num_embeddings),
+            device=device,
+        )
+
+        yield {"input": indices_2d, "weight": weight},
+
+
+def embedding_backward_input_fn(shape, dtype, device):
+    for forward_args in embedding_input_fn(shape, dtype, device):
+        input = forward_args[0]["input"]
+        weight = forward_args[0]["weight"]
+
+        weight.requires_grad_(True)
+        # import pudb; pudb.set_trace()
+        # output = torch.nn.functional.embedding(input, weight)
+        # grad_output = torch.randn_like(output)
+        yield input, weight
+
+
+@pytest.mark.embedding
+def test_embedding():
+    # Note(Zhengzekang): triton do not support bfloat16 atomic add which is used in embedding grad.
+    bench = EmbeddingBenchmark(
+        input_fn=embedding_input_fn,
+        op_name="embedding",
+        torch_op=torch.nn.functional.embedding,
+        dtypes=[
+            torch.float32,
+            torch.float16,
+        ],
+    )
+    bench.run()
+
+
+@pytest.mark.embedding_backward
+def test_embedding_backward():
+    # Note(Zhengzekang): triton do not support bfloat16 atomic add which is used in embedding grad.
+    bench = EmbeddingBenchmark(
+        input_fn=embedding_backward_input_fn,
+        op_name="embedding_backward",
+        torch_op=torch.nn.functional.embedding,
+        dtypes=[
+            torch.float32,
+            torch.float16,
+        ],
+        is_backward=True,
+    )
+    bench.run()

benchmark/test_embedding_dense_backward.py

@@ -0,0 +1,49 @@
+import pytest
+import torch
+
+import flag_gems
+
+from . import attri_util as attr_utils
+from . import performance_utils as utils
+
+
+class EmbeddingDenseBackwardBenchmark(utils.GenericBenchmark):
+    def set_shapes(self, shape_file_path=None):
+        self.shapes = [
+            (32, 2048, 128, 8192),
+            (16, 2048, 256, 16384),
+            (8, 4096, 256, 32768),
+        ]
+
+
+def _input_fn(shape, dtype, device):
+    B, M, D, num_weights = shape
+
+    grad_output = torch.randn((B, M, D), device=device, dtype=dtype)
+    indices = torch.randint(0, num_weights, (B, M), device=device, dtype=torch.long)
+
+    def inject_padding_idx(cur_indices: torch.Tensor, padding_idx: int) -> torch.Tensor:
+        if padding_idx < 0:
+            return cur_indices
+        mask = torch.rand((B, M), device=device) < 0.25
+        return torch.where(mask, torch.full_like(cur_indices, padding_idx), cur_indices)
+
+    test_cases = [(-1, False), (0, True), (5, False)]
+    for padding_idx, scale_grad_by_freq in test_cases:
+        cur_indices = inject_padding_idx(indices, padding_idx)
+        yield grad_output, cur_indices, num_weights, padding_idx, scale_grad_by_freq
+
+
+@pytest.mark.skipif(
+    (not torch.cuda.is_available()) or (flag_gems.device != "cuda"),
+    reason="CUDA backend is not available for this benchmark.",
+)
+@pytest.mark.embedding_dense_backward
+def test_embedding_dense_backward():
+    bench = EmbeddingDenseBackwardBenchmark(
+        input_fn=_input_fn,
+        op_name="embedding_dense_backward",
+        torch_op=torch.ops.aten.embedding_dense_backward,
+        dtypes=attr_utils.FLOAT_DTYPES,
+    )
+    bench.run()

benchmark/test_functional_sym_constrain_range_for_size.py

@@ -0,0 +1,21 @@
+import pytest
+import torch
+
+from . import attri_util as attr_utils
+from . import performance_utils as utils
+
+
+def _input_fn(shape, cur_dtype, device):
+    dep_token = utils.generate_tensor_input(shape, cur_dtype, device)
+    yield 5, 1, 10, dep_token
+
+
+@pytest.mark.functional_sym_constrain_range_for_size
+def test_functional_sym_constrain_range_for_size():
+    bench = utils.GenericBenchmark(
+        op_name="functional_sym_constrain_range_for_size",
+        torch_op=torch.ops.aten._functional_sym_constrain_range_for_size,
+        dtypes=attr_utils.FLOAT_DTYPES,
+        input_fn=_input_fn,
+    )
+    bench.run()