[FlagGems Operator Development Competition] Added roll operator

[YoihenBachu]

PR Category

• Operator

Type of Change

• New Feature

Description

Implements roll operator for FlagGems using Triton.

This PR adds support for the roll operation with:

• Forward operation: torch.roll(input, shifts, dims=None)
• Support for all dtypes (float, int, bool)
• Single or multiple dimension rolling
• Positive and negative shift values
• Dynamic code generation per tensor rank for optimal performance

The implementation follows the FlagGems operator development pattern using dynamic kernel code generation (similar to tile/repeat operators) with proper index remapping for circular shifts.

Issue

• Associated with FlagGems Operator Development Competition
• Implements roll operator as part of layout manipulation operations

Progress

• Change is properly reviewed (1 reviewer required, 2 recommended).
• Change is responded to an issue.
• Change is fully covered by a UT.

Performance

Performance benchmarks show competitive results against PyTorch native implementation:

• Tested across multiple tensor shapes: 1D to 5D tensors
• Tested across all dtypes: float, int, bool
• Speedup ≥0.9x compared to PyTorch
• Efficient memory access with single load/store per element

Test Coverage

• Basic functionality: Multiple shapes (1D-5D), all dtypes
• Single dimension rolling: Positive/negative shifts, boundary cases
• Multiple dimension rolling: 2D, 3D simultaneous shifts
• Edge cases: Zero shifts, shifts larger than dimension size, empty tensors
• Flatten mode: dims=None (rolls over flattened tensor)
• Negative indices: dims=-1, dims=-2
• Non-contiguous tensors: Transposed, strided inputs
• Shift normalization: Modulo behavior for large shifts

Run benchmarks with:

pytest benchmark/test_roll_perf.py

NOTE One performance test failed due to a lack of GPU hardware capacity. Cannot be resolved.

Run tests with:

pytest tests/test_unary_pointwise_ops.py -m roll

[tengqm]

LGTM

[YoihenBachu]

LGTM
Hi @tengqm, thanks for the review!

I’ve rerun pre-commit and pushed a new PR addressing the workflow issue. Could you please take another look and approve the workflows when you have a moment?

If possible, I’d also appreciate a quick review of the workflows for the other operators I submitted earlier, as getting early feedback would help catch any blockers sooner and give me some relief.

Thanks!

[YoihenBachu]

Is this error related to my development, or is it something else?
@zhangpeiyang1

[tengqm]

@YoihenBachu Double checked the gate status, the job failure is related to the implmentation of mm operator, irrelevant to your contribution.

[w1120029931-bit]

please resolve conflicts

[YoihenBachu]

please resolve conflicts

I am confused that there is already a roll.py implemented inside the ops. I believe that is exactly where I have to place my implementation and the conflict cannot be resolved if there is already a roll.py. Is this operator closed for competition?

[YoihenBachu]

This operator is not closed for competition. After evaluation, we have decided to adopt your implementation for this operator, as it better fits our requirements. The existing implementation will be replaced accordingly.

I will be pleased to resolve the merge conflicts and update it. How do I move forward from here? Do I replace the existing implementation with my own?

[douxetpur]

Do I replace the existing implementation with my own?

YES!

[YoihenBachu]

Done. Please review

[douxetpur]

Hi, thank you for your participation in the FlagOS Open Computing Competition 🙏

After review, we have chosen to move forward with a different implementation for this operator. As a result, this PR will be closed.

We appreciate your time and contribution, and hope to see more of your submissions in the future.