【KernelGen】Add einsum operator (flagos-ai#1731)

* feat: add einsum operator with tests and benchmark

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

* Apply suggestions from code review

Co-authored-by: Qiming Teng <tengqm@outlook.com>
Signed-off-by: Zang Peiyu <166481866+factnn@users.noreply.github.com>

* fix: use consts.FLOAT_DTYPES and base.Benchmark in einsum benchmark

* fix: fix einsum benchmark - use correct base classes and shapes

---------

Signed-off-by: Zang Peiyu <166481866+factnn@users.noreply.github.com>
Co-authored-by: factnn <1050552884@qq.com>
Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
Co-authored-by: Zang Peiyu <166481866+factnn@users.noreply.github.com>
Co-authored-by: Qiming Teng <tengqm@outlook.com>

Author: 4 people

benchmark/test_einsum.py

@@ -0,0 +1,192 @@
+from typing import Generator
+
+import pytest
+import torch
+
+from . import base, consts
+
+
+class EinsumBenchmark(base.Benchmark):
+    DEFAULT_METRICS = consts.DEFAULT_METRICS[:] + ["tflops"]
+    DEFAULT_SHAPES = [(1, 512, 512, 512), (1, 1024, 1024, 1024), (16, 512, 512, 512)]
+
+    def __init__(self, *args, batched=False, input_fn=None, **kwargs):
+        self.batched = batched
+        super().__init__(*args, **kwargs)
+
+    def set_more_shapes(self):
+        return []
+
+    def set_shapes(self, *args, **kwargs):
+        self.shapes = self.DEFAULT_SHAPES
+
+    def get_input_iter(self, dtype) -> Generator:
+        for b, m, n, k in self.shapes:
+            if self.batched:
+                inp1 = torch.randn([b, m, k], dtype=dtype, device=self.device)
+                inp2 = torch.randn([b, k, n], dtype=dtype, device=self.device)
+            else:
+                inp1 = torch.randn([m, k], dtype=dtype, device=self.device)
+                inp2 = torch.randn([k, n], dtype=dtype, device=self.device)
+            yield inp1, inp2
+
+    def get_tflops(self, op, *args, **kwargs):
+        A, B = args[0], args[1]
+        if self.batched:
+            return A.shape[0] * A.shape[1] * B.shape[2] * A.shape[2] * 2
+        return A.shape[0] * B.shape[1] * A.shape[1] * 2
+
+
+class EinsumGenericBenchmark(base.GenericBenchmark):
+    def set_shapes(self, *args, **kwargs):
+        pass  # keep shapes set by caller
+
+
+def dot_input_fn(shape, dtype, device):
+    (n,) = shape
+    yield torch.randn(n, dtype=dtype, device=device), torch.randn(
+        n, dtype=dtype, device=device
+    )
+
+
+def outer_input_fn(shape, dtype, device):
+    m, n = shape
+    yield torch.randn(m, dtype=dtype, device=device), torch.randn(
+        n, dtype=dtype, device=device
+    )
+
+
+def unary_2d_input_fn(shape, dtype, device):
+    m, n = shape
+    yield (torch.randn(m, n, dtype=dtype, device=device),)
+
+
+def unary_3d_input_fn(shape, dtype, device):
+    m, n, k = shape
+    yield (torch.randn(m, n, k, dtype=dtype, device=device),)
+
+
+def ellipsis_input_fn(shape, dtype, device):
+    b, h, m, k, n = shape
+    yield torch.randn(b, h, m, k, dtype=dtype, device=device), torch.randn(
+        b, h, k, n, dtype=dtype, device=device
+    )
+
+
+@pytest.mark.einsum
+def test_einsum_matmul():
+    bench = EinsumBenchmark(
+        input_fn=None,
+        op_name="einsum",
+        torch_op=lambda A, B: torch.einsum("ij,jk->ik", A, B),
+        dtypes=consts.FLOAT_DTYPES,
+    )
+    bench.run()
+
+
+@pytest.mark.einsum
+def test_einsum_bmm():
+    bench = EinsumBenchmark(
+        input_fn=None,
+        op_name="einsum",
+        torch_op=lambda A, B: torch.einsum("bij,bjk->bik", A, B),
+        dtypes=consts.FLOAT_DTYPES,
+        batched=True,
+    )
+    bench.run()
+
+
+@pytest.mark.einsum
+def test_einsum_dot():
+    bench = EinsumGenericBenchmark(
+        input_fn=dot_input_fn,
+        op_name="einsum",
+        torch_op=lambda A, B: torch.einsum("i,i->", A, B),
+        dtypes=consts.FLOAT_DTYPES,
+    )
+    bench.shapes = [(1024,), (4096,), (65536,)]
+    bench.run()
+
+
+@pytest.mark.einsum
+def test_einsum_outer():
+    bench = EinsumGenericBenchmark(
+        input_fn=outer_input_fn,
+        op_name="einsum",
+        torch_op=lambda A, B: torch.einsum("i,j->ij", A, B),
+        dtypes=consts.FLOAT_DTYPES,
+    )
+    bench.shapes = [(1024, 1024), (4096, 4096)]
+    bench.run()
+
+
+@pytest.mark.einsum
+def test_einsum_trace():
+    bench = EinsumGenericBenchmark(
+        input_fn=unary_2d_input_fn,
+        op_name="einsum",
+        torch_op=lambda A: torch.einsum("ii->", A),
+        dtypes=consts.FLOAT_DTYPES,
+    )
+    bench.shapes = [(1024, 1024), (4096, 4096)]
+    bench.run()
+
+
+@pytest.mark.einsum
+def test_einsum_diagonal():
+    bench = EinsumGenericBenchmark(
+        input_fn=unary_2d_input_fn,
+        op_name="einsum",
+        torch_op=lambda A: torch.einsum("ii->i", A),
+        dtypes=consts.FLOAT_DTYPES,
+    )
+    bench.shapes = [(1024, 1024), (4096, 4096)]
+    bench.run()
+
+
+@pytest.mark.einsum
+def test_einsum_transpose():
+    bench = EinsumGenericBenchmark(
+        input_fn=unary_2d_input_fn,
+        op_name="einsum",
+        torch_op=lambda A: torch.einsum("ij->ji", A),
+        dtypes=consts.FLOAT_DTYPES,
+    )
+    bench.shapes = [(1024, 1024), (4096, 4096)]
+    bench.run()
+
+
+@pytest.mark.einsum
+def test_einsum_sum_all():
+    bench = EinsumGenericBenchmark(
+        input_fn=unary_3d_input_fn,
+        op_name="einsum",
+        torch_op=lambda A: torch.einsum("ijk->", A),
+        dtypes=consts.FLOAT_DTYPES,
+    )
+    bench.shapes = [(64, 64, 64), (128, 128, 128)]
+    bench.run()
+
+
+@pytest.mark.einsum
+def test_einsum_sum_dim():
+    bench = EinsumGenericBenchmark(
+        input_fn=unary_3d_input_fn,
+        op_name="einsum",
+        torch_op=lambda A: torch.einsum("ijk->j", A),
+        dtypes=consts.FLOAT_DTYPES,
+    )
+    bench.shapes = [(64, 64, 64), (128, 128, 128)]
+    bench.run()
+
+
+@pytest.mark.einsum
+def test_einsum_ellipsis():
+    bench = EinsumGenericBenchmark(
+        input_fn=ellipsis_input_fn,
+        op_name="einsum",
+        torch_op=lambda A, B: torch.einsum("...ij,...jk->...ik", A, B),
+        dtypes=consts.FLOAT_DTYPES,
+    )
+    bench.shapes = [(2, 4, 64, 64, 128), (2, 8, 128, 128, 256)]
+    bench.run()

src/flag_gems/__init__.py

@@ -193,6 +193,7 @@ def torch_ge(v):
     ("divide_.Tensor", true_divide_),
     ("divide_.Tensor_mode", div_mode_),
     ("dot", dot),
+    ("einsum", einsum),
     ("elu", elu),
     ("elu_", elu_),
     ("elu_backward", elu_backward),

src/flag_gems/ops/__init__.py

@@ -108,6 +108,7 @@
 )
 from flag_gems.ops.dot import dot
 from flag_gems.ops.dropout import dropout, dropout_backward
+from flag_gems.ops.einsum import einsum
 from flag_gems.ops.elu import elu, elu_, elu_backward
 from flag_gems.ops.embedding import embedding, embedding_backward
 from flag_gems.ops.embedding_dense_backward import embedding_dense_backward
@@ -483,6 +484,7 @@
     "dot",
     "dropout",
     "dropout_backward",
+    "einsum",
     "elu",
     "elu_",
     "elu_backward",