【KernelGen】Add fmod operator

benchmark/test_fmod.py

@@ -0,0 +1,62 @@
+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)
+    inp2 = torch.where(inp2 == 0, torch.ones_like(inp2), inp2)
+    yield inp1, inp2
+
+
+def _scalar_input_fn(shape, dtype, device):
+    inp1 = utils.generate_tensor_input(shape, dtype, device)
+    yield inp1, 0.5
+
+
+@pytest.mark.fmod_tensor
+def test_fmod_tensor():
+    bench = base.GenericBenchmark(
+        input_fn=_tensor_input_fn,
+        op_name="fmod_tensor",
+        torch_op=torch.fmod,
+        dtypes=consts.FLOAT_DTYPES,
+    )
+    bench.run()
+
+
+@pytest.mark.fmod_scalar
+def test_fmod_scalar():
+    bench = base.GenericBenchmark(
+        input_fn=_scalar_input_fn,
+        op_name="fmod_scalar",
+        torch_op=torch.fmod,
+        dtypes=consts.FLOAT_DTYPES,
+    )
+    bench.run()
+
+
+@pytest.mark.fmod_tensor_
+def test_fmod_tensor_():
+    bench = base.GenericBenchmark(
+        input_fn=_tensor_input_fn,
+        op_name="fmod_tensor_",
+        torch_op=torch.Tensor.fmod_,
+        dtypes=consts.FLOAT_DTYPES,
+        inplace=True,
+    )
+    bench.run()
+
+
+@pytest.mark.fmod_scalar_
+def test_fmod_scalar_():
+    bench = base.GenericBenchmark(
+        input_fn=_scalar_input_fn,
+        op_name="fmod_scalar_",
+        torch_op=torch.Tensor.fmod_,
+        dtypes=consts.FLOAT_DTYPES,
+        inplace=True,
+    )
+    bench.run()

src/flag_gems/__init__.py

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

src/flag_gems/ops/__init__.py

@@ -132,6 +132,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_scalar, fmod_scalar_, fmod_tensor, fmod_tensor_
 from flag_gems.ops.fp8_matmul import fp8_matmul
 from flag_gems.ops.full import full
 from flag_gems.ops.full_like import full_like
@@ -524,6 +525,10 @@
     "floor_divide_",
     "fmin",
     "fmin_out",
+    "fmod_scalar",
+    "fmod_scalar_",
+    "fmod_tensor",
+    "fmod_tensor_",
     "full",
     "full_like",
     "gather",

src/flag_gems/ops/fmod.py

@@ -0,0 +1,51 @@
+import logging
+
+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_tensor(A, B):
+    logger.debug("GEMS FMOD_TENSOR")
+    return fmod_func(A, B)
+
+
+def fmod_scalar(A, B):
+    logger.debug("GEMS FMOD_SCALAR")
+    return fmod_func_tensor_scalar(A, B)
+
+
+def fmod_tensor_(A, B):
+    logger.debug("GEMS FMOD_TENSOR_")
+    return fmod_func(A, B, out0=A)
+
+
+def fmod_scalar_(A, B):
+    logger.debug("GEMS FMOD_SCALAR_")
+    return fmod_func_tensor_scalar(A, B, out0=A)

tests/test_fmod.py

@@ -0,0 +1,72 @@
+import pytest
+import torch
+
+import flag_gems
+
+from .accuracy_utils import (
+    FLOAT_DTYPES,
+    POINTWISE_SHAPES,
+    SCALARS,
+    gems_assert_close,
+    to_reference,
+)
+
+
+@pytest.mark.fmod_tensor
+@pytest.mark.parametrize("shape", POINTWISE_SHAPES)
+@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
+def test_fmod_tensor(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_scalar
+@pytest.mark.parametrize("shape", POINTWISE_SHAPES)
+@pytest.mark.parametrize("scalar", SCALARS)
+@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
+def test_fmod_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.fmod_tensor_
+@pytest.mark.parametrize("shape", POINTWISE_SHAPES)
+@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
+def test_fmod_tensor_inplace(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.fmod_scalar_
+@pytest.mark.parametrize("shape", POINTWISE_SHAPES)
+@pytest.mark.parametrize("scalar", SCALARS)
+@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
+def test_fmod_scalar_inplace(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)