【KernelGen】Add index_copy_ operator

benchmark/test_index_copy_perf.py

@@ -0,0 +1,53 @@
+import pytest
+import torch
+
+from . import base, consts, utils
+
+
+class IndexCopyBenchmark(base.GenericBenchmark2DOnly):
+    def get_gbps(self, bench_fn_args, latency):
+        index = bench_fn_args[2]
+        src = bench_fn_args[3]
+        io_amount = sum([utils.size_in_bytes(item) for item in [index, src, src]])
+        return io_amount * 1e-9 / (latency * 1e-3)
+
+
+def _tensor_input_fn(shape, dtype, device):
+    inp = torch.randn(shape, dtype=dtype, device=device)
+    dim = 0 if len(shape) == 1 else 1
+    src_shape = list(inp.shape)
+    index_max = src_shape[dim]
+    index_len = index_max // 2 if index_max >= 2 else 1
+    index = torch.randperm(index_len, device=device)
+    src_shape[dim] = index_len
+    src = torch.randn(src_shape, dtype=dtype, device=device)
+    yield inp, dim, index, src
+
+
+def _inplace_input_fn(shape, dtype, device):
+    yield from _tensor_input_fn(shape, dtype, device)
+
+
+@pytest.mark.index_copy
+def test_index_copy():
+    bench = base.GenericBenchmark2DOnly(
+        input_fn=_tensor_input_fn,
+        op_name="index_copy",
+        torch_op=torch.index_copy,
+        dtypes=consts.FLOAT_DTYPES,
+        get_gbps=IndexCopyBenchmark.get_gbps,
+    )
+    bench.run()
+
+
+@pytest.mark.index_copy_
+def test_index_copy_():
+    bench = base.GenericBenchmark2DOnly(
+        input_fn=_tensor_input_fn,
+        op_name="index_copy_",
+        torch_op=torch.Tensor.index_copy_,
+        dtypes=consts.FLOAT_DTYPES,
+        get_gbps=IndexCopyBenchmark.get_gbps,
+        inplace=True,
+    )
+    bench.run()

conf/operators.yaml

@@ -2696,6 +2696,30 @@ ops:
       - Tensor
     stages:
       - stable: '4.0'
+  - id: index_copy
+    description: |
+      Copies the elements from `source` into `input` at the positions specified by
+      `index` along the given `dim`.
+    for:
+      - index_copy
+    labels:
+      - aten
+      - KernelGen
+    kind:
+      - Tensor
+    stages:
+      - beta: '5.1'
+  - id: index_copy_
+    description: The in-place version of `index_copy()`.
+    for:
+      - index_copy_
+    labels:
+      - aten
+      - KernelGen
+    kind:
+      - Tensor
+    stages:
+      - beta: '5.1'
   - id: index_put
     description: |
       Puts values from the tensor `values` into the tensor `input` using the indices specified
@@ -4961,6 +4985,30 @@ ops:
       - Activation
     stages:
       - stable: '2.0'
+  - id: silu_and_mul_with_clamp
+    description: A custom operator in vLLM as activation function for SwiGLU.
+    for:
+      - silu_and_mul_with_clamp
+    labels:
+      - fused
+      - pointwise
+      - vLLM
+    kind:
+      - Activation
+    stages:
+      - stable: '5.1'
+  - id: silu_and_mul_with_clamp_out
+    description: A variant of `silu_and_mul_with_clamp` with an extra `out` argument.
+    for:
+      - silu_and_mul_with_clamp.out
+    labels:
+      - fused
+      - pointwise
+      - vLLM
+    kind:
+      - Activation
+    stages:
+      - stable: '5.1'
   - id: silu_backward
     description: A variant of `silu()` for backward case.
     for:

src/flag_gems/__init__.py

@@ -259,6 +259,8 @@ def torch_ge(v):
     ("index.Tensor", index),
     ("index_add", index_add),
     ("index_add_", index_add_),
+    ("index_copy", index_copy),
+    ("index_copy_", index_copy_),
     ("index_put", index_put),
     ("index_put_", index_put_),
     ("index_select", index_select),

src/flag_gems/ops/__init__.py

@@ -157,6 +157,7 @@
 from flag_gems.ops.i0_ import i0_
 from flag_gems.ops.index import index
 from flag_gems.ops.index_add import index_add, index_add_
+from flag_gems.ops.index_copy_ import index_copy, index_copy_
 from flag_gems.ops.index_put import _index_put_impl_, index_put, index_put_
 from flag_gems.ops.index_select import index_select
 from flag_gems.ops.isclose import allclose, isclose
@@ -552,6 +553,8 @@
     "index",
     "index_add",
     "index_add_",
+    "index_copy",
+    "index_copy_",
     "index_put",
     "index_put_",
     "index_select",