# Copyright 2024 The JAX Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Utilities for code generator."""

from collections.abc import Iterator, Sequence
import contextlib
import dataclasses
import enum
import functools
import itertools
import logging
import math
import os
from typing import Any, Literal, cast, overload

import jax
from jax import numpy as jnp
from jax._src.lib import mosaic_gpu_dialect as dialect  # noqa: F401
from jax.extend import backend as jex_backend
from jax.interpreters import mlir
from jaxlib.mlir import ir
from jaxlib.mlir.dialects import arith
from jaxlib.mlir.dialects import builtin
from jaxlib.mlir.dialects import gpu
from jaxlib.mlir.dialects import llvm
from jaxlib.mlir.dialects import memref
from jaxlib.mlir.dialects import nvvm
from jaxlib.mlir.dialects import scf
from jaxlib.mlir.dialects import vector
import numpy as np

logger = logging.getLogger(__name__)

WARP_SIZE: int = 32
WARPGROUP_SIZE: int = 128
WARPS_IN_WARPGROUP: int = WARPGROUP_SIZE // WARP_SIZE
DYNAMIC = -9223372036854775808
DYNAMIC32 = -2147483648
MBARRIER_BYTES = 8


def dump_to_file_or_stdout(
    content: str, name: str, path: str
) -> None:
  """Dumps content to path/name if path is non-empty, else to stdout."""
  if not name:
    raise ValueError("name must be non-empty")
  if not path:
    print(content)
    return
  filepath = os.path.join(path, name)
  try:
    with open(filepath, "w", encoding="utf-8") as f:
      f.write(content)
      f.write("\n")
  except OSError as e:
    logger.error("Failed to write output to %s: %s", filepath, e)
    # TODO(bchetioui): revisit whether this default of writing to stdout is the
    # right one. If we change it, we will have to change the corresponding C++
    # implementation as well.
    logger.error("Output will be written to stdout instead.")
    print(content)


def _ptx_constraint(t: ir.Type) -> str:
  """Infers the LLVM inline asm constraint string for a given MLIR type in PTX."""
  if isinstance(t, ir.IntegerType):
    width = ir.IntegerType(t).width
    return {1: "b", 16: "h", 32: "r", 64: "l"}[width]
  if isinstance(t, ir.FloatType):
    width = ir.FloatType(t).width
    return {16: "h", 32: "f"}[width]
  if isinstance(t, llvm.PointerType):
    # Treat global memory pointers as 64-bit, and others as 32-bit. This is not
    # entirely correct, since SMEM pointers may also be using 64 bits unless in
    # short address mode (see computeNVPTXDataLayout in
    # llvm/llvm-project/llvm/lib/TargetParser/TargetDataLayout.cpp), but it is
    # practically the most convenient mapping---as PTX intrinsics mostly expect
    # these pointers to be provided as 32-bit values anyway.
    return "l" if t.address_space in [0, 1] else "r"
  raise NotImplementedError(f"Unknown PTX constraint for type: {t}")


@overload
def inline_ptx(
    ptx: str,
    *args: ir.Value,
    result_types: ir.Type,
    predicate: None = ...,
    has_side_effects: bool = ...,
    convergent: bool = ...,
) -> ir.Value:
  ...


@overload
def inline_ptx(
    ptx: str,
    *args: ir.Value,
    result_types: Sequence[ir.Type],
    predicate: None = ...,
    has_side_effects: bool = ...,
    convergent: bool = ...,
) -> tuple[ir.Value, ...]:
  ...


@overload
def inline_ptx(
    ptx: str,
    *args: ir.Value,
    result_types: None = ...,
    predicate: ir.Value | None = ...,
    has_side_effects: Literal[True] = ...,
    convergent: bool = ...,
) -> None:
  ...


def inline_ptx(
    ptx: str,
    *args: ir.Value,
    result_types: Sequence[ir.Type] | ir.Type | None = None,
    predicate: ir.Value | None = None,
    has_side_effects: bool = False,
    convergent: bool = False,
) -> ir.Value | tuple[ir.Value, ...] | None:
  """Emits an LLVM inline assembly operation targeting PTX.

  Args:
    ptx: The PTX assembly template string, in the format expected by
      `llvm.inline_asm`.
    *args: `ir.Value`s passed as operands to the assembly.
    result_types: The output type, sequence of output types, or `None` if void.
    predicate: An optional predicate to prepend to each PTX instruction.
      This must be `None` if `result_types` is not `None`.
    has_side_effects: Whether the inline asm has side effects. If the assembly
      does not return any result, this must be `True`.
    convergent: Whether the inline asm is convergent, i.e. whether it interacts
      with a dynamically-determined set of threads (barriers, shuffles, and
      other warp/warpgroup collectives). This prevents passes from moving or
      merging it across divergent control flow. `has_side_effects` does not
      imply this.

  Returns:
    `None` if `result_types` is `None`, a single `ir.Value` if `result_types` is
    a single `ir.Type`, or a tuple of `ir.Value`s if `result_types` is a
    sequence of types.
  """
  if isinstance(result_types, ir.Type):
    normalized_types = (result_types,)
  elif result_types is None:
    normalized_types = ()
    # This is consistent with the overload we define---`has_side_effects` is
    # defined to be `True` when `result_types` is `None`.
    has_side_effects = True
  else:
    normalized_types = tuple(result_types)

  # An inline asm with no results and no side effects is a no-op, which is
  # almost certainly a mistake (this covers `result_types=()`).
  if not normalized_types and not has_side_effects:
    raise ValueError(
        "`result_types` must be specified if `has_side_effects` is False."
    )

  if not normalized_types:
    asm_ret_type = ir.Type.parse("!llvm.void")
  elif len(normalized_types) == 1:
    asm_ret_type = normalized_types[0]
  else:
    asm_ret_type = llvm.StructType.get_literal(normalized_types)

  if predicate is not None:
    if normalized_types:
      raise ValueError("predicate must be None if result_types is not None.")

    if ptx.count(";") != 1:
      raise NotImplementedError(
          "predicate is not yet supported for multi-line PTX: " + ptx
      )

    ptx = f"@${len(args)} {ptx}"
    args = (*args, predicate)

  out_constraints = [f"={_ptx_constraint(t)}" for t in normalized_types]
  in_constraints = [_ptx_constraint(arg.type) for arg in args]
  constraints = ",".join(out_constraints + in_constraints)

  result = llvm.inline_asm(
      asm_ret_type,
      args,
      ptx,
      constraints,
      has_side_effects=has_side_effects,
      convergent=convergent,
  )
  if result_types is None:
    return None
  assert isinstance(result, ir.Value)
  if isinstance(result_types, ir.Type):
    return result
  assert isinstance(result_types, Sequence)
  if len(result_types) == 1:
    return (result,)
  return tuple(
      llvm.extractvalue(t, result, [i])
      for i, t in enumerate(normalized_types)
  )


def gpu_address_space_to_nvptx(address_space: gpu.AddressSpace) -> int:
  match address_space:
    case gpu.AddressSpace.Global:
      return 1
    case gpu.AddressSpace.Workgroup:
      return 3
    case _:
      raise NotImplementedError(f"address_space not supported: {address_space}")


WORKGROUP_NVPTX_ADDRESS_SPACE = gpu_address_space_to_nvptx(
    gpu.AddressSpace.Workgroup
)


def memref_descriptor_type(memref_ty: ir.MemRefType) -> llvm.StructType:
  i64 = ir.IntegerType.get_signless(64)
  rank = len(memref_ty.shape)
  address_space = get_memref_llvm_address_space(memref_ty)
  ptr_ty = llvm.PointerType.get(address_space)
  desc_ty_fields = [ptr_ty, ptr_ty, i64]
  if rank > 0:
    desc_ty_fields += [llvm.ArrayType.get(i64, rank)] * 2
  return llvm.StructType.get_literal(desc_ty_fields)


def to_memref_descriptor(memref_arg: ir.Value) -> ir.Value:
  """Casts a memref to its underlying LLVM struct descriptor.

  The returned struct has the following fields (matching MLIR's MemRefToLLVM
  lowering convention):
    [0]: allocated pointer (!llvm.ptr)
    [1]: aligned base pointer (!llvm.ptr)
    [2]: offset in elements (i64)
    [3]: sizes in elements (!llvm.array<rank x i64>, only if rank > 0)
    [4]: strides in elements (!llvm.array<rank x i64>, only if rank > 0)
  """
  memref_ty = ir.MemRefType(memref_arg.type)
  desc_ty = memref_descriptor_type(memref_ty)
  desc = builtin.unrealized_conversion_cast([desc_ty], [memref_arg])
  assert isinstance(desc, ir.Value)
  return desc


def ptr_as_memref(ptr, memref_ty: ir.MemRefType):
  ptr_ty = llvm.PointerType(ptr.type)
  if ptr_ty.address_space != (get_memref_llvm_address_space(memref_ty) or 0):
    raise ValueError(
        f"Pointer address space {ptr_ty.address_space} does not match "
        f"memref memory space {memref_ty.memory_space}."
    )

  strides, offset = memref_ty.get_strides_and_offset()
  if offset != 0:
    raise ValueError("Non-zero offset is not supported for ptr_as_memref")
  i64 = ir.IntegerType.get_signless(64)
  rank = len(memref_ty.shape)
  desc_ty = memref_descriptor_type(memref_ty)
  desc = llvm.mlir_undef(desc_ty)
  desc = llvm.insertvalue(desc, ptr, [0])  # Allocation
  desc = llvm.insertvalue(desc, ptr, [1])  # Aligned Base
  desc = llvm.insertvalue(desc, c(0, i64), [2])
  if rank > 0:
    for i, s in enumerate(memref_ty.shape):
      desc = llvm.insertvalue(desc, c(s, i64), [3, i])
    for i, s in enumerate(strides):
      desc = llvm.insertvalue(desc, c(s, i64), [4, i])
  result = builtin.unrealized_conversion_cast([memref_ty], [desc])
  assert isinstance(result, ir.Value)
  return result


def pack_array(values):
  if not values:
    raise ValueError("Empty array")
  elem_ty = values[0].type
  i64 = ir.IntegerType.get_signless(64)
  ptr_ty = ir.Type.parse("!llvm.ptr")
  arr_ptr = llvm.alloca(ptr_ty, c(len(values), i64), elem_ty)
  for i, v in enumerate(values):
    elem_ptr = getelementptr(arr_ptr, [i], elem_ty)
    llvm.store(v, elem_ptr)
  return arr_ptr


def get_contiguous_strides(xs):
  strides_ret = []
  stride = 1
  for x in xs[::-1]:
    strides_ret.append(stride)
    stride *= x
  return strides_ret[::-1]


def c(val: int | float, ty):
  if isinstance(ty, (ir.IntegerType, ir.IndexType)):
    if not isinstance(val, (int, np.integer)):
      raise TypeError(type(val))
    attr = ir.IntegerAttr.get(ty, val)
  elif isinstance(ty, ir.FloatType):
    attr = ir.FloatAttr.get(ty, val)
  elif isinstance(ty, ir.VectorType):
    return vector.broadcast(ty, c(val, ir.VectorType(ty).element_type))
  else:
    raise NotImplementedError(ty)
  return arith.constant(ty, attr)


def _debug_scalar_ty_format(arg):
  if isinstance(arg.type, ir.IndexType):
    return "%llu", arg
  if isinstance(arg.type, ir.IntegerType):
    if ir.IntegerType(arg.type).width < 64:
      arg = arith.extui(ir.IntegerType.get_signless(64), arg)
    return "%llu", arg
  if isinstance(arg.type, ir.F32Type):
    return "%f", arg
  if isinstance(arg.type, ir.Float8E8M0FNUType):
    return "%u", arith.extui(
        ir.IntegerType.get_signless(32),
        arith.bitcast(ir.IntegerType.get_signless(8), arg),
    )
  if isinstance(arg.type, (ir.BF16Type, ir.F16Type)):
    arg = arith.extf(ir.F32Type.get(), arg)
    return "%f", arg
  raise NotImplementedError(f"Can't print the type {arg.type}")


def debug_print(fmt, *args, uniform=True, scope=None):
  if not uniform and scope is not None:
    raise ValueError("Cannot specify scope to a non-uniform debug_print.")
  if scope is None:
    scope = ThreadSubset.WARPGROUP
  type_formats = []
  new_args = []
  for arg in args:
    if isinstance(arg.type, ir.VectorType):
      vec_ty = ir.VectorType(arg.type)
      if len(vec_ty.shape) > 1:
        raise NotImplementedError(
            f"2D+ vectors are not supported in debug_print: {vec_ty}"
        )
      vec_args = [
          vector.extract(
              arg,
              dynamic_position=[],
              static_position=ir.DenseI64ArrayAttr.get([i]),
          )
          for i in range(vec_ty.shape[0])
      ]
      ty_formats, args = zip(*map(_debug_scalar_ty_format, vec_args))
      ty_format = f"[{','.join(ty_formats)}]"
      new_args += args
    else:
      ty_format, arg = _debug_scalar_ty_format(arg)
      new_args.append(arg)

    if ty_format is None:
      raise NotImplementedError(arg.type)
    type_formats.append(ty_format)
  ctx = (
      functools.partial(single_thread, scope=scope)
      if uniform
      else contextlib.nullcontext
  )
  with ctx():
    gpu.printf(fmt.format(*type_formats) + "\n", *new_args)


