【KernelGen】Add roll operator (#1759)

* Add roll operator implementation, tests and benchmark

* fix: codestyle

* Rewrite roll with Triton kernel

---------

Co-authored-by: factnn <1050552884@qq.com>

Author: Schopenhauer-loves-Hegel

benchmark/test_generic_pointwise_perf.py

@@ -21,6 +21,14 @@ def flip_input_fn(shape, cur_dtype, device):
         yield inp, {"dims": (0,)}
 
 
+def roll_input_fn(shape, cur_dtype, device):
+    inp = generate_tensor_input(shape, cur_dtype, device)
+    if len(shape) > 1:
+        yield inp, {"shifts": (1, 2), "dims": (0, 1)}
+    else:
+        yield inp, {"shifts": 1, "dims": 0}
+
+
 def where_input_fn(shape, cur_dtype, device):
     inp1 = generate_tensor_input(shape, cur_dtype, device)
     inp2 = generate_tensor_input(shape, cur_dtype, device)
@@ -108,6 +116,13 @@ def addcdiv_input_fn(shape, cur_dtype, device):
             FLOAT_DTYPES + INT_DTYPES,
             marks=pytest.mark.flip,
         ),
+        pytest.param(
+            "roll",
+            torch.roll,
+            roll_input_fn,
+            FLOAT_DTYPES + INT_DTYPES,
+            marks=pytest.mark.roll,
+        ),
         pytest.param(
             "where", torch.where, where_input_fn, FLOAT_DTYPES, marks=pytest.mark.where
         ),

src/flag_gems/__init__.py

@@ -376,6 +376,7 @@ def torch_ge(v):
     ("resolve_conj", resolve_conj),
     ("resolve_neg", resolve_neg),
     ("rms_norm", rms_norm),
+    ("roll", roll),
     ("round", round),
     ("round.out", round_out),
     ("round_", round_),

src/flag_gems/ops/__init__.py

@@ -254,6 +254,7 @@
 from flag_gems.ops.resolve_conj import resolve_conj
 from flag_gems.ops.resolve_neg import resolve_neg
 from flag_gems.ops.rms_norm import rms_norm, rms_norm_backward, rms_norm_forward
+from flag_gems.ops.roll import roll
 from flag_gems.ops.round import round, round_, round_out
 from flag_gems.ops.rrelu_with_noise_backward import rrelu_with_noise_backward
 from flag_gems.ops.rsqrt import rsqrt, rsqrt_
@@ -638,6 +639,7 @@
     "rms_norm",
     "rms_norm_backward",
     "rms_norm_forward",
+    "roll",
     "round",
     "round_",
     "round_out",

src/flag_gems/ops/roll.py

@@ -0,0 +1,105 @@
+import logging
+
+import torch
+import triton
+import triton.language as tl
+
+from flag_gems.utils import libentry
+
+logger = logging.getLogger(__name__)
+
+
+@libentry()
+@triton.jit
+def roll_kernel(
+    inp_ptr,
+    out_ptr,
+    N,
+    dim_size,
+    shift,
+    inner_size,
+    BLOCK_SIZE: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    offset = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE).to(tl.int64)
+    mask = offset < N
+
+    # Decompose flat index into (outer, dim_idx, inner)
+    outer_stride = dim_size * inner_size
+    outer_idx = offset // outer_stride
+    remainder = offset % outer_stride
+    dim_idx = remainder // inner_size
+    inner_idx = remainder % inner_size
+
+    # Apply roll: source_dim_idx = (dim_idx - shift) % dim_size
+    source_dim_idx = (dim_idx - shift + dim_size) % dim_size
+
+    # Reconstruct source flat index
+    source_offset = outer_idx * outer_stride + source_dim_idx * inner_size + inner_idx
+
+    val = tl.load(inp_ptr + source_offset, mask=mask, other=0.0)
+    tl.store(out_ptr + offset, val, mask=mask)
+
+
+def _roll_single_dim(inp: torch.Tensor, shift: int, dim: int) -> torch.Tensor:
+    size = inp.size(dim)
+    if size == 0:
+        return inp.clone()
+
+    shift = shift % size
+    if shift == 0:
+        return inp.clone()
+
+    inp_contig = inp.contiguous()
+    out = torch.empty_like(inp_contig)
+
+    inner_size = 1
+    for i in range(dim + 1, inp.ndim):
+        inner_size *= inp.size(i)
+
+    N = inp.numel()
+    BLOCK_SIZE = 512
+    grid = lambda meta: (triton.cdiv(N, BLOCK_SIZE),)
+
+    roll_kernel[grid](
+        inp_contig,
+        out,
+        N,
+        size,
+        shift,
+        inner_size,
+        BLOCK_SIZE=BLOCK_SIZE,
+    )
+    return out
+
+
+def roll(inp: torch.Tensor, shifts, dims=None) -> torch.Tensor:
+    logger.debug("GEMS ROLL")
+
+    if inp.numel() == 0:
+        return inp.clone()
+
+    if dims is None:
+        if isinstance(shifts, (list, tuple)):
+            shift = shifts[0] if len(shifts) == 1 else sum(shifts)
+        else:
+            shift = shifts
+        original_shape = inp.shape
+        flat = inp.contiguous().reshape(-1)
+        out_flat = _roll_single_dim(flat, shift, 0)
+        return out_flat.reshape(original_shape)
+
+    if isinstance(dims, int):
+        dims = [dims]
+    if isinstance(shifts, int):
+        shifts = [shifts]
+
+    assert len(shifts) == len(dims), "shifts and dims must have the same length"
+
+    ndim = inp.ndim
+    dims = [d % ndim for d in dims]
+
+    result = inp
+    for shift, dim in zip(shifts, dims):
+        result = _roll_single_dim(result, shift, dim)
+    return result

