# 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 __future__ import annotations from collections.abc import Callable, Iterable, Sequence import dataclasses import enum import functools import itertools import math from typing import Any, Literal, Protocol, TypeAlias, assert_never, cast, overload, runtime_checkable import jax import jax.experimental.mosaic.gpu as mgpu from jaxlib.mlir import ir from jaxlib.mlir.dialects import arith from jaxlib.mlir.dialects import gpu from jaxlib.mlir.dialects import llvm from jaxlib.mlir.dialects import math as mlir_math from jaxlib.mlir.dialects import memref from jaxlib.mlir.dialects import nvvm from jaxlib.mlir.dialects import vector import numpy as np from . import utils try: from jax._src.lib import mosaic_gpu as mgpu_lib except ImportError: mgpu_lib = None WARPGROUP_SIZE = utils.WARPGROUP_SIZE WARP_SIZE = 32 WARPS_IN_WARPGROUP = WARPGROUP_SIZE // WARP_SIZE SMEM_BANKS = 32 SMEM_BANK_BYTES = 4 c = utils.c ReductionKind = nvvm.ReductionKind @dataclasses.dataclass(frozen=True) class Tiling: """A tiling expression describing a permutation of elements of an nd-array. To apply one level of tiling to an array, each of the trailing dimensions (up to the rank of the tile) is unfolded into two dimensions: first equal to the ratio of the dimension size and the tile size, and second equal to the tile size. Then, all newly unfolded minor dimensions are transposed to appear at the end. This expression describes multi-level tiling, by applying each element of `tiles` in sequence to the array. See https://openxla.org/xla/tiled_layout for a more detailed explanation. """ tiles: tuple[tuple[int, ...], ...] def __post_init__(self): if not self.tiles: return last_tile_rank = len(self.tiles[0]) for tile in self.tiles: if len(tile) > last_tile_rank: raise ValueError("Tiles must have a decreasing rank") if not tile: raise ValueError("Tiles must not be empty") if any(d <= 0 for d in tile): raise ValueError(f"Tile shape must only have positive sizes, got: {self.tiles}") last_tile_rank = len(tile) def __str__(self): return f"Tiling({''.join(map(str, self.tiles))})" def tile_shape(self, shape: tuple[int, ...]) -> tuple[int, ...]: """Computes the shape of an array after tiling.""" orig_shape = shape def fail(): raise ValueError(f"Tiling {self.tiles} does not apply to shape {orig_shape}") for tile in self.tiles: if len(tile) > len(shape): fail() untiled_dims, tiled_dims = shape[:-len(tile)], shape[-len(tile):] if any(s % t != 0 for s, t in zip(tiled_dims, tile)): fail() shape = (*untiled_dims, *(d // t for d, t in zip(tiled_dims, tile)), *tile) return shape def untile_shape(self, shape: tuple[int, ...]) -> tuple[int, ...]: """Computes the shape of an array before tiling from its tiled shape.""" orig_shape = shape def fail(): raise ValueError( f"shape {orig_shape} is not a valid result of applying tiling {self}." ) for tile in reversed(self.tiles): if len(tile) > len(shape): fail() untiled_dims = shape[:-2 * len(tile)] tiled_dims = shape[-2 * len(tile):-len(tile)] tiling_dims = shape[-len(tile):] if tiling_dims != tile: fail() shape = (*untiled_dims, *(d * t for d, t in zip(tiled_dims, tile))) return shape def canonicalize(self) -> Tiling: """Returns a canonicalized version of the tiling. We define a tiling to be canonical if, at each step (except the first one, which defines the base tile shape): 1. The tiling partitions at least one dimension in more than 1 tile. For example, the tiling `(8, 8)(8, 8)` is not canonical, as applying it yields a shape `(1, 1, 8, 8)`. We canonicalize it to `(8, 8)`, which allows getting rid of the unnecessary `1` dimensions. 2. The leading dimensions of each tile are not `1`. If canonicalizing a tile in this way leads to an empty tile, then the tile is given shape `(1,)`---which is still a meaningful (final) tile. For example, the tiling `(8, 8)(1, 4)` is not canonical, as applying it yields a shape `(8, 2, 1, 4)`. We canonicalize it to `(8, 8)(4,)`, which allows getting rid of the unnecessary `1` dimension, and yields a shape `(8, 2, 4)`. """ if len(self.tiles) <= 1: return self shape = self.tiles[0] new_tiling = [self.tiles[0]] for tile in self.tiles[1:]: for i, d in enumerate(tile): if d != 1: canonical_tile = tile[i:] break else: canonical_tile = (1,) tiled_dims = shape[-len(canonical_tile):] if tiled_dims == canonical_tile: continue shape = canonical_tile new_tiling.append(canonical_tile) return Tiling(tuple(new_tiling)) def tile_strides(self, strides: tuple[int, ...]) -> tuple[int, ...]: """Computes the strides of an array after tiling.""" for tile in self.tiles: untiled, tiled = strides[:-len(tile)], strides[-len(tile):] strides = (*untiled, *(s * t for s, t in zip(tiled, tile)), *tiled) return strides def tile_dimension(self, dim: int) -> tuple[bool, ...]: """Result is True whenever the tiled dim originated from the given input dim.""" tiling_rank = len(self.tiles[0]) if dim < 0 or dim >= tiling_rank: raise ValueError(f"Invalid dimension {dim} for tiling {self}") strides = [1] * tiling_rank strides[dim] = 0 return tuple(s == 0 for s in self.tile_strides(tuple(strides))) def remove_dimension(self, dim: int) -> Tiling: """Returns a tiling with the given dimension removed.""" tiling_rank = len(self.tiles[0]) if dim < 0 or dim >= tiling_rank: raise ValueError(f"Invalid dimension {dim} for tiling {self}") dim_in_tile = dim tiles = [] last_tile_rank = len(self.tiles[0]) for t in self.tiles: assert last_tile_rank >= len(t) dim_in_tile -= last_tile_rank - len(t) last_tile_rank = len(t) if dim_in_tile >= 0: t = t[:dim_in_tile] + t[dim_in_tile + 1:] if not t: # If this tile is empty, all other tiles will be empty too. break tiles.append(t) return Tiling(tuple(tiles)) def tile_nested_shape_strides( self, shape: tuple[tuple[int, ...], ...], strides: tuple[tuple[int, ...], ...], ) -> tuple[tuple[tuple[int, ...], ...], tuple[tuple[int, ...], ...]]: """A fused version of `tile_shape` and `tile_strides` for nested shapes. By nested shape we mean that each logical dimension (i.e. each element of shape/strides) is actually composed out of multiple physical dimensions. For example, a row-major array of logical shape (128, 128) that is tiled into (64, 64) tiles would have a nested shape ((2, 64), (2, 64)) (i.e. each dim is split into two sub-dims) and nested strides of ((2 * 64 * 64, 64), (64 * 64, 1)). """ if len(shape) != len(strides): raise ValueError( f"Shape {shape} and strides {strides} must have the same length" ) def fail_if(cond, shape=shape): # Capture shape now. if cond: raise ValueError(f"Tiling {self.tiles} does not apply to shape {shape}") for tile in self.tiles: fail_if(len(tile) > len(shape)) untiled_shape, tiled_shape = shape[:-len(tile)], shape[-len(tile):] untiled_strides, tiled_strides = strides[:-len(tile)], strides[-len(tile):] major_dim_shapes, major_dim_strides = [], [] minor_dim_shapes, minor_dim_strides = [], [] for t, dim_shape, dim_strides in zip(tile, tiled_shape, tiled_strides): major_dim_shape_rev, major_dim_stride_rev = [], [] minor_dim_shape_rev, minor_dim_stride_rev = [], [] for d, s in zip(reversed(dim_shape), reversed(dim_strides), strict=True): if d < t: # We will need to tile more dims fail_if(t % d != 0) t //= d minor_dim_shape_rev.append(d) minor_dim_stride_rev.append(s) elif t != 1: # Last dim to tile! fail_if(d % t != 0) minor_dim_shape_rev.append(t) minor_dim_stride_rev.append(s) if d != t: # No need to insert singleton dims. major_dim_shape_rev.append(d // t) major_dim_stride_rev.append(s * t) t = 1 else: # Done tiling! major_dim_shape_rev.append(d) major_dim_stride_rev.append(s) fail_if(t != 1) major_dim_shapes.append(major_dim_shape_rev[::-1]) minor_dim_shapes.append(minor_dim_shape_rev[::-1]) major_dim_strides.append(major_dim_stride_rev[::-1]) minor_dim_strides.append(minor_dim_stride_rev[::-1]) shape = (*untiled_shape, *major_dim_shapes, *minor_dim_shapes) # type: ignore[arg-type] strides = (*untiled_strides, *major_dim_strides, *minor_dim_strides) # type: ignore[arg-type] return ( tuple(tuple(d) if d else (1,) for d in shape), tuple(tuple(d) if d else (1,) for d in strides), ) def tile_indices(self, indices: tuple[int, ...]) -> tuple[int, ...]: for tile in self.tiles: untiled, tiled = indices[:-len(tile)], indices[-len(tile):] indices = ( *untiled, *(i // t for i, t in zip(tiled, tile)), *(i % t for i, t in zip(tiled, tile)), ) return indices def untile_indices(self, indices: tuple[int, ...]) -> tuple[int, ...]: for tile in reversed(self.tiles): untiled = indices[:-2 * len(tile)] outer = indices[-2 * len(tile):-len(tile)] inner = indices[-len(tile):] indices = (*untiled, *(o * t + i for o, i, t in zip(outer, inner, tile))) return indices class Rounding(enum.Enum): TO_NEAREST_EVEN = enum.auto() TO_POSITIVE_INFINITY = enum.auto() TO_ZERO = enum.auto() @property def ptx(self) -> str: match self: case Rounding.TO_NEAREST_EVEN: return "rn" case Rounding.TO_POSITIVE_INFINITY: return "rp" case Rounding.TO_ZERO: return "rz" @property def arith(self) -> arith.RoundingMode: match self: case Rounding.TO_NEAREST_EVEN: return arith.RoundingMode.to_nearest_even case Rounding.TO_POSITIVE_INFINITY: return arith.RoundingMode.upward case Rounding.TO_ZERO: return arith.RoundingMode.toward_zero def enumerate_negative[T](elems: Sequence[T]) -> Iterable[tuple[int, T]]: """Like built-in enumerate, but returns negative indices into the sequence.""" offset = len(elems) for i, e in enumerate(elems): yield i - offset, e @dataclasses.dataclass(frozen=True) class Replicated: times: int @dataclasses.dataclass(frozen=True) class Unreduced: """Indicates a dimension that has not yet been reduced across warps or lanes. Attributes: times: The factor by which the dimension is unreduced across warps or lanes. """ times: int @dataclasses.dataclass(frozen=True) class TiledLayout: """A FragmentedArray layout derived from a tiling expression. A logical array is transformed according to the tiling expression, and then split across warps (within a warpgroup), lanes, and vectorized according to the dimension indices. All dimension indices must be negative and should refer to the dimensions after tiling is applied. To better understand this layout, consider the example of WGMMA-related tiling from https://docs.nvidia.com/cuda/parallel-thread-execution/#wgmma-64n16-d as applied to a 128x128 array. The corresponding TiledLayout has a tiling of: (64, 8)(16, 8)(8, 8)(1, 2) and warp_dims=(-8,), lane_dims=(-4, -3), vector_dim=-1. We begin by applying the tiling (note that it always applies to a suffix): Tiled shape Remaining tiling actions =========================================================================== 128 128 (64, 8)(16, 8)(8, 8)(1, 2) 2 16 64 8 (16, 8)(8, 8)(1, 2) 2 16 4 1 16 8 (8, 8)(1, 2) 2 16 4 1 2 1 8 8 (1, 2) 2 16 4 1 2 1 8 4 1 2 The last expression is our final shape. At this stage, we're ready to partition the dimensions: warp_dims=(-8,) means that the 8-th dimension from the end is partitioned over 4 warps in a warpgroup (and so it must be of size 4). lane_dims=(-4, -3) indicate that those two dimensions are partitioned over the lanes within a warp (their product must be equal to 32, i.e. warp size). Finally, vector_dim=-1 indicates that each (logical) register is a vector containing 2 elements (there are no shape restrictions here). Given the above, the shape of the (logical) register array used to represent the array in each thread is: (2, 16, 1, 1, 2, 1, 1, 1, 1, 1). We have set all the dimensions above to 1, since each thread is a member of a single warp, a single lane, and the elements along the vectorized dimension are represented by a single (logical) register. """ tiling: Tiling warp_dims: tuple[int | Replicated | Unreduced, ...] # major-to-minor lane_dims: tuple[int | Replicated | Unreduced, ...] # major-to-minor vector_dim: int unreduced_operation: str | None = None # Whether to enforce that the layout is canonical. Users of `TiledLayout` # should not set this to `False`, but it is helpful to be able to construct # non-canonical layouts as an intermediate state when implementing layout # transformations. _check_canonical: dataclasses.InitVar[bool] = True def __post_init__(self, _check_canonical: bool): if not self.tiling.tiles: raise ValueError("Tiling must have at least one tile") min_shape = self.tiling.tiles[0] min_tiled_shape = self.tiling.tile_shape(min_shape) dims_set = { *self.partitioned_warp_dims, *self.partitioned_lane_dims, self.vector_dim, } if len(dims_set) != len(self.partitioned_warp_dims) + len(self.partitioned_lane_dims) + 1: raise ValueError("Duplicate partitioning dimensions") for d in dims_set: if d >= 0: raise ValueError("All dimensions must be negative") if d < -(len(min_tiled_shape) - len(min_shape)): raise ValueError("Dimension out of range") warp_dims_prod = math.prod( min_tiled_shape[d] if isinstance(d, int) else d.times for d in self.warp_dims ) if warp_dims_prod != WARPS_IN_WARPGROUP: raise ValueError( "The product of warp dims does not equal the number of warps in a" " warpgroup" ) lane_dims_prod = math.prod( min_tiled_shape[d] if isinstance(d, int) else d.times for d in self.lane_dims ) if lane_dims_prod != WARP_SIZE: raise ValueError("The product of lane dims does not equal the warp size") has_unreduced = any( isinstance(d, Unreduced) for d in itertools.chain(self.warp_dims, self.lane_dims) ) if self.unreduced_operation is not None: if not has_unreduced: raise ValueError( "unreduced_operation is set but layout has no unreduced dimensions" ) else: if has_unreduced: raise ValueError( "layout has unreduced dimensions but unreduced_operation is None" ) if _check_canonical: canonical_layout = self.canonicalize() if self != canonical_layout: raise ValueError(f"{self} is not canonical.") @property def has_unreduced_dims(self) -> bool: return self.unreduced_operation is not None @functools.cached_property def partitioned_warp_dims(self) -> tuple[int, ...]: return tuple(d for d in self.warp_dims if isinstance(d, int)) @functools.cached_property def partitioned_lane_dims(self) -> tuple[int, ...]: return tuple(d for d in self.lane_dims if isinstance(d, int)) def thread_idxs(self, shape: tuple[int, ...]) -> Iterable[tuple[ir.Value, ...]]: # We first find the linear index and then divide by the shape to # get the index. i32 = ir.IntegerType.get_signless(32) index = ir.IndexType.get() contig_strides = tuple(utils.get_contiguous_strides(shape)) tile_strides = self.tiling.tile_strides(contig_strides) dyn_tile_strides = [c(s, i32) for s in tile_strides[-self.tiled_tiling_rank:]] warp_offset = utils.dyn_dot(self.warp_indices(), dyn_tile_strides) lane_offset = utils.dyn_dot(self.lane_indices(), dyn_tile_strides) dyn_offset = arith.addi(warp_offset, lane_offset) register_shape = self.registers_shape(shape) for tile_idx in np.ndindex(register_shape): tile_lin_idx = sum(i * s for i, s in zip(tile_idx, tile_strides)) dyn_lin_idx = arith.addi(dyn_offset, c(tile_lin_idx, i32)) idx = [] for stride in contig_strides: idx.append(arith.index_castui(index, arith.divui(dyn_lin_idx, c(stride, i32)))) dyn_lin_idx = arith.remui(dyn_lin_idx, c(stride, i32)) yield tuple(idx) @property def base_tile_shape(self) -> tuple[int, ...]: """The shape of the first tile in the tiling expression. This tile acts as the divisibility constraint for a suffix of arrays to which this layout applies. """ return self.tiling.tiles[0] @functools.cached_property def tiled_tiling_shape(self) -> tuple[int, ...]: """The shape of the suffix of the array after tiling. We only allow our repeated tiling actions to further subdivide the dimensions created by previous tiling actions (except for the first one), so the tiled shape always ends with this suffix, no matter what array shape it's applied to. """ base_tile_shape = self.base_tile_shape return self.tiling.tile_shape(base_tile_shape)[len(base_tile_shape):] @functools.cached_property def tiled_tiling_rank(self) -> int: return len(self.tiled_tiling_shape) @property def vector_length(self) -> int: return self.tiled_tiling_shape[self.vector_dim] def registers_element_type(self, t: ir.Type) -> ir.Type: return ir.VectorType.get((self.vector_length,), t) def registers_shape(self, shape: tuple[int, ...]) -> tuple[int, ...]: """Returns the shape of the register array needed to represent an array of the given logical shape.""" tiled_shape = list(self.tiling.tile_shape(shape)) for d in self.partitioned_warp_dims: tiled_shape[d] = 1 for d in self.partitioned_lane_dims: tiled_shape[d] = 1 tiled_shape[self.vector_dim] = 1 return tuple(tiled_shape) def shape_from_registers_shape(self, shape: tuple[int, ...]) -> tuple[int, ...]: """Returns the logical shape of an array given its register array shape. Inverse to `registers_shape`. """ tiled_tiling = self.tiled_tiling_shape shape: list[int] = list(shape) # pyrefly: ignore[redefinition] for d in self.partitioned_warp_dims: shape[d] = tiled_tiling[d] for d in self.partitioned_lane_dims: shape[d] = tiled_tiling[d] shape[self.vector_dim] = tiled_tiling[self.vector_dim] return self.tiling.untile_shape(tuple(shape)) def _delinearize_index( self, idx: ir.Value, dims: tuple[int | Replicated | Unreduced, ...] ) -> tuple[ir.Value, ...]: i32 = ir.IntegerType.get_signless(32) tiled_shape = self.tiled_tiling_shape dims_shape = tuple( tiled_shape[d] if isinstance(d, int) else d.times for d in dims ) dims_strides = utils.get_contiguous_strides(dims_shape) dims_indices = tuple( arith.remui(arith.divui(idx, c(stride, i32)), c(size, i32)) for stride, size in zip(dims_strides, dims_shape) ) full_indices = [arith.constant(i32, 0)] * len(tiled_shape) for d, i in zip(dims, dims_indices): if isinstance(d, int): full_indices[d] = i return tuple(full_indices) def lane_indices(self, lane_idx: ir.Value | None = None) -> tuple[ir.Value, ...]: i32 = ir.IntegerType.get_signless(32) if lane_idx is None: lane_idx = arith.remui(utils.thread_idx(), c(WARP_SIZE, i32)) assert lane_idx.type == i32 return self._delinearize_index(lane_idx, self.lane_dims) def warp_indices(self) -> tuple[ir.Value, ...]: i32 = ir.IntegerType.get_signless(32) warp_idx = arith.remui( arith.divui(utils.thread_idx(), c(WARP_SIZE, i32)), c(WARPS_IN_WARPGROUP, i32), ) return self._delinearize_index(warp_idx, self.warp_dims) def remove_dimension( self, dim: int, *, local_only: bool = False, op: str | None = None, ) -> TiledLayout: if dim < 0 or dim >= len(self.tiling.tiles[0]): raise ValueError(f"Dimension {dim} is out of range for {self.tiling}") new_tiling = self.tiling.remove_dimension(dim) tiled_shape = self.tiled_tiling_shape removed_dim = self.tiling.tile_dimension(dim) dim_offsets = np.cumsum(removed_dim[::-1])[::-1].tolist() if removed_dim[self.vector_dim]: new_tiling = Tiling((*new_tiling.tiles, (1,))) new_vector_dim = -1 dim_offsets = [o - 1 for o in dim_offsets] # We inserted an extra dim. else: new_vector_dim = self.vector_dim + dim_offsets[self.vector_dim] def replace_dim(d: int | Replicated | Unreduced): if isinstance(d, (Replicated, Unreduced)): return d elif removed_dim[d]: size = tiled_shape[d] return Unreduced(size) if local_only else Replicated(size) else: return d + dim_offsets[d] new_warp_dims = tuple(replace_dim(d) for d in self.warp_dims) new_lane_dims = tuple(replace_dim(d) for d in self.lane_dims) has_unreduced = any( isinstance(d, Unreduced) for d in itertools.chain(new_warp_dims, new_lane_dims) ) new_unreduced_op = ( (op or self.unreduced_operation) if has_unreduced else None ) return TiledLayout( new_tiling, new_warp_dims, new_lane_dims, new_vector_dim, unreduced_operation=new_unreduced_op, _check_canonical=False, ).canonicalize() def reduce( self, axes: Sequence[int], *, local_only: bool = False, op: str | None = None, ) -> TiledLayout: """Returns the layout obtained by reducing `axes` out of this layout. Args: axes: The tiled axes to reduce. local_only: If true, dimensions distributed over warps or lanes become `Unreduced` rather than `Replicated`, leaving the cross-thread part of the reduction to a later `reduce`. op: The reduction operation. Required when `local_only` is true, since it must be recorded in the layout to be applied later. """ if local_only and op is None: raise ValueError("op must be specified for local-only reductions") if ( local_only and self.unreduced_operation is not None and self.unreduced_operation != op ): raise ValueError( f"Cannot partially reduce layout with op {op!r} when it is already" f" unreduced with respect to {self.unreduced_operation!r}" ) reduced_layout = self for a in sorted(axes, reverse=True): reduced_layout = reduced_layout.remove_dimension( a, local_only=local_only, op=op ) if not local_only and reduced_layout.has_unreduced_dims: # A full reduction also completes whatever reduction was left pending by # earlier local-only reductions. reduced_layout = dataclasses.replace( reduced_layout, lane_dims=tuple( Replicated(d.times) if isinstance(d, Unreduced) else d for d in reduced_layout.lane_dims ), warp_dims=tuple( Replicated(d.times) if isinstance(d, Unreduced) else d for d in reduced_layout.warp_dims ), unreduced_operation=None, ) return reduced_layout def canonicalize(self) -> TiledLayout: """Returns a version of this layout where tiling is canonical.""" canonical_tiling = self.tiling.canonicalize() s = self.base_tile_shape tiled_tiling_shape = self.tiled_tiling_shape canonical_tiled_tiling_shape = canonical_tiling.tile_shape(s)[len(s):] offset = len(canonical_tiled_tiling_shape) - 1 rev_removed_dims = [] # Iterate starting from the end in order to eliminate leading dimensions, # whenever possible. For instance, say we have # # shape=(4, 32, 1, 1, 1, 1, 1) # warp_dims=(-7,), # lane_dims=(-6,) # vector_dim=-1 # # and we want to canonicalize this to # # shape=(4, 32, 1) # warp_dims=(-3,), # lane_dims=(-2,) # vector_dim=-1. # # After the loop below, we end up with # # rev_removed_dims=[False, True, True, True, True, False, False] # # which will yield offsets `4` for `warp_dims[0]`, `4` for `lane_dims[0]`, # and `0` for `vector_dim`. for d in reversed(tiled_tiling_shape): if offset >= 0 and d == canonical_tiled_tiling_shape[offset]: rev_removed_dims.append(False) offset -= 1 else: rev_removed_dims.append(True) assert offset == -1 dim_offsets = np.cumsum(rev_removed_dims)[::-1].tolist() def replace_tiled_dim(d: int | Replicated | Unreduced): return d + dim_offsets[d] if isinstance(d, int) else d def is_nontrivial(d: int | Replicated | Unreduced): return not isinstance(d, int) or tiled_tiling_shape[d] != 1 return TiledLayout( canonical_tiling, tuple(replace_tiled_dim(d) for d in self.warp_dims if is_nontrivial(d)), tuple(replace_tiled_dim(d) for d in self.lane_dims if is_nontrivial(d)), replace_tiled_dim(self.vector_dim), unreduced_operation=self.unreduced_operation, _check_canonical=False, ) def find_dim_in_tiling(self, dim: int) -> tuple[int, int]: """Returns the index of the tile containing the given dimension and its position within the tile.""" assert dim < 0 dims = 0 tiling = self.tiling.tiles for dim_tile in range(1, len(tiling) + 1): dims -= len(tiling[-dim_tile]) if dim >= dims: dim_within_tile = dim - dims return -dim_tile, dim_within_tile raise ValueError(f"Dimension {dim} not found in tiling") def subdivide_tile(self, tile_idx: int, new_tile: tuple[int, ...]) -> TiledLayout: """Inserts new_tile right after tile_idx in the tiling.""" assert tile_idx <= -1 new_tiles = list(self.tiling.tiles) stable_dim_suffix = sum(map(len, new_tiles[tile_idx:])) assert len(new_tile) == len(new_tiles[tile_idx]) if tile_idx == -1: new_tiles.append(new_tile) else: new_tiles.insert(tile_idx + 1, new_tile) def adjust_dim(d: int | Replicated | Unreduced): if isinstance(d, int): if d < -stable_dim_suffix: yield d - len(new_tile) return elif d < -stable_dim_suffix + len(new_tile): yield d - len(new_tile) yield d return yield d if self.vector_dim < -stable_dim_suffix + len(new_tile): raise NotImplementedError("Multiple vector dimensions not supported.") return TiledLayout( Tiling(tuple(new_tiles)), tuple(itertools.chain.from_iterable(map(adjust_dim, self.warp_dims))), tuple(itertools.chain.from_iterable(map(adjust_dim, self.lane_dims))), self.vector_dim, unreduced_operation=self.unreduced_operation, _check_canonical=False, ).canonicalize() @property def replication_factor(self) -> int: replication_factor = 1 for dim in self.warp_dims: if isinstance(dim, Replicated): replication_factor *= dim.times for dim in self.lane_dims: if isinstance(dim, Replicated): replication_factor *= dim.times return replication_factor def _tiled_wgmma_layout(shape: tuple[int, ...]): """Returns the tiled layout relevant for WGMMA operations. The tiled layout is equivalent to one described here in PTX documentation: https://docs.nvidia.com/cuda/parallel-thread-execution/#wgmma-64n16-d """ if len(shape) != 2: raise ValueError(f"Shape {shape} is not 2D") if shape[0] % 64 != 0 or shape[1] % 8 != 0: raise ValueError(f"Shape {shape} is not a multiple of 64x8") return WGMMA_LAYOUT @dataclasses.dataclass(frozen=True) class WGSplatFragLayout: """A fragmented array where all the values are equal represented as a register per thread. FragmentedArrays in this layout can be are always the result of a splat, each thread in the warpgroup has a single copy of the value, while the FragmentedArray pretends it has whatever shape the user wants. This means we can trivially broadcast, reshape and do elementwise operations with all other layouts. Examples: To load a value in ``` FragmentedArray.splat(memref.load(ref_1d, [1]), (10,20,2)) ``` A shape is always provided for sanity check reasons. """ shape: tuple[int, ...] = () def can_broadcast_to(self, shape) -> bool: """Check that the shape can be broadcast. All source dimensions must match the target's trailing dimensions by equality or being set to 1 (i.e. we can broadcast 1-sized dimensions or create new leading dimensions). """ return len(self.shape) <= len(shape) and all( dim1 == dim2 or dim1 == 1 for dim1, dim2 in zip(self.shape[::-1], shape[::-1]) ) def registers_element_type(self, t: ir.Type) -> ir.Type: return t def registers_shape(self, shape: tuple[int, ...]) -> tuple[int, ...]: """Returns the shape of the register array needed to represent an array of the given logical shape.""" del shape # Unused. return () def shape_from_registers_shape( self, shape: tuple[int, ...] ) -> tuple[int, ...]: del shape # Unused. return self.shape def thread_idxs(self, shape): assert shape == self.shape raise NotImplementedError @dataclasses.dataclass(frozen=True) class WGStridedFragLayout: """Convert the array to 1D and then shard across threads.""" shape: tuple[int, ...] vec_size: int def __post_init__(self): if np.prod(self.shape) % (self.vec_size * WARPGROUP_SIZE) != 0: raise ValueError((self, WARPGROUP_SIZE)) @classmethod def from_shaped_type(cls, shaped_ty: ir.Type) -> WGStridedFragLayout | None: """Returns a WGStridedFragLayout for the given shaped type. Return None if the shaped type cannot have a strided layout. """ if not isinstance(shaped_ty, ir.ShapedType): raise TypeError(shaped_ty) shaped_ty = ir.ShapedType(shaped_ty) bitwidth = utils.bitwidth(shaped_ty.element_type) size = math.prod(shaped_ty.shape) if size % WARPGROUP_SIZE != 0: return None max_vec_size = size // WARPGROUP_SIZE vec_size = min(64 // bitwidth, max_vec_size) while vec_size > 0 and size % (vec_size * WARPGROUP_SIZE) != 0: vec_size //= 2 if vec_size == 0 or (vec_size * bitwidth) % 8 != 0: return None return cls(shape=tuple(shaped_ty.shape), vec_size=vec_size) def registers_element_type(self, t: ir.Type) -> ir.Type: return ir.VectorType.get((self.vec_size,), t) def registers_shape(self, shape: tuple[int, ...]) -> tuple[int, ...]: """Returns the shape of the register array needed to represent an array of the given logical shape.""" if shape != self.shape: raise ValueError(f"Shape {shape} is not compatible with {self}") return (math.prod(self.shape) // (WARPGROUP_SIZE * self.vec_size),) def shape_from_registers_shape( self, shape: tuple[int, ...] ) -> tuple[int, ...]: del shape # Unused. return self.shape def thread_idxs(self, shape): assert shape == self.shape index = ir.IndexType.get() for v in self.linear_thread_idxs(): res = [] for dim in reversed(self.shape): dim = c(dim, index) res.append(arith.remui(v, dim)) v = arith.divui(v, dim) res.reverse() yield res def linear_thread_idxs(self): """The indexes to be used for vector load/store WGStridedFragLayout. Yields: The indices of the vector that correspond to the current thread. """ index = ir.IndexType.get() cardinality = np.prod(self.shape) assert cardinality % (WARPGROUP_SIZE * self.vec_size) == 0 reg_num = cardinality // (WARPGROUP_SIZE * self.vec_size) tidx = arith.remui(gpu.thread_id(gpu.Dimension.x), c(WARPGROUP_SIZE, index)) off = arith.muli(tidx, c(self.vec_size, tidx.type)) for i in range(reg_num): yield arith.addi(off, c(i * WARPGROUP_SIZE * self.vec_size, tidx.type)) FragmentedLayout: TypeAlias = ( WGSplatFragLayout | WGStridedFragLayout | TiledLayout ) WGMMA_COL_LAYOUT = TiledLayout( Tiling(((8,), (2,))), warp_dims=(Replicated(4),), lane_dims=(Replicated(8), -2), vector_dim=-1, ) WGMMA_ROW_LAYOUT = TiledLayout( Tiling(((64,), (16,), (8,), (1,))), warp_dims=(-4,), lane_dims=(-2, Replicated(4)), vector_dim=-1, ) # The tiled layout is equivalent to one described here in PTX documentation: # https://docs.nvidia.com/cuda/parallel-thread-execution/#wgmma-64n16-d # In this layout, we partition the 64x8 tiles over 4 warps into 16x8 tiles. # Then, we further split the 16x8 tiles into 8x8 submatrices which are the unit # of data that is split across a warp. Since 8*8 = 64, but a warp has only 32 # threads, we vectorize pairs of elements along columns. # The assignment of elements to warp lanes is as follows: # # 0 0 1 1 2 2 3 3 # 4 4 5 5 6 6 7 7 # 8 8 9 9 10 10 11 11 # 12 12 13 13 14 14 15 15 # ... WGMMA_LAYOUT = TiledLayout( Tiling(((64, 8), (16, 8), (8, 8), (2,))), warp_dims=(-7,), lane_dims=(-3, -2), vector_dim=-1, ) # This is the same as WGMMA_LAYOUT, only with a vector length of 1. LLVM now # treats <2 x float> as a native PTX type and uses 64-bit registers to store # them. This, in turn, means that we have to explode them into 32-bit registers # right before WGMMA, which makes ptxas very unhappy and causes it to insert # lots of WGMMA waits that absolutely tank the performance. As a workaround, # we use this layout when 32-bit data with WGMMA_LAYOUT is used to initialize # a WGMMAAccumulator, to ensure that the LLVM accumulator registers will always # be represented as 32-bit PTX registers. WGMMA_LAYOUT_ACC_32BIT = TiledLayout( Tiling(((64, 8), (16, 8), (8, 8), (2,), (1,))), warp_dims=(-8,), lane_dims=(-4, -3), vector_dim=-1, ) # The tiled layout is equivalent to one described here in PTX documentation: # https://docs.nvidia.com/cuda/parallel-thread-execution/#wgmma-64n32-a # In this layout, we partition the 64x16 tiles over 4 warps into 16x16 tiles. # Then, we further split the 16x16 tiles into 8x16 submatrices which are the unit # of data that is split across a warp. Since 8*16 = 128, but a warp has only 32 # threads, we vectorize quadruplets of elements along columns. # The assignment of elements to warp lanes is as follows: # # 0 0 0 0 1 1 1 1 2 2 2 2 3 3 3 3 # 4 4 4 4 5 5 5 5 6 6 6 6 7 7 7 7 # 8 8 8 8 9 9 9 9 10 10 10 10 11 11 11 11 # 12 12 12 12 13 13 13 13 14 14 14 14 15 15 15 15 # ... WGMMA_LAYOUT_8BIT = TiledLayout( Tiling(((64, 16), (16, 16), (8, 16), (4,))), warp_dims=(-7,), lane_dims=(-3, -2), vector_dim=-1, ) # This tiled layout is similar to the WGMMA layout, only the unit at which we # assign submatrices to warps grows from 8x8 to 8x16. The elements within each # submatrix are assigned to threads in the following way: # # 0 0 0 0 2 2 2 2 1 1 1 1 3 3 3 3 # 4 4 4 4 6 6 6 6 5 5 5 5 7 7 7 7 # ... # # Our vector length is twice the size of that of WGMMA_LAYOUT, which lets us use # 32-bit SMEM loads/stores when dealing with 8-bit values. The conversion # to the WGMMA layout only requires communication between with index differing # in their 2 bit (i.e. 0 and 1, 2 and 4), so the conversion to WGMMA_LAYOUT # only requires a single warp shuffle (plus permutes local to each thread). WGMMA_LAYOUT_UPCAST_2X = TiledLayout( Tiling(((64, 16), (16, 16), (8, 16), (8,), (4,))), warp_dims=(-8,), lane_dims=(-4, -2, -3), vector_dim=-1, ) # This layout should be used when upcasting 4-bit elements to 16-bit, for the # purpose of passing them into WGMMA later. The core matrices stored by a warp # are 8x32, because each of the 4 threads in a row holds 8 elements in a single # vector. Note that unlike WGMMA_LAYOUT_UPCAST_2X, we assign columns to each # group of 4 threads in order (as opposed to the swapping between 1 and 2, # 5 and 6, etc. that WGMMA_LAYOUT_UPCAST_2X does). WGMMA_LAYOUT_UPCAST_4X = TiledLayout( Tiling(((64, 32), (16, 32), (8, 32), (8,))), warp_dims=(-7,), lane_dims=(-3, -2), vector_dim=-1, ) # This tiled layout is similar to WGMMA_LAYOUT. There, each warp stores a 8x8 # submatrix in the following way (we only show the first 4 rows for brevity): # # 0 0 1 1 2 2 3 3 # 4 4 5 5 6 6 7 7 # 8 8 9 9 10 10 11 11 # 12 12 13 13 14 14 15 15 # ... # # This tiled layout stores the same 8x8 submatrix in the following way: # # 0 4 1 5 2 6 3 7 # 0 4 1 5 2 6 3 7 # 8 12 9 13 10 14 11 15 # 8 12 9 13 10 14 11 15 # ... # # You can see that we have taken 2x2 submatrices from the above layout and # transposed them. The assignment of lanes to elements is such that in both # layouts the same two lanes map to a single 2x2 submatrix, making the transpose # very cheap (one shuffle and permute suffices to change between those layouts). WGMMA_TRANSPOSED_LAYOUT = TiledLayout( Tiling(((64, 8), (16, 8), (8, 8), (2, 2), (2, 1))), warp_dims=(-10,), lane_dims=(-6, -3, -5), vector_dim=-2, ) # Like WGMMA_LAYOUT, only each warp holds a 32xN strip instead of 16xN. TCGEN05_LAYOUT = TiledLayout( Tiling(((128, 8), (32, 8), (8, 8), (2,))), warp_dims=(-7,), lane_dims=(-3, -2), vector_dim=-1, ) # Like WGMMA_TRANSPOSED_LAYOUT, only each warp holds a 32xN strip instead of 16xN. TCGEN05_TRANSPOSED_LAYOUT = TiledLayout( Tiling(((128, 8), (32, 8), (8, 8), (2, 2), (2, 1))), warp_dims=(-10,), lane_dims=(-6, -3, -5), vector_dim=-2, ) # TCGEN05_ROW_LAYOUT is to TCGEN05_LAYOUT as WGMMA_ROW_LAYOUT is to # WGMMA_LAYOUT. TCGEN05_ROW_LAYOUT = TiledLayout( Tiling(tiles=((128,), (32,), (8,), (1,))), warp_dims=(-4,), lane_dims=(-2, Replicated(times=4)), vector_dim=-1, ) # TCGEN05_COL_LAYOUT is to TCGEN05_LAYOUT as WGMMA_COL_LAYOUT is to # WGMMA_LAYOUT. TCGEN05_COL_LAYOUT = TiledLayout( Tiling(tiles=((8,), (2,))), warp_dims=(Replicated(times=4),), lane_dims=(Replicated(times=8), -2), vector_dim=-1, ) def tmem_native_layout(vector_length: int): """A layout resembling the logical organization of TMEM. The 128 rows in a tile are assigned to 128 lanes in the warpgroup. Useful when the result needs to be processed in registers and then stored back into TMEM. Usually shouldn't be used if the result is to be written back to SMEM, as there is no good way to store it without bank conflicts, but it still sometimes pays off. """ return TiledLayout( Tiling(((128, vector_length), (32, vector_length))), warp_dims=(-4,), lane_dims=(-2,), vector_dim=-1, ) # We use a vector_dim of 2, to be able to make sure that the vectors are always # a multiple of 32-bits, even when the data is 16-bits. TMEM_NATIVE_LAYOUT = tmem_native_layout(2) # A layout for the row indices used by TMA gather4/scatter4 instructions. # Index 00 01 02 03 04 05 06 07 08 09 10 11 12 13 14 15 16 17 ... # Warp <--- 0 ---> <--- 1 ---> <--- 2 ---> <--- 3 ---> <--- 0 -- TMA_INDICES_LAYOUT = TiledLayout( Tiling(((16,), (4,))), warp_dims=(-2,), lane_dims=(Replicated(32),), vector_dim=-1, ) # A replicated layout for TMA indices where every thread holds the same indices. TMA_INDICES_4_LAYOUT = TiledLayout( Tiling(((4,),)), warp_dims=(Replicated(4),), lane_dims=(Replicated(32),), vector_dim=-1, ) def can_relayout_wgmma_4x_to_wgmma_2x(bitwidth: int) -> bool: return bitwidth == 4 def can_relayout_wgmma_2x_to_wgmma(bitwidth: int) -> bool: return bitwidth in {4, 8, 16} def _int_pow(x: ir.Value, n: int) -> ir.Value: if n < 0: raise ValueError("Negative exponent not supported for integers") result = c(1, x.type) base = x while n > 0: if n % 2 == 1: result = arith.muli(result, base) base = arith.muli(base, base) n //= 2 return result # Any non-decreasing unary function `f` satisfies `f(max(a, b)) == # max(f(a), f(b))` and `f(min(a, b)) == min(f(a), f(b))`. Note that this # requires `f` to be defined everywhere: e.g. `log(max(-1, 1)) == 0` while # `max(log(-1), log(1))` is NaN, since our `max` propagates NaNs. _NON_DECREASING_DISTRIBUTES_OVER = ("max", "min") def _can_splat_unreduced(value: ir.Value, op: str) -> bool: # Each unreduced thread holds a partial result, so splatting `value` # across them would contribute it once per thread when the reduction is # completed, unless `op` is idempotent (`min`, `max`) or `value` is the # neutral element of the unreduced operation. if op in ("min", "max"): return True if not isinstance(value, ir.OpResult): return False const = value.owner.opview if not isinstance(const, arith.ConstantOp): return False if isinstance(value.type, (ir.FloatType, ir.IntegerType)): match op: case "add": return const.literal_value == 0 case "prod": return const.literal_value == 1 return False @jax.tree_util.register_pytree_node_class @dataclasses.dataclass(init=False, frozen=True, slots=True) class FragmentedArray: # An array of ir.Value, see checks in init for shapes. registers: np.ndarray = dataclasses.field(repr=False) layout: FragmentedLayout is_signed: bool | None def __init__( self, *, _registers: np.ndarray, _layout: FragmentedLayout, _is_signed: bool | None, ): """Initializes a fragmented array. This is a low-level API. Prefer using classmethods to construct fragmented arrays instead. """ # We need to use ``object.__setattr__`` here because of ``frozen=True``. object.__setattr__(self, "registers", _registers) object.__setattr__(self, "layout", _layout) object.__setattr__(self, "is_signed", _is_signed) if (_is_signed is not None) != isinstance(self.mlir_dtype, ir.IntegerType): raise TypeError( "is_signed must be non-None if and only if the MLIR type is an" f" integer type, got {_is_signed=} for {self.mlir_dtype}" ) match self.layout: # Registers are flat case WGStridedFragLayout(shape): [reg_size] = ir.VectorType(_registers.flat[0].type).shape if ( math.prod(shape) != math.prod(_registers.shape) * WARPGROUP_SIZE * reg_size ): raise ValueError( f"Invalid register array shape: math.prod({_registers.shape}) *" f" {WARPGROUP_SIZE} * {reg_size}, want: math.prod({shape})" ) # Just a single register case WGSplatFragLayout(): if _registers.size != 1: raise ValueError(f"Invalid register array shape: {_registers.shape}") case TiledLayout(): try: self.layout.shape_from_registers_shape(_registers.shape) except ValueError: raise ValueError( "Register array shape does not match the tiled layout" ) from None [vec_len] = self.registers.flat[0].type.shape assert vec_len == self.layout.vector_length case _: raise NotImplementedError @classmethod def load_strided( cls, ref: ir.Value, *, is_signed: bool | None = None, vec_size: int | None = None, ) -> FragmentedArray: index = ir.IndexType.get() i64 = ir.IntegerType.get_signless(64) if not isinstance(ref.type, ir.MemRefType): raise TypeError(ref.type) ref_ty = ir.MemRefType(ref.type) shape = tuple(ref_ty.shape) if vec_size is None: layout = WGStridedFragLayout.from_shaped_type(ref_ty) if layout is None: raise ValueError( f"{ref_ty} must have a number of elements that is a multiple of" f" {WARPGROUP_SIZE} (got {math.prod(shape)})" ) else: layout = WGStridedFragLayout(shape=shape, vec_size=vec_size) registers = np.empty(layout.registers_shape(shape), dtype=object) vec_ty = ir.VectorType.get((layout.vec_size,), ref_ty.element_type) # vector.load lowering doesn't work for element types with < 8 bits if (bitwidth := utils.bitwidth(ref_ty.element_type)) < 8: total_bits = layout.vec_size * bitwidth if total_bits % 8 != 0: raise NotImplementedError("Vector length should be a multiple of byte size") int_ty = ir.IntegerType.get_signless(bitwidth) int_vec_ty = ir.VectorType.get((layout.vec_size,), int_ty) for _get, update, transfer_ref, idx in cls.transfer_strided(ref, layout.vec_size): base_ptr = utils.memref_ptr(transfer_ref) strides, _ = transfer_ref.type.get_strides_and_offset() elem_flat_idx = utils.dyn_dot(idx, [arith.constant(index, s) for s in strides]) vec_flat_idx = arith.divui(elem_flat_idx, arith.constant(index, layout.vec_size)) vec_flat_idx = arith.index_cast(i64, vec_flat_idx) ptr = utils.getelementptr(base_ptr, [vec_flat_idx], int_vec_ty) loaded_val = llvm.load(int_vec_ty, ptr) if not isinstance(ref_ty.element_type, ir.IntegerType): loaded_val = vector.bitcast(vec_ty, loaded_val) update(registers, loaded_val) else: # We only keep this branch because it emits smaller IR. for _get, update, transfer_ref, idx in cls.transfer_strided(ref, layout.vec_size): update(registers, vector.load(vec_ty, transfer_ref, idx)) return cls(_registers=registers, _layout=layout, _is_signed=is_signed) @classmethod def splat( cls, value: ir.Value, shape: tuple[int, ...], layout: FragmentedLayout | None = None, *, is_signed: bool | None = None, ) -> FragmentedArray: layout = layout or WGSplatFragLayout(shape) if isinstance(layout, TiledLayout) and layout.has_unreduced_dims: assert layout.unreduced_operation is not None if not _can_splat_unreduced(value, layout.unreduced_operation): raise ValueError( "Cannot splat a non-neutral value into a layout with unreduced dims" ) match layout: case WGSplatFragLayout(): pass case WGStridedFragLayout() | TiledLayout(): value = vector.broadcast( layout.registers_element_type(value.type), value ) case _: raise NotImplementedError(layout) return cls( _registers=np.full(layout.registers_shape(shape), value, dtype=object), _layout=layout, _is_signed=is_signed, ) @staticmethod def broadcasted_iota( dtype: ir.Type, shape: tuple[int, ...], dimension: int, layout: FragmentedLayout | None = None, *, is_signed: bool | None = None, ) -> FragmentedArray: """Creates a broadcasted iota array along the specified dimension.""" if dimension >= len(shape): raise ValueError( "`dimension` must be smaller than the rank of the array." ) def cast(idx: ir.Value) -> ir.Value: if isinstance(dtype, ir.FloatType): i32 = ir.IntegerType.get_signless(32) return arith.uitofp(dtype, arith.index_cast(i32, idx)) return arith.index_cast(dtype, idx) return mgpu.FragmentedArray.splat( llvm.mlir_undef(dtype), shape, layout, is_signed=is_signed, ).foreach( lambda _, idx: cast(idx[dimension]), create_array=True, is_signed=is_signed, ) @property def shape(self) -> tuple[int, ...]: match self.layout: case WGStridedFragLayout(shape): return shape case WGSplatFragLayout(shape=shape): return shape case TiledLayout(): return self.layout.shape_from_registers_shape(self.registers.shape) case _: raise NotImplementedError @property def mlir_dtype(self) -> ir.Type: reg_ty = self.registers.flat[0].type match self.layout: case WGStridedFragLayout() | TiledLayout(): return ir.VectorType(reg_ty).element_type case WGSplatFragLayout(): return reg_ty case _: raise NotImplementedError @property def is_unreduced(self) -> bool: return ( isinstance(self.layout, TiledLayout) and self.layout.has_unreduced_dims ) def to_layout(self, new_layout: FragmentedLayout) -> FragmentedArray: """Converts the fragmented array to the given layout.""" if self.is_unreduced: raise NotImplementedError( "to_layout not supported on partially reduced FragmentedArray" ) i32 = ir.IntegerType.get_signless(32) c = lambda x: arith.constant(i32, x) if self.layout == new_layout: return self shape = self.shape bitwidth = utils.bitwidth(self.mlir_dtype) transpose_pairs = ( (WGMMA_LAYOUT, WGMMA_TRANSPOSED_LAYOUT), (TCGEN05_LAYOUT, TCGEN05_TRANSPOSED_LAYOUT), ) if bitwidth in {16, 32} and ( (self.layout, new_layout) in transpose_pairs or (new_layout, self.layout) in transpose_pairs ): is_even_row = arith.cmpi( arith.CmpIPredicate.eq, arith.remui(arith.divui(utils.thread_idx(), c(4)), c(2)), c(0), ) perm = arith.select(is_even_row, c(0x5410), c(0x3276)) tmp_new_regs = [] for reg in self.registers.flat: reg_ty = reg.type if bitwidth == 16: reg = utils.bitcast(reg, i32) reg_shfl = utils.shfl_bfly(reg, 4) new_reg = utils.prmt(reg, reg_shfl, perm) elif bitwidth == 32: i32_vec = ir.VectorType.get((1,), i32) regs = [ utils.bitcast(utils.vector_slice(reg, slice(i, i + 1)), i32) for i in range(2) ] reg_to_shfl = arith.select(is_even_row, regs[1], regs[0]) reg_shfl = utils.shfl_bfly(reg_to_shfl, 4) new_reg_low = arith.select(is_even_row, regs[0], reg_shfl) new_reg_high = arith.select(is_even_row, reg_shfl, regs[1]) new_reg_i32 = utils.vector_concat([ utils.bitcast(new_reg_low, i32_vec), utils.bitcast(new_reg_high, i32_vec), ]) new_reg = utils.bitcast(new_reg_i32, reg_ty) else: raise ValueError(f"Unsupported bitwidth: {bitwidth}") tmp_new_regs.append(utils.bitcast(new_reg, reg_ty)) new_regs = np.asarray( tmp_new_regs, dtype=object ).reshape(new_layout.registers_shape(shape)) return FragmentedArray( _registers=new_regs, _layout=new_layout, _is_signed=self.is_signed ) if ( isinstance(self.layout, TiledLayout) and isinstance(new_layout, TiledLayout) and self.layout == tmem_native_layout(self.layout.vector_length) and new_layout == tmem_native_layout(new_layout.vector_length) ): new_registers = np.empty(new_layout.registers_shape(shape), dtype=object) if self.layout.vector_length > new_layout.vector_length: ratio = self.layout.vector_length // new_layout.vector_length new_length = new_layout.vector_length for idx, reg in np.ndenumerate(self.registers): for i in range(ratio): new_reg = utils.vector_slice( reg, slice(i * new_length, (i + 1) * new_length) ) new_registers[(idx[0], idx[1] * ratio + i, *idx[2:])] = new_reg elif self.layout.vector_length < new_layout.vector_length: ratio = new_layout.vector_length // self.layout.vector_length for idx in np.ndindex(new_registers.shape): new_reg = utils.vector_concat([ cast(ir.Value, self.registers[idx[0], idx[1] * ratio + i, *idx[2:]]) for i in range(ratio) ]) new_registers[idx] = new_reg return FragmentedArray( _registers=new_registers, _layout=new_layout, _is_signed=self.is_signed, ) if self.layout == WGMMA_LAYOUT_ACC_32BIT and new_layout == WGMMA_LAYOUT: new_regs_shape = new_layout.registers_shape(shape) assert new_regs_shape[-1] == 1 assert self.registers.shape == (*new_regs_shape[:-1], 2, 1) new_regs = np.empty(new_regs_shape, dtype=object) for idx in np.ndindex(new_regs_shape[:-1]): new_regs[(*idx, 0)] = utils.vector_concat([ self.registers[*idx, i, 0] for i in range(2) ]) return FragmentedArray( _registers=new_regs, _layout=new_layout, _is_signed=self.is_signed, ) if self.layout == WGMMA_LAYOUT and new_layout == WGMMA_LAYOUT_ACC_32BIT: new_regs_shape = new_layout.registers_shape(shape) assert self.registers.shape[-1] == 1 assert new_regs_shape == (*self.registers.shape[:-1], 2, 1) new_regs = np.empty(new_regs_shape, dtype=object) for idx, reg in np.ndenumerate(self.registers): for i in range(2): new_regs[(*idx[:-1], i, 0)] = utils.vector_slice(reg, slice(i, i + 1)) return FragmentedArray( _registers=new_regs, _layout=new_layout, _is_signed=self.is_signed, ) dtype_bitwidth = utils.bitwidth(self.mlir_dtype) if (self.layout, new_layout) in ( (WGMMA_LAYOUT_UPCAST_2X, WGMMA_LAYOUT), (WGMMA_LAYOUT, WGMMA_LAYOUT_UPCAST_2X), ) and can_relayout_wgmma_2x_to_wgmma(dtype_bitwidth): assert isinstance(self.layout, TiledLayout) assert isinstance(new_layout, TiledLayout) assert shape[-1] % 16 == 0 # Should be implied by the layout new_registers = np.empty(new_layout.registers_shape(shape), dtype=object) is_even = arith.cmpi( arith.CmpIPredicate.eq, arith.remui(utils.thread_idx(), c(2)), c(0) ) col_dim = len(shape) - 1 registers = self.registers reg_bitwidth = self.layout.vector_length * dtype_bitwidth target_bitwidth = max( 32, WGMMA_LAYOUT_UPCAST_2X.vector_length * dtype_bitwidth ) pack_factor = target_bitwidth // reg_bitwidth if pack_factor > 1: if registers.shape[col_dim] % pack_factor: raise NotImplementedError( "This relayout implementation requires the number of column" f" tiles to be divisible by {pack_factor} (to pack them for" " efficiency)" ) # We pair up the consecutive column tiles, so each register is 32-bit # (or 64-bit for 16-bit elements). # If this layout originated from a WGMMA_LAYOUT_UPCAST_4X layout, # LLVM will realize that the paired up vectors actually came from the # same 32-bit register and it will become a no-op. col_minor_registers = np.moveaxis(registers, col_dim, -1) flat_registers = [ utils.vector_concat(group) for group in zip(*( col_minor_registers.flat[i::pack_factor] for i in range(pack_factor) )) ] registers = np.asarray(flat_registers, dtype=object).reshape( *col_minor_registers.shape[:-1], col_minor_registers.shape[-1] // pack_factor, ) registers = np.moveaxis(registers, -1, col_dim) out_vec_len = new_layout.vector_length for idx, reg in np.ndenumerate(registers): if dtype_bitwidth == 16: assert reg.type.shape == [4] # A single vector is 64-bits, but shuffles are only 32-bit wide. # We only shuffle the half that needs to go to other thread. low = utils.vector_slice(reg, slice(0, 2)) high = utils.vector_slice(reg, slice(2, 4)) to_exchange = arith.select(is_even, high, low) # Exchange values between even and odd threads. exchanged = utils.shfl_bfly(to_exchange, 1) low = arith.select(is_even, low, exchanged) high = arith.select(is_even, exchanged, high) out_regs = ( (low, high) if new_layout == WGMMA_LAYOUT else (utils.vector_concat((low, high)),) ) elif dtype_bitwidth == 8: assert reg.type.shape == [4] # The vector is 32-bits, so we just shuffle the whole thing and # use prmt to blend it with the local register. exchanged = utils.shfl_bfly(reg, 1) # Consider lanes 0 and 1, because the situation is symmetric for # each pair. If we feed reg[lane] and exchanged[lane] (which is # really the same as reg of the other lane) to prmt, we can index # the elements of the result using the following indices: # reg[0]: 0 1 2 3 reg[1]: 8 9 10 11 # prmt[0]: 0 1 2 3 4 5 6 7 # prmt[1]: 4 5 6 7 0 1 2 3 # The expected outputs and their respective permutations are: # out[0]: 0 1 8 9 out[1]: 2 3 10 11 # prmt[0]: 0 1 4 5 prmt[1]: 6 7 2 3 # Note that the patterns still need to be flipped, since we listed # bytes with LSB on the left, which is the opposite of how the # numeric constants are spelled in Python (LSB on the right). perm = arith.select(is_even, c(0x5410), c(0x3276)) blend = utils.prmt(reg, exchanged, perm) out_regs = ( tuple( utils.vector_slice(blend, slice(i * 2, i * 2 + 2)) for i in range(2) ) if new_layout == WGMMA_LAYOUT else (blend,) ) else: assert dtype_bitwidth == 4 assert reg.type.shape == [8] # We paired up the registers above. exchanged = utils.shfl_bfly(reg, 1) # See comment above for a more complete explanation. # reg[0]: 0 1 2 3 16 17 18 19 reg[1]: 8 9 10 11 24 25 26 27 # prmt[0]: -0- -1- --2-- --3-- -4- --5-- --6-- --7-- # prmt[1]: -4- -5- --6-- --7-- -0- --1-- --2-- --3-- # The expected outputs and their respective permutations are: # out[0]: 0 1 8 9 16 17 24 25 out[1]: 2 3 10 11 18 19 26 27 # prmt[0]: -0- -4- --2-- --6-- prmt[1]: -5- --1-- --7-- --3-- perm = arith.select(is_even, c(0x6240), c(0x3715)) blend = utils.prmt(reg, exchanged, perm) out_regs = tuple( utils.vector_slice( blend, slice(i * out_vec_len, (i + 1) * out_vec_len) ) for i in range(8 // out_vec_len) ) # WGMMA_LAYOUT_UPCAST_2X has one more tiled dimension of size 1. rest_idx = ( idx[col_dim + 1 : -1] if new_layout == WGMMA_LAYOUT else (*idx[col_dim + 1 :], 0) ) for i, out_reg in enumerate(out_regs): new_registers[( *idx[:col_dim], idx[col_dim] * len(out_regs) + i, *rest_idx, )] = out_reg assert all(r is not None for r in new_registers) return FragmentedArray( _registers=new_registers, _layout=new_layout, _is_signed=self.is_signed, ) if ( self.layout == WGMMA_LAYOUT_UPCAST_4X and new_layout == WGMMA_LAYOUT_UPCAST_2X and can_relayout_wgmma_4x_to_wgmma_2x(dtype_bitwidth) ): assert shape[0] % 64 == 0 # Should be implied by the layout assert shape[1] % 32 == 0 # Should be implied by the layout new_registers = np.empty(new_layout.registers_shape(shape), dtype=object) i32 = ir.IntegerType.get_signless(32) c = lambda x: arith.constant(i32, x) is_01 = arith.cmpi( arith.CmpIPredicate.ult, arith.remui(utils.thread_idx(), c(4)), c(2) ) for idx, reg in np.ndenumerate(self.registers): assert ir.VectorType(reg.type).shape == [8] # The vector is 32-bits, so we just shuffle the whole thing and # use prmt to blend it with the local register. exchanged = utils.shfl_bfly(reg, 2) # See comments above for conventions. Here we exchange data between # threads with lane index related by flipping 2nd bit (e.g. 0 and 2). # reg[0]: 0 1 2 3 4 5 6 7 reg[2]: 16 17 18 19 20 21 22 23 # prmt[0]: -0- -1- -2- -3- --4-- --5-- --6-- --7-- # prmt[1]: -4- -5- -6- -7- --0-- --1-- --2-- --3-- # The expected outputs and their respective permutations are: # out[0]: 0 1 2 3 16 17 18 19 out[2]: 4 5 6 7 20 21 22 23 # prmt[0]: -0- -1- --4-- --5-- prmt[2]: -6- -7- --2-- --3-- perm = arith.select(is_01, c(0x5410), c(0x3276)) blend = utils.prmt(reg, exchanged, perm) for i in range(2): reg = utils.vector_slice(blend, slice(i * 4, i * 4 + 4)) new_registers[(idx[0], idx[1] * 2 + i, *idx[2:-1])] = reg assert all(r is not None for r in new_registers) return FragmentedArray( _registers=new_registers, _layout=new_layout, _is_signed=self.is_signed, ) if self.layout == WGMMA_LAYOUT_UPCAST_4X and new_layout == WGMMA_LAYOUT: return self.to_layout(WGMMA_LAYOUT_UPCAST_2X).to_layout(new_layout) if not isinstance(self.layout, WGSplatFragLayout): raise NotImplementedError( f"Cannot convert from {self.layout} to {new_layout}" ) return type(self).splat( self.registers.item(), self.shape, new_layout, is_signed=self.is_signed ) def _pointwise( self, op, *other, output_is_signed: bool | None = None, restrict_bitwidth: bool = True, distributes_over: Sequence[str] = (), ) -> FragmentedArray: """Applies `op` elementwise to `self` and `other`. Args: op: The operation to apply to each register. *other: The other operands. Scalars and splat operands (including `self`) are broadcast to the layout of the non-splat operand. output_is_signed: The signedness of the result, if it is an integer. restrict_bitwidth: Whether to reject sub-16-bit operands (except 1-bit). distributes_over: The reductions (binary associative operations, e.g. `"add"`) that `op` distributes over, i.e. the reductions `r` such that for all `a`, `b`, `c`, `d`: `op(r(a, b), r(c, d)) == r(op(a, c), op(b, d))` (and likewise for other arities). `op` can only be applied to an unreduced array if it distributes over the pending reduction, since it is then applied to each partial result before the reduction is completed. """ if restrict_bitwidth: if (bitwidth := utils.bitwidth(self.mlir_dtype)) <= 8 and bitwidth != 1: raise NotImplementedError( f"Pointwise operations on {bitwidth}-bit types are unsupported" " (except bitwise operations). Upcast to a 16- or 32-bit type" " before performing the operation." ) if self.is_unreduced: assert isinstance(self.layout, TiledLayout) if self.layout.unreduced_operation not in distributes_over: raise ValueError( "Pointwise operations on partially reduced FragmentedArrays must" " distribute over the pending reduction" f" {self.layout.unreduced_operation!r}" ) # If our layout is a splat, then we should either dispatch to a non-splat # layout, or broadcast ourselves to the output shape first. if isinstance(self.layout, WGSplatFragLayout): output_shape = self.shape for i, o in enumerate(other): if not isinstance(o, FragmentedArray): continue elif not isinstance(o.layout, WGSplatFragLayout): return o._pointwise( lambda o, this, *args: op(this, *args[:i], o, *args[i:]), self, *other[:i], *other[i + 1 :], output_is_signed=output_is_signed, distributes_over=distributes_over, ) else: output_shape = np.broadcast_shapes(output_shape, o.shape) # If we get here then we haven't found any non-splat layout. if self.shape != output_shape: return self.broadcast(output_shape)._pointwise( op, *other, output_is_signed=output_is_signed, distributes_over=distributes_over, ) other_arrs = [] for o in other: if not isinstance(o, FragmentedArray): if isinstance(o, (float, int)): o = utils.c(o, self.mlir_dtype) elif not isinstance(o, ir.Value): raise NotImplementedError(o) o = FragmentedArray.splat( o, shape=self.shape, layout=self.layout, is_signed=self.is_signed ) if isinstance(o.layout, WGSplatFragLayout): if not o.layout.can_broadcast_to(self.shape): raise ValueError( f"Cannot broadcast shape {self.shape} to layout {o.layout}") o = FragmentedArray.splat( o.registers.flat[0], shape=self.shape, layout=self.layout, is_signed=o.is_signed, ) else: if self.layout != o.layout: raise ValueError("Incompatible FragmentedArray layouts") if self.registers.shape != o.registers.shape: raise ValueError("Incompatible FragmentedArray shapes") other_arrs.append(o) new_regs = np.empty_like(self.registers) for idx, reg in np.ndenumerate(self.registers): new_regs[idx] = op(reg, *(o.registers[idx] for o in other_arrs)) reg_ty = new_regs.flat[0].type if isinstance(reg_ty, ir.VectorType): reg_ty = ir.VectorType(reg_ty).element_type if output_is_signed is None and isinstance(reg_ty, ir.IntegerType): output_is_signed = self.is_signed return FragmentedArray( _registers=new_regs, _layout=self.layout, _is_signed=output_is_signed ) def __pos__(self): return self def __neg__(self): if isinstance(self.mlir_dtype, ir.FloatType): # -(a + b) = (-a) + (-b) return self._pointwise(arith.negf, distributes_over=("add",)) elif isinstance(self.mlir_dtype, ir.IntegerType): return 0 - self else: return NotImplemented def __add__(self, other): if isinstance(self.mlir_dtype, ir.FloatType): return self._pointwise(addf, other, distributes_over=("add",)) elif isinstance(self.mlir_dtype, ir.IntegerType): return self._pointwise(arith.addi, other, distributes_over=("add",)) else: return NotImplemented def __radd__(self, other): return self + other def __mul__(self, other): if isinstance(self.mlir_dtype, ir.FloatType): return self._pointwise(mulf, other, distributes_over=("prod",)) elif isinstance(self.mlir_dtype, ir.IntegerType): return self._pointwise(arith.muli, other, distributes_over=("prod",)) else: return NotImplemented def __rmul__(self, other): return self * other def __sub__(self, other): # (a + b) - (c + d) = (a - c) + (b - d) if isinstance(self.mlir_dtype, ir.FloatType): return self._pointwise(subf, other, distributes_over=("add",)) elif isinstance(self.mlir_dtype, ir.IntegerType): return self._pointwise(arith.subi, other, distributes_over=("add",)) else: return NotImplemented def __rsub__(self, other): if isinstance(self.mlir_dtype, ir.FloatType): return self._pointwise( lambda s, o: subf(o, s), other, distributes_over=("add",) ) elif isinstance(self.mlir_dtype, ir.IntegerType): return self._pointwise( lambda s, o: arith.subi(o, s), other, distributes_over=("add",) ) else: return NotImplemented def _e8m0_reciprocal(self): def e8m0_inv(x): if not isinstance(x.type, ir.VectorType): raise NotImplementedError(x.type) [vec_len] = ir.VectorType(x.type).shape i8 = ir.IntegerType.get_signless(8) i8_vec = ir.VectorType.get((vec_len,), i8) c254 = vector.broadcast(i8_vec, arith.constant(i8, 254)) return utils.bitcast(arith.subi(c254, utils.bitcast(x, i8_vec)), x.type) return self._pointwise(e8m0_inv, restrict_bitwidth=False) @staticmethod def _is_e8m0_constant_one(val): if isinstance(val, (int, float)): return val == 1 if not isinstance(val, FragmentedArray): return False if not isinstance(val.layout, WGSplatFragLayout): return False if not isinstance(val.mlir_dtype, ir.Float8E8M0FNUType): return False [reg] = val.registers.flat return ( isinstance(reg, ir.OpResult) and isinstance(reg.owner.opview, arith.ConstantOp) and float(ir.FloatAttr(reg.owner.opview.value)) == 1.0 ) def __truediv__(self, other): if not isinstance(self.mlir_dtype, ir.FloatType): return NotImplemented if ( isinstance(other, FragmentedArray) and isinstance(other.mlir_dtype, ir.Float8E8M0FNUType) and FragmentedArray._is_e8m0_constant_one(self) ): return other._e8m0_reciprocal() # (a * b) / (c * d) = (a / c) * (b / d) return self._pointwise(arith.divf, other, distributes_over=("prod",)) def __rtruediv__(self, other): if not isinstance(self.mlir_dtype, ir.FloatType): return NotImplemented if isinstance( self.mlir_dtype, ir.Float8E8M0FNUType ) and FragmentedArray._is_e8m0_constant_one(other): return self._e8m0_reciprocal() return self._pointwise( lambda s, o: arith.divf(o, s), other, distributes_over=("prod",) ) def __floordiv__(self, other): if isinstance(self.mlir_dtype, ir.FloatType): return self._pointwise( lambda s, o: mlir_math.floor(arith.divf(s, o)), other ) elif isinstance(self.mlir_dtype, ir.IntegerType): if self.is_signed: return self._pointwise(arith.floordivsi, other) else: return self._pointwise(arith.divui, other) else: return NotImplemented def trunc_div(self, other): if not isinstance(self.mlir_dtype, ir.IntegerType): raise ValueError( f"trunc_div supports only integer types, got {self.mlir_dtype}" ) if self.is_signed: return self._pointwise(arith.divsi, other) else: return self._pointwise(arith.divui, other) def __rfloordiv__(self, other): if isinstance(self.mlir_dtype, ir.FloatType): return self._pointwise( lambda s, o: mlir_math.floor(arith.divf(o, s)), other ) elif isinstance(self.mlir_dtype, ir.IntegerType): if self.is_signed: return self._pointwise(lambda s, o: arith.floordivsi(o, s), other) else: return self._pointwise(lambda s, o: arith.divui(o, s), other) else: return NotImplemented def __mod__(self, other): if not isinstance(self.mlir_dtype, ir.IntegerType): return NotImplemented if self.is_signed: return self._pointwise(arith.remsi, other) else: return self._pointwise(arith.remui, other) def __rmod__(self, other): if not isinstance(self.mlir_dtype, ir.IntegerType): return NotImplemented if self.is_signed: return self._pointwise(lambda s, o: arith.remsi(o, s), other) else: return self._pointwise(lambda s, o: arith.remui(o, s), other) def __pow__(self, other): if isinstance(self.mlir_dtype, ir.IntegerType): return self._pointwise(lambda x: _int_pow(x, other)) if not isinstance(self.mlir_dtype, ir.FloatType): return NotImplemented return self._pointwise(mlir_math.powf, other) def __rpow__(self, other): if not isinstance(self.mlir_dtype, ir.FloatType): return NotImplemented return self._pointwise(lambda s, o: mlir_math.powf(o, s), other) def __invert__(self): if not isinstance(self.mlir_dtype, ir.IntegerType): return NotImplemented return self ^ ~0 def __or__(self, other): if not isinstance(self.mlir_dtype, ir.IntegerType): return NotImplemented return self._pointwise(arith.ori, other, restrict_bitwidth=False) def __ror__(self, other): return self | other def __and__(self, other): if not isinstance(self.mlir_dtype, ir.IntegerType): return NotImplemented return self._pointwise(arith.andi, other, restrict_bitwidth=False) def __rand__(self, other): return self & other def __xor__(self, other): if not isinstance(self.mlir_dtype, ir.IntegerType): return NotImplemented return self._pointwise(arith.xori, other, restrict_bitwidth=False) def __rxor__(self, other): return self ^ other def __lshift__(self, other): if not isinstance(self.mlir_dtype, ir.IntegerType): return NotImplemented return self._pointwise(arith.shli, other) def __rshift__(self, other): if not isinstance(self.mlir_dtype, ir.IntegerType): return NotImplemented return self._pointwise( arith.shrsi if self.is_signed else arith.shrui, other, ) def __eq__(self, other): return self._compare( other, f_pred=arith.CmpFPredicate.OEQ, si_pred=arith.CmpIPredicate.eq, ui_pred=arith.CmpIPredicate.eq, restrict_bitwidth=False, ) def __ne__(self, other): return self._compare( other, f_pred=arith.CmpFPredicate.UNE, si_pred=arith.CmpIPredicate.ne, ui_pred=arith.CmpIPredicate.ne, restrict_bitwidth=False, ) def __lt__(self, other): return self._compare( other, f_pred=arith.CmpFPredicate.OLT, si_pred=arith.CmpIPredicate.slt, ui_pred=arith.CmpIPredicate.ult, ) def __le__(self, other): return self._compare( other, f_pred=arith.CmpFPredicate.OLE, si_pred=arith.CmpIPredicate.sle, ui_pred=arith.CmpIPredicate.ule, ) def __gt__(self, other): return self._compare( other, f_pred=arith.CmpFPredicate.OGT, si_pred=arith.CmpIPredicate.sgt, ui_pred=arith.CmpIPredicate.ugt, ) def __ge__(self, other): return self._compare( other, f_pred=arith.CmpFPredicate.OGE, si_pred=arith.CmpIPredicate.sge, ui_pred=arith.CmpIPredicate.uge, ) def _compare( self, other, *, f_pred, si_pred, ui_pred, restrict_bitwidth=True ): if isinstance(self.mlir_dtype, ir.FloatType): pred = functools.partial(arith.cmpf, f_pred) elif isinstance(self.mlir_dtype, ir.IntegerType): if self.is_signed: pred = functools.partial(arith.cmpi, si_pred) else: pred = functools.partial(arith.cmpi, ui_pred) else: return NotImplemented return self._pointwise( pred, other, output_is_signed=False, restrict_bitwidth=restrict_bitwidth ) def max(self, other) -> FragmentedArray: if isinstance(self.mlir_dtype, ir.FloatType): maximumf = arith.maximumf if isinstance(self.mlir_dtype, ir.F32Type): maximumf = self._lift_fast_instr("max.NaN.f32") elif isinstance(self.mlir_dtype, ir.F16Type): maximumf = self._lift_fast_packed_instr("max.NaN.f16x2", "max.NaN.f16") elif isinstance(self.mlir_dtype, ir.BF16Type): maximumf = self._lift_fast_packed_instr("max.NaN.bf16x2", "max.NaN.bf16") return self._pointwise(maximumf, other, distributes_over=("max",)) elif isinstance(self.mlir_dtype, ir.IntegerType): width = utils.bitwidth(self.mlir_dtype) if width == 16: sign = "s" if self.is_signed else "u" instr = self._lift_fast_packed_instr(f"max.{sign}16x2", f"max.{sign}16") return self._pointwise(instr, other, distributes_over=("max",)) return self._pointwise( arith.maxsi if self.is_signed else arith.maxui, other, distributes_over=("max",), ) else: raise NotImplementedError def min(self, other) -> FragmentedArray: if isinstance(self.mlir_dtype, ir.FloatType): minimumf = arith.minimumf if isinstance(self.mlir_dtype, ir.F32Type): minimumf = self._lift_fast_instr("min.NaN.f32") elif isinstance(self.mlir_dtype, ir.F16Type): minimumf = self._lift_fast_packed_instr("min.NaN.f16x2", "min.NaN.f16") elif isinstance(self.mlir_dtype, ir.BF16Type): minimumf = self._lift_fast_packed_instr("min.NaN.bf16x2", "min.NaN.bf16") return self._pointwise(minimumf, other, distributes_over=("min",)) elif isinstance(self.mlir_dtype, ir.IntegerType): width = utils.bitwidth(self.mlir_dtype) if width == 16: sign = "s" if self.is_signed else "u" instr = self._lift_fast_packed_instr(f"min.{sign}16x2", f"min.{sign}16") return self._pointwise(instr, other, distributes_over=("min",)) return self._pointwise( arith.minsi if self.is_signed else arith.minui, other, distributes_over=("min",), ) else: raise NotImplementedError def copysign(self, other: FragmentedArray) -> FragmentedArray: if not isinstance(self.mlir_dtype, ir.FloatType): raise NotImplementedError return self._pointwise(mlir_math.copysign, other) def exp(self, *, approx: bool = False) -> FragmentedArray: if not isinstance(self.mlir_dtype, ir.FloatType): raise NotImplementedError if approx: dtype = self.mlir_dtype log2e = arith.constant(dtype, ir.FloatAttr.get(dtype, 1.4426950408889634)) # Multiplying by a positive constant is non-decreasing. return self._pointwise( mulf, log2e, distributes_over=_NON_DECREASING_DISTRIBUTES_OVER ).exp2(approx=True) return self._pointwise( mlir_math.exp, distributes_over=_NON_DECREASING_DISTRIBUTES_OVER ) def exp2(self, *, approx: bool = False) -> FragmentedArray: if not isinstance(self.mlir_dtype, ir.FloatType): raise NotImplementedError exp2f = mlir_math.exp2 if approx: if not isinstance(self.mlir_dtype, ir.F32Type): raise NotImplementedError(self.mlir_dtype) exp2f = self._lift_fast_instr("ex2.approx.ftz.f32") return self._pointwise( exp2f, distributes_over=_NON_DECREASING_DISTRIBUTES_OVER ) def log(self, *, approx: bool = False) -> FragmentedArray: if not isinstance(self.mlir_dtype, ir.FloatType): raise NotImplementedError if approx: dtype = self.mlir_dtype ln2 = arith.constant(dtype, ir.FloatAttr.get(dtype, 0.6931471805599453)) return self.log2(approx=True) * ln2 return self._pointwise(mlir_math.log) def log2(self, *, approx: bool = False) -> FragmentedArray: if not isinstance(self.mlir_dtype, ir.FloatType): raise NotImplementedError(self.mlir_dtype) if approx: if not isinstance(self.mlir_dtype, ir.F32Type): raise NotImplementedError(self.mlir_dtype) return self._pointwise(self._lift_fast_instr("lg2.approx.ftz.f32")) return self._pointwise(mlir_math.log2) def sin(self, *, approx: bool = False) -> FragmentedArray: if not isinstance(self.mlir_dtype, ir.FloatType): raise NotImplementedError if approx and self.mlir_dtype != ir.F32Type.get(): raise NotImplementedError return self._pointwise( self._lift_fast_instr("sin.approx.f32") if approx else mlir_math.sin ) def cos(self, *, approx: bool = False) -> FragmentedArray: if not isinstance(self.mlir_dtype, ir.FloatType): raise NotImplementedError if approx and self.mlir_dtype != ir.F32Type.get(): raise NotImplementedError return self._pointwise( self._lift_fast_instr("cos.approx.f32") if approx else mlir_math.cos ) def tanh(self, *, approx: bool = False) -> FragmentedArray: if not isinstance(self.mlir_dtype, ir.FloatType): raise NotImplementedError tanhf = mlir_math.tanh if approx: if isinstance(self.mlir_dtype, ir.F32Type): tanhf = self._lift_fast_instr("tanh.approx.f32") elif isinstance(self.mlir_dtype, ir.F16Type): tanhf = self._lift_fast_packed_instr( "tanh.approx.f16x2", "tanh.approx.f16" ) elif isinstance(self.mlir_dtype, ir.BF16Type): tanhf = self._lift_fast_packed_instr( "tanh.approx.bf16x2", "tanh.approx.bf16" ) else: raise NotImplementedError(self.mlir_dtype) return self._pointwise( tanhf, distributes_over=_NON_DECREASING_DISTRIBUTES_OVER ) def rsqrt(self, *, approx: bool = False) -> FragmentedArray: if not isinstance(self.mlir_dtype, ir.FloatType): raise NotImplementedError if approx and self.mlir_dtype != ir.F32Type.get(): raise NotImplementedError return self._pointwise( self._lift_fast_instr("rsqrt.approx.f32") if approx else mlir_math.rsqrt ) def sqrt(self, *, approx: bool = False) -> FragmentedArray: if not isinstance(self.mlir_dtype, ir.FloatType): raise NotImplementedError if approx and self.mlir_dtype != ir.F32Type.get(): raise NotImplementedError return self._pointwise( self._lift_fast_instr("sqrt.approx.f32") if approx else mlir_math.sqrt ) def abs(self) -> FragmentedArray: if isinstance(self.mlir_dtype, ir.FloatType): absf = mlir_math.absf if isinstance(self.mlir_dtype, ir.F32Type): absf = self._lift_fast_instr("abs.f32") elif isinstance(self.mlir_dtype, ir.F16Type): absf = self._lift_fast_packed_instr("abs.f16x2", "abs.f16") elif isinstance(self.mlir_dtype, ir.BF16Type): absf = self._lift_fast_packed_instr("abs.bf16x2", "abs.bf16") elif isinstance(self.mlir_dtype, (ir.Float8E4M3FNType, ir.Float8E5M2Type, ir.Float4E2M1FNType)): bitwidth = utils.bitwidth(self.mlir_dtype) itype = ir.IntegerType.get_signless(bitwidth) i_arr = self.bitcast(itype, output_is_signed=False) mask = (1 << (bitwidth - 1)) - 1 res_i = i_arr & mask return res_i.bitcast(self.mlir_dtype) return self._pointwise(absf) elif isinstance(self.mlir_dtype, ir.IntegerType): if self.is_signed: return self._pointwise(mlir_math.absi) else: return self else: raise NotImplementedError def round(self) -> FragmentedArray: """Same as `lax.round(..., AWAY_FROM_ZERO)`.""" if not isinstance(self.mlir_dtype, ir.FloatType): raise NotImplementedError return self._pointwise( mlir_math.round, distributes_over=_NON_DECREASING_DISTRIBUTES_OVER ) def round_even(self) -> FragmentedArray: """Same as `lax.round(..., TO_NEAREST_EVEN)`.""" if not isinstance(self.mlir_dtype, ir.FloatType): raise NotImplementedError return self._pointwise( mlir_math.roundeven, distributes_over=_NON_DECREASING_DISTRIBUTES_OVER ) def erf(self) -> FragmentedArray: if not isinstance(self.mlir_dtype, ir.FloatType): raise NotImplementedError(self.mlir_dtype) return self._pointwise( mlir_math.erf, distributes_over=_NON_DECREASING_DISTRIBUTES_OVER ) def atan2(self, other: FragmentedArray) -> FragmentedArray: if not isinstance(self.mlir_dtype, ir.FloatType): raise NotImplementedError(self.mlir_dtype) return self._pointwise(mlir_math.atan2, other) @staticmethod def _lift_fast_instr( instr: str | Callable[[ir.Value], ir.Value], ) -> Callable[[ir.Value, ir.Value], ir.Value]: def fast_instr(*args): f32 = ir.F32Type.get() arg_ty = args[0].type assert all(a.type == arg_ty for a in args) if arg_ty == f32: if isinstance(instr, str): args_ptx = ", ".join(f"${i}" for i in range(len(args) + 1)) return llvm.inline_asm( f32, args, f"{instr} {args_ptx};", "=f" + ",f" * len(args) ) else: return instr(*args) elif isinstance(arg_ty, ir.VectorType): result = llvm.mlir_undef(arg_ty) [vec_len] = ir.VectorType(arg_ty).shape for i in range(vec_len): vs = [ vector.extract( a, dynamic_position=[], static_position=ir.DenseI64ArrayAttr.get([i]), ) for a in args ] vr = fast_instr(*vs) result = vector.insert( vr, result, dynamic_position=[], static_position=ir.DenseI64ArrayAttr.get([i]), ) return result else: raise NotImplementedError(arg_ty) return fast_instr @staticmethod def _lift_fast_packed_instr( packed_instr: str, single_instr: str, ) -> Callable[[ir.Value, ir.Value], ir.Value]: def fast_instr(*args): arg_ty = original_arg_ty = args[0].type assert all(a.type == arg_ty for a in args) if not isinstance(arg_ty, ir.VectorType): args = [vector.broadcast(ir.VectorType.get((1,), arg_ty), a) for a in args] arg_ty = ir.VectorType(args[0].type) [vec_len] = arg_ty.shape vec_bitwidth = vec_len * utils.bitwidth(arg_ty.element_type) if vec_len == 1 or vec_bitwidth == 32: assert vec_bitwidth.bit_count() == 1 if vec_bitwidth == 32: cstr = "r" elif vec_bitwidth == 16: cstr = "h" else: raise NotImplementedError(vec_bitwidth) int_ty = ir.IntegerType.get_signless(vec_bitwidth) args_ptx = ", ".join(f"${i}" for i in range(len(args) + 1)) args_int = [utils.bitcast(a, int_ty) for a in args] result_int = llvm.inline_asm( int_ty, args_int, f"{single_instr if vec_len == 1 else packed_instr} {args_ptx};", f"={cstr}" + f",{cstr}" * len(args) ) assert isinstance(result_int, ir.Value) return utils.bitcast(result_int, original_arg_ty) else: assert vec_bitwidth > 32 slice_len = 32 // utils.bitwidth(arg_ty.element_type) offset = 0 slices = [] while offset < vec_len: slice_end = min(offset + slice_len, vec_len) args_slice = [utils.vector_slice(a, slice(offset, slice_end)) for a in args] slices.append(fast_instr(*args_slice)) offset = slice_end return utils.vector_concat(slices) return fast_instr def bitcast( self, elt: ir.Type, *, output_is_signed: bool | None = None ) -> FragmentedArray: if (output_is_signed is not None) != isinstance(elt, ir.IntegerType): raise TypeError( "output_is_signed must be non-None if and only if the MLIR type is an" f" integer type, got {output_is_signed=} for {elt}" ) if elt == self.mlir_dtype: return self if utils.bitwidth(elt) != utils.bitwidth(self.mlir_dtype): raise ValueError("Only bitcast between types of the same bitwidth supported") reg_type = self.registers.flat[0].type if isinstance(reg_type, ir.VectorType): reg_shape = ir.VectorType(reg_type).shape ty = ir.VectorType.get(reg_shape, elt) else: ty = elt return self._pointwise( lambda x: arith.bitcast(ty, x), output_is_signed=output_is_signed, restrict_bitwidth=False ) def __getitem__(self, idx) -> FragmentedArray: base_idx, slice_shape, is_squeezed = utils.parse_indices(idx, self.shape) if isinstance(self.layout, WGSplatFragLayout): shape = tuple(d for d, s in zip(slice_shape, is_squeezed) if not s) return self.splat(self.registers.item(), shape, is_signed=self.is_signed) if not isinstance(self.layout, TiledLayout): raise NotImplementedError("Only arrays with tiled layouts can be sliced") if any(isinstance(idx, ir.Value) for idx in base_idx): raise ValueError("Only slicing with static indices allowed") base_idx = cast(tuple[int, ...], base_idx) base_tile_shape = self.layout.base_tile_shape untiled_rank = len(self.shape) - len(base_tile_shape) if any(is_squeezed[untiled_rank:]): raise NotImplementedError( "Integer indexing not implemented for tiled dimensions (only slicing" " allowed)" ) if untiled_rank: base_tile_shape = (1,) * untiled_rank + base_tile_shape if any(b % t for b, t in zip(base_idx, base_tile_shape, strict=True)): raise ValueError( "Base indices of array slices must be aligned to the beginning of a" f" tile. The array uses a tiling of {base_tile_shape}, but your base" f" indices are {base_idx}. Consider using a different array layout." ) if any(l % t for l, t in zip(slice_shape, base_tile_shape, strict=True)): raise ValueError( "The slice shape must be a multiple of the tile shape. The array" f" uses a tiling of {base_tile_shape}, but your slice shape is" f" {slice_shape}. Consider using a different array layout." ) register_slices = tuple( b if sq else slice(b // t, (b + l) // t) for b, l, t, sq in zip( base_idx, slice_shape, base_tile_shape, is_squeezed, strict=True ) ) new_regs = self.registers[register_slices] return FragmentedArray( _registers=new_regs, _layout=self.layout, _is_signed=self.is_signed ) def __setitem__(self, idx: object, value: FragmentedArray) -> None: if not isinstance(value, FragmentedArray): raise ValueError(f"Expected a FragmentedArray, got: {value}") if not isinstance(self.layout, TiledLayout): raise NotImplementedError("Only arrays with tiled layouts can be sliced") base_idx, slice_shape, is_squeezed = utils.parse_indices(idx, self.shape) if any(isinstance(idx, ir.Value) for idx in base_idx): raise ValueError("Only slicing with static indices allowed") base_idx = cast(tuple[int, ...], base_idx) if any(is_squeezed): raise NotImplementedError("Integer indexing not implemented (only slicing allowed)") if value.shape != tuple(slice_shape): raise ValueError( f"Slice has shape {tuple(slice_shape)}, but assigned array has shape" f" {value.shape}" ) if value.mlir_dtype != self.mlir_dtype: raise ValueError( f"Array has dtype {value.mlir_dtype}, but assigned array has dtype" f" {self.mlir_dtype}" ) if value.layout != self.layout: raise ValueError( f"Array has layout {value.layout}, but assigned array has layout" f" {self.layout}" ) base_tile_shape = self.layout.base_tile_shape if len(base_tile_shape) != len(self.shape): raise NotImplementedError("Tiling has different rank than array") if any( b % t or l % t for b, l, t in zip(base_idx, slice_shape, base_tile_shape, strict=True) ): raise NotImplementedError("Only tile aligned slicing supported") register_slices = tuple( slice(b // t, (b + l) // t) for b, l, t in zip(base_idx, slice_shape, base_tile_shape, strict=True) ) assert self.registers[register_slices].shape == value.registers.shape self.registers[register_slices] = value.registers def copy(self) -> FragmentedArray: return FragmentedArray( _registers=np.copy(self.registers), _layout=self.layout, _is_signed=self.is_signed, ) # TODO(apaszke): Support JAX dtypes here as well? def astype( self, new_dtype: ir.Type, *, is_signed: bool | None = None, rounding: Rounding | None = None, ) -> FragmentedArray: if self.is_unreduced: raise NotImplementedError i4 = ir.IntegerType.get_signless(4) i8 = ir.IntegerType.get_signless(8) i16 = ir.IntegerType.get_signless(16) i32 = ir.IntegerType.get_signless(32) f32 = ir.F32Type.get() f16 = ir.F16Type.get() bf16 = ir.BF16Type.get() f8e4m3fn = ir.Float8E4M3FNType.get() f8e5m2 = ir.Float8E5M2Type.get() f8e8m0fnu = ir.Float8E8M0FNUType.get() f4e2m1fn = ir.Float4E2M1FNType.get() cur_dtype = self.mlir_dtype if cur_dtype == new_dtype: if self.is_signed == is_signed: return self return FragmentedArray( _registers=self.registers, _layout=self.layout, _is_signed=is_signed ) ptx_isa_version = 0 if ( min(utils.bitwidth(cur_dtype), utils.bitwidth(new_dtype)) <= 8 and mgpu_lib is not None ): ptx_isa_version = mgpu_lib._mosaic_gpu_ext._get_ptxas_isa_version() # type: ignore any_reg = self.registers.flat[0] reg_type = any_reg.type is_vector_reg = isinstance(reg_type, ir.VectorType) reg_shape = tuple(ir.VectorType(reg_type).shape) if is_vector_reg else (1,) [vector_len] = reg_shape # This is meant to be a 1D assertion. if (new_reg_bitwidth := utils.bitwidth(new_dtype) * vector_len) % 8: raise ValueError( "Register bitwidth in target type must be divisible by 8, got" f" {new_reg_bitwidth}" ) # If the vector originates from a slice (common after relayouts), we # can fuse the slicing into the conversion and reuse many # preprocessing ops (shifts, prmts) accross different vectors. regs_from_32bit_slice = ( isinstance( _slice_op := getattr(any_reg.owner, "opview", None), vector.ExtractStridedSliceOp, ) and utils.bitwidth(_slice_op.source.type) == 32 # pyrefly: ignore[missing-attribute] and _slice_op.strides[0].value == 1 ) def packed_registers( dst_vector_len: int, *, if_not_sliced: bool ) -> Iterable[tuple[Sequence[tuple[int, ...]], ir.Value]]: """Tries to pack registers up to destination vector length.""" if regs_from_32bit_slice and if_not_sliced: for idx, reg in np.ndenumerate(self.registers): yield [idx], reg return generator = np.ndenumerate(self.registers) indices = [] regs = [] while True: try: for _ in range(max(dst_vector_len // vector_len, 1)): idx, reg = next(generator) indices.append(idx) regs.append(reg) yield indices, utils.vector_concat(regs) regs.clear() indices.clear() except StopIteration: break if regs: yield indices, utils.vector_concat(regs) if cur_dtype == i4 and new_dtype == f8e4m3fn: # The algorithm here is taken from CUTLASS's `NumericArrayConverter` # specialization for int4 -> f8e4m3, available at # https://github.com/NVIDIA/cutlass/blob/5c6bca04414e06ce74458ab0a2018e2b8272701c/include/cutlass/numeric_conversion.h#L4982 # and uint4 -> f8e4m3, available at # https://github.com/NVIDIA/cutlass/blob/5c6bca04414e06ce74458ab0a2018e2b8272701c/include/cutlass/numeric_conversion.h#L5196. # Each call to the function below will upcast 4 contiguous nibbles of # the input 32-bit register, and whether to select the 4 low nibbles or # the 4 high nibbles is determined by the `part` argument. def upcast_to_f8e4m3fn(reg: ir.Value, part: int): if self.is_signed: lut = [ 0x44403800, # [0, 1, 2, 3] encoded as f8e4m3fn 0x4E4C4A48, # [4, 5, 6, 7] encoded as f8e4m3fn 0xCACCCED0, # [-8, -7, -6, -5] encoded as f8e4m3fn 0xB8C0C4C8, # [-4, -3, -2, -1] encoded as f8e4m3fn ] else: lut = [ 0x44403800, # [0, 1, 2, 3] encoded as f8e4m3fn 0x4E4C4A48, # [4, 5, 6, 7] encoded as f8e4m3fn 0x53525150, # [8, 9, 10, 11] encoded as f8e4m3fn 0x57565554, # [12, 13, 14, 15] encoded as f8e4m3fn ] sign = arith.shrui(arith.andi(reg, c(0x88888888, i32)), c(1, i32)) # Ignore the sign when indexing into the LUT. lut_idx = arith.andi(reg, c(0x77777777, i32)) assert 0 <= part < 2 if part == 1: lut_idx = arith.shrui(lut_idx, c(16, i32)) sign = arith.shrui(sign, c(16, i32)) prmt_sign_pattern = arith.ori(sign, c(0x32103210, i32)) return llvm.inline_asm( i32, [lut_idx, prmt_sign_pattern], f""" {{ .reg .b32 pos_f8s, neg_f8s; prmt.b32 pos_f8s, {lut[0]}, {lut[1]}, $1; prmt.b32 neg_f8s, {lut[2]}, {lut[3]}, $1; prmt.b32 $0, pos_f8s, neg_f8s, $2; }} """, "=r,r,r", ) new_registers = np.empty_like(self.registers) # TODO(apaszke,bchetioui): Using 8 helps some (but not all) cases. # TODO(apaszke,bchetioui): Add the slice optimization here. packing_width = 8 if vector_len == 2 else 4 for indices, reg in packed_registers(packing_width, if_not_sliced=False): [group_size] = ir.VectorType(reg.type).shape assert group_size % vector_len == 0 int_ty = ir.IntegerType.get_signless(group_size * 4) reg_as_i32 = utils.bitcast(reg, int_ty) if int_ty != i32: reg_as_i32 = arith.extsi(i32, reg_as_i32) out_i32_regs = [ upcast_to_f8e4m3fn(reg_as_i32, part=part) for part in range(max(group_size // 4, 1)) ] out_vec_int = utils.vector_concat([ vector.broadcast(ir.VectorType.get((1,), i32), out_i32_reg) for out_i32_reg in out_i32_regs ]) out_vector_len = len(out_i32_regs) * 4 # Bitcast to i8 first to allow slicing as necessary, since LLVM chokes # on f8 types. out_vec = utils.bitcast( out_vec_int, ir.VectorType.get((out_vector_len,), i8) ) offset = 0 for idx in indices: sliced_out_vec = utils.vector_slice( out_vec, slice(offset, offset + vector_len) ) new_registers[idx] = utils.bitcast( sliced_out_vec, ir.VectorType.get((vector_len,), f8e4m3fn) ) offset += vector_len return FragmentedArray( _registers=new_registers, _layout=self.layout, _is_signed=None ) if cur_dtype == i4 and self.is_signed and new_dtype == bf16 and vector_len % 2 == 0: new_registers = np.empty_like(self.registers) out_vec_ty = ir.VectorType.get((vector_len,), new_dtype) # We use packed_registers for consistency, even though the packing is not # really profitable here: the PTX below begins by an op dependent on the # extracted part and so there are no ops that can be shared across packed # parts. for indices, reg in packed_registers(2, if_not_sliced=True): def upcast_i4_to_bf16(reg: ir.Value, part: int): assert 0 <= part < 4 # `cvt` with `.s2f6x2` instruction type introduced in PTX ISA v9.1 if utils.get_arch().major >= 10 and ptx_isa_version >= 91: int_reg = llvm.inline_asm( i32, [reg], f""" {{ .reg .b32 evens, odds, lo, hi; .reg .b16 part<4>; .reg .b16 scale; and.b32 evens, $1, 0x0F0F0F0F; and.b32 odds, $1, 0xF0F0F0F0; shl.b32 evens, evens, 4; prmt.b32 lo, evens, odds, 0x5140; prmt.b32 hi, evens, odds, 0x7362; mov.b32 {{part0, part1}}, lo; mov.b32 {{part2, part3}}, hi; mov.b16 scale, 0x8181; cvt.rn.scaled::n2::ue8m0.bf16x2.s2f6x2 $0, part{part}, scale; }} """, "=r,r", ) else: # The algorithm here is largely the same as CUTLASS's # NumericArrayConverter specialization for int4 -> bf16 casts. # We modify it slightly, because we only extract 2 values. # We first shift the value by 4 bits, to put the high int4 in low bits. # The prmt then blends the two values together, by putting them into the # low bits of each 16-bit subword of our register. Then, we use the lop3 # to zero any bits that don't belong to our int4s, and finally use the # XOR to: (1) set the exponent bits to 0x43 (at which point the mantissa # represents integer increments) and (2) flip the sign bit. If we # interpret the 4 bits as uint4 after the flip, then we'll see that # positive int4s will end up larger than negative int4s, with a bias of # 8. Use use the sub to subtract the base (our initial exponent) and the # bias coming from flipping the sign bit which is 136 (0x4308 as bits). int_reg = llvm.inline_asm( i32, [reg, arith.shrui(reg, c(4, i32))], f""" {{ .reg .b32 s<4>; prmt.b32 s1, $1, $2, 0xF{part + 4}F{part}; lop3.b32 s2, s1, 0x000F000F, 0x43084308, (0xf0 & 0xcc) ^ 0xaa; mov.b32 s3, 0x43084308; sub.bf16x2 $0, s2, s3; }} """, "=r,r,r", ) assert isinstance(int_reg, ir.Value) return utils.bitcast(int_reg, ir.VectorType.get((2,), bf16)) [group_size] = ir.VectorType(reg.type).shape assert group_size % vector_len == 0 assert group_size * 4 <= 32 int_ty = ir.IntegerType.get_signless(group_size * 4) # If the vector originates from a slice (common after relayouts), we # can fuse the slicing into the conversion and prevent LLVM from # generating a bunch of shifts to align the vector data to the LSB. # This also lets us share the right shift among more vectors. out_int_regs: list[ir.Value] = [] if regs_from_32bit_slice: slice_op: Any = reg.owner slice_offset = slice_op.offsets[0].value reg_int = utils.bitcast(slice_op.source, i32) assert slice_offset % 2 == 0 out_int_regs.extend( upcast_i4_to_bf16(reg_int, part=slice_offset // 2 + part) for part in range(group_size // 2) ) else: reg_slice_int = utils.bitcast(reg, int_ty) if int_ty != i32: reg_slice_int = arith.extsi(i32, reg_slice_int) out_int_regs.extend( upcast_i4_to_bf16(reg_slice_int, part=part) for part in range(group_size // 2) ) out_reg = utils.vector_concat(out_int_regs) offset = 0 for idx in indices: new_registers[idx] = new_reg = utils.vector_slice( out_reg, slice(offset, offset + vector_len) ) offset += vector_len assert new_reg.type == out_vec_ty return FragmentedArray( _registers=new_registers, _layout=self.layout, _is_signed=None ) if cur_dtype == i4 and self.is_signed and new_dtype == i8 and is_signed: new_registers = np.empty_like(self.registers) out_vec_ty = ir.VectorType.get((vector_len,), new_dtype) for indices, reg in packed_registers(8, if_not_sliced=True): def upcast_i4_to_i8(reg: ir.Value, first_valid_nibble: int = 0): # When first_valid_nibble is >0, then only the nibbles in the range # [first_valid_nibble, 8) will be upcast and placed in the low # elements of the output vector. All high entries are undefined. assert first_valid_nibble % 2 == 0 low_prmt = "".join(str(min(first_valid_nibble // 2 + i, 7)) for i in [5, 1, 4, 0]) high_prmt = "".join(str(min(first_valid_nibble // 2 + i, 7)) for i in [7, 3, 6, 2]) # Note: (0xf0 & 0xaa) | (0xcc & ~0xaa) = 0xe4. lop3 acts as a blend. # Below xN means the value of nibble N, sN means that all 4 bits are # equal to the sign bit of nibble N, and 00 means an all 0 nibble. out_struct = llvm.inline_asm( llvm.StructType.get_literal([i32, i32]), [reg], f""" {{ .reg .b32 high_even; // $2 is high_odd .reg .b32 low_odd; // $2 is low_even .reg .b32 sign_even, sign_odd; .reg .b32 i8_odd, i8_even; shl.b32 high_even, $2, 4; // x6x5x4x3x2x1x000 prmt.b32 sign_even, high_even, high_even, 0xba98; // s6s6s4s4s2s2s0s0 prmt.b32 sign_odd, $2, $2, 0xba98; // s7s7s5s5s3s3s1s1 shr.u32 low_odd, $2, 4; // 00x7x6x5x4x3x2x1 lop3.b32 i8_odd, sign_odd, low_odd, 0xf0f0f0f0, 0xe4; // s7x7s5x5s3x3s1x1 lop3.b32 i8_even, sign_even, $2, 0xf0f0f0f0, 0xe4; // s6x6s4x4s2x2s0x0 prmt.b32 $0, i8_even, i8_odd, 0x{low_prmt}; // s3x3s2x2s1x2s0x0 prmt.b32 $1, i8_even, i8_odd, 0x{high_prmt}; // s7x7s6x5s4x4s3x3 }} """, "=r,=r,r", ) assert isinstance(out_struct, ir.Value) i8_vec = ir.VectorType.get((4,), i8) return utils.vector_concat([ utils.bitcast(llvm.extractvalue(i32, out_struct, (i,)), i8_vec) for i in range(2) ]) [group_size] = ir.VectorType(reg.type).shape assert group_size % vector_len == 0 assert group_size * 4 <= 32 int_ty = ir.IntegerType.get_signless(group_size * 4) if regs_from_32bit_slice: slice_op: Any = reg.owner slice_offset = slice_op.offsets[0].value reg_int = utils.bitcast(slice_op.source, i32) reg_i8 = upcast_i4_to_i8(reg_int, first_valid_nibble=slice_offset) else: reg_slice_int = utils.bitcast(reg, int_ty) if int_ty != i32: reg_slice_int = arith.extsi(i32, reg_slice_int) reg_i8 = upcast_i4_to_i8(reg_slice_int) # distribute packed registers to original indices offset = 0 for idx in indices: new_registers[idx] = new_reg = utils.vector_slice( reg_i8, slice(offset, offset + vector_len) ) offset += vector_len assert new_reg.type == out_vec_ty return FragmentedArray( _registers=new_registers, _layout=self.layout, _is_signed=is_signed ) if ( cur_dtype == i8 and self.is_signed and new_dtype == bf16 and (vector_len == 2 or vector_len % 4 == 0) ): new_registers = np.empty_like(self.registers) def upcast_i8_to_bf16(reg, high): # We first embed the s8 into a bf16 with the exponent equal to # bias + mantissa bits. Then, we zero the msb that didn't fit into the # mantissa, zero out all bits other than msb, and subtract the last # two values from each other. This takes advantage of the fact that the # lsb of the exponent (msb of the second byte) is zero, which allows us # to losslesly pack the msb there. When 1, it doubles the value of s2, # making the result negative. return llvm.inline_asm( i32, [reg], f""" {{ .reg .b32 s<3>; prmt.b32 s0, $1, 0x43, {0x4342 if high else 0x4140}; and.b32 s1, s0, 0xff7fff7f; and.b32 s2, s0, 0xff80ff80; sub.bf16x2 $0, s1, s2; }} """, "=r,r", ) empty_vec_32 = llvm.mlir_undef(ir.VectorType.get((vector_len // 2,), i32)) pad_vec_16 = llvm.mlir_undef(ir.VectorType.get((1,), i16)) for idx, reg in np.ndenumerate(self.registers): if vector_len == 2: reg_16 = vector.bitcast(ir.VectorType.get((1,), i16), reg) reg_32 = utils.vector_concat([reg_16, pad_vec_16]) new_reg_32 = upcast_i8_to_bf16(reg_32, high=False) new_vec_32 = llvm.insertelement(empty_vec_32, new_reg_32, c(0, i32)) else: assert vector_len % 4 == 0 vec_32 = vector.bitcast(ir.VectorType.get((vector_len // 4,), i32), reg) new_vec_32 = empty_vec_32 for i in range(vector_len // 4): reg_32 = vector.extract(vec_32, [], [i]) low = upcast_i8_to_bf16(reg_32, high=False) high = upcast_i8_to_bf16(reg_32, high=True) new_vec_32 = llvm.insertelement(new_vec_32, low, c(2 * i, i32)) new_vec_32 = llvm.insertelement(new_vec_32, high, c(2 * i + 1, i32)) new_registers[idx] = vector.bitcast( ir.VectorType.get((vector_len,), new_dtype), new_vec_32 ) return FragmentedArray( _registers=new_registers, _layout=self.layout, _is_signed=is_signed ) # Most f8 casts are done by converting two elements at a time. def pairwise_convert(ptx_instr: str): src_bitwidth = utils.bitwidth(cur_dtype) tgt_bitwidth = utils.bitwidth(new_dtype) src_int_ty = ir.IntegerType.get_signless(src_bitwidth) tgt_int_ty = ir.IntegerType.get_signless(tgt_bitwidth) even_vector_len = vector_len + (vector_len % 2) new_registers = np.empty_like(self.registers) def do_convert(vec, convert_vec_len): assert convert_vec_len.bit_count() == 1 # Construct the PTX assembly out of ptx_instr. # The complication here is that some of the f4x2 instructions take or # return .b8 registers which are not supported by the NVPTX LLVM backend, # so we need to convert them to 16-bit at the boundaries. assert 4 <= src_bitwidth <= 32 assert 4 <= tgt_bitwidth <= 16 src_vec_bitwidth = src_bitwidth * convert_vec_len tgt_vec_bitwidth = tgt_bitwidth * convert_vec_len def get_ptx_constraint(bitwidth): if bitwidth <= 16: return "h", 1 return "r", bitwidth // 32 src_ptx_constraint, src_regs = get_ptx_constraint(src_vec_bitwidth) tgt_ptx_constraint, tgt_regs = get_ptx_constraint(tgt_vec_bitwidth) ptx_constraints = ["=" + tgt_ptx_constraint] * tgt_regs ptx_constraints += [src_ptx_constraint] * src_regs ptx_constraints = ",".join(ptx_constraints) ptx_lines = ["{"] src_packing = 32 // src_bitwidth if src_bitwidth == 32: # No unpacking necessary. get_ptx_operands = lambda vec: [ vector.extract(vec, [], [i]) for i in range(convert_vec_len) ] # Operands are flipped in 32-bit instructions. ptx_operands = [f"${tgt_regs + i + 1}, ${tgt_regs + i}" for i in range(0, convert_vec_len, 2)] elif src_bitwidth == 16: get_ptx_operands = lambda vec: [ utils.bitcast(utils.vector_slice(vec, slice(i, i + 2)), i32) for i in range(0, convert_vec_len, 2) ] ptx_operands = [f"${tgt_regs + i}" for i in range(convert_vec_len // 2)] elif convert_vec_len == 2: # Single narrow pair if src_bitwidth == 8: get_ptx_operands = lambda vec: [utils.bitcast(vec, i16)] ptx_operands = ["$1"] else: # NVPTX inline_asm has no support for 8-bit registers... assert src_bitwidth == 4 ptx_lines.append(".reg .b8 source_pair;") ptx_lines.append("mov.b16 {source_pair, _}, $1;") get_ptx_operands = lambda vec: [arith.extui(i16, utils.bitcast(vec, i8))] ptx_operands = ["source_pair"] else: # Multiple narrow pairs assert 4 <= src_bitwidth <= 8 assert convert_vec_len > 2 ptx_operands = [f"source_pair{i}" for i in range(convert_vec_len // 2)] ptx_lines.append( f".reg .b{src_bitwidth * 2} source_pair<{convert_vec_len // 2}>;") # 4xf4 is still less than 32 bits... pairs_per_src_reg = min(32, src_vec_bitwidth) // (src_bitwidth * 2) for i in range(src_regs): source_pairs = ", ".join(f"source_pair{i * pairs_per_src_reg + j}" for j in range(pairs_per_src_reg)) ptx_lines.append( f"mov.b{min(32, src_vec_bitwidth)} {{ {source_pairs} }}, ${tgt_regs + i};" ) if src_vec_bitwidth < 32: assert src_vec_bitwidth == 16 get_ptx_operands = lambda vec: [utils.bitcast(vec, i16)] else: get_ptx_operands = lambda vec: [ utils.bitcast(utils.vector_slice(vec, slice(i, i + src_packing)), i32) for i in range(0, convert_vec_len, src_packing) ] ptx_lines.append( f".reg .b{tgt_bitwidth * 2} result_pair<{convert_vec_len // 2}>;" ) ptx_targets = [f"result_pair{i}" for i in range(convert_vec_len // 2)] for tgt, op in zip(ptx_targets, ptx_operands, strict=True): ptx_lines.append(f"{ptx_instr} {tgt}, {op};") tgt_packing = 32 // tgt_bitwidth if tgt_vec_bitwidth > 32: ptx_result_ty = llvm.StructType.get_literal( [i32] * (convert_vec_len // tgt_packing) ) def process_ptx_result(ptx_result): vec_32 = ir.VectorType.get((tgt_packing,), tgt_int_ty) elements = [ utils.bitcast(llvm.extractvalue(i32, ptx_result, [i]), vec_32) for i in range(convert_vec_len // tgt_packing) ] return utils.vector_concat(elements) assert tgt_packing >= 2 ptx_targets_per_out = tgt_packing // 2 for out_vec_idx in range(convert_vec_len // tgt_packing): out_targets = ptx_targets[ out_vec_idx * ptx_targets_per_out:(out_vec_idx + 1) * ptx_targets_per_out ] if len(out_targets) == 1: [mov_src] = out_targets else: mov_src = f"{{ {', '.join(out_targets) } }}" ptx_lines.append(f"mov.b32 ${out_vec_idx}, {mov_src};") elif tgt_vec_bitwidth >= 16: ptx_result_ty = ir.IntegerType.get_signless(tgt_vec_bitwidth) def process_ptx_result(ptx_result): return utils.bitcast( ptx_result, ir.VectorType.get((convert_vec_len,), tgt_int_ty) ) if len(ptx_targets) == 1: [mov_src] = ptx_targets else: mov_src = f"{{ {', '.join(ptx_targets) } }}" ptx_lines.append(f"mov.b{tgt_vec_bitwidth} $0, {mov_src};") elif tgt_vec_bitwidth == 8: # NVPTX inline_asm has no support for 8-bit registers... ptx_result_ty = i16 def process_ptx_result(ptx_result): ptx_result = arith.trunci(i8, ptx_result) return utils.bitcast( ptx_result, ir.VectorType.get((convert_vec_len,), tgt_int_ty) ) [ptx_target] = ptx_targets ptx_lines.append(f"mov.b16 $0, {{ {ptx_target}, {ptx_target} }};") else: raise AssertionError("tgt_bitwidth too small") ptx_lines.append("}") ptx = "\t" + "\n\t".join(ptx_lines) ptx_result = llvm.inline_asm( ptx_result_ty, get_ptx_operands(vec), ptx, ptx_constraints ) result = process_ptx_result(ptx_result) assert result.type == ir.VectorType.get((convert_vec_len,), tgt_int_ty) return result longest_useful_vector = 32 // min(src_bitwidth, tgt_bitwidth) # We query the ptxas isa version as a proxy for PTX version. Old ptxas # binaries miscompile some of the patterns we generate here. for idx, reg in np.ndenumerate(self.registers): reg = utils.bitcast(reg, ir.VectorType.get((vector_len,), src_int_ty)) if vector_len % 2: reg = utils.vector_concat([reg, llvm.mlir_undef(ir.VectorType.get((1,), src_int_ty))]) convert_vec_len = longest_useful_vector base_idx = 0 result_vecs = [] while convert_vec_len >= 2: if cur_dtype == f4e2m1fn and convert_vec_len >= 4 and ptx_isa_version < 90: convert_vec_len //= 2 # ptxas miscompiles >=4xfp4 on CUDA 12.8... continue while (next_base_idx := base_idx + convert_vec_len) <= even_vector_len: vec = utils.vector_slice(reg, slice(base_idx, next_base_idx)) new_vec = do_convert(vec, convert_vec_len) assert new_vec.type == ir.VectorType.get((convert_vec_len,), tgt_int_ty) result_vecs.append(new_vec) base_idx = next_base_idx convert_vec_len //= 2 new_reg = utils.vector_concat(result_vecs) if vector_len % 2: new_reg = utils.vector_slice(new_reg, slice(0, vector_len)) new_registers[idx] = vector.bitcast(ir.VectorType.get((vector_len,), new_dtype), new_reg) return FragmentedArray( _registers=new_registers, _layout=self.layout, _is_signed=is_signed ) # Here we handle all conversions involving f8 types. # TODO(apaszke): Figure out proper satfinite control. supported_f8_f16: dict[ir.Type, tuple[ir.Type, ...]] = { f8e4m3fn: (f16,), f8e5m2: (f16,), f8e8m0fnu: (bf16,) } if ptx_isa_version >= 92: # Technically those downcasts are already supported in 91, but the casts # only become symmetric in 92 so that's what we use as the cutoff. supported_f8_f16[f8e4m3fn] += (bf16,) supported_f8_f16[f8e5m2] += (bf16,) f8_ptx_names: dict[ir.Type, str] = { f8e4m3fn: "e4m3", f8e5m2: "e5m2", f8e8m0fnu: "ue8m0" } f16_ptx_names: dict[ir.Type, str] = {f16: "f16", bf16: "bf16"} f8_types = f8_ptx_names.keys() f16_types = f16_ptx_names.keys() if f8e8m0fnu in {cur_dtype, new_dtype} and utils.get_arch().major < 10: raise ValueError( "f8e8m0fnu type only supported on Blackwell and newer GPUs" ) def check_supported_rounding(allowed_roundings, hw_limitation: bool = True): if rounding is None: return if rounding not in allowed_roundings: err_type = ValueError if hw_limitation else NotImplementedError raise err_type( f"Only {' and '.join(map(lambda x: x.name, allowed_roundings))}" f" rounding mode{'' if len(allowed_roundings) == 1 else 's'} are" f" supported for {cur_dtype} -> {new_dtype}, but got {rounding}" ) def get_fp8_rounding(fp8_type: ir.Type) -> str: assert fp8_type in f8_types allowed_rounding: tuple[Rounding, ...] if fp8_type == f8e8m0fnu: if rounding is None: return Rounding.TO_POSITIVE_INFINITY.ptx allowed_rounding = (Rounding.TO_ZERO, Rounding.TO_POSITIVE_INFINITY) else: if rounding is None: return Rounding.TO_NEAREST_EVEN.ptx allowed_rounding = (Rounding.TO_NEAREST_EVEN,) check_supported_rounding(allowed_rounding) return rounding.ptx # f8 <-> f32 if cur_dtype == f32 and new_dtype in f8_types: name_8 = f8_ptx_names[new_dtype] ptx_round = get_fp8_rounding(new_dtype) return pairwise_convert(f"cvt.{ptx_round}.satfinite.{name_8}x2.f32") # No f8 type supports direct conversion to f32, so we go via 16-bit floats. if cur_dtype in f8_types and new_dtype == f32: # We can pick any of the supported f16 types. return self.astype(supported_f8_f16[cur_dtype][0]).astype(f32) # f8 <-> f16 if new_dtype in f8_types and cur_dtype in supported_f8_f16[new_dtype]: name_16 = f16_ptx_names[cur_dtype] name_8 = f8_ptx_names[new_dtype] ptx_round = get_fp8_rounding(new_dtype) return pairwise_convert(f"cvt.{ptx_round}.satfinite.{name_8}x2.{name_16}x2") if cur_dtype in f8_types and new_dtype in supported_f8_f16[cur_dtype]: name_8 = f8_ptx_names[cur_dtype] name_16 = f16_ptx_names[new_dtype] return pairwise_convert(f"cvt.rn.{name_16}x2.{name_8}x2") # We don't emulate the unsupported f8 <-> f16 conversions, but rather force # the user to go via f32 to let them know it's expensive. if (new_dtype in f8_types and cur_dtype in f16_types) or ( new_dtype in f16_types and cur_dtype in f8_types ): # Remap the 16-bit type to the supported one. supported_16 = supported_f8_f16[cur_dtype if cur_dtype in f8_types else new_dtype] ok_new_dtype = " or ".join(map(str, supported_16)) ptx_hint = "" if ptx_isa_version < 92 and f8e8m0fnu not in {cur_dtype, new_dtype}: ptx_hint = ( " This conversion would be supported if a newer CUDA (ptxas)" " version was available." ) raise NotImplementedError( f"Hardware has no support for converting from {cur_dtype} to" f" {new_dtype} (only cast from {cur_dtype} to {ok_new_dtype} is" " supported). Cast to f32 first and then to the target type" " (expensive, but sufficient)." + ptx_hint ) # Repack through a shared 16-bit type. if cur_dtype in f8_types and new_dtype in f8_types: # We don't use set intersection to avoid introducing non-determinism. common_f16 = next( (e for e in supported_f8_f16[cur_dtype] if e in supported_f8_f16[new_dtype]), None ) if common_f16 is not None: return self.astype(common_f16).astype(new_dtype) raise NotImplementedError( f"Conversion from {cur_dtype} to {new_dtype} must go through f32," " which is expensive. Cast to f32 explicitly if you really want it." ) # f4e2m1fn casts if f4e2m1fn in {cur_dtype, new_dtype} and utils.get_arch().major < 10: raise ValueError( "f4e2m1fn casts only supported on Blackwell and newer GPUs" ) if cur_dtype == f4e2m1fn and new_dtype in {bf16, f16}: if new_dtype == bf16 and ptx_isa_version < 92: return self.astype(f32).astype(bf16) return pairwise_convert(f"cvt.rn.{f16_ptx_names[new_dtype]}x2.e2m1x2") if (cur_dtype == f4e2m1fn and new_dtype in f8_ptx_names) or ( new_dtype == f4e2m1fn and cur_dtype in f8_ptx_names ): f8_type = new_dtype if cur_dtype == f4e2m1fn else cur_dtype if f16 in supported_f8_f16[f8_type]: return self.astype(f16).astype(new_dtype) if ptx_isa_version >= 92: assert bf16 in supported_f8_f16[f8_type] return self.astype(bf16).astype(new_dtype) return self.astype(f32).astype(new_dtype) if cur_dtype == f4e2m1fn and new_dtype == f32: return self.astype(f16).astype(f32) if new_dtype == f4e2m1fn and cur_dtype in {bf16, f16}: if ptx_isa_version >= 91: return pairwise_convert(f"cvt.rn.satfinite.e2m1x2.{f16_ptx_names[cur_dtype]}x2") else: return self.astype(f32).astype(f4e2m1fn) if new_dtype == f4e2m1fn and cur_dtype == f32: return pairwise_convert("cvt.rn.satfinite.e2m1x2.f32") if new_dtype == f4e2m1fn or cur_dtype == f4e2m1fn: raise NotImplementedError(f"Unimplemented cast from {cur_dtype} to {new_dtype}") # Generic path. from_float = isinstance(cur_dtype, ir.FloatType) to_float = isinstance(new_dtype, ir.FloatType) from_integer = isinstance(cur_dtype, ir.IntegerType) to_integer = isinstance(new_dtype, ir.IntegerType) from_narrow_float = from_float and utils.bitwidth(cur_dtype) <= 8 to_narrow_float = to_float and utils.bitwidth(new_dtype) <= 8 if from_narrow_float or to_narrow_float: raise NotImplementedError( f"Unsupported conversion involving narrow float types: {cur_dtype}" f" to {new_dtype}" ) arith_rounding = rounding.arith if rounding else None if from_float and to_float: cur_ty_width = ir.FloatType(cur_dtype).width new_ty_width = ir.FloatType(new_dtype).width if cur_ty_width == new_ty_width: # There is no instruction to perform conversions between two float types # of the same width. Go through the next-larger standard type. # TODO(bchetioui): support conversions between float types of width 8. # Which larger type to pick will depend on the number of bits in the # smallest exponent. if cur_ty_width != 16: raise NotImplementedError( "Conversion between float types of width other than 16 not" " supported" ) larger_ty = ir.F32Type.get() match self.layout: case WGStridedFragLayout() | TiledLayout(): shape = ir.VectorType(self.registers.flat[0].type).shape upcast_ty = ir.VectorType.get(shape, larger_ty) case WGSplatFragLayout(): upcast_ty = larger_ty case _: raise NotImplementedError(f"Unsupported layout {self.layout}") convert = lambda ty, x: arith.truncf( ty, arith.extf(upcast_ty, x), roundingmode=arith_rounding ) elif cur_ty_width > new_ty_width: convert = functools.partial(arith.truncf, roundingmode=arith_rounding) else: convert = arith.extf elif from_integer and to_integer: if ir.IntegerType(cur_dtype).width > ir.IntegerType(new_dtype).width: convert = arith.trunci else: convert = arith.extsi if self.is_signed else arith.extui elif from_integer and to_float: # TODO(apaszke): Support other rounding modes. check_supported_rounding((Rounding.TO_NEAREST_EVEN,), hw_limitation=False) convert = arith.sitofp if self.is_signed else arith.uitofp elif from_float and to_integer: # TODO(apaszke): Support other rounding modes. check_supported_rounding((Rounding.TO_ZERO,), hw_limitation=False) convert = arith.fptosi if is_signed else arith.fptoui else: raise NotImplementedError(f"Unsupported conversion {cur_dtype} -> {new_dtype}") new_registers = np.empty_like(self.registers) match self.layout: case WGStridedFragLayout() | TiledLayout(): shape = ir.VectorType(self.registers.flat[0].type).shape new_reg_ty = ir.VectorType.get(shape, new_dtype) case WGSplatFragLayout(): new_reg_ty = new_dtype case _: raise NotImplementedError(f"Unsupported layout {self.layout}") for idx, reg in np.ndenumerate(self.registers): new_registers[idx] = convert(new_reg_ty, reg) return FragmentedArray( _registers=new_registers, _layout=self.layout, _is_signed=is_signed ) def reduce( self, op: str | Callable[[ir.Value, ir.Value], ir.Value], axis: int | Sequence[int], scratch: ir.Value | None = None, *, acc_ilp: int | None = None, target_layout: TiledLayout | None = None, ) -> FragmentedArray: """Reduces `axis` out of the array. Args: op: The reduction operation. axis: The axes to reduce. May be empty if the array is unreduced, in which case only the pending cross-thread reduction is performed. scratch: SMEM scratch used to exchange data across warps. Only required when the reduction spans warps. acc_ilp: The number of independent accumulators to use. target_layout: The layout of the result. Defaults to fully reducing `axis`. Each warp or lane dimension being reduced may instead be left `Unreduced`, which defers that part of the reduction to a later `reduce`. """ if not isinstance(self.layout, TiledLayout) and target_layout is not None: raise NotImplementedError( "A target layout can only be specified when reducing an array with" " a TiledLayout" ) if isinstance(axis, int): axis = (axis,) else: axis = tuple(axis) if not axis and not self.is_unreduced: return self if self.is_unreduced: assert isinstance(self.layout, TiledLayout) if not isinstance(op, str) or op != self.layout.unreduced_operation: raise ValueError( "reduce on an unreduced FragmentedArray requires the same" f" operation, got {op!r}, expected" f" {self.layout.unreduced_operation!r}" ) if target_layout is not None: assert isinstance(self.layout, TiledLayout) untiled_rank = len(self.shape) - len(self.layout.base_tile_shape) tiled_axis = tuple(a - untiled_rank for a in axis if a >= untiled_rank) unreduced_op = target_layout.unreduced_operation expected_layout = self.layout.reduce( tiled_axis, local_only=unreduced_op is not None, op=unreduced_op ) def has_compatible_dims(expected, target) -> bool: if len(expected) != len(target): return False for d1, d2 in zip(expected, target, strict=True): if d1 == d2: continue if ( isinstance(d1, Unreduced) and isinstance(d2, Replicated) and d1.times == d2.times ): continue return False return True if ( unreduced_op not in (None, op) or expected_layout.tiling != target_layout.tiling or not has_compatible_dims( expected_layout.warp_dims, target_layout.warp_dims ) or not has_compatible_dims( expected_layout.lane_dims, target_layout.lane_dims ) or expected_layout.vector_dim != target_layout.vector_dim ): raise ValueError( "target_layout must match self.layout.reduce(axis), up to warp and" " lane dimensions left unreduced" ) i32 = ir.IntegerType.get_signless(32) splat_op = None redux_op = None default_acc_ilp = 1 # TODO(apaszke): For associative reductions that reduce both inside and # across warps, we could just have everyone use SMEM atomics instead of # performing an explicit warp reduction in registers. if isinstance(op, str): match op: case "add": reduced_elems = math.prod(self.shape[a] for a in axis) if isinstance(self.mlir_dtype, ir.FloatType): op = addf # TODO(cjfj): Consider bumping to 16 for f16/bf16. default_acc_ilp = 8 splat_op = lambda x: arith.mulf(x, c(reduced_elems, x.type)) # TODO(apaszke): Use redux.sync on Blackwell for f32. elif isinstance(self.mlir_dtype, ir.IntegerType): op = arith.addi splat_op = lambda x: arith.muli(x, c(reduced_elems, x.type)) if utils.bitwidth(self.mlir_dtype) == 32: redux_op = functools.partial(utils.redux, kind=ReductionKind.ADD) else: raise NotImplementedError(self.mlir_dtype) case "max": default_acc_ilp = 4 if isinstance(self.mlir_dtype, ir.F32Type): op = self._lift_fast_instr("max.NaN.f32") if utils.get_arch().major == 10: redux_op = functools.partial(utils.redux, kind=ReductionKind.FMAX) elif isinstance(self.mlir_dtype, ir.F16Type): op = self._lift_fast_packed_instr("max.NaN.f16x2", "max.NaN.f16") elif isinstance(self.mlir_dtype, ir.BF16Type): op = self._lift_fast_packed_instr("max.NaN.bf16x2", "max.NaN.bf16") elif isinstance(self.mlir_dtype, ir.FloatType): op = arith.maximumf elif isinstance(self.mlir_dtype, ir.IntegerType): op = arith.maxsi if self.is_signed else arith.maxui if utils.bitwidth(self.mlir_dtype) == 32: kind = ReductionKind.MAX if self.is_signed else ReductionKind.UMAX redux_op = functools.partial(utils.redux, kind=kind) else: raise NotImplementedError(self.mlir_dtype) splat_op = lambda x: x case "min": default_acc_ilp = 4 if isinstance(self.mlir_dtype, ir.F32Type): op = self._lift_fast_instr("min.NaN.f32") if utils.get_arch().major == 10: redux_op = functools.partial(utils.redux, kind=ReductionKind.FMIN) elif isinstance(self.mlir_dtype, ir.F16Type): op = self._lift_fast_packed_instr("min.NaN.f16x2", "min.NaN.f16") elif isinstance(self.mlir_dtype, ir.BF16Type): op = self._lift_fast_packed_instr("min.NaN.bf16x2", "min.NaN.bf16") elif isinstance(self.mlir_dtype, ir.FloatType): op = arith.minimumf elif isinstance(self.mlir_dtype, ir.IntegerType): op = arith.minsi if self.is_signed else arith.minui else: raise NotImplementedError(self.mlir_dtype) splat_op = lambda x: x case "prod": reduced_elems = math.prod(self.shape[a] for a in axis) if isinstance(self.mlir_dtype, ir.FloatType): op = arith.mulf default_acc_ilp = 8 # For splat, prod(x, x, ..., x) = x^n splat_op = lambda x: mlir_math.powf( x, c(float(reduced_elems), x.type) ) elif isinstance(self.mlir_dtype, ir.IntegerType): op = arith.muli # For splat, use repeated squaring to compute x^n splat_op = functools.partial(_int_pow, n=reduced_elems) else: raise NotImplementedError(self.mlir_dtype) case _: raise ValueError(f"Unrecognized reduction operator: {op}") if acc_ilp is None: acc_ilp = default_acc_ilp assert not isinstance(op, str) match self.layout: case WGStridedFragLayout(shape=_, vec_size=vec_size): if set(axis) != set(range(len(self.shape))): raise NotImplementedError( "Warpgroup strided layout only support reductions along all axes" ) # We reinterpret the data as a tiled layout. We're reducing it all anyway. layout = TiledLayout( tiling=Tiling(((128 * vec_size,), (32 * vec_size,), (vec_size,))), warp_dims=(-3,), lane_dims=(-2,), vector_dim=-1, ) return FragmentedArray( _registers=self.registers.reshape( layout.registers_shape((math.prod(self.shape),)) ), _layout=layout, _is_signed=self.is_signed, ).reduce(op, 0, scratch) case WGSplatFragLayout(): if splat_op is None: raise NotImplementedError( "Splat reductions only supported when the operator is a string" ) assert not self.registers.shape return FragmentedArray( _registers=np.asarray( splat_op(self.registers.item()), dtype=object ), _layout=WGSplatFragLayout( tuple(d for a, d in enumerate(self.shape) if a not in axis) ), _is_signed=self.is_signed, ) case TiledLayout(): pass case _: raise NotImplementedError(self.layout) # Silence type checker complaints. assert isinstance(self.layout, TiledLayout) layout = self.layout untiled_rank = len(self.shape) - len(layout.base_tile_shape) tiled_tiling_shape = layout.tiled_tiling_shape tiled_axes = tuple(a - untiled_rank for a in axis if a >= untiled_rank) if target_layout is None: target_layout = layout.reduce(tiled_axes) tiled_reduced_dims = (False,) * (len(layout.base_tile_shape) + len(tiled_tiling_shape)) for a in tiled_axes: tiled_reduced_dims = tuple( r or d for r, d in zip(tiled_reduced_dims, layout.tiling.tile_dimension(a), strict=True) ) reduced_dims = (*(a in axis for a in range(untiled_rank)), *tiled_reduced_dims) regs_shape = self.registers.shape reduced_shape = tuple( d if r else 1 for r, d in zip(reduced_dims, regs_shape, strict=True) ) remaining_shape = tuple( 1 if r else d for r, d in zip(reduced_dims, regs_shape) ) out_regs = np.empty(remaining_shape, dtype=object) index = ir.IndexType.get() def apply_op(a: ir.Value | None, b: ir.Value | None) -> ir.Value | None: if a is None: return b if b is None: return a return op(a, b) def should_reduce( d: int | Replicated | Unreduced, target_d: int | Replicated | Unreduced, ) -> bool: return isinstance(target_d, Replicated) and not isinstance(d, Replicated) def reduce_within_warp(out_idx): # Compute partial reductions, breaking the dependency between subsequent # element-wise operations. [vec_len] = ir.VectorType(self.registers.flat[0].type).shape num_reductions = math.prod(remaining_shape) num_parts = max(1, acc_ilp // vec_len // num_reductions) part_regs: list[ir.Value | None] = [None] * num_parts for i, red_idx in enumerate(np.ndindex(reduced_shape)): src_idx = tuple(o + r for o, r in zip(out_idx, red_idx)) slot = i % num_parts part_regs[slot] = apply_op(part_regs[slot], self.registers[src_idx]) out_reg = functools.reduce(apply_op, part_regs) assert out_reg is not None # Reduce within the vector dimension, if necessary. if reduced_dims[layout.vector_dim]: [vec_len] = ir.VectorType(out_reg.type).shape scalar_out_reg: ir.Value | None = None for i in range(vec_len): scalar = vector.extract( out_reg, dynamic_position=[], static_position=ir.DenseI64ArrayAttr.get([i]), ) scalar_out_reg = ( scalar if scalar_out_reg is None else op(scalar_out_reg, scalar) ) assert scalar_out_reg is not None out_reg = vector.broadcast( ir.VectorType.get((1,), out_reg.type.element_type), scalar_out_reg ) # Reduce across warp lanes, if necessary (using warp shuffles). lane_pairs = tuple( zip(layout.lane_dims, target_layout.lane_dims, strict=True) ) if any(should_reduce(d, td) for d, td in lane_pairs): all_lanes = all(should_reduce(d, td) for d, td in lane_pairs) # It doesn't make sense to use redux unless we reduce across all lanes. # The instruction seems to have a uniform register output. if redux_op is not None and all_lanes: out_reg = redux_op(out_reg, arith.constant(i32, 0xffffffff)) else: lane_stride = 1 for d, td in reversed(lane_pairs): # Iterate minor-to-major size = tiled_tiling_shape[d] if isinstance(d, int) else d.times if not should_reduce(d, td): lane_stride *= size else: assert lane_stride.bit_count() == 1 reduction_size = size while reduction_size > 1: other_out_reg = utils.shfl_bfly(out_reg, lane_stride) out_reg = op(out_reg, other_out_reg) lane_stride *= 2 reduction_size //= 2 assert lane_stride == WARP_SIZE, lane_stride return out_reg def swizzle_warp_idx_fn(lane_idx: ir.Value, vec_len: int): bitwidth = utils.bitwidth(self.mlir_dtype) num_banks = bank_bitwidth = 32 bitwidth_per_store = vec_len * bitwidth num_banks_per_output = WARPS_IN_WARPGROUP * bitwidth_per_store // bank_bitwidth # This range supports vector types from vector<1xi8> to # {vector<32xi8>, vector<16xi16>, vector<8xi32>}, which should be plenty # for realistic use cases. Other cases are not guaranteed to work, but # even if they do, their performance hasn't been evaluated. As a result, # we prefer failing explicitly, to make sure that we don't end up # emitting reductions with poor performance. # # Note: this implementation is batch invariant, because we use a XOR # swizzle, and a binary tree to perform the final reduction. The XOR # swizzle only permutes the left and right subtrees, which allows us to # always recover the same result (since floating-point operations are # commutative). if ( WARPS_IN_WARPGROUP * bitwidth_per_store < bank_bitwidth or bitwidth_per_store > 128 or num_banks % num_banks_per_output != 0 ): raise NotImplementedError( "Unoptimized configuration for cross-warp reduction: " f"{self.mlir_dtype} with {vec_len=}" ) # Define one row to be 128 bytes (32 banks of 4 bytes). For a given lane # index, we want to store the data coming from all 4 warps # contiguously in order to enable vectorized loads later on. If we # simply store the data in order of thread_idx naively, this will # result in bank conflicts, since each warp will hit only a quarter of # the shared memory banks. # # In order to avoid these bank conflicts, we swizzle the data # manually, such that across every four rows of 128 bytes, each warp # will hit all the shared memory banks. lanes_per_row = num_banks // num_banks_per_output num_rows = WARP_SIZE // lanes_per_row row_idx = arith.divui(lane_idx, c(lanes_per_row, i32)) match num_banks_per_output: case 1: assert num_rows == 1, num_rows # Here, each lane stores to a different bank, so we don't need to # swizzle the warp index at all. swizzle_warp_idx = lambda widx: widx case 2: assert num_rows == 2, num_rows # In this case, each lane stores 16-bit in a single bank, and the # stores look like: # # | Bank 0 | Bank 1 | ... # | 16-bit | 16-bit | 16-bit | 16-bit | ... # r0: | warp0 lane0 | warp1 lane0 | warp2 lane0 | warp3 lane0 | ... # r1: | warp2 lane16 | warp3 lane16 | warp0 lane16 | warp1 lane16 | ... # # such that the first 16 lanes of each warp are mapped to row 0, and # the last 16 lanes of each warp are mapped to row 1, and the # relative ordering of elements coming from different warps is # always the same in each row. # # We avoid bank conflicts on the two rows by xoring the warp index # with 2 on row 1, such that lanes in warp0 and warp1 hit even banks # on row 0 and odd banks on row 1 (and the opposite for warp2 and # warp3). lane_xor = arith.shli(row_idx, c(1, i32)) swizzle_warp_idx = lambda widx: arith.xori(widx, lane_xor) # As long as we use a multiple of 4 banks, we can use the same # formulation to swizzle the order of the 4 warps. case x if x % 4 == 0: # In that case, each lane stores a multiple of 32-bit. The following # shows how the stores look like in the 32-bit case: # # r0: | warp0 lane0 | warp1 lane0 | warp2 lane0 | warp3 lane0 | ... # r1: | warp1 lane8 | warp0 lane8 | warp3 lane8 | warp2 lane8 | ... # r2: | warp2 lane16 | warp3 lane16 | warp0 lane16 | warp1 lane16 | ... # r3: | warp3 lane24 | warp2 lane24 | warp1 lane24 | warp0 lane24 | ... # # Lanes 0-8 are mapped to row 0, lanes 8-16 to row 1, lanes 16-24 to # row 2, and lanes 24-32 to row 3. In each row, the index of the # warp is swizzled by xoring it with the index of the row. rhs = row_idx if x == 4 else arith.andi(row_idx, c(3, i32)) swizzle_warp_idx = lambda widx: arith.xori(widx, rhs) case _: raise NotImplementedError(num_banks_per_output) return swizzle_warp_idx def store_swizzled( reg: ir.Value, step_idx: int, lane_idx: ir.Value, scratch: ir.Value, swizzle_warp_idx: Callable[[ir.Value], ir.Value] ): [vec_len] = ir.VectorType(reg.type).shape thread_idx = utils.thread_idx() warp_idx = arith.divui( arith.remui(thread_idx, c(WARPGROUP_SIZE, i32)), c(WARP_SIZE, i32) ) step_base_scratch_idx = c(step_idx * WARPGROUP_SIZE, i32) lane_base_scratch_idx = arith.addi( step_base_scratch_idx, arith.muli(lane_idx, c(WARPS_IN_WARPGROUP, i32)) ) store_idx = arith.addi(lane_base_scratch_idx, swizzle_warp_idx(warp_idx)) as_index = lambda x: arith.index_cast(index, x) # TODO(bchetioui): investigate whether adding predication here can yield # additional performance improvements. In the case where we have a # `Replicated` dimension in there, we will repeat the same store # multiple times. Maybe the memory controller resolves this conflict # automatically, but we should investigate. vector.store( reg, scratch, [as_index(arith.muli(store_idx, c(vec_len, i32)))]) def reduce_stored( reg_ty: ir.VectorType, step_idx: int, lane_idx: ir.Value, swizzle_warp_idx: Callable[[ir.Value], ir.Value] ): assert scratch is not None [vec_len] = ir.VectorType(reg_ty).shape out_reg = None step_base_scratch_idx = c(step_idx * WARPGROUP_SIZE, i32) lane_base_scratch_idx = arith.addi( step_base_scratch_idx, arith.muli(lane_idx, c(WARPS_IN_WARPGROUP, i32)) ) # warp_idx & warp_group_mask gives you the reduction group of the current warp. warp_pairs = tuple( zip(layout.warp_dims, target_layout.warp_dims, strict=True) ) if all(should_reduce(d, td) for d, td in warp_pairs): # When we load all the data that we have stored, we can omit swizzling # the warp index without any loss of correctness or determinism. By # relying on the properties of XOR and using a "tree reduction" # to reduce the data, we also maintain batch invariance. # # Without this manual optimization, LLVM can fail to recognize that # it can use wider load instructions, leading to worse performance # (presumably due to scheduling pressure) and sometimes also due to # unnecessary bank conflicts. # # TODO(bchetioui): there are still load conflicts in the case where # `vec_len * WARPS_IN_WARPGROUP * bitwidth` exceeds 128 bits. To # avoid bank conflicts in that case, we need to swizzle the loads as # well, using a similar pattern as above for groups of 128 bits. A # little more care will have to be taken to uphold batch invariance. load_ty = ir.VectorType.get((vec_len * WARPS_IN_WARPGROUP,), reg_ty.element_type) load_idx = arith.muli(lane_base_scratch_idx, c(vec_len, i32)) parts = vector.load( load_ty, scratch, [arith.index_cast(index, load_idx)] ) parts = [utils.vector_slice(parts, slice(i * vec_len, (i + 1) * vec_len)) for i in range(WARPS_IN_WARPGROUP)] out_reg = op(op(parts[0], parts[1]), op(parts[2], parts[3])) else: # 4 has only two non-trivial prime factors: 2 and 2. assert len(layout.warp_dims) == 2 (wd0, twd0), (wd1, twd1) = warp_pairs # TODO(bchetioui): these paths are optimizable. The above logic is # well-suited for loads of values stored by all 4 warps, but we should # adapt the store logic to also account for these cases where we only # load the value stored by every other warp. In this case, we should # use a different swizzle function, in order to make sure we can # always get vectorized loads! if should_reduce(wd0, twd0): warp_offsets, warp_group_mask = [0, 2], 1 else: assert should_reduce(wd1, twd1) warp_offsets, warp_group_mask = [0, 1], 2 thread_idx = utils.thread_idx() warp_idx = arith.divui( arith.remui(thread_idx, c(WARPGROUP_SIZE, i32)), c(WARP_SIZE, i32) ) warp_reduction_group = arith.andi(warp_idx, arith.constant(i32, warp_group_mask)) for warp_offset in warp_offsets: reduced_warp = arith.addi(warp_reduction_group, c(warp_offset, i32)) load_idx = arith.addi(lane_base_scratch_idx, swizzle_warp_idx(reduced_warp)) part = vector.load( reg_ty, scratch, [arith.index_cast(index, arith.muli(load_idx, c(vec_len, i32)))] ) out_reg = part if out_reg is None else op(out_reg, part) return out_reg # Note that we will infer a splat layout if we reduce all dimensions, but in # that case the .vector_length we get here will be 1 anyway. vec_len = target_layout.vector_length thread_idx = utils.thread_idx() lane_idx = arith.remui(thread_idx, c(WARP_SIZE, i32)) do_reduce_across_warps = any( should_reduce(d, td) for d, td in zip(layout.warp_dims, target_layout.warp_dims, strict=True) ) if do_reduce_across_warps: if scratch is None: raise ValueError( "scratch must be provided when cross-warp reduction is required" ) scratch_ty = ir.MemRefType(scratch.type) if scratch_ty.rank != 1: raise ValueError(f"Expected rank 1 for scratch, got {scratch_ty.rank}") if scratch_ty.element_type != self.mlir_dtype: raise ValueError( f"Expected element type {self.mlir_dtype} for scratch, got" f" {scratch_ty.element_type}" ) # TODO(apaszke): All lanes that replicate data can share the same scratch. # For now we treat the complete reduction as a special case. reduces_all_dims = set(axis) == set(range(len(self.shape))) unique_lanes = 1 if reduces_all_dims else 32 scratch_ty = ir.MemRefType(scratch.type) scratch_elems_per_register = WARPS_IN_WARPGROUP * unique_lanes * vec_len if scratch_ty.shape[0] < scratch_elems_per_register: available_bytes = scratch_ty.shape[0] * utils.bitwidth(scratch_ty.element_type) // 8 required_bytes = scratch_elems_per_register * utils.bitwidth(scratch_ty.element_type) // 8 raise ValueError( f"Required reduction scratch size ({required_bytes} bytes) is " f"larger than the available scratch size ({available_bytes} bytes)" ) if scratch_ty.get_strides_and_offset()[0] != [1]: raise ValueError("Expected scratch to be contiguous") num_concurrent_cross_warp_reductions = scratch_ty.shape[0] // scratch_elems_per_register if reduces_all_dims: lane_idx = c(0, i32) swizzle_warp_idx = swizzle_warp_idx_fn(lane_idx, vec_len) else: lane_idx = num_concurrent_cross_warp_reductions = swizzle_warp_idx = None unreduced_indices: list[tuple[int, ...]] = [] for out_idx in np.ndindex(remaining_shape): out_reg = reduce_within_warp(out_idx) reg_ty = ir.VectorType(out_reg.type) if do_reduce_across_warps: # TODO(bchetioui): explore pipelining computer+store and loads+reduce # by double buffering the scratch. This could offer more # instruction-level parallelism. assert scratch is not None assert lane_idx is not None and swizzle_warp_idx is not None step = len(unreduced_indices) store_swizzled(out_reg, step, lane_idx, scratch, swizzle_warp_idx) unreduced_indices.append(out_idx) if len(unreduced_indices) == num_concurrent_cross_warp_reductions: utils.warpgroup_barrier() for i, unreduced_index in enumerate(unreduced_indices): out_regs[unreduced_index] = reduce_stored( reg_ty, i, lane_idx, swizzle_warp_idx ) unreduced_indices = [] utils.warpgroup_barrier() else: out_regs[out_idx] = out_reg if unreduced_indices: utils.warpgroup_barrier() for i, unreduced_index in enumerate(unreduced_indices): assert lane_idx is not None and swizzle_warp_idx is not None out_regs[unreduced_index] = reduce_stored( # pyrefly: ignore[unbound-name] reg_ty, i, lane_idx, swizzle_warp_idx ) utils.warpgroup_barrier() del unreduced_indices # Infer the output layout and reshape the registers accordingly. reduced_logical_shape = list(self.shape) for a in sorted(axis, reverse=True): del reduced_logical_shape[a] if not reduced_logical_shape: if target_layout.has_unreduced_dims: raise NotImplementedError( "Partial reductions to a scalar are not supported" ) # Complete reduction results in a splat. assert out_regs.size == 1 out_reg = out_regs.flat[0] assert ir.VectorType(out_reg.type).shape == [1] out_reg = vector.extract( out_reg, dynamic_position=[], static_position=ir.DenseI64ArrayAttr.get([0]), ) out_regs = np.asarray(out_reg, dtype=object) return FragmentedArray( _registers=out_regs, _layout=WGSplatFragLayout(()), _is_signed=self.is_signed, ) else: out_regs = out_regs.reshape( target_layout.registers_shape(tuple(reduced_logical_shape)) ) return FragmentedArray( _registers=out_regs, _layout=target_layout, _is_signed=self.is_signed, ) def broadcast(self, shape: tuple[int, ...]) -> FragmentedArray: new_layout: FragmentedLayout if isinstance(self.layout, WGStridedFragLayout): new_layout = WGStridedFragLayout(shape, self.layout.vec_size) elif isinstance(self.layout, WGSplatFragLayout): new_layout = WGSplatFragLayout(shape) else: raise NotImplementedError(self.layout) dims = tuple(range(len(shape) - len(self.shape), len(shape))) return self.broadcast_in_dim(shape, dims, new_layout) def reshape(self, shape: tuple[int, ...]) -> FragmentedArray: if self.is_unreduced: raise NotImplementedError if self.shape == shape: return self if math.prod(shape) != math.prod(self.shape): raise ValueError(f"Can't reshape {self.shape} to {shape}") match self.layout: case WGSplatFragLayout() | WGStridedFragLayout(): new_layout = dataclasses.replace(self.layout, shape=shape) return FragmentedArray( _registers=self.registers, _layout=new_layout, _is_signed=self.is_signed, ) case TiledLayout(): base_tile_shape = self.layout.base_tile_shape assert base_tile_shape old_shape_suffix = self.shape[-len(base_tile_shape):] new_shape_suffix = shape[-len(base_tile_shape):] # We already know that old_shape_suffix[0] is divisible by # base_tile_shape[0]. if ( old_shape_suffix[1:] != new_shape_suffix[1:] or new_shape_suffix[0] % base_tile_shape[0] ): raise ValueError( f"Can't reshape {self.shape} to {shape} with a tiled layout with" f" base tile of {base_tile_shape}" ) new_registers_shape = self.layout.registers_shape(shape) return FragmentedArray( _registers=self.registers.reshape(new_registers_shape), _layout=self.layout, _is_signed=self.is_signed, ) case _: raise NotImplementedError(self.layout) def broadcast_minor(self, n: int) -> FragmentedArray: if self.is_unreduced: raise NotImplementedError if len(self.shape) != 1: raise ValueError("Broadcast minor is only supported for 1D arrays") if n % 8: raise ValueError(f"The broadcast dimension must be a multiple of 8, got {n}") if self.layout == WGMMA_ROW_LAYOUT: new_layout = WGMMA_LAYOUT elif self.layout == TCGEN05_ROW_LAYOUT: new_layout = TCGEN05_LAYOUT else: raise NotImplementedError(self.layout) return self.broadcast_in_dim((self.shape[0], n), (0,), new_layout) def broadcast_in_dim( self, shape: tuple[int, ...], source_dimensions: tuple[int, ...], layout: FragmentedLayout, ) -> FragmentedArray: if self.is_unreduced: raise NotImplementedError if len(source_dimensions) != len(self.shape): raise ValueError( f"The number of source dimensions ({len(source_dimensions)}) must" f" match the rank of the array ({len(self.shape)})" ) if not all(0 <= d < len(shape) for d in source_dimensions): raise ValueError( f"{source_dimensions=} must be within the range [0, {len(shape)})" ) if any( d1 >= d2 for d1, d2 in zip(source_dimensions, source_dimensions[1:]) ): raise ValueError( f"{source_dimensions=} must be strictly increasing and unique" ) for i, target_dim in enumerate(source_dimensions): if self.shape[i] != shape[target_dim] and self.shape[i] != 1: raise ValueError( f"Dimension {i} has size {self.shape[i]} in source shape and" f" {shape[target_dim]} in shape after broadcast" ) if isinstance(self.layout, WGSplatFragLayout): return type(self).splat( self.registers.item(), shape, layout, is_signed=self.is_signed ) if isinstance(self.layout, WGStridedFragLayout) and isinstance(layout, WGStridedFragLayout): if not is_supported_strided_layout_broadcast(self.layout, layout, source_dimensions): raise NotImplementedError( "Unsupported broadcast for WGStridedFragLayout" ) assert layout.shape == shape, (layout.shape, shape) return FragmentedArray( _registers=np.tile( self.registers, math.prod(shape) // math.prod(self.shape), ), _layout=layout, _is_signed=self.is_signed, ) if not isinstance(self.layout, TiledLayout) or not isinstance(layout, TiledLayout): raise NotImplementedError(self.layout, layout) if len(layout.base_tile_shape) != len(shape): raise NotImplementedError( "Tiling rank different than broadcast result rank, " f"{layout.base_tile_shape} vs {shape}" ) new_dimensions = sorted(set(range(len(shape))) - set(source_dimensions)) expected_layout = layout.reduce(new_dimensions) if expected_layout != self.layout: raise ValueError( "Source and destination layouts aren't compatible for a broadcast" ) new_registers_shape = layout.registers_shape(shape) pre_broadcast_registers_shape = list(new_registers_shape) for new_dim in new_dimensions: for i, is_new in enumerate(layout.tiling.tile_dimension(new_dim)): if is_new: pre_broadcast_registers_shape[i] = 1 # The broadcast for all dims but the vector_dim amounts to repeating the # registers along the new dimensions. Along the vector_dim, we actually need # to extend the vector length to change the type of the registers. if layout.vector_length != self.layout.vector_length: assert self.layout.vector_length == 1 registers = np.empty_like(self.registers) for idx, reg in np.ndenumerate(self.registers): registers[idx] = utils.vector_concat([reg] * layout.vector_length) else: registers = self.registers new_registers = np.broadcast_to( registers.reshape(pre_broadcast_registers_shape), new_registers_shape, ) return FragmentedArray( _registers=new_registers, _layout=layout, _is_signed=self.is_signed, ) def select(self, on_true, on_false): if ( not isinstance(self.mlir_dtype, ir.IntegerType) or ir.IntegerType(self.mlir_dtype).width != 1 ): raise NotImplementedError # We change the receiver here, because the return type is defined by # `on_true` and `on_false` and not the predicate `self`. return on_true._pointwise( lambda t, p, f: arith.select(p, t, f), self, on_false, ) @classmethod def build( cls, shape: tuple[int, ...], layout: FragmentedLayout, fn: Callable[..., ir.Value], # ir.Value varargs, one for each dim *, is_signed: bool | None = None, ) -> FragmentedArray: undef = llvm.mlir_undef(ir.IntegerType.get_signless(32)) dummy = cls.splat(undef, shape, layout, is_signed=False) return dummy.foreach( lambda _, idx: fn(*idx), create_array=True, is_signed=is_signed ) def foreach( self, fn: Callable[[ir.Value, tuple[ir.Value, ...]], ir.Value | None], *, create_array=False, is_signed=None, ): """Call a function for each value and index.""" if self.is_unreduced: raise NotImplementedError index = ir.IndexType.get() new_regs = None orig_fn = fn del fn def wrapped_fn(*args): nonlocal new_regs result = orig_fn(*args) old_reg_type = self.registers.flat[0].type # Lazily create new_regs once we know the desired output type. if create_array and new_regs is None: assert result is not None if isinstance(old_reg_type, ir.VectorType): new_reg_type = ir.VectorType.get(old_reg_type.shape, result.type) else: new_reg_type = result.type new_regs = np.full_like(self.registers, llvm.mlir_undef(new_reg_type)) return result for mlir_idx, reg_idx in zip(self.layout.thread_idxs(self.shape), np.ndindex(self.registers.shape), strict=True): reg = cast(ir.Value, self.registers[reg_idx]) assert len(mlir_idx) == len(self.shape), (mlir_idx, self.shape) if isinstance(reg.type, ir.VectorType): [elems] = ir.VectorType(reg.type).shape for i in range(elems): c_i = c(i, index) val = wrapped_fn( vector.extract( reg, dynamic_position=[], static_position=ir.DenseI64ArrayAttr.get([i]), ), (*mlir_idx[:-1], arith.addi(mlir_idx[-1], c_i)), ) if create_array: assert new_regs is not None new_regs[reg_idx] = vector.insert( val, new_regs[reg_idx], dynamic_position=[], static_position=ir.DenseI64ArrayAttr.get([i]), ) else: val = wrapped_fn(reg, mlir_idx) if create_array: assert new_regs is not None new_regs[reg_idx] = val if create_array: assert new_regs is not None return FragmentedArray(_registers=new_regs, _layout=self.layout, _is_signed=is_signed) def debug_print(self, fmt: str) -> None: idx_fmt = ", ".join(["{}"] * len(self.shape)) @self.foreach def _(val, idx): fmt_str = fmt.format(f"[{idx_fmt}]: {{}}") utils.debug_print(fmt_str, *idx, val, uniform=False) def store_untiled( self, ref: ir.Value | utils.MultimemRef, *, swizzle: int = 16, optimized: bool = True, atomic: Literal["add", "max", "min", "and", "or", "xor"] | None = None, ) -> None: if self.is_unreduced: raise ValueError("Cannot store unreduced FragmentedArray") index = ir.IndexType.get() i64 = ir.IntegerType.get_signless(64) if not isinstance(ref.type, ir.MemRefType): raise ValueError(ref) match self.layout: case WGSplatFragLayout(): if isinstance(ref, utils.MultimemRef): raise NotImplementedError("Splat layout does not support multimem") if atomic is not None: raise NotImplementedError( "Atomic stores not supported for splat layout" ) # All values are the same so swizzle does not affect anything here. self._store_untiled_splat(ref) case WGStridedFragLayout(): if swizzle != 16: raise ValueError("Only TiledLayouts support swizzling") assert isinstance(self.layout, WGStridedFragLayout) vec_size = self.layout.vec_size bitwidth = utils.bitwidth(self.mlir_dtype) total_bits = vec_size * bitwidth if total_bits % 8 != 0: raise NotImplementedError("Vector length should be a multiple of byte size") # pyrefly: ignore[bad-argument-type] for get, _update, transfer_ref, idx in self.transfer_strided(ref, vec_size): if isinstance(transfer_ref, utils.MultimemRef): ptr = utils.memref_ptr(utils.memref_slice(transfer_ref.ref, tuple(idx))) if atomic is not None: self._store_register_atomic( ptr, get(self.registers), atomic, is_smem=False, multimem=True, ) else: utils.multimem_store(ptr, get(self.registers)) elif atomic is not None: is_smem = utils.is_smem_ref(transfer_ref) base_ptr = utils.memref_ptr( utils.memref_slice(transfer_ref, tuple(idx)) ) self._store_register_atomic( base_ptr, get(self.registers), atomic, is_smem, ) elif bitwidth >= 8: # vector.store lowering doesn't work for element types with < 8 bits vector.store(get(self.registers), transfer_ref, idx) else: int_ty = ir.IntegerType.get_signless(bitwidth) int_vec_ty = ir.VectorType.get((vec_size,), int_ty) base_ptr = utils.memref_ptr(transfer_ref) strides, _ = transfer_ref.type.get_strides_and_offset() elem_flat_idx = utils.dyn_dot( idx, [arith.constant(index, s) for s in strides] ) vec_flat_idx = arith.divui( elem_flat_idx, arith.constant(index, vec_size) ) vec_flat_idx = arith.index_cast(i64, vec_flat_idx) ptr = utils.getelementptr(base_ptr, [vec_flat_idx], int_vec_ty) stored_val = get(self.registers) if not isinstance(self.mlir_dtype, ir.IntegerType): stored_val = vector.bitcast(int_vec_ty, stored_val) llvm.store(stored_val, ptr) case TiledLayout(): ref_shape = ir.MemRefType(ref.type).shape ref = utils.memref_reshape(ref, (*(1 for _ in ref_shape), *ref_shape)) self.store_tiled(ref, swizzle=swizzle, optimized=optimized, atomic=atomic) case _: raise NotImplementedError(self.layout) @classmethod def load_reduce_untiled( cls, ref: utils.MultimemRef, layout: TiledLayout | WGStridedFragLayout, reduction: utils.MultimemReductionOp, swizzle: int = 16, is_signed: bool | None = None, ): ref_ty = ir.MemRefType(ref.type) shape = tuple(ref_ty.shape) if isinstance(layout, WGStridedFragLayout): if swizzle != 16: raise ValueError("Only TiledLayouts support swizzling") registers = np.empty(layout.registers_shape(shape), dtype=object) vec_ty = ir.VectorType.get((layout.vec_size,), ref_ty.element_type) # pyrefly: ignore[bad-argument-type] for _get, update, ref, idx in cls.transfer_strided(ref, layout.vec_size): ptr = utils.memref_ptr(utils.memref_slice(ref.ref, tuple(idx))) update(registers, utils.multimem_load_reduce(vec_ty, ptr, reduction, is_signed)) return cls(_registers=registers, _layout=layout, _is_signed=is_signed) ref = utils.memref_reshape(ref, (*(1 for _ in shape), *shape)) return cls.load_tiled( ref.ref, swizzle=swizzle, is_signed=is_signed, layout=layout, optimized=False, # multimem refs are always GMEM refs. _load_fun=functools.partial( utils.multimem_load_reduce, reduction=reduction, is_signed=is_signed ), # multimem_load_reduce supports vectors of narrow floats, so we don't # need to do any casting. _narrow_float_as_int=False, ) @classmethod def load_untiled( cls, ref: ir.Value, *, layout: TiledLayout, swizzle: int = 16, is_signed: bool | None = None, optimized: bool = True, ) -> FragmentedArray: ref_ty = ir.MemRefType(ref.type) ref = utils.memref_reshape(ref, (*(1 for _ in ref_ty.shape), *ref_ty.shape)) return cls.load_tiled( ref, swizzle=swizzle, is_signed=is_signed, layout=layout, optimized=optimized ) def _store_untiled_splat(self, ref: ir.Value): if math.prod(self.shape) == 1: c0 = c(0, ir.IndexType.get()) memref.store( self.registers.flat[0], ref, [c0] * len(ir.MemRefType(ref.type).shape) ) return vec_size = 64 // mgpu.bitwidth(self.mlir_dtype) if np.prod(self.shape) < vec_size * WARPGROUP_SIZE: vec_size = 1 if np.prod(self.shape) % WARPGROUP_SIZE * vec_size: raise NotImplementedError( "Arrays with the splat layout can only be stored when they have a" f" single element or a multiple of {WARPGROUP_SIZE} elements" ) fa = FragmentedArray.splat( self.registers.flat[0], self.shape, layout=WGStridedFragLayout(shape=self.shape, vec_size=vec_size), is_signed=self.is_signed, ) fa.store_untiled(ref) def store_tiled_async( self, ref: ir.Value, barrier: utils.BarrierRef, cluster_dim: gpu.Dimension, cluster_idx: ir.Value, swizzle: int | None, optimized: bool = True, tiling_rank: int | None = None, atomic: Literal["add", "max", "min", "and", "or", "xor"] | None = None, ): if self.is_unreduced: raise ValueError("Cannot store unreduced FragmentedArray") i32 = ir.IntegerType.get_signless(32) i64 = ir.IntegerType.get_signless(64) if isinstance(ref, utils.MultimemRef): raise ValueError("Multimem refs are not supported in store_tiled_async") layout, shape = self.layout, self.shape if not isinstance(layout, TiledLayout): raise NotImplementedError(self.layout) if any( isinstance(d, Replicated) for d in itertools.chain(layout.warp_dims, layout.lane_dims) ): raise NotImplementedError("Replicated dimensions are not supported") full_cluster_idx: list[ir.Value] = [ gpu.cluster_block_id(d) for d in gpu.Dimension ] full_cluster_idx[cluster_dim] = cluster_idx lin_cluster_idx = arith.index_cast( i32, utils.cluster_idx(tuple(gpu.Dimension), full_cluster_idx) ) cluster_barrier_ptr = utils.get_cluster_ptr( barrier.get_ptr(), lin_cluster_idx, generic=False ) cluster_ref = utils.get_cluster_ref( ref, cluster_dim, cluster_idx, generic=False ) stores = self.transfer_tiled( cluster_ref, swizzle, layout, shape, optimized, ref_tiling_rank=tiling_rank ) if atomic is not None: for get, _update, _idx, cluster_ptr in stores: self._store_register_atomic( cluster_ptr, get(self.registers), atomic, is_smem=False, cluster_barrier_ptr=cluster_barrier_ptr, ) return for get, _update, _idx, cluster_ptr in stores: reg = get(self.registers) reg_ty = ir.VectorType(reg.type) element_bitwidth = utils.bitwidth(reg_ty.element_type) if ( isinstance(reg_ty.element_type, ir.FloatType) and element_bitwidth <= 8 ): narrow_int = ir.IntegerType.get_signless(element_bitwidth) reg = vector.bitcast(ir.VectorType.get(reg_ty.shape, narrow_int), reg) reg_bitwidth = utils.bitwidth(reg_ty) if reg_bitwidth == 32: ptx_constraint = "r" ptx_type = "b32" reg = utils.bitcast(reg, i32) elif reg_bitwidth == 64: ptx_constraint = "l" ptx_type = "b64" reg = utils.bitcast(reg, i64) else: raise NotImplementedError(f"Unsupported register bitwidth: {reg_bitwidth}") llvm.inline_asm( ir.Type.parse("!llvm.void"), [cluster_ptr, reg, cluster_barrier_ptr], f"st.async.cluster.shared::cluster.mbarrier::complete_tx::bytes.{ptx_type} [$0], $1, [$2];", f"l,{ptx_constraint},l", has_side_effects=True, ) def _store_register_atomic( self, base_ptr: ir.Value, vreg: ir.Value, atomic: Literal["add", "max", "min", "and", "or", "xor"], is_smem: bool, multimem: bool = False, cluster_barrier_ptr: ir.Value | None = None, ): i32 = ir.IntegerType.get_signless(32) if cluster_barrier_ptr is not None: assert not is_smem and not multimem ptr_ty = llvm.PointerType.get(address_space=7) assert cluster_barrier_ptr.type == ptr_ty, cluster_barrier_ptr.type red = "red.async" scope = "cluster.shared::cluster.mbarrier::complete_tx::bytes" space = "" ptr_constraint = "l" elif multimem: assert not is_smem red = "multimem.red" scope = "sys" space = ".global" ptr_constraint = "l" else: red = "red" scope = "cta" if is_smem else "gpu" space = ".shared::cta" if is_smem else ".global" ptr_constraint = "r" if is_smem else "l" if not is_smem and base_ptr.type.address_space != 1: if base_ptr.type.address_space != 0: raise ValueError( f"base_ptr should be a generic pointer, but got {base_ptr.type.address_space}" ) base_ptr = llvm.addrspacecast( llvm.PointerType.get(address_space=1), base_ptr ) element_type = self.mlir_dtype element_bitwidth = utils.bitwidth(element_type) noftz = "" if isinstance(element_type, ir.F32Type): if cluster_barrier_ptr is not None: raise NotImplementedError("f32 not supported for async atomics") if atomic != "add": raise NotImplementedError(f"f32 only supports add atomics, got {atomic}") ptx_type = "f32" elif isinstance(element_type, ir.IntegerType) and element_bitwidth == 32: if atomic in ("and", "or", "xor"): ptx_type = "b32" else: ptx_type = "s32" if self.is_signed else "u32" elif isinstance(element_type, (ir.F16Type, ir.BF16Type)): if cluster_barrier_ptr is not None: raise NotImplementedError("f16/bf16 not supported for async atomics") if atomic not in ("add", "min", "max"): raise NotImplementedError( f"f16/bf16 only supports add, min, max atomics, got {atomic}" ) if (is_smem or multimem) and atomic != "add": raise NotImplementedError( f"f16/bf16 SMEM/multimem atomics only support add, got {atomic}" ) ptx_type = f"{element_type}x2" noftz = "" if multimem else ".noftz" else: raise NotImplementedError( f"Unsupported element type for atomic stores: {element_type}" ) [vec_len] = vreg.type.shape if element_bitwidth == 16: if vec_len % 2 != 0: raise NotImplementedError( f"f16/bf16 atomic stores require even vector length," f" got {vec_len}" ) i32_vec_len = vec_len * element_bitwidth // 32 vreg = utils.bitcast(vreg, ir.VectorType.get( (i32_vec_len,), i32, )) regs = [ llvm.extractelement(vreg, arith.constant(i32, i)) for i in range(i32_vec_len) ] width = 1 if not is_smem and element_bitwidth == 16: for vec_width in (4, 2): if i32_vec_len % vec_width == 0: width = vec_width break for start in range(0, i32_vec_len, width): regs_slice = regs[start:start + width] ptr = utils.getelementptr(base_ptr, [start], i32) if element_bitwidth == 16 and not is_smem: # Annoyingly, the global atomics don't support v1, so we have to # use unpacked registers with .v2 and hope that ptxas optimizes it. if width == 1: i16 = ir.IntegerType.get_signless(16) vec2xi16 = ir.VectorType.get((2,), i16) [reg] = regs_slice pair = llvm.bitcast(vec2xi16, reg) llvm.inline_asm( ir.Type.parse("!llvm.void"), [ ptr, llvm.extractelement(pair, arith.constant(i32, 0)), llvm.extractelement(pair, arith.constant(i32, 1)), ], f"{red}{space}.relaxed.{scope}.{atomic}{noftz}.v2.{element_type} [$0], {{$1, $2}};", f"{ptr_constraint},h,h", has_side_effects=True, ) else: operands = "{" + ", ".join(f"${i + 1}" for i in range(width)) + "}" llvm.inline_asm( ir.Type.parse("!llvm.void"), [ptr, *regs_slice], f"{red}{space}.relaxed.{scope}.{atomic}{noftz}.v{width}.{element_type}x2 [$0], {operands};", f"{ptr_constraint}" + ",r" * width, has_side_effects=True, ) else: ptx_t = f"{element_type}x2" if element_bitwidth == 16 else ptx_type [reg] = regs_slice operands = [ptr, reg] barrier_op = "" constraints = f"{ptr_constraint},r" if cluster_barrier_ptr is not None: operands.append(cluster_barrier_ptr) barrier_op = ", [$2]" constraints += ",l" llvm.inline_asm( ir.Type.parse("!llvm.void"), operands, f"{red}{space}.relaxed.{scope}.{atomic}{noftz}.{ptx_t} [$0], $1{barrier_op};", constraints, has_side_effects=True, ) def store_tiled( self, ref: ir.Value | utils.MultimemRef, swizzle: int | None, optimized: bool = True, tiling_rank: int | None = None, atomic: Literal["add", "max", "min", "and", "or", "xor"] | None = None, ): if self.is_unreduced: raise ValueError("Cannot store unreduced FragmentedArray") i32 = ir.IntegerType.get_signless(32) if not isinstance(self.layout, TiledLayout): raise NotImplementedError(self.layout) layout, shape = self.layout, self.shape if atomic is not None: if any(isinstance(d, Replicated) for d in layout.warp_dims + layout.lane_dims): raise NotImplementedError( "Atomic stores not supported for layouts with replicated dims" ) multimem = isinstance(ref, utils.MultimemRef) is_smem = not multimem and utils.is_smem_ref(ref) stores = self.transfer_tiled( ref.ref if multimem else ref, swizzle, layout, shape, optimized, ref_tiling_rank=tiling_rank, ) for get, _update, _idx, base_ptr in stores: self._store_register_atomic( base_ptr, get(self.registers), atomic, is_smem, multimem=multimem, ) return # Note that the loop below will "race" for layouts that replicate data. # However, in that case all of the racing writes store the same data, which # is ok in the CUDA memory model. if isinstance(ref, utils.MultimemRef): stores = self.transfer_tiled( ref.ref, swizzle, layout, shape, optimized, ref_tiling_rank=tiling_rank ) for get, _update, _idx, ptr in stores: utils.multimem_store(ptr, get(self.registers)) else: try: if utils.get_arch().major < 9: raise TxMatrixIneligible("TxMatrix not supported on pre-Hopper GPUs") stores = self.transfer_tiled( ref, swizzle, layout, shape, optimized, ref_tiling_rank=tiling_rank, use_txmatrix=True, ) bitwidth = utils.bitwidth(self.mlir_dtype) for gets, _updates, _idxs, ptr, tx_layout in stores: if tx_layout == nvvm.MMALayout.col and bitwidth == 8: s = 16 elt_type = nvvm.LdStMatrixEltType.B8 else: s = 8 elt_type = nvvm.LdStMatrixEltType.B16 nvvm.stmatrix( ptr, [utils.bitcast(get(self.registers), i32) for get in gets], tx_layout, ir.Attribute.parse(f"#nvvm.ld_st_matrix_shape"), elt_type, ) return except TxMatrixIneligible: pass stores = self.transfer_tiled( ref, swizzle, layout, shape, optimized, ref_tiling_rank=tiling_rank ) for get, _update, _idx, ptr in stores: reg = get(self.registers) reg_ty = ir.VectorType(reg.type) element_bitwidth = utils.bitwidth(reg_ty.element_type) if ( isinstance(reg_ty.element_type, ir.FloatType) and element_bitwidth <= 8 ): narrow_int = ir.IntegerType.get_signless(element_bitwidth) reg = vector.bitcast(ir.VectorType.get(reg_ty.shape, narrow_int), reg) llvm.store(reg, ptr) @classmethod def load_tiled( cls, ref, swizzle: int | None, *, is_signed: bool | None = None, layout: FragmentedLayout = WGMMA_LAYOUT, optimized: bool = True, tiling_rank: int | None = None, _load_fun: Callable[[ir.VectorType, ir.Value], ir.Value] = llvm.load, _narrow_float_as_int: bool = True, ) -> FragmentedArray: i32 = ir.IntegerType.get_signless(32) if not isinstance(layout, TiledLayout): raise NotImplementedError(layout) if layout.has_unreduced_dims: raise ValueError("Cannot load into a layout with unreduced dims") ref_ty = ir.MemRefType(ref.type) dtype = ref_ty.element_type tiled_shape = ref_ty.shape if tiling_rank is None: if len(tiled_shape) % 2: raise ValueError("Tiled reference must have even rank") if len(tiled_shape) < 2: raise ValueError("Tiled reference must have at least two dimensions") tiling_rank = len(tiled_shape) // 2 else: if tiling_rank > len(tiled_shape) // 2: raise ValueError( f"Tiling rank for reference of shape {tiled_shape} must be at most" f" {len(tiled_shape) // 2}" ) if not tiling_rank: raise ValueError( "Tiling rank for reference of shape {tiled_shape} must be at least" " 1" ) tiling = Tiling((tuple(tiled_shape[-tiling_rank:]),)) shape = tiling.untile_shape(tuple(tiled_shape)) reg_ty = ir.VectorType.get((layout.vector_length,), dtype) zero = vector.broadcast(reg_ty, c(0, dtype)) registers = np.full(layout.registers_shape(shape), zero, dtype=object) bitwidth = utils.bitwidth(dtype) is_narrow_float = ( isinstance(dtype, ir.FloatType) and bitwidth <= 8 ) narrow_int = ir.IntegerType.get_signless(bitwidth) # Narrow floats are not supported by LLVM, so we need to transfer them as # narrow ints and bitcast back to the desired type. transfer_ty = ir.VectorType.get( (layout.vector_length,), narrow_int if is_narrow_float and _narrow_float_as_int else dtype ) if _load_fun is llvm.load: loads = cls.transfer_tiled( ref, swizzle, layout, shape, optimized, ref_tiling_rank=tiling_rank, use_txmatrix=True ) try: for _gets, updates, _idxs, ptr, tx_layout in loads: if tx_layout == nvvm.MMALayout.col and bitwidth == 8: num_scale = 2 elt_type = nvvm.LdStMatrixEltType.B8 else: num_scale = 1 elt_type = nvvm.LdStMatrixEltType.B16 s = 8 * num_scale loaded_regs_value = nvvm.ldmatrix( ptr, layout=tx_layout, num=len(updates) // num_scale, shape=ir.Attribute.parse(f"#nvvm.ld_st_matrix_shape"), elt_type=elt_type, ) # ldmatrix returns a single i32 or a struct of i32s. if len(updates) == 1: loaded_regs = [loaded_regs_value] else: loaded_regs = [ llvm.extractvalue(i32, loaded_regs_value, [i]) for i in range(len(updates)) ] for loaded_reg, update in zip(loaded_regs, updates, strict=True): # type: ignore update(registers, utils.bitcast(loaded_reg, reg_ty)) return cls(_registers=registers, _layout=layout, _is_signed=is_signed) except TxMatrixIneligible: pass loads = cls.transfer_tiled( ref, swizzle, layout, shape, optimized, ref_tiling_rank=tiling_rank ) for _get, update, _idx, ptr in loads: loaded_reg = _load_fun(transfer_ty, ptr) if is_narrow_float and _narrow_float_as_int: loaded_reg = vector.bitcast(reg_ty, loaded_reg) update(registers, loaded_reg) return cls(_registers=registers, _layout=layout, _is_signed=is_signed) @classmethod def transfer_strided(self, ref: ir.Value, vec_size: int): ref_ty = ir.MemRefType(ref.type) layout = WGStridedFragLayout(shape=tuple(ref_ty.shape), vec_size=vec_size) try: # Flattening the reference potentially produces simpler PTX but # if the ref is not already 1D and has strided dimensions # flattening won't work. ref = mgpu.memref_fold(ref, 0, len(ref_ty.shape)) except ValueError: if vec_size > 1: ref_ty = ir.MemRefType(ref.type) shape = ref_ty.shape strides, _ = ref_ty.get_strides_and_offset() # Try to fold contiguous dimension pairs. for i in reversed(range(len(shape) - 1)): if strides[i] == shape[i+1] * strides[i+1]: ref = mgpu.memref_fold(ref, i, 2) ref_ty = ir.MemRefType(ref.type) shape = ref_ty.shape strides, _ = ref_ty.get_strides_and_offset() has_contiguous_dim = False for size, stride in zip(shape, strides): if stride == 1: has_contiguous_dim = True if size % vec_size != 0: raise ValueError( "The contiguous dimension of the reference must be a" f" multiple of the layout's vector size (got {size} and" f" vector size {vec_size})" ) from None elif size > 1: if stride % vec_size != 0: raise ValueError( "Non-contiguous dimension of the reference must have strides" " that are multiples of the layout's vector size (got" f" {stride} and vector size {vec_size})" ) from None if not has_contiguous_dim: raise ValueError( "The reference must have a contiguous dimension when vec_size > 1" ) layout = WGStridedFragLayout(shape=tuple(ref_ty.shape), vec_size=vec_size) idx_gen = layout.thread_idxs(tuple(ref_ty.shape)) else: idx_gen = map(lambda x: [x], layout.linear_thread_idxs()) for i, vec_idx in enumerate(idx_gen): def update(registers, reg, _i=i): registers[_i] = reg def get(registers, _i=i): return registers[_i] yield get, update, ref, vec_idx @staticmethod def transfer_tiled( ref: ir.Value, swizzle: int | None, layout: TiledLayout, shape: tuple[int, ...], optimized: bool = True, ref_tiling_rank: int | None = None, use_txmatrix: bool = False, ): """Generate a transfer schedule for a tiled layout. Given a ref with one level tiling applied to it (we assume all dimensions have been tiled), this function generates an iterable describing a good schedule for swizzled SMEM loads/stores. At each step, the iterable yields a tuple of three values: * a function that takes a register array and returns the register to be stored at the current address * a function that takes a register array and a register loaded from the current address, and updates the register array with that register * the current address for load/store instructions """ c = lambda x: arith.constant(ir.IntegerType.get_signless(32), x) i32 = ir.IntegerType.get_signless(32) tiling = layout.tiling ref_ty = ir.MemRefType(ref.type) dtype = ref_ty.element_type if ref_tiling_rank is None: if len(ref_ty.shape) % 2: raise ValueError("Tiled reference must have an even rank when its tiling rank is not specified") ref_tiling_rank = ref_ty.rank // 2 if ref_tiling_rank > len(ref_ty.shape) // 2: raise ValueError( f"Tiling rank for reference of shape {ref_ty.shape} must be at most" f" {len(ref_ty.shape) // 2}" ) assert ref_tiling_rank and ref_ty.rank > ref_tiling_rank ref_logical_rank = ref_ty.rank - ref_tiling_rank ref_tiling_shape = tuple(ref_ty.shape[ref_logical_rank:]) ref_tiling = Tiling((ref_tiling_shape,)) ref_strides, _ = ref_ty.get_strides_and_offset() if (ref_logical_shape := ref_tiling.untile_shape(tuple(ref_ty.shape))) != shape: raise ValueError( f"The reference has untiled shape of {ref_logical_shape} while the" f" register array has shape {shape}" ) first_tiled_dim = ref_logical_rank - ref_tiling_rank nested_ref_shape = tuple( (ref_ty.shape[i], ref_ty.shape[i + ref_tiling_rank]) if i >= first_tiled_dim and ref_ty.shape[i + ref_tiling_rank] != 1 else (ref_ty.shape[i],) for i in range(ref_logical_rank) ) nested_ref_strides = tuple( (ref_strides[i], ref_strides[i + ref_tiling_rank]) if i >= first_tiled_dim and ref_ty.shape[i + ref_tiling_rank] != 1 else (ref_strides[i],) for i in range(ref_logical_rank) ) tiled_nested_shape, tiled_nested_strides = tiling.tile_nested_shape_strides( nested_ref_shape, nested_ref_strides ) if use_txmatrix: minor_lane_dim = layout.lane_dims[-1] major_lane_dim = layout.lane_dims[0] is_txmatrix_reg_layout = ( utils.bitwidth(dtype) * layout.vector_length == 32 and isinstance(minor_lane_dim, int) and math.prod(tiled_nested_shape[minor_lane_dim]) % 4 == 0 ) is_row_txmatrix_mem_layout = ( is_txmatrix_reg_layout # We need the data of each group of 4 consecutive lanes to map to a # contiguous 16-byte chunk of memory. and tiled_nested_shape[minor_lane_dim][-1] % 4 == 0 and tiled_nested_strides[minor_lane_dim][-1] == layout.vector_length # The stride along vector_dim will be checked below. # Strides along other lane dims are irrelevant. ) is_col_txmatrix_mem_layout = ( is_txmatrix_reg_layout and utils.bitwidth(dtype) in {8, 16} and isinstance(major_lane_dim, int) and len(layout.lane_dims) == 2 and tiled_nested_shape[major_lane_dim] == (8,) # At this point, considering is_txmatrix_reg_layout, we know that lane # dims represent a 8x4 matrix of vectors. and tiled_nested_strides[major_lane_dim] == (1,) ) tx_pair_tiled_dim = None # For 8-bit types, each tile is actually 8x16. However, the only shape # supported by the transposed txmatrix is 16x16. This means that each # transfer actually deals with two tiles and we need to find the tiled # dimension that iterates over tile pairs. One side of the matrix is then # formed by pair_tile_dim and major_lane_dim, while the other is formed by # minor_lane_dim and vector_dim. if is_col_txmatrix_mem_layout and utils.bitwidth(dtype) == 8: if utils.get_arch().major < 10: raise TxMatrixIneligible( "Transposed 8-bit loads and stores are not supported on" " pre-Blackwell architectures" ) shape_strides = list(zip(tiled_nested_shape, tiled_nested_strides)) for tx_pair_tiled_dim, (shapes, strides) in enumerate_negative(shape_strides): if ( tx_pair_tiled_dim in layout.lane_dims or tx_pair_tiled_dim in layout.warp_dims or tx_pair_tiled_dim == layout.vector_dim ): continue if strides[-1] == 8 and shapes[-1] % 2 == 0: break else: # If we don't find a dimension that is a contiguous extension of the # major_lane_dim in the ref (or has an uneven length), then we can't # use txmatrix. is_col_txmatrix_mem_layout = False can_use_txmatrix = ( is_row_txmatrix_mem_layout or is_col_txmatrix_mem_layout ) and utils.is_smem_ref(ref) if not can_use_txmatrix: raise TxMatrixIneligible("Cannot use txmatrix for this layout") load_vector_dim = ( layout.vector_dim if is_row_txmatrix_mem_layout else major_lane_dim ) assert isinstance(load_vector_dim, int) tx_layout = nvvm.MMALayout.row if is_row_txmatrix_mem_layout else nvvm.MMALayout.col # We remap pair_tile_dim to a valid index in tiles_shape (computed below) if tx_pair_tiled_dim is not None: # We need tx_pair_nested_tiled_dim to be positive, since that's what the # lane quadrant logic expects. if tx_pair_tiled_dim == -1: tx_pair_nested_tiled_dim = sum(map(len, tiled_nested_shape)) - 1 else: tx_pair_nested_tiled_dim = sum(len(s) for s in tiled_nested_shape[:tx_pair_tiled_dim + 1]) - 1 else: tx_pair_tiled_dim = tx_pair_nested_tiled_dim = None else: load_vector_dim = layout.vector_dim tx_layout = None tx_pair_tiled_dim = tx_pair_nested_tiled_dim = None # Not sure if this is strictly required for all data types, but it certainly # is for sub-byte types (else we might not increment the pointer by whole bytes). if any( any(s % layout.vector_length and d != 1 for s, d in zip(ss, ds)) for i, (ss, ds) in enumerate_negative( list(zip(tiled_nested_strides, tiled_nested_shape)) ) if i != load_vector_dim ): raise UnsupportedTransferError( "Tiled strides must be a multiple of the vector length, except for" " the load vectorized dimension" ) if tiled_nested_strides[load_vector_dim] != (1,): raise UnsupportedTransferError( "Vectorized dimension should not require further tiling and have a" " stride of 1" ) tiles_shape = list(tiled_nested_shape) tiles_strides = list(tiled_nested_strides) for d in (*layout.partitioned_warp_dims, *layout.partitioned_lane_dims, layout.vector_dim): # We could avoid repeating the singleton dimensions, but it simplifies the # code below that computes the register index for a given tile. tiles_shape[d] = (1,) * len(tiles_shape[d]) tiles_strides[d] = (0,) * len(tiles_strides[d]) tiles_shape = list(itertools.chain.from_iterable(tiles_shape)) tiles_strides = list(itertools.chain.from_iterable(tiles_strides)) vector_length = layout.vector_length element_bits = mgpu.bitwidth(dtype) if (vector_length * element_bits) % 8 != 0: raise ValueError( f"Vector length ({vector_length}) must be a multiple of bytes," f" but has {vector_length * element_bits} bits" ) transfer_bytes = (vector_length * element_bits) // 8 if swizzle not in {16, 32, 64, 128}: raise ValueError("Only swizzled transfers supported") # We will be computing the offsets in units of vectors, not elements, # to better support sub-byte types. swizzle_tile_transfers = 16 // transfer_bytes swizzle_group_transfers = 128 // transfer_bytes swizzle_groups_per_block = swizzle // 16 swizzle_block_transfers = swizzle_groups_per_block * swizzle_group_transfers if isinstance(dtype, ir.FloatType) and element_bits <= 8: narrow_int = ir.IntegerType.get_signless(element_bits) transfer_dtype = ir.VectorType.get((vector_length,), narrow_int) else: transfer_dtype = ir.VectorType.get((vector_length,), dtype) if ref_ty.memory_space is None: llvm_memory_space = None elif utils.is_smem_ref(ref_ty): llvm_memory_space = 3 elif ref_ty.memory_space == ir.IntegerAttr.get(i32, 7): # Cluster SMEM llvm_memory_space = 7 else: raise ValueError(f"Unsupported memory space: {ref_ty.memory_space}") plan = TrivialTransferPlan() if optimized and llvm_memory_space is None: plan = plan_gmem_transfer( nested_ref_shape, nested_ref_strides, layout, element_bits, ) elif optimized: if llvm_memory_space != 3 and llvm_memory_space != 7: raise NotImplementedError("Only optimized transfers to SMEM supported") mem_layout = layout if tx_layout == nvvm.MMALayout.col: major_lane_dim, _ = layout.lane_dims assert isinstance(major_lane_dim, int) if element_bits == 16: # First, find the tile that contains the major lane dim and split it major_lane_dim_tile, major_lane_dim_within_tile = layout.find_dim_in_tiling(major_lane_dim) # Now, we add a new tile that splits the 8-sized dimension into 4 and 2. # It will be inserted right after the tile containing major_lane_dim, # in a way such that after tiling the original tile containing # major_lane_dim will contain a single 4-sized dim in place of # major_lane_dim with all other elements being 1. # For example, if we had a tiling of (8, 8)(2,) with lane_dims=(-3, -2) # and vector_dim=-1, we will transform it into: # tiling=(8, 8)(2, 8)(2,) (tiled base shape of (4, 1, 2, 4, 2)). new_tile = list(layout.tiling.tiles[major_lane_dim_tile]) new_tile[major_lane_dim_within_tile] = layout.vector_length sub_layout = layout.subdivide_tile(major_lane_dim_tile, tuple(new_tile)) # We now adjust the lane and vector dimensions. The 8-sized lane dim # is now composed of the minor lane dim and the original vector dim, # the minor 4-sized lane dim corresponds to the major portion of the # 4x2 split of the original 8-sized major lane dim, with the remaining # 2 becoming the new vector dim. sub_row_major, sub_row_minor, sub_col_major = sub_layout.lane_dims assert isinstance(sub_row_minor, int) sub_col_minor = sub_layout.vector_dim mem_layout = TiledLayout( sub_layout.tiling, sub_layout.warp_dims, (sub_col_major, sub_col_minor, sub_row_major), sub_row_minor, ) elif element_bits == 8: assert tx_pair_tiled_dim is not None *_, minor_lane_dim = layout.lane_dims assert isinstance(minor_lane_dim, int) # If the order was different, we'd have to subdivide the layout in a # different order, or adjust dimension indices. if not (tx_pair_tiled_dim < major_lane_dim < minor_lane_dim): raise NotImplementedError("Unsupported layout for 8-bit txmatrix") # First, we make sure that the tx_pair dimension will be equal to 2. tx_pair_dim_tile, tx_pair_dim_within_tile = layout.find_dim_in_tiling( tx_pair_tiled_dim ) new_tile = list(layout.tiling.tiles[tx_pair_dim_tile]) new_tile[tx_pair_dim_within_tile] = 16 sub_layout = layout.subdivide_tile(tx_pair_dim_tile, tuple(new_tile)) # Now, similarly like in the 16-bit case, we partition the major dim, # only this time into 2x4 (major_lane_dim remains 2). major_lane_dim_tile, major_lane_dim_within_tile = sub_layout.find_dim_in_tiling(major_lane_dim) new_major_lane_tile = list(sub_layout.tiling.tiles[major_lane_dim_tile]) new_major_lane_tile[major_lane_dim_within_tile] = layout.vector_length sub_layout = sub_layout.subdivide_tile(major_lane_dim_tile, tuple(new_major_lane_tile)) # Finally, we partition the minor lane dim into 2x2. minor_lane_dim_tile, minor_lane_dim_within_tile = sub_layout.find_dim_in_tiling(minor_lane_dim) new_minor_lane_tile = list(sub_layout.tiling.tiles[minor_lane_dim_tile]) new_minor_lane_tile[minor_lane_dim_within_tile] = 2 * layout.vector_length sub_layout = sub_layout.subdivide_tile(minor_lane_dim_tile, tuple(new_minor_lane_tile)) # Note that we skip the majormost col dim. A 16x16 matrix requires # two SMEM wavefronts so we model sequential 16x8 access instead. sub_row_major, sub_row_minor, _, sub_col_major = sub_layout.lane_dims # 2x4x2x2 sub_col_minor = sub_layout.vector_dim # 4 assert isinstance(sub_row_minor, int) # The lane dims are adjusted automatically by subdivide_tile, but we # need to manually adjust the tx_pair dim. sub_tx_pair_tiled_dim = ( tx_pair_tiled_dim - len(new_major_lane_tile) - len(new_minor_lane_tile) ) # The transfer layout now loads 8 columns (in majormost lanes), and # then uses the remaining minor 4 lanes to load 4-element vectors # along rows. mem_layout = TiledLayout( sub_layout.tiling, sub_layout.warp_dims, (sub_col_major, sub_col_minor, sub_tx_pair_tiled_dim, sub_row_major), sub_row_minor, ) else: raise NotImplementedError( f"Unsupported element bits for txmatrix: {element_bits}" ) plan = plan_tiled_transfer( nested_ref_shape, nested_ref_strides, mem_layout, element_bits, swizzle, ) if tx_layout is not None and not isinstance(plan, TrivialTransferPlan): raise TxMatrixIneligible("txmatrix requires a trivial transfer plan") tiles_strides_transfer = [s // vector_length for s in tiles_strides] # Technically we should keep the vector_dim stride set to 1, but its shape # is 1 so it does not matter. dyn_tiled_strides = [ c(s // vector_length) for s in itertools.chain.from_iterable( tiled_nested_strides[-layout.tiled_tiling_rank :] ) ] # This expands a tiled index into a finer-grained index that accounts for # the fact that some tiled dims are tiled further in the nested shape. def expand_nested_dims(idxs: Sequence[ir.Value]) -> list[ir.Value]: assert len(idxs) == layout.tiled_tiling_rank new_idxs = [] for idx, dim_shape in zip(idxs, tiled_nested_shape[-layout.tiled_tiling_rank :]): if dim_shape == (1,): new_idxs.append(idx) continue dim_strides = utils.get_contiguous_strides(dim_shape) for i, (size, stride) in enumerate(zip(dim_shape, dim_strides)): new_idx = arith.divui(idx, c(stride)) if i != 0: # No need to apply rem to the first dim. new_idx = arith.remui(new_idx, c(size)) new_idxs.append(new_idx) assert len(new_idxs) == sum(map(len, tiled_nested_shape[-layout.tiled_tiling_rank :])) return new_idxs # All offsets are in units of transfer_dtype. offset_lane_idx = None if tx_layout == nvvm.MMALayout.col: minor_lane_dim = layout.lane_dims[-1] assert isinstance(minor_lane_dim, int) col_nested_shape = (*tiled_nested_shape[minor_lane_dim], *tiled_nested_shape[layout.vector_dim]) col_nested_strides = (*tiled_nested_strides[minor_lane_dim], *tiled_nested_strides[layout.vector_dim]) lane_offset = c(0) remaining_lane_idx = arith.remui(utils.thread_idx(), c(math.prod(col_nested_shape))) for size, stride in reversed(list(zip(col_nested_shape, col_nested_strides))): lane_offset = arith.addi( lane_offset, arith.muli(arith.remui(remaining_lane_idx, c(size)), c(stride // vector_length)), ) remaining_lane_idx = arith.divui(remaining_lane_idx, c(size)) else: if tx_layout == nvvm.MMALayout.row: offset_lane_idx = arith.muli(arith.remui(utils.thread_idx(), c(8)), c(4)) lane_offset = utils.dyn_dot( expand_nested_dims(layout.lane_indices(offset_lane_idx)), dyn_tiled_strides, ) warp_offset = utils.dyn_dot(expand_nested_dims(layout.warp_indices()), dyn_tiled_strides) dyn_offset = arith.addi(lane_offset, warp_offset) ptr = utils.memref_ptr(ref) _as_consts = lambda consts: [c(const) for const in consts.tolist()] # This has bits set only for the offset bits that influence swizzling. swizzle_mask = swizzle_block_transfers - swizzle_tile_transfers def get_tile_transfer(tile_idx): indices = np.asarray([f(tile_idx) for f in plan.tile_index_transforms]) const_offset = np.dot(indices, tiles_strides_transfer) # We split the offset into a part that interacts with swizzling and a # part that doesn't. This lets us generate better code because constant # offsets can be fused into load and store instructions. const_offset_swizzle = const_offset & swizzle_mask const_offset_no_swizzle = const_offset - const_offset_swizzle offset_pre_swizzle = arith.addi( dyn_offset, plan.select(_as_consts(const_offset_swizzle)) ) swizzle_group = arith.remui( arith.divui(offset_pre_swizzle, c(swizzle_group_transfers)), c(swizzle_groups_per_block), ) swizzle_bits = arith.muli(swizzle_group, c(swizzle_tile_transfers)) offset = arith.xori(offset_pre_swizzle, swizzle_bits) reg_ptr = utils.getelementptr(ptr, [offset], transfer_dtype) offset_no_swizzle = plan.select(_as_consts(const_offset_no_swizzle)) reg_ptr = utils.getelementptr(reg_ptr, [offset_no_swizzle], transfer_dtype) # Here, registers are organized in an array with shape obtained by tiling # the logical data bounds. But, the reference was tiled and so each # logical tiled dimension can map to multiple dims in tiled_shape. # The transform below maps this potentially higher-rank representation # back to the lower-rank representation used by the register arrays. def mem_idx_to_reg_idx(idx): reg_tiled_idx = [] base_idx = 0 for dim_shape in tiled_nested_shape: dim_strides = utils.get_contiguous_strides(dim_shape) dim_idxs = idx[base_idx:base_idx + len(dim_shape)] base_idx += len(dim_shape) reg_tiled_idx.append(sum(i * s for i, s in zip(dim_idxs, dim_strides))) return tuple(reg_tiled_idx) reg_idxs = [mem_idx_to_reg_idx(idx) for idx in indices.tolist()] def get_register(regs, reg_idxs=reg_idxs): # f8 data types are not handled by the LLVM dialect, so we need to # transfer them as i8 and bitcast them back to f8. return plan.select([regs[reg_idx] for reg_idx in reg_idxs]) def update_registers(regs, new, reg_idxs=reg_idxs): # TODO(apaszke): If the staggering forms a permutation with a small # cycle length, then instead of blending at each step we could construct # a small routing network (kind of like a sorting network) to fix up # each cycle separately after all the loads are performed. # This would be especially useful for dims that are powers of two and # staggered by another power of 2, since all cycles are of length 2 (and # we could save half the selects). for i, reg_idx in enumerate(reg_idxs): regs[reg_idx] = plan.select_if_group(i, regs[reg_idx], new) def get_base_index(): if not isinstance(plan, TrivialTransferPlan): raise NotImplementedError( "Base index computation only supported for trivial transfer plans" ) if any(len(t) != 1 for t in tiled_nested_shape): raise NotImplementedError("Tiling too complicated") return tiling.untile_indices(indices.tolist()[0]) return get_register, update_registers, get_base_index, reg_ptr if not use_txmatrix: yield from map(get_tile_transfer, np.ndindex(*tiles_shape)) return # Below we implement a partitoning scheme for the iteration space that # allows us to use as many high-num transfers as possible. # We use two properties of iteration spaces: # 1. Factorization, where e.g. we can take a 2x3 space and turn it into # 2*(1x3) spaces. # 2. Splitting, where e.g. the 2x3 space can be rewritten as 2x(2+1), which # is equivalent to taking a sum of two spaces: 2x2 + 2x1. # We factorize the grid, trying to harvest factors of 2 to get to 4 if # possible. If we can't find enough, we try applying the splitting rule to # split the iteration space into a large even part and a smaller remainder. # # Consider a 2x3 space. The best transfer we could probably derive is to # first perform factorization to get 2*(1x3). At this point we can't # factorize any further, so we split that into 2*(1x2 + 1x1). Using # distributivity we get 2*(1x2) + 2*(1x1) which can be further factored # into 4*(1x1) + 2*(1x1). This gives us a good schedule of a single num=4 # transfer and a single num=2 transfer, which is ideal given that we had 6 # tiles overall. # # The scheme outlined above is fine for most layouts. However, if the loaded # arrays are used to feed MMA instructions directly, we need to ensure that # the order of registers produced by ldmatrix matches up the order in which # the MMA instructions consume them. Otherwise, ptxas emits a bunch of extra # MOV instructions that destroy the performance. from .mma import MMALayouts if any(layout == MMALayouts(dtype, m_warps=w).lhs for w in (1, 2, 4)): assert tiled_nested_shape[-6] == (2,) assert tiled_nested_shape[-5] == (2,) major_dim = len(tiles_shape) - sum(len(s) for s in tiled_nested_shape[-6:]) minor_dim = major_dim + 1 factored_quadrant_dims = [(major_dim, 2), (minor_dim, 2)] elif any(layout == MMALayouts(dtype, m_warps=w).rhs for w in (1, 2, 4)): assert tiled_nested_shape[-6] == (2,) dim = len(tiles_shape) - sum(len(s) for s in tiled_nested_shape[-6:]) factored_quadrant_dims = [(dim, 2)] elif tx_pair_nested_tiled_dim: # For 8-bit values, each tile is 8x16. But, the transposed instruction # uses 16x16 tiles, which is equivalent to pairing up two 8x16 tiles along # tx_pair_nested_tiled_dim, so the transfer implicitly already has a num of 2. # XXX: We must keep tx_pair_nested_tiled_dim minor! if tiles_shape[tx_pair_nested_tiled_dim] % 4 == 0: factored_quadrant_dims = [(tx_pair_nested_tiled_dim, 4)] else: assert tiles_shape[tx_pair_nested_tiled_dim] % 2 == 0 candidate_dim = next( ( quadrant_dim for quadrant_dim, d in enumerate(tiles_shape) if d % 2 == 0 and quadrant_dim != tx_pair_nested_tiled_dim ), None, ) if candidate_dim is not None: factored_quadrant_dims = [ (candidate_dim, 2), (tx_pair_nested_tiled_dim, 2), ] else: factored_quadrant_dims = [(tx_pair_nested_tiled_dim, 2)] else: for quadrant_dim, d in enumerate(tiles_shape): if d % 4 == 0: factored_quadrant_dims = [(quadrant_dim, 4)] break else: factored_quadrant_dims = [ (quadrant_dim, 2) for quadrant_dim, d in enumerate(tiles_shape) if d % 2 == 0 ][:2] @dataclasses.dataclass(frozen=True) class TxMatrixTransfer: quadrant_dims: tuple[tuple[int, int], ...] tile_groups_shape: tuple[int, ...] tile_groups_offset: tuple[int, ...] @property def num(self) -> int: return math.prod(d[1] for d in self.quadrant_dims) for dim, size in factored_quadrant_dims: assert tiles_shape[dim] % size == 0 tiles_shape[dim] //= size offsets = [0] * len(tiles_shape) if (factored_num := math.prod(d[1] for d in factored_quadrant_dims)) == 4: transfers = [ TxMatrixTransfer( tuple(factored_quadrant_dims), tuple(tiles_shape), tuple(offsets) ) ] else: transfers = [] missing_factors = (2,) if factored_num == 2 else (4, 2) for factor in missing_factors: for i, size in enumerate(tiles_shape): if size > factor: if any(d[0] == i for d in factored_quadrant_dims): # More than 1 factor => we have num=4 and we wouldn't be here. assert factored_quadrant_dims == [(i, 2)] tx_quadrant_dims = ((i, 4),) total_dim_factor = 4 else: tx_quadrant_dims = ((i, factor), *factored_quadrant_dims) total_dim_factor = factor dim_steps = size // factor transfers.append( TxMatrixTransfer( tx_quadrant_dims, (*tiles_shape[:i], dim_steps, *tiles_shape[i + 1 :]), tuple(offsets), ) ) offsets[i] += dim_steps * total_dim_factor tiles_shape[i] %= factor transfers.append( TxMatrixTransfer( tuple(factored_quadrant_dims), tuple(tiles_shape), tuple(offsets), ) ) lane_quadrant = arith.remui(arith.divui(utils.thread_idx(), c(WARP_SIZE // 4)), c(4)) lane_half = arith.remui(arith.divui(utils.thread_idx(), c(WARP_SIZE // 2)), c(2)) base_dyn_offset = dyn_offset for tx in transfers: assert tx.num in (1, 2, 4) lane_tile_offset = arith.constant(i32, 0) if tx_pair_nested_tiled_dim: if len(tx.quadrant_dims) == 1: assert tx.quadrant_dims[0][0] == tx_pair_nested_tiled_dim if tx.quadrant_dims[0][1] == 4: lane_tile_offset = arith.muli(lane_half, c(tiles_strides_transfer[tx_pair_nested_tiled_dim] * 2)) else: assert len(tx.quadrant_dims) == 2 assert tx.quadrant_dims[-1] == (tx_pair_nested_tiled_dim, 2), tx.quadrant_dims dim = tx.quadrant_dims[-2][0] assert dim != tx_pair_nested_tiled_dim lane_tile_offset = arith.muli(lane_half, c(tiles_strides_transfer[dim])) else: lane_quadrant_remaining = lane_quadrant for dim, size in tx.quadrant_dims[::-1]: idx = arith.remui(lane_quadrant_remaining, c(size)) lane_tile_offset = arith.addi(lane_tile_offset, arith.muli(idx, c(tiles_strides_transfer[dim]))) lane_quadrant_remaining = arith.divui(lane_quadrant_remaining, c(size)) def get_tile_idx(tile_group_idx, num_i): tile_group_idx = list(tile_group_idx) for dim, size in tx.quadrant_dims[::-1]: tile_group_idx[dim] = size * tile_group_idx[dim] + (num_i % size) num_i //= size return tuple(i + o for i, o in zip(tile_group_idx, tx.tile_groups_offset)) # Note that this will affect the call to get_tile_transfer below. dyn_offset = arith.addi(base_dyn_offset, lane_tile_offset) for tile_group_idx in np.ndindex(*tx.tile_groups_shape): reg_transfers = [] for i in range(tx.num): reg_transfers.append(get_tile_transfer(get_tile_idx(tile_group_idx, i))) transfers_t = list(zip(*reg_transfers)) # The first address is the one our quadrant is responsible for providing. yield (*transfers_t[:-1], transfers_t[-1][0], tx_layout) def tree_flatten(self): aux = self.layout, self.registers.shape, self.is_signed return list(self.registers.flat), aux @classmethod def tree_unflatten(cls, aux, flat_registers): layout, reg_shape, is_signed = aux registers = np.asarray(flat_registers, dtype=object).reshape(reg_shape) return cls(_registers=registers, _layout=layout, _is_signed=is_signed) class TxMatrixIneligible(Exception): pass class UnsupportedTransferError(Exception): pass IndexTransform: TypeAlias = Callable[[tuple[int, ...]], tuple[int, ...]] @runtime_checkable class TransferPlan(Protocol): @property def tile_index_transforms(self) -> tuple[IndexTransform, ...]: raise NotImplementedError def select(self, group_elems: Sequence[ir.Value], /) -> ir.Value: """Selects the value corresponding to the group of the current thread. The argument must be of the same length as tile_index_transforms. """ raise NotImplementedError def select_if_group(self, group_idx: int, old: ir.Value, new: ir.Value, /) -> ir.Value: """Returns `new` if the current thread belongs to the given group and `old` otherwise. group_idx must be between 0 and len(tile_index_transforms) - 1. """ raise NotImplementedError @dataclasses.dataclass(frozen=True) class TrivialTransferPlan(TransferPlan): @property def tile_index_transforms(self): return (lambda x: x,) def select(self, group_elems: Sequence[ir.Value]) -> ir.Value: assert len(group_elems) == 1 return group_elems[0] def select_if_group(self, group_idx: int, old: ir.Value, new: ir.Value) -> ir.Value: assert group_idx == 0 return new @dataclasses.dataclass(frozen=True) class StaggeredTransferPlan(TransferPlan): stagger: int dim: int size: int group_stride: int @functools.cached_property def group_pred(self) -> ir.Value: c32 = lambda x: arith.constant(ir.IntegerType.get_signless(32), x) lane_idx = arith.remui(utils.thread_idx(), c32(WARP_SIZE)) group_idx = arith.remui(arith.divui(lane_idx, c32(self.group_stride)), c32(2)) return arith.cmpi(arith.CmpIPredicate.ne, group_idx, c32(0)) @property def tile_index_transforms(self): dim = self.dim def rotate(idx: tuple[int, ...]) -> tuple[int, ...]: return ( *idx[:dim], (idx[dim] + self.stagger) % self.size, *idx[dim + 1 :], ) return (lambda x: x, rotate) def select(self, group_elems: Sequence[ir.Value]) -> ir.Value: assert len(group_elems) == 2 return arith.select(self.group_pred, group_elems[1], group_elems[0]) def select_if_group(self, group_idx: int, old: ir.Value, new: ir.Value) -> ir.Value: assert 0 <= group_idx <= 1 sides = [old, new] if group_idx == 0 else [new, old] return arith.select(self.group_pred, *sides) class TransferPlanDerivationError(Exception): """Raised when a transfer plan cannot be derived due to a known limitation.""" pass def plan_gmem_transfer( nested_ref_shape: Sequence[Sequence[int]], nested_ref_strides: Sequence[Sequence[int]], layout: TiledLayout, element_bits: int, ) -> TransferPlan: """Plans the tiled transfer for GMEM and verifies sector coalescing.""" tiled_nested_shape, tiled_nested_strides = ( layout.tiling.tile_nested_shape_strides( tuple(tuple(x) for x in nested_ref_shape), tuple(tuple(x) for x in nested_ref_strides), ) ) tiles_shape = list(tiled_nested_shape) tiles_strides = list(tiled_nested_strides) for d in ( *layout.partitioned_warp_dims, *layout.partitioned_lane_dims, layout.vector_dim, ): tiles_shape[d] = (1,) * len(tiles_shape[d]) tiles_strides[d] = (0,) * len(tiles_strides[d]) tiles_shape = list(itertools.chain.from_iterable(tiles_shape)) tiles_strides = list(itertools.chain.from_iterable(tiles_strides)) warp_shape = list( itertools.chain.from_iterable( tiled_nested_shape[d] if isinstance(d, int) else (d.times,) for d in layout.warp_dims ) ) warp_strides = list( itertools.chain.from_iterable( tiled_nested_strides[d] if isinstance(d, int) else (0,) for d in layout.warp_dims ) ) lane_shape = list( itertools.chain.from_iterable( tiled_nested_shape[d] if isinstance(d, int) else (d.times,) for d in layout.lane_dims ) ) lane_strides = list( itertools.chain.from_iterable( tiled_nested_strides[d] if isinstance(d, int) else (0,) for d in layout.lane_dims ) ) vector_length = layout.vector_length transfer_bytes = (vector_length * element_bits) // 8 # The logic below relies on transfer_bytes being <= 16, but larger transfers # are inefficient anyway (they are uncoalesced). if transfer_bytes > 16: raise TransferPlanDerivationError( "GMEM transfer with per-thread vector length > 16 bytes is inefficient" ) if transfer_bytes.bit_count() != 1: raise TransferPlanDerivationError( "GMEM transfer with non-power-of-2 byte length is inefficient" ) lane_elem_offsets = np.dot(list(np.ndindex(*lane_shape)), lane_strides) lane_byte_offsets = lane_elem_offsets * element_bits // 8 # Find the fastest changing tile dimension. If its stride is equal to vector # length, then it's likely that LLVM will auto-vectorize consecutive accesses. # TODO(apaszke): We could also just change the layout to guarantee this. minor_tile_dim = min( ( d for d, (sz, st) in enumerate(zip(tiles_shape, tiles_strides)) if sz > 1 and st > 0 ), key=lambda d: (tiles_strides[d], -d), default=None, ) if minor_tile_dim is not None and tiles_strides[minor_tile_dim] == vector_length: max_vector_factor = max(1, 16 // transfer_bytes) # If the dimension is not divisible, we consider the epilogue since it will # have the shortest transfer. llvm_vector_factor = min( max_vector_factor, # Largest power of 2 that divides the tiled dimension. tiles_shape[minor_tile_dim] & -tiles_shape[minor_tile_dim], ) # Vectorization can only happen if all transfers are aligned. while llvm_vector_factor > 1: if np.all(lane_byte_offsets % (transfer_bytes * llvm_vector_factor) == 0): tiles_shape[minor_tile_dim] //= llvm_vector_factor tiles_strides[minor_tile_dim] *= llvm_vector_factor transfer_bytes *= llvm_vector_factor break llvm_vector_factor //= 2 # Transfers narrower than 4 bytes frequently don't perform well. if transfer_bytes < 4: raise TransferPlanDerivationError( f"GMEM transfer width ({transfer_bytes} bytes/thread) is" " smaller than 4 bytes" ) # We want the transfer to cover a whole number of entire sectors. # We split each sector into slots that are exactly transfer_bytes wide. SECTOR_BYTES = 32 slots_per_sector = SECTOR_BYTES // transfer_bytes lane_slot_offsets = lane_byte_offsets // transfer_bytes num_unique_slots = len(np.unique(lane_slot_offsets)) unique_lane_sectors = np.unique(lane_slot_offsets // slots_per_sector) if num_unique_slots != len(unique_lane_sectors) * slots_per_sector: raise TransferPlanDerivationError( "GMEM transfer does not access a whole number of complete 32 byte sectors" ) # Each warp-wide transfer must be aligned to 32-byte sectors. if any( ((s * element_bits) // 8) % SECTOR_BYTES != 0 for d, s in zip( itertools.chain(tiles_shape, warp_shape), itertools.chain(tiles_strides, warp_strides), ) if d > 1 and s != 0 ): raise TransferPlanDerivationError( "GMEM transfer stride is not aligned to 32B sector size" ) # Across the whole warpgroup (all tiles and warps), we want the transfer to # cover a whole number of complete 128-byte L2 cache lines. CACHE_LINE_BYTES = 128 SECTORS_PER_CACHE_LINE = CACHE_LINE_BYTES // SECTOR_BYTES wg_sector_offsets = unique_lane_sectors for d, s in zip( itertools.chain(tiles_shape, warp_shape), itertools.chain(tiles_strides, warp_strides), ): if d <= 1 or s == 0: continue assert (s * element_bits) % (8 * SECTOR_BYTES) == 0 s_sectors = (s * element_bits) // (8 * SECTOR_BYTES) wg_sector_offsets = ( wg_sector_offsets[:, None] + np.arange(d) * s_sectors ).ravel() num_unique_wg_sectors = len(np.unique(wg_sector_offsets)) num_unique_cache_lines = len( np.unique(wg_sector_offsets // SECTORS_PER_CACHE_LINE) ) if num_unique_wg_sectors != num_unique_cache_lines * SECTORS_PER_CACHE_LINE: raise TransferPlanDerivationError( "GMEM transfer does not access a whole number of complete 128 byte" " cache lines" ) return TrivialTransferPlan() def plan_tiled_transfer( nested_ref_shape: Sequence[Sequence[int]], nested_ref_strides: Sequence[Sequence[int]], layout: TiledLayout, element_bits: int, swizzle: int, ) -> TransferPlan: """Plans the tiled transfer in a way that avoids SMEM bank conflicts. Arguments: nested_ref_shape: The nested shape of the reference. For a logical ref with shape (2, 8, 16) and tiling (2, 8), this would be ((2,), (4, 2), (2, 8)). nested_ref_strides: The strides associated with the `nested_ref_shape`. layout: The layout of the value in registers being transferred. element_bits: Element bitwidth. swizzle: The swizzle pattern length. """ tiling = layout.tiling tiled_nested_shape, tiled_nested_strides = tiling.tile_nested_shape_strides( tuple(tuple(x) for x in nested_ref_shape), tuple(tuple(x) for x in nested_ref_strides), ) tiles_shape = list(tiled_nested_shape) tiles_strides = list(tiled_nested_strides) for d in (*layout.partitioned_warp_dims, *layout.partitioned_lane_dims, layout.vector_dim): tiles_shape[d] = (1,) * len(tiles_shape[d]) tiles_strides[d] = (0,) * len(tiles_strides[d]) tiles_shape = list(itertools.chain.from_iterable(tiles_shape)) tiles_strides = list(itertools.chain.from_iterable(tiles_strides)) warp_shape = list( itertools.chain.from_iterable( tiled_nested_shape[d] if isinstance(d, int) else (d.times,) for d in layout.warp_dims ) ) warp_strides = list( itertools.chain.from_iterable( tiled_nested_strides[d] if isinstance(d, int) else (0,) for d in layout.warp_dims ) ) lane_shape = list( itertools.chain.from_iterable( tiled_nested_shape[d] if isinstance(d, int) else (d.times,) for d in layout.lane_dims ) ) lane_strides = list( itertools.chain.from_iterable( tiled_nested_strides[d] if isinstance(d, int) else (0,) for d in layout.lane_dims ) ) vector_length = layout.vector_length # TODO(apaszke): Rewrite this function in terms of transfer_bytes (that we get # from the caller). swizzle_tile_elems = (16 * 8) // element_bits swizzle_group_elems = (128 * 8) // element_bits # Should be checked at the call site. assert vector_length * element_bits % 8 == 0 transfer_bytes = (vector_length * element_bits) // 8 # Below, all calculations are in elements, not in bytes, since it should # generalize better to sub-byte types. # Here, we verify two conditions: # 1. Each vector transfer only accesses addresses that fall within a single # swizzle tile (if not we'd need to split it and swizzle parts differently). chain = itertools.chain transfer_alignment = math.gcd(*( s for (s, d) in zip( chain(tiles_strides, warp_strides, lane_strides), chain(tiles_shape, warp_shape, lane_shape), ) if d > 1 )) if ( swizzle_tile_elems % transfer_alignment and vector_length <= transfer_alignment ): raise TransferPlanDerivationError( "Failed to prove that vector transfers don't cross swizzle tile" " boundaries. This check is incomplete, and does not guarantee that" f" this is a user error, but it might be. {transfer_alignment=}" ) # 2. The transfer pattern does not cause bank conflicts. # TODO(apaszke): For now, when performing transfers narrower than a bank, # we simply narrow each bank to the transfer width. The truth is more likely # that bank conflicts only don't occur if the addresses mapping to the same # bank are contiguous, but that's a more complicated check to perform. if transfer_bytes > SMEM_BANK_BYTES * 4: raise TransferPlanDerivationError( f"{transfer_bytes=} > {SMEM_BANK_BYTES * 4} not implemented" ) if element_bits > SMEM_BANK_BYTES * 8: raise TransferPlanDerivationError( f"{element_bits=} > {SMEM_BANK_BYTES * 8} not implemented" ) smem_bank_bytes = min(SMEM_BANK_BYTES, transfer_bytes) num_banks = SMEM_BANKS * (SMEM_BANK_BYTES // smem_bank_bytes) elems_per_bank = (smem_bank_bytes * 8) // element_bits num_wavefronts = max(transfer_bytes // smem_bank_bytes, 1) wavefront_lanes = WARP_SIZE // num_wavefronts lane_mask = np.full(lane_shape, False) lane_mask[tuple(slice(0, 1) if s == 0 else slice(None) for s in lane_strides)] = True wavefront_mask = lane_mask.reshape(num_wavefronts, wavefront_lanes) lane_offsets_in_tile = np.dot(list(np.ndindex(*lane_shape)), lane_strides) def has_bank_conflicts(tile_idx_transform): num_tiles = math.prod(tiles_shape) tile_idxs = np.unravel_index(np.arange(num_tiles), tiles_shape) tile_idxs = np.expand_dims(np.stack(tile_idxs, 1), 1) # [#tiles, 1, #dims] lane_tile_idx = tile_idx_transform(tile_idxs) # [#tiles, #lanes/1, #dims] assert lane_tile_idx.shape[1] in {1, WARP_SIZE} lane_tile_offsets = np.dot(lane_tile_idx, tiles_strides) offsets = lane_tile_offsets + lane_offsets_in_tile # [#tiles, #lanes] assert offsets.shape[-1] == WARP_SIZE swizzle_groups = (offsets // swizzle_group_elems) % (swizzle // 16) swizzle_bits = swizzle_groups * swizzle_tile_elems lane_banks = ((offsets ^ swizzle_bits) // elems_per_bank) % num_banks wavefront_banks = lane_banks.reshape(-1, num_wavefronts, wavefront_lanes) # We step over wavefronts since they might have a different number of lanes. wavefront_banks = wavefront_banks.swapaxes(0, 1) for banks, mask in zip(wavefront_banks, wavefront_mask): banks = banks[:, mask] # Order of threads within the wavefront is unimportant. banks = np.sort(banks, axis=-1) # There are no conflicts if each wavefront only contains unique banks. repeats = np.any(banks[..., 1:] == banks[..., :-1]) if repeats: return True return False # We don't need any special treatment if there are no conflicts when each lane # transfers the same tile at a time. if not has_bank_conflicts(lambda tile_idx: tile_idx): return TrivialTransferPlan() # Otherwise, we will try to partition the lanes into two groups and have # each group store to different tile. The only tile dimensions that can help # us with bank conflicts are those that have multiple elements and a stride # that's not a multiple of the number of banks. # # Note that the code is set up so that we could also consider partitioning # the lanes into more groups, but the selects will become more expensive if # we do that. It's a possibility we have if we need it. candidate_dims = ( i for i, (s, d) in enumerate(zip(tiles_strides, tiles_shape)) if d > 1 and s % (SMEM_BANKS * elems_per_bank) ) for dim in candidate_dims: for group_stride in (1, 2, 4, 8, 16): # We change the group assignment each group_stride lanes. lane_id = np.arange(WARP_SIZE)[:, None] lane_group = (lane_id // group_stride) % 2 # We only consider a transformation where the second group stores to a # tile that's a constant offset (modulo dim size) from the first one. for stagger in range(1, tiles_shape[dim]): offset = np.zeros(len(tiles_shape), np.int64) offset[dim] = stagger transform = lambda idx: (idx + offset * lane_group) % tiles_shape if not has_bank_conflicts(transform): # We've found a strategy that avoids bank conflicts! return StaggeredTransferPlan( stagger, dim, tiles_shape[dim], group_stride ) raise TransferPlanDerivationError( "Failed to synthesize a transfer pattern that avoids bank conflicts" ) # We allow contractions, to potentially take advantage of FMA instructions. # They can change the results, but the precision should only increase. def addf(a: ir.Value, b: ir.Value): return arith.addf(a, b, fastmath=arith.FastMathFlags.contract) def subf(a: ir.Value, b: ir.Value): return arith.subf(a, b, fastmath=arith.FastMathFlags.contract) def mulf(a: ir.Value, b: ir.Value): return arith.mulf(a, b, fastmath=arith.FastMathFlags.contract) @overload def optimization_barrier( a: mgpu.FragmentedArray, b: mgpu.FragmentedArray, /, *arrays: mgpu.FragmentedArray, ) -> Sequence[mgpu.FragmentedArray]: ... @overload def optimization_barrier(a: mgpu.FragmentedArray, /) -> mgpu.FragmentedArray: ... def optimization_barrier(*arrays): """Acts as an optimization barrier for LLVM. Passing arrays through this function will make sure that they are computed before any side-effecting operations that follow this barrier. """ i32 = ir.IntegerType.get_signless(32) def _repack(regs_it, reg_ty): if not isinstance(reg_ty, ir.VectorType): result_reg = next(regs_it) assert result_reg.type == reg_ty return result_reg num_i32_regs = utils.bitwidth(reg_ty) // 32 i32_reg_ty = ir.VectorType.get((num_i32_regs,), i32) reg = llvm.mlir_undef(i32_reg_ty) for i_elem in range(num_i32_regs): val = llvm.bitcast(i32, next(regs_it)) reg = llvm.insertelement(reg, val, arith.constant(i32, i_elem)) return vector.bitcast(reg_ty, reg) regs = [] reg_dtypes = [] reg_constraints = [] # We unpack each array into a flat list of registers, and prepare the # functions that invert the transform in repack_fns. for array in arrays: reg_ty = array.registers.flat[0].type dtype = array.mlir_dtype if isinstance(dtype, ir.F32Type) or dtype == i32: if isinstance(reg_ty, ir.VectorType): [vec_len] = ir.VectorType(reg_ty).shape array_regs = [ vector.extract( reg, dynamic_position=[], static_position=ir.DenseI64ArrayAttr.get([pos]), ) for reg in array.registers.flat for pos in range(vec_len) ] else: array_regs = list(array.registers.flat) reg_constraint = "r" if dtype == i32 else "f" elif utils.bitwidth(dtype) < 32: reg_packing = 4 // utils.bytewidth(dtype) if not isinstance(reg_ty, ir.VectorType): raise NotImplementedError(array.mlir_dtype) [vec_len] = ir.VectorType(reg_ty).shape if vec_len % reg_packing: raise NotImplementedError(vec_len) num_i32_regs = vec_len // reg_packing i32_reg_ty = ir.VectorType.get((num_i32_regs,), i32) array_regs = [ vector.extract( vector.bitcast(i32_reg_ty, reg), dynamic_position=[], static_position=ir.DenseI64ArrayAttr.get([i]), ) for i in range(num_i32_regs) for reg in array.registers.flat ] reg_constraint = "r" else: raise NotImplementedError(array.mlir_dtype) regs += array_regs reg_dtypes += [array_regs[0].type] * len(array_regs) reg_constraints += [reg_constraint] * len(array_regs) ptx = "" all_reg_constraints = ",".join( [*("=" + c for c in reg_constraints), *map(str, range(len(reg_constraints)))] ) if len(reg_dtypes) == 1: # The InlineAsm::verify() function doesn't allow a struct output when there # is only one element (even though that seems to work for the case below). result_elem = llvm.inline_asm( reg_dtypes[0], regs, ptx, all_reg_constraints, asm_dialect=0, has_side_effects=True, ) regs = [result_elem] else: struct_ty = llvm.StructType.get_literal(reg_dtypes) result_struct = llvm.inline_asm( struct_ty, regs, ptx, all_reg_constraints, asm_dialect=0, has_side_effects=True, ) assert isinstance(result_struct, ir.Value) regs = [ llvm.extractvalue(dtype, result_struct, [i]) for i, dtype in enumerate(reg_dtypes) ] i32 = ir.IntegerType.get_signless(32) results = [] regs_it = iter(regs) for array in arrays: num_regs = array.registers.size reg_ty = array.registers.flat[0].type if isinstance(reg_ty, ir.VectorType): reg_ty = ir.VectorType(reg_ty) new_registers = np.empty((num_regs,), dtype=object) for i_vreg in range(num_regs): reg = _repack(regs_it, reg_ty) assert reg.type == reg_ty, (reg.type, reg_ty) new_registers[i_vreg] = reg results.append( FragmentedArray( _registers=new_registers.reshape(array.registers.shape), _layout=array.layout, _is_signed=array.is_signed, ) ) return results[0] if len(arrays) == 1 else results def tiled_copy_smem_gmem_layout( row_tiles: int, col_tiles: int, swizzle: int, bitwidth: int ) -> TiledLayout: swizzle_elems = 8 * swizzle // bitwidth if row_tiles % 4 == 0: warp_row_tiles, warp_col_tiles = 4, 1 elif row_tiles % 2 == 0: if col_tiles % 2: raise NotImplementedError("Number of tiles is not a multiple of 4") warp_row_tiles, warp_col_tiles = 2, 2 else: if col_tiles % 4: raise NotImplementedError("Number of tiles is not a multiple of 4") warp_row_tiles, warp_col_tiles = 1, 4 row_tiles //= warp_row_tiles col_tiles //= warp_col_tiles bytes_per_thread = min(16, 8 * swizzle // WARP_SIZE) lane_row_tiles = lane_col_tiles = 1 if bytes_per_thread < 16: # Try to splread multiple tiles over a warp. max_scale_up = 16 // bytes_per_thread while max_scale_up > 1 and col_tiles % 2 == 0: max_scale_up //= 2 lane_col_tiles *= 2 col_tiles //= 2 while max_scale_up > 1 and row_tiles % 2 == 0: max_scale_up //= 2 lane_row_tiles *= 2 row_tiles //= 2 bytes_per_thread *= lane_row_tiles * lane_col_tiles if 8 * bytes_per_thread < bitwidth: raise NotImplementedError("Element types with bitwidth so large aren't supported") vector_length = bytes_per_thread * 8 // bitwidth assert swizzle_elems % vector_length == 0 # How many steps of vector transfers are needed to transfer a single tile? if vector_length * WARP_SIZE > 8 * swizzle_elems: steps_per_tile = 1 else: steps_per_tile = 8 * swizzle_elems // (vector_length * WARP_SIZE) tile_rows_per_step = 8 // steps_per_tile # There are two cases to consider here: either a single transfer fits within # a single tile (lane_row_tiles == lane_col_tiles == 1), which is the case # for large swizzles, or it spans multiple tiles. The layout below ensures # that consecutive lanes first traverse the columns within a tile, followed # by rows within a tile, columns across tiles, and then rows across tiles. # This ensures we never end up with bank conflicts, and yields well # coalesced GMEM accesses. return TiledLayout( Tiling( ( (warp_row_tiles * lane_row_tiles * 8, warp_col_tiles * lane_col_tiles * swizzle_elems), (lane_row_tiles * 8, lane_col_tiles * swizzle_elems), (8, swizzle_elems), (tile_rows_per_step, swizzle_elems), (vector_length,) ) ), warp_dims=(-9, -8), lane_dims=(-7, -6, -3, -2), vector_dim=-1, _check_canonical=False, ).canonicalize() def copy_tiled(src: ir.Value, dst: ir.Value, swizzle: int = 16): """Copy the data from the src reference to the dst reference. Exactly one of src/dst should be in SMEM, while the other should be in GMEM. The SMEM reference is expected to be tiled into (8, swizzle_elems) (as it would for MMA), and so should have a rank larger by 2 than the GMEM ref. """ src_ty = ir.MemRefType(src.type) dst_ty = ir.MemRefType(dst.type) if math.prod(src_ty.shape) != math.prod(dst_ty.shape): raise ValueError( "Source and destination must have the same number of elements, but got" f" source shape {src_ty.shape} and destination shape {dst_ty.shape}" ) if src_ty.element_type != dst_ty.element_type: raise ValueError( "Source and destination must have the same element type, but got" f" source type {src_ty.element_type} and destination type" f" {dst_ty.element_type}" ) bitwidth = utils.bitwidth(src_ty.element_type) # Signedness doesn't matter, but we need to specify something for the # intermediate arrays. is_signed = False if isinstance(src_ty.element_type, ir.IntegerType) else None if utils.is_smem_ref(src_ty) != utils.is_smem_ref(dst_ty): if utils.is_smem_ref(src_ty): smem_ty, gmem_ty = src_ty, dst_ty else: smem_ty, gmem_ty = dst_ty, src_ty if smem_ty.rank != gmem_ty.rank + 2: raise ValueError( "SMEM reference must have a rank larger by 2 than the destination" f" reference (due to 2D tiling), but got SMEM rank {smem_ty.rank} and" f" destination rank {gmem_ty.rank}." ) swizzle_elems = 8 * swizzle // bitwidth if smem_ty.shape[-2:] != [8, swizzle_elems]: raise NotImplementedError( f"For {swizzle=}, expected SMEM tiling to be (8, {swizzle_elems})" ) expected_src_shape = utils.tile_shape(gmem_ty.shape, (8, swizzle_elems)) if tuple(smem_ty.shape) != expected_src_shape: raise ValueError( f"Expected SMEM reference to have shape {expected_src_shape} (tiling" f" {gmem_ty.shape} by (8, {swizzle_elems})), but got {smem_ty.shape}" ) layout = tiled_copy_smem_gmem_layout( *smem_ty.shape[-4:-2], swizzle, bitwidth # pyrefly: ignore[bad-argument-count] ) if utils.is_smem_ref(src_ty): regs = FragmentedArray.load_tiled(src, swizzle, is_signed=is_signed, layout=layout) regs.store_untiled(dst, optimized=False) else: regs = FragmentedArray.load_untiled(src, is_signed=is_signed, layout=layout, optimized=False) regs.store_tiled(dst, swizzle) return raise NotImplementedError(f"Unsupported copy: {src.type} -> {dst.type}") def is_supported_strided_layout_broadcast( src: WGStridedFragLayout, dst: WGStridedFragLayout, dims: tuple[int, ...], ) -> bool: """We only support broadcasting of leading dimensions.""" if src.vec_size != dst.vec_size: return False # Check if input maps exactly to the end (prevents trailing dims). if dims != tuple(range(len(dst.shape) - len(src.shape), len(dst.shape))): return False # Identify input indices that are expanded vs. those that are preserved # Expansion: input is 1, output is > 1. # Preserved: input is > 1. exp_indices, pre_indices = [], [] for i, dim in enumerate(src.shape): if dim == 1 and dst.shape[dims[i]] > 1: exp_indices.append(i) if dim > 1: pre_indices.append(i) # If both exist, all expansions must happen before all preserved # dimensions. if exp_indices and pre_indices and max(exp_indices) >= min(pre_indices): return False return True def concatenate( arrays: Sequence[FragmentedArray], axis: int = 0, ) -> FragmentedArray: """Concatenates fragmented arrays along the specified axis.""" if not arrays: raise ValueError("Need at least one array to concatenate") arr0 = arrays[0] rank = len(arr0.shape) if not -rank <= axis < rank: raise ValueError(f"{axis=} is out of bounds for array of {rank=}") if axis < 0: axis += rank if len(arrays) == 1: return arr0 new_shape = list(arr0.shape) for i, arr in enumerate(arrays[1:], start=1): if len(arr.shape) != rank: raise ValueError( f"All arrays must have the same rank, got {len(arr.shape)} at index" f" {i} (expected {rank})" ) if arr.mlir_dtype != arr0.mlir_dtype: raise ValueError( f"All arrays must have the same dtype, got {arr.mlir_dtype} at" f" index {i} (expected {arr0.mlir_dtype})" ) if arr.is_signed != arr0.is_signed: raise ValueError( f"All arrays must have the same signedness, got {arr.is_signed} at" f" index {i} (expected {arr0.is_signed})" ) for d in range(rank): if d != axis and arr.shape[d] != arr0.shape[d]: raise ValueError( "All arrays must have matching shapes along non-concatenated" f" dimensions, got shape {arr.shape} at index {i} (expected dim" f" {d} to be {arr0.shape[d]})" ) new_shape[axis] += arr.shape[axis] new_shape = tuple(new_shape) match arr0.layout: case TiledLayout(): for i, arr in enumerate(arrays[1:], start=1): if arr.layout != arr0.layout: raise ValueError( f"All arrays must have the same layout, got {arr.layout} at" f" index {i} (expected {arr0.layout})" ) new_regs = np.concatenate([arr.registers for arr in arrays], axis=axis) return FragmentedArray( _registers=new_regs, _layout=arr0.layout, _is_signed=arr0.is_signed ) case WGStridedFragLayout(vec_size=vec_size): if axis != 0: raise NotImplementedError( "Concatenating arrays with strided layout is only supported along" " axis 0" ) for i, arr in enumerate(arrays[1:], start=1): if not isinstance(arr.layout, WGStridedFragLayout): raise ValueError( f"Expected WGStridedFragLayout, got {arr.layout} at index {i}" ) if arr.layout.vec_size != vec_size: raise ValueError( "All WGStridedFragLayout arrays must have the same vec_size," f" got {arr.layout.vec_size} at index {i} (expected {vec_size})" ) new_layout = WGStridedFragLayout(shape=new_shape, vec_size=vec_size) new_regs = np.concatenate([arr.registers for arr in arrays], axis=0) return FragmentedArray( _registers=new_regs, _layout=new_layout, _is_signed=arr0.is_signed ) case WGSplatFragLayout(): raise NotImplementedError( "Concatenating arrays with splat layout is not supported." ) case layout: assert_never(layout)