[KernelGen][metax] Add mvlgamma_ vendor specialization with constexpr P optimization (#233)

Co-authored-by: yzw1128 <yzw1128@users.noreply.github.com>
Co-authored-by: Dongxu-H <dxhan@baai.ac.cn>

Author: 2 people

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

@@ -28,6 +28,7 @@
 from .matmul_int8 import matmul_int8
 from .min import min, min_dim
 from .mm import mm, mm_out
+from .mvlgamma_ import mvlgamma_
 from .new_ones import new_ones
 from .nonzero import nonzero
 from .nonzero_numpy import nonzero_numpy
@@ -88,6 +89,7 @@
     "masked_fill_",
     "min_dim",
     "min",
+    "mvlgamma_",
     "mm",
     "mm_out",
     "new_ones",

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

@@ -0,0 +1,102 @@
+# 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
+
+from flag_gems.runtime import torch_device_fn
+from flag_gems.utils import tl_extra_shim
+
+logger = logging.getLogger(__name__)
+
+
+@triton.jit
+def mvlgamma_kernel(
+    x_ptr,
+    n_elements,
+    BLOCK_SIZE: tl.constexpr,
+    P: tl.constexpr,
+):
+    """Muxi-optimized 1D kernel for mvlgamma_.
+
+    P is a tl.constexpr so the Triton compiler generates only the
+    needed number of lgamma calls via tl.static_range(P), eliminating
+    dead computation for small p values.
+    """
+    pid = tl.program_id(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)
+    x_f32 = x.to(tl.float32)
+
+    LOG_PI = 1.1447298858494002
+    LOG_PI_OVER_4 = LOG_PI / 4.0
+    constant_term = P * (P - 1) * LOG_PI_OVER_4
+
+    # Compile-time loop: exactly P iterations, zero dead lgamma calls.
+    sum_term = tl.zeros_like(x_f32)
+    for k in tl.static_range(P):
+        sum_term += tl_extra_shim.lgamma(x_f32 - 0.5 * k)
+
+    result = constant_term + sum_term
+    tl.store(x_ptr + offsets, result.to(x.dtype), mask=mask)
+
+
+def mvlgamma_(*args, **kwargs):
+    """In-place multivariate log-gamma function, optimized for Metax GPUs.
+
+    Uses a hand-crafted 1D contiguous Triton kernel with P as a
+    constexpr, so the compiler generates exactly p lgamma calls per
+    element with zero dead-code overhead.
+    """
+    logger.debug("GEMS_METAX MVLGAMMA_")
+
+    x = args[0]
+    p = args[1] if len(args) > 1 else kwargs.get("p", 1)
+
+    if not isinstance(x, torch.Tensor):
+        raise TypeError("mvlgamma_ expects a torch.Tensor as the first argument")
+
+    if not isinstance(p, int) or p < 1:
+        raise ValueError("p must be a positive integer")
+
+    if p > 12:
+        raise ValueError("p must be <= 12 for this implementation")
+
+    n_elements = x.numel()
+    if n_elements == 0:
+        return x
+
+    # For non-contiguous tensors, create a contiguous copy, operate on
+    # it, and copy back to preserve in-place semantics.
+    if not x.is_contiguous():
+        y = x.contiguous()
+        grid = (triton.cdiv(n_elements, 1024),)
+        with torch_device_fn.device(y.device):
+            mvlgamma_kernel[grid](y, n_elements, BLOCK_SIZE=1024, P=p)
+        x.copy_(y)
+        return x
+
+    grid = (triton.cdiv(n_elements, 1024),)
+    with torch_device_fn.device(x.device):
+        mvlgamma_kernel[grid](x, n_elements, BLOCK_SIZE=1024, P=p)
+
+    return x