Skip to content

Commit 6bdad3f

Browse files
Schopenhauer-loves-Hegelfactnnclaudebin913tengqm
committed
【KernelGen】Add fmod operator (#1737)
* feat: add fmod operator with tests and benchmark Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * Apply suggestion from @tengqm Co-authored-by: Qiming Teng <tengqm@outlook.com> Signed-off-by: bin913 <842884726@qq.com> * Apply suggestions from code review Co-authored-by: Qiming Teng <tengqm@outlook.com> Signed-off-by: bin913 <842884726@qq.com> * Apply suggestions from code review Co-authored-by: Qiming Teng <tengqm@outlook.com> Signed-off-by: Zang Peiyu <166481866+factnn@users.noreply.github.com> * Apply suggestions from code review Co-authored-by: Qiming Teng <tengqm@outlook.com> Signed-off-by: bin913 <842884726@qq.com> --------- Signed-off-by: bin913 <842884726@qq.com> Signed-off-by: Zang Peiyu <166481866+factnn@users.noreply.github.com> Co-authored-by: factnn <1050552884@qq.com> Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com> Co-authored-by: bin913 <842884726@qq.com> Co-authored-by: Qiming Teng <tengqm@outlook.com> Co-authored-by: Zang Peiyu <166481866+factnn@users.noreply.github.com>
1 parent caf0fac commit 6bdad3f

5 files changed

Lines changed: 194 additions & 0 deletions

File tree

‎benchmark/test_fmod.py‎

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,62 @@
1+
import pytest
2+
import torch
3+
4+
from . import base, consts, utils
5+
6+
7+
def _tensor_input_fn(shape, dtype, device):
8+
inp1 = utils.generate_tensor_input(shape, dtype, device)
9+
inp2 = utils.generate_tensor_input(shape, dtype, device)
10+
inp2 = torch.where(inp2 == 0, torch.ones_like(inp2), inp2)
11+
yield inp1, inp2
12+
13+
14+
def _scalar_input_fn(shape, dtype, device):
15+
inp1 = utils.generate_tensor_input(shape, dtype, device)
16+
yield inp1, 0.5
17+
18+
19+
@pytest.mark.fmod_tensor
20+
def test_fmod_tensor():
21+
bench = base.GenericBenchmark(
22+
input_fn=_tensor_input_fn,
23+
op_name="fmod_tensor",
24+
torch_op=torch.fmod,
25+
dtypes=consts.FLOAT_DTYPES,
26+
)
27+
bench.run()
28+
29+
30+
@pytest.mark.fmod_scalar
31+
def test_fmod_scalar():
32+
bench = base.GenericBenchmark(
33+
input_fn=_scalar_input_fn,
34+
op_name="fmod_scalar",
35+
torch_op=torch.fmod,
36+
dtypes=consts.FLOAT_DTYPES,
37+
)
38+
bench.run()
39+
40+
41+
@pytest.mark.fmod_tensor_
42+
def test_fmod_tensor_():
43+
bench = base.GenericBenchmark(
44+
input_fn=_tensor_input_fn,
45+
op_name="fmod_tensor_",
46+
torch_op=torch.Tensor.fmod_,
47+
dtypes=consts.FLOAT_DTYPES,
48+
inplace=True,
49+
)
50+
bench.run()
51+
52+
53+
@pytest.mark.fmod_scalar_
54+
def test_fmod_scalar_():
55+
bench = base.GenericBenchmark(
56+
input_fn=_scalar_input_fn,
57+
op_name="fmod_scalar_",
58+
torch_op=torch.Tensor.fmod_,
59+
dtypes=consts.FLOAT_DTYPES,
60+
inplace=True,
61+
)
62+
bench.run()

‎src/flag_gems/__init__.py‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -234,6 +234,10 @@ def torch_ge(v):
234234
("floor_divide_.Tensor", floor_divide_),
235235
("fmin", fmin),
236236
("fmin.out", fmin_out),
237+
("fmod.Scalar", fmod_scalar),
238+
("fmod.Tensor", fmod_tensor),
239+
("fmod_.Scalar", fmod_scalar_),
240+
("fmod_.Tensor", fmod_tensor_),
237241
("full", full),
238242
("full_like", full_like),
239243
("gather", gather),

‎src/flag_gems/ops/__init__.py‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -133,6 +133,7 @@
133133
from flag_gems.ops.flip import flip
134134
from flag_gems.ops.floor_ import floor_
135135
from flag_gems.ops.fmin import fmin, fmin_out
136+
from flag_gems.ops.fmod import fmod_scalar, fmod_scalar_, fmod_tensor, fmod_tensor_
136137
from flag_gems.ops.fp8_matmul import fp8_matmul
137138
from flag_gems.ops.full import full
138139
from flag_gems.ops.full_like import full_like
@@ -529,6 +530,10 @@
529530
"floor_divide_",
530531
"fmin",
531532
"fmin_out",
533+
"fmod_scalar",
534+
"fmod_scalar_",
535+
"fmod_tensor",
536+
"fmod_tensor_",
532537
"full",
533538
"full_like",
534539
"gather",

‎src/flag_gems/ops/fmod.py‎

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,51 @@
1+
import logging
2+
3+
import triton
4+
import triton.language as tl
5+
6+
from flag_gems.utils import pointwise_dynamic
7+
from flag_gems.utils.triton_lang_extension import fmod as _fmod
8+
9+
logger = logging.getLogger(__name__)
10+
11+
12+
@pointwise_dynamic(promotion_methods=[(0, 1, "DEFAULT")])
13+
@triton.jit
14+
def fmod_func(x, y):
15+
# Convert to float32 for computation to avoid libdevice float16/bfloat16 issues
16+
dtype = x.dtype
17+
x_fp32 = x.to(tl.float32)
18+
y_fp32 = y.to(tl.float32)
19+
result = _fmod(x_fp32, y_fp32)
20+
return result.to(dtype)
21+
22+
23+
@pointwise_dynamic(is_tensor=[True, False], promotion_methods=[(0, 1, "DEFAULT")])
24+
@triton.jit
25+
def fmod_func_tensor_scalar(x, y):
26+
# Convert to float32 for computation to avoid libdevice float16/bfloat16 issues
27+
dtype = x.dtype
28+
x_fp32 = x.to(tl.float32)
29+
y_fp32 = y.to(tl.float32)
30+
result = _fmod(x_fp32, y_fp32)
31+
return result.to(dtype)
32+
33+
34+
def fmod_tensor(A, B):
35+
logger.debug("GEMS FMOD_TENSOR")
36+
return fmod_func(A, B)
37+
38+
39+
def fmod_scalar(A, B):
40+
logger.debug("GEMS FMOD_SCALAR")
41+
return fmod_func_tensor_scalar(A, B)
42+
43+
44+
def fmod_tensor_(A, B):
45+
logger.debug("GEMS FMOD_TENSOR_")
46+
return fmod_func(A, B, out0=A)
47+
48+
49+
def fmod_scalar_(A, B):
50+
logger.debug("GEMS FMOD_SCALAR_")
51+
return fmod_func_tensor_scalar(A, B, out0=A)

‎tests/test_fmod.py‎

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,72 @@
1+
import pytest
2+
import torch
3+
4+
import flag_gems
5+
6+
from .accuracy_utils import (
7+
FLOAT_DTYPES,
8+
POINTWISE_SHAPES,
9+
SCALARS,
10+
gems_assert_close,
11+
to_reference,
12+
)
13+
14+
15+
@pytest.mark.fmod_tensor
16+
@pytest.mark.parametrize("shape", POINTWISE_SHAPES)
17+
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
18+
def test_fmod_tensor(shape, dtype):
19+
inp1 = torch.randn(shape, dtype=dtype, device=flag_gems.device)
20+
inp2 = torch.randn(shape, dtype=dtype, device=flag_gems.device)
21+
inp2 = torch.where(inp2 == 0, torch.ones_like(inp2), inp2)
22+
ref_inp1 = to_reference(inp1, True)
23+
ref_inp2 = to_reference(inp2, True)
24+
ref_out = torch.fmod(ref_inp1, ref_inp2)
25+
with flag_gems.use_gems():
26+
res_out = torch.fmod(inp1, inp2)
27+
gems_assert_close(res_out, ref_out, dtype)
28+
29+
30+
@pytest.mark.fmod_scalar
31+
@pytest.mark.parametrize("shape", POINTWISE_SHAPES)
32+
@pytest.mark.parametrize("scalar", SCALARS)
33+
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
34+
def test_fmod_scalar(shape, scalar, dtype):
35+
inp1 = torch.randn(shape, dtype=dtype, device=flag_gems.device)
36+
inp2 = scalar if scalar != 0 else 1.0
37+
ref_inp1 = to_reference(inp1, True)
38+
ref_out = torch.fmod(ref_inp1, inp2)
39+
with flag_gems.use_gems():
40+
res_out = torch.fmod(inp1, inp2)
41+
atol = 1e-3 if abs(scalar) < 0.01 else 1e-4
42+
gems_assert_close(res_out, ref_out, dtype, atol=atol)
43+
44+
45+
@pytest.mark.fmod_tensor_
46+
@pytest.mark.parametrize("shape", POINTWISE_SHAPES)
47+
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
48+
def test_fmod_tensor_inplace(shape, dtype):
49+
inp1 = torch.randn(shape, dtype=dtype, device=flag_gems.device)
50+
inp2 = torch.randn(shape, dtype=dtype, device=flag_gems.device)
51+
inp2 = torch.where(inp2 == 0, torch.ones_like(inp2), inp2)
52+
ref_inp1 = to_reference(inp1.clone(), True)
53+
ref_inp2 = to_reference(inp2, True)
54+
ref_out = ref_inp1.fmod_(ref_inp2)
55+
with flag_gems.use_gems():
56+
res_out = inp1.fmod_(inp2)
57+
gems_assert_close(res_out, ref_out, dtype)
58+
59+
60+
@pytest.mark.fmod_scalar_
61+
@pytest.mark.parametrize("shape", POINTWISE_SHAPES)
62+
@pytest.mark.parametrize("scalar", SCALARS)
63+
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
64+
def test_fmod_scalar_inplace(shape, scalar, dtype):
65+
inp1 = torch.randn(shape, dtype=dtype, device=flag_gems.device)
66+
inp2 = scalar if scalar != 0 else 1.0
67+
ref_inp1 = to_reference(inp1.clone(), True)
68+
ref_out = ref_inp1.fmod_(inp2)
69+
with flag_gems.use_gems():
70+
res_out = inp1.fmod_(inp2)
71+
atol = 1e-3 if abs(scalar) < 0.01 else 1e-4
72+
gems_assert_close(res_out, ref_out, dtype, atol=atol)

0 commit comments

Comments
 (0)