# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# See LICENSE for license information.
"""
Triton utilities for JAX primitives.

This module provides utility functions for integrating Triton kernels into
JAX primitives. Triton is only imported when this module is used.

Triton Package Compatibility:
    There are two Triton packages that can be used:

    1. 'triton' (from OpenAI/PyPI): Standard package, works with JAX out of the box.
       Install with: pip install triton

    2. 'pytorch-triton' (from PyTorch's index): Bundled with PyTorch, includes
       PyTorch-specific patches. Version format: "3.0.0+<commit_sha>"

       IMPORTANT: The 'pytorch-triton' package on PyPI (version 0.0.1) is a
       placeholder that will NOT work. The real pytorch-triton is only available
       from PyTorch's package index and is auto-installed with PyTorch:
           pip install torch --index-url https://download.pytorch.org/whl/cu121

       pytorch-triton has been tested to work with JAX Triton kernels.

Environment Variables:
    NVTE_USE_PYTORCH_TRITON: If set to "1", explicitly acknowledge using
        pytorch-triton for JAX Triton kernels (suppresses warnings). This is
        useful when both JAX and PyTorch are installed in the same environment.
        Default is "0".
    NVTE_JAX_ENFORCE_TRITON_AUTOTUNING: If set to "1", raise a RuntimeError when
        the installed JAX is too old to safely run TritonAutotunedKernelCall
        (jax-ml/jax#35218) instead of silently falling back to non-autotuned
        dispatch. Useful for CI or debugging to ensure autotuning is active.
        Default is "0" (silent compatibility fallback).
"""

import hashlib
import os
import warnings
from typing import Any, Callable, Mapping
import zlib

from packaging import version

from jax import core
from jaxlib.mlir import ir
import jax
import jax.numpy as jnp

from ..version_utils import (
    TRITON_AUTOTUNED_INPUT_OUTPUT_ALIAS_MIN_JAX_VERSION,
    TRITON_EXTENSION_CUDA_GRAPH_MIN_JAX_VERSION,
    TRITON_EXTENSION_MIN_JAX_VERSION,
    is_triton_autotuned_alias_safe,
    is_triton_extension_supported,
    jax_version_meet_requirement,
)


# Placeholder package version on PyPI that should never be used
_PYTORCH_TRITON_PLACEHOLDER_VERSION = "0.0.1"


def _detect_triton_package():
    """Detect which Triton package is installed and validate compatibility.

    Returns:
        tuple: (triton_version: str or None, is_pytorch_triton: bool, is_placeholder: bool)

    The function detects:
    - None: Triton not installed
    - Standard triton from OpenAI (versions like "3.1.0")
    - Real pytorch-triton from PyTorch's index (versions like "3.0.0+45fff310c8")
    - Placeholder pytorch-triton from PyPI (version "0.0.1" - broken, raises RuntimeError)
    """
    try:
        import triton

        triton_version = getattr(triton, "__version__", "unknown")
    except ImportError:
        return None, False, False
    except RuntimeError as e:
        # The placeholder pytorch-triton package from PyPI raises:
        # RuntimeError: "Should never be installed"
        if "Should never be installed" in str(e):
            return _PYTORCH_TRITON_PLACEHOLDER_VERSION, False, True
        raise

    # Check for placeholder package (version 0.0.1 from PyPI)
    is_placeholder = triton_version == _PYTORCH_TRITON_PLACEHOLDER_VERSION

    # Real pytorch-triton versions have a commit SHA suffix like "3.0.0+45fff310c8"
    is_pytorch_triton = "+" in triton_version and len(triton_version.split("+")[-1]) >= 8

    return triton_version, is_pytorch_triton, is_placeholder


