feat: add fmod operator implementation, tests and benchmark

Author: factnn

benchmark/test_binary_pointwise_perf.py

@@ -76,6 +76,7 @@ def get_tflops(self, op, *args, **kwargs):
             ("bitwise_or", torch.bitwise_or, INT_DTYPES + BOOL_DTYPES),
             ("div", torch.div, FLOAT_DTYPES + COMPLEX_DTYPES),
             ("dunder_or", lambda a, b: a | b, INT_DTYPES + BOOL_DTYPES),
+            # Comparison operations
             ("eq", torch.eq, FLOAT_DTYPES),
             ("equal", torch.equal, FLOAT_DTYPES),
             ("floor_divide", torch.floor_divide, INT_DTYPES),

benchmark/test_fmod_perf.py

@@ -0,0 +1,39 @@
+import pytest
+import torch
+
+from benchmark.attri_util import FLOAT_DTYPES
+from benchmark.performance_utils import GenericBenchmark, generate_tensor_input
+
+
+def fmod_input_fn(shape, cur_dtype, device):
+    inp1 = generate_tensor_input(shape, cur_dtype, device)
+    inp2 = generate_tensor_input(shape, cur_dtype, device)
+    inp2 = torch.where(inp2 == 0, torch.ones_like(inp2), inp2)
+    yield inp1, inp2
+
+
+def fmod_scalar_input_fn(shape, cur_dtype, device):
+    inp1 = generate_tensor_input(shape, cur_dtype, device)
+    yield inp1, 0.5
+
+
+@pytest.mark.fmod
+def test_perf_fmod():
+    bench = GenericBenchmark(
+        input_fn=fmod_input_fn,
+        op_name="fmod",
+        torch_op=torch.fmod,
+        dtypes=FLOAT_DTYPES,
+    )
+    bench.run()
+
+
+@pytest.mark.fmod
+def test_perf_fmod_scalar():
+    bench = GenericBenchmark(
+        input_fn=fmod_scalar_input_fn,
+        op_name="fmod.Scalar",
+        torch_op=torch.fmod,
+        dtypes=FLOAT_DTYPES,
+    )
+    bench.run()

src/flag_gems/__init__.py

@@ -213,6 +213,10 @@ def torch_ge(v):
     ("floor_divide_.Tensor", floor_divide_),
     ("fmin", fmin),
     ("fmin.out", fmin_out),
+    ("fmod.Scalar", fmod),
+    ("fmod.Tensor", fmod),
+    ("fmod_.Scalar", fmod_),
+    ("fmod_.Tensor", fmod_),
     ("full", full),
     ("full_like", full_like),
     ("gather", gather),

src/flag_gems/ops/__init__.py

@@ -124,6 +124,7 @@
 from flag_gems.ops.flip import flip
 from flag_gems.ops.floor_ import floor_
 from flag_gems.ops.fmin import fmin, fmin_out
+from flag_gems.ops.fmod import fmod, fmod_
 from flag_gems.ops.full import full
 from flag_gems.ops.full_like import full_like
 from flag_gems.ops.gather import gather, gather_backward
@@ -476,6 +477,8 @@
     "floor_divide_",
     "fmin",
     "fmin_out",
+    "fmod",
+    "fmod_",
     "full",
     "full_like",
     "gather",

src/flag_gems/ops/fmod.py

@@ -0,0 +1,51 @@
+import logging
+
+import torch
+import triton
+import triton.language as tl
+
+from flag_gems.utils import pointwise_dynamic
+from flag_gems.utils.triton_lang_extension import fmod as _fmod
+
+logger = logging.getLogger(__name__)
+
+
+@pointwise_dynamic(promotion_methods=[(0, 1, "DEFAULT")])
+@triton.jit
+def fmod_func(x, y):
+    # Convert to float32 for computation to avoid libdevice float16/bfloat16 issues
+    dtype = x.dtype
+    x_fp32 = x.to(tl.float32)
+    y_fp32 = y.to(tl.float32)
+    result = _fmod(x_fp32, y_fp32)
+    return result.to(dtype)
+
+
+@pointwise_dynamic(is_tensor=[True, False], promotion_methods=[(0, 1, "DEFAULT")])
+@triton.jit
+def fmod_func_tensor_scalar(x, y):
+    # Convert to float32 for computation to avoid libdevice float16/bfloat16 issues
+    dtype = x.dtype
+    x_fp32 = x.to(tl.float32)
+    y_fp32 = y.to(tl.float32)
+    result = _fmod(x_fp32, y_fp32)
+    return result.to(dtype)
+
+
+def fmod(A, B):
+    logger.debug("GEMS FMOD")
+    if isinstance(A, torch.Tensor) and isinstance(B, torch.Tensor):
+        return fmod_func(A, B)
+    elif isinstance(A, torch.Tensor):
+        return fmod_func_tensor_scalar(A, B)
+    else:
+        # Both scalar - fallback to PyTorch
+        return torch.fmod(torch.tensor(A), B)
+
+
+def fmod_(A, B):
+    logger.debug("GEMS FMOD_")
+    if isinstance(B, torch.Tensor):
+        return fmod_func(A, B, out0=A)
+    else:
+        return fmod_func_tensor_scalar(A, B, out0=A)