@dataclasses.dataclass(frozen=True)
class MultimemRef:
  ref: ir.Value[ir.MemRefType]

  @property
  def type(self) -> ir.Type:
    return ir.MemRefType(self.ref.type)

  def store(self, value: ir.Value, indices: Sequence[ir.Value]):
    ptr = memref_ptr(memref_slice(self.ref, tuple(indices)))
    multimem_store(ptr, value)


def multimem_store(ptr: ir.Value, value: ir.Value):
  i32 = ir.IntegerType.get_signless(32)
  if (bw := bitwidth(value.type)) not in {32, 64, 128}:
    raise ValueError("Only 32-, 64- and 128-bit stores are supported")
  vector_length = bw // 32
  value = bitcast(value, ir.VectorType.get((vector_length,), i32))
  regs = [
      llvm.extractelement(value, arith.constant(i32, i))
      for i in range(vector_length)
  ]
  if vector_length == 1:
    vec_ptx = "$1"
    vec_mod = ""
  else:
    vec_ptx = f"{{{','.join(f'${i}' for i in range(1, vector_length + 1))}}}"
    vec_mod = ".v" + str(vector_length)
  # It's unclear to me why, but at least according to PTX docs, we have to use
  # the floating-point instructions here to be able to store vectors.
  inline_ptx(
      f"multimem.st.relaxed.sys.global{vec_mod}.f32 [$0], {vec_ptx};",
      ptr,
      *regs,
      has_side_effects=True,
  )


MultimemReductionOp = Literal["add", "min", "max", "and", "or", "xor"]


def multimem_load_reduce(
    ty: ir.Type,
    ptr: ir.Value,
    reduction: MultimemReductionOp,
    is_signed: bool | None = None,
):
  i32 = ir.IntegerType.get_signless(32)
  if bitwidth(ty) not in {32, 64, 128}:
    raise ValueError("Only 32-, 64- and 128-bit loads are supported")
  if isinstance(ty, ir.VectorType):
    vty = ir.VectorType(ty)
    if len(vty.shape) > 1:
      raise ValueError("Only 1D vectors are supported")
    vector_length = vty.shape[0]
    vector_i32_length = vector_length * bitwidth(vty.element_type) // 32
    if isinstance(vty.element_type, ir.IntegerType):
      if vector_length != 1:
        results = []
        elem_ty = vty.element_type
        for i in range(vector_length):
          elem_ptr = getelementptr(ptr, [i], elem_ty)
          v1_ty = ir.VectorType.get((1,), elem_ty)
          elem_res = multimem_load_reduce(
              v1_ty, elem_ptr, reduction, is_signed=is_signed
          )
          results.append(elem_res)
        return vector_concat(results)
      if bitwidth(vty.element_type) not in {32, 64}:
        raise NotImplementedError(
            "Only 32-bit and 64-bit integer operations are supported"
        )
      if reduction in {"and", "or", "xor"}:
        ptx_ty = f"b{bitwidth(vty.element_type)}"
      elif reduction in {"min", "max", "add"}:
        if is_signed is None:
          raise ValueError(
              "Signedness must be specified for integer min, max and add"
              " reductions"
          )
        ptx_ty = f"{'s' if is_signed else 'u'}{bitwidth(vty.element_type)}"
      else:
        raise ValueError(f"Unsupported reduction operation: {reduction}")
    elif isinstance(vty.element_type, ir.FloatType):
      if reduction not in {"add", "min", "max"}:
        raise ValueError("Only add, min and max are supported for floats")
      if isinstance(vty.element_type, ir.F32Type):
        if reduction != "add":
          raise ValueError("Only add is supported for f32")
        ptx_ty = "f32"
      elif isinstance(vty.element_type, ir.BF16Type):
        ptx_ty = "bf16x2"
      elif isinstance(vty.element_type, ir.F16Type):
        ptx_ty = "f16x2"
      elif isinstance(vty.element_type, ir.Float8E5M2Type):
        ptx_ty = "e5m2x4"
      elif isinstance(vty.element_type, ir.Float8E4M3FNType):
        ptx_ty = "e4m3x4"
      else:
        raise NotImplementedError(vty.element_type)
    else:
      raise NotImplementedError(vty.element_type)
  else:
    raise NotImplementedError(ty)
  if vector_i32_length == 1:
    vec_ptx = "$0"
    vec_mod = ""
  else:
    vec_ptx = f"{{{','.join(f'${i}' for i in range(vector_i32_length))}}}"
    vec_mod = ".v" + str(vector_i32_length)
  # It's unclear to me why, but at least according to PTX docs, we have to use
  # the floating-point instructions here to be able to store vectors.
  acc_prec = ""
  out_regs = inline_ptx(
      f"multimem.ld_reduce.relaxed.sys.global.{reduction}{acc_prec}{vec_mod}.{ptx_ty}"
      f" {vec_ptx}, [${vector_i32_length}];",
      ptr,
      result_types=[i32] * vector_i32_length,
      has_side_effects=True,
  )
  if vector_i32_length == 1:
    return bitcast(out_regs[0], ty)
  vec_i32_ty = ir.VectorType.get((1,), i32)
  return bitcast(
      vector_concat([bitcast(out_reg, vec_i32_ty) for out_reg in out_regs]),
      ty,
  )


@dataclasses.dataclass(frozen=True)
class ForResult:
  op: scf.ForOp
  results: tuple[Any, ...]

  @property
  def result(self):
    if len(self.results) != 1:
      raise ValueError
    return self.results[0]


def fori(bound, carrys):
  unwrap = False
  if not isinstance(carrys, (list, tuple)):
    carrys = [carrys]
    unwrap = True
  flat_carrys, carry_treedef = jax.tree.flatten(carrys)

  def wrapper(f):
    c0 = arith.constant(bound.type, 0)
    c1 = arith.constant(bound.type, 1)
    for_op = scf.ForOp(c0, bound, c1, flat_carrys)
    with ir.InsertionPoint(for_op.body):
      i = for_op.induction_variable
      inner_carrys = jax.tree.unflatten(carry_treedef, for_op.inner_iter_args)
      if unwrap:
        [inner_carrys] = inner_carrys
      new_carrys = f(i, inner_carrys)
      if unwrap:
        new_carrys = [new_carrys]
      new_flat_carrys, new_carry_treedef = jax.tree.flatten(new_carrys)
      if new_carry_treedef != carry_treedef:
        raise ValueError(new_carry_treedef, carry_treedef)
      scf.YieldOp(new_flat_carrys)
    final_flat_carrys = for_op.results
    return ForResult(
        for_op, jax.tree.unflatten(carry_treedef, final_flat_carrys)
    )

  return wrapper


@contextlib.contextmanager
def when(cond):
  with ir.InsertionPoint(scf.IfOp(cond).then_block):
    yield
    scf.yield_([])


def _3d_to_1d_idx(dim_idx_fn, dim_size_fn):
  i32 = ir.IntegerType.get_signless(32)
  as_i32 = lambda x: arith.index_cast(i32, x)
  idx = as_i32(dim_idx_fn(gpu.Dimension.x))
  stride = as_i32(dim_size_fn(gpu.Dimension.x))
  for dim in (gpu.Dimension.y, gpu.Dimension.z):
    idx = arith.addi(idx, arith.muli(as_i32(dim_idx_fn(dim)), stride))
    stride = arith.muli(stride, as_i32(dim_size_fn(dim)))
  return idx


thread_idx = functools.partial(_3d_to_1d_idx, gpu.thread_id, gpu.block_dim)
block_idx = functools.partial(_3d_to_1d_idx, gpu.block_id, gpu.grid_dim)


def _warp_bcast(val, lane_idx=0):
  i32 = ir.IntegerType.get_signless(32)
  mask = c(0xFFFFFFFF, i32)
  return nvvm.shfl_sync(
      mask, val, c(lane_idx, i32), c(0x1F, i32), nvvm.ShflKind.idx
  )


def warp_idx(sync=True):
  if get_arch().major < 9:  # Unlikely to have any benefit.
    sync = False
  i32 = ir.IntegerType.get_signless(32)
  warp_idx = arith.shrui(thread_idx(), c(5, i32))
  # Performing a warp broadcast improves performance as compiler understands
  # that the value is uniform across the warp.
  return _warp_bcast(warp_idx) if sync else warp_idx


def warpgroup_idx(sync=True):
  major = get_arch().major
  if major < 9 or major == 12:  # Unlikely to have any benefit.
    sync = False
  i32 = ir.IntegerType.get_signless(32)
  wg_idx = arith.shrui(thread_idx(), c(7, i32))
  # Performing a warp broadcast improves performance as compiler understands
  # that the value is uniform across the warp.
  return _warp_bcast(wg_idx) if sync else wg_idx


class ThreadSubset(enum.IntEnum):
  WARP = enum.auto()
  WARPGROUP = enum.auto()
  BLOCK = enum.auto()


# True within `once()` contexts.
_ONCE_PER: ThreadSubset | None = None


def single_thread_predicate(scope: ThreadSubset = ThreadSubset.BLOCK):
  """Returns a predicate that selects a single thread.

  Args:
    scope: What level of the thread hierarchy to select a thread from. For
      example, if the scope is BLOCK, only one thread per block will be
      selected.
  """
  if get_arch().major < 9:
    tidx = thread_idx()
    match scope:
      case ThreadSubset.WARP:
        tidx = arith.remui(tidx, c(WARP_SIZE, tidx.type))
      case ThreadSubset.WARPGROUP:
        tidx = arith.remui(tidx, c(WARPGROUP_SIZE, tidx.type))
      case ThreadSubset.BLOCK:
        pass
    return arith.cmpi(arith.CmpIPredicate.eq, tidx, c(0, tidx.type))
  elected = nvvm.elect_sync()
  if scope == ThreadSubset.WARP:
    return elected
  warp = warp_idx()
  if scope is not ThreadSubset.BLOCK:
    warp = arith.remui(warp, c(4, warp.type))
  first_warp = arith.cmpi(arith.CmpIPredicate.eq, warp, c(0, warp.type))
  return arith.andi(first_warp, elected)


@contextlib.contextmanager
def single_thread(scope: ThreadSubset = ThreadSubset.BLOCK):
  """Runs the context only from a single thread.

  Args:
    scope: What level of the thread hierarchy to select a thread from. For
      example, if the scope is BLOCK, only one thread per block will be
      selected.
  """
  global _ONCE_PER
  # If we're already in a single-thread context, we don't have to do anything.
  if _ONCE_PER is not None and _ONCE_PER >= scope:
    yield
    return

  prev_scope = _ONCE_PER
  _ONCE_PER = scope
  try:
    if_op = scf.IfOp(single_thread_predicate(scope))
    with ir.InsertionPoint(if_op.then_block):
      yield
      scf.YieldOp([])
  finally:
    _ONCE_PER = prev_scope


def clock():
  i32 = ir.IntegerType.get_signless(32)
  # `has_side_effects=True` prevents the compiler reordering this instruction.
  return inline_ptx(
      "mov.u32 $0, %clock;", result_types=i32, has_side_effects=True
  )


def smid():
  i32 = ir.IntegerType.get_signless(32)
  # `has_side_effects=True` prevents the compiler reordering this instruction.
  return inline_ptx(
      "mov.u32 $0, %smid;", result_types=i32, has_side_effects=True
  )


def globaltimer(kind: Literal["low", "high"] | None = None):
  # `has_side_effects=True` prevents the compiler reordering this instruction.
  if kind is None:
    i64 = ir.IntegerType.get_signless(64)
    return inline_ptx(
        "mov.u64 $0, %globaltimer;", result_types=i64, has_side_effects=True
    )
  i32 = ir.IntegerType.get_signless(32)
  return inline_ptx(
      f"mov.u32  $0,%globaltimer_{kind[:2]};",
      result_types=i32,
      has_side_effects=True,
  )


def bytewidth(ty: ir.Type):
  bw = bitwidth(ty)
  assert bw % 8 == 0, ty
  return bw // 8


