# Copyright 2023 The JAX Authors.
#
# 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
#
#     https://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.

"""Module for lowering JAX primitives to Triton IR."""

from __future__ import annotations

from collections.abc import Callable, Hashable, Sequence
import dataclasses
import functools
import math
from typing import Any, TypeVar

import jax
from jax import lax
from jax import tree_util
from jax._src import ad_checkpoint
from jax._src import ad_util
from jax._src import api_util
from jax._src import config
from jax._src import core as jax_core
from jax._src import custom_derivatives
from jax._src import debugging
from jax._src import linear_util as lu
from jax._src import literals
from jax._src import pjit
from jax._src import source_info_util
from jax._src import state
from jax._src import util
from jax._src.interpreters import mlir
from jax._src.interpreters import partial_eval as pe
from jax._src.lib.mlir import ir
from jax._src.lib.mlir.dialects import arith as arith_dialect
from jax._src.lib.mlir.dialects import math as math_dialect
from jax._src.lib.mlir.dialects import scf as scf_dialect
from jax._src.lib.triton import dialect as tt_dialect
from jax._src.pallas import core as pallas_core
from jax._src.pallas import primitives
from jax._src.pallas import utils as pallas_utils
from jax._src.state import indexing
from jax._src.state import primitives as sp
from jax._src.util import foreach
from jax._src.util import split_list
import jax.numpy as jnp
import numpy as np

_T = TypeVar("_T")

map, unsafe_map = util.safe_map, map
zip, unsafe_zip = util.safe_zip, zip

NDIndexer = indexing.NDIndexer
GridMapping = pallas_core.GridMapping
BlockMapping = pallas_core.BlockMapping
Blocked = pallas_core.Blocked


# # General lowering logic
@dataclasses.dataclass
class ModuleContext:
  name: str
  grid_mapping: GridMapping
  program_ids: Sequence[ir.Value]
  traceback_caches: mlir.TracebackCaches = dataclasses.field(repr=False)
  platform: str
  compute_capability: int | None


@dataclasses.dataclass
class BlockInfo:
  full_shape_dtype: jax_core.ShapedArray
  start_indices: Sequence[Any]
  start_indices_alignment: Sequence[int]
  block_shape: tuple[int | pallas_core.Squeezed, ...]


@dataclasses.dataclass
class LoweringRuleContext:
  context: ModuleContext
  avals_in: Sequence[jax_core.ShapedArray]
  avals_out: Sequence[jax_core.ShapedArray]
  block_infos: Sequence[BlockInfo | None]

  replace = dataclasses.replace


@dataclasses.dataclass
class LoweringResult:
  """Keeps python objects alive."""

  module: ir.Module
  grid: tuple[int, ...]


class LoweringError(Exception):
  pass


def _eval_index_map(
    ctx: ModuleContext, idx, block_mapping: BlockMapping
):
  block_indices = lower_jaxpr_to_triton_ir(
      ctx, block_mapping.index_map_jaxpr.jaxpr, None, *idx
  )
  block_indices = tuple(
      _ensure_ir_value(i, jax_core.ShapedArray((), jnp.int32))
      for i in block_indices
  )
  block_indices = tree_util.tree_unflatten(
      block_mapping.index_map_out_tree, block_indices)
  if block_mapping.pipeline_mode is not None:
    raise NotImplementedError(
        "Pipeline mode is not supported in Triton lowering."
    )
  if any(
      isinstance(b, pallas_core.Element) and b.padding != (0, 0)
      for b in block_mapping.block_shape
  ):
    raise NotImplementedError(
        "Unblocked indexing with padding is not supported in Triton lowering."
    )
  def _get_start_index(i, b):
    match b:
      case pallas_core.Squeezed() | pallas_core.Element():
        return i
      case pallas_core.Blocked():
        return _mul(i, _ir_constant(b.block_size, i.type))
      case _:
        raise ValueError(f"Unsupported block dim type: {type(b)}")
  return tuple(
      _get_start_index(i, b) for i, b in
      zip(block_indices, block_mapping.block_shape)
  )


def _get_index_alignment(block_mapping: BlockMapping) -> tuple[int, ...]:
  def _get_bdim_alignment(b: pallas_core.BlockDim):
    match b:
      case pallas_core.Squeezed() | pallas_core.Element():
        return 1
      case pallas_core.Blocked():
        return b.block_size
  return tuple(_get_bdim_alignment(b) for b in block_mapping.block_shape)


def _bcast_to(a: ir.Value, shape: Sequence[int]) -> ir.Value:
  if not isinstance(a.type, ir.RankedTensorType):
    if not shape:
      return a
    return tt_dialect.splat(ir.RankedTensorType.get(shape, a.type), a)
  else:
    a_type = ir.RankedTensorType(a.type)
    if a_type.shape == [*shape]:
      return a
    for _ in range(a_type.rank, len(shape)):
      a = _expand_dims(a, 0)
    if a_type.rank != len(shape) or not all(
        a_type.shape[i] in (dim, 1) for i, dim in enumerate(shape)
    ):
      raise ValueError(f"Cannot broadcast from {a_type.shape} to {[*shape]}")
    return tt_dialect.broadcast(
        ir.RankedTensorType.get(shape, a_type.element_type, a_type.encoding), a
    )


def _bcast(
    x: ir.Value,
    y: ir.Value,
    x_aval: jax_core.ShapedArray,
    y_aval: jax_core.ShapedArray,
    out_aval: jax_core.ShapedArray,
) -> tuple[ir.Value, ir.Value]:
  if isinstance(
      x, (np.ndarray, np.number, int, float, literals.TypedNdArray)
  ):
    x_dtype = x_aval.dtype
    if x_aval.weak_type:
      x_dtype = y_aval.dtype
    x = _ir_constant(x, _dtype_to_ir_type(x_dtype))
  if isinstance(
      y, (np.ndarray, np.number, int, float, literals.TypedNdArray)
  ):
    y_dtype = y_aval.dtype
    if y_aval.weak_type:
      y_dtype = x_aval.dtype
    y = _ir_constant(y, _dtype_to_ir_type(y_dtype))
  if x_aval.shape != out_aval.shape:
    x = _bcast_to(x, out_aval.shape)
  if y_aval.shape != out_aval.shape:
    y = _bcast_to(y, out_aval.shape)
  return x, y


triton_lowering_rules = {}


def register_lowering(primitive: jax_core.Primitive) -> Callable[[_T], _T]:
  def wrapper(fn):
    triton_lowering_rules[primitive] = fn
    return fn
  return wrapper


def _process_grid_to_3d_grid(grid_mapping: GridMapping):
  launch_grid: list[int] = []
  launch_grid_to_pallas_grid = []

  # Preserve grid order provided to pallas_call
  for i, s in enumerate(grid_mapping.grid):
    assert isinstance(s, int)
    if i not in grid_mapping.vmapped_dims:
      launch_grid.append(s)
      launch_grid_to_pallas_grid.append(i)

  # For mapped dims, iterate from inner to outer. This follows the pallas_call
  # batching rule that prepends the vmapped dimension.
  for dim in reversed(grid_mapping.vmapped_dims):
    s = grid_mapping.grid[dim]
    assert isinstance(s, int)
    launch_grid.append(s)
    launch_grid_to_pallas_grid.append(dim)

  num_collapse = len(launch_grid[:-2])

  cuda_yz_limit = 2**16 - 1

  # Check Z and then Y launch dims to make sure they're within CUDA bounds
  if (num_collapse + 1 < len(launch_grid) and
      launch_grid[num_collapse + 1] > cuda_yz_limit):
    num_collapse += 2
  elif (num_collapse < len(launch_grid) and
        launch_grid[num_collapse] > cuda_yz_limit):
    num_collapse += 1

  collapse_dims = launch_grid[:num_collapse]
  prog_id_dims = launch_grid[num_collapse:]

  if len(collapse_dims) == 0:
    prog_ids: list[ir.Value | None] = [None] * len(prog_id_dims)
    for i in range(len(prog_id_dims)):
      prog_ids[launch_grid_to_pallas_grid[i]] = _program_id(i, prog_id_dims)

    return prog_id_dims, prog_ids

  new_grid = [math.prod(collapse_dims), *prog_id_dims]

  assert new_grid[0] < 2**31 - 1, \
          "Cannot fix pallas kernel launch grid within CUDA limits"

  out_indices: list[ir.Value | None] = [None] * len(grid_mapping.grid)

  grid0 = _program_id(0, new_grid)
  for i, s in enumerate(collapse_dims):
    out_idx = launch_grid_to_pallas_grid[i]
    s = _i32_constant(s)
    out_indices[out_idx] = _mod(grid0, s, signed=False)
    grid0 = _floordiv(grid0, s, signed=False)

  for i in range(len(prog_id_dims)):
    out_idx = launch_grid_to_pallas_grid[num_collapse + i]
    out_indices[out_idx] = _program_id(i + 1, new_grid)

  assert len(out_indices) == len(grid_mapping.grid)
  return new_grid, out_indices


def _new_ir_context() -> ir.Context:
  ctx = mlir.JaxIrContext()
  ctx.append_dialect_registry(mlir.upstream_dialects)
  tt_dialect.register_dialect(ctx)
  ctx.load_all_available_dialects()
  return ctx

# Many Trion operations require that their inputs and outputs have sizes that
# are a power of 2 (they are defined to have TensorSizeTrait that enforces
# this). This check is only needed to obtain a nicer error message; the
# Triton lowering will fail anyway but it will crash with a C++ exception.
# We currently apply this check only to load/store operations.
def _check_tensor_size(shape: tuple[int | pallas_core.Squeezed, ...]):
  size = math.prod(1 if isinstance(d, pallas_core.Squeezed) else d
                   for d in shape)
  power_of_2 = (size & (size - 1)) == 0
  if not power_of_2:
    raise ValueError(
        "The Pallas Triton lowering currently requires that all "
        "operations have array arguments and results whose size "
        "is a power of 2. Encountered an array of "
        f"shape {shape}")


def lower_jaxpr_to_triton_module(
    jaxpr: jax_core.Jaxpr,
    grid_mapping: GridMapping,
    platform: str,
    compute_capability: int | None,
) -> LoweringResult:
  debug_info = jaxpr.debug_info
  if grid_mapping.num_dynamic_grid_bounds:
    raise NotImplementedError(
        "dynamic grid bounds not supported in the Triton backend"
    )
  if grid_mapping.num_index_operands:
    raise NotImplementedError(
        "scalar prefetch not implemented in the Triton backend"
    )
  if jaxpr.invars[grid_mapping.slice_scratch_ops]:
    raise NotImplementedError(
        "scratch memory not implemented in the Triton backend"
    )
  with _new_ir_context(), ir.Location.unknown():
    module = ir.Module.create()
    attrs = module.operation.attributes
    module_name = mlir.sanitize_name(debug_info.func_name)
    attrs["sym_name"] = ir.StringAttr.get(module_name)
    param_types = [
        # pyrefly: ignore[missing-attribute]
        tt_dialect.PointerType.get(_dtype_to_ir_type(var.aval.dtype), 1)
        for var in jaxpr.invars
    ]
    assert len(jaxpr.outvars) == 0
    fn_type = ir.FunctionType.get(param_types, [])
    fn = tt_dialect.FuncOp(
        module_name,
        ir.TypeAttr.get(fn_type),
        sym_visibility="public",
        res_attrs=ir.DictAttr.get(dict(noinline=ir.BoolAttr.get(False))),
        ip=ir.InsertionPoint.at_block_begin(module.body),
    )
    fn.arg_attrs = ir.ArrayAttr.get(
        [ir.DictAttr.get({"tt.divisibility": mlir.i32_attr(32)})]
        * len(param_types)
    )
    fn.body.blocks.append(*fn_type.inputs)
    [entry] = fn.body.blocks
    with ir.InsertionPoint(entry):
      new_grid, program_ids = _process_grid_to_3d_grid(grid_mapping)
      local_program_ids = [
          pid
          for i, pid in enumerate(program_ids)
          if i not in grid_mapping.vmapped_dims
      ]
      ctx = ModuleContext(
          mlir.sanitize_name(debug_info.func_name),
          grid_mapping,
          local_program_ids,
          mlir.TracebackCaches(),
          platform,
          compute_capability,
      )
      block_infos = [
          BlockInfo(
              block_mapping.array_aval,
              _eval_index_map(ctx, program_ids, block_mapping),
              _get_index_alignment(block_mapping),
              tuple(pallas_core.squeezed if isinstance(b, pallas_core.Squeezed)
                    else pallas_core._get_block_dim_size(b)
                    for b in block_mapping.block_shape),
          )
          for block_mapping in grid_mapping.block_mappings
      ]
      () = lower_jaxpr_to_triton_ir(ctx, jaxpr, block_infos, *entry.arguments)
      tt_dialect.return_([])
    return LoweringResult(module, new_grid)


def lower_jaxpr_to_triton_ir(
    ctx: ModuleContext,
    jaxpr: jax_core.Jaxpr,
    block_infos: Sequence[BlockInfo | None] | None,
    *args,
) -> Sequence[Any]:
  env = {}
  block_info_env = {}

  def read_env(atom: jax_core.Atom):
    return atom.val if isinstance(atom, jax_core.Literal) else env[atom]

  def read_block_info_env(atom: jax_core.Atom):
    if isinstance(atom, jax_core.Literal):
      return None
    return block_info_env.get(atom)

  def write_env(var: jax_core.Var, val):
    env[var] = val

  if block_infos is not None:
    for invar, block_info in zip(jaxpr.invars, block_infos):
      if block_info is not None:
        block_info_env[invar] = block_info

  foreach(write_env, jaxpr.invars, args)

  for eqn in jaxpr.eqns:
    invals = map(read_env, eqn.invars)
    if eqn.primitive not in triton_lowering_rules:
      raise NotImplementedError(
          "Unimplemented primitive in Pallas GPU lowering: "
          f"{eqn.primitive.name}. "
          "Please file an issue on https://github.com/jax-ml/jax/issues.")
    rule = triton_lowering_rules[eqn.primitive]
    avals_in = [v.aval for v in eqn.invars]
    avals_out = [v.aval for v in eqn.outvars]
    eqn_block_infos = map(read_block_info_env, eqn.invars)
    loc = mlir.source_info_to_location(
        ctx, eqn.primitive, eqn.source_info.name_stack,
        eqn.source_info.traceback)
    rule_ctx = LoweringRuleContext(ctx, avals_in, avals_out, eqn_block_infos)  # pyrefly: ignore[bad-argument-type]  # pyrefly#2385
    try:
      with source_info_util.user_context(eqn.source_info.traceback), loc:
        outvals: Any = rule(rule_ctx, *invals, **eqn.params)
    except LoweringError:
      raise  # We only add the extra info to the innermost exception.
    except Exception as e:
      if not config.jax_pallas_verbose_errors.value:
        raise
      inval_types = map(lambda t: getattr(t, "type", None), invals)
      raise LoweringError(
          f"Exception while lowering eqn:\n  {eqn}\nWith context:\n "
          f" {rule_ctx}\nWith inval types={inval_types}\nIn jaxpr:\n{jaxpr}\n"
          f"msg={e}"
      ) from e
    if eqn.primitive.multiple_results:
      if len(eqn.outvars) != len(outvals):
        raise ValueError(
            f"Primitive {eqn.primitive.name} has multiple results, but "
            f"{len(eqn.outvars)} outvars do not match {len(outvals)} outvals."
        )
      foreach(write_env, eqn.outvars, outvals)
    else:
      write_env(eqn.outvars[0], outvals)

  return map(read_env, jaxpr.outvars)