def _check_triton_compatibility():
    """Check Triton package compatibility and emit warnings if necessary.

    This function handles the case where both JAX and PyTorch may be installed,
    each expecting different Triton packages:
    - JAX typically uses the standard 'triton' package from OpenAI
    - PyTorch uses 'pytorch-triton' which is versioned with commit SHAs

    The NVTE_USE_PYTORCH_TRITON environment variable can be used to explicitly
    acknowledge using pytorch-triton with JAX (suppresses warnings).

    Raises:
        ImportError: If triton is not installed or the placeholder package is detected.
    """
    triton_version, is_pytorch_triton, is_placeholder = _detect_triton_package()

    # Handle placeholder package from PyPI
    if is_placeholder:
        raise ImportError(
            "Detected the placeholder 'pytorch-triton' package (version 0.0.1) from PyPI.\n"
            "This is NOT a functional Triton installation.\n\n"
            "The placeholder package exists to prevent namespace conflicts. To fix this:\n\n"
            "Option 1 - Use standard Triton (recommended for JAX-only environments):\n"
            "    pip uninstall pytorch-triton triton\n"
            "    pip install triton\n\n"
            "Option 2 - Use real pytorch-triton (for mixed JAX+PyTorch environments):\n"
            "    pip uninstall pytorch-triton triton\n"
            "    pip install torch --index-url https://download.pytorch.org/whl/cu121\n"
            "    # pytorch-triton is automatically installed as a torch dependency\n\n"
            "Note: Do NOT run 'pip install pytorch-triton' directly - this installs\n"
            "the broken placeholder. The real pytorch-triton only comes from PyTorch's index."
        )

    if triton_version is None:
        raise ImportError(
            "Triton is required for transformer_engine.jax.triton_extensions.\n\n"
            "Option 1 - Install standard Triton (recommended for JAX-only):\n"
            "    pip install triton\n\n"
            "Option 2 - Install PyTorch with pytorch-triton (for mixed environments):\n"
            "    pip install torch --index-url https://download.pytorch.org/whl/cu121\n\n"
            "If you don't need Triton, use transformer_engine.jax.cpp_extensions instead."
        )

    val = os.environ.get("NVTE_USE_PYTORCH_TRITON", "0")
    try:
        use_pytorch_triton_explicit = bool(int(val))
    except ValueError as e:
        raise ValueError(
            f"NVTE_USE_PYTORCH_TRITON must be an integer (0 or 1), got: {val!r}"
        ) from e

    if is_pytorch_triton:
        if use_pytorch_triton_explicit:
            # User explicitly opted in - just log info (no warning)
            pass  # Silent acknowledgment, no warning needed
        else:
            # pytorch-triton detected but user didn't explicitly opt in
            warnings.warn(
                f"Detected pytorch-triton package (version {triton_version}) instead of the"
                " standard 'triton' package from OpenAI. This typically happens when PyTorch is"
                " installed alongside JAX.\n\npytorch-triton is compatible with JAX Triton"
                " kernels. To suppress this warning, set:\n    export"
                " NVTE_USE_PYTORCH_TRITON=1\n\nAlternatively, for a JAX-only environment:\n  - Use"
                " separate virtual environments for JAX and PyTorch, or\n  - Use"
                " transformer_engine.jax.cpp_extensions instead (CUDA-based, no Triton needed)",
                category=UserWarning,
                stacklevel=3,
            )

    return triton_version, is_pytorch_triton


# Perform compatibility check and get triton info
_TRITON_VERSION, _IS_PYTORCH_TRITON = _check_triton_compatibility()

# Enforce minimum JAX version before importing gpu_triton.  The segfault on old
# jaxlib occurs at Triton kernel dispatch time, not at import time, so gpu_triton
# itself is safe to import on older jaxlib.  The guard is placed here (before the
# import) as a belt-and-suspenders measure so that if the import behaviour ever
# changes, we still fail fast with a clear error rather than a cryptic crash.
if not is_triton_extension_supported():
    raise RuntimeError(
        f"JAX >= {TRITON_EXTENSION_MIN_JAX_VERSION} required for "
        "transformer_engine.jax.triton_extensions. "
        "Triton kernel dispatch segfaults with older jaxlib. "
        f"Current jax version: {jax.__version__}. "
        "Please upgrade: pip install --upgrade jax jaxlib. "
        "If you don't need Triton, use transformer_engine.jax.cpp_extensions instead."
    )

try:
    from jax._src.lib import gpu_triton
    from triton.compiler import compiler as tc
    from triton.backends.nvidia import compiler as cb
    from triton.runtime import autotuner
except ImportError as e:
    raise ImportError(
        "Triton is required for transformer_engine.jax.triton_extensions. "
        "Install with: pip install triton\n"
        "If you don't need Triton, use transformer_engine.jax.cpp_extensions instead."
    ) from e


__all__ = ["triton_call_lowering", "get_triton_info"]

# Triton kernel cache (module-level, shared across all kernels)
_TRITON_KERNEL_CACHE = {}