def bitwidth_impl(ty: ir.Type):
  # The actual width of TF32 is 19 bits. However, we need to treat it as
  # 32 bits for compatibility reasons. TF32 used to be 32 bits wide in upstream
  # MLIR, but it changed in
  # https://github.com/llvm/llvm-project/commit/67a1fdb014790a38a205d28e1748634de34471dd.
  if isinstance(ty, ir.FloatTF32Type):
    return 32
  if isinstance(ty, ir.IntegerType):
    return ir.IntegerType(ty).width
  if isinstance(ty, ir.FloatType):
    return ir.FloatType(ty).width
  if isinstance(ty, dialect.BarrierType):
    return MBARRIER_BYTES * 8
  if isinstance(ty, dialect.B6x16P32Type):
    return 128
  if isinstance(ty, dialect.P2B6Type):
    return 8
  if isinstance(ty, ir.VectorType):
    vty = ir.VectorType(ty)
    return math.prod(vty.shape) * bitwidth(vty.element_type)
  raise NotImplementedError(ty)


def bitwidth(ty: ir.Type):
  result = bitwidth_impl(ty)
  if result.bit_count() != 1:
    raise ValueError(f"Only power of 2 bitwidths are supported, got: {result}")
  return result


@dataclasses.dataclass(frozen=True)
class DynamicSlice:
  base: ir.Value | int
  length: int

  def __post_init__(self):
    if isinstance(self.base, int) and self.base < 0:
      raise ValueError(f"base must be non-negative, got {self.base}")
    if self.length < 0:
      raise ValueError(f"length must be non-negative, got {self.length}")


ds = DynamicSlice


def memref_slice(ref: ir.Value[ir.MemRefType], index) -> ir.Value[ir.MemRefType]:
  ref_ty = ir.MemRefType(ref.type)
  base_indices, slice_shape, is_squeezed = parse_indices(index, ref_ty.shape)
  # TODO(apaszke): Check that slice is within the memref (indices might be
  # dynamic, but we can at least catch some OOB slices).

  memref_strides, offset = ref_ty.get_strides_and_offset()
  dynamic_offset = ir.ShapedType.get_dynamic_stride_or_offset()
  new_offset = offset
  if new_offset != dynamic_offset:
    for idx, stride in zip(base_indices, memref_strides):
      if isinstance(idx, int):
        new_offset += idx * stride
      else:
        new_offset = dynamic_offset
        break
  new_strides = [
      s for s, squeeze in zip(memref_strides, is_squeezed) if not squeeze
  ]
  new_shape = [s for s, squeeze in zip(slice_shape, is_squeezed) if not squeeze]
  new_layout = ir.StridedLayoutAttr.get(new_offset, new_strides)

  ref_slice = memref.subview(
      ref,
      base_indices,
      slice_shape,
      [1] * len(ref_ty.shape),
      result_type=ir.MemRefType.get(
          new_shape, ref_ty.element_type, new_layout, ref_ty.memory_space
      ),
  )
  return ref_slice


def _is_contiguous_shape_slice(
    ref_ty: ir.MemRefType, dim_slice: slice = slice(None)
):
  # If it's not a strided layout then we are definitely contiguous.
  if not isinstance(ref_ty.layout, ir.StridedLayoutAttr):
    return True

  strides = ir.StridedLayoutAttr(ref_ty.layout).strides[dim_slice]
  shape = ref_ty.shape[dim_slice]

  # Check that each dimension fits exactly it the immediately larger stride.
  for (prev_stride, _), (stride, shape) in itertools.pairwise(zip(strides, shape)):
    if stride * shape != prev_stride:
      return False

  return True


def _reshape(ref: ir.Value, sh0: list[int], sh1: list[int]):
  """Reshapes using only "parallel" folds/unfolds.

  This function uses folds/unfolds that are "parallel" in that they
  only act on original dimensions, i.e. they won't fold into an
  intermediate dimension that they will then unfold.
  """

  i0, i1 = 0, 0

  def fold_until(shape, off, target) -> tuple[int, int]:
    assert shape[off] < target
    dim = 1
    for to in range(off, len(shape)):
      dim *= shape[to]
      if dim == target:
        return to + 1, dim
      if dim > target:
        # TODO(cperivol): Implement dependent fold-unfolds for subsections
        # of the shape eg (..., 4,5,5, ...) -> (..., 10,10, ...) could be
        # supported without touching any other dimensions.
        raise NotImplementedError(
            f"Can't reshape {sh0} to {sh1} by composing independent"
            " folds/unfolds."
        )

    raise AssertionError(
        f"Unreachable: number of elements don't match in each shape ({sh0} ans"
        f" {sh1})"
    )

  while i0 < len(sh0) and i1 < len(sh1):
    if sh0[i0] > sh1[i1]:
      # How many dimensions following i1 should we unfold i0 into.
      idx, _ = fold_until(sh1, i1, sh0[i0])
      ref = memref_unfold(ref, i0, sh1[i1:idx])
      sh0[i0 : i0 + 1] = sh1[i1:idx]
      i0 += idx - i1
      i1 = idx
    elif sh0[i0] < sh1[i1]:
      # How many dimensions after i0 should we fold to make dim at i1.
      idx, dim = fold_until(sh0, i0, sh1[i1])
      sh0[i0:idx] = [dim]
      ref = memref_fold(ref, i0, idx - i0)
      i0 += 1
      i1 += 1
    else:
      i0 += 1
      i1 += 1

  # Fold the trailing ones
  if i0 < len(sh0):
    assert i1 == len(sh1)
    ref = memref_fold(ref, i0 - 1, len(sh0) - i0 + 1)

  if i1 < len(sh1):
    assert i0 == len(sh0)
    ref = memref_unfold(ref, i0 - 1, [sh0[i0 - 1]] + [1] * (len(sh1) - i1))

  return ref


@overload
def memref_reshape(ref: ir.Value, shape: tuple[int, ...]) -> ir.Value:
  ...


@overload
def memref_reshape(ref: MultimemRef, shape: tuple[int, ...]) -> MultimemRef:
  ...


def memref_reshape(
    ref: ir.Value | MultimemRef, shape: tuple[int, ...]
) -> ir.Value | MultimemRef:
  """Reshape by means of folding and unfolding.

  The use of memref fold/unfold may avoid some possible issues with
  strided memrefs.
  """

  if isinstance(ref, MultimemRef):
    return MultimemRef(memref_reshape(ref.ref, shape))

  ref_ty = ir.MemRefType(ref.type)
  if math.prod(ref_ty.shape) != math.prod(shape):
    raise ValueError(
        f"Cannot reshape to a different size. Ref shape: {ref_ty.shape} (size:"
        f" {math.prod(ref_ty.shape)}), new shape: {shape} (size:"
        f" {math.prod(shape)})"
    )
  if not all(dim > 0 for dim in shape):
    raise ValueError(
        "Shapes must havbe only positive dimensions (no -1 or 0 dimensions"
        f" allowed) {shape}"
    )

  src_shape = list(ref_ty.shape)
  dst_shape = list(shape)
  if src_shape == dst_shape:
    return ref
  if not src_shape:
    _, offset = ref_ty.get_strides_and_offset()
    identity = ir.AffineMapAttr.get(ir.AffineMap.get_identity(0))
    if ref_ty.layout == identity:
      new_layout = ir.AffineMapAttr.get(
          ir.AffineMap.get_identity(len(dst_shape))
      )
    else:
      new_layout = ir.StridedLayoutAttr.get(offset, [1] * len(dst_shape))
    result_ty = ir.MemRefType.get(
        dst_shape, ref_ty.element_type, new_layout, ref_ty.memory_space
    )
    return memref.expand_shape(result_ty, ref, [], [], dst_shape)
  if not dst_shape:
    _, offset = ref_ty.get_strides_and_offset()
    identity = ir.AffineMapAttr.get(ir.AffineMap.get_identity(ref_ty.rank))
    contig_strided_1d = ir.Attribute.parse("strided<[1]>")
    if ref_ty.layout == identity or ref_ty.layout == contig_strided_1d:
      new_layout = ir.AffineMapAttr.get(ir.AffineMap.get_identity(0))
    else:
      new_layout = ir.StridedLayoutAttr.get(offset, [])
    result_ty = ir.MemRefType.get(
        (), ref_ty.element_type, new_layout, ref_ty.memory_space
    )
    return memref.collapse_shape(result_ty, ref, [])
  try:
    return _reshape(ref, src_shape, dst_shape)
  except NotImplementedError:
    # For contiguous refs we can do arbitrary reshapes easily.
    if _is_contiguous_shape_slice(ref_ty):
      return memref_unfold(memref_fold(ref, 0, ref_ty.rank), 0, shape)
    raise


@overload
def memref_fold(ref: ir.Value, dim, fold_rank) -> ir.Value:
  ...


@overload
def memref_fold(ref: MultimemRef, dim, fold_rank) -> MultimemRef:
  ...


def memref_fold(
    ref: ir.Value | MultimemRef, dim, fold_rank
) -> ir.Value | MultimemRef:
  if isinstance(ref, MultimemRef):
    return MultimemRef(memref_fold(ref.ref, dim, fold_rank))

  ref_ty = ir.MemRefType(ref.type)
  new_shape = list(ref_ty.shape)
  if dim < 0:
    raise ValueError(f"Dimension {dim} is negative")
  if dim + fold_rank > len(new_shape):
    raise ValueError(
        f"Folding {fold_rank} dimensions starting from {dim} is out of bounds"
        f" for shape {new_shape}"
    )
  new_shape[dim : dim + fold_rank] = [
      math.prod(new_shape[dim : dim + fold_rank])
  ]
  identity = ir.AffineMapAttr.get(ir.AffineMap.get_identity(ref_ty.rank))
  contig_strided_1d = ir.Attribute.parse("strided<[1]>")
  # Not sure why but MLIR expects the strided 1D layout to disappear in this op.
  if ref_ty.layout == identity or ref_ty.layout == contig_strided_1d:
    new_layout = ir.AffineMapAttr.get(
        ir.AffineMap.get_identity(ref_ty.rank - fold_rank + 1)
    )
  elif _is_contiguous_shape_slice(ref_ty, slice(dim, dim + fold_rank)):
    new_strides, offset = ref_ty.get_strides_and_offset()
    new_strides[dim : dim + fold_rank] = [new_strides[dim + fold_rank - 1]]
    new_layout = ir.StridedLayoutAttr.get(offset, new_strides)
  else:
    raise ValueError(
        f"strides={ref_ty.get_strides_and_offset()[0]}, {ref_ty.shape=},"
        f" {dim=}, {fold_rank=}"
    )

  new_ty = ir.MemRefType.get(
      new_shape, ref_ty.element_type, new_layout, ref_ty.memory_space
  )
  assoc = [[d] for d in range(dim)]
  assoc.append([dim + i for i in range(fold_rank)])
  assoc.extend([d] for d in range(dim + fold_rank, ref_ty.rank))
  assert len(assoc) == new_ty.rank
  return memref.collapse_shape(new_ty, ref, assoc)


def memref_unfold(ref: ir.Value, dim, factors) -> ir.Value:
  """Unfolds dim into two dimensions, the size of leading one given be major_factor."""
  ref_ty = ir.MemRefType(ref.type)
  new_shape = list(ref_ty.shape)
  if sum(f is None for f in factors) > 1:
    raise ValueError("Can only infer one dimension")
  known_factor_prod = np.prod([f for f in factors if f is not None])
  if new_shape[dim] % known_factor_prod:
    raise ValueError("Non-divisible unfold:", new_shape[dim], factors)
  factors = tuple(
      new_shape[dim] // known_factor_prod if f is None else f for f in factors
  )
  new_shape[dim : dim + 1] = factors
  identity = ir.AffineMapAttr.get(ir.AffineMap.get_identity(ref_ty.rank))
  contig_strided_1d = ir.Attribute.parse("strided<[1]>")
  if ref_ty.layout == identity or ref_ty.layout == contig_strided_1d:
    new_layout = ir.AffineMapAttr.get(
        ir.AffineMap.get_identity(ref_ty.rank + len(factors) - 1)
    )
  else:
    new_strides, offset = ref_ty.get_strides_and_offset()
    prev_stride = new_strides[dim]
    inserted_strides = []
    for f in reversed(factors):
      inserted_strides.append(prev_stride)
      prev_stride *= f
    new_strides[dim : dim + 1] = reversed(inserted_strides)
    new_layout = ir.StridedLayoutAttr.get(offset, new_strides)
  new_ty = ir.MemRefType.get(
      new_shape, ref_ty.element_type, new_layout, ref_ty.memory_space
  )
  if dim == ref_ty.rank:
    assoc = [[d] for d in range(ref_ty.rank)]
    assoc[-1].extend(range(ref_ty.rank, ref_ty.rank + len(factors) - 1))
  else:
    assoc = [[d] for d in range(dim)]
    assoc.append(list(range(dim, dim + len(factors))))
    assoc.extend([d + len(factors) - 1] for d in range(dim + 1, ref_ty.rank))
  assert len(assoc) == ref_ty.rank
  return memref.expand_shape(new_ty, ref, assoc, [], new_ty.shape)


