feat: add rsub operator with tests and benchmark (flagos-ai#1761)

Co-authored-by: factnn <1050552884@qq.com>
Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>

Author: 3 people

benchmark/test_rsub.py

@@ -0,0 +1,37 @@
+import pytest
+import torch
+
+from . import base, consts, utils
+
+
+def _tensor_input_fn(shape, dtype, device):
+    inp1 = utils.generate_tensor_input(shape, dtype, device)
+    inp2 = utils.generate_tensor_input(shape, dtype, device)
+    yield inp1, inp2
+
+
+def _scalar_input_fn(shape, dtype, device):
+    inp1 = utils.generate_tensor_input(shape, dtype, device)
+    yield inp1, 0.5
+
+
+@pytest.mark.rsub_tensor
+def test_rsub_tensor():
+    bench = base.GenericBenchmark(
+        input_fn=_tensor_input_fn,
+        op_name="rsub.Tensor",
+        torch_op=torch.rsub,
+        dtypes=consts.FLOAT_DTYPES,
+    )
+    bench.run()
+
+
+@pytest.mark.rsub_scalar
+def test_rsub_scalar():
+    bench = base.GenericBenchmark(
+        input_fn=_scalar_input_fn,
+        op_name="rsub.Scalar",
+        torch_op=torch.rsub,
+        dtypes=consts.FLOAT_DTYPES,
+    )
+    bench.run()

src/flag_gems/__init__.py

@@ -410,6 +410,8 @@ def torch_ge(v):
     ("rrelu_with_noise_backward", rrelu_with_noise_backward),
     ("rsqrt", rsqrt),
     ("rsqrt_", rsqrt_),
+    ("rsub.Scalar", rsub_scalar),
+    ("rsub.Tensor", rsub_tensor),
     ("scaled_softmax_backward", scaled_softmax_backward),
     ("scaled_softmax_forward", scaled_softmax_forward),
     ("scatter.reduce", scatter),

src/flag_gems/ops/__init__.py

@@ -277,6 +277,7 @@
 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_
+from flag_gems.ops.rsub import rsub_scalar, rsub_tensor
 from flag_gems.ops.scaled_softmax import scaled_softmax_backward, scaled_softmax_forward
 from flag_gems.ops.scatter import scatter, scatter_
 from flag_gems.ops.scatter_add_ import scatter_add_
@@ -697,6 +698,8 @@
     "rrelu_with_noise_backward",
     "rsqrt",
     "rsqrt_",
+    "rsub_scalar",
+    "rsub_tensor",
     "scaled_dot_product_attention",
     "scaled_dot_product_attention_backward",
     "scaled_dot_product_attention_forward",

src/flag_gems/ops/rsub.py

@@ -0,0 +1,31 @@
+import logging
+
+import triton
+
+from flag_gems.utils import pointwise_dynamic
+
+logger = logging.getLogger(__name__)
+
+
+@pointwise_dynamic(is_tensor=[True, True, False], promotion_methods=[(0, 1, "DEFAULT")])
+@triton.jit
+def rsub_func(x, y, alpha):
+    return y - x * alpha
+
+
+@pointwise_dynamic(
+    is_tensor=[True, False, False], promotion_methods=[(0, 1, "DEFAULT")]
+)
+@triton.jit
+def rsub_func_tensor_scalar(x, y, alpha):
+    return y - x * alpha
+
+
+def rsub_tensor(A, B, *, alpha=1):
+    logger.debug("GEMS RSUB_TENSOR")
+    return rsub_func(A, B, alpha)
+
+
+def rsub_scalar(A, B, alpha=1):
+    logger.debug("GEMS RSUB_SCALAR")
+    return rsub_func_tensor_scalar(A, B, alpha)

tests/test_rsub.py

@@ -0,0 +1,42 @@
+import pytest
+import torch
+
+import flag_gems
+
+from .accuracy_utils import (
+    FLOAT_DTYPES,
+    POINTWISE_SHAPES,
+    gems_assert_close,
+    to_reference,
+)
+
+
+@pytest.mark.rsub_tensor
+@pytest.mark.parametrize("shape", POINTWISE_SHAPES)
+@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
+def test_rsub_tensor(shape, dtype):
+    inp1 = torch.randn(shape, dtype=dtype, device=flag_gems.device)
+    inp2 = torch.randn(shape, dtype=dtype, device=flag_gems.device)
+    ref_inp1 = to_reference(inp1)
+    ref_inp2 = to_reference(inp2)
+
+    ref_out = torch.rsub(ref_inp1, ref_inp2)
+    with flag_gems.use_gems():
+        res_out = torch.rsub(inp1, inp2)
+
+    gems_assert_close(res_out, ref_out, dtype)
+
+
+@pytest.mark.rsub_scalar
+@pytest.mark.parametrize("shape", POINTWISE_SHAPES)
+@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
+def test_rsub_scalar(shape, dtype):
+    inp1 = torch.randn(shape, dtype=dtype, device=flag_gems.device)
+    ref_inp1 = to_reference(inp1)
+    inp2 = 0.5
+
+    ref_out = torch.rsub(ref_inp1, inp2)
+    with flag_gems.use_gems():
+        res_out = torch.rsub(inp1, inp2)
+
+    gems_assert_close(res_out, ref_out, dtype)