def get_triton_info():
    """Get information about the installed Triton package.

    Returns:
        dict: Dictionary containing:
            - version (str): Triton version string (e.g., "3.1.0" or "3.0.0+45fff310c8")
            - is_pytorch_triton (bool): True if using real pytorch-triton from PyTorch's index
            - is_openai_triton (bool): True if using standard triton from OpenAI/PyPI
            - env_acknowledged (bool): True if NVTE_USE_PYTORCH_TRITON=1 is set
            - source (str): "pytorch" or "openai" indicating the package source

    Example:
        from transformer_engine.jax.triton_extensions import get_triton_info
        info = get_triton_info()
        print(f"Triton version: {info['version']} (from {info['source']})")
        if info['is_pytorch_triton']:
             print("Using pytorch-triton - compatible with both PyTorch and JAX")
    """
    val = os.environ.get("NVTE_USE_PYTORCH_TRITON", "0")
    try:
        env_acknowledged = bool(int(val))
    except ValueError as e:
        raise ValueError(
            f"NVTE_USE_PYTORCH_TRITON must be an integer (0 or 1), got: {val!r}"
        ) from e

    return {
        "version": _TRITON_VERSION,
        "is_pytorch_triton": _IS_PYTORCH_TRITON,
        "is_openai_triton": not _IS_PYTORCH_TRITON,
        "env_acknowledged": env_acknowledged and _IS_PYTORCH_TRITON,
        "source": "pytorch" if _IS_PYTORCH_TRITON else "openai",
    }


def get_triton_dtype(aval):
    """Convert JAX dtype to Triton type string.

    Args:
        aval: JAX ShapedArray

    Returns:
        Triton type string (e.g., "*fp32" for pointer, "i32" for scalar)
    """
    dtype_map = {
        jnp.dtype("bfloat16"): "bf16",
        jnp.dtype("float64"): "fp64",
        jnp.dtype("float32"): "fp32",
        jnp.dtype("float16"): "fp16",
        jnp.dtype("float8_e4m3fn"): "fp8e4nv",
        jnp.dtype("float8_e5m2"): "fp8e5",
        jnp.dtype("int64"): "i64",
        jnp.dtype("int32"): "i32",
        jnp.dtype("int16"): "i16",
        jnp.dtype("int8"): "i8",
        jnp.dtype("uint64"): "u64",
        jnp.dtype("uint32"): "u32",
        jnp.dtype("uint16"): "u16",
        jnp.dtype("uint8"): "u8",
        jnp.dtype("bool"): "i1",
    }

    assert isinstance(aval, core.ShapedArray), "aval must be a JAX ShapedArray"
    return f"*{dtype_map[aval.dtype]}"


def compile_triton(
    kernel_fn: Callable,
    signature: Mapping[str, str],
    constants: Mapping[str, Any],
    num_warps: int,
    num_stages: int,
    num_ctas: int,
    compute_capability: int,
    enable_fp_fusion: bool = False,
):
    """Compile a Triton kernel to PTX.

    Kernels are cached to avoid recompilation.

    Args:
        kernel_fn: Triton kernel function (decorated with @triton.jit)
        signature: Dict mapping arg names to types (e.g., {"x_ptr": "*fp32", "n": "i32"})
        constants: Dict of compile-time constants
        num_warps: Number of warps per block
        num_stages: Number of pipeline stages
        num_ctas: Number of CTAs (cooperative thread arrays)
        compute_capability: CUDA compute capability
        enable_fp_fusion: Enable FP fusion optimizations (default False for accuracy)

    Returns:
        TritonKernel object for JAX
    """
    # Create cache key
    cache_key = hashlib.md5(
        str(
            (
                kernel_fn.__name__,
                tuple(sorted(signature.items())),
                tuple(sorted(constants.items())),
                num_warps,
                num_stages,
                num_ctas,
                enable_fp_fusion,
                compute_capability,
            )
        ).encode()
    ).hexdigest()

    if cache_key in _TRITON_KERNEL_CACHE:
        return _TRITON_KERNEL_CACHE[cache_key]

    # Compile kernel
    cuda_option_kwargs = {}
    if version.parse(_TRITON_VERSION) < version.parse("3.6.0"):
        cuda_option_kwargs["cluster_dims"] = (1, 1, 1)
    options = cb.CUDAOptions(
        num_warps=num_warps,
        num_stages=num_stages,
        num_ctas=num_ctas,
        debug=False,
        enable_fp_fusion=enable_fp_fusion,
        **cuda_option_kwargs,
    )

    # Mark constants as constexpr in signature
    signature_with_constexpr = dict(signature)
    for const_name in constants.keys():
        if const_name in signature_with_constexpr:
            signature_with_constexpr[const_name] = "constexpr"

    src = tc.ASTSource(
        fn=kernel_fn,
        constexprs=constants,
        signature=signature_with_constexpr,
    )

    compiled = tc.compile(
        src,
        target=tc.GPUTarget("cuda", compute_capability, 32),
        options=options.__dict__,
    )

    # Create kernel object for JAX
    # From jax/jaxlib/gpu/triton_kernels.cc:
    if version.parse(jax.__version__) >= version.parse("0.8.2"):
        kernel = gpu_triton.TritonKernel(
            compiled.name,  # arg0: kernel_name (str)
            num_warps,  # arg1: num_warps (int)
            num_ctas,  # arg2: num_ctas (int)
            compiled.metadata.shared,  # arg3: shared_mem_bytes (int)
            compiled.asm["ptx"],  # arg4: ptx (str)
            "",  # arg5: ttir (str) - empty
            compute_capability,  # arg6: compute_capability (int)
        )
    else:
        kernel = gpu_triton.TritonKernel(
            compiled.name,
            num_warps,
            compiled.metadata.shared,
            compiled.asm["ptx"],
            "",  # ttir
            compute_capability,
            1,
            1,
            1,
        )

    _TRITON_KERNEL_CACHE[cache_key] = kernel
    return kernel