def lower_fun(
    fun: Callable[..., Any], *, multiple_results: bool
) -> Callable[..., Any]:
  fn = fun if multiple_results else lambda *args, **kw: (fun(*args, **kw),)

  def f_lowered(ctx: LoweringRuleContext, *args, **params):
    wrapped_fun = lu.wrap_init(
        fn, params,
        debug_info=api_util.debug_info("pallas triton lower_fun", fun,
                                       args, params))
    jaxpr, _, consts = pe.trace_to_jaxpr_dynamic(wrapped_fun, ctx.avals_in)
    jaxpr = jax_core.ClosedJaxpr(jaxpr, consts)
    out = _closed_call_lowering_rule(ctx, *args, call_jaxpr=jaxpr)
    return out if multiple_results else out[0]

  return f_lowered


# # Primitive lowering rules
# ## Programming model primitives


def _program_id(axis: int, launch_grid: Sequence[int]) -> ir.Value:
  if axis not in range(3):
    raise ValueError(f"axis must be in [0, 3), but got: {axis}")
  if launch_grid[axis] == 1:
    return _i32_constant(0)
  return tt_dialect.get_program_id(axis)


@register_lowering(primitives.program_id_p)
def _program_id_lowering_rule(ctx: LoweringRuleContext, *, axis):
  return ctx.context.program_ids[axis]


@register_lowering(primitives.num_programs_p)
def _num_programs_lowering_rule(ctx: LoweringRuleContext, *, axis):
  if axis not in range(3):
    raise ValueError(f"axis must be in [0, 3), but got: {axis}")
  return tt_dialect.get_num_programs(axis)

def _atomic_rmw(
    op: tt_dialect.RMWOp,
    ptr: ir.Value,
    val: ir.Value,
    mask: ir.Value | None = None,
    semantic: tt_dialect.MemSemantic = tt_dialect.MemSemantic.ACQUIRE_RELEASE,
    sync_scope: tt_dialect.MemSyncScope = tt_dialect.MemSyncScope.GPU,
) -> ir.Value:
  if isinstance(ptr.type, ir.RankedTensorType):
    ptr_type = ir.RankedTensorType(ptr.type)
    element_type = tt_dialect.PointerType(ptr_type.element_type)
    result_type = ir.RankedTensorType.get(
        ptr_type.shape, element_type.pointee_type, ptr_type.encoding
    )
  else:
    result_type = tt_dialect.PointerType(ptr.type).pointee_type
  return tt_dialect.atomic_rmw(
      result_type, op, ptr, val, mask=mask, sem=semantic, scope=sync_scope
  )

def _associative_scan_lowering(body, ctx: LoweringRuleContext, args, axes):
  flat_args = tree_util.tree_leaves(args)
  (axis,) = axes
  dtype = ctx.avals_in[0].dtype
  in_avals = [
      jax_core.ShapedArray((), dtype=dtype),
      jax_core.ShapedArray((), dtype=dtype),
  ]
  in_tree = tree_util.tree_structure((args, args))
  flat_fun, out_tree_thunk = api_util.flatten_fun_nokwargs(
      lu.wrap_init(
          body,
          debug_info=api_util.debug_info("pallas triton associative_scan",
                                         body, (args, args), {})),
      in_tree
  )
  combine_jaxpr, _, consts = pe.trace_to_jaxpr_dynamic(
      flat_fun, in_avals
  )
  out_tree = out_tree_thunk()
  del out_tree  # Not needed
  if consts:
    raise NotImplementedError("Associative scan with constants not supported.")
  element_types = [_element_type(arg.type) for arg in flat_args]
  scan_op = tt_dialect.ScanOp(flat_args, axis)
  param_types = element_types * 2
  entry = scan_op.regions[0].blocks.append(*param_types)
  with ir.InsertionPoint.at_block_begin(entry):
    results = lower_jaxpr_to_triton_ir(
        ctx.context, combine_jaxpr, None, *entry.arguments
    )
    tt_dialect.scan_return(results)
  scan_op.verify()
  return list(scan_op.result)


@register_lowering(lax.cumsum_p)
def _cumsum_lowering_rule(
    ctx: LoweringRuleContext, x, *, axis: int, reverse: bool
):
  if reverse:
    raise NotImplementedError("Reverse cumsum is not supported.")
  return _associative_scan_lowering(jnp.add, ctx, x, (axis,))[0]


@register_lowering(lax.not_p)
def _not_lowering_rule(ctx: LoweringRuleContext, x):
  [x_aval] = ctx.avals_in
  return arith_dialect.xori(x, _full(x.type, ~x_aval.dtype.type(0)))


@dataclasses.dataclass(frozen=True)
class _Extern:
  arg_types: Sequence[jax.typing.DTypeLike]
  symbol: str
  result_type: str

  def matches(self, avals: Sequence[jax_core.ShapedArray]) -> bool:
    if len(avals) != len(self.arg_types):
      return False
    return all(
        aval.dtype == jnp.dtype(arg_type)
        or (aval.weak_type and aval.dtype.kind == jnp.dtype(arg_type).kind)
        for aval, arg_type in zip(avals, self.arg_types)
    )

  def lower(self, ctx: LoweringRuleContext, *args: ir.Value):
    [out_aval] = ctx.avals_out
    bcast_args = []
    for aval, arg, arg_type in zip(ctx.avals_in, args, self.arg_types):
      bcast_arg = _bcast_to(_ensure_ir_value(arg, aval), out_aval.shape)
      if aval.weak_type and aval.dtype != jnp.dtype(arg_type):
        bcast_arg = _cast(bcast_arg, aval.dtype, jnp.dtype(arg_type))
      bcast_args.append(bcast_arg)

    result_type = _dtype_to_ir_type(jnp.dtype(self.result_type))
    if out_aval.shape:
      result_type = ir.RankedTensorType.get(out_aval.shape, result_type)
    return tt_dialect.extern_elementwise(
        result_type,
        bcast_args,
        libname="",
        libpath="",
        symbol=self.symbol,
        pure=True,
    )


@dataclasses.dataclass(frozen=True)
class _Fallback:
  arg_classes: Sequence[jax.typing.DTypeLike]
  op: Callable[..., ir.Value]

  def matches(self, avals: Sequence[jax_core.ShapedArray]) -> bool:
    if len(avals) != len(self.arg_classes):
      return False
    return all(
        jnp.issubdtype(aval.dtype, arg_class)
        for aval, arg_class in zip(avals, self.arg_classes)
    )

  def lower(self, ctx: LoweringRuleContext, *args: ir.Value):
    [out_aval] = ctx.avals_out
    bcast_args = []
    for aval, arg in zip(ctx.avals_in, args):
      bcast_args.append(_bcast_to(_ensure_ir_value(arg, aval), out_aval.shape))
    return self.op(*args)


def _make_dispatch_table(
    name: str, **tables: Sequence[_Extern | _Fallback]
) -> Callable[..., ir.Value]:

  def inner(
      ctx: LoweringRuleContext, *args: ir.Value, **_
  ) -> ir.Value:
    table = tables[ctx.context.platform]
    h = next((e for e in table if e.matches(ctx.avals_in)), None)
    if h is None:
      arg_aval_dtypes = tuple(aval.dtype for aval in ctx.avals_in)
      raise NotImplementedError(
          f"unsupported types for {name}: {arg_aval_dtypes}"
      )
    return h.lower(ctx, *args)

  return inner


abs_dispatch_table = _make_dispatch_table(
    "abs",
    cuda=[
        _Extern([jnp.int32], "__nv_abs", jnp.int32),
        _Extern([jnp.int64], "__nv_llabs", jnp.int64),
        _Extern([jnp.float32], "__nv_fabsf", jnp.float32),
        _Extern([jnp.float64], "__nv_fabs", jnp.float64),
        _Fallback([jnp.integer], math_dialect.absi),
        _Fallback([jnp.floating], math_dialect.absf),
    ],
    rocm=[
        _Extern([jnp.float32], "__ocml_fabs_f32", jnp.float32),
        _Extern([jnp.float64], "__ocml_fabs_f64", jnp.float64),
        _Fallback([jnp.integer], math_dialect.absi),
        _Fallback([jnp.floating], math_dialect.absf),
    ],
)

ceil_dispatch_table = _make_dispatch_table(
    "ceil",
    cuda=[
        _Extern([jnp.float32], "__nv_ceilf", jnp.float32),
        _Extern([jnp.float64], "__nv_ceil", jnp.float64),
        _Fallback([jnp.floating], math_dialect.ceil),
    ],
    rocm=[
        _Extern([jnp.float32], "__ocml_ceil_f32", jnp.float32),
        _Extern([jnp.float64], "__ocml_ceil_f64", jnp.float64),
        _Fallback([jnp.floating], math_dialect.ceil),
    ],
)

floor_dispatch_table = _make_dispatch_table(
    "floor",
    cuda=[
        _Extern([jnp.float32], "__nv_floorf", jnp.float32),
        _Extern([jnp.float64], "__nv_floor", jnp.float64),
        _Fallback([jnp.floating], math_dialect.floor),
    ],
    rocm=[
        _Extern([jnp.float32], "__ocml_floor_f32", jnp.float32),
        _Extern([jnp.float64], "__ocml_floor_f64", jnp.float64),
        _Fallback([jnp.floating], math_dialect.floor),
    ],
)

exp_dispatch_table = _make_dispatch_table(
    "exp",
    cuda=[
        _Extern([jnp.float32], "__nv_expf", jnp.float32),
        _Extern([jnp.float64], "__nv_exp", jnp.float64),
        _Fallback([jnp.floating], math_dialect.exp),
    ],
    rocm=[
        _Extern([jnp.float32], "__ocml_exp_f32", jnp.float32),
        _Extern([jnp.float64], "__ocml_exp_f64", jnp.float64),
        _Fallback([jnp.floating], math_dialect.exp),
    ],
)

exp2_dispatch_table = _make_dispatch_table(
    "exp2",
    cuda=[
        _Extern([jnp.float32], "__nv_exp2f", jnp.float32),
        _Extern([jnp.float64], "__nv_exp2", jnp.float64),
        _Fallback([jnp.floating], math_dialect.exp2),
    ],
    rocm=[
        _Extern([jnp.float32], "__ocml_exp2_f32", jnp.float32),
        _Extern([jnp.float64], "__ocml_exp2_f64", jnp.float64),
        _Fallback([jnp.floating], math_dialect.exp2),
    ],
)

expm1_dispatch_table = _make_dispatch_table(
    "expm1",
    cuda=[
        _Extern([jnp.float32], "__nv_expm1f", jnp.float32),
        _Extern([jnp.float64], "__nv_expm1", jnp.float64),
        _Fallback([jnp.floating], math_dialect.expm1),
    ],
    rocm=[
        _Extern([jnp.float32], "__ocml_expm1_f32", jnp.float32),
        _Extern([jnp.float64], "__ocml_expm1_f64", jnp.float64),
        _Fallback([jnp.floating], math_dialect.expm1),
    ],
)

log_dispatch_table = _make_dispatch_table(
    "log",
    cuda=[
        _Extern([jnp.float32], "__nv_logf", jnp.float32),
        _Extern([jnp.float64], "__nv_log", jnp.float64),
        _Fallback([jnp.floating], math_dialect.log),
    ],
    rocm=[
        _Extern([jnp.float32], "__ocml_log_f32", jnp.float32),
        _Extern([jnp.float64], "__ocml_log_f64", jnp.float64),
        _Fallback([jnp.floating], math_dialect.log),
    ],
)

log1p_dispatch_table = _make_dispatch_table(
    "log1p",
    cuda=[
        _Extern([jnp.float32], "__nv_log1pf", jnp.float32),
        _Extern([jnp.float64], "__nv_log1p", jnp.float64),
        _Fallback([jnp.floating], math_dialect.log1p),
    ],
    rocm=[
        _Extern([jnp.float32], "__ocml_log1p_f32", jnp.float32),
        _Extern([jnp.float64], "__ocml_log1p_f64", jnp.float64),
        _Fallback([jnp.floating], math_dialect.log1p),
    ],
)

sqrt_dispatch_table = _make_dispatch_table(
    "sqrt",
    cuda=[
        _Extern([jnp.float32], "__nv_sqrtf", jnp.float32),
        _Extern([jnp.float64], "__nv_sqrt", jnp.float64),
        _Fallback([jnp.floating], math_dialect.sqrt),
    ],
    rocm=[
        _Extern([jnp.float32], "__ocml_sqrt_f32", jnp.float32),
        _Extern([jnp.float64], "__ocml_sqrt_f64", jnp.float64),
        _Fallback([jnp.floating], math_dialect.sqrt),
    ],
)