def memref_unsqueeze(ref: ir.Value, dim) -> ir.Value:
  """Inserts a singleton dimension."""
  ref_ty = ir.MemRefType(ref.type)
  if dim == ref_ty.rank:
    new_shape = list(ref_ty.shape)
    new_shape.append(1)
    identity = ir.AffineMapAttr.get(ir.AffineMap.get_identity(ref_ty.rank))
    if ref_ty.layout == identity:
      new_layout = ir.AffineMapAttr.get(
          ir.AffineMap.get_identity(ref_ty.rank + 1)
      )
    else:
      new_strides, offset = ref_ty.get_strides_and_offset()
      new_strides.append(1)
      new_layout = ir.StridedLayoutAttr.get(offset, new_strides)
    new_ty = ir.MemRefType.get(
        new_shape, ref_ty.element_type, new_layout, ref_ty.memory_space
    )
    assoc = [[d] for d in range(ref_ty.rank)]
    assoc[-1].append(ref_ty.rank)
    return memref.expand_shape(new_ty, ref, assoc, [], new_ty.shape)
  else:
    return memref_unfold(ref, dim, (1, None))


def is_memref_transposed(ref: ir.MemRefType) -> bool:
  strides, _ = ref.get_strides_and_offset()
  prev_stride = math.inf
  for stride in strides:
    if stride > prev_stride:
      return True
    prev_stride = stride
  return False


def memref_transpose(ref: ir.Value, permutation: Sequence[int]) -> ir.Value:
  ref_ty = ir.MemRefType(ref.type)
  strides, offset = ref_ty.get_strides_and_offset()
  new_strides = [strides[p] for p in permutation]
  new_shape = [ref_ty.shape[p] for p in permutation]
  new_layout = ir.StridedLayoutAttr.get(offset, new_strides)
  new_ty = ir.MemRefType.get(
      new_shape, ref_ty.element_type, new_layout, ref_ty.memory_space
  )
  return memref.transpose(
      new_ty, ref, ir.AffineMap.get_permutation(permutation)
  )


def parse_indices(
    index,
    shape: Sequence[int],
    *,
    check_oob: bool = True,
    allow_negative: bool = False,
) -> tuple[list[ir.Value | int], list[int], list[bool]]:
  if not isinstance(index, tuple):
    index = (index,)
  if trailing_dims := len(shape) - len(index):
    index += (slice(None),) * trailing_dims
  base_indices: list[ir.Value | int] = []
  slice_shape = []
  is_squeezed = []
  for axis, (idx, bound) in enumerate(zip(index, shape)):
    if isinstance(idx, (ir.Operation, ir.OpView)):
      idx = idx.result
    if isinstance(idx, int):
      if not allow_negative and idx < 0:
        raise NotImplementedError(
            f"Index {idx} along axis {axis} has negative bounds"
        )
      if check_oob and idx >= bound:
        raise IndexError(
            f"Index {idx} along axis {axis} is out of bounds for shape {shape}"
        )
      base_indices.append(idx)
      slice_shape.append(1)
      is_squeezed.append(True)
    elif isinstance(idx, slice):
      if any(isinstance(v, ir.Value) for v in (idx.start, idx.stop, idx.step)):
        raise NotImplementedError(
            f"Slice {idx} along axis {axis} has dynamic bounds. Use"
            " ds(base, length) to slice at a dynamic offset with a"
            " static length."
        )
      if idx.step is not None and idx.step != 1:
        raise NotImplementedError("Strided slices not implemented")
      if not allow_negative and any(
          v is not None and v < 0 for v in (idx.start, idx.stop)
      ):
        raise NotImplementedError(
            f"Slice {idx} along axis {axis} has negative bounds"
        )
      start = idx.start or 0
      stop = idx.stop or bound
      if check_oob and (start >= bound or stop > bound):
        raise IndexError(
            f"Slice {idx} along axis {axis} is out of bounds for shape {shape}"
        )
      base_indices.append(start)
      slice_shape.append(stop - start)
      is_squeezed.append(False)
    elif isinstance(idx, DynamicSlice):
      if check_oob and (
          isinstance(idx.base, int) and idx.base + idx.length > bound
      ):
        raise IndexError(
            f"Slice {idx} along axis {axis} is out of bounds for shape {shape}"
        )
      base_indices.append(idx.base)
      slice_shape.append(idx.length)
      is_squeezed.append(False)
    elif isinstance(idx, ir.Value):
      if not isinstance(idx.type, ir.IndexType):
        raise ValueError("Expected an index-typed index")
      base_indices.append(idx)
      slice_shape.append(1)
      is_squeezed.append(True)
    else:
      raise NotImplementedError(type(idx))
  assert len(base_indices) == len(slice_shape) == len(is_squeezed) == len(shape)
  return base_indices, slice_shape, is_squeezed


def commit_shared():
  nvvm.fence_proxy(
      nvvm.ProxyKind.async_shared, space=nvvm.SharedSpace.shared_cta
  )
  warpgroup_barrier()


def warpgroup_barrier_idx(sync: bool = True) -> ir.Value[ir.IntegerType]:
  # gpu.barrier() uses barrier number 0, and it would be unsafe to reuse it,
  # so we shift the warpgroup index by 1.
  i32 = ir.IntegerType.get_signless(32)
  return arith.addi(warpgroup_idx(sync=sync), c(1, i32))


def warpgroup_barrier():
  wg_idx = warpgroup_barrier_idx(sync=False)
  inline_ptx(f"bar.sync $0, {WARPGROUP_SIZE};", wg_idx, has_side_effects=True)


def warp_barrier():
  nvvm.bar_warp_sync(c(0xFFFFFFFF, ir.IntegerType.get_signless(32)))


def before_thread_sync(
    *,
    sync_threads: bool,
    scope: ThreadSubset = ThreadSubset.WARPGROUP,
):
  nvvm.tcgen05_fence(nvvm.Tcgen05FenceKind.BEFORE_THREAD_SYNC)
  if sync_threads:
    if scope == ThreadSubset.WARPGROUP:
      warpgroup_barrier()
    elif scope == ThreadSubset.WARP:
      warp_barrier()
    else:
      raise ValueError(f"Unsupported scope: {scope}")


def prefetch_tensormap(
    desc_ptr: ir.Value[llvm.PointerType],
    predicate: ir.Value[ir.IntegerType] | None = None,
) -> None:
  """Prefetches a TMA descriptor into the cache using PTX `prefetch.tensormap`.

  Args:
    desc_ptr: A pointer to the 128-byte aligned TMA descriptor.
    predicate: An optional i1 predicate value.
  """
  inline_ptx(
      "prefetch.tensormap [$0];",
      desc_ptr,
      predicate=predicate,
      has_side_effects=True,
  )


@dataclasses.dataclass(frozen=True)
class BarrierRef:
  base_address: ir.Value
  offset: ir.Value
  phases: ir.Value
  num_barriers: int

  @staticmethod
  def initialize(
      barrier_memref: ir.Value, arrival_count: int = 1
  ) -> "BarrierRef":
    barrier_ty = ir.MemRefType(barrier_memref.type)
    [num_barriers] = barrier_ty.shape
    if num_barriers > 32:
      raise NotImplementedError("Only up to 32 barriers per group supported")
    i32 = ir.IntegerType.get_signless(32)
    i64 = ir.IntegerType.get_signless(64)
    address = memref_ptr(barrier_memref)
    phases = memref.alloca(ir.MemRefType.get((), i32), [], [])
    memref.store(c(0, i32), phases, [])
    predicate = single_thread_predicate(scope=ThreadSubset.BLOCK)
    for i in range(num_barriers):
      nvvm.mbarrier_init(
          getelementptr(address, [i], i64),
          c(arrival_count, i32),
          predicate=predicate,
      )
    return BarrierRef(address, c(0, i32), phases, num_barriers)

  def __iter__(self) -> Iterator["BarrierRef"]:
    if self.num_barriers == 1:
      yield self
    else:
      for offset in range(self.num_barriers):
        yield self[offset]

  def __getitem__(self, offset: ir.Value | int) -> "BarrierRef":
    i32 = ir.IntegerType.get_signless(32)
    if isinstance(offset, int):
      if offset >= self.num_barriers:
        raise IndexError(f"Barrier offset {offset} is out of bounds")
      offset = c(offset, i32)
    elif isinstance(offset.type, ir.IndexType):
      offset = arith.index_castui(i32, offset)
    elif offset.type != i32:
      raise ValueError(f"Expected a dynamic index or an integer, got {offset}")
    return BarrierRef(
        self.base_address,
        arith.addi(self.offset, offset),
        self.phases,
        1,
    )

  @property
  def _ptx_scope(self) -> str:
    if self.base_address.type == ir.Type.parse("!llvm.ptr<7>"):
      return "cluster"
    return "cta"

  def test_parity(
      self,
      parity,
      orders_tensor_core: bool = False,
      scope: ThreadSubset = ThreadSubset.WARPGROUP,
  ) -> ir.Value:
    i1 = ir.IntegerType.get_signless(1)
    i32 = ir.IntegerType.get_signless(32)
    parity = arith.extui(i32, parity)
    wait_complete = nvvm.mbarrier_test_wait(self.get_ptr(), parity)

    if scope == ThreadSubset.WARPGROUP:
      wait_complete = inline_ptx(
          f"bar.red.or.pred $0, $1, {WARPGROUP_SIZE}, $2;",
          warpgroup_barrier_idx(sync=False),
          wait_complete,
          result_types=i1,
          has_side_effects=True,
      )
      wait_complete = cast(ir.OpResult[ir.IntegerType], wait_complete)
    elif scope == ThreadSubset.WARP:
      wait_complete = nvvm.vote_sync(c(0xFFFFFFFF, i32), wait_complete, "any")
    else:
      raise ValueError(f"Unsupported scope: {scope}")

    if orders_tensor_core:
      nvvm.tcgen05_fence(nvvm.Tcgen05FenceKind.AFTER_THREAD_SYNC)
    return wait_complete

  def test(
      self,
      orders_tensor_core: bool = False,
      scope: ThreadSubset = ThreadSubset.WARPGROUP,
  ) -> ir.Value:
    parities = memref.load(self.phases, [])
    parity, new_parities = self.update_parities(parities)
    wait_complete = self.test_parity(parity, orders_tensor_core, scope)
    with when(wait_complete):
      memref.store(new_parities, self.phases, [])
    return wait_complete

  def wait_parity(self, parity, orders_tensor_core: bool = False):
    if self._ptx_scope != "cta":
      raise ValueError("Can only await on CTA-local barriers")
    i32 = ir.IntegerType.get_signless(32)
    parity = arith.extui(i32, parity)
    if get_arch().major < 9:
      # TODO(apaszke): consider using a single lane + barrier for waiting
      i1 = ir.IntegerType.get_signless(1)
      while_op = scf.WhileOp([], [])
      before_block = while_op.before.blocks.append()
      with ir.InsertionPoint.at_block_begin(before_block):
        wait_complete = nvvm.mbarrier_test_wait(self.get_ptr(), parity)
        wait_incomplete = arith.xori(wait_complete, c(1, i1))
        scf.condition(wait_incomplete, [])
      after_block = while_op.after.blocks.append()
      with ir.InsertionPoint.at_block_begin(after_block):
        scf.yield_([])
    else:
      ticks = arith.constant(i32, 10000000)
      nvvm.mbarrier_try_wait_parity(self.get_ptr(), parity, ticks)
    if orders_tensor_core:
      nvvm.tcgen05_fence(nvvm.Tcgen05FenceKind.AFTER_THREAD_SYNC)

  def wait(self, orders_tensor_core: bool = False):
    parities = memref.load(self.phases, [])
    parity, new_parities = self.update_parities(parities)
    memref.store(new_parities, self.phases, [])
    self.wait_parity(parity, orders_tensor_core)

  def update_parities(self, parities: ir.Value) -> tuple[ir.Value, ir.Value]:
    i32 = ir.IntegerType.get_signless(32)
    bitmask = arith.shli(c(1, i32), self.offset)
    parity = arith.cmpi(
        arith.CmpIPredicate.ne, arith.andi(parities, bitmask), c(0, i32)
    )
    return parity, arith.xori(parities, bitmask)

  def arrive(
      self,
      arrival_count: int = 1,
      can_complete: bool = True,
      orders_tensor_core: bool = False,
      predicate: ir.Value | None = None,
      scope: ThreadSubset = ThreadSubset.WARPGROUP,
  ):
    if orders_tensor_core:
      # We need to synchronize the threads after `::before_thread_sync` if a
      # predicate is used, as not all threads arrive on the barrier.
      before_thread_sync(sync_threads=predicate is not None, scope=scope)

    ptx_scope = self._ptx_scope
    if can_complete or ptx_scope != "cta":
      count_ptx = f", {arrival_count}"
      if get_arch().major < 9:
        if arrival_count != 1:
          raise ValueError(
              "Only single-thread arrival is supported on pre-Hopper hardware"
          )
        count_ptx = ""
      inline_ptx(
          f"mbarrier.arrive.release.{ptx_scope}.shared::{ptx_scope}.b64 _, [$0]{count_ptx};",
          self.get_ptr(),
          predicate=predicate,
          has_side_effects=True,
      )
    else:
      if predicate is not None:
        raise NotImplementedError(
            "Predicate not supported for no-complete arrive"
        )
      count = c(arrival_count, ir.IntegerType.get_signless(32))
      nvvm.mbarrier_arrive_nocomplete(self.get_ptr(), count)

  def arrive_expect_tx(
      self,
      tx_count: int | ir.Value,
      predicate: ir.Value | None = None,
      tensor_core_order_scope: ThreadSubset | None = None,
  ):
    if get_arch().major < 9:
      raise NotImplementedError("arrive_expect_tx is only supported on Hopper+ hardware")

    if tensor_core_order_scope is not None:
      before_thread_sync(
          sync_threads=predicate is not None, scope=tensor_core_order_scope
      )

    i32 = ir.IntegerType.get_signless(32)
    if isinstance(tx_count, int):
      tx_count = c(tx_count, i32)
    elif isinstance(tx_count.type, ir.IndexType):
      tx_count = arith.index_cast(i32, tx_count)
    ptx_scope = self._ptx_scope
    inline_ptx(
        f"mbarrier.arrive.expect_tx.release.{ptx_scope}.shared::{ptx_scope}.b64 _, [$0], $1;",
        self.get_ptr(),
        tx_count,
        predicate=predicate,
        has_side_effects=True,
        convergent=True
    )

  def complete_tx(
      self, tx_count: int | ir.Value, predicate: ir.Value | None = None
  ):
    if get_arch().major < 9:
      raise NotImplementedError("complete_tx is only supported on Hopper+ hardware")

    i32 = ir.IntegerType.get_signless(32)
    if isinstance(tx_count, int):
      tx_count = c(tx_count, i32)
    elif isinstance(tx_count.type, ir.IndexType):
      tx_count = arith.index_cast(i32, tx_count)

    inline_ptx(
        f"mbarrier.complete_tx.shared::{self._ptx_scope}.b64 [$0], $1;",
        self.get_ptr(),
        tx_count,
        predicate=predicate,
        has_side_effects=True,
    )

  def get_ptr(self):
    if self.num_barriers != 1:
      raise ValueError(
          f"Operation is only valid for a single barrier, but barrier ref "
          f"contains {self.num_barriers}. Individual barriers can be accessed "
          f"by indexing into the ref."
      )
    i64 = ir.IntegerType.get_signless(64)
    return getelementptr(self.base_address, [self.offset], i64)

  def remap_to_cluster(self, dim: gpu.Dimension, idx: ir.Value):
    if self._ptx_scope != "cta":
      # I think this should be fine, but we are not testing it.
      raise ValueError("Remapping a non-local ref unsupported.")
    i32 = ir.IntegerType.get_signless(32)
    idxs: list[ir.Value] = [gpu.cluster_block_id(d) for d in gpu.Dimension]
    idxs[dim] = idx
    flat_block = arith.index_cast(i32, cluster_idx(dim_idx=idxs))
    cptr = get_cluster_ptr(self.base_address, flat_block, generic=False)
    return BarrierRef(cptr, self.offset, self.phases, self.num_barriers)


