Merge branch 'master' into pr/mthreads-reflection_pad3d_backward

Author: Dongxu-H

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

@@ -23,6 +23,11 @@
 from .mm import mm, mm_out
 from .repeat import repeat
 from .scatter_add import scatter_add_
+from .special_chebyshev_polynomial_w import (
+    special_chebyshev_polynomial_w,
+    special_chebyshev_polynomial_w_out,
+)
+from .special_gammainc import special_gammainc
 from .tile import tile
 from .var import var, var_correction, var_dim
 
@@ -39,6 +44,9 @@
     "matmul_int8",
     "repeat",
     "scatter_add_",
+    "special_chebyshev_polynomial_w",
+    "special_chebyshev_polynomial_w_out",
+    "special_gammainc",
     "tile",
     "var",
     "var_correction",

src/flag_gems/runtime/backend/_iluvatar/ops/special_chebyshev_polynomial_w.py

@@ -0,0 +1,217 @@
+# 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 chebyshev_polynomial_w_scalar_n_kernel(
+    x_ptr,
+    n_ptr,
+    out_ptr,
+    n_elements,
+    BLOCK_SIZE: tl.constexpr,
+    MAX_DEGREE: tl.constexpr,
+):
+    """Optimized kernel for scalar n.
+
+    Reads n from device memory (no host-device sync needed).
+    Uses fixed MAX_DEGREE unrolling with tl.where to select result.
+    """
+    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)
+    # Load scalar n (single element) - avoids host-device sync
+    n_val = tl.load(n_ptr).to(tl.int32)
+
+    x_f32 = x.to(tl.float32)
+
+    # W_0(x) = 1
+    w0 = 1.0 + 0.0 * x_f32
+    # W_1(x) = 2x + 1
+    w1 = 2.0 * x_f32 + 1.0
+
+    result = w0
+    result = tl.where(n_val >= 1, w1, result)
+
+    # Recurrence: W_k(x) = 2*x*W_{k-1}(x) - W_{k-2}(x)
+    w_km2 = w0
+    w_km1 = w1
+    for k in tl.static_range(2, MAX_DEGREE):
+        w_k = 2.0 * x_f32 * w_km1 - w_km2
+        result = tl.where(n_val >= k, w_k, result)
+        w_km2 = w_km1
+        w_km1 = w_k
+
+    tl.store(out_ptr + offsets, result, mask=mask)
+
+
+@triton.jit
+def chebyshev_polynomial_w_kernel(
+    x_ptr,
+    n_ptr,
+    out_ptr,
+    n_elements,
+    BLOCK_SIZE: tl.constexpr,
+    MAX_DEGREE: tl.constexpr,
+):
+    """General kernel when n varies per element."""
+    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)
+    n = tl.load(n_ptr + offsets, mask=mask, other=0)
+
+    x_f32 = x.to(tl.float32)
+    n_i32 = n.to(tl.int32)
+
+    # W_0(x) = 1
+    w0 = 1.0 + 0.0 * x_f32
+    # W_1(x) = 2x + 1
+    w1 = 2.0 * x_f32 + 1.0
+
+    result = w0
+    result = tl.where(n_i32 >= 1, w1, result)
+
+    # Compute W_k for k >= 2 using recurrence
+    w_km2 = w0
+    w_km1 = w1
+    for k in tl.static_range(2, MAX_DEGREE):
+        w_k = 2.0 * x_f32 * w_km1 - w_km2
+        result = tl.where(n_i32 >= k, w_k, result)
+        w_km2 = w_km1
+        w_km1 = w_k
+
+    tl.store(out_ptr + offsets, result, mask=mask)
+
+
+# Fixed unroll depth: 21 covers most practical polynomial degrees.
+# Avoids host-device sync to read n value at runtime.
+_DEFAULT_MAX_DEGREE = 21
+
+
+def _launch_chebyshev_w(out: torch.Tensor, x: torch.Tensor, n: torch.Tensor):
+    assert (
+        x.device.type == flag_gems.device
+        and n.device.type == flag_gems.device
+        and out.device.type == flag_gems.device
+    ), f"All tensors must be {flag_gems.device} tensors"
+
+    x_in = x
+    if not x_in.is_floating_point():
+        x_in = x_in.to(torch.get_default_dtype())
+    if x_in.dtype != out.dtype:
+        x_in = x_in.to(out.dtype)
+
+    x_contig = x_in.contiguous()
+    out_was_noncontig = not out.is_contiguous()
+    out_contig = out.contiguous() if out_was_noncontig else out
+    n_elements = out_contig.numel()
+    BLOCK_SIZE = 1024
+    grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),)
+
+    if n.dim() == 0 or n.numel() == 1:
+        # Scalar n: read n in kernel from device memory, no sync.
+        n_contig = n.contiguous()
+        chebyshev_polynomial_w_scalar_n_kernel[grid](
+            x_contig,
+            n_contig,
+            out_contig,
+            n_elements,
+            BLOCK_SIZE=BLOCK_SIZE,
+            MAX_DEGREE=_DEFAULT_MAX_DEGREE,
+            num_warps=4,
+            num_stages=2,
+        )
+    else:
+        # General case: n varies per element, need broadcast
+        x_in_b, n_in_b = torch.broadcast_tensors(x_in, n)
+        x_contig = x_in_b.contiguous()
+        n_contig = n_in_b.contiguous()
+
+        chebyshev_polynomial_w_kernel[grid](
+            x_contig,
+            n_contig,
+            out_contig,
+            n_elements,
+            BLOCK_SIZE=BLOCK_SIZE,
+            MAX_DEGREE=_DEFAULT_MAX_DEGREE,
+            num_warps=4,
+            num_stages=2,
+        )
+
+    if out_was_noncontig:
+        out.copy_(out_contig)
+    return out
+
+
+def special_chebyshev_polynomial_w(x, n):
+    logger.debug("GEMS_ILUVATAR SPECIAL_CHEBYSHEV_POLYNOMIAL_W")
+    if not isinstance(x, torch.Tensor):
+        x = torch.tensor(x, dtype=torch.float32)
+    if x.device.type != flag_gems.device:
+        raise ValueError(
+            "special_chebyshev_polynomial_w: "
+            f"input x must be on {flag_gems.device} device"
+        )
+    if x.dtype not in (torch.float32, torch.float64):
+        raise ValueError(
+            "special_chebyshev_polynomial_w only supports "
+            f"float32/float64, got {x.dtype}"
+        )
+    if not isinstance(n, torch.Tensor):
+        n = torch.tensor(n, dtype=torch.int64, device=x.device)
+    if n.device.type != flag_gems.device:
+        n = n.to(x.device)
+
+    out = torch.empty_like(x)
+    _launch_chebyshev_w(out, x, n)
+    return out
+
+
+def special_chebyshev_polynomial_w_out(x, n, out):
+    logger.debug("GEMS_ILUVATAR SPECIAL_CHEBYSHEV_POLYNOMIAL_W_OUT")
+    if not isinstance(x, torch.Tensor):
+        x = torch.tensor(x, dtype=torch.float32)
+    if x.device.type != flag_gems.device:
+        raise ValueError(
+            "special_chebyshev_polynomial_w: "
+            f"input x must be on {flag_gems.device} device"
+        )
+    if x.dtype not in (torch.float32, torch.float64):
+        raise ValueError(
+            "special_chebyshev_polynomial_w only supports "
+            f"float32/float64, got {x.dtype}"
+        )
+    if not isinstance(n, torch.Tensor):
+        n = torch.tensor(n, dtype=torch.int64, device=x.device)
+    if n.device.type != flag_gems.device:
+        n = n.to(x.device)
+
+    _launch_chebyshev_w(out, x, n)
+    return out