Merge branch 'master' into auto-gen/index_copy_

Author: tengqm

.github/workflows/command.yaml

@@ -22,6 +22,9 @@ jobs:
         with:
           command: "/test"
           allowed_contexts: pull_request
+          allow_forks: "true"
+          skip_reviews: "true"
+          fork_review_bypass: "true"
           permissions: "write,admin"
 
       - id: parse-command
@@ -32,7 +35,7 @@ jobs:
           op=$(echo $params | cut -d ':' -f 1)
           runner=$(echo $params | cut -d ':' -f 2)
           echo "op=${op}" >> $GITHUB_OUTPUT
-          echo "runner=${op}" >> $GITHUB_OUTPUT
+          echo "runner=${runner}" >> $GITHUB_OUTPUT
 
   test-operator:
     needs: process-command

benchmark/core_shapes.yaml

@@ -197,6 +197,37 @@ GenericBenchmarkExcluse1D:
     - [64, 512, 512]
     - [1024, 1024, 1024]
 
+tril:
+  shapes:
+    - [64, 64]
+    - [1024, 1024]
+    - [4096, 4096]
+    - [64, 512, 512]
+    - [1024, 1024, 1024]
+    - [1024, 16, 16]
+    - [512, 32, 32]
+    - [128, 64, 64]
+  shape_desc: "(B), M, N"
+
+tril_extreme_diagonal:
+  shapes:
+    - [4096, 4096]
+  shape_desc: "M, N"
+
+tril_out_transposed:
+  shapes:
+    - [1024, 1024]
+    - [16, 128, 128]
+    - [512, 32, 32]
+  shape_desc: "(B), M, N"
+
+tril_out_sliced:
+  shapes:
+    - [1024, 1024]
+    - [16, 128, 128]
+    - [512, 32, 32]
+  shape_desc: "(B), M, N"
+
 GenericBenchmarkExcluse3D:
   shapes:
     - [1048576] # 1024 * 1024

benchmark/test_concatenate.py

@@ -0,0 +1,22 @@
+import pytest
+import torch
+
+from . import base, consts, utils
+
+
+def _input_fn(shape, dtype, device):
+    inp1 = utils.generate_tensor_input(shape, dtype, device)
+    inp2 = utils.generate_tensor_input(shape, dtype, device)
+    inp3 = utils.generate_tensor_input(shape, dtype, device)
+    yield [inp1, inp2, inp3], {"dim": 0}
+
+
+@pytest.mark.concatenate
+def test_concatenate():
+    bench = base.GenericBenchmark(
+        op_name="concatenate",
+        torch_op=torch.concatenate,
+        input_fn=_input_fn,
+        dtypes=consts.FLOAT_DTYPES + consts.INT_DTYPES,
+    )
+    bench.run()

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()

benchmark/test_feature_dropout.py

@@ -0,0 +1,31 @@
+import pytest
+import torch
+
+from . import base, consts, utils
+
+
+def _input_fn(shape, dtype, device):
+    inp = utils.generate_tensor_input(shape, dtype, device)
+    yield inp, 0.5, True
+
+
+@pytest.mark.feature_dropout
+def test_feature_dropout():
+    bench = base.GenericBenchmarkExcluse1D(
+        input_fn=_input_fn,
+        op_name="feature_dropout",
+        torch_op=torch.feature_dropout,
+        dtypes=consts.FLOAT_DTYPES,
+    )
+    bench.run()
+
+
+@pytest.mark.feature_dropout_
+def test_feature_dropout_():
+    bench = base.GenericBenchmarkExcluse1D(
+        input_fn=_input_fn,
+        op_name="feature_dropout_",
+        torch_op=torch.feature_dropout_,
+        dtypes=consts.FLOAT_DTYPES,
+    )
+    bench.run()