[KernelGen][thead] Add linalg_eigvals vendor specialization

Author: chx7514

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

@@ -15,10 +15,12 @@
 from .broadcast_to import broadcast_to
 from .index_copy_ import index_copy, index_copy_
 from .lcm import lcm, lcm_
+from .linalg_eigvals import _linalg_eigvals
 from .log_normal_ import log_normal_
 from .special_chebyshev_polynomial_u import special_chebyshev_polynomial_u
 
 __all__ = [
+    "_linalg_eigvals",
     "adaptive_max_pool3d_backward",
     "broadcast_to",
     "index_copy",

src/flag_gems/runtime/backend/_thead/ops/linalg_eigvals.py

@@ -0,0 +1,81 @@
+# 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
+
+logger = logging.getLogger(__name__)
+
+
+def _linalg_eigvals(inp):
+    """Compute the eigenvalues of a square matrix.
+
+    Thead specialization: the PPU's native CUDA key holds a cuSOLVER-backed
+    implementation (``cusolverDnXgeev``) that is unsupported on this hardware,
+    and FlagGems' generic registrar cannot override it (no ``allow_override``).
+    The generic Python body itself offloads to CPU LAPACK
+    (``torch.linalg.eigvals(inp.cpu())``), which is correct but is never
+    reached because dispatch is intercepted by the native cuSOLVER kernel.
+
+    This implementation keeps the proven CPU-offload path: eigenvalues are
+    computed on CPU via LAPACK and moved back to the original device. This
+    trades PCIe round-trip latency for correctness and device enablement —
+    a pure-Triton general eigensolver (QR iteration with shifts for
+    non-symmetric matrices) is out of scope here.
+    """
+    logger.debug(
+        "GEMS_THEAD _LINALG_EIGVALS, shape: %s, dtype: %s", inp.shape, inp.dtype
+    )
+
+    if inp.ndim < 2 or inp.shape[-2] != inp.shape[-1]:
+        raise ValueError(
+            "_linalg_eigvals: input must be a square matrix or batch of square matrices"
+        )
+
+    # cuSOLVER path supports float32/complex64/complex128; keep the same contract.
+    if inp.dtype not in (torch.float32, torch.complex64, torch.complex128):
+        raise TypeError(
+            f"_linalg_eigvals only supports float32/complex64/complex128, got {inp.dtype}"
+        )
+
+    # CPU uses LAPACK (geev), which is supported and accurate.
+    return torch.linalg.eigvals(inp.cpu()).to(inp.device)
+
+
+def _dispatched__linalg_eigvals(inp):
+    """CUDA-key dispatcher for ``aten::_linalg_eigvals``.
+
+    Force-overrides the cuSOLVER-occupied CUDA key: when FlagGems is active
+    (``use_gems`` context entered, indicated by ``current_work_registrar``) it
+    routes to the thead specialization above; otherwise it falls back to the
+    same CPU-LAPACK offload so the op stays usable in eager mode (the native
+    cuSOLVER kernel is broken on this hardware, so there is no device-side
+    baseline to preserve).
+    """
+    import flag_gems
+
+    if getattr(flag_gems, "current_work_registrar", None) is not None:
+        return _linalg_eigvals(inp)
+    # The native CUDA implementation (cusolverDnXgeev) is unsupported on PPU,
+    # and redispatch to a CUDA-resident keyset just loops back into this
+    # dispatcher. Fall back to the same CPU-LAPACK path directly.
+    logger.debug("GEMS_THEAD _LINALG_EIGVALS eager fallback (CPU offload)")
+    return _linalg_eigvals(inp)
+
+
+# Persistent registration that force-overrides the cuSOLVER-occupied CUDA key.
+_linalg_eigvals_lib = torch.library.Library("aten", "IMPL")
+_linalg_eigvals_lib.impl("_linalg_eigvals", _dispatched__linalg_eigvals, "CUDA")