Add fmod operator implementation, tests and benchmark

Author: Schopenhauer-loves-Hegel

benchmark/test_binary_pointwise_perf.py

@@ -49,6 +49,7 @@ def get_tflops(self, op, *args, **kwargs):
             ("pow", torch.pow, FLOAT_DTYPES),
             ("polar", torch.polar, [torch.float32]),
             ("floor_divide", torch.floor_divide, INT_DTYPES),
+            ("fmod", torch.fmod, FLOAT_DTYPES),
             ("remainder", torch.remainder, INT_DTYPES),
             ("logical_or", torch.logical_or, INT_DTYPES + BOOL_DTYPES),
             ("logical_and", torch.logical_and, INT_DTYPES + BOOL_DTYPES),

src/flag_gems/__init__.py

@@ -165,6 +165,10 @@ def torch_ge(v):
     ("fill_.Scalar", fill_scalar_),
     ("fill_.Tensor", fill_tensor_),
     ("flip", flip),
+    ("fmod.Scalar", fmod),
+    ("fmod.Tensor", fmod),
+    ("fmod_.Scalar", fmod_),
+    ("fmod_.Tensor", fmod_),
     ("floor_divide", floor_divide),
     ("floor_divide.Scalar", floor_divide),
     ("floor_divide_.Scalar", floor_divide_),

src/flag_gems/ops/__init__.py

@@ -92,6 +92,7 @@
 from flag_gems.ops.eye_m import eye_m
 from flag_gems.ops.fill import fill_scalar, fill_scalar_, fill_tensor, fill_tensor_
 from flag_gems.ops.flip import flip
+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
@@ -350,6 +351,8 @@
     "flash_attention_forward",
     "flash_attn_varlen_func",
     "flip",
+    "fmod",
+    "fmod_",
     "floor_divide",
     "floor_divide_",
     "full",

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

@@ -2166,3 +2166,79 @@ def test_accuracy_addcdiv(shape, dtype):
         res_out = torch.addcdiv(res_inp, t1, t2, value=v)
 
     gems_assert_close(res_out, ref_out, dtype)
+
+
+@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)
+    # Avoid division by zero
+    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)
+    # Avoid division by zero
+    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)
+
+    # Use larger tolerance for fmod with small divisors due to float32 precision limits
+    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)
+    # Avoid division by zero
+    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)
+    # Avoid division by zero
+    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)
+
+    # Use larger tolerance for fmod with small divisors due to float32 precision limits
+    atol = 1e-3 if abs(scalar) < 0.01 else 1e-4
+    gems_assert_close(res_out, ref_out, dtype, atol=atol)