Add rsub operator implementation, tests and benchmark

Author: factnn

benchmark/test_binary_pointwise_perf.py

@@ -75,6 +75,7 @@ def get_tflops(self, op, *args, **kwargs):
             ("div", torch.div, FLOAT_DTYPES + COMPLEX_DTYPES),
             ("mul", torch.mul, FLOAT_DTYPES + COMPLEX_DTYPES),
             ("sub", torch.sub, FLOAT_DTYPES + COMPLEX_DTYPES),
+            ("rsub", torch.rsub, FLOAT_DTYPES),
             ("pow", torch.pow, FLOAT_DTYPES),
             ("polar", torch.polar, [torch.float32]),
             ("floor_divide", torch.floor_divide, INT_DTYPES),

src/flag_gems/__init__.py

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

@@ -256,6 +256,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_
@@ -640,6 +641,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_binary_pointwise_ops.py

@@ -1864,6 +1864,40 @@ def test_accuracy_sub_scalar_scalar(dtype):
         gems_assert_close(res_out, ref_out, dtype)
 
 
+@pytest.mark.rsub
+@pytest.mark.parametrize("shape", POINTWISE_SHAPES)
+@pytest.mark.parametrize("alpha", SCALARS)
+@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
+def test_accuracy_rsub(shape, alpha, 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, True)
+    ref_inp2 = to_reference(inp2, True)
+
+    ref_out = torch.rsub(ref_inp1, ref_inp2, alpha=alpha)
+    with flag_gems.use_gems():
+        res_out = torch.rsub(inp1, inp2, alpha=alpha)
+
+    gems_assert_close(res_out, ref_out, dtype)
+
+
+@pytest.mark.rsub
+@pytest.mark.parametrize("shape", POINTWISE_SHAPES)
+@pytest.mark.parametrize("scalar", SCALARS)
+@pytest.mark.parametrize("alpha", SCALARS)
+@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
+def test_accuracy_rsub_tensor_scalar(shape, scalar, alpha, dtype):
+    inp1 = torch.randn(shape, dtype=dtype, device=flag_gems.device)
+    inp2 = scalar
+    ref_inp1 = to_reference(inp1, True)
+
+    ref_out = torch.rsub(ref_inp1, inp2, alpha=alpha)
+    with flag_gems.use_gems():
+        res_out = torch.rsub(inp1, inp2, alpha=alpha)
+
+    gems_assert_close(res_out, ref_out, dtype)
+
+
 @pytest.mark.where
 @pytest.mark.parametrize("shape", POINTWISE_SHAPES)
 @pytest.mark.parametrize("dtype", FLOAT_DTYPES)