def triton_call_lowering(
    ctx,
    kernel_fn: Callable,
    *array_args,
    grid,
    input_output_aliases: Mapping[int, int] = None,
    constexprs: Mapping[str, Any] = None,
):
    """Helper for MLIR lowering that calls a Triton kernel.

    Use this in your primitive's lowering method to call Triton kernels.

    Args:
        ctx: MLIR lowering context
        kernel_fn: Triton kernel function
        *array_args: Input arrays (from ctx)
        grid: Grid dimensions. May be either:
            - an int or tuple (fixed grid for every config), or
            - a callable ``meta -> int|tuple`` (evaluated per autotune config).

            Use the callable form for autotuned kernels whose grid depends on
            ``BLOCK_SIZE`` (or any other autotuned constexpr); otherwise the
            launch grid will not match the autotuner-selected config and the
            kernel will either over-launch (waste) or under-cover. ``meta`` is
            the merged dict ``{**constexprs, **config.kwargs}`` for the chosen
            config — the same convention as jax-triton's ``triton_call``.
        input_output_aliases: Mapping of input to output aliases
        constexprs: Compile-time constants for the kernel. This includes both
                    tl.constexpr arguments AND scalar runtime arguments (like
                    num_tokens, strides) that are known at JAX trace time.

    Returns:
        MLIR lowering result

    Example:
        @staticmethod
        def lowering(ctx, x, *, block_size):
            from ..triton_extensions import triton_call_lowering
            n = ctx.avals_in[0].size

            def grid(meta):
                return (triton.cdiv(n, meta["BLOCK_SIZE"]),)

            return triton_call_lowering(
                ctx, my_kernel, x, grid=grid, constexprs={"n_elements": n},
            )
    """
    # Get compute capability using gpu_triton
    compute_capability = gpu_triton.get_compute_capability(0)  # device 0

    # Build signature dict: map arg names to types
    # Get arg names from kernel function
    if isinstance(kernel_fn, autotuner.Autotuner):
        arg_names = kernel_fn.fn.arg_names
    else:
        arg_names = kernel_fn.arg_names

    # Build signature for tensor arguments only (inputs + outputs)
    # Scalar arguments should be passed via constexprs and will be
    # specialized into the kernel at compile time
    all_avals = list(ctx.avals_in) + list(ctx.avals_out)
    constexpr_names = set(constexprs.keys()) if constexprs else set()
    tensor_arg_names = [n for n in arg_names if n not in constexpr_names]
    signature = {n: get_triton_dtype(a) for n, a in zip(tensor_arg_names, all_avals)}

    assert callable(grid) or isinstance(grid, tuple), (
        "Argument 'grid' must be a tuple or a callable but received: "
        f"type={type(grid)}, value={grid}"
    )

    # Normalize grid to 3D. When `grid` is a callable, defer evaluation until
    # we know the per-config meta (so each autotune config gets its own grid,
    # matching jax-triton's behavior).
    def _normalize_grid(grid_tuple):
        if isinstance(grid_tuple, int):
            return (grid_tuple, 1, 1)
        if len(grid_tuple) == 1:
            return (grid_tuple[0], 1, 1)
        if len(grid_tuple) == 2:
            return (grid_tuple[0], grid_tuple[1], 1)
        return tuple(grid_tuple[:3])

    grid_callable = grid if callable(grid) else None
    if grid_callable is None:
        grid_tuple = _normalize_grid(grid)
    else:
        grid_tuple = None  # evaluated per-config below

    # Default kernel launch parameters. These apply to non-autotuned kernels
    # and as a fallback when an autotuned config doesn't specify them. Values
    # match Triton's own `triton.Config` defaults (num_warps=4, num_stages=3,
    # num_ctas=1) and jax-triton's `get_or_create_triton_kernel`. Using a
    # larger default (e.g. num_warps=32) over-provisions threads per block,
    # which slashes SM occupancy on non-autotuned kernels — measured as an 8×
    # slowdown on `_make_chunk_sort_map_kernel` vs jax-triton.
    actual_kernel_fn = kernel_fn
    num_warps = 4
    num_stages = 3
    num_ctas = 1
    kernel_constexprs = constexprs if constexprs is not None else {}

    # Handle autotuned kernels - compile all configs.
    # On JAX < TRITON_AUTOTUNED_INPUT_OUTPUT_ALIAS_MIN_JAX_VERSION the save/restore
    # loop in TritonAutotunedKernelCall is buggy (jax-ml/jax#35218).  Fall back to a
    # single non-autotuned dispatch for compatibility.  Set
    # NVTE_JAX_ENFORCE_TRITON_AUTOTUNING=1 to raise an error instead, prompting the
    # user to upgrade JAX for improved performance.
    is_autotuned = isinstance(kernel_fn, autotuner.Autotuner)
    used_autotuned_launch = False
    if is_autotuned and not is_triton_autotuned_alias_safe():
        val = os.environ.get("NVTE_JAX_ENFORCE_TRITON_AUTOTUNING", "0")
        try:
            enforce = bool(int(val))
        except ValueError as e:
            raise ValueError(
                f"NVTE_JAX_ENFORCE_TRITON_AUTOTUNING must be an integer (0 or 1), got: {val!r}"
            ) from e
        if enforce:
            raise RuntimeError(
                "NVTE_JAX_ENFORCE_TRITON_AUTOTUNING=1 requires JAX >= "
                f"{TRITON_AUTOTUNED_INPUT_OUTPUT_ALIAS_MIN_JAX_VERSION} (stable) or a "
                "post-2026-03-17 nightly for safe Triton autotuning (jax-ml/jax#35218). "
                f"Current JAX version: {jax.__version__}. "
                "Upgrade: pip install --upgrade jax jaxlib"
            )
        # Compatibility fallback: disable autotuning on old JAX to avoid
        # CUDA_ERROR_INVALID_VALUE from the unfixed save/restore loop.
        is_autotuned = False

    if is_autotuned:
        # Compile all configs for runtime selection
        kernel_calls = []
        actual_kernel_fn = kernel_fn.fn

        for config in kernel_fn.configs:
            # Extract parameters from config
            config_num_warps = config.num_warps if config.num_warps is not None else num_warps
            config_num_stages = config.num_stages if config.num_stages is not None else num_stages
            config_num_ctas = config.num_ctas if config.num_ctas is not None else num_ctas

            # Config kwargs (e.g. BLOCK_SIZE) take priority over caller constexprs so that
            # each autotuning candidate actually compiles with its own BLOCK_SIZE rather than
            # having the caller-supplied grid BLOCK_SIZE override every config.
            config_constexprs = {**(constexprs if constexprs else {}), **config.kwargs}

            # Compile this config
            config_kernel = compile_triton(
                actual_kernel_fn,
                signature,
                config_constexprs,
                config_num_warps,
                config_num_stages,
                config_num_ctas,
                compute_capability,
                enable_fp_fusion=False,
            )

            # Create kernel call for this config
            config_params = []
            for _ in list(ctx.avals_in) + list(ctx.avals_out):
                config_params.append(gpu_triton.create_array_parameter(0, 16))

            # Per-config grid: evaluate `grid(meta)` if grid is a callable so
            # the launch shape matches this config's BLOCK_SIZE (etc.).
            if grid_callable is not None:
                config_grid = _normalize_grid(grid_callable(config_constexprs))
            else:
                config_grid = grid_tuple

            config_call = gpu_triton.TritonKernelCall(
                config_kernel,
                config_grid[0],
                config_grid[1],
                config_grid[2],
                config_params,
            )

            kernel_calls.append((config_call, str(config)))

        input_output_aliases_with_sizes = ()
        if input_output_aliases:
            # JAX version is guaranteed >= TRITON_AUTOTUNED_INPUT_OUTPUT_ALIAS_MIN_JAX_VERSION
            # here — verified by the upfront check that set is_autotuned.
            num_inputs = len(ctx.avals_in)
            aliases = []
            for input_idx, output_idx in input_output_aliases.items():
                aval = ctx.avals_in[input_idx]
                size_bytes = aval.size * jnp.dtype(aval.dtype).itemsize
                # AutotunedKernelCall expects buffer indices (inputs + outputs).
                buffer_output_idx = num_inputs + output_idx
                aliases.append((input_idx, buffer_output_idx, size_bytes))
            input_output_aliases_with_sizes = tuple(aliases)

        kernel_call = gpu_triton.TritonAutotunedKernelCall(
            f"{actual_kernel_fn.__name__}_autotuned",
            kernel_calls,
            input_output_aliases_with_sizes,
        )
        used_autotuned_launch = True

    else:
        # Regular kernel: compile single config.
        # If the kernel is an Autotuner but JAX is too old for safe autotuning, unwrap
        # it and use the first config's kwargs (user constexprs take priority via dict merge).
        if isinstance(kernel_fn, autotuner.Autotuner):
            actual_kernel_fn = kernel_fn.fn
            if kernel_fn.configs:
                first_cfg = kernel_fn.configs[0]
                # user constexprs override config kwargs (so stride / size scalars win)
                kernel_constexprs = {**first_cfg.kwargs, **(constexprs or {})}
                num_warps = first_cfg.num_warps if first_cfg.num_warps is not None else num_warps
                num_stages = (
                    first_cfg.num_stages if first_cfg.num_stages is not None else num_stages
                )
                num_ctas = first_cfg.num_ctas if first_cfg.num_ctas is not None else num_ctas

        kernel = compile_triton(
            actual_kernel_fn,
            signature,
            kernel_constexprs,
            num_warps,
            num_stages,
            num_ctas,
            compute_capability,
            enable_fp_fusion=False,
        )

        kernel_params = []
        for _ in list(ctx.avals_in) + list(ctx.avals_out):
            kernel_params.append(gpu_triton.create_array_parameter(0, 16))

        # Non-autotuned dispatch: evaluate `grid(meta)` once with the merged
        # constexprs (which already reflect the single config we'll launch).
        if grid_callable is not None:
            single_grid = _normalize_grid(grid_callable(kernel_constexprs))
        else:
            single_grid = grid_tuple

        kernel_call = gpu_triton.TritonKernelCall(
            kernel,
            single_grid[0],
            single_grid[1],
            single_grid[2],
            kernel_params,
        )

    serialized_metadata = b""
    call_proto = kernel_call.to_proto(actual_kernel_fn.__name__, serialized_metadata)

    if input_output_aliases:
        ffi_operand_output_aliases = input_output_aliases
    else:
        ffi_operand_output_aliases = None

    compressed_call_proto = zlib.compress(call_proto)
    if not used_autotuned_launch and jax_version_meet_requirement(
        TRITON_EXTENSION_CUDA_GRAPH_MIN_JAX_VERSION
    ):
        rule = jax.ffi.ffi_lowering(
            "triton_kernel_call_ffi",
            backend_config={"opaque": ir.StringAttr.get(compressed_call_proto)},
            operand_output_aliases=ffi_operand_output_aliases,
        )
    else:
        rule = jax.ffi.ffi_lowering(
            "triton_kernel_call",  # Custom call target registered in gpu_triton.py
            api_version=2,
            backend_config=compressed_call_proto,
            operand_output_aliases=ffi_operand_output_aliases,
        )

    return rule(ctx, *array_args)
