【KernelGen】Add index_copy_ operator

benchmark/test_index_copy.py

@@ -0,0 +1,50 @@
+import pytest
+import torch
+
+from . import base, consts, utils
+
+
+class IndexCopyBenchmark(base.GenericBenchmark):
+    def set_more_shapes(self):
+        return [(1, 2), (4096, 256), (200, 40999, 3)]
+
+    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
+
+
+@pytest.mark.index_copy
+def test_index_copy():
+    bench = IndexCopyBenchmark(
+        input_fn=_tensor_input_fn,
+        op_name="index_copy",
+        torch_op=torch.index_copy,
+        dtypes=consts.FLOAT_DTYPES,
+    )
+    bench.run()
+
+
+@pytest.mark.index_copy_
+def test_index_copy_():
+    bench = IndexCopyBenchmark(
+        input_fn=_tensor_input_fn,
+        op_name="index_copy_",
+        torch_op=torch.Tensor.index_copy_,
+        dtypes=consts.FLOAT_DTYPES,
+        inplace=True,
+    )
+    bench.run()

conf/operators.yaml

@@ -2817,6 +2817,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

src/flag_gems/__init__.py

@@ -277,6 +277,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

@@ -170,6 +170,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
@@ -593,6 +594,8 @@
     "index",
     "index_add",
     "index_add_",
+    "index_copy",
+    "index_copy_",
     "index_put",
     "index_put_",
     "index_select",