[KernelGen][Nvidia] Add special_i1e operator with Triton kernel (flagos-ai#5276)

Co-authored-by: shawn <chanceying@foxmail.com>

Author: CoyeCHEN

benchmark/test_special_i1e.py

@@ -0,0 +1,31 @@
+# Copyright 2026 FlagOS Contributors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+import torch
+
+import flag_gems
+
+from . import base, consts
+
+
+@pytest.mark.special_i1e
+@pytest.mark.skipif(
+    flag_gems.vendor_name == "tsingmicro", reason="Issue #4131: not working"
+)
+def test_special_i1e():
+    bench = base.UnaryPointwiseBenchmark(
+        op_name="special_i1e", torch_op=torch.special.i1e, dtypes=consts.FLOAT_DTYPES
+    )
+    bench.run()

conf/operators.yaml

@@ -9674,6 +9674,20 @@ ops:
     stages:
       - beta: '5.0'
       - stable: '5.3'
+  - id: special_i1e
+    description: |
+      Computes the exponentially scaled modified Bessel function of the first kind
+      of order 1, i1e(x) = i1(x) * exp(-|x|), for each element in the input tensor.
+    for:
+      - special_i1e
+    labels:
+      - aten
+      - pointwise
+      - KernelGen
+    kind:
+      - Math
+    stages:
+      - alpha: '5.4'
   - id: special_legendre_polynomial_p
     description: |
       Computes the Legendre polynomial P_n(x) of degree n for each element in the input tensor,

src/flag_gems/__init__.py

@@ -888,6 +888,7 @@ def torch_ge(v):
     ("special_i0e_out", special_i0e_out),
     ("special_i1", special_i1),
     ("special_i1_out", special_i1_out),
+    ("special_i1e", special_i1e),
     ("special_legendre_polynomial_p", special_legendre_polynomial_p),
     ("special_log1p", special_log1p),
     ("special_log1p.out", special_log1p_out),

src/flag_gems/ops/__init__.py

@@ -646,6 +646,7 @@
 from flag_gems.ops.special_hermite_polynomial_h import special_hermite_polynomial_h
 from flag_gems.ops.special_i0e import special_i0e, special_i0e_out
 from flag_gems.ops.special_i1 import special_i1, special_i1_out
+from flag_gems.ops.special_i1e import special_i1e
 from flag_gems.ops.special_legendre_polynomial_p import special_legendre_polynomial_p
 from flag_gems.ops.special_log1p import special_log1p, special_log1p_out
 from flag_gems.ops.special_log_softmax import special_log_softmax
@@ -1504,6 +1505,7 @@
     "special_i0e_out",
     "special_i1",
     "special_i1_out",
+    "special_i1e",
     "special_legendre_polynomial_p",
     "special_log_softmax",
     "special_log1p",

src/flag_gems/ops/special_i1e.py

@@ -0,0 +1,106 @@
+# Copyright 2026, The FlagOS Contributors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# Generated by KernelGen: https://github.com/flagos-ai/KernelGen
+import logging
+
+import torch
+import triton
+import triton.language as tl
+
+import flag_gems
+
+logger = logging.getLogger(__name__)
+
+
+@triton.jit
+def _special_i1e_kernel(x_ptr, out_ptr, n_elements, BLOCK_SIZE: tl.constexpr):
+    pid = tl.program_id(axis=0)
+    block_start = pid * BLOCK_SIZE
+    offsets = block_start + tl.arange(0, BLOCK_SIZE)
+    mask = offsets < n_elements
+
+    x = tl.load(x_ptr + offsets, mask=mask, other=0.0)
+
+    # Compute in fp32 for accuracy/stability.
+    # i1e(x) = i1(x) * exp(-|x|). Fold the exponential scaling directly into
+    # the Cephes I1 approximation so the large-|x| branch never overflows.
+    xf = x.to(tl.float32)
+    ax = tl.abs(xf)
+
+    # Small region: |x| <= 3.75. I1(x) = x * P(y^2), y = x / 3.75.
+    y = xf / 3.75
+    y2 = y * y
+    p = 0.00032411
+    p = 0.00301532 + y2 * p
+    p = 0.02658733 + y2 * p
+    p = 0.15084934 + y2 * p
+    p = 0.51498869 + y2 * p
+    p = 0.87890594 + y2 * p
+    p = 0.5 + y2 * p
+    small = xf * p * tl.exp(-ax)
+
+    # Large region: |x| > 3.75. I1(x) ~ exp(|x|)/sqrt(|x|) * Q(3.75/|x|).
+    # Multiplying by exp(-|x|) cancels the exponential prefactor.
+    t = 3.75 / tl.maximum(ax, 1e-20)
+    q = -0.00420059
+    q = 0.01787654 + t * q
+    q = -0.02895312 + t * q
+    q = 0.02282967 + t * q
+    q = -0.01031555 + t * q
+    q = 0.00163801 + t * q
+    q = -0.00362018 + t * q
+    q = -0.03988024 + t * q
+    q = 0.39894228 + t * q
+    large = q / tl.sqrt(tl.maximum(ax, 1e-20))
+    # I1 is odd, so i1e is odd as well.
+    large = tl.where(xf < 0, -large, large)
+
+    y_out = tl.where(ax <= 3.75, small, large)
+    tl.store(out_ptr + offsets, y_out.to(x.dtype), mask=mask)
+
+
+def _run_special_i1e_kernel(x: torch.Tensor, out: torch.Tensor):
+    if x.device.type != flag_gems.device or out.device.type != flag_gems.device:
+        raise ValueError(f"Tensors must be {flag_gems.device} tensors")
+    assert x.dtype in (
+        torch.float16,
+        torch.bfloat16,
+        torch.float32,
+        torch.float64,
+    ), "Unsupported dtype"
+    assert out.dtype == x.dtype, "Output dtype must match input dtype"
+
+    x_c = x.contiguous()
+    out_c = out.contiguous()
+
+    n_elements = out_c.numel()
+    if n_elements == 0:
+        return out
+
+    grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),)
+    _special_i1e_kernel[grid](x_c, out_c, n_elements, BLOCK_SIZE=1024)
+
+    if out_c.data_ptr() != out.data_ptr():
+        out.copy_(out_c)
+    return out
+
+
+def special_i1e(x: torch.Tensor):
+    """
+    ATen wrapper: special_i1e(Tensor self) -> Tensor
+    """
+    logger.debug("GEMS SPECIAL_I1E")
+    out = torch.empty_like(x)
+    return _run_special_i1e_kernel(x, out)

tests/test_special_i1e.py

@@ -0,0 +1,34 @@
+# Copyright 2026 FlagOS Contributors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+import torch
+
+import flag_gems
+
+from . import accuracy_utils as utils
+
+
+@pytest.mark.special_i1e
+@pytest.mark.parametrize("shape", utils.POINTWISE_SHAPES)
+@pytest.mark.parametrize("dtype", utils.FLOAT_DTYPES)
+def test_special_i1e(shape, dtype, caplog):
+    inp = torch.randn(shape, dtype=dtype, device=flag_gems.device)
+    ref_inp = utils.to_reference(inp, True)
+    ref_out = torch.special.i1e(ref_inp)
+    with caplog.at_level("DEBUG", logger="flag_gems.ops.special_i1e"):
+        with flag_gems.use_gems():
+            res_out = torch.special.i1e(inp)
+    assert "GEMS SPECIAL_I1E" in caplog.text
+    utils.gems_assert_close(res_out, ref_out, dtype)