pow_dispatch_table = _make_dispatch_table(
    "pow",
    cuda=[
        _Extern([jnp.float32, jnp.int32], "__nv_powif", jnp.float32),
        _Extern([jnp.float64, jnp.int32], "__nv_powi", jnp.float64),
        _Fallback(
            [jnp.floating, jnp.integer],
            math_dialect.fpowi
        ),
        _Extern([jnp.float32, jnp.float32], "__nv_powf", jnp.float32),
        _Extern([jnp.float64, jnp.float64], "__nv_pow", jnp.float64),
        _Fallback(
            [jnp.floating, jnp.floating],
            math_dialect.powf
        ),
    ],
    rocm=[
        _Extern([jnp.float32, jnp.int32], "__ocml_pown_f32", jnp.float32),
        _Extern([jnp.float64, jnp.int32], "__ocml_pown_f64", jnp.float64),
        _Fallback(
            [jnp.floating, jnp.integer],
            math_dialect.fpowi
        ),
        _Extern([jnp.float32, jnp.float32], "__ocml_pow_f32", jnp.float32),
        _Extern([jnp.float64, jnp.float64], "__ocml_pow_f64", jnp.float64),
        _Fallback(
            [jnp.floating, jnp.floating],
            math_dialect.powf
        ),
    ],
)

cbrt_dispatch_table = _make_dispatch_table(
    "cbrt",
    cuda=[
        _Extern([jnp.float32], "__nv_cbrtf", jnp.float32),
        _Extern([jnp.float64], "__nv_cbrt", jnp.float64),
        _Fallback([jnp.floating], math_dialect.cbrt),
    ],
    rocm=[
        _Extern([jnp.float32], "__ocml_cbrt_f32", jnp.float32),
        _Extern([jnp.float64], "__ocml_cbrt_f64", jnp.float64),
        _Fallback([jnp.floating], math_dialect.cbrt),
    ],
)

rsqrt_dispatch_table = _make_dispatch_table(
    "rsqrt",
    cuda=[
        _Extern([jnp.float32], "__nv_rsqrtf", jnp.float32),
        _Extern([jnp.float64], "__nv_rsqrt", jnp.float64),
        _Fallback([jnp.floating], math_dialect.rsqrt),
    ],
    rocm=[
        _Extern([jnp.float32], "__ocml_rsqrt_f32", jnp.float32),
        _Extern([jnp.float64], "__ocml_rsqrt_f64", jnp.float64),
        _Fallback([jnp.floating], math_dialect.rsqrt),
    ],
)

sin_dispatch_table = _make_dispatch_table(
    "sin",
    cuda=[
        _Extern([jnp.float32], "__nv_sinf", jnp.float32),
        _Extern([jnp.float64], "__nv_sin", jnp.float64),
        _Fallback([jnp.floating], math_dialect.sin),
    ],
    rocm=[
        _Extern([jnp.float32], "__ocml_sin_f32", jnp.float32),
        _Extern([jnp.float64], "__ocml_sin_f64", jnp.float64),
        _Fallback([jnp.floating], math_dialect.sin),
    ],
)

cos_dispatch_table = _make_dispatch_table(
    "cos",
    cuda=[
        _Extern([jnp.float32], "__nv_cosf", jnp.float32),
        _Extern([jnp.float64], "__nv_cos", jnp.float64),
        _Fallback([jnp.floating], math_dialect.cos),
    ],
    rocm=[
        _Extern([jnp.float32], "__ocml_cos_f32", jnp.float32),
        _Extern([jnp.float64], "__ocml_cos_f64", jnp.float64),
        _Fallback([jnp.floating], math_dialect.cos),
    ],
)

tan_dispatch_table = _make_dispatch_table(
    "tan",
    cuda=[
        _Extern([jnp.float32], "__nv_tanf", jnp.float32),
        _Extern([jnp.float64], "__nv_tan", jnp.float64),
        _Fallback([jnp.floating], math_dialect.tan),
    ],
    rocm=[
        _Extern([jnp.float32], "__ocml_tan_f32", jnp.float32),
        _Extern([jnp.float64], "__ocml_tan_f64", jnp.float64),
        _Fallback([jnp.floating], math_dialect.tan),
    ],
)

asin_dispatch_table = _make_dispatch_table(
    "asin",
    cuda=[
        _Extern([jnp.float32], "__nv_asinf", jnp.float32),
        _Extern([jnp.float64], "__nv_asin", jnp.float64),
        _Fallback([jnp.floating], math_dialect.asin),
    ],
    rocm=[
        _Extern([jnp.float32], "__ocml_asin_f32", jnp.float32),
        _Extern([jnp.float64], "__ocml_asin_f64", jnp.float64),
        _Fallback([jnp.floating], math_dialect.asin),
    ],
)

acos_dispatch_table = _make_dispatch_table(
    "acos",
    cuda=[
        _Extern([jnp.float32], "__nv_acosf", jnp.float32),
        _Extern([jnp.float64], "__nv_acos", jnp.float64),
        _Fallback([jnp.floating], math_dialect.acos),
    ],
    rocm=[
        _Extern([jnp.float32], "__ocml_acos_f32", jnp.float32),
        _Extern([jnp.float64], "__ocml_acos_f64", jnp.float64),
        _Fallback([jnp.floating], math_dialect.acos),
    ],
)

atan_dispatch_table = _make_dispatch_table(
    "atan",
    cuda=[
        _Extern([jnp.float32], "__nv_atanf", jnp.float32),
        _Extern([jnp.float64], "__nv_atan", jnp.float64),
        _Fallback([jnp.floating], math_dialect.atan),
    ],
    rocm=[
        _Extern([jnp.float32], "__ocml_atan_f32", jnp.float32),
        _Extern([jnp.float64], "__ocml_atan_f64", jnp.float64),
        _Fallback([jnp.floating], math_dialect.atan),
    ],
)

atan2_dispatch_table = _make_dispatch_table(
    "atan2",
    cuda=[
        _Extern([jnp.float32, jnp.float32], "__nv_atan2f", jnp.float32),
        _Extern([jnp.float64, jnp.float64], "__nv_atan2", jnp.float64),
        _Fallback([jnp.floating, jnp.floating], math_dialect.atan2),
    ],
    rocm=[
        _Extern([jnp.float32, jnp.float32], "__ocml_atan2_f32", jnp.float32),
        _Extern([jnp.float64, jnp.float64], "__ocml_atan2_f64", jnp.float64),
        _Fallback([jnp.floating, jnp.floating], math_dialect.atan2),
    ],
)

sinh_dispatch_table = _make_dispatch_table(
    "sinh",
    cuda=[
        _Extern([jnp.float32], "__nv_sinhf", jnp.float32),
        _Extern([jnp.float64], "__nv_sinh", jnp.float64),
        _Fallback([jnp.floating], math_dialect.sinh),
    ],
    rocm=[
        _Extern([jnp.float32], "__ocml_sinh_f32", jnp.float32),
        _Extern([jnp.float64], "__ocml_sinh_f64", jnp.float64),
        _Fallback([jnp.floating], math_dialect.sinh),
    ],
)

cosh_dispatch_table = _make_dispatch_table(
    "cosh",
    cuda=[
        _Extern([jnp.float32], "__nv_coshf", jnp.float32),
        _Extern([jnp.float64], "__nv_cosh", jnp.float64),
        _Fallback([jnp.floating], math_dialect.cosh),
    ],
    rocm=[
        _Extern([jnp.float32], "__ocml_cosh_f32", jnp.float32),
        _Extern([jnp.float64], "__ocml_cosh_f64", jnp.float64),
        _Fallback([jnp.floating], math_dialect.cosh),
    ],
)

tanh_dispatch_table = _make_dispatch_table(
    "tanh",
    cuda=[
        _Extern([jnp.float32], "__nv_tanhf", jnp.float32),
        _Extern([jnp.float64], "__nv_tanh", jnp.float64),
        _Fallback([jnp.floating], math_dialect.tanh),
    ],
    rocm=[
        _Extern([jnp.float32], "__ocml_tanh_f32", jnp.float32),
        _Extern([jnp.float64], "__ocml_tanh_f64", jnp.float64),
        _Fallback([jnp.floating], math_dialect.tanh),
    ],
)

asinh_dispatch_table = _make_dispatch_table(
    "asinh",
    cuda=[
        _Extern([jnp.float32], "__nv_asinhf", jnp.float32),
        _Extern([jnp.float64], "__nv_asinh", jnp.float64),
        _Fallback([jnp.floating], math_dialect.asinh),
    ],
    rocm=[
        _Extern([jnp.float32], "__ocml_asinh_f32", jnp.float32),
        _Extern([jnp.float64], "__ocml_asinh_f64", jnp.float64),
        _Fallback([jnp.floating], math_dialect.asinh),
    ],
)

acosh_dispatch_table = _make_dispatch_table(
    "acosh",
    cuda=[
        _Extern([jnp.float32], "__nv_acoshf", jnp.float32),
        _Extern([jnp.float64], "__nv_acosh", jnp.float64),
        _Fallback([jnp.floating], math_dialect.acosh),
    ],
    rocm=[
        _Extern([jnp.float32], "__ocml_acosh_f32", jnp.float32),
        _Extern([jnp.float64], "__ocml_acosh_f64", jnp.float64),
        _Fallback([jnp.floating], math_dialect.acosh),
    ],
)

atanh_dispatch_table = _make_dispatch_table(
    "atanh",
    cuda=[
        _Extern([jnp.float32], "__nv_atanhf", jnp.float32),
        _Extern([jnp.float64], "__nv_atanh", jnp.float64),
        _Fallback([jnp.floating], math_dialect.atanh),
    ],
    rocm=[
        _Extern([jnp.float32], "__ocml_atanh_f32", jnp.float32),
        _Extern([jnp.float64], "__ocml_atanh_f64", jnp.float64),
        _Fallback([jnp.floating], math_dialect.atanh),
    ],
)

population_count_dispatch_table = _make_dispatch_table(
    "population_count",
    cuda=[
        _Extern([jnp.int32], "__nv_popc", jnp.int32),
        _Extern([jnp.int64], "__nv_popcll", jnp.int32),
        _Fallback([jnp.integer], math_dialect.ctpop),
    ],
    rocm=[
        _Extern([jnp.int32], "__ockl_popcount_u32", jnp.int32),
        _Extern([jnp.int64], "__ockl_popcount_u64", jnp.int64),
        _Fallback([jnp.integer], math_dialect.ctpop),
    ],
)

clz_dispatch_table = _make_dispatch_table(
    "clz",
    cuda=[
        _Extern([jnp.int32], "__nv_clz", jnp.int32),
        _Extern([jnp.int64], "__nv_clzll", jnp.int32),
        _Fallback([jnp.integer], math_dialect.ctlz),
    ],
    rocm=[
        _Extern([jnp.int32], "__ockl_clz_u32", jnp.int32),
        _Extern([jnp.int64], "__ockl_clz_u64", jnp.int64),
        _Fallback([jnp.integer], math_dialect.ctlz),
    ],
)

nextafter_dispatch_table = _make_dispatch_table(
    "nextafter",
    cuda=[
        _Extern([jnp.float32, jnp.float32], "__nv_nextafterf", jnp.float32),
        _Extern([jnp.float64, jnp.float64], "__nv_nextafter", jnp.float64),
    ],
    rocm=[
        _Extern(
            [jnp.float32, jnp.float32], "__ocml_nextafter_f32", jnp.float32
        ),
        _Extern(
            [jnp.float64, jnp.float64], "__ocml_nextafter_f64", jnp.float64
        ),
    ],
)

triton_lowering_rules.update({
    lax.abs_p: abs_dispatch_table,
    lax.neg_p: lambda ctx, x: _minus(x),
    lax.ceil_p: ceil_dispatch_table,
    lax.floor_p: floor_dispatch_table,
    lax.exp_p: exp_dispatch_table,
    lax.exp2_p: exp2_dispatch_table,
    lax.expm1_p: expm1_dispatch_table,
    lax.log_p: log_dispatch_table,
    lax.log1p_p: log1p_dispatch_table,
    lax.sqrt_p: sqrt_dispatch_table,
    lax.square_p: lambda ctx, x: _mul(x, x),
    lax.pow_p: pow_dispatch_table,
    lax.cbrt_p: cbrt_dispatch_table,
    lax.rsqrt_p: rsqrt_dispatch_table,
    lax.sin_p: sin_dispatch_table,
    lax.cos_p: cos_dispatch_table,
    lax.tan_p: tan_dispatch_table,
    lax.asin_p: asin_dispatch_table,
    lax.acos_p: acos_dispatch_table,
    lax.atan_p: atan_dispatch_table,
    lax.atan2_p: atan2_dispatch_table,
    lax.sinh_p: sinh_dispatch_table,
    lax.cosh_p: cosh_dispatch_table,
    lax.tanh_p: tanh_dispatch_table,
    lax.asinh_p: asinh_dispatch_table,
    lax.acosh_p: acosh_dispatch_table,
    lax.atanh_p: atanh_dispatch_table,
    lax.population_count_p: population_count_dispatch_table,
    lax.clz_p: clz_dispatch_table,
    lax.nextafter_p: nextafter_dispatch_table,
})


def _is_triton_pointer_type(t):
  if hasattr(tt_dialect.PointerType, "isinstance"):
    return tt_dialect.PointerType.isinstance(t)
  return isinstance(t, tt_dialect.PointerType)

def _fp_bits_type(t: ir.Type) -> ir.Type:
  if isinstance(t, ir.RankedTensorType):
    t_type = ir.RankedTensorType(t)
    return ir.RankedTensorType.get(
      t_type.shape, _fp_bits_type(t_type.element_type), t_type.encoding
    )
  elif _is_triton_pointer_type(t):
    ptr_type = tt_dialect.PointerType(t)
    return tt_dialect.PointerType.get(
      _fp_bits_type(ptr_type.pointee_type), ptr_type.address_space
    )
  else:
    assert isinstance(t, ir.FloatType)
    return ir.IntegerType.get_signless(t.width)

def _minus(x: ir.Value) -> ir.Value:
  if _is_triton_pointer_type(x.type):
    raise NotImplementedError(f"unsupported type: {x.type}")
  return _sub(_zeros_like(x), x)


def _add(x: ir.Value, y: ir.Value):
  x_element_type = _element_type(x.type)
  y_element_type = _element_type(y.type)

  if _is_triton_pointer_type(x_element_type):
    assert not _is_triton_pointer_type(y_element_type)
    return tt_dialect.addptr(x.type, x, y)
  if _is_triton_pointer_type(y_element_type):
    return tt_dialect.addptr(y.type, y, x)

  assert x.type == y.type, (str(x.type), str(y.type))
  if isinstance(x_element_type, ir.IntegerType):
    return arith_dialect.addi(x, y)
  if isinstance(x_element_type, ir.FloatType):
    return arith_dialect.addf(x, y)
  raise NotImplementedError(f"unsupported dtypes: {x.type} and {y.type}")