tests/test_binary_pointwise_ops.py

@@ -2722,3 +2722,5 @@ def test_dunder_ior_scalar(shape, dtype):
         inp1 |= inp2
 
     gems_assert_equal(inp1, ref_inp1)
+
+

tests/test_fmod.py

@@ -0,0 +1,74 @@
+import pytest
+import torch
+
+import flag_gems
+
+from .accuracy_utils import (
+    FLOAT_DTYPES,
+    POINTWISE_SHAPES,
+    SCALARS,
+    gems_assert_close,
+    to_reference,
+)
+
+
+@pytest.mark.fmod
+@pytest.mark.parametrize("shape", POINTWISE_SHAPES)
+@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
+def test_accuracy_fmod(shape, dtype):
+    inp1 = torch.randn(shape, dtype=dtype, device=flag_gems.device)
+    inp2 = torch.randn(shape, dtype=dtype, device=flag_gems.device)
+    inp2 = torch.where(inp2 == 0, torch.ones_like(inp2), inp2)
+    ref_inp1 = to_reference(inp1, True)
+    ref_inp2 = to_reference(inp2, True)
+    ref_out = torch.fmod(ref_inp1, ref_inp2)
+    with flag_gems.use_gems():
+        res_out = torch.fmod(inp1, inp2)
+    gems_assert_close(res_out, ref_out, dtype)
+
+
+@pytest.mark.fmod
+@pytest.mark.parametrize("shape", POINTWISE_SHAPES)
+@pytest.mark.parametrize("scalar", SCALARS)
+@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
+def test_accuracy_fmod_tensor_scalar(shape, scalar, dtype):
+    inp1 = torch.randn(shape, dtype=dtype, device=flag_gems.device)
+    inp2 = scalar if scalar != 0 else 1.0
+    ref_inp1 = to_reference(inp1, True)
+    ref_out = torch.fmod(ref_inp1, inp2)
+    with flag_gems.use_gems():
+        res_out = torch.fmod(inp1, inp2)
+    atol = 1e-3 if abs(scalar) < 0.01 else 1e-4
+    gems_assert_close(res_out, ref_out, dtype, atol=atol)
+
+
+@pytest.mark.inplace
+@pytest.mark.fmod_
+@pytest.mark.parametrize("shape", POINTWISE_SHAPES)
+@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
+def test_accuracy_fmod_(shape, dtype):
+    inp1 = torch.randn(shape, dtype=dtype, device=flag_gems.device)
+    inp2 = torch.randn(shape, dtype=dtype, device=flag_gems.device)
+    inp2 = torch.where(inp2 == 0, torch.ones_like(inp2), inp2)
+    ref_inp1 = to_reference(inp1.clone(), True)
+    ref_inp2 = to_reference(inp2, True)
+    ref_out = ref_inp1.fmod_(ref_inp2)
+    with flag_gems.use_gems():
+        res_out = inp1.fmod_(inp2)
+    gems_assert_close(res_out, ref_out, dtype)
+
+
+@pytest.mark.inplace
+@pytest.mark.fmod_
+@pytest.mark.parametrize("shape", POINTWISE_SHAPES)
+@pytest.mark.parametrize("scalar", SCALARS)
+@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
+def test_accuracy_fmod_tensor_scalar_(shape, scalar, dtype):
+    inp1 = torch.randn(shape, dtype=dtype, device=flag_gems.device)
+    inp2 = scalar if scalar != 0 else 1.0
+    ref_inp1 = to_reference(inp1.clone(), True)
+    ref_out = ref_inp1.fmod_(inp2)
+    with flag_gems.use_gems():
+        res_out = inp1.fmod_(inp2)
+    atol = 1e-3 if abs(scalar) < 0.01 else 1e-4
+    gems_assert_close(res_out, ref_out, dtype, atol=atol)