[KernelGen][Metax] Add special_gammainc operator for Metax backend (#98)

Co-authored-by: yzw1128 <yzw1128@users.noreply.github.com>

Author: yzw1128

src/flag_gems/runtime/backend/_metax/ops/__init__.py

@@ -44,6 +44,7 @@
 from .repeat_interleave import repeat_interleave_self_tensor
 from .resolve_conj import resolve_conj
 from .sigmoid import sigmoid
+from .special_gammainc import special_gammainc
 from .tanh import tanh
 from .unique import _unique2
 from .upsample_nearest2d import upsample_nearest2d
@@ -95,6 +96,7 @@
     "repeat_interleave_self_tensor",
     "resolve_conj",
     "sigmoid",
+    "special_gammainc",
     "tanh",
     "upsample_nearest2d",
     "zeros",

src/flag_gems/runtime/backend/_metax/ops/lt_.py

@@ -16,6 +16,7 @@
 """Metax-specific inplace wrapper for lt_ and lt_scalar_ with fused kernels
 that avoid uitofp i1->bf16 by using tl.where to produce 1.0/0.0 float
 values."""
+
 import logging
 
 import triton

src/flag_gems/runtime/backend/_metax/ops/masked_scatter.py

@@ -25,6 +25,7 @@
 inter-block synchronization.  num_warps is also made device-aware to
 respect Metax's 512-thread-per-block limit (warp_size=64 -> max 8 warps).
 """
+
 import logging
 
 import torch

src/flag_gems/runtime/backend/_metax/ops/special_gammainc.py

@@ -0,0 +1,183 @@
+# 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
+from flag_gems.utils import tl_extra_shim
+
+logger = logging.getLogger(__name__)
+
+
+@triton.jit
+def gammainc_kernel_metax(a_ptr, 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
+
+    a = tl.load(a_ptr + offsets, mask=mask, other=0.0)
+    x = tl.load(x_ptr + offsets, mask=mask, other=0.0)
+
+    # Compute in float32 for better precision
+    a_f32 = a.to(tl.float32)
+    x_f32 = x.to(tl.float32)
+
+    # Handle edge cases
+    # P(a, 0) = 0 for a > 0; NaN for a <= 0 or x < 0
+    result = tl.where((a_f32 > 0.0) & (x_f32 >= 0.0), 0.0, float("nan"))
+
+    # Determine which method to use based on x and a
+    # Series expansion is better when x < a + 1
+    use_series = x_f32 < (a_f32 + 1.0)
+
+    # =========================================================================
+    # Series expansion: P(a, x) = exp(-x) * x^a * sum_{n=0} x^n / Gamma(a+n+1)
+    # Term recurrence: term_0 = 1/a, term_n = term_{n-1} * x / (a + n)
+    # Use tl.where to zero out the term once converged, so subsequent
+    # multiplications produce 0 and the sum is not affected.
+    # 60 iterations is sufficient for float32 (machine eps ~1.2e-7).
+    # =========================================================================
+    series_sum = 0.0
+    term = 1.0 / a_f32
+    series_sum = term
+    for i in range(1, 60):
+        term = term * x_f32 / (a_f32 + tl.cast(i, tl.float32))
+        # Zero out converged lanes to avoid unnecessary accumulation
+        converged = tl.abs(term) < tl.abs(series_sum) * 1e-8
+        term = tl.where(converged, 0.0, term)
+        series_sum = series_sum + term
+
+    # Divide by Gamma(a) to get the regularized value P(a, x)
+    log_gamma_a = tl_extra_shim.lgamma(a_f32)
+    series_result = tl.exp(-x_f32 + a_f32 * tl.log(x_f32) - log_gamma_a) * series_sum
+
+    # =========================================================================
+    # Lentz's continued fraction for Q(a,x) = Gamma(a,x)/Gamma(a)
+    # for the large-x regime (x >= a + 1).
+    #
+    # CF = b_0 + a_1/(b_1 + a_2/(b_2 + ...))
+    #   b_0 = x + 1 - a
+    #   a_n = n(a - n),  b_n = x + 2n + 1 - a   (n >= 1)
+    # Then Q = e^{-x} * x^a / (Gamma(a) * CF) and P = 1 - Q.
+    #
+    # 100 iterations is sufficient for float32 convergence.
+    # Once |delta - 1| < eps, freeze delta at 1.0 so f_val stops changing.
+    # =========================================================================
+    tiny = 1e-30
+    b0 = x_f32 + 1.0 - a_f32
+    f_val = b0
+    C_val = b0
+    D_val = 0.0 * x_f32
+    for i_val in range(1, 100):
+        i_f = tl.cast(i_val, tl.float32)
+        an = i_f * (a_f32 - i_f)
+        bn = x_f32 + 2.0 * i_f + 1.0 - a_f32
+
+        D_val = bn + an * D_val
+        D_val = tl.where(tl.abs(D_val) < tiny, tiny, D_val)
+
+        C_val = bn + an / C_val
+        C_val = tl.where(tl.abs(C_val) < tiny, tiny, C_val)
+
+        D_val = 1.0 / D_val
+        delta = C_val * D_val
+        # Freeze converged lanes: once delta ~ 1, multiplying by 1 is a no-op
+        cf_converged = tl.abs(delta - 1.0) < 1e-8
+        delta = tl.where(cf_converged, 1.0, delta)
+        f_val = f_val * delta
+
+    log_gamma_a2 = tl_extra_shim.lgamma(a_f32)
+    log_q = a_f32 * tl.log(x_f32) - x_f32 - log_gamma_a2 - tl.log(f_val)
+    q_val = tl.exp(log_q)
+    q_val = tl.where(q_val > 1.0, 1.0, tl.where(q_val < 0.0, 0.0, q_val))
+    frac_result = 1.0 - q_val
+
+    # Combine results
+    result = tl.where(
+        (a_f32 > 0.0) & (x_f32 > 0.0),
+        tl.where(use_series, series_result, frac_result),
+        result,
+    )
+
+    # Store result
+    tl.store(out_ptr + offsets, result, mask=mask)
+
+
+def _launch_gammainc_metax(out: torch.Tensor, a: torch.Tensor, x: torch.Tensor):
+    assert (
+        a.device.type == flag_gems.device
+        and x.device.type == flag_gems.device
+        and out.device.type == flag_gems.device
+    ), f"All tensors must be {flag_gems.device} tensors"
+    assert (
+        out.numel() == a.numel() == x.numel()
+    ), "All tensors must have the same number of elements"
+    assert out.device == a.device == x.device, "All tensors must be on the same device"
+
+    # Ensure floating point compute
+    a_in = a
+    x_in = x
+    out_in = out
+
+    if not a_in.is_floating_point():
+        a_in = a_in.to(torch.get_default_dtype())
+    if not x_in.is_floating_point():
+        x_in = x_in.to(torch.get_default_dtype())
+
+    # Cast input to match the desired output dtype if needed
+    if a_in.dtype != out_in.dtype:
+        a_in = a_in.to(out_in.dtype)
+    if x_in.dtype != out_in.dtype:
+        x_in = x_in.to(out_in.dtype)
+
+    a_contig = a_in.contiguous()
+    x_contig = x_in.contiguous()
+    out_was_noncontig = not out_in.is_contiguous()
+    out_contig = out_in.contiguous() if out_was_noncontig else out_in
+
+    n_elements = out_contig.numel()
+    # 1024 provides good occupancy for element-wise gammainc kernel
+    BLOCK_SIZE = 1024
+    grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),)
+
+    gammainc_kernel_metax[grid](
+        a_contig, x_contig, out_contig, n_elements, BLOCK_SIZE=BLOCK_SIZE
+    )
+
+    if out_was_noncontig:
+        out_in.copy_(out_contig)
+    return out_in
+
+
+def special_gammainc(a: torch.Tensor, x: torch.Tensor, *, out: torch.Tensor = None):
+    logger.debug("GEMS_METAX SPECIAL_GAMMAINC")
+    if a.device.type != flag_gems.device:
+        raise ValueError(
+            f"gammainc: first input tensor must be on {flag_gems.device} device"
+        )
+    if x.device.type != flag_gems.device:
+        raise ValueError(
+            f"gammainc: second input tensor must be on {flag_gems.device} device"
+        )
+
+    if out is None:
+        out = torch.empty_like(a)
+    _launch_gammainc_metax(out, a, x)
+    return out