tests/test_unary_pointwise_ops.py

@@ -2656,6 +2656,38 @@ def test_accuracy_floor_(shape, dtype):
     gems_assert_equal(res_out, ref_out)
 
 
+ROLL_SHIFTS_DIMS = [
+    (1, 0),
+    (-1, 0),
+    (2, -1),
+    (3, 1),
+]
+
+
+@pytest.mark.roll
+@pytest.mark.parametrize("shape", POINTWISE_SHAPES)
+@pytest.mark.parametrize("dtype", FLOAT_DTYPES + ALL_INT_DTYPES)
+@pytest.mark.parametrize("shifts_dims", ROLL_SHIFTS_DIMS)
+def test_accuracy_roll_single_dim(shape, dtype, shifts_dims):
+    shifts, dims = shifts_dims
+    ndim = len(shape)
+    # Adjust dims if it's out of range for this shape
+    if dims >= ndim or dims < -ndim:
+        pytest.skip(f"dims {dims} out of range for shape {shape}")
+
+    if dtype in ALL_FLOAT_DTYPES:
+        inp = torch.randn(shape, dtype=dtype, device=flag_gems.device)
+    else:
+        inp = torch.randint(-1000, 1000, shape, device=flag_gems.device).to(dtype)
+    ref_inp = to_reference(inp, False)
+
+    ref_out = torch.roll(ref_inp, shifts, dims)
+    with flag_gems.use_gems():
+        res_out = torch.roll(inp, shifts, dims)
+
+    gems_assert_equal(res_out, ref_out)
+
+
 @pytest.mark.special_i0e
 @pytest.mark.parametrize("shape", [(2, 3), (128, 256), (512, 512)])
 @pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@@ -2690,3 +2722,72 @@ def test_accuracy_special_i0e_out(shape, dtype):
         act_out = torch.ops.aten.special_i0e.out(x, out=out_act)
     gems_assert_close(act_out, ref_out, dtype)
     gems_assert_close(out_act, out_ref, dtype)
+
+
+ROLL_MULTI_DIMS = [
+    ((1, 2), (0, 1)),
+    ((-1, 1), (0, -1)),
+    ((2, -2), (-2, -1)),
+]
+
+
+@pytest.mark.roll
+@pytest.mark.parametrize("shape", POINTWISE_SHAPES)
+@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
+@pytest.mark.parametrize("shifts_dims", ROLL_MULTI_DIMS)
+def test_accuracy_roll_multi_dims(shape, dtype, shifts_dims):
+    shifts, dims = shifts_dims
+    ndim = len(shape)
+    # Check all dims are valid for this shape
+    for d in dims:
+        if d >= ndim or d < -ndim:
+            pytest.skip(f"dims {d} out of range for shape {shape}")
+
+    inp = torch.randn(shape, dtype=dtype, device=flag_gems.device)
+    ref_inp = to_reference(inp, False)
+
+    ref_out = torch.roll(ref_inp, shifts, dims)
+    with flag_gems.use_gems():
+        res_out = torch.roll(inp, shifts, dims)
+
+    gems_assert_equal(res_out, ref_out)
+
+
+ROLL_FLATTEN_SHIFTS = [1, -1, 5, -3]
+
+
+@pytest.mark.roll
+@pytest.mark.parametrize("shape", POINTWISE_SHAPES)
+@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
+@pytest.mark.parametrize("shifts", ROLL_FLATTEN_SHIFTS)
+def test_accuracy_roll_flatten(shape, dtype, shifts):
+    inp = torch.randn(shape, dtype=dtype, device=flag_gems.device)
+    ref_inp = to_reference(inp, False)
+
+    # Roll without specifying dims (flatten case)
+    ref_out = torch.roll(ref_inp, shifts)
+    with flag_gems.use_gems():
+        res_out = torch.roll(inp, shifts)
+
+    gems_assert_equal(res_out, ref_out)
+
+
+@pytest.mark.roll
+@pytest.mark.parametrize("shape", POINTWISE_SHAPES)
+@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
+def test_accuracy_roll_with_non_dense_input(shape, dtype):
+    if len(shape) < 2:
+        pytest.skip("Need at least 2D for non-dense test")
+
+    shape_dilated = tuple(item * 2 for item in shape)
+    inp = torch.randn(shape_dilated, dtype=dtype, device=flag_gems.device)[::2, ::2]
+    ref_inp = to_reference(inp, False)
+
+    shifts = 2
+    dims = 0
+
+    ref_out = torch.roll(ref_inp, shifts, dims)
+    with flag_gems.use_gems():
+        res_out = torch.roll(inp, shifts, dims)
+
+    gems_assert_equal(res_out, ref_out)