@dataclasses.dataclass(frozen=True)
class DialectBarrierRef:
  barrier_ref: BarrierRef
  orders_tensor_core: bool

  @staticmethod
  def initialize(
      barrier_memref: ir.Value,
      arrival_count: int,
      orders_tensor_core: bool,
  ) -> "DialectBarrierRef":
    barrier_ty = ir.MemRefType(barrier_memref.type)
    [num_barriers] = barrier_ty.shape
    if num_barriers > 32:
      raise NotImplementedError("Only up to 32 barriers per group supported")

    address = memref_ptr(barrier_memref)
    dialect.initialize_barrier(
        address, arrival_count, num_barriers, orders_tensor_core
    )
    i32 = ir.IntegerType.get_signless(32)
    phases = memref.alloca(ir.MemRefType.get((), i32), [], [])
    memref.store(c(0, i32), phases, [])
    return DialectBarrierRef(
        barrier_ref=BarrierRef(address, c(0, i32), phases, num_barriers),
        orders_tensor_core=orders_tensor_core,
    )

  def __iter__(self) -> Iterator["DialectBarrierRef"]:
    if self.barrier_ref.num_barriers == 1:
      yield self
    else:
      for offset in range(self.barrier_ref.num_barriers):
        yield self[offset]

  def __getitem__(self, offset: ir.Value | int) -> "DialectBarrierRef":
    return DialectBarrierRef(self.barrier_ref[offset], self.orders_tensor_core)

  def test_parity(
      self,
      parity,
      orders_tensor_core: bool = False,
      scope: ThreadSubset = ThreadSubset.WARPGROUP,
  ) -> ir.Value:
    assert self.orders_tensor_core == orders_tensor_core
    return self.barrier_ref.test_parity(parity, orders_tensor_core, scope=scope)

  def test(
      self,
      orders_tensor_core: bool = False,
      scope: ThreadSubset = ThreadSubset.WARPGROUP,
  ) -> ir.Value:
    assert self.orders_tensor_core == orders_tensor_core
    assert self.barrier_ref.phases is not None
    return self.barrier_ref.test(orders_tensor_core, scope=scope)

  def wait_parity(self, parity, orders_tensor_core: bool = False):
    assert self.orders_tensor_core == orders_tensor_core
    self.barrier_ref.wait_parity(parity, orders_tensor_core)

  def wait(self, orders_tensor_core: bool = False):
    assert self.orders_tensor_core == orders_tensor_core
    assert self.barrier_ref.phases is not None
    self.barrier_ref.wait(orders_tensor_core)

  def update_parities(self, parities: ir.Value) -> tuple[ir.Value, ir.Value]:
    return self.barrier_ref.update_parities(parities)

  def arrive(
      self,
      *,
      orders_tensor_core: bool = False,
      predicate: ir.Value | None = None,
  ):
    assert self.orders_tensor_core == orders_tensor_core
    dialect.ArriveOp(
        self.as_barrier_memref(),
        orders_tensor_core=orders_tensor_core,
        predicate=predicate,
    )

  def arrive_expect_tx(self, tx_count: int | ir.Value):
    if isinstance(tx_count, int):
      tx_count = c(tx_count, ir.IntegerType.get_signless(32))
    # pyrefly: ignore[bad-argument-type]
    dialect.ArriveExpectTxOp(barrier=self.as_barrier_memref(), expect_tx=tx_count)

  def get_ptr(self):
    return self.barrier_ref.get_ptr()

  def as_barrier_memref(self) -> ir.Value:
    num_barriers = self.barrier_ref.num_barriers
    shape = () if num_barriers == 1 else (num_barriers,)
    barrier_type = dialect.BarrierType.get(self.orders_tensor_core)
    ptr_type = llvm.PointerType(self.get_ptr().type)
    assert ptr_type.address_space == WORKGROUP_NVPTX_ADDRESS_SPACE
    memref_type = ir.MemRefType.get(shape, barrier_type, memory_space=smem())
    result = builtin.unrealized_conversion_cast([memref_type], [self.get_ptr()])
    assert isinstance(result, ir.Value)
    return result

  @classmethod
  def from_barrier_memref(cls, barrier: ir.Value):
    """Creates a DialectBarrierRef from a memref of a dialect barrier."""
    memref_type = ir.MemRefType(barrier.type)
    if memref_type.rank > 1 or not isinstance(
        memref_type.element_type, dialect.BarrierType
    ):
      raise ValueError(
          "Expected a memref with rank 0 or 1 and element type "
          f"!mosaic_gpu.barrier, but got {barrier.type}"
      )

    ptr_type = llvm.PointerType.get(get_memref_llvm_address_space(memref_type))
    addr = builtin.unrealized_conversion_cast([ptr_type], [barrier])
    assert isinstance(addr, ir.Value)
    return cls(
        barrier_ref=BarrierRef(
            base_address=addr,
            offset=c(0, ir.IntegerType.get_signless(64)),
            # TODO(slebedev): Why is it safe to use None here?
            phases=None,  # pyrefly: ignore[bad-argument-type]
            num_barriers=(1 if memref_type.rank == 0 else memref_type.shape[0]),
        ),
        orders_tensor_core=memref_type.element_type.orders_tensor_core,
    )