def _sub(x: ir.Value, y: ir.Value) -> ir.Value:
  x_element_type = _element_type(x.type)
  y_element_type = _element_type(y.type)
  if _is_triton_pointer_type(x_element_type):
    return tt_dialect.addptr(x.type, x, _minus(y))
  elif not _is_triton_pointer_type(y_element_type):
    assert x.type == y.type, (str(x.type), str(y.type))
    if isinstance(x_element_type, ir.IntegerType):
      return arith_dialect.subi(x, y)
    elif isinstance(x_element_type, ir.FloatType):
      return arith_dialect.subf(x, y)
  raise NotImplementedError(f"unsupported dtype: {y.type}")


def _mul(x: ir.Value, y: ir.Value, *, out_dtype=None) -> ir.Value:
  if out_dtype is not None:
    raise NotImplementedError("out_dtype is not supported..")
  assert x.type == y.type, (str(x.type), str(y.type))
  x_element_type = _element_type(x.type)
  if isinstance(x_element_type, ir.IntegerType):
    return arith_dialect.muli(x, y)
  elif isinstance(x_element_type, ir.FloatType):
    return arith_dialect.mulf(x, y)
  raise NotImplementedError(f"unsupported types: {x.type} and {y.type}")


def _floordiv(x: ir.Value, y: ir.Value, *, signed: bool) -> ir.Value:
  assert x.type == y.type, (str(x.type), str(y.type))
  x_element_type = _element_type(x.type)
  if isinstance(x_element_type, (ir.F32Type, ir.F64Type)):
    return arith_dialect.divf(x, y)
  if not isinstance(x_element_type, ir.IntegerType):
    raise NotImplementedError(f"unsupported types: {x.type} and {y.type}")
  if signed:
    return arith_dialect.divsi(x, y)
  else:
    return arith_dialect.divui(x, y)


def _truediv(x: ir.Value, y: ir.Value, *, signed: bool) -> ir.Value:
  assert x.type == y.type, (str(x.type), str(y.type))
  x_element_type = _element_type(x.type)
  if isinstance(x_element_type, ir.IntegerType):
    x_element_type = ir.F32Type.get()
    x = _int_float_cast(x, x_element_type, signed=signed)
    y = _int_float_cast(y, x_element_type, signed=signed)
  if isinstance(x_element_type, (ir.F32Type, ir.F64Type)):
    return arith_dialect.divf(x, y)
  raise NotImplementedError(f"unsupported types: {x.type} and {y.type}")


def _mod(x: ir.Value, y: ir.Value, *, signed: bool) -> ir.Value:
  assert x.type == y.type, (str(x.type), str(y.type))
  x_element_type = _element_type(x.type)
  if isinstance(x_element_type, ir.FloatType):
    return arith_dialect.remf(x, y)
  if not isinstance(x_element_type, ir.IntegerType):
    raise NotImplementedError(f"unsupported types: {x.type} and {y.type}")
  if signed:
    return arith_dialect.remsi(x, y)
  else:
    return arith_dialect.remui(x, y)


def _cmp(
    x: ir.Value,
    y: ir.Value,
    si_pred: arith_dialect.CmpIPredicate,
    ui_pred: arith_dialect.CmpIPredicate,
    f_pred: arith_dialect.CmpFPredicate,
    *,
    signed: bool,
) -> ir.Value:
  assert x.type == y.type, (str(x.type), str(y.type))
  x_element_type = _element_type(x.type)
  if isinstance(x_element_type, ir.IntegerType):
    return arith_dialect.cmpi(si_pred if signed else ui_pred, x, y)
  elif isinstance(x_element_type, ir.FloatType):
    return arith_dialect.cmpf(f_pred, x, y)
  else:
    raise NotImplementedError(f"unsupported types: {x.type} and {y.type}")


_equal = functools.partial(
    _cmp,
    si_pred=arith_dialect.CmpIPredicate.eq,
    ui_pred=arith_dialect.CmpIPredicate.eq,
    f_pred=arith_dialect.CmpFPredicate.OEQ,
)
_not_equal = functools.partial(
    _cmp,
    si_pred=arith_dialect.CmpIPredicate.ne,
    ui_pred=arith_dialect.CmpIPredicate.ne,
    f_pred=arith_dialect.CmpFPredicate.UNE,
)
_less_than = functools.partial(
    _cmp,
    si_pred=arith_dialect.CmpIPredicate.slt,
    ui_pred=arith_dialect.CmpIPredicate.ult,
    f_pred=arith_dialect.CmpFPredicate.OLT,
)
_less_equal = functools.partial(
    _cmp,
    si_pred=arith_dialect.CmpIPredicate.sle,
    ui_pred=arith_dialect.CmpIPredicate.ule,
    f_pred=arith_dialect.CmpFPredicate.OLE,
)
_greater_than = functools.partial(
    _cmp,
    si_pred=arith_dialect.CmpIPredicate.sgt,
    ui_pred=arith_dialect.CmpIPredicate.ugt,
    f_pred=arith_dialect.CmpFPredicate.OGT,
)
_greater_equal = functools.partial(
    _cmp,
    si_pred=arith_dialect.CmpIPredicate.sge,
    ui_pred=arith_dialect.CmpIPredicate.uge,
    f_pred=arith_dialect.CmpFPredicate.OGE,
)


def _is_nan(x: ir.Value) -> ir.Value:
  return arith_dialect.cmpf(arith_dialect.CmpFPredicate.UNO, x, x)


_JAX_TO_TRITON_BINARY = {
    lax.add_p: _add,
    lax.sub_p: _sub,
    lax.mul_p: _mul,
    lax.and_p: arith_dialect.andi,
    lax.or_p: arith_dialect.ori,
    lax.xor_p: arith_dialect.xori,
    lax.shift_left_p: arith_dialect.shli,
    lax.shift_right_arithmetic_p: arith_dialect.shrsi,
    lax.shift_right_logical_p: arith_dialect.shrui,
    ad_util.add_any_p: _add,
}

for prim, fn in _JAX_TO_TRITON_BINARY.items():

  def signless_rule(ctx: LoweringRuleContext, x, y, fn=fn, **kwargs):
    x, y = _bcast(x, y, *ctx.avals_in, *ctx.avals_out)
    return fn(x, y, **kwargs)

  triton_lowering_rules[prim] = signless_rule


_JAX_TO_TRITON_SIGNED_BINARY = {
    lax.rem_p: _mod,
    lax.eq_p: _equal,
    lax.ne_p: _not_equal,
    lax.gt_p: _greater_than,
    lax.ge_p: _greater_equal,
    lax.lt_p: _less_than,
    lax.le_p: _less_equal,
}

for prim, fn in _JAX_TO_TRITON_SIGNED_BINARY.items():

  def signed_rule(ctx: LoweringRuleContext, x, y, fn=fn):
    x_aval, _ = ctx.avals_in
    x, y = _bcast(x, y, *ctx.avals_in, *ctx.avals_out)
    return fn(x, y, signed=jnp.issubdtype(x_aval.dtype, jnp.signedinteger))

  triton_lowering_rules[prim] = signed_rule


@register_lowering(debugging.debug_print_p)
def debug_print_lowering_rule(
    ctx: LoweringRuleContext,
    *args: ir.Value,
    fmt: str,
    ordered,
    partitioned,
    in_tree,
    static_args,
    np_printoptions,
    has_placeholders,
    logging_record,
):
  del partitioned, np_printoptions
  if ordered:
    raise NotImplementedError("Ordered debug_print is not supported on Pallas.")
  if has_placeholders:
    raise ValueError(
        "pl.debug_print() does not support placeholders when lowering to Triton"
    )
  args, kwargs = debugging.merge_callback_args(in_tree, args, static_args)
  if kwargs:
    raise ValueError(
        "Only positional arguments are supported by debug_print on Pallas."
    )

  tt_dialect.print_(
      f" {fmt} ",
      hex=False,
      args=args,
      is_signed=ir.DenseI32ArrayAttr.get([
          jnp.issubdtype(aval.dtype, jnp.signedinteger) for aval in ctx.avals_in
      ]),
  )
  return ()


def _set_attr(v: ir.Value, name: str, attr: ir.Attribute) -> None:
  if not isinstance(v, ir.BlockArgument):
    v.owner.attributes[name] = attr  # pyrefly: ignore[missing-attribute]
    return

  arg = ir.BlockArgument(v)
  name += f"_arg{arg.arg_number}"
  owner = arg.owner
  is_entry = owner.region.blocks[0] == owner
  if not is_entry:
    return
  if (op := owner.owner.operation) and not isinstance(op, tt_dialect.FuncOp):
    op.attributes[name] = attr


@register_lowering(primitives.multiple_of_p)
def _multiple_of_rule(ctx: LoweringRuleContext, x, values: Sequence[int]):
  [x_aval] = ctx.avals_in
  assert max(1, len(x_aval.shape)) == len(values)
  _set_attr(
      x,
      "tt.divisibility",
      ir.DenseIntElementsAttr.get(np.asarray(values, dtype=np.int32)),  # pyrefly: ignore[no-matching-overload]
  )
  return x


@register_lowering(sp.broadcast_to_p)
def _broadcast_to_rule(ctx: LoweringRuleContext, x, shape: Sequence[int]):
  (x_aval,) = ctx.avals_in
  return _bcast_to(_ensure_ir_value(x, x_aval), shape)


@register_lowering(lax.integer_pow_p)
def _integer_pow_rule(ctx: LoweringRuleContext, x, *, y: int):
  if y == 0:
    return _ones_like(x)

  is_reciprocal = y < 0
  if is_reciprocal:
    y = -y

  acc: ir.Value | None = None
  while y > 0:
    y, mod = divmod(y, 2)
    if mod:
      acc = x if acc is None else _mul(acc, x)
    if y > 0:
      x = _mul(x, x)
  assert acc is not None

  [x_aval] = ctx.avals_in
  [out_aval] = ctx.avals_out
  acc = _cast(acc, x_aval.dtype, out_aval.dtype)
  if is_reciprocal:
    signed = jnp.issubdtype(out_aval.dtype, jnp.signedinteger)
    return  _truediv(_ones_like(acc), acc, signed=signed)
  else:
    return acc


_JAX_FN_MAPPING = {
    lax.clamp_p: lambda min, a, max: jnp.minimum(jnp.maximum(min, a), max),
    lax.logistic_p: lambda a, accuracy: 1 / (1 + jnp.exp(-a)),
    lax.is_finite_p: lambda x: jnp.logical_and(~jnp.isnan(x), ~jnp.isinf(x)),
}

for prim, fn in _JAX_FN_MAPPING.items():
  triton_lowering_rules[prim] = lower_fun(fn, multiple_results=False)


@register_lowering(lax.min_p)
def _min_lowering_rule(ctx: LoweringRuleContext, x, y):
  # TODO(slebedev): Consider allowing customizing nan behavior.
  x_aval, y_aval = ctx.avals_in
  x, y = _bcast(x, y, *ctx.avals_in, *ctx.avals_out)
  if jnp.issubdtype(x_aval.dtype, jnp.floating):
    # TODO(slebedev): Triton promotes bfloat16 to float32 and back here.
    return arith_dialect.minnumf(x, y)
  if not jnp.issubdtype(x_aval.dtype, jnp.integer):
    raise NotImplementedError(
        f"unsupported dtypes: {x_aval.dtype} and {y_aval.dtype}"
    )
  if jnp.issubdtype(x_aval.dtype, jnp.signedinteger):
    return arith_dialect.minsi(x, y)
  else:
    return arith_dialect.minui(x, y)


@register_lowering(lax.max_p)
def _max_lowering_rule(ctx: LoweringRuleContext, x, y):
  # TODO(slebedev): Consider allowing customizing nan behavior.
  x_aval, y_aval = ctx.avals_in
  x, y = _bcast(x, y, *ctx.avals_in, *ctx.avals_out)
  if jnp.issubdtype(x_aval.dtype, jnp.floating):
    # TODO(slebedev): Triton promotes bfloat16 to float32 and back here.
    return arith_dialect.maxnumf(x, y)
  if not jnp.issubdtype(x_aval.dtype, jnp.integer):
    raise NotImplementedError(
        f"unsupported dtypes: {x_aval.dtype} and {y_aval.dtype}"
    )
  if jnp.issubdtype(x_aval.dtype, jnp.signedinteger):
    return arith_dialect.maxsi(x, y)
  else:
    return arith_dialect.maxui(x, y)


@register_lowering(lax.div_p)
def _div_lowering_rule(ctx: LoweringRuleContext, x, y):
  x_aval, y_aval = ctx.avals_in
  x, y = _bcast(x, y, *ctx.avals_in, *ctx.avals_out)
  signed = jnp.issubdtype(x_aval.dtype, jnp.signedinteger) or jnp.issubdtype(
      y_aval.dtype, jnp.signedinteger
  )
  if jnp.issubdtype(x_aval.dtype, np.floating) or jnp.issubdtype(
      y_aval.dtype, np.floating
  ):
    return _truediv(x, y, signed=signed)
  return _floordiv(x, y, signed=signed)


register_lowering(lax.sign_p)(
    lower_fun(pallas_utils.sign_lowering_helper, multiple_results=False)
)


register_lowering(lax.erf_inv_p)(
    lower_fun(pallas_utils.erf_inv_lowering_helper, multiple_results=False)
)


@register_lowering(lax.iota_p)
def _iota_lowering_rule(ctx: LoweringRuleContext, *, dtype, shape, dimension,
                        sharding):
  iota = _make_range(0, shape[dimension])
  iota = _cast(iota, jnp.int32, dtype)
  for i in range(len(shape)):
    if i != dimension:
      iota = _expand_dims(iota, i)
  return _bcast_to(iota, shape)


def _element_type(t: ir.Type) -> ir.Type:
  if isinstance(t, ir.RankedTensorType):
    return ir.RankedTensorType(t).element_type
  else:
    return t


def _make_range(start: int, end: int) -> ir.Value:
  if end <= start:
    raise ValueError(
        f"end must be greater than start, but got: {end} <= {start}"
    )
  if max(start, end) >= 2**32:
    raise ValueError("start and end must fit in int32")
  return tt_dialect.make_range(
      ir.RankedTensorType.get([end - start], ir.IntegerType.get_signless(32)),
      start,
      end,
  )