@dataclasses.dataclass(frozen=True)
class CollectiveBarrierRef:
  barrier: BarrierRef
  cluster_mask: ir.Value | None
  leader_tracked: bool = False

  @staticmethod
  def initialize(
      barrier_memref: ir.Value,
      arrival_count: int,
      dims: Sequence[gpu.Dimension | Sequence[gpu.Dimension]],
      cluster_shape: tuple[int, int, int],
      leader_tracked: bool = False,
  ) -> "CollectiveBarrierRef":
    if leader_tracked:
      # TODO(bchetioui): expand for multiple collective dimensions. Not sure
      # what useful semantics should be.
      if len(dims) != 1:
        raise NotImplementedError(
            "Only leader tracking for a single collective dimension is supported."
        )
      # TODO(bchetioui): adapt or delete case of dims containing sequences
      # of dimensions.
      if not isinstance(dims[0], gpu.Dimension):
        raise NotImplementedError(
            "Only leader tracking for ungrouped collective dimensions is supported."
        )
    i32 = ir.IntegerType.get_signless(32)
    # With the exception of the current device, each pair of slices along
    # collective dims is disjoint. Since the current device is overcounted,
    # we must decrease the arrival count a little.
    dims_shape = [
        cluster_shape[d]
        if isinstance(d, gpu.Dimension)
        else math.prod(cluster_shape[dd] for dd in d)
        for d in dims
    ]
    cluster_arrival_count = sum(dims_shape) - len(dims) + 1
    if cluster_arrival_count == 1:
      assert all(s == 1 for s in dims_shape)
      cluster_mask = None
    else:
      cluster_mask = c(0, i32)
      for d, size in zip(dims, dims_shape):
        if size == 1:
          # Only the current device is in this mask, but it will also be
          # present in one of the non-trivial cluster dims.
          continue
        cluster_mask = arith.ori(
            cluster_mask, cluster_collective_mask(cluster_shape, d, leader_tracked)
        )
    barrier = BarrierRef.initialize(
        barrier_memref, arrival_count=arrival_count * cluster_arrival_count
    )
    return CollectiveBarrierRef(barrier, cluster_mask, leader_tracked)

  def __iter__(self):
    for b in self.barrier:
      yield CollectiveBarrierRef(b, self.cluster_mask, self.leader_tracked)

  def __getitem__(self, offset):
    return CollectiveBarrierRef(
        self.barrier[offset], self.cluster_mask, self.leader_tracked
    )

  def arrive(
      self,
      *,
      orders_tensor_core: bool = False,
      predicate: ir.Value | None = None,
  ):
    """Arrives on a barrier in one or several blocks in a cluster.

    Specifically,
    - If ``self.leader_tracked`` is set, arrive only the barrier mapped in the
      leader block (block 0 along the collective dimension, with otherwise the
      same coordinates as the current block);
    - Otherwise, arrive in all blocks whose coordinates vary in at most one
      collective dimension.

    Note that unlike in arrive, each warpgroup arrives once.
    """
    if self.barrier.num_barriers != 1:
      raise ValueError("Can only arrive on a single barrier")

    if self.cluster_mask is None:
      pred = single_thread_predicate(ThreadSubset.WARPGROUP)
      if predicate is not None:
        pred = arith.andi(predicate, pred)
      return self.barrier.arrive(
          predicate=pred, orders_tensor_core=orders_tensor_core
      )

    if orders_tensor_core:
      # We need to synchronize the threads after `::before_thread_sync`, as
      # not all threads arrive on the barrier.
      before_thread_sync(sync_threads=True)

    i32 = ir.IntegerType.get_signless(32)
    thread_in_warpgroup = arith.remui(thread_idx(), c(WARPGROUP_SIZE, i32))
    signaled_block = arith.divui(
        thread_in_warpgroup, c(WARPGROUP_SIZE // 16, i32)
    )
    is_collective_block = arith.cmpi(
        arith.CmpIPredicate.ne,
        arith.andi(self.cluster_mask, arith.shli(c(1, i32), signaled_block)),
        c(0, i32),
    )
    is_signaling_thread = arith.cmpi(
        arith.CmpIPredicate.eq,
        arith.remui(thread_in_warpgroup, c(WARPGROUP_SIZE // 16, i32)),
        c(0, i32),
    )
    should_arrive = arith.andi(is_collective_block, is_signaling_thread)
    if predicate is not None:
      should_arrive = arith.andi(predicate, should_arrive)
    inline_ptx(
        """
    {
        .reg .b32 mapped_addr;
        @$0 mapa.shared::cluster.u32 mapped_addr, $1, $2;
        @$0 mbarrier.arrive.shared::cluster.b64 _, [mapped_addr];
    }""",
        should_arrive,
        self.barrier.get_ptr(),
        signaled_block,
        has_side_effects=True,
    )

  def wait(self, *args, **kwargs):
    self.barrier.wait(*args, **kwargs)

  def wait_parity(self, *args, **kwargs):
    self.barrier.wait_parity(*args, **kwargs)


@dataclasses.dataclass(frozen=True)
class CtaBarrierRef:
  base_id: int
  offset: ir.Value
  arrival_count: int
  num_barriers: int = 1

  @staticmethod
  def initialize(
      base_id: int,
      arrival_count: int = 1,
      num_barriers: int = 1,
  ) -> "CtaBarrierRef":
    i32 = ir.IntegerType.get_signless(32)
    return CtaBarrierRef(base_id, c(0, i32), arrival_count, num_barriers)

  def __getitem__(self, offset: ir.Value | int) -> "CtaBarrierRef":
    i32 = ir.IntegerType.get_signless(32)
    if isinstance(offset, int):
      if offset < 0 or offset >= self.num_barriers:
        raise IndexError(
            f"Barrier offset {offset} is out of bounds for"
            f" num_barriers={self.num_barriers}"
        )
      offset = c(offset, i32)
    elif isinstance(offset.type, ir.IndexType):
      offset = arith.index_castui(i32, offset)
    elif offset.type != i32:
      raise ValueError(f"Expected a dynamic index or an integer, got {offset}")
    return CtaBarrierRef(
        self.base_id, arith.addi(self.offset, offset), self.arrival_count, 1
    )

  def get_barrier_id(self) -> ir.Value:
    if self.num_barriers != 1:
      raise ValueError(
          "Operation is only valid for a single barrier, but barrier ref"
          f" contains {self.num_barriers}. Individual barriers can be accessed"
          " by indexing into the ref."
      )
    i32 = ir.IntegerType.get_signless(32)
    return arith.addi(c(self.base_id, i32), self.offset)

  def arrive(self, *, predicate: ir.Value | None = None) -> None:
    inline_ptx(
        f"bar.arrive $0, {self.arrival_count};",
        self.get_barrier_id(),
        predicate=predicate,
        has_side_effects=True,
    )

  def arrive_and_wait(self) -> None:
    inline_ptx(
        f"bar.sync $0, {self.arrival_count};",
        self.get_barrier_id(),
        has_side_effects=True,
    )


@dataclasses.dataclass(frozen=True)
class SemaphoreRef:
  ptr: ir.Value

  def signal(
      self,
      value: ir.Value | int,
      predicate: ir.Value | None = None,
      relaxed: bool = False,
      memory_scope: Literal["sys", "gpu"] = "sys",
  ):
    i32 = ir.IntegerType.get_signless(32)
    if not isinstance(value, ir.Value):
      value = c(value, i32)
    elif value.type != i32:
      raise ValueError(f"Expected a i32 value, got {value.type}")
    if predicate is None:
      predicate = single_thread_predicate(ThreadSubset.WARPGROUP)
    if memory_scope not in ("sys", "gpu"):
      raise ValueError(f"Unsupported memory_scope: {memory_scope}")

    semantics = "relaxed" if relaxed else "release"
    inline_ptx(
        f"red.{semantics}.{memory_scope}.global.add.u32 [$0], $1;",
        self.ptr,
        value,
        predicate=predicate,
        has_side_effects=True,
    )

  @staticmethod
  def signal_multimem(ptr, value, predicate: ir.Value | None = None):
    i32 = ir.IntegerType.get_signless(32)
    if not isinstance(value, ir.Value):
      value = c(value, i32)
    elif value.type != i32:
      raise ValueError(f"Expected a i32 value, got {value.type}")
    if predicate is None:
      predicate = single_thread_predicate(ThreadSubset.WARPGROUP)
    inline_ptx(
        """{
            @$2 multimem.red.release.sys.global.add.u32 [$0], $1;
            fence.proxy.alias;
        }
        """,
        ptr,
        value,
        predicate,
        has_side_effects=True,
    )

  def wait(
      self,
      value: ir.Value | int = 1,
      *,
      decrement: bool = True,
      scope: ThreadSubset = ThreadSubset.WARPGROUP,
      memory_scope: Literal["sys", "gpu"] = "sys",
  ):
    i32 = ir.IntegerType.get_signless(32)
    if not isinstance(value, ir.Value):
      value = c(value, i32)
    elif value.type != i32:
      raise ValueError(f"Expected a i32 value, got {value.type}")
    if memory_scope not in ("sys", "gpu"):
      raise ValueError(f"Unsupported memory_scope: {memory_scope}")

    with single_thread(scope=scope):
      while_op = scf.WhileOp([i32], [value])
      before_block = while_op.before.blocks.append(i32)
      with ir.InsertionPoint.at_block_begin(before_block):
        [expected_in_memory] = before_block.arguments
        if decrement:
          in_memory = inline_ptx(
              f"atom.relaxed.{memory_scope}.global.cas.b32 $0, [$1], $2, $3;",
              self.ptr,
              expected_in_memory,
              arith.subi(expected_in_memory, value),
              result_types=i32,
              has_side_effects=True,
          )
          assert isinstance(in_memory, ir.Value)
          ne_pred = arith.CmpIPredicate.ne
          comparison = arith.cmpi(ne_pred, in_memory, expected_in_memory)
          new_expected_in_memory = arith.maxui(in_memory, value)
        else:
          in_memory = inline_ptx(
              f"ld.relaxed.{memory_scope}.global.b32 $0, [$1];",
              self.ptr,
              result_types=i32,
              has_side_effects=True,
          )
          lt_pred = arith.CmpIPredicate.ult
          comparison = arith.cmpi(lt_pred, in_memory, value)
          new_expected_in_memory = expected_in_memory
        scf.condition(comparison, [new_expected_in_memory])
      after_block = while_op.after.blocks.append(i32)
      with ir.InsertionPoint.at_block_begin(after_block):
        scf.yield_(after_block.arguments)
      inline_ptx(f"fence.acquire.{memory_scope};", has_side_effects=True)
    if scope == ThreadSubset.WARPGROUP:
      warpgroup_barrier()
    elif scope == ThreadSubset.WARP:
      warp_barrier()
    else:
      raise ValueError(f"Unsupported scope: {scope}")


def fence_release_sys():
  inline_ptx("fence.release.sys;", has_side_effects=True)


class Partition:
  source_bounds: tuple[int, ...]
  target_bounds: tuple[int, ...]
  partition: tuple[int | None, ...]
  base_offset: tuple[ir.Value, ...] | None

  def __init__(
      self,
      elements: tuple[int, ...],
      *,
      partition: tuple[int | None, ...],
      base_offset: tuple[ir.Value, ...] | None = None,
      num_chunks: tuple[int, ...] | None = None,
      chunk_size: tuple[int, ...] | None = None,
  ):
    self.target_bounds = elements
    self.partition = partition
    self.base_offset = base_offset
    if len(self.target_bounds) != len(self.partition):
      raise ValueError
    if num_chunks is None == chunk_size is None:
      raise ValueError(
          "Exactly one of num_chunks and chunk_size must be specified"
      )
    if num_chunks is not None:
      self.source_bounds = num_chunks
    else:
      assert chunk_size is not None
      if len(chunk_size) != len(self.target_bounds):
        raise ValueError
      source_bounds = []
      for els, chunk in zip(elements, chunk_size):
        if els % chunk:
          raise ValueError("Non-divisible partition", elements, chunk_size)
        source_bounds.append(els // chunk)
      self.source_bounds = tuple(source_bounds)

    seen_dims = set()
    for p in self.partition:
      if p is None:
        continue
      if not (0 <= p < len(self.source_bounds)):
        raise ValueError
      if p in seen_dims:
        raise ValueError
      seen_dims.add(p)
    for tb, p in zip(self.target_bounds, self.partition):
      if p is not None and tb % self.source_bounds[p]:
        raise ValueError("Non-divisible partitioning")

  @property
  def num_chunks(self) -> tuple[int, ...]:
    return self.source_bounds

  @property
  def target_block_shape(self):
    return tuple(
        tb if p is None else tb // self.source_bounds[p]
        for tb, p in zip(self.target_bounds, self.partition)
    )

  def get_base(self, *source_coords: ir.Value | int) -> list[ir.Value]:
    coords = []
    index = ir.IndexType.get()
    for i, (tbs, p) in enumerate(zip(self.target_block_shape, self.partition)):
      if p is None:
        dim_base = c(0, index)
      else:
        assert isinstance(source_coord := source_coords[p], ir.Value)
        dim_base = arith.muli(c(tbs, index), source_coord)
      if self.base_offset is not None:
        dim_base = arith.addi(self.base_offset[i], dim_base)
      coords.append(dim_base)
    return coords


class Partition1D:
  partition: Partition

  def __init__(
      self,
      elements: int,
      *,
      base_offset: ir.Value | None = None,
      num_chunks: int | None = None,
      chunk_size: int | None = None,
  ):
    self.base_offset = base_offset
    if num_chunks is None == chunk_size is None:
      raise ValueError(
          "Exactly one of num_chunks and chunk_size must be specified"
      )
    common_kwargs: dict[str, Any] = dict(elements=(elements,), partition=(0,))
    if base_offset is not None:
      common_kwargs["base_offset"] = (base_offset,)
    if num_chunks is not None:
      self.partition = Partition(num_chunks=(num_chunks,), **common_kwargs)
    else:
      assert chunk_size is not None
      self.partition = Partition(chunk_size=(chunk_size,), **common_kwargs)

  @property
  def num_chunks(self) -> int:
    return self.partition.source_bounds[0]

  def get_base(self, source_coords: ir.Value) -> ir.Value:
    return self.partition.get_base(source_coords)[0]

  def refine(
      self,
      *,
      chunk: ir.Value | None = None,
      num_chunks: int | None = None,
      chunk_size: int | None = None,
  ):
    return Partition1D(
        self.partition.target_block_shape[0],
        num_chunks=num_chunks,
        chunk_size=chunk_size,
        base_offset=self.get_base(chunk) if chunk is not None else None,
    )


def tile_shape(shape, tiling):
  if len(tiling) > len(shape):
    raise ValueError(
        "Expected tiling to be at most rank of shape. Got tiling:"
        f" {tiling} (rank: {len(tiling)}) and shape {shape} (rank:"
        f" {len(shape)})."
    )
  if not tiling:
    return shape
  tiling_rank = len(tiling)
  for s, t in zip(shape[-tiling_rank:], tiling):
    if s % t:
      raise ValueError("Non-divisible tiling:", shape, tiling)
  return (
      *shape[:-tiling_rank],
      *(s // t for s, t in zip(shape[-tiling_rank:], tiling)),
      *tiling,
  )


def warp_tree_reduce(value, op, group_size):
  """Reduce a value across the warpgroup."""
  assert bytewidth(value.type) == 4
  assert 32 % group_size == 0 and group_size <= 32
  i32 = ir.IntegerType.get_signless(32)
  result = value
  iters = np.log2(group_size)
  if not iters.is_integer():
    raise ValueError(
        f"Warp reduction group size should be a power of 2 (got {group_size})"
    )
  iters = int(iters)
  for i in range(iters):
    other_result = nvvm.shfl_sync(
        c(0xFFFFFFFF, i32),
        result,
        c(1 << i, i32),
        c(0x1F, i32),
        nvvm.ShflKind.bfly
    )
    result = op(result, other_result)

  return result


_MEMORY_SPACES = {f"#gpu.address_space<{str(x)}>": x for x in gpu.AddressSpace}


def get_memref_llvm_address_space(memref_ty: ir.MemRefType) -> int | None:
  if (memory_space := memref_ty.memory_space) is None:
    return None
  if isinstance(memory_space, ir.IntegerAttr):
    return memory_space.value
  return gpu_address_space_to_nvptx(_MEMORY_SPACES[str(memory_space)])


def memref_ptr(memref_arg) -> ir.Value:
  i64 = ir.IntegerType.get_signless(64)
  memref_ty = ir.MemRefType(memref_arg.type)
  address_space = get_memref_llvm_address_space(memref_ty)
  ptr_ty = llvm.PointerType.get(address_space)
  desc = to_memref_descriptor(memref_arg)
  aligned_ptr = llvm.extractvalue(ptr_ty, desc, [1])
  offset_elems = llvm.extractvalue(i64, desc, [2])

  elem_bitwidth = bitwidth(memref_ty.element_type)
  if elem_bitwidth < 8:
    *_, static_offset = memref_ty.get_strides_and_offset()
    if static_offset != ir.ShapedType.get_dynamic_stride_or_offset():
      assert elem_bitwidth.bit_count() == 1
      packing = 8 // elem_bitwidth
      if static_offset % packing != 0:
        raise ValueError(
            f"{memref_ty} {static_offset=} is not divisible by {packing=}`"
        )
      offset_bytes = c(static_offset // packing, i64)
    else:
      offset_bits = llvm.mul(
          offset_elems,
          c(elem_bitwidth, i64),
          overflow_flags=llvm.IntegerOverflowFlags.none,
      )
      offset_bytes = llvm.udiv(offset_bits, c(8, i64))
  else:
    assert elem_bitwidth % 8 == 0
    offset_bytes = llvm.mul(
        offset_elems,
        c(elem_bitwidth // 8, i64),
        overflow_flags=llvm.IntegerOverflowFlags.none,
    )
  return llvm.inttoptr(
      ptr_ty,
      llvm.add(
          llvm.ptrtoint(i64, aligned_ptr),
          offset_bytes,
          overflow_flags=llvm.IntegerOverflowFlags.none,
      ),
  )


def cluster_collective_mask(
    cluster_shape: tuple[int, int, int],
    collective: Sequence[gpu.Dimension] | gpu.Dimension,
    leader_tracked: bool = False,
):
  if isinstance(collective, gpu.Dimension):
    collective = (collective,)
  # We first compute the linearized index of the slice along the collective
  # dim that contains the current block. Then, the mask is a sequence of 1s
  # strided by the position of the collective dim, shifted left by the linear
  # slice index.
  # TODO(apaszke): Make sure this gets hoisted outside of any loops.
  # If not, we might need to do it manually.
  i32 = ir.IntegerType.get_signless(32)
  mask_shift = c(0, i32)
  # NOTE: GPU dimensions are minor-to-major.
  cluster_strides = get_contiguous_strides(cluster_shape[::-1])[::-1]
  for stride, cluster_dim in zip(cluster_strides, gpu.Dimension):
    if cluster_dim in collective:
      continue
    if cluster_shape[cluster_dim] != 1:  # Constant-fold multiply by 0.
      dim_idx = arith.index_castui(i32, gpu.cluster_block_id(cluster_dim))
      mask_shift = arith.addi(
          mask_shift,
          arith.muli(dim_idx, c(stride, i32)),
      )
  if leader_tracked:
    mask_unshifted = 1
  else:
    mask_unshifted = 0
    collective_strides = [cluster_strides[d] for d in collective]
    collective_shape = tuple(cluster_shape[d] for d in collective)
    for idx in np.ndindex(collective_shape):
      mask_unshifted |= 1 << sum(i * s for i, s in zip(idx, collective_strides))
  return arith.shli(c(mask_unshifted, i32), mask_shift)


def dtype_to_ir_type(dtype: jax.typing.DTypeLike) -> ir.Type:
  dtype = jnp.dtype(dtype)
  if jnp.issubdtype(dtype, jnp.integer):
    # All integer types in Mosaic GPU are signless.
    return ir.IntegerType.get_signless(jnp.iinfo(dtype).bits)
  return mlir.dtype_to_ir_type(dtype)


def is_signed(dtype: jax.typing.DTypeLike) -> bool | None:
  if jnp.issubdtype(dtype, jnp.bool_):
    return False
  elif jnp.issubdtype(dtype, jnp.integer):
    return jnp.issubdtype(dtype, jnp.signedinteger)
  return None


def getelementptr(
    ptr: ir.Value, indices: Sequence[ir.Value | int], dtype: ir.Type
) -> ir.Value:
  static_indices = [i if isinstance(i, int) else DYNAMIC32 for i in indices]
  dyn_indices = [i for i in indices if not isinstance(i, int)]
  return llvm.getelementptr(
      ptr.type,
      ptr,
      dyn_indices,
      static_indices,
      dtype,
      llvm.GEPNoWrapFlags.none,
  )


def dyn_dot(x, y):
  assert len(x) == len(y)
  return functools.reduce(arith.addi, (arith.muli(a, b) for a, b in zip(x, y)))


def shfl_bfly(x: ir.Value, distance: int | ir.Value):
  i32 = ir.IntegerType.get_signless(32)
  if isinstance(distance, int):
    distance = c(distance, i32)
  if (result_type := x.type) != i32:
    if (x_bitwidth := bitwidth(x.type)) < 32:  # Pad to 32-bits if necessary.
      assert 32 % x_bitwidth == 0
      x = bitcast(x, ir.IntegerType.get_signless(x_bitwidth))
      empty32 = llvm.mlir_undef(ir.VectorType.get((32 // x_bitwidth,), x.type))
      x = vector.insert(
          x,
          empty32,
          dynamic_position=[],
          static_position=ir.DenseI64ArrayAttr.get([0]),
      )
    elif x_bitwidth > 32:
      assert x_bitwidth % 32 == 0
      num_words = x_bitwidth // 32
      xs_vec = bitcast(x, ir.VectorType.get((num_words,), i32))
      y = llvm.mlir_undef(xs_vec.type)
      for i in range(num_words):
        x_elem = vector.extract(
            xs_vec,
            dynamic_position=[],
            static_position=ir.DenseI64ArrayAttr.get([i]),
        )
        y_elem = shfl_bfly(x_elem, distance)
        y = vector.insert(
            y_elem,
            y,
            dynamic_position=[],
            static_position=ir.DenseI64ArrayAttr.get([i]),
        )
      return bitcast(y, result_type)
    x = bitcast(x, i32)
  y = nvvm.shfl_sync(
      c(0xFFFFFFFF, i32),
      x,
      distance,
      c(0x1F, i32),
      nvvm.ShflKind.bfly
  )
  if (x_bitwidth := bitwidth(result_type)) < 32:
    bits_ty = ir.IntegerType.get_signless(x_bitwidth)
    y_vec = bitcast(y, ir.VectorType.get((32 // x_bitwidth,), bits_ty))
    y = vector.extract(
        y_vec,
        dynamic_position=[],
        static_position=ir.DenseI64ArrayAttr.get([0]),
    )
  return bitcast(y, result_type)


ReductionKind = nvvm.ReductionKind


def redux(x: ir.Value, mask: ir.Value, kind: ReductionKind):
  i32 = ir.IntegerType.get_signless(32)
  if isinstance(vec_ty := x.type, ir.VectorType):
    if bitwidth(vec_ty.element_type) != 32:
      raise ValueError("Only 32-bit types supported")
    [vec_len] = vec_ty.shape
    result = llvm.mlir_undef(x.type)
    for i in range(vec_len):
      xi = llvm.extractelement(x, arith.constant(i32, i))
      yi = redux(xi, mask, kind)
      result = llvm.insertelement(result, yi, arith.constant(i32, i))
    return result
  if bitwidth(x.type) != 32:
    raise ValueError("Only 32-bit scalar types supported")
  if isinstance(x.type, ir.IntegerType):
    pass
  elif isinstance(x.type, ir.F32Type):
    if get_arch().major != 10:
      raise ValueError("F32 redux only supported on Blackwell GPUs")
  else:
    raise NotImplementedError(x.type)
  assert mask.type == i32
  extra_kwargs: dict[str, Any] = {}
  if kind == ReductionKind.FMAX or kind == ReductionKind.FMIN:
    extra_kwargs = dict(nan=True)
  return nvvm.redux_sync(x, kind, mask, **extra_kwargs)


def prmt(high: ir.Value, low: ir.Value, permutation: ir.Value):
  i32 = ir.IntegerType.get_signless(32)
  if (result_type := high.type) != low.type:
    raise ValueError(f"Types must match, got {high.type} and {low.type}")
  if high.type != i32:
    high = bitcast(high, i32)
  if low.type != i32:
    low = bitcast(low, i32)
  if permutation.type != i32:
    permutation = bitcast(permutation, i32)
  result = inline_ptx(
      "prmt.b32 $0, $1, $2, $3;", high, low, permutation, result_types=i32
  )
  return bitcast(result, result_type)


def bitcast(x: ir.Value, new_type: ir.Type):
  if x.type == new_type:
    return x
  if (x_bw := bitwidth(x.type)) != (new_bw := bitwidth(new_type)):
    raise ValueError(
        f"Can't bitcast {x.type} (of bitwidth {x_bw}) to {new_type} (of"
        f" bitwidth {new_bw})"
    )
  if isinstance(x.type, ir.VectorType) and isinstance(new_type, ir.IntegerType):
    new_type = ir.IntegerType(new_type)
    x_ty = ir.VectorType(x.type)
    assert new_type.width == bitwidth(x_ty.element_type) * math.prod(x_ty.shape)
    return vector.extract(
        vector.bitcast(ir.VectorType.get((1,), new_type), x),
        dynamic_position=[],
        static_position=ir.DenseI64ArrayAttr.get([0]),
    )
  if isinstance(x.type, ir.IntegerType) and isinstance(new_type, ir.VectorType):
    new_type = ir.VectorType(new_type)
    x_ty = ir.IntegerType(x.type)
    assert x_ty.width == bitwidth(new_type.element_type) * math.prod(
        new_type.shape
    )
    return vector.bitcast(
        new_type, vector.broadcast(ir.VectorType.get((1,), x_ty), x)
    )
  if isinstance(x.type, ir.VectorType) and isinstance(new_type, ir.VectorType):
    x_ty = ir.VectorType(x.type)
    new_ty = ir.VectorType(new_type)
    if bitwidth(x_ty) != bitwidth(new_ty):
      raise ValueError(f"Can't bitcast {x.type} to {new_type}")
    return vector.bitcast(new_type, x)
  if isinstance(x.type, ir.IntegerType) and isinstance(new_type, ir.FloatType):
    return arith.bitcast(new_type, x)
  if isinstance(x.type, ir.FloatType) and isinstance(new_type, ir.IntegerType):
    return arith.bitcast(new_type, x)
  if isinstance(x.type, ir.FloatType) and isinstance(new_type, ir.FloatType):
    return arith.bitcast(new_type, x)
  raise ValueError(f"Can't bitcast {x.type} to {new_type}")


def ceil_div(x: int, y: int):
  return (x + y - 1) // y


def vector_slice(v: ir.Value, s: slice):
  v_ty = ir.VectorType(v.type)
  if len(v_ty.shape) != 1:
    raise NotImplementedError(f"Only 1D vectors are supported {v_ty}")
  [v_len] = v_ty.shape
  slice_length = len(range(v_len)[s])
  return vector.extract_strided_slice(
      ir.VectorType.get((slice_length,), v_ty.element_type),
      v,
      [s.start or 0],
      [slice_length],
      [1],
  )


def vector_concat(
    vectors: Sequence[ir.Value[ir.VectorType]],
) -> ir.Value[ir.VectorType]:
  if not vectors:
    raise ValueError("Cannot concatenate an empty list of vectors")
  vty = vectors[0].type
  if not isinstance(vty, ir.VectorType):
    raise ValueError("Cannot concatenate non-vector values")
  vty = ir.VectorType(vty)
  if vty.rank != 1:
    raise NotImplementedError("Only 1D vectors are supported")
  for v in vectors:
    if v.type.element_type != vty.element_type:
      raise ValueError("Cannot concatenate vectors of different element types")
    if v.type.rank != 1:
      raise ValueError("Can only concatenate 1D vectors")
  return _vector_concat_rec(vectors)


def _vector_concat_rec(
    vectors: Sequence[ir.Value[ir.VectorType]],
) -> ir.Value[ir.VectorType]:
  match vectors:
    case [v]:
      return v
    case [v, w]:
      [v_len] = ir.VectorType(v.type).shape
      [w_len] = ir.VectorType(w.type).shape
      mask = ir.DenseI64ArrayAttr.get(list(range(v_len + w_len)))
      return vector.shuffle(*vectors, mask=mask)
    case _:
      assert vectors
      l = _vector_concat_rec(vectors[: len(vectors) // 2])
      r = _vector_concat_rec(vectors[len(vectors) // 2 :])
      return _vector_concat_rec([l, r])


def is_known_divisible(value: ir.Value, divisor: int, max_depth=10) -> bool:
  """Returns True if the value is statically known to be divisible by the divisor."""
  if divisor == 1:
    return True
  if max_depth < 0:
    return False

  new_depth = max_depth - 1
  def_op = value.owner

  match def_op:
    case ir.Block() as block:
      op = block.owner
      if isinstance(op, dialect.WarpMapOp):
        arg_index = list(block.arguments).index(value)  # pyrefly: ignore[bad-argument-type]
        operand = op.operands[arg_index]
        return is_known_divisible(operand, divisor, new_depth)
      return False
    case dialect.AssumeMultipleOp() as op2:
      return ir.IntegerAttr(
          op2.multiple
      ).value % divisor == 0 or is_known_divisible(op2.value, divisor, new_depth)
    case arith.IndexCastOp():
      return is_known_divisible(def_op.in_, divisor, max_depth - 1)
    case arith.ConstantOp():
      return def_op.literal_value % divisor == 0
    case arith.MulIOp():
      # Only cover the case where one operand is divisible. It's still possible
      # that the final product is divisible, but we don't check that here.
      return is_known_divisible(
          def_op.lhs, divisor, new_depth
      ) or is_known_divisible(def_op.rhs, divisor, new_depth)
    case arith.SelectOp():
      return is_known_divisible(
          def_op.true_value, divisor, new_depth
      ) and is_known_divisible(def_op.false_value, divisor, new_depth)
    case arith.MaxSIOp() | arith.MinSIOp() | arith.MaxUIOp() | arith.MinUIOp():
      return is_known_divisible(
          def_op.lhs, divisor, new_depth
      ) and is_known_divisible(def_op.rhs, divisor, new_depth)
    case arith.AddIOp() | arith.SubIOp() | arith.RemUIOp() | arith.RemSIOp():
      # Only cover the common case where both operands are divisible.
      return is_known_divisible(
          def_op.lhs, divisor, new_depth
      ) and is_known_divisible(def_op.rhs, divisor, new_depth)
    case arith.AndIOp():
      # Only cover the specific case where the divisor is a power of two.
      return divisor.bit_count() == 1 and (
          is_known_divisible(def_op.lhs, divisor, new_depth)
          or is_known_divisible(def_op.rhs, divisor, new_depth)
      )
    case arith.TruncIOp():
      # Only cover the specific case where the divisor is a power of two.
      # trunci(a, bitwidth) = a % 2**bitwidth = a - k * 2**bitwidth for some k.
      # When the divisor is a power of 2, there are two cases:
      #   1. the divisor is smaller than 2**bitwidth. In this case, the divisor
      #      divides 2**bitwidth; if it divides a, it must thus also divide
      #      the truncated value a - k*2**bitwidth for any k;
      #   2. the divisor is larger than 2**bitwidth. In this case, if the
      #      divisor divides a, then we can conclude that a % 2**bitwidth == 0,
      #      and thus that the divisor also divides the truncated value.
      return (divisor.bit_count() == 1 and
              is_known_divisible(def_op.in_, divisor, new_depth))

  logger.debug("Unsupported defining operation %s when "
               "checking divisibility of %s", def_op, value)

  return False


def smem() -> ir.Attribute:
  """Returns the attribute for the SMEM memory space."""
  return ir.Attribute.parse("#gpu.address_space<workgroup>")


def tmem() -> ir.Attribute:
  """Returns the attribute for the TMEM memory space."""
  return ir.Attribute.parse("#mosaic_gpu.tmem")


def smem_cluster() -> ir.Attribute:
  """Returns the attribute for the cluster SMEM memory space."""
  return ir.Attribute.parse("#mosaic_gpu.smem_cluster")


def is_smem_ref(ref: ir.Value | ir.Type) -> bool:
  """Returns true if the input mem ref or memref type points to SMEM.

  If the input is not at all of a memref type, raises a ValueError.
  """
  if isinstance(ref, ir.Value):
    ref = ref.type
  if not isinstance(ref, ir.MemRefType):
    raise ValueError(f"Expected a memref type but got {ref}")
  ref = ir.MemRefType(ref)
  return ref.memory_space is not None and ref.memory_space == smem()


def is_tmem_ref(ref: ir.Value | ir.Type) -> bool:
  """Returns true if the input mem ref or memref type points to TMEM.

  If the input is not at all of a memref type, raises a ValueError.
  """
  if isinstance(ref, ir.Value):
    ref = ref.type
  if not isinstance(ref, ir.MemRefType):
    raise ValueError(f"Expected a memref type but got {ref}")
  ref = ir.MemRefType(ref)
  return ref.memory_space is not None and ref.memory_space == tmem()


def is_cluster_smem_ref(ref: ir.Value | ir.Type) -> bool:
  """Returns true if the input mem ref or memref type points to cluster SMEM."""
  if isinstance(ref, ir.Value):
    ref = ref.type
  if not isinstance(ref, ir.MemRefType):
    raise ValueError(f"Expected a memref type but got {ref}")
  ref = ir.MemRefType(ref)
  return ref.memory_space is not None and ref.memory_space == smem_cluster()


def try_cluster_cancel(
    result_ref,
    barrier: BarrierRef,
    predicate: ir.Value | None = None,
):
  """Atomically cancels a pending cluster launch.

  The response is stored in a opaque 128-bit value containing the CTA id of the
  first CTA in the canceled cluster.
  """
  if predicate is None:
    predicate = single_thread_predicate(ThreadSubset.BLOCK)
  inline_ptx(
      "clusterlaunchcontrol.try_cancel.async.shared::cta.mbarrier::complete_tx::bytes.multicast::cluster::all.b128 [$0], [$1];",
      memref_ptr(result_ref),
      barrier.get_ptr(),
      predicate=predicate,
      has_side_effects=True,
  )


def query_cluster_cancel(
    result_ref,
) -> tuple[ir.Value, ir.Value, ir.Value, ir.Value]:
  """Decodes the response of `try_cluster_cancel`.

  It checks if the cancellation was successful, and if yes, it also extracts
  the CTA ID of the first CTA in the canceled cluster.
  """

  i32 = ir.IntegerType.get_signless(32)
  i1 = ir.IntegerType.get_signless(1)
  struct_ty = llvm.StructType.get_literal([i32, i32, i32, i1])
  desc = llvm.inline_asm(
      struct_ty,
      [memref_ptr(result_ref)],
      """
    {
        .reg .b128 handle;
        ld.shared.b128 handle, [$4];
        clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 $3, handle;
        @$3 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {$0, $1, $2, _},  handle;
    }""",
      "=r,=r,=r,=b,r",
      has_side_effects=True,
  )
  assert isinstance(desc, ir.Value)
  cta_id_x = llvm.extractvalue(i32, desc, [0])
  cta_id_y = llvm.extractvalue(i32, desc, [1])
  cta_id_z = llvm.extractvalue(i32, desc, [2])
  cancelled_launch = llvm.extractvalue(i1, desc, [3])
  return cta_id_x, cta_id_y, cta_id_z, cancelled_launch


def nanosleep(nanos: ir.Value):
  """Sleeps the current thread for the given number of nanoseconds."""
  inline_ptx("nanosleep.u32 $0;", nanos, has_side_effects=True)


def cluster_idx(
    dim: gpu.Dimension | Sequence[gpu.Dimension] | None = None,
    dim_idx: ir.Value | Sequence[ir.Value] | None = None,
) -> ir.Value:
  """Returns the linear index of a block within a subset of the cluster spanned by the given dimensions.

  dim_idx can be used to specify the index of another block along the selected
  dimensions. If not provided, the current block's index is used.
  """
  if dim is None:
    dim = tuple(gpu.Dimension)
  elif isinstance(dim, gpu.Dimension):
    dim = (dim,)
  if dim_idx is None:
    dim_idx = [gpu.cluster_block_id(d) for d in dim]
  elif isinstance(dim_idx, ir.Value):
    if len(dim) != 1:
      raise ValueError(
          "Expected a single dimension when passing a single index"
      )
    dim_idx = [dim_idx]
  index = ir.IndexType.get()
  stride = c(1, index)
  lin_idx = c(0, index)
  for d, idx in sorted(zip(dim, dim_idx, strict=True), key=lambda x: x[0]):
    lin_idx = arith.addi(lin_idx, arith.muli(idx, stride))
    stride = arith.muli(stride, gpu.cluster_dim_blocks(d))
  return lin_idx


def get_cluster_ptr(
    ptr: ir.Value, cluster_block: ir.Value, generic: bool = True
):
  i32 = ir.IntegerType.get_signless(32)
  assert cluster_block.type == i32, cluster_block.type
  assert ptr.type == llvm.PointerType.get(3), ptr.type
  mapped_smem_ptr = inline_ptx(
      "mapa.shared::cluster.u32 $0, $1, $2;",
      ptr,
      cluster_block,
      result_types=llvm.PointerType.get(7),
  )
  if not generic:
    return mapped_smem_ptr
  return llvm.addrspacecast(llvm.PointerType.get(), mapped_smem_ptr)


def get_cluster_ref(
    ref: ir.Value, dim: gpu.Dimension, idx: ir.Value, generic: bool = True
):
  i32 = ir.IntegerType.get_signless(32)
  i64 = ir.IntegerType.get_signless(64)
  ref_ty = ir.MemRefType(ref.type)
  if not is_smem_ref(ref_ty):
    raise ValueError(f"Expected SMEM but got: {ref_ty.memory_space}")
  result_type = ir.MemRefType.get(
      ref_ty.shape,
      ref_ty.element_type,
      ref_ty.layout,
      None if generic else ir.IntegerAttr.get(i32, 7),
  )
  idxs: list[ir.Value] = [gpu.cluster_block_id(d) for d in gpu.Dimension]
  idxs[dim] = idx
  flat_block = arith.index_cast(i32, cluster_idx(dim_idx=idxs))
  desc = to_memref_descriptor(ref)
  aligned_ptr = llvm.extractvalue(
      llvm.PointerType.get(WORKGROUP_NVPTX_ADDRESS_SPACE), desc, [1]
  )
  cluster_ptr = get_cluster_ptr(aligned_ptr, flat_block, generic)
  result_desc = llvm.mlir_undef(memref_descriptor_type(result_type))
  result_desc = llvm.insertvalue(result_desc, cluster_ptr, [0])
  result_desc = llvm.insertvalue(result_desc, cluster_ptr, [1])
  result_desc = llvm.insertvalue(
      result_desc, llvm.extractvalue(i64, desc, [2]), [2]
  )
  if ref_ty.rank > 0:
    array_ty = llvm.ArrayType.get(i64, ref_ty.rank)
    result_desc = llvm.insertvalue(
        result_desc, llvm.extractvalue(array_ty, desc, [3]), [3]
    )
    result_desc = llvm.insertvalue(
        result_desc, llvm.extractvalue(array_ty, desc, [4]), [4]
    )
  result = builtin.unrealized_conversion_cast([result_type], [result_desc])
  assert isinstance(result, ir.Value)
  return result


def elements_to_bytes(offset: ir.Value, element_bitwidth: int) -> ir.Value:
  """Convert an element-based linear offset to a byte-based offset."""
  index_ty = offset.type

  if element_bitwidth > 8:
    return arith.muli(offset, c(element_bitwidth // 8, index_ty))
  elif element_bitwidth < 8:
    return arith.divsi(offset, c(8 // element_bitwidth, index_ty))
  else:
    return offset


@dataclasses.dataclass(frozen=True)
class Arch:
  major: int
  minor: int


def _infer_arch() -> tuple[int, int]:
  device: Any = jax.sharding.get_abstract_mesh().abstract_device
  default_device = jex_backend.get_default_device()
  if device is None:
    device = default_device
  elif (
      hasattr(default_device, "compute_capability")
      and device.device_kind == default_device.device_kind
  ):
    device = default_device
  if not hasattr(device, "compute_capability"):
    return (9, 0)  # TODO(apaszke): Remove this once we figure out the export story.
  arch_name = device.compute_capability
  # Handle ROCm devices that return architecture strings like "gfxXXX".
  if arch_name.startswith("gfx"):
    raise ValueError(
        f"Mosaic GPU does not yet support AMD ROCm devices. "
        f"Got compute_capability: {arch_name}"
    )
  return tuple(map(int, arch_name.split(".")))  # pyrefly: ignore[bad-return]


def get_arch() -> Arch:
  try:
    ip = ir.InsertionPoint.current
  except ValueError:
    # Infer the architecture; this is needed when architecture-sensitive code
    # is called from test setup code and there is no active module.
    return Arch(*_infer_arch())
  block = ip.block
  op = block.owner
  while op is not None:
    if op.name == "builtin.module":
      arch_major = op.attributes["mosaic_gpu.arch_major"]
      arch_minor = op.attributes["mosaic_gpu.arch_minor"]
      assert isinstance(arch_major, ir.IntegerAttr)
      assert isinstance(arch_minor, ir.IntegerAttr)
      return Arch(arch_major.value, arch_minor.value)
    op = op.parent
  raise ValueError("Cannot retrieve the architecture: no module found")


def reduce_shape(
    shape: Sequence[int], axes: Sequence[int], keep_dims: bool = False
) -> tuple[int, ...]:
  res = []
  for i, dim in enumerate(shape):
    if i in axes:
      if keep_dims:
        res.append(1)
    else:
      res.append(dim)
  return tuple(res)