def _full(t: ir.Type, v: Any) -> ir.Value:
  element_type = _element_type(t)
  if isinstance(element_type, ir.IntegerType):
    result = arith_dialect.constant(element_type, int(v))
  elif isinstance(element_type, ir.FloatType):
    result = arith_dialect.constant(element_type, float(v))
  else:
    raise NotImplementedError

  if isinstance(t, ir.RankedTensorType):
    return tt_dialect.splat(t, result)
  else:
    return result


def _zeros(t: ir.Type) -> ir.Value:
  return _full(t, 0)


def _zeros_like(x: ir.Value) -> ir.Value:
  return _full(x.type, 0)


def _ones(t: ir.Type) -> ir.Value:
  return _full(t, 1)


def _ones_like(x: ir.Value) -> ir.Value:
  return _full(x.type, 1)


def _splat(x: ir.Value, shape: Sequence[int]) -> ir.Value:
  if isinstance(x.type, ir.RankedTensorType):
    raise TypeError("cannot splat a tensor")
  if not shape:
    return x
  return tt_dialect.splat(ir.RankedTensorType.get(shape, x.type), x)


def _expand_dims(x: ir.Value, axis: int) -> ir.Value:
  if not isinstance(x.type, ir.RankedTensorType):
    shape = list(ir.RankedTensorType(x.type).shape)
    shape.insert(axis, 1)
    return _splat(x, shape)
  return tt_dialect.expand_dims(x, axis)


def _float_float_cast(src: ir.Value, dst_type: ir.Type) -> ir.Value:
  src_element_type = ir.FloatType(_element_type(src.type))
  dst_element_type = ir.FloatType(_element_type(dst_type))
  if src_element_type.width == 8 or dst_element_type.width == 8:
    rounding = (
        tt_dialect.RoundingMode.RTNE if src_element_type.width > 8 else None
    )
    return tt_dialect.fp_to_fp(dst_type, src, rounding=rounding)
  if src_element_type.width > dst_element_type.width:
    return arith_dialect.truncf(dst_type, src)
  elif src_element_type.width < dst_element_type.width:
    return arith_dialect.extf(dst_type, src)
  else:
    raise NotImplementedError


def _int_int_cast(src: ir.Value, dst_type: ir.Type, signed: bool) -> ir.Value:
  src_element_type = ir.IntegerType(_element_type(src.type))
  dst_element_type = ir.IntegerType(_element_type(dst_type))
  assert src_element_type != dst_element_type
  if dst_element_type.width == 1:
    return _not_equal(src, _zeros_like(src), signed=signed)

  if src_element_type.width == dst_element_type.width:
    return arith_dialect.bitcast(dst_type, src)
  elif src_element_type.width > dst_element_type.width:
    return arith_dialect.trunci(dst_type, src)
  elif signed and src_element_type.width != 1:
    return arith_dialect.extsi(dst_type, src)
  else:
    return arith_dialect.extui(dst_type, src)


def _float_int_cast(
    src: ir.Value, dst_type: ir.Type, *, signed: bool
) -> ir.Value:
  src_element_type = _element_type(src.type)
  if not isinstance(src_element_type, (ir.BF16Type, ir.F16Type, ir.F32Type, ir.F64Type)):
    raise NotImplementedError(f"cannot cast {src} tp {dst_type}")
  dst_element_type = ir.IntegerType(_element_type(dst_type))
  if dst_element_type.width == 1:
    return _not_equal(src, _zeros_like(src), signed=signed)
  else:
    # We clamp the float value to the min/max integer destination value
    # in order to match JAX/XLA casting behavior. Note that this differs
    # from numpy casting behavior.
    if signed:
      maxint = 2**(dst_element_type.width-1) - 1
      minint = -2**(dst_element_type.width-1)
    else:
      maxint = 2**dst_element_type.width - 1
      minint = 0
    src = arith_dialect.minimumf(src, _full(src.type, maxint))
    src = arith_dialect.maximumf(src, _full(src.type, minint))
    if signed:
      return arith_dialect.fptosi(dst_type, src)
    else:
      return arith_dialect.fptoui(dst_type, src)


def _int_float_cast(
    src: ir.Value, dst_type: ir.Type, *, signed: bool
) -> ir.Value:
  src_element_type = ir.IntegerType(_element_type(src.type))
  dst_element_type = _element_type(dst_type)
  if not isinstance(
      dst_element_type, (ir.BF16Type, ir.F16Type, ir.F32Type, ir.F64Type)
  ):
    raise NotImplementedError(f"cannot cast {src} tp {dst_type}")
  if src_element_type.width == 1 or not signed:
    return arith_dialect.uitofp(dst_type, src)
  else:
    return arith_dialect.sitofp(dst_type, src)


def _cast(
    src: ir.Value,
    src_type: jax.typing.DTypeLike,
    dst_type: jax.typing.DTypeLike,
    *,
    compute_capability: int | None = None,
) -> ir.Value:
  return _ir_cast(
      src,
      _dtype_to_ir_type(dst_type),
      signed=jnp.issubdtype(src_type, jnp.signedinteger),
      dst_signed=jnp.issubdtype(dst_type, jnp.signedinteger),
      compute_capability=compute_capability,
  )


def _is_float8_e4m3fn_cast_supported(compute_capability: int | None) -> bool:
  return compute_capability is None or compute_capability >= 89


_UNSUPPORTED_CAST_DTYPES = (
    (ir.Float8E4M3FNType, "float8_e4m3fn", _is_float8_e4m3fn_cast_supported),
    # TODO(slebedev): Check the CUDA version and raise conditionally.
    (ir.Float8E4M3FNUZType, "float8_e4m3fnuz", lambda _: False),
)


def _ir_cast(
    src: ir.Value,
    dst_type: ir.Type,
    *,
    signed: bool,
    dst_signed: bool = False,
    compute_capability: int | None = None,
) -> ir.Value:
  if isinstance(src.type, ir.RankedTensorType) and not isinstance(
      dst_type, ir.RankedTensorType
  ):
    src_type = ir.RankedTensorType(src.type)
    dst_type = ir.RankedTensorType.get(
        src_type.shape,
        dst_type,
        src_type.encoding,
    )
  if src.type == dst_type:
    return src

  src_element_type = _element_type(src.type)
  dst_element_type = _element_type(dst_type)

  for dtype, dtype_name, is_supported in _UNSUPPORTED_CAST_DTYPES:
    if isinstance(src_element_type, dtype):
      if not is_supported(compute_capability):
        raise NotImplementedError(f"cannot cast from `{dtype_name}`")
    if isinstance(dst_element_type, dtype):
      if not is_supported(compute_capability):
        raise NotImplementedError(f"cannot cast to `{dtype_name}`")

  if isinstance(src_element_type, (ir.F16Type, ir.BF16Type)) and not isinstance(
      dst_element_type, ir.F32Type
  ):
    return _ir_cast(
        _ir_cast(src, ir.F32Type.get(), signed=False),
        dst_type, signed=False, dst_signed=dst_signed
    )

  if isinstance(src_element_type, ir.FloatType) and isinstance(
      dst_element_type, ir.FloatType
  ):
    return _float_float_cast(src, dst_type)

  if isinstance(src_element_type, ir.IntegerType) and isinstance(
      dst_element_type, ir.IntegerType
  ):
    return _int_int_cast(src, dst_type, signed=signed)

  if isinstance(src_element_type, ir.FloatType) and isinstance(
      dst_element_type, ir.IntegerType
  ):
    return _float_int_cast(src, dst_type, signed=dst_signed)
  if isinstance(src_element_type, ir.IntegerType) and isinstance(
      dst_element_type, ir.FloatType
  ):
    return _int_float_cast(src, dst_type, signed=signed)

  if _is_triton_pointer_type(src_element_type) and isinstance(
      dst_element_type, ir.IntegerType
  ):
    if dst_element_type.width == 64:
      return tt_dialect.ptr_to_int(dst_type, src)
    elif dst_element_type.width == 1:
      x = _ir_cast(src, ir.IntegerType.get_signless(64), signed=signed)
      zero = _zeros_like(x)
      return _ir_cast(_not_equal(x, zero, signed=signed), dst_type, signed=signed)
  if isinstance(src_element_type, ir.IntegerType) and _is_triton_pointer_type(
      dst_element_type
  ):
    return tt_dialect.int_to_ptr(dst_type, src)
  if _is_triton_pointer_type(src_element_type) and _is_triton_pointer_type(
      dst_element_type
  ):
    return tt_dialect.bitcast(dst_type, src)

  raise NotImplementedError(f"cannot cast {src} to {dst_type}")


@register_lowering(lax.convert_element_type_p)
def _convert_element_type_lowering_rule(
    ctx: LoweringRuleContext, x, *, new_dtype, weak_type, sharding
):
  [x_aval] = ctx.avals_in
  x = _ensure_ir_value(x, x_aval)
  if new_dtype == x_aval.dtype:
    return x
  cc = ctx.context.compute_capability
  return _cast(x, x_aval.dtype, new_dtype, compute_capability=cc)


@register_lowering(lax.select_n_p)
def select_n_lowering_rule(ctx: LoweringRuleContext, pred, x, y):
  pred_aval, a_aval, b_aval = ctx.avals_in
  [out_aval] = ctx.avals_out
  pred, x = _bcast(pred, x, pred_aval, a_aval, out_aval)
  pred, y = _bcast(pred, y, pred_aval, b_aval, out_aval)
  return arith_dialect.select(pred, y, x)


@register_lowering(lax.broadcast_in_dim_p)
def _broadcast_in_dim_lowering_rule(
    ctx: LoweringRuleContext, x, *, broadcast_dimensions, shape, sharding
):
  del sharding
  x = _ensure_ir_value(x, *ctx.avals_in)
  if not isinstance(x.type, ir.RankedTensorType):
    return _bcast_to(x, shape)
  expand_dims = [i for i in range(len(shape)) if i not in broadcast_dimensions]
  for dim in expand_dims:
    x = _expand_dims(x, dim)
  return _bcast_to(x, shape)


@register_lowering(lax.squeeze_p)
def _squeeze_lowering_rule(ctx: LoweringRuleContext, a, *, dimensions):
  del dimensions
  return _reshape_lowering_rule(ctx, a, new_sizes=None, dimensions=None, sharding=None)


@register_lowering(lax.reshape_p)
def _reshape_lowering_rule(
    ctx: LoweringRuleContext, a, *, new_sizes, dimensions, sharding,
):
  del new_sizes  # Unused.
  if dimensions is not None:
    return ValueError("`dimensions` is not supported.")

  a = _ensure_ir_value(a, *ctx.avals_in)
  [a_aval] = ctx.avals_in
  [out_aval] = ctx.avals_out
  # Triton Reshape doesn't support scalar result types (only 0d tensors).
  if out_aval.ndim == 0:
    return _reduce_lowering(jnp.add, ctx, a, axes=tuple(range(a_aval.ndim)))
  return _reshape(a, out_aval.shape)


def _reshape(a: ir.Value, shape: Sequence[int]) -> ir.Value:
  if not isinstance(a.type, ir.RankedTensorType):
    assert all(dim_size == 1 for dim_size in shape)
    return _splat(a, shape)

  ty = ir.RankedTensorType(a.type)
  return tt_dialect.reshape(
      ir.RankedTensorType.get(shape, ty.element_type, ty.encoding),
      a,
      allow_reorder=False,
  )


def get_join_type(old_type: ir.RankedTensorType):
  shape = old_type.shape
  shape.append(2)
  return ir.RankedTensorType.get(shape, old_type.element_type, old_type.encoding)


@register_lowering(lax.concatenate_p)
def _concatenate_lowering_rule(ctx: LoweringRuleContext, *args, dimension):
  if len(args) != 2:
    raise NotImplementedError("Only 2-argument concatenate is supported.")
  x_aval, y_aval = ctx.avals_in
  x, y = args
  if dimension != x_aval.ndim-1:
    raise NotImplementedError(
        "Only concatenate along the last dimension is supported."
    )
  if x_aval.shape[-1] != 1 or y_aval.shape[-1] != 1:
    raise NotImplementedError(
        "Only arguments with shape [..., 1] are supported."
    )
  lhs = _reshape(x, x_aval.shape[:-1])
  rhs = _reshape(y, y_aval.shape[:-1])
  ret_type = get_join_type(ir.RankedTensorType(rhs.type))
  return tt_dialect.join(ret_type, lhs, rhs)


@register_lowering(lax.split_p)
def _split_lowering_rule(ctx: LoweringRuleContext, x, *, sizes, axis):
  pass
  # TODO(cjfj): Add support for larger powers of 2.
  num_parts = len(sizes)
  if num_parts != pallas_utils.next_power_of_2(num_parts):
    raise NotImplementedError("Only power-of-2 num parts supported.")
  if any(size != sizes[0] for size in sizes):
    raise NotImplementedError("Only equal-sized splits are supported.")

  def split_into_2(x):
    shape = ir.RankedTensorType(x.type).shape
    x = _reshape(x, shape[:axis] + [2, shape[axis] // 2] + shape[axis + 1 :])
    permutation = tuple(d for d in range(len(shape) + 1) if d != axis) + (axis,)
    return tuple(tt_dialect.split(tt_dialect.trans(x, permutation)))

  x_parts: tuple[ir.Value, ...] = (x,)
  while len(x_parts) < num_parts:
    x_parts = sum(map(split_into_2, x_parts), ())
  return x_parts


def _compute_offsets_from_indices(
    block_info: BlockInfo, nd_indexer: NDIndexer
) -> ir.Value:
  full_shape = block_info.full_shape_dtype.shape
  num_squeezed_dims = sum(isinstance(b, pallas_core.Squeezed)
                          for b in block_info.block_shape)
  strides = pallas_utils.strides_from_shape(full_shape)
  indexer_shape = nd_indexer.get_indexer_shape_static()
  int_indexer_shape = nd_indexer.int_indexer_shape
  _check_tensor_size(indexer_shape)
  indices = nd_indexer.indices
  other_shape = indexer_shape[len(int_indexer_shape) :]
  other_shape_idx = 0
  assert len(indices) + num_squeezed_dims == len(full_shape)
  assert len(block_info.start_indices) == len(full_shape)

  array_dtype = jnp.dtype(block_info.full_shape_dtype.dtype)
  full_size = math.prod(full_shape) * array_dtype.itemsize
  # Use 64-bit indexing when offset might be >= 2**32 bytes.
  offset_eltype = ir.IntegerType.get_signless(64 if full_size > 2**32 else 32)
  if indexer_shape:
    offsets = _zeros(ir.RankedTensorType.get(indexer_shape, offset_eltype))
  else:
    offsets = _ir_constant(0, offset_eltype)

  indexer_iter = iter(indices)
  for dim_stride, dim_block_size, start_offset in zip(
      strides, block_info.block_shape, block_info.start_indices
  ):
    match dim_block_size:
      case pallas_core.Squeezed():
        index = _ir_constant(0, offset_eltype)
      case int():
        index = next(indexer_iter)
      case _:
        raise ValueError(f"Unexpected dim_block_size: {dim_block_size}")

    if isinstance(index, slice):
      index = primitives.Slice.from_slice(
          index, pallas_core.get_block_size(dim_block_size)
      )

    if isinstance(index, primitives.Slice):
      if index.is_dynamic_start or (index.stride != 1):
        if not index.is_dynamic_start:
          start = _ir_constant(index.start, offset_eltype)
        else:
          assert isinstance(index.start, ir.Value)
          start = index.start
        start = _ir_cast(start, offset_eltype, signed=False)

        iota = _ir_cast(
            _make_range(0, int(index.size)), offset_eltype, signed=False
        )
        if index.stride != 1:
          iota = _mul(iota, _full(iota.type, index.stride))
        dim_offsets = _add(_bcast_to(start, (int(index.size),)), iota)
      else:
        iota = _make_range(int(index.start), int(index.start + index.size))
        dim_offsets = _ir_cast(iota, offset_eltype, signed=False)

      other_shape_idx += 1
      for _ in other_shape[other_shape_idx:]:
        rank = ir.RankedTensorType(dim_offsets.type).rank
        dim_offsets = _expand_dims(dim_offsets, rank)
    else:
      # indexer is either a *scalar* or an array of size `int_indexer_shape`
      dim_offsets = index
      if not isinstance(dim_offsets, ir.Value):
        dim_offsets = _ir_constant(dim_offsets, offset_eltype)
      dim_offsets = _ir_cast(dim_offsets, offset_eltype, signed=False)

      if isinstance(dim_offsets.type, ir.RankedTensorType):
        for _ in other_shape:
          rank = ir.RankedTensorType(dim_offsets.type).rank
          dim_offsets = _expand_dims(dim_offsets, rank)

    if isinstance(dim_offsets.type, ir.RankedTensorType):
      rank = ir.RankedTensorType(dim_offsets.type).rank
      for _ in range(len(indexer_shape) - rank):
        dim_offsets = _expand_dims(dim_offsets, 0)
    dim_offsets = _bcast_to(dim_offsets, indexer_shape)

    if start_offset is not None:
      start_offset = _ir_cast(start_offset, offset_eltype, signed=False)
      dim_offsets = _add(dim_offsets, _bcast_to(start_offset, indexer_shape))

    dim_offsets = _mul(dim_offsets, _full(dim_offsets.type, dim_stride))
    offsets = _add(offsets, dim_offsets)

  return offsets


def _compute_pointers_from_indices(
    root_ptr: ir.Value, block_info: BlockInfo, nd_indexer: NDIndexer
) -> ir.Value:
  offsets = _compute_offsets_from_indices(block_info, nd_indexer)
  shape = nd_indexer.get_indexer_shape_static()
  return _add(_bcast_to(root_ptr, shape), offsets)


@register_lowering(sp.get_p)
def _get_lowering_rule(ctx: LoweringRuleContext, ptr, *idx, tree):
  indexers = tree_util.tree_unflatten(tree, idx)
  if not _is_triton_pointer_type(ptr.type):
    assert len(indexers) == 0
    return ptr
  args_flat, args_tree = tree_util.tree_flatten((ptr, indexers, None, None))
  return _masked_load_lowering_rule(
      ctx,
      *args_flat,
      args_tree=args_tree,
      eviction_policy=None,
      cache_modifier=None,
      is_volatile=False,
  )


_STR_TO_EVICTION_POLICY = {str(e): e for e in tt_dialect.EvictionPolicy}
_STR_TO_CACHE_MODIFIER = {str(c): c for c in tt_dialect.CacheModifier}


def _load(
    ptr: ir.Value,
    mask: ir.Value | None = None,
    other: ir.Value | None = None,
    *,
    cache_modifier: str | None = None,
    eviction_policy: str | None = None,
    is_volatile: bool = False,
) -> ir.Value:
  if cache_modifier is None:
    cache = tt_dialect.CacheModifier.NONE
  elif cache_modifier == ".ca" or cache_modifier == ".cg":
    cache = _STR_TO_CACHE_MODIFIER[cache_modifier]
  else:
    raise ValueError(f"unsupported cache modifier: {cache_modifier}")
  if eviction_policy is None:
    evict = tt_dialect.EvictionPolicy.NORMAL
  else:
    try:
      evict = _STR_TO_EVICTION_POLICY[eviction_policy]
    except KeyError:
      raise ValueError(
          f"unsupported eviction policy: {eviction_policy}"
      ) from None

  if _is_triton_pointer_type(ptr.type):
    ptr_type = tt_dialect.PointerType(ptr.type)
    if isinstance(ptr_type.pointee_type, ir.RankedTensorType):
      raise NotImplementedError("loading from a block pointer is not supported")

  ptr_type = _element_type(ptr.type)
  if not _is_triton_pointer_type(ptr_type):
    raise ValueError(f"unsupported pointer type: {ptr_type}")
  ptr_type = tt_dialect.PointerType(ptr_type)
  if other is not None and mask is None:
    raise ValueError("other requires mask to be provided")
  if not isinstance(ptr.type, ir.RankedTensorType):
    if other is not None and isinstance(other.type, ir.RankedTensorType):
      raise ValueError("other cannot be a block if pointer is not a block")
    if mask is not None and isinstance(mask.type, ir.RankedTensorType):
      raise ValueError("mask cannot be a block if pointer is not a block")

  pointee_type = ptr_type.pointee_type
  is_int1 = isinstance(pointee_type, ir.IntegerType) and pointee_type.width == 1
  if is_int1:
    pointee_type = ir.IntegerType.get_signless(8)
    ptr = _ir_cast(
        ptr,
        tt_dialect.PointerType.get(pointee_type, ptr_type.address_space),
        signed=False,
    )

  if other is not None:
    other = _ir_cast(other, pointee_type, signed=False)

  result = tt_dialect.load(
      ptr,
      mask=mask,
      other=other,
      cache=cache,
      evict=evict,
      is_volatile=is_volatile,
  )
  return (
      result
      if not is_int1
      else _ir_cast(result, ir.IntegerType.get_signless(1), signed=False)
  )


def _is_contiguous_int4(block_info: BlockInfo, nd_indexer: NDIndexer) -> bool:
  """Returns True if the block is contiguous in the last dimension."""
  # In order to loaded as `uint8` the index must be an aligned slice.
  return (
      block_info.full_shape_dtype.dtype in (jnp.int4, jnp.uint4)
      and bool(block_info.start_indices_alignment)
      and (block_info.start_indices_alignment[-1] % 2 == 0)
      and isinstance(slc := nd_indexer.indices[-1], indexing.Slice)
      and isinstance(slc.start, int)
      and isinstance(slc.size, int)
      and (slc.start % 2 == 0)
      and (slc.size % 2 == 0)
      and (slc.stride == 1)
  )


def _reinterpret_int4_as_uint8(
    block_info: BlockInfo, nd_indexer: NDIndexer
) -> tuple[BlockInfo, NDIndexer]:
  """Returns a new block info and indexer that reads `int4` as `uint8`."""
  last_idx = nd_indexer.indices[-1]
  # pyrefly: ignore[missing-attribute]
  new_last_idx = indexing.Slice(last_idx.start // 2, last_idx.size // 2)
  new_indices = (*nd_indexer.indices[:-1], new_last_idx)
  new_shape = (*nd_indexer.shape[:-1], nd_indexer.shape[-1] // 2)
  idx = dataclasses.replace(nd_indexer, indices=new_indices, shape=new_shape)

  full_shape = block_info.full_shape_dtype.shape
  new_full_shape = (*full_shape[:-1], full_shape[-1] // 2)
  start_idx = block_info.start_indices[-1]
  new_start_idx = _floordiv(start_idx, _full(start_idx.type, 2), signed=False)
  new_start_indices = (*block_info.start_indices[:-1], new_start_idx)
  block_info = dataclasses.replace(
      block_info,
      full_shape_dtype=jax_core.ShapedArray(new_full_shape, jnp.uint8),
      start_indices=new_start_indices,
  )
  return block_info, idx


@register_lowering(primitives.load_p)
def _masked_load_lowering_rule(
    ctx: LoweringRuleContext,
    *args_flat,
    args_tree,
    eviction_policy,
    cache_modifier,
    is_volatile,
):
  block_info, *_ = ctx.block_infos
  assert block_info is not None
  ptr, indexers, mask, other = args_tree.unflatten(args_flat)
  *_, mask_aval, other_aval = args_tree.unflatten(ctx.avals_in)
  if len(indexers) > 1:
    raise NotImplementedError("No support for multiple indexers yet.")
  indexers = list(indexers)
  if not indexers:
    ref_aval = state.transform_type(indexers, ctx.avals_in[0])
    assert isinstance(ref_aval, state.AbstractRef)
    idx = NDIndexer.make_trivial_indexer(ref_aval.shape)
  else:
    idx = indexers[0]
  if not _is_triton_pointer_type(ptr.type):
    assert len(ctx.avals_in) == 1
    return ptr

  is_int4 = block_info.full_shape_dtype.dtype in (jnp.int4, jnp.uint4)
  is_contiguous_int4 = _is_contiguous_int4(block_info, idx)

  if is_contiguous_int4:
    # If the load reads contiguously in the last dimension, we can reinterpret
    # the `int4` block as `uint8`. This generates much more efficient code. The
    # more generic `int4` code below has offsets like `0, 0, 1, 1, ...`, which
    # Triton doesn't optimize as well.
    block_info, idx = _reinterpret_int4_as_uint8(block_info, idx)

  offsets = _compute_offsets_from_indices(block_info, idx)
  ptr_offsets = offsets

  if is_int4 and not is_contiguous_int4:
    ptr_offsets = _floordiv(offsets, _full(offsets.type, 2), signed=False)

  shape = idx.get_indexer_shape_static()
  ptr = _add(_bcast_to(ptr, shape), ptr_offsets)
  if mask is not None:
    mask = _bcast_to(_ensure_ir_value(mask, mask_aval), shape)
  if other is not None:
    other = _bcast_to(_ensure_ir_value(other, other_aval), shape)
  values = _load(
      ptr,
      mask=mask,
      other=other,
      cache_modifier=cache_modifier,
      is_volatile=is_volatile,
      eviction_policy=eviction_policy,
  )

  if not is_int4:
    return values

  if is_contiguous_int4:
    msb_values = arith_dialect.shrui(values, _full(values.type, 4))
    join_type = get_join_type(ir.RankedTensorType(values.type))
    values = tt_dialect.join(join_type, values, msb_values)
    shape = ir.RankedTensorType(values.type).shape
    values = _reshape(values, (*shape[:-2], shape[-2] * shape[-1]))
  else:
    offsets = _ir_cast(offsets, ir.IntegerType.get_signless(32), signed=False)
    in_msb = _mod(offsets, _full(offsets.type, 2), signed=False)
    shift = _mul(in_msb, _full(in_msb.type, 4))
    shift = _ir_cast(shift, values.type, signed=False)
    values = arith_dialect.shrui(values, shift)
  return _ir_cast(values, ir.IntegerType.get_signless(4), signed=False)


@register_lowering(sp.swap_p)
def _swap_lowering_rule(ctx: LoweringRuleContext, ptr, value, *idx, tree):
  indexers = tree_util.tree_unflatten(tree, idx)
  if not _is_triton_pointer_type(ptr.type):
    assert len(indexers) == 0
    return ptr
  if len(indexers) > 1:
    raise NotImplementedError("No support for multiple indexers yet.")
  args_flat, args_tree = tree_util.tree_flatten((ptr, indexers, value, None))
  return _masked_swap_lowering_rule(
      ctx, *args_flat, args_tree=args_tree, eviction_policy=None
  )


def _store(
    ptr: ir.Value,
    value: ir.Value,
    mask: ir.Value | None = None,
    *,
    cache_modifier: str | None = None,
    eviction_policy: str | None = None,
) -> None:
  if cache_modifier is None:
    cache = tt_dialect.CacheModifier.NONE
  elif cache_modifier != ".ca":
    cache = _STR_TO_CACHE_MODIFIER[cache_modifier]
  else:
    raise ValueError(f"unsupported cache modifier: {cache_modifier}")
  if eviction_policy is None:
    evict = tt_dialect.EvictionPolicy.NORMAL
  else:
    try:
      evict = _STR_TO_EVICTION_POLICY[eviction_policy]
    except KeyError:
      raise ValueError(
          f"unsupported eviction policy: {eviction_policy}"
      ) from None

  if _is_triton_pointer_type(ptr.type):
    ptr_type = tt_dialect.PointerType(ptr.type)
    if isinstance(ptr_type.pointee_type, ir.RankedTensorType):
      raise NotImplementedError("loading from a block pointer is not supported")

  ptr_type = _element_type(ptr.type)
  if not _is_triton_pointer_type(ptr_type):
    raise ValueError(f"unsupported pointer type: {ptr_type}")
  ptr_type = tt_dialect.PointerType(ptr_type)
  if not isinstance(ptr.type, ir.RankedTensorType):
    if isinstance(value.type, ir.RankedTensorType):
      raise ValueError("value cannot be a block if pointer is not a block")
    if mask is not None and isinstance(mask.type, ir.RankedTensorType):
      raise ValueError("mask cannot be a block if pointer is not a block")

  pointee_type = ptr_type.pointee_type
  if isinstance(pointee_type, ir.IntegerType) and pointee_type.width == 1:
    pointee_type = ir.IntegerType.get_signless(8)
    ptr = _ir_cast(
        ptr,
        tt_dialect.PointerType.get(pointee_type, ptr_type.address_space),
        signed=False,
    )

  value = _ir_cast(value, pointee_type, signed=False)
  tt_dialect.store(ptr, value, mask=mask, cache=cache, evict=evict)


@register_lowering(primitives.swap_p)
def _masked_swap_lowering_rule(
    ctx: LoweringRuleContext, *args_flat, args_tree, eviction_policy
):
  block_info, *_ = ctx.block_infos
  assert block_info is not None
  ptr, indexers, value, mask = args_tree.unflatten(args_flat)
  ref_aval, _, value_aval, mask_aval = args_tree.unflatten(ctx.avals_in)
  if len(indexers) > 1:
    raise NotImplementedError("No support for multiple indexers yet.")
  if not indexers:
    idx = NDIndexer.make_trivial_indexer(ref_aval.shape)
  else:
    idx = indexers[0]
  ptr = _compute_pointers_from_indices(ptr, block_info, idx)
  other = None
  if value is not None:
    value = _ensure_ir_value(value, value_aval)
  if mask is not None:
    shape = idx.get_indexer_shape_static()
    mask = _bcast_to(_ensure_ir_value(mask, mask_aval), shape)
    if value is not None:
      other = _bcast_to(value, shape)

  old_value = _load(ptr, mask=mask, other=other)
  _store(ptr, value, mask=mask, eviction_policy=eviction_policy)  # pyrefly: ignore[bad-argument-type]
  return old_value


@register_lowering(sp.addupdate_p)
def _addupdate_lowering_rule(ctx: LoweringRuleContext, ptr, value, *idx, tree):
  block_info, *_ = ctx.block_infos
  assert block_info is not None
  indexers = tree_util.tree_unflatten(tree, idx)
  if not _is_triton_pointer_type(ptr.type):
    assert len(indexers) == 0
    return ptr
  if len(indexers) > 1:
    raise NotImplementedError("No support for multiple indexers yet.")
  indexer = indexers[0]
  ptr = _compute_pointers_from_indices(ptr, block_info, indexer)
  op = tt_dialect.RMWOp.FADD
  if isinstance(_element_type(value.type), ir.IntegerType):
    op = tt_dialect.RMWOp.ADD
  _atomic_rmw(op, ptr, value)
  return []


@register_lowering(lax.transpose_p)
def _transpose_lowering(ctx: LoweringRuleContext, x, *, permutation):
  return tt_dialect.trans(x, permutation)


_TF32_PRECISIONS = (lax.Precision.HIGH, lax.Precision.DEFAULT)


def _as_bf16(x):
  return _ir_cast(x, _dtype_to_ir_type(jnp.bfloat16), signed=False)


def _as_f32(x):
  return _ir_cast(x, _dtype_to_ir_type(jnp.float32), signed=False)


@register_lowering(lax.dot_general_p)
def _dot_general_lowering(
    ctx: LoweringRuleContext,
    a,
    b,
    *,
    dimension_numbers,
    out_sharding,
    precision,
    preferred_element_type,
):
  for aval in ctx.avals_in:
    if jnp.issubdtype(aval.dtype, jnp.unsignedinteger):
      raise NotImplementedError(
          f"Unsigned integer dtype {aval.dtype} is not supported for"
          " dot_general (matmul) on the Pallas Triton GPU backend because"
          " dot_general interprets all integer inputs as signed. Consider"
          " casting to a signed type before the dot operation."
      )
  del preferred_element_type, out_sharding  # Unused.
  ((a_contract_dim,), (b_contract_dim,)), batch_dims = dimension_numbers
  assert batch_dims == ((), ())

  if a_contract_dim == 0:
    a = tt_dialect.trans(a, (1, 0))
  if b_contract_dim == 1:
    b = tt_dialect.trans(b, (1, 0))

  a_aval, b_aval = ctx.avals_in
  [out_aval] = ctx.avals_out

  if precision is None or (precision == lax.DotAlgorithmPreset.DEFAULT):
    precision = (lax.Precision.DEFAULT, lax.Precision.DEFAULT)

  if isinstance(precision, lax.DotAlgorithmPreset):
    match precision:
      case lax.DotAlgorithmPreset.TF32_TF32_F32:
        input_precision = tt_dialect.InputPrecision.TF32
      case lax.DotAlgorithmPreset.TF32_TF32_F32_X3:
        input_precision = tt_dialect.InputPrecision.TF32x3
      case lax.DotAlgorithmPreset.F32_F32_F32:
        input_precision = tt_dialect.InputPrecision.IEEE
      case (
          lax.DotAlgorithmPreset.F16_F16_F16
          | lax.DotAlgorithmPreset.F16_F16_F32
          | lax.DotAlgorithmPreset.BF16_BF16_BF16
          | lax.DotAlgorithmPreset.BF16_BF16_F32
          | lax.DotAlgorithmPreset.BF16_BF16_F32_X3
          | lax.DotAlgorithmPreset.BF16_BF16_F32_X6
          | lax.DotAlgorithmPreset.BF16_BF16_F32_X9
      ):
        input_precision = None
      case _:
        raise NotImplementedError(f"Unsupported dot algorithm: {precision}.")

    assert precision.supported_lhs_types is not None
    assert precision.supported_rhs_types is not None
    a = _cast(a, a_aval.dtype, precision.supported_lhs_types[0])
    b = _cast(b, b_aval.dtype, precision.supported_rhs_types[0])
    acc_dtype = precision.accumulation_type
  elif isinstance(precision, tuple):
    a_precision, b_precision = precision
    if a_precision in _TF32_PRECISIONS or b_precision in _TF32_PRECISIONS:
      input_precision = tt_dialect.InputPrecision.TF32
    elif a_aval.dtype == jnp.float32:
      input_precision = tt_dialect.InputPrecision.IEEE
    else:
      input_precision = None

    acc_dtype = out_aval.dtype
    if acc_dtype not in (jnp.int32, jnp.float16, jnp.float64):
      acc_dtype = jnp.float32
  else:
    raise NotImplementedError(f"Unsupported dot precision: {precision}.")

  a_type = ir.RankedTensorType(a.type)
  b_type = ir.RankedTensorType(b.type)
  if len(a_type.shape) != 2 or len(b_type.shape) != 2:
    raise ValueError("a and b must be 2D, but got:"
                     f" {a_type.shape} and {b_type.shape}")

  m, k = a_type.shape
  _, n = b_type.shape
  if a_type.element_type == ir.F64Type.get():
    # Triton's MMAv2 fp64 path uses the m8n8k4 PTX instruction but aggregates
    # it with NumRegisters={m:2, n:1, k:4}, producing an effective m16n8k16
    # per-warp tile.  Blocks smaller than these minimums cause repM/repN/repK
    # to round to zero, corrupting the ValueTable and segfaulting the compiler.
    #   M >= 16  (2 × instrM=8)
    #   N >=  8  (1 × instrN=8)
    #   K >= 16  (4 × instrK=4)
    errors = []
    if m < 16:
      errors.append(f"M={m} < 16")
    if n < 8:
      errors.append(f"N={n} < 8")
    if k < 16:
      errors.append(f"K={k} < 16")
    if errors:
      raise ValueError(
          f"float64 dot requires M>=16, N>=8, K>=16 per warp tile "
          f"(Triton MMAv2 m8n8k4 layout); got {', '.join(errors)}"
      )

  if a_type.element_type != b_type.element_type:
    raise ValueError(
        "a and b must have the same element type, but got:"
        f" {a_type.element_type} and {b_type.element_type}"
    )

  assert acc_dtype is not None
  acc = _zeros(ir.RankedTensorType.get([m, n], _dtype_to_ir_type(acc_dtype)))

  if precision in (
      lax.DotAlgorithmPreset.BF16_BF16_F32_X3,
      lax.DotAlgorithmPreset.BF16_BF16_F32_X6,
      lax.DotAlgorithmPreset.BF16_BF16_F32_X9,
  ):
    a_bf16 = _as_bf16(a)
    b_bf16 = _as_bf16(b)
    a_err0 = _sub(a, _as_f32(a_bf16))
    b_err0 = _sub(b, _as_f32(b_bf16))
    a_err0_bf16 = _as_bf16(a_err0)
    b_err0_bf16 = _as_bf16(b_err0)
    a_err1_bf16 = _as_bf16(_sub(a_err0, _as_f32(a_err0_bf16)))
    b_err1_bf16 = _as_bf16(_sub(b_err0, _as_f32(b_err0_bf16)))
    # Accumulate the smallest values first to reduce the numeric error.
    if precision == lax.DotAlgorithmPreset.BF16_BF16_F32_X9:
      acc = tt_dialect.dot(a_err1_bf16, b_err0_bf16, acc)
      acc = tt_dialect.dot(a_err1_bf16, b_err1_bf16, acc)
      acc = tt_dialect.dot(a_err0_bf16, b_err1_bf16, acc)
    if precision in (
        lax.DotAlgorithmPreset.BF16_BF16_F32_X6,
        lax.DotAlgorithmPreset.BF16_BF16_F32_X9,
    ):
      acc = tt_dialect.dot(a_err1_bf16, b_bf16, acc)
      acc = tt_dialect.dot(a_bf16, b_err1_bf16, acc)
      acc = tt_dialect.dot(a_err0_bf16, b_err0_bf16, acc)
    acc = tt_dialect.dot(a_err0_bf16, b_bf16, acc)
    acc = tt_dialect.dot(a_bf16, b_err0_bf16, acc)
    # If `a` rounding error is zero and `b` is `inf` then `acc` may contain
    # `NaN`s (as `0 * inf = NaN`), and vice versa.
    acc = arith_dialect.select(_is_nan(acc), _zeros_like(acc), acc)
    a, b = a_bf16, b_bf16

  acc = tt_dialect.dot(a, b, acc, input_precision=input_precision)
  return _cast(acc, acc_dtype, out_aval.dtype)


def _reduction_lowering(body, ctx: LoweringRuleContext, a, axes):
  flat_args = tree_util.tree_leaves(a)
  (axis,) = axes
  mapped_avals = [jax_core.ShapedArray((), aval.dtype) for aval in ctx.avals_in]
  in_tree = tree_util.tree_structure((a, a))
  flat_fun, out_tree_thunk = api_util.flatten_fun_nokwargs(
      lu.wrap_init(
          body,
          debug_info=api_util.debug_info("pallas triton reduction",
                                         body, (a, a), {})),
      in_tree
  )
  combine_jaxpr, _, consts = pe.trace_to_jaxpr_dynamic(
      flat_fun, [*mapped_avals, *mapped_avals]
  )
  out_tree = out_tree_thunk()
  del out_tree  # Not needed
  if consts:
    raise NotImplementedError("Reductions with constants not supported.")
  element_types = [_element_type(arg.type) for arg in flat_args]
  reduce_op = tt_dialect.ReduceOp(flat_args, axis)
  param_types = element_types * 2
  entry = reduce_op.regions[0].blocks.append(*param_types)
  with ir.InsertionPoint.at_block_begin(entry):
    results = lower_jaxpr_to_triton_ir(
        ctx.context, combine_jaxpr, None, *entry.arguments
    )
    tt_dialect.reduce_return(results)
  reduce_op.verify()
  return list(reduce_op.result)


def _reduce_lowering(body, ctx: LoweringRuleContext, a, *, axes, **kwargs):
  assert isinstance(axes, tuple)
  if not axes:
    return a
  while len(axes) > 1:
    axis = max(axes)
    dst_avals = tuple(v.update(shape=v.shape[:axis] + v.shape[axis + 1:])
                      for v in ctx.avals_in)
    a = _reduce_lowering(
        body, ctx.replace(avals_out=dst_avals), a, axes=(axis,))
    # Adding an intervening -(-reduce(.)) introduces a convert_layout between
    # reduces, which seems necessary for correctness.
    # TODO(bjp): Get rid of the double negation.
    #     https://github.com/openai/triton/issues/1776
    a = _minus(_minus(a))
    ctx = ctx.replace(avals_in=dst_avals)
    axes = tuple(ax for ax in axes if ax != axis)
  return _reduction_lowering(body, ctx, a, axes=axes)[0]


triton_lowering_rules[lax.reduce_max_p] = functools.partial(
    _reduce_lowering, jnp.maximum
)
triton_lowering_rules[lax.reduce_min_p] = functools.partial(
    _reduce_lowering, jnp.minimum
)
triton_lowering_rules[lax.reduce_sum_p] = functools.partial(
    _reduce_lowering, jnp.add
)


def _argreduce_lowering(
    body, ctx: LoweringRuleContext, a, *, axes, index_dtype
):
  if index_dtype != jnp.int32:
    raise ValueError("`index_type` must be i32.")
  if len(axes) != 1:
    raise ValueError("`pallas` reduce operations only support one reduce axis.")
  [axis] = axes
  [a_aval] = ctx.avals_in
  index = _make_range(0, a_aval.shape[axis])
  if len(a_aval.shape) > 1:
    # Broadcast index across the non-reduced axes
    for i in range(len(a_aval.shape)):
      if i != axis:
        index = _expand_dims(index, i)
    index = _bcast_to(index, a_aval.shape)
  ctx = ctx.replace(avals_in=[a_aval, a_aval.update(dtype=jnp.dtype(jnp.int32))])
  _, indices = _reduction_lowering(body, ctx, (a, index), axes=axes)
  return indices


def _reduce_argmax_combine(left, right):
  value1, index1 = left
  value2, index2 = right
  gt = value1 > value2
  lt = value1 < value2
  index_min = jnp.minimum(index1, index2)
  index_ret = jnp.where(gt, index1, jnp.where(lt, index2, index_min))
  value_ret = jnp.maximum(value1, value2)
  return value_ret, index_ret


triton_lowering_rules[lax.argmax_p] = functools.partial(
    _argreduce_lowering, _reduce_argmax_combine
)


def _reduce_argmin_combine(left, right):
  value1, index1 = left
  value2, index2 = right
  gt = value1 > value2
  lt = value1 < value2
  index_min = jnp.minimum(index1, index2)
  index_ret = jnp.where(lt, index1, jnp.where(gt, index2, index_min))
  value_ret = jnp.minimum(value1, value2)
  return value_ret, index_ret


triton_lowering_rules[lax.argmin_p] = functools.partial(
    _argreduce_lowering, _reduce_argmin_combine
)


@register_lowering(pjit.jit_p)
def _pjit_lowering_rule(ctx: LoweringRuleContext, *args, jaxpr, **_):
  if jaxpr.consts:
    raise NotImplementedError
  return lower_jaxpr_to_triton_ir(
      ctx.context, jaxpr.jaxpr, ctx.block_infos, *args
  )


@register_lowering(pjit.reshard_p)
def _reshard_lowering_rule(ctx, x, *, dst_sharding, concrete_mesh):
  return x


@register_lowering(jax_core.closed_call_p)
@register_lowering(custom_derivatives.custom_jvp_call_p)
def _closed_call_lowering_rule(
    ctx: LoweringRuleContext, *args, call_jaxpr, **_
):
  jaxpr, consts = call_jaxpr.jaxpr, call_jaxpr.consts
  if consts:
    raise NotImplementedError
  return lower_jaxpr_to_triton_ir(ctx.context, jaxpr, ctx.block_infos, *args)


@register_lowering(ad_checkpoint.remat_p)
def _remat_lowering_rule(ctx: LoweringRuleContext, *args, jaxpr, **_):
  return lower_jaxpr_to_triton_ir(ctx.context, jaxpr, ctx.block_infos, *args)


triton_lowering_rules[ad_util.stop_gradient_p] = lambda _, x: x


@register_lowering(lax.axis_index_p)
def _axis_index_rule(ctx: LoweringRuleContext, *, axis_name: Hashable):
  grid_names = ctx.context.grid_mapping.grid_names
  if grid_names is not None and axis_name in grid_names:
    # We are querying a named axis corresponding to a grid dimension.
    return _program_id_lowering_rule(ctx, axis=grid_names.index(axis_name))
  raise LookupError(f"Axis name {axis_name} not found in grid.")


def _lower_jaxpr_to_for_loop(
    ctx: LoweringRuleContext,
    jaxpr: jax_core.Jaxpr,
    lower_bound,
    upper_bound,
    consts,
    *args,
    has_loop_index: bool,
    step: int = 1,
    bound_type: ir.IntegerType | None = None,
):
  if step != 1:
    raise NotImplementedError
  if bound_type is None or bound_type.width == 32:
    step_val = _i32_constant(step)
  else:
    step_val = _i64_constant(step)

  for_op = scf_dialect.ForOp(lower_bound, upper_bound, step_val, args)
  with ir.InsertionPoint.at_block_begin(for_op.body):
    loop_index = for_op.induction_variable
    for_body_args = [for_op.body.arguments[i + 1] for i, _ in enumerate(args)]
    if has_loop_index:
      jaxpr_args = [*consts, loop_index, *for_body_args]
    else:
      jaxpr_args = [*consts, *for_body_args]
    all_out = lower_jaxpr_to_triton_ir(
        ctx.context, jaxpr, ctx.block_infos, *jaxpr_args
    )
    scf_dialect.yield_(all_out)

  return list(for_op.results_)


@register_lowering(lax.scan_p)
def _scan_lowering_rule(
    ctx: LoweringRuleContext,
    *args,
    jaxpr,
    length,
    reverse,
    unroll,
    num_consts,
    num_carry,
):
  # Only implements fori_loop-like scans
  if reverse: raise NotImplementedError
  if unroll != 1: raise NotImplementedError
  del unroll, reverse

  jaxpr, jaxpr_consts = jaxpr.jaxpr, jaxpr.consts
  if jaxpr_consts: raise NotImplementedError
  del jaxpr_consts

  jaxpr, has_loop_index = (
      pallas_utils.pattern_match_scan_to_fori_loop(jaxpr, num_consts, num_carry)
  )
  args = map(_ensure_ir_value, args, ctx.avals_in)
  consts, args = util.split_list(args, [num_consts])
  if has_loop_index:
    lower_bound, *args = args
    upper_bound = _add(lower_bound, _ir_constant(length, lower_bound.type))
    bound_type = lower_bound.type
    assert isinstance(bound_type, ir.IntegerType)
  else:
    lower_bound = _i32_constant(0)
    upper_bound = _i32_constant(length)
    bound_type = ir.IntegerType.get_signless(32)
  for_out = _lower_jaxpr_to_for_loop(
      ctx, jaxpr, lower_bound, upper_bound, consts, *args,
      has_loop_index=has_loop_index, step=1, bound_type=bound_type)
  if has_loop_index:
    # Need to return the final loop index value if the outer scan expects
    # it as an output
    return [upper_bound, *for_out]
  return for_out


def _maybe_pattern_match_fori_loop(
    ctx: LoweringRuleContext,
    *args,
    cond_nconsts,
    cond_jaxpr,
    body_nconsts,
    body_jaxpr,
):
  if cond_nconsts:
    return None
  _, cond_invars = split_list(cond_jaxpr.jaxpr.invars, [cond_nconsts])
  cond_in_avals = [v.aval for v in cond_invars]
  if len(cond_in_avals) < 2:
    return None
  # Check that the first two carry values are scalar ints
  a1, a2 = cond_in_avals[:2]
  if a1.shape != () or a1.dtype not in (jnp.int32, jnp.int64):
    return None
  if a2.shape != () or a2.dtype not in (jnp.int32, jnp.int64):
    return None
  # Check that the only eqn in the cond checks the loop index condition
  v1, v2 = cond_invars[:2]
  outvar = cond_jaxpr.jaxpr.outvars[0]
  assert outvar.aval.dtype == jnp.bool_
  if len(cond_jaxpr.jaxpr.eqns) != 1:
    return None
  eqn = cond_jaxpr.jaxpr.eqns[0]
  if eqn.primitive != lax.lt_p:
    return None
  if eqn.outvars != [outvar]:
    return None
  if eqn.invars != [v1, v2]:
    return None
  # Check that the carry is updated in the body appropriately
  _, body_invars = split_list(body_jaxpr.jaxpr.invars, [body_nconsts])
  v1, v2 = body_invars[:2]
  vo1, vo2 = body_jaxpr.jaxpr.outvars[:2]
  # Upper bound should be constant
  if v2 is not vo2:
    return None
  # Check that we increment the loop index in the body
  for i, eqn in enumerate(body_jaxpr.jaxpr.eqns):
    if eqn.primitive is lax.add_p:
      if eqn.invars[0] is v1:
        if isinstance(eqn.invars[1], jax_core.Literal):
          if eqn.invars[1].val == 1:
            if eqn.outvars[0] == vo1:
              eqn_index = i
              break
  else:
    return None
  jaxpr = body_jaxpr.jaxpr
  new_invars = (*jaxpr.invars[:body_nconsts],
                jaxpr.invars[body_nconsts],
                *jaxpr.invars[body_nconsts + 2:])
  new_outvars = tuple(jaxpr.outvars[2:])
  jaxpr = jaxpr.replace(
      eqns=jaxpr.eqns[:eqn_index] + jaxpr.eqns[eqn_index + 1:],
      invars=new_invars,
      outvars=new_outvars,
      debug_info=jaxpr.debug_info.with_unknown_names())
  _, body_consts, carry = split_list(args, [cond_nconsts, body_nconsts])
  (lb, ub), args = carry[:2], carry[2:]
  const_block_infos, args_block_infos = split_list(ctx.block_infos,
                                                   [body_nconsts])
  ctx = ctx.replace(block_infos=[*const_block_infos, None,
                                 *args_block_infos[2:]])
  for_out = _lower_jaxpr_to_for_loop(
      ctx,
      jaxpr,
      lb,
      ub,
      body_consts,
      *args,
      has_loop_index=True,
      step=1,
      bound_type=lb.type,
  )
  return [ub, ub, *for_out]


@register_lowering(lax.while_p)
def _while_lowering_rule(
    ctx: LoweringRuleContext,
    *args,
    cond_nconsts,
    cond_jaxpr,
    body_nconsts,
    body_jaxpr,
):
  args = map(_ensure_ir_value, args, ctx.avals_in)

  # First, try to pattern match to fori_loop and lower to scf.for if possible
  # TODO(slebedev): Use `pallas_utils.pattern_match_while_to_fori_loop`.
  result = _maybe_pattern_match_fori_loop(ctx, *args, cond_nconsts=cond_nconsts,
                                          body_nconsts=body_nconsts, cond_jaxpr=cond_jaxpr,
                                          body_jaxpr=body_jaxpr)
  if result is not None:
    return result
  # Fall back to default while lowering
  cond_consts, body_consts, carry = util.split_list(
      args, [cond_nconsts, body_nconsts]
  )
  cond_const_block_infos, body_const_block_infos, carry_block_infos = (
      util.split_list(ctx.block_infos, [cond_nconsts, body_nconsts])
  )
  cond_const_types = [a.type for a in cond_consts]
  body_const_types = [a.type for a in body_consts]
  carry_types = [a.type for a in carry]
  all_types = [*cond_const_types, *body_const_types, *carry_types]
  while_op = scf_dialect.WhileOp(all_types, args)

  before_block = while_op.before.blocks.append(*all_types)
  cond_consts_, _, carry_ = util.split_list(
      before_block.arguments,
      [cond_nconsts, body_nconsts],
  )
  cond_args = [*cond_consts_, *carry_]
  with ir.InsertionPoint.at_block_begin(before_block):
    [cond] = lower_jaxpr_to_triton_ir(
        ctx.context,
        cond_jaxpr.jaxpr,
        [*cond_const_block_infos, *carry_block_infos],
        *cond_args,
    )
    scf_dialect.condition(cond, before_block.arguments)

  after_block = while_op.after.blocks.append(*all_types)
  cond_consts_, body_consts_, carry_ = util.split_list(
      after_block.arguments,
      [cond_nconsts, body_nconsts],
  )
  all_args = [*cond_consts_, *body_consts_, *carry_]
  cond_const_args, body_const_args, carry_args = util.split_list(
      all_args, [cond_nconsts, body_nconsts]
  )
  with ir.InsertionPoint.at_block_begin(after_block):
    loop_out = lower_jaxpr_to_triton_ir(
        ctx.context,
        body_jaxpr.jaxpr,
        [*body_const_block_infos, *carry_block_infos],
        *body_const_args,
        *carry_args
    )
    all_handles = [*cond_const_args, *body_const_args, *loop_out]
    if all_handles:
      scf_dialect.yield_(all_handles)

  all_out = list(while_op.results_)
  return all_out[cond_nconsts + body_nconsts :]


@register_lowering(lax.cond_p)
def _cond_lowering_rule(
    ctx: LoweringRuleContext,
    index,
    *args,  # *consts, *ops
    branches,  # tuple(jaxprs)
):
  block_infos = ctx.block_infos

  def to_type(out_aval):
    element_type = _dtype_to_ir_type(out_aval.dtype)
    if not out_aval.shape:
      return element_type
    return ir.RankedTensorType.get(out_aval.shape, element_type)

  out_types = [to_type(out) for out in ctx.avals_out]

  use_branch0 = _equal(index, _ir_constant(0, index.type), signed=False)
  # TODO(bjp): Switch to scf.index_switch once exposed in triton.cc
  if_op = scf_dialect.IfOp(use_branch0, out_types, has_else=True)
  with ir.InsertionPoint.at_block_begin(if_op.then_block):
    outs0 = lower_jaxpr_to_triton_ir(
        ctx.context,
        branches[0].jaxpr,
        block_infos[1:],
        *args)
    scf_dialect.yield_(outs0)
  assert if_op.else_block is not None
  with ir.InsertionPoint.at_block_begin(if_op.else_block):
    # TODO(bjp): Instead of linear nest of 'if's, partition into halves.
    if len(branches) > 2:
      outs1 = _cond_lowering_rule(
          ctx,
          _sub(index, _ir_constant(1, index.type)),
          *args,
          branches=branches[1:],
      )
    else:
      outs1 = lower_jaxpr_to_triton_ir(
          ctx.context,
          branches[1].jaxpr,
          block_infos[1:],
          *args)
    scf_dialect.yield_(outs1)

  return list(if_op.results_)


def _ensure_ir_value(x: object, aval: jax_core.ShapedArray) -> ir.Value:
  if isinstance(x, ir.Value):
    return x
  elif isinstance(
      x, (np.number, np.ndarray, int, float, literals.TypedNdArray)
  ):
    return _ir_constant(x, _dtype_to_ir_type(aval.dtype))
  raise NotImplementedError


def _ir_constant(v: object, t: ir.Type) -> ir.Value:
  if isinstance(
      v, (np.number, np.ndarray, int, float, literals.TypedNdArray)
  ):
    if isinstance(t, ir.IntegerType):
      v = int(v)
    else:
      assert isinstance(t, ir.FloatType)
      v = float(v)
    return arith_dialect.constant(t, v)
  raise NotImplementedError


def _i32_constant(v: int) -> ir.Value:
  return arith_dialect.constant(ir.IntegerType.get_signless(32), v)


def _i64_constant(v: int) -> ir.Value:
  return arith_dialect.constant(ir.IntegerType.get_signless(64), v)


def _dtype_to_ir_type(dtype: jax.typing.DTypeLike) -> ir.Type:
  dtype = jnp.dtype(dtype)
  if jnp.issubdtype(dtype, np.integer):
    # All integer types in Triton are signless.
    return ir.IntegerType.get_signless(dtype.itemsize * 8)
  return mlir.dtype_to_ir_type(dtype)


@register_lowering(lax.bitcast_convert_type_p)
def _bitcast_convert_type_lowering_rule(
    ctx: LoweringRuleContext, operand: ir.Value, *, new_dtype
) -> ir.Value:
  # TODO(petebu) Handle case where src and dst types have different bitwidths
  src_elem_type = _element_type(operand.type)
  dst_elem_type = _element_type(_dtype_to_ir_type(new_dtype))
  assert isinstance(src_elem_type, (ir.IntegerType, ir.FloatType))
  assert isinstance(dst_elem_type, (ir.IntegerType, ir.FloatType))
  if src_elem_type.width != dst_elem_type.width:
    raise NotImplementedError(
        f"cannot cast {operand} to {new_dtype} because of different widths"
    )
  if isinstance(operand.type, ir.RankedTensorType):
    shape = ir.RankedTensorType(operand.type).shape
    result_type = ir.RankedTensorType.get(shape, dst_elem_type)
  else:
    result_type = dst_elem_type
  return tt_dialect.bitcast(result_type, operand)
