From 52b8d50fc30cd5a8b290474344c1e931c08c45c2 Mon Sep 17 00:00:00 2001 From: Pavel Belevich Date: Sun, 23 Aug 2026 13:14:06 -0400 Subject: [PATCH] Version 16: Add CUDA register tiles --- src/mytriton/block_shapes.py | 196 ++++++- src/mytriton/cuda_codegen.py | 669 ++++++++++++++++++++---- tests/test_block_shapes.py | 186 +++++++ tests/test_shared_memory.py | 979 ++++++++++++++++++++++++++++++++++- 4 files changed, 1885 insertions(+), 145 deletions(-) diff --git a/src/mytriton/block_shapes.py b/src/mytriton/block_shapes.py index f28d865..1edf2c6 100644 --- a/src/mytriton/block_shapes.py +++ b/src/mytriton/block_shapes.py @@ -129,6 +129,86 @@ def coordinates(self, linear_index: int) -> tuple[int, ...]: return tuple(coordinates) +@dataclass(frozen=True) +class CudaRegisterTileLayout: + """Distribution of a logical rank-2 tile across threads and registers.""" + + logical_shape: tuple[int, ...] + thread_shape: tuple[int, ...] + + def __post_init__(self) -> None: + for name, shape in ( + ("logical tile", self.logical_shape), + ("CUDA thread", self.thread_shape), + ): + if len(shape) != 2: + raise ValueError(f"{name} shape must have rank 2, got {shape}") + + if any(type(dim) is not int or dim <= 0 for dim in shape): + raise ValueError( + f"{name} dimensions must be positive integers, got {shape}" + ) + + for logical_dim, thread_dim in zip( + self.logical_shape, + self.thread_shape, + strict=True, + ): + if logical_dim % thread_dim != 0: + raise ValueError( + "register tile requires logical dimensions divisible " + "by CUDA thread dimensions, " + f"got {self.logical_shape} and {self.thread_shape}" + ) + + @property + def register_shape(self) -> tuple[int, ...]: + return tuple( + logical_dim // thread_dim + for logical_dim, thread_dim in zip( + self.logical_shape, + self.thread_shape, + strict=True, + ) + ) + + @property + def registers_per_thread(self) -> int: + return prod(self.register_shape) + + def logical_coordinate( + self, + *, + thread_coordinate: tuple[int, ...], + register_coordinate: tuple[int, ...], + ) -> tuple[int, ...]: + for name, coordinate, shape in ( + ("thread", thread_coordinate, self.thread_shape), + ("register", register_coordinate, self.register_shape), + ): + if len(coordinate) != 2 or any( + type(index) is not int or index < 0 or index >= dim + for index, dim in zip( + coordinate, + shape, + strict=True, + ) + ): + raise ValueError( + f"invalid {name} coordinate {coordinate} for shape {shape}" + ) + + return tuple( + thread_index + register_index * thread_dim + for thread_index, register_index, thread_dim in zip( + thread_coordinate, + register_coordinate, + self.thread_shape, + strict=True, + ) + ) + + @dataclass(frozen=True) class CudaKernelLayout: """Logical output tile and physical CUDA thread organization.""" @@ -177,6 +257,12 @@ def is_rank2(self) -> bool: def threads_per_block(self) -> int: return prod(self.thread_shape) + def register_tile_layout(self) -> CudaRegisterTileLayout: + return CudaRegisterTileLayout( + logical_shape=self.output_tile_shape, + thread_shape=self.thread_shape, + ) + def tile_layout( self, logical_shape: tuple[int, ...], @@ -307,6 +393,30 @@ def reduction_block_shapes(ssa_ops: list[SSAItem]) -> list[tuple[int, ...]]: return shapes +def dot_result_shapes( + ssa_ops: list[SSAItem], +) -> list[tuple[int, ...]]: + """Collect logical output shapes produced by tl.dot.""" + + from .ssa import SSAForRange, SSAValue + + shapes = [] + + for op in ssa_ops: + if isinstance(op, SSAForRange): + shapes.extend(dot_result_shapes(op.body)) + continue + + if ( + op.opcode == "dot" + and isinstance(op.result, SSAValue) + and isinstance(op.result.ty, BlockType) + ): + shapes.append(op.result.ty.shape) + + return shapes + + def _infer_cuda_kernel_tile_shape(ssa_ops: list[SSAItem]) -> tuple[int, ...]: shapes = store_block_shapes(ssa_ops) @@ -345,38 +455,81 @@ def _infer_cuda_kernel_tile_shape(ssa_ops: list[SSAItem]) -> tuple[int, ...]: return (next(iter(widths)),) +CUDA_DOT_MAX_THREADS = 32 + + +def _divisors(dim: int) -> tuple[int, ...]: + return tuple(candidate for candidate in range(1, dim + 1) if dim % candidate == 0) + + +def _infer_cuda_dot_thread_shape( + output_tile_shape: tuple[int, ...], +) -> tuple[int, ...]: + if len(output_tile_shape) != 2: + raise ValueError(f"CUDA dot output must have rank 2, got {output_tile_shape}") + + rows, columns = output_tile_shape + candidates = [ + (thread_rows, thread_columns) + for thread_rows in _divisors(rows) + for thread_columns in _divisors(columns) + if (thread_rows * thread_columns <= CUDA_DOT_MAX_THREADS) + ] + + if not candidates: + raise ValueError( + f"cannot fit CUDA dot tile {output_tile_shape} " + f"into {CUDA_DOT_MAX_THREADS} threads" + ) + + return max( + candidates, + key=lambda shape: ( + prod(shape), + min(shape), + shape[1], + ), + ) + + def _infer_cuda_thread_shape( output_tile_shape: tuple[int, ...], reduction_shapes: list[tuple[int, ...]], + dot_shapes: list[tuple[int, ...]], ) -> tuple[int, ...]: - if not reduction_shapes: - return output_tile_shape - - if any(len(shape) != 1 for shape in reduction_shapes): - rendered = ", ".join(str(shape) for shape in reduction_shapes) - raise ValueError(f"CUDA reductions require rank-1 inputs, got {rendered}") + if reduction_shapes: + if any(len(shape) != 1 for shape in reduction_shapes): + rendered = ", ".join(str(shape) for shape in reduction_shapes) + raise ValueError(f"CUDA reductions require rank-1 inputs, got {rendered}") + + reduction_widths = {shape[0] for shape in reduction_shapes} + if len(reduction_widths) != 1: + rendered = ", ".join(str(width) for width in sorted(reduction_widths)) + raise ValueError( + f"CUDA reductions require one thread width, got: {rendered}" + ) - reduction_widths = {shape[0] for shape in reduction_shapes} - if len(reduction_widths) != 1: - rendered = ", ".join(str(width) for width in sorted(reduction_widths)) - raise ValueError(f"CUDA reductions require one thread width, got: {rendered}") + reduction_width = next(iter(reduction_widths)) - reduction_width = next(iter(reduction_widths)) + if len(output_tile_shape) == 1: + output_width = output_tile_shape[0] + if output_width not in (1, reduction_width): + raise ValueError( + f"reduction width {reduction_width} does not match " + f"output tile width {output_width}" + ) + return (reduction_width,) - if len(output_tile_shape) == 1: - output_width = output_tile_shape[0] - if output_width not in (1, reduction_width): + if prod(output_tile_shape) != reduction_width: raise ValueError( f"reduction width {reduction_width} does not match " - f"output tile width {output_width}" + f"output tile size {prod(output_tile_shape)}" ) - return (reduction_width,) - if prod(output_tile_shape) != reduction_width: - raise ValueError( - f"reduction width {reduction_width} does not match " - f"output tile size {prod(output_tile_shape)}" - ) + return output_tile_shape + + if len(output_tile_shape) == 2 and output_tile_shape in dot_shapes: + return _infer_cuda_dot_thread_shape(output_tile_shape) return output_tile_shape @@ -386,6 +539,7 @@ def cuda_kernel_layout(ssa_ops: list[SSAItem]) -> CudaKernelLayout: thread_shape = _infer_cuda_thread_shape( output_tile_shape, reduction_block_shapes(ssa_ops), + dot_result_shapes(ssa_ops), ) return CudaKernelLayout( diff --git a/src/mytriton/cuda_codegen.py b/src/mytriton/cuda_codegen.py index d98b37d..6c7773b 100644 --- a/src/mytriton/cuda_codegen.py +++ b/src/mytriton/cuda_codegen.py @@ -2,7 +2,11 @@ from dataclasses import dataclass from typing import ClassVar -from .block_shapes import CudaKernelLayout, cuda_kernel_layout +from .block_shapes import ( + CudaKernelLayout, + CudaRegisterTileLayout, + cuda_kernel_layout, +) from .cuda_dot_staging import ( CudaDotSharedBuffers, CudaDotStagingAnalysis, @@ -44,6 +48,63 @@ def width(self) -> int: return self.end - self.start +@dataclass(frozen=True) +class CudaRegisterTileRef: + """CUDA registers owned by one thread for a logical rank-2 tile.""" + + base: str + layout: CudaRegisterTileLayout + broadcast_axes: tuple[int, ...] = () + + def __post_init__(self) -> None: + if tuple(sorted(set(self.broadcast_axes))) != self.broadcast_axes: + raise ValueError( + f"broadcast axes must be unique and sorted, got {self.broadcast_axes}" + ) + + if any(axis not in (0, 1) for axis in self.broadcast_axes): + raise ValueError(f"invalid broadcast axes {self.broadcast_axes}") + + @property + def storage_shape(self) -> tuple[int, ...]: + return tuple( + 1 if axis in self.broadcast_axes else dim + for axis, dim in enumerate(self.layout.register_shape) + ) + + def storage_coordinates(self) -> tuple[tuple[int, int], ...]: + return tuple( + (register_row, register_column) + for register_row in range(self.storage_shape[0]) + for register_column in range(self.storage_shape[1]) + ) + + def element(self, register_coordinate: tuple[int, ...]) -> str: + if len(register_coordinate) != 2 or any( + type(index) is not int or index < 0 or index >= dim + for index, dim in zip( + register_coordinate, + self.layout.register_shape, + strict=True, + ) + ): + raise ValueError( + f"invalid register coordinate {register_coordinate} " + f"for {self.layout.register_shape}" + ) + + storage_coordinate = tuple( + 0 if axis in self.broadcast_axes else index + for axis, index in enumerate(register_coordinate) + ) + + if all(dim == 1 for dim in self.storage_shape): + return self.base + + row, column = storage_coordinate + return f"{self.base}_{row}_{column}" + + class SSACUDACodegen: MAX_SHARED_MEMORY_BYTES: ClassVar[int] = 48 * 1024 @@ -58,7 +119,10 @@ class SSACUDACodegen: def __init__(self): self.lines: list[str] = [] - self.values: dict[int, str | CudaPtrRef | CudaArangeRef] = {} + self.values: dict[ + int, + str | CudaPtrRef | CudaArangeRef | CudaRegisterTileRef, + ] = {} self.layout = CudaKernelLayout( output_tile_shape=(1,), thread_shape=(1,), @@ -104,7 +168,9 @@ def literal(self, value: object) -> str: raise TypeError(f"Unsupported CUDA literal: {value!r}") - def operand(self, operand: SSAOperand) -> str | CudaPtrRef | CudaArangeRef | None: + def operand( + self, operand: SSAOperand + ) -> str | CudaPtrRef | CudaArangeRef | CudaRegisterTileRef | None: if operand is None: return None if isinstance(operand, SSAValue): @@ -133,6 +199,61 @@ def expression_operand(self, operand: SSAOperand) -> str: raise TypeError(f"Expected CUDA scalar expression, got {value}") return value + def register_expression_operand( + self, + operand: SSAOperand, + register_coordinate: tuple[int, ...], + expected_layout: CudaRegisterTileLayout, + ) -> str: + value = self.operand(operand) + + if isinstance(value, CudaRegisterTileRef): + if value.layout != expected_layout: + raise TypeError( + "incompatible CUDA register tile layouts: " + f"{value.layout} and {expected_layout}" + ) + + return value.element(register_coordinate) + + if isinstance(value, str): + return value + + raise TypeError( + f"expected CUDA scalar or register tile expression, got {value}" + ) + + def register_pointer_operand( + self, + operand: SSAOperand, + register_coordinate: tuple[int, ...], + expected_layout: CudaRegisterTileLayout, + ) -> CudaPtrRef: + value = self.operand(operand) + + if isinstance(value, CudaRegisterTileRef): + if value.layout != expected_layout: + raise TypeError( + "incompatible CUDA register tile layouts: " + f"{value.layout} and {expected_layout}" + ) + + return CudaPtrRef( + base=value.element(register_coordinate), + index="0", + ) + + if isinstance(value, CudaPtrRef): + return value + + if isinstance(value, str): + return CudaPtrRef( + base=value, + index="0", + ) + + raise TypeError(f"expected CUDA pointer or register pointer tile, got {value}") + def pointer_operand(self, operand: SSAOperand) -> CudaPtrRef: value = self.operand(operand) if isinstance(value, str): @@ -307,13 +428,22 @@ def emit_dot_from_shared_memory( f"got {result.ty.shape}" ) - if result.ty.shape != self.layout.thread_shape: + if result.ty.shape != self.layout.output_tile_shape: raise TypeError( - "CUDA-core dot currently requires one CUDA thread per " - f"result element, got result {result.ty.shape} and " - f"thread layout {self.layout.thread_shape}" + "CUDA-core dot result must match the kernel output tile, " + f"got result {result.ty.shape} and " + f"output tile {self.layout.output_tile_shape}" ) + try: + register_layout = self.layout.register_tile_layout() + except ValueError as error: + raise TypeError( + "CUDA-core dot cannot distribute result tile " + f"{result.ty.shape} across CUDA threads " + f"{self.layout.thread_shape}" + ) from error + if ( result.ty.element != F32 or buffers.lhs.element_ty != F32 @@ -321,33 +451,49 @@ def emit_dot_from_shared_memory( ): raise TypeError("CUDA-core dot currently supports only f32") - result_name = f"v{result.id}" + result_ref = CudaRegisterTileRef( + base=f"v{result.id}", + layout=register_layout, + ) reduction_index = f"dot_k_{result.id}" - row = self.thread_coordinate(0) - column = self.thread_coordinate(1) - lhs_element = buffers.lhs.element( - row, - reduction_index, - ) - rhs_element = buffers.rhs.element( - reduction_index, - column, - ) + register_coordinates = self.register_coordinates(register_layout) - self.lines.extend( - [ - f" float {result_name} = 0.0f;", - ( - f" for (int {reduction_index} = 0; " - f"{reduction_index} < {buffers.reduction_size}; " - f"++{reduction_index}) {{" - ), - (f" {result_name} += {lhs_element} * {rhs_element};"), - " }", - ] + for register_coordinate in register_coordinates: + result_element = result_ref.element(register_coordinate) + self.lines.append(f" float {result_element} = 0.0f;") + + self.lines.append( + f" for (int {reduction_index} = 0; " + f"{reduction_index} < {buffers.reduction_size}; " + f"++{reduction_index}) {{" ) - self.values[result.id] = result_name + + for register_coordinate in register_coordinates: + row, column = self.register_logical_coordinates( + register_layout, + register_coordinate, + ) + result_element = result_ref.element(register_coordinate) + lhs_element = buffers.lhs.element( + row, + reduction_index, + ) + rhs_element = buffers.rhs.element( + reduction_index, + column, + ) + + self.lines.append( + f" {result_element} += {lhs_element} * {rhs_element};" + ) + + self.lines.append(" }") + + if register_layout.registers_per_thread == 1: + self.values[result.id] = result_ref.element((0, 0)) + else: + self.values[result.id] = result_ref # All threads must finish reading the current tiles before another # runtime K-loop iteration overwrites the shared buffers. @@ -457,6 +603,38 @@ def thread_coordinate(self, thread_axis: int) -> str: return coordinates[thread_axis] + def register_coordinates( + self, + layout: CudaRegisterTileLayout, + ) -> tuple[tuple[int, int], ...]: + return tuple( + (register_row, register_column) + for register_row in range(layout.register_shape[0]) + for register_column in range(layout.register_shape[1]) + ) + + def register_logical_coordinates( + self, + layout: CudaRegisterTileLayout, + register_coordinate: tuple[int, ...], + ) -> tuple[str, ...]: + offsets = layout.logical_coordinate( + thread_coordinate=(0, 0), + register_coordinate=register_coordinate, + ) + + return tuple( + thread_coordinate if offset == 0 else f"{thread_coordinate} + {offset}" + for thread_coordinate, offset in zip( + ( + self.thread_coordinate(0), + self.thread_coordinate(1), + ), + offsets, + strict=True, + ) + ) + def emit_rank2_prologue(self) -> None: if not self.is_rank2_kernel(): return @@ -560,7 +738,7 @@ def emit_for_range(self, loop: SSAForRange) -> None: index_name = f"v{loop.index.id}" - carried_names = [] + carried_values: list[str | CudaRegisterTileRef] = [] for carried_input, carried_arg, result in zip( loop.carried_inputs, @@ -568,17 +746,52 @@ def emit_for_range(self, loop: SSAForRange) -> None: loop.results, strict=True, ): - init = self.expression_operand(carried_input) - name = f"v{result.id}" + register_layout: CudaRegisterTileLayout | None = None + + if ( + isinstance(result.ty, BlockType) + and result.ty.rank == 2 + and result.ty.shape == self.layout.output_tile_shape + ): + candidate_layout = self.layout.register_tile_layout() + + if candidate_layout.registers_per_thread > 1: + register_layout = candidate_layout + + if register_layout is None: + init = self.expression_operand(carried_input) + name = f"v{result.id}" + cuda_ty = self.cuda_type(result.ty) + + self.lines.append(f" {cuda_ty} {name} = {init};") + self.values[carried_arg.id] = name + self.values[result.id] = name + carried_values.append(name) + continue + + result_ref = CudaRegisterTileRef( + base=f"v{result.id}", + layout=register_layout, + ) cuda_ty = self.cuda_type(result.ty) - self.lines.append(f" {cuda_ty} {name} = {init};") - self.values[carried_arg.id] = name - self.values[result.id] = name - carried_names.append(name) + for register_coordinate in self.register_coordinates(register_layout): + init = self.register_expression_operand( + carried_input, + register_coordinate, + register_layout, + ) + name = result_ref.element(register_coordinate) + self.lines.append(f" {cuda_ty} {name} = {init};") + + self.values[carried_arg.id] = result_ref + self.values[result.id] = result_ref + carried_values.append(result_ref) self.lines.append( - f" for (int {index_name} = {start}; {index_name} < {stop}; {index_name} += {step}) {{" + f" for (int {index_name} = {start}; " + f"{index_name} < {stop}; " + f"{index_name} += {step}) {{" ) self.values[loop.index.id] = index_name @@ -591,20 +804,173 @@ def emit_for_range(self, loop: SSAForRange) -> None: elif not self.is_staging_only(body_op): self.emit(body_op) - for yielded, carried_name in zip(loop.yields, carried_names, strict=True): - value = self.expression_operand(yielded) - self.lines.append(f" {carried_name} = {value};") + for yielded, carried_value in zip( + loop.yields, + carried_values, + strict=True, + ): + if isinstance(carried_value, CudaRegisterTileRef): + for register_coordinate in self.register_coordinates( + carried_value.layout + ): + value = self.register_expression_operand( + yielded, + register_coordinate, + carried_value.layout, + ) + name = carried_value.element(register_coordinate) + self.lines.append(f" {name} = {value};") + else: + value = self.expression_operand(yielded) + self.lines.append(f" {carried_value} = {value};") body_lines = self.lines[body_start:] self.lines[body_start:] = [f" {line}" for line in body_lines] self.lines.append(" }") - def emit(self, op: SSAOp) -> None: - if op.opcode == "store": - ptr = self.pointer_operand(op.operands[0]) - value = self.expression_operand(op.operands[1]) - mask_operand = op.operands[2] + def register_broadcast_axes( + self, + ty: Type, + ) -> tuple[int, ...] | None: + if not isinstance(ty, BlockType) or ty.rank != 2: + return None + + broadcast_axes = [] + + for axis, (result_dim, output_dim) in enumerate( + zip( + ty.shape, + self.layout.output_tile_shape, + strict=True, + ) + ): + if result_dim == output_dim: + continue + + if result_dim == 1: + broadcast_axes.append(axis) + continue + + return None + + return tuple(broadcast_axes) + + def emit_binary(self, op: SSAOp, result: SSAValue) -> None: + symbol = self.BINARY_OPS[op.opcode] + broadcast_axes = self.register_broadcast_axes(result.ty) + + if broadcast_axes is None: + lhs = self.expression_operand(op.operands[0]) + rhs = self.expression_operand(op.operands[1]) + self.assign(result, f"({lhs} {symbol} {rhs})") + return + + register_layout = self.layout.register_tile_layout() + + if register_layout.registers_per_thread == 1: + lhs = self.expression_operand(op.operands[0]) + rhs = self.expression_operand(op.operands[1]) + self.assign(result, f"({lhs} {symbol} {rhs})") + return + + result_ref = CudaRegisterTileRef( + base=f"v{result.id}", + layout=register_layout, + broadcast_axes=broadcast_axes, + ) + cuda_ty = self.cuda_type(result.ty) + + for register_coordinate in result_ref.storage_coordinates(): + lhs = self.register_expression_operand( + op.operands[0], + register_coordinate, + register_layout, + ) + rhs = self.register_expression_operand( + op.operands[1], + register_coordinate, + register_layout, + ) + result_element = result_ref.element(register_coordinate) + + self.lines.append( + f" {cuda_ty} {result_element} = ({lhs} {symbol} {rhs});" + ) + + self.values[result.id] = result_ref + + def emit_addptr(self, op: SSAOp, result: SSAValue) -> None: + broadcast_axes = self.register_broadcast_axes(result.ty) + + if broadcast_axes is not None: + register_layout = self.layout.register_tile_layout() + + if register_layout.registers_per_thread > 1: + result_ref = CudaRegisterTileRef( + base=f"v{result.id}", + layout=register_layout, + broadcast_axes=broadcast_axes, + ) + cuda_ty = self.cuda_type(result.ty) + + for register_coordinate in result_ref.storage_coordinates(): + base = self.register_pointer_operand( + op.operands[0], + register_coordinate, + register_layout, + ) + offset = self.register_expression_operand( + op.operands[1], + register_coordinate, + register_layout, + ) + + if base.index != "0": + offset = f"({base.index} + {offset})" + + result_element = result_ref.element(register_coordinate) + self.lines.append( + f" {cuda_ty} {result_element} = {base.base} + {offset};" + ) + + self.values[result.id] = result_ref + return + + scalar_base = self.operand(op.operands[0]) + offset = self.expression_operand(op.operands[1]) + + if isinstance(scalar_base, CudaPtrRef): + if scalar_base.index != "0": + offset = f"({scalar_base.index} + {offset})" + scalar_base = scalar_base.base + + if not isinstance(scalar_base, str): + raise TypeError(f"addptr expects pointer base, got {scalar_base}") + + self.values[result.id] = CudaPtrRef( + scalar_base, + offset, + ) + + def emit_store(self, op: SSAOp) -> None: + pointer_operand, value_operand, mask_operand = op.operands + + operand_values = tuple( + self.operand(operand) for operand in op.operands if operand is not None + ) + register_ref = next( + ( + value + for value in operand_values + if isinstance(value, CudaRegisterTileRef) + ), + None, + ) + + if register_ref is None: + ptr = self.pointer_operand(pointer_operand) + value = self.expression_operand(value_operand) mask = ( None if mask_operand is None else self.expression_operand(mask_operand) ) @@ -621,6 +987,145 @@ def emit(self, op: SSAOp) -> None: ) return + register_layout = register_ref.layout + + for register_coordinate in self.register_coordinates(register_layout): + ptr = self.register_pointer_operand( + pointer_operand, + register_coordinate, + register_layout, + ) + value = self.register_expression_operand( + value_operand, + register_coordinate, + register_layout, + ) + mask = ( + None + if mask_operand is None + else self.register_expression_operand( + mask_operand, + register_coordinate, + register_layout, + ) + ) + + if mask is None: + self.lines.append(f" {ptr.base}[{ptr.index}] = {value};") + else: + self.lines.extend( + [ + f" if ({mask}) {{", + f" {ptr.base}[{ptr.index}] = {value};", + " }", + ] + ) + + def emit_expand_dims( + self, + op: SSAOp, + result: SSAValue, + ) -> None: + if not self.is_rank2_kernel(): + raise TypeError( + "CUDA expand_dims lowering currently requires rank-2 kernel" + ) + + operand = op.operands[0] + if not isinstance(operand, SSAValue): + raise TypeError(f"expand_dims expects SSA operand, got {operand}") + + arange_ref = self.operand(operand) + if not isinstance(arange_ref, CudaArangeRef): + raise TypeError( + "CUDA expand_dims MVP supports only direct arange expansion, " + f"got {arange_ref}" + ) + + axis = op.attrs.get("axis") + if type(axis) is not int: + raise TypeError(f"expand_dims axis must be an integer, got {axis}") + + assert isinstance(axis, int) + + if not isinstance(result.ty, BlockType): + raise TypeError(f"expand_dims expects block result, got {result.ty}") + + result_shape = result.ty.shape + register_layout = self.layout.register_tile_layout() + + if register_layout.registers_per_thread > 1: + expected_shape = list(self.layout.output_tile_shape) + expected_shape[axis] = 1 + + if result_shape != tuple(expected_shape): + raise TypeError( + f"cannot map expand_dims result {result.ty} into " + f"CUDA output tile {self.layout.output_tile_shape}" + ) + + result_ref = CudaRegisterTileRef( + base=f"v{result.id}", + layout=register_layout, + broadcast_axes=(axis,), + ) + cuda_ty = self.cuda_type(result.ty) + source_axis = 1 - axis + + for register_coordinate in result_ref.storage_coordinates(): + logical_coordinates = self.register_logical_coordinates( + register_layout, + register_coordinate, + ) + coordinate = logical_coordinates[source_axis] + expression = ( + coordinate + if arange_ref.start == 0 + else f"({arange_ref.start} + {coordinate})" + ) + result_element = result_ref.element(register_coordinate) + + self.lines.append(f" {cuda_ty} {result_element} = {expression};") + + self.values[result.id] = result_ref + return + + try: + tile_layout = self.layout.tile_layout( + result_shape, + broadcast_axes=(axis,), + ) + except ValueError as error: + raise TypeError( + f"cannot map expand_dims result {result.ty} into CUDA tile " + f"shape {self.layout.thread_shape}" + ) from error + + mapped_axes = [ + thread_axis + for thread_axis in tile_layout.thread_axes + if thread_axis is not None + ] + + if len(mapped_axes) != 1: + raise TypeError( + "expanded arange must map to exactly one CUDA thread axis, " + f"got {tile_layout}" + ) + + coordinate = self.thread_coordinate(mapped_axes[0]) + expression = ( + coordinate + if arange_ref.start == 0 + else f"({arange_ref.start} + {coordinate})" + ) + self.assign(result, expression) + + def emit(self, op: SSAOp) -> None: + if op.opcode == "store": + self.emit_store(op) + return + result = op.result if result is None: raise TypeError(f"SSA opcode {op.opcode!r} requires a result") @@ -667,20 +1172,9 @@ def emit(self, op: SSAOp) -> None: buffers, ) elif op.opcode in self.BINARY_OPS: - lhs = self.expression_operand(op.operands[0]) - rhs = self.expression_operand(op.operands[1]) - symbol = self.BINARY_OPS[op.opcode] - self.assign(result, f"({lhs} {symbol} {rhs})") + self.emit_binary(op, result) elif op.opcode == "addptr": - base = self.operand(op.operands[0]) - offset = self.expression_operand(op.operands[1]) - if isinstance(base, CudaPtrRef): - if base.index != "0": - offset = f"({base.index} + {offset})" - base = base.base - if not isinstance(base, str): - raise TypeError(f"addptr expects pointer base, got {base}") - self.values[result.id] = CudaPtrRef(base, offset) + self.emit_addptr(op, result) elif op.opcode == "load": ptr = self.pointer_operand(op.operands[0]) mask_operand = op.operands[1] @@ -728,62 +1222,7 @@ def emit(self, op: SSAOp) -> None: elif op.opcode in ("sum", "max", "min"): self.emit_reduction(op) elif op.opcode == "expand_dims": - if not self.is_rank2_kernel(): - raise TypeError( - "CUDA expand_dims lowering currently requires rank-2 kernel" - ) - - operand = op.operands[0] - if not isinstance(operand, SSAValue): - raise TypeError(f"expand_dims expects SSA operand, got {operand}") - - arange_ref = self.operand(operand) - if not isinstance(arange_ref, CudaArangeRef): - raise TypeError( - "CUDA expand_dims MVP supports only direct arange expansion, " - f"got {arange_ref}" - ) - - axis = op.attrs.get("axis") - if type(axis) is not int: - raise TypeError(f"expand_dims axis must be an integer, got {axis}") - - assert isinstance(axis, int) - - if not isinstance(result.ty, BlockType): - raise TypeError(f"expand_dims expects block result, got {result.ty}") - - result_shape = result.ty.shape - - try: - tile_layout = self.layout.tile_layout( - result_shape, - broadcast_axes=(axis,), - ) - except ValueError as error: - raise TypeError( - f"cannot map expand_dims result {result.ty} into CUDA tile " - f"shape {self.layout.thread_shape}" - ) from error - - mapped_axes = [ - thread_axis - for thread_axis in tile_layout.thread_axes - if thread_axis is not None - ] - - if len(mapped_axes) != 1: - raise TypeError( - "expanded arange must map to exactly one CUDA thread axis, " - f"got {tile_layout}" - ) - - coord = self.thread_coordinate(mapped_axes[0]) - - expression = ( - coord if arange_ref.start == 0 else f"({arange_ref.start} + {coord})" - ) - self.assign(result, expression) + self.emit_expand_dims(op, result) else: raise TypeError(f"Unsupported SSA opcode: {op.opcode}") diff --git a/tests/test_block_shapes.py b/tests/test_block_shapes.py index 2c2a17d..342db5c 100644 --- a/tests/test_block_shapes.py +++ b/tests/test_block_shapes.py @@ -3,6 +3,7 @@ from mytriton.block_shapes import ( CudaCooperativeTileLayout, CudaKernelLayout, + CudaRegisterTileLayout, CudaTileLayout, cuda_kernel_layout, ) @@ -230,6 +231,50 @@ def test_cuda_kernel_layout_separates_output_tile_from_threads() -> None: assert layout.threads_per_block == 256 +def test_register_tile_layout_maps_multiple_results_per_thread() -> None: + layout = CudaRegisterTileLayout( + logical_shape=(8, 8), + thread_shape=(4, 4), + ) + + assert layout.register_shape == (2, 2) + assert layout.registers_per_thread == 4 + + assert layout.logical_coordinate( + thread_coordinate=(2, 3), + register_coordinate=(0, 0), + ) == (2, 3) + assert layout.logical_coordinate( + thread_coordinate=(2, 3), + register_coordinate=(0, 1), + ) == (2, 7) + assert layout.logical_coordinate( + thread_coordinate=(2, 3), + register_coordinate=(1, 0), + ) == (6, 3) + assert layout.logical_coordinate( + thread_coordinate=(2, 3), + register_coordinate=(1, 1), + ) == (6, 7) + + logical_coordinates = { + layout.logical_coordinate( + thread_coordinate=(thread_row, thread_column), + register_coordinate=(register_row, register_column), + ) + for thread_row in range(layout.thread_shape[0]) + for thread_column in range(layout.thread_shape[1]) + for register_row in range(layout.register_shape[0]) + for register_column in range(layout.register_shape[1]) + } + + assert logical_coordinates == { + (row, column) + for row in range(layout.logical_shape[0]) + for column in range(layout.logical_shape[1]) + } + + def test_cooperative_layout_distributes_a_tile_across_threads() -> None: kernel_layout = CudaKernelLayout( output_tile_shape=(4, 8), @@ -394,3 +439,144 @@ def test_cuda_layouts_reject_more_than_1024_threads() -> None: threads_per_block=1025, order=(1, 0), ) + + +@pytest.mark.parametrize( + ("logical_shape", "thread_shape"), + [ + ((8,), (4,)), + ((8, 8), (4,)), + ((8, 0), (4, 4)), + ((8, 8), (0, 4)), + ((7, 8), (4, 4)), + ((4, 8), (8, 4)), + ], +) +def test_register_tile_layout_rejects_invalid_shapes( + logical_shape: tuple[int, ...], + thread_shape: tuple[int, ...], +) -> None: + with pytest.raises(ValueError): + CudaRegisterTileLayout(logical_shape, thread_shape) + + +def test_register_tile_layout_rejects_invalid_coordinates() -> None: + layout = CudaRegisterTileLayout( + logical_shape=(8, 8), + thread_shape=(4, 4), + ) + + with pytest.raises(ValueError, match="invalid thread coordinate"): + layout.logical_coordinate( + thread_coordinate=(4, 0), + register_coordinate=(0, 0), + ) + + with pytest.raises(ValueError, match="invalid thread coordinate"): + layout.logical_coordinate( + thread_coordinate=(-1, 0), + register_coordinate=(0, 0), + ) + + with pytest.raises(ValueError, match="invalid register coordinate"): + layout.logical_coordinate( + thread_coordinate=(0, 0), + register_coordinate=(2, 0), + ) + + with pytest.raises(ValueError, match="invalid register coordinate"): + layout.logical_coordinate( + thread_coordinate=(0, 0), + register_coordinate=(0, -1), + ) + + +def test_cuda_kernel_layout_builds_output_register_tile_layout() -> None: + kernel_layout = CudaKernelLayout( + output_tile_shape=(8, 8), + thread_shape=(4, 4), + ) + + assert kernel_layout.register_tile_layout() == CudaRegisterTileLayout( + logical_shape=(8, 8), + thread_shape=(4, 4), + ) + + +@pytest.mark.parametrize( + ("output_shape", "expected_thread_shape"), + [ + ((4, 8), (4, 8)), + ((8, 8), (4, 8)), + ((16, 32), (4, 8)), + ], +) +def test_dot_kernel_layout_uses_one_warp_register_tile( + output_shape: tuple[int, ...], + expected_thread_shape: tuple[int, ...], +) -> None: + rows, columns = output_shape + + lhs = SSAValue( + id=0, + ty=BlockType((rows, 16), F32), + ) + rhs = SSAValue( + id=1, + ty=BlockType((16, columns), F32), + ) + dot = SSAValue( + id=2, + ty=BlockType(output_shape, F32), + ) + pointers = SSAValue( + id=3, + ty=BlockType(output_shape, PTR_F32), + ) + + ssa_ops: list[SSAItem] = [ + SSAOp( + opcode="dot", + operands=(lhs, rhs), + result=dot, + ), + SSAOp( + opcode="store", + operands=(pointers, dot, None), + ), + ] + + layout = cuda_kernel_layout(ssa_ops) + + assert layout.output_tile_shape == output_shape + assert layout.thread_shape == expected_thread_shape + assert layout.threads_per_block <= 32 + + register_layout = layout.register_tile_layout() + assert register_layout.logical_shape == output_shape + assert register_layout.thread_shape == expected_thread_shape + + +def test_non_dot_kernel_keeps_one_thread_per_result() -> None: + shape = (8, 8) + pointers = SSAValue( + id=0, + ty=BlockType(shape, PTR_F32), + ) + values = SSAValue( + id=1, + ty=BlockType(shape, F32), + ) + + layout = cuda_kernel_layout( + [ + SSAOp( + opcode="store", + operands=(pointers, values, None), + ) + ] + ) + + assert layout.output_tile_shape == shape + assert layout.thread_shape == shape + assert layout.register_tile_layout().register_shape == (1, 1) diff --git a/tests/test_shared_memory.py b/tests/test_shared_memory.py index 1cf86f2..c52fef3 100644 --- a/tests/test_shared_memory.py +++ b/tests/test_shared_memory.py @@ -6,8 +6,15 @@ import mytriton as triton import mytriton.language as tl from mytriton.ast_frontend import trace as trace_ast -from mytriton.block_shapes import CudaKernelLayout -from mytriton.cuda_codegen import SSACUDACodegen +from mytriton.block_shapes import ( + CudaKernelLayout, + cuda_kernel_layout, +) +from mytriton.cuda_codegen import ( + CudaArangeRef, + CudaRegisterTileRef, + SSACUDACodegen, +) from mytriton.cuda_dot_staging import ( CudaDotOperandMatcher, CudaDotSharedBuffers, @@ -28,6 +35,7 @@ SSAValue, ) from mytriton.trace import ( + BOOL, F32, I32, PTR_F32, @@ -1226,7 +1234,7 @@ def test_cuda_core_dot_rejects_wrong_result_shape() -> None: assert codegen.values == {} -def test_cuda_core_dot_requires_one_thread_per_result_element() -> None: +def test_cuda_core_dot_emits_multiple_accumulators_per_thread() -> None: codegen = SSACUDACodegen() codegen.layout = CudaKernelLayout( output_tile_shape=(4, 8), @@ -1250,14 +1258,967 @@ def test_cuda_core_dot_requires_one_thread_per_result_element() -> None: ), ) + codegen.emit_dot_from_shared_memory(result, buffers) + + assert codegen.lines == [ + " float v7_0_0 = 0.0f;", + " float v7_1_0 = 0.0f;", + " for (int dot_k_7 = 0; dot_k_7 < 16; ++dot_k_7) {", + ( + " v7_0_0 += " + "dot_lhs_7[(tile_i) * 16 + (dot_k_7)] * " + "dot_rhs_7[(dot_k_7) * 8 + (tile_j)];" + ), + ( + " v7_1_0 += " + "dot_lhs_7[(tile_i + 2) * 16 + (dot_k_7)] * " + "dot_rhs_7[(dot_k_7) * 8 + (tile_j)];" + ), + " }", + " __syncthreads();", + ] + + value = codegen.values[result.id] + assert isinstance(value, CudaRegisterTileRef) + assert value.element((0, 0)) == "v7_0_0" + assert value.element((1, 0)) == "v7_1_0" + + +def test_cuda_register_tile_ref_names_multiple_registers() -> None: + register_tile = CudaRegisterTileRef( + base="v7", + layout=CudaKernelLayout( + output_tile_shape=(8, 8), + thread_shape=(4, 4), + ).register_tile_layout(), + ) + + assert register_tile.element((0, 0)) == "v7_0_0" + assert register_tile.element((0, 1)) == "v7_0_1" + assert register_tile.element((1, 0)) == "v7_1_0" + assert register_tile.element((1, 1)) == "v7_1_1" + + +def test_cuda_register_tile_ref_preserves_scalar_name() -> None: + register_tile = CudaRegisterTileRef( + base="v7", + layout=CudaKernelLayout( + output_tile_shape=(4, 8), + thread_shape=(4, 8), + ).register_tile_layout(), + ) + + assert register_tile.element((0, 0)) == "v7" + + +def test_cuda_register_tile_ref_rejects_invalid_coordinate() -> None: + register_tile = CudaRegisterTileRef( + base="v7", + layout=CudaKernelLayout( + output_tile_shape=(8, 8), + thread_shape=(4, 4), + ).register_tile_layout(), + ) + + with pytest.raises(ValueError, match="invalid register coordinate"): + register_tile.element((2, 0)) + + +def test_cuda_codegen_emits_register_wise_binary_operation() -> None: + codegen = SSACUDACodegen() + codegen.layout = CudaKernelLayout( + output_tile_shape=(4, 8), + thread_shape=(2, 8), + ) + + ty = BlockType((4, 8), F32) + accumulator = SSAValue(id=0, ty=ty) + dot = SSAValue(id=1, ty=ty) + result = SSAValue(id=2, ty=ty) + + register_layout = codegen.layout.register_tile_layout() + codegen.values[accumulator.id] = "acc" + codegen.values[dot.id] = CudaRegisterTileRef( + base="v1", + layout=register_layout, + ) + + codegen.emit( + SSAOp( + opcode="add", + operands=(accumulator, dot), + result=result, + ) + ) + + assert codegen.lines == [ + " float v2_0_0 = (acc + v1_0_0);", + " float v2_1_0 = (acc + v1_1_0);", + ] + + value = codegen.values[result.id] + assert isinstance(value, CudaRegisterTileRef) + assert value.element((0, 0)) == "v2_0_0" + assert value.element((1, 0)) == "v2_1_0" + + +def test_cuda_for_range_carries_register_tile() -> None: + codegen = SSACUDACodegen() + codegen.layout = CudaKernelLayout( + output_tile_shape=(4, 8), + thread_shape=(2, 8), + ) + + ty = BlockType((4, 8), F32) + initial = SSAValue(id=0, ty=ty) + increment = SSAValue(id=1, ty=ty) + index = SSAValue(id=2, ty=I32) + carried_arg = SSAValue(id=3, ty=ty) + updated = SSAValue(id=4, ty=ty) + result = SSAValue(id=5, ty=ty) + + register_layout = codegen.layout.register_tile_layout() + codegen.values[initial.id] = "initial" + codegen.values[increment.id] = CudaRegisterTileRef( + base="increment", + layout=register_layout, + ) + + codegen.emit_for_range( + SSAForRange( + index=index, + start=Const(0), + stop=Const(2), + step=Const(1), + carried_inputs=(initial,), + carried_args=(carried_arg,), + body=[ + SSAOp( + opcode="add", + operands=(carried_arg, increment), + result=updated, + ) + ], + yields=(updated,), + results=(result,), + ) + ) + + assert codegen.lines == [ + " float v5_0_0 = initial;", + " float v5_1_0 = initial;", + " for (int v2 = 0; v2 < 2; v2 += 1) {", + " float v4_0_0 = (v5_0_0 + increment_0_0);", + " float v4_1_0 = (v5_1_0 + increment_1_0);", + " v5_0_0 = v4_0_0;", + " v5_1_0 = v4_1_0;", + " }", + ] + + carried_value = codegen.values[carried_arg.id] + result_value = codegen.values[result.id] + + assert isinstance(carried_value, CudaRegisterTileRef) + assert isinstance(result_value, CudaRegisterTileRef) + assert carried_value == result_value + assert result_value.element((0, 0)) == "v5_0_0" + assert result_value.element((1, 0)) == "v5_1_0" + + +def test_cuda_codegen_emits_register_wise_addptr() -> None: + codegen = SSACUDACodegen() + codegen.layout = CudaKernelLayout( + output_tile_shape=(4, 8), + thread_shape=(2, 8), + ) + + offsets = SSAValue( + id=0, + ty=BlockType((4, 8), I32), + ) + pointers = SSAValue( + id=1, + ty=BlockType((4, 8), PTR_F32), + ) + + register_layout = codegen.layout.register_tile_layout() + codegen.values[offsets.id] = CudaRegisterTileRef( + base="offset", + layout=register_layout, + ) + + codegen.emit( + SSAOp( + opcode="addptr", + operands=(Param("out", PTR_F32), offsets), + result=pointers, + ) + ) + + assert codegen.lines == [ + " float* v1_0_0 = out + offset_0_0;", + " float* v1_1_0 = out + offset_1_0;", + ] + + value = codegen.values[pointers.id] + assert isinstance(value, CudaRegisterTileRef) + assert value.element((0, 0)) == "v1_0_0" + assert value.element((1, 0)) == "v1_1_0" + + +def test_cuda_codegen_emits_register_wise_masked_store() -> None: + codegen = SSACUDACodegen() + codegen.layout = CudaKernelLayout( + output_tile_shape=(4, 8), + thread_shape=(2, 8), + ) + + pointers = SSAValue( + id=0, + ty=BlockType((4, 8), PTR_F32), + ) + values = SSAValue( + id=1, + ty=BlockType((4, 8), F32), + ) + mask = SSAValue( + id=2, + ty=BlockType((4, 8), BOOL), + ) + + register_layout = codegen.layout.register_tile_layout() + codegen.values[pointers.id] = CudaRegisterTileRef( + base="pointer", + layout=register_layout, + ) + codegen.values[values.id] = CudaRegisterTileRef( + base="value", + layout=register_layout, + ) + codegen.values[mask.id] = CudaRegisterTileRef( + base="mask", + layout=register_layout, + ) + + codegen.emit( + SSAOp( + opcode="store", + operands=(pointers, values, mask), + ) + ) + + assert codegen.lines == [ + " if (mask_0_0) {", + " pointer_0_0[0] = value_0_0;", + " }", + " if (mask_1_0) {", + " pointer_1_0[0] = value_1_0;", + " }", + ] + + +def test_cuda_codegen_rejects_store_with_incompatible_register_layouts() -> None: + codegen = SSACUDACodegen() + codegen.layout = CudaKernelLayout( + output_tile_shape=(4, 8), + thread_shape=(2, 8), + ) + + pointers = SSAValue( + id=0, + ty=BlockType((4, 8), PTR_F32), + ) + values = SSAValue( + id=1, + ty=BlockType((4, 8), F32), + ) + + codegen.values[pointers.id] = CudaRegisterTileRef( + base="pointer", + layout=codegen.layout.register_tile_layout(), + ) + codegen.values[values.id] = CudaRegisterTileRef( + base="value", + layout=CudaKernelLayout( + output_tile_shape=(4, 8), + thread_shape=(4, 4), + ).register_tile_layout(), + ) + with pytest.raises( TypeError, - match="requires one CUDA thread per result element", + match="incompatible CUDA register tile layouts", ): - codegen.emit_dot_from_shared_memory( - result, - buffers, + codegen.emit( + SSAOp( + opcode="store", + operands=(pointers, values, None), + ) ) - assert codegen.lines == [] - assert codegen.values == {} + +def test_cuda_register_tile_ref_broadcasts_register_axes() -> None: + register_layout = CudaKernelLayout( + output_tile_shape=(8, 8), + thread_shape=(4, 4), + ).register_tile_layout() + + rows = CudaRegisterTileRef( + base="rows", + layout=register_layout, + broadcast_axes=(1,), + ) + columns = CudaRegisterTileRef( + base="columns", + layout=register_layout, + broadcast_axes=(0,), + ) + + assert rows.storage_shape == (2, 1) + assert rows.storage_coordinates() == ( + (0, 0), + (1, 0), + ) + assert rows.element((0, 0)) == "rows_0_0" + assert rows.element((0, 1)) == "rows_0_0" + assert rows.element((1, 0)) == "rows_1_0" + assert rows.element((1, 1)) == "rows_1_0" + + assert columns.storage_shape == (1, 2) + assert columns.storage_coordinates() == ( + (0, 0), + (0, 1), + ) + assert columns.element((0, 0)) == "columns_0_0" + assert columns.element((1, 0)) == "columns_0_0" + assert columns.element((0, 1)) == "columns_0_1" + assert columns.element((1, 1)) == "columns_0_1" + + +def test_cuda_codegen_expands_aranges_across_register_tile() -> None: + codegen = SSACUDACodegen() + codegen.layout = CudaKernelLayout( + output_tile_shape=(8, 8), + thread_shape=(4, 4), + ) + + row_arange = SSAValue( + id=0, + ty=BlockType((8,), I32), + ) + rows = SSAValue( + id=1, + ty=BlockType((8, 1), I32), + ) + column_arange = SSAValue( + id=2, + ty=BlockType((8,), I32), + ) + columns = SSAValue( + id=3, + ty=BlockType((1, 8), I32), + ) + + codegen.values[row_arange.id] = CudaArangeRef( + start=0, + end=8, + ) + codegen.values[column_arange.id] = CudaArangeRef( + start=0, + end=8, + ) + + codegen.emit( + SSAOp( + opcode="expand_dims", + operands=(row_arange,), + result=rows, + attrs={"axis": 1}, + ) + ) + codegen.emit( + SSAOp( + opcode="expand_dims", + operands=(column_arange,), + result=columns, + attrs={"axis": 0}, + ) + ) + + assert codegen.lines == [ + " int v1_0_0 = tile_i;", + " int v1_1_0 = tile_i + 4;", + " int v3_0_0 = tile_j;", + " int v3_0_1 = tile_j + 4;", + ] + + row_value = codegen.values[rows.id] + column_value = codegen.values[columns.id] + + assert isinstance(row_value, CudaRegisterTileRef) + assert isinstance(column_value, CudaRegisterTileRef) + + assert row_value.broadcast_axes == (1,) + assert row_value.element((1, 1)) == "v1_1_0" + + assert column_value.broadcast_axes == (0,) + assert column_value.element((1, 1)) == "v3_0_1" + + +def test_cuda_binary_operations_preserve_and_expand_register_broadcasts() -> None: + codegen = SSACUDACodegen() + codegen.layout = CudaKernelLayout( + output_tile_shape=(8, 8), + thread_shape=(4, 4), + ) + + rows = SSAValue( + id=0, + ty=BlockType((8, 1), I32), + ) + scaled_rows = SSAValue( + id=1, + ty=BlockType((8, 1), I32), + ) + columns = SSAValue( + id=2, + ty=BlockType((1, 8), I32), + ) + offsets = SSAValue( + id=3, + ty=BlockType((8, 8), I32), + ) + + register_layout = codegen.layout.register_tile_layout() + codegen.values[rows.id] = CudaRegisterTileRef( + base="rows", + layout=register_layout, + broadcast_axes=(1,), + ) + codegen.values[columns.id] = CudaRegisterTileRef( + base="columns", + layout=register_layout, + broadcast_axes=(0,), + ) + + codegen.emit( + SSAOp( + opcode="mul", + operands=(rows, Param("stride", I32)), + result=scaled_rows, + ) + ) + codegen.emit( + SSAOp( + opcode="add", + operands=(scaled_rows, columns), + result=offsets, + ) + ) + + assert codegen.lines == [ + " int v1_0_0 = (rows_0_0 * stride);", + " int v1_1_0 = (rows_1_0 * stride);", + " int v3_0_0 = (v1_0_0 + columns_0_0);", + " int v3_0_1 = (v1_0_0 + columns_0_1);", + " int v3_1_0 = (v1_1_0 + columns_0_0);", + " int v3_1_1 = (v1_1_0 + columns_0_1);", + ] + + scaled_value = codegen.values[scaled_rows.id] + offset_value = codegen.values[offsets.id] + + assert isinstance(scaled_value, CudaRegisterTileRef) + assert isinstance(offset_value, CudaRegisterTileRef) + + assert scaled_value.broadcast_axes == (1,) + assert scaled_value.element((1, 1)) == "v1_1_0" + + assert offset_value.broadcast_axes == () + assert offset_value.element((1, 1)) == "v3_1_1" + + +def test_cuda_binary_operations_build_register_wise_mask() -> None: + codegen = SSACUDACodegen() + codegen.layout = CudaKernelLayout( + output_tile_shape=(8, 8), + thread_shape=(4, 4), + ) + + rows = SSAValue( + id=0, + ty=BlockType((8, 1), I32), + ) + columns = SSAValue( + id=1, + ty=BlockType((1, 8), I32), + ) + row_mask = SSAValue( + id=2, + ty=BlockType((8, 1), BOOL), + ) + column_mask = SSAValue( + id=3, + ty=BlockType((1, 8), BOOL), + ) + mask = SSAValue( + id=4, + ty=BlockType((8, 8), BOOL), + ) + + register_layout = codegen.layout.register_tile_layout() + codegen.values[rows.id] = CudaRegisterTileRef( + base="rows", + layout=register_layout, + broadcast_axes=(1,), + ) + codegen.values[columns.id] = CudaRegisterTileRef( + base="columns", + layout=register_layout, + broadcast_axes=(0,), + ) + + codegen.emit( + SSAOp( + opcode="cmp_lt", + operands=(rows, Param("M", I32)), + result=row_mask, + ) + ) + codegen.emit( + SSAOp( + opcode="cmp_lt", + operands=(columns, Param("N", I32)), + result=column_mask, + ) + ) + codegen.emit( + SSAOp( + opcode="and", + operands=(row_mask, column_mask), + result=mask, + ) + ) + + assert codegen.lines == [ + " bool v2_0_0 = (rows_0_0 < M);", + " bool v2_1_0 = (rows_1_0 < M);", + " bool v3_0_0 = (columns_0_0 < N);", + " bool v3_0_1 = (columns_0_1 < N);", + " bool v4_0_0 = (v2_0_0 && v3_0_0);", + " bool v4_0_1 = (v2_0_0 && v3_0_1);", + " bool v4_1_0 = (v2_1_0 && v3_0_0);", + " bool v4_1_1 = (v2_1_0 && v3_0_1);", + ] + + +def test_cuda_addptr_preserves_and_expands_register_broadcasts() -> None: + codegen = SSACUDACodegen() + codegen.layout = CudaKernelLayout( + output_tile_shape=(8, 8), + thread_shape=(4, 4), + ) + + row_offsets = SSAValue( + id=0, + ty=BlockType((8, 1), I32), + ) + column_offsets = SSAValue( + id=1, + ty=BlockType((1, 8), I32), + ) + row_pointers = SSAValue( + id=2, + ty=BlockType((8, 1), PTR_F32), + ) + output_pointers = SSAValue( + id=3, + ty=BlockType((8, 8), PTR_F32), + ) + + register_layout = codegen.layout.register_tile_layout() + codegen.values[row_offsets.id] = CudaRegisterTileRef( + base="rows", + layout=register_layout, + broadcast_axes=(1,), + ) + codegen.values[column_offsets.id] = CudaRegisterTileRef( + base="columns", + layout=register_layout, + broadcast_axes=(0,), + ) + + codegen.emit( + SSAOp( + opcode="addptr", + operands=(Param("out", PTR_F32), row_offsets), + result=row_pointers, + ) + ) + codegen.emit( + SSAOp( + opcode="addptr", + operands=(row_pointers, column_offsets), + result=output_pointers, + ) + ) + + assert codegen.lines == [ + " float* v2_0_0 = out + rows_0_0;", + " float* v2_1_0 = out + rows_1_0;", + " float* v3_0_0 = v2_0_0 + columns_0_0;", + " float* v3_0_1 = v2_0_0 + columns_0_1;", + " float* v3_1_0 = v2_1_0 + columns_0_0;", + " float* v3_1_1 = v2_1_0 + columns_0_1;", + ] + + row_pointer_value = codegen.values[row_pointers.id] + output_pointer_value = codegen.values[output_pointers.id] + + assert isinstance(row_pointer_value, CudaRegisterTileRef) + assert isinstance(output_pointer_value, CudaRegisterTileRef) + + assert row_pointer_value.broadcast_axes == (1,) + assert row_pointer_value.element((1, 1)) == "v2_1_0" + + assert output_pointer_value.broadcast_axes == () + assert output_pointer_value.element((1, 1)) == "v3_1_1" + + +def test_ast_frontend_lowers_multi_result_dot_to_cuda( + monkeypatch, +) -> None: + monkeypatch.setenv("MYTRITON_FRONTEND", "ast") + monkeypatch.setenv("MYTRITON_BACKEND", "cuda") + + M = N = K = 8 + BM = BK = BN = 8 + + a = np.zeros((M, K), dtype=np.float32) + b = np.zeros((K, N), dtype=np.float32) + out = np.zeros((M, N), dtype=np.float32) + + single_tile_dot_kernel.clear_cache() + + _, ssa_ops, cuda_src = single_tile_dot_kernel[(1, 1)]( + a, + b, + out, + M, + N, + K, + 0, + BM=BM, + BK=BK, + BN=BN, + ) + + layout = cuda_kernel_layout(ssa_ops) + + assert layout.output_tile_shape == (8, 8) + assert layout.thread_shape == (4, 8) + assert layout.threads_per_block == 32 + assert layout.register_tile_layout().register_shape == (2, 1) + + dot = next(op for op in ssa_ops if isinstance(op, SSAOp) and op.opcode == "dot") + store = next(op for op in ssa_ops if isinstance(op, SSAOp) and op.opcode == "store") + + assert dot.result is not None + dot_id = dot.result.id + + pointer = store.operands[0] + value = store.operands[1] + mask = store.operands[2] + + assert isinstance(pointer, SSAValue) + assert isinstance(value, SSAValue) + assert isinstance(mask, SSAValue) + assert value == dot.result + + assert " int tile_i = threadIdx.x / 8;" in cuda_src + assert " int tile_j = threadIdx.x % 8;" in cuda_src + + assert f" float v{dot_id}_0_0 = 0.0f;" in cuda_src + assert f" float v{dot_id}_1_0 = 0.0f;" in cuda_src + + assert ( + f" v{dot_id}_0_0 += " + f"dot_lhs_{dot_id}[(tile_i) * 8 + (dot_k_{dot_id})] * " + f"dot_rhs_{dot_id}[(dot_k_{dot_id}) * 8 + (tile_j)];" + ) in cuda_src + + assert ( + f" v{dot_id}_1_0 += " + f"dot_lhs_{dot_id}[(tile_i + 4) * 8 + (dot_k_{dot_id})] * " + f"dot_rhs_{dot_id}[(dot_k_{dot_id}) * 8 + (tile_j)];" + ) in cuda_src + + assert ( + f" if (v{mask.id}_0_0) {{\n" + f" v{pointer.id}_0_0[0] = v{dot_id}_0_0;\n" + " }" + ) in cuda_src + + assert ( + f" if (v{mask.id}_1_0) {{\n" + f" v{pointer.id}_1_0[0] = v{dot_id}_1_0;\n" + " }" + ) in cuda_src + + +@pytest.mark.execution +def test_shared_memory_dot_executes_multiple_results_per_thread( + cp, + monkeypatch, +) -> None: + monkeypatch.setenv("MYTRITON_FRONTEND", "ast") + monkeypatch.setenv("MYTRITON_BACKEND", "cuda") + + M, N, K = 13, 11, 7 + BM, BK, BN = 8, 8, 8 + + a_host = np.arange(M * K, dtype=np.float32).reshape(M, K) % 9 - 4 + b_host = np.arange(K * N, dtype=np.float32).reshape(K, N) % 7 - 3 + + a = cp.asarray(a_host) + b = cp.asarray(b_host) + out = cp.zeros((M, N), dtype=cp.float32) + + grid = ( + (M + BM - 1) // BM, + (N + BN - 1) // BN, + ) + + single_tile_dot_kernel.clear_cache() + single_tile_dot_kernel[grid]( + a, + b, + out, + M, + N, + K, + 0, + BM=BM, + BK=BK, + BN=BN, + ) + cp.cuda.Stream.null.synchronize() + + expected = a_host @ b_host + cp.testing.assert_allclose( + out, + cp.asarray(expected), + rtol=1e-5, + atol=1e-5, + ) + + +def test_runtime_k_loop_carries_multiple_results_per_thread( + monkeypatch, +) -> None: + monkeypatch.setenv("MYTRITON_FRONTEND", "ast") + monkeypatch.setenv("MYTRITON_BACKEND", "cuda") + + M, N, K = 8, 8, 16 + BM, BK, BN = 8, 8, 8 + + a = np.zeros((M, K), dtype=np.float32) + b = np.zeros((K, N), dtype=np.float32) + out = np.zeros((M, N), dtype=np.float32) + + tiled_matmul_kernel.clear_cache() + + _, ssa_ops, cuda_src = tiled_matmul_kernel[(1, 1)]( + a, + b, + out, + M, + N, + K, + BM=BM, + BK=BK, + BN=BN, + ) + + layout = cuda_kernel_layout(ssa_ops) + + assert layout.output_tile_shape == (8, 8) + assert layout.thread_shape == (4, 8) + assert layout.register_tile_layout().register_shape == (2, 1) + + loop = next(op for op in ssa_ops if isinstance(op, SSAForRange)) + dot = next(op for op in loop.body if isinstance(op, SSAOp) and op.opcode == "dot") + accumulation = next( + op + for op in loop.body + if (isinstance(op, SSAOp) and op.opcode == "add" and dot.result in op.operands) + ) + + assert dot.result is not None + assert accumulation.result is not None + assert len(loop.carried_inputs) == 1 + assert len(loop.results) == 1 + + initial = loop.carried_inputs[0] + assert isinstance(initial, SSAValue) + + dot_id = dot.result.id + accumulation_id = accumulation.result.id + accumulator_id = loop.results[0].id + index_id = loop.index.id + + assert (f" float v{accumulator_id}_0_0 = v{initial.id};") in cuda_src + assert (f" float v{accumulator_id}_1_0 = v{initial.id};") in cuda_src + + outer_loop = ( + f" for (int v{index_id} = 0; v{index_id} < K; v{index_id} += {BK}) {{" + ) + assert outer_loop in cuda_src + + assert f" float v{dot_id}_0_0 = 0.0f;" in cuda_src + assert f" float v{dot_id}_1_0 = 0.0f;" in cuda_src + + assert ( + f" v{dot_id}_1_0 += " + f"dot_lhs_{dot_id}[(tile_i + 4) * {BK} + " + f"(dot_k_{dot_id})] * " + f"dot_rhs_{dot_id}[(dot_k_{dot_id}) * {BN} + " + "(tile_j)];" + ) in cuda_src + + expected_update = ( + f" float v{accumulation_id}_0_0 = " + f"(v{accumulator_id}_0_0 + v{dot_id}_0_0);\n" + f" float v{accumulation_id}_1_0 = " + f"(v{accumulator_id}_1_0 + v{dot_id}_1_0);\n" + f" v{accumulator_id}_0_0 = " + f"v{accumulation_id}_0_0;\n" + f" v{accumulator_id}_1_0 = " + f"v{accumulation_id}_1_0;" + ) + assert expected_update in cuda_src + + assert cuda_src.count("__syncthreads();") == 2 + + +@pytest.mark.execution +def test_tiled_matmul_executes_multiple_results_per_thread( + cp, + monkeypatch, +) -> None: + monkeypatch.setenv("MYTRITON_FRONTEND", "ast") + monkeypatch.setenv("MYTRITON_BACKEND", "cuda") + + M, N, K = 13, 11, 19 + BM, BK, BN = 8, 8, 8 + + a_host = np.arange(M * K, dtype=np.float32).reshape(M, K) % 9 - 4 + b_host = np.arange(K * N, dtype=np.float32).reshape(K, N) % 7 - 3 + + a = cp.asarray(a_host) + b = cp.asarray(b_host) + out = cp.zeros((M, N), dtype=cp.float32) + + grid = ( + (M + BM - 1) // BM, + (N + BN - 1) // BN, + ) + + tiled_matmul_kernel.clear_cache() + tiled_matmul_kernel[grid]( + a, + b, + out, + M, + N, + K, + BM=BM, + BK=BK, + BN=BN, + ) + cp.cuda.Stream.null.synchronize() + + expected = a_host @ b_host + cp.testing.assert_allclose( + out, + cp.asarray(expected), + rtol=1e-5, + atol=1e-5, + ) + + +@pytest.mark.parametrize( + "broadcast_axes", + [ + (1, 0), + (0, 0), + (-1,), + (2,), + ], +) +def test_cuda_register_tile_ref_rejects_invalid_broadcast_axes( + broadcast_axes: tuple[int, ...], +) -> None: + layout = CudaKernelLayout( + output_tile_shape=(8, 8), + thread_shape=(4, 4), + ).register_tile_layout() + + with pytest.raises(ValueError, match="broadcast axes"): + CudaRegisterTileRef( + base="value", + layout=layout, + broadcast_axes=broadcast_axes, + ) + + +@pytest.mark.execution +def test_tiled_matmul_executes_two_dimensional_register_tile( + cp, + monkeypatch, +) -> None: + monkeypatch.setenv("MYTRITON_FRONTEND", "ast") + monkeypatch.setenv("MYTRITON_BACKEND", "cuda") + + M, N, K = 17, 19, 15 + BM, BK, BN = 16, 8, 16 + + a_host = np.arange(M * K, dtype=np.float32).reshape(M, K) % 9 - 4 + b_host = np.arange(K * N, dtype=np.float32).reshape(K, N) % 7 - 3 + + a = cp.asarray(a_host) + b = cp.asarray(b_host) + out = cp.zeros((M, N), dtype=cp.float32) + + grid = ( + (M + BM - 1) // BM, + (N + BN - 1) // BN, + ) + + tiled_matmul_kernel.clear_cache() + _, ssa_ops, _ = tiled_matmul_kernel[grid]( + a, + b, + out, + M, + N, + K, + BM=BM, + BK=BK, + BN=BN, + ) + cp.cuda.Stream.null.synchronize() + + layout = cuda_kernel_layout(ssa_ops) + + assert layout.output_tile_shape == (16, 16) + assert layout.thread_shape == (4, 8) + assert layout.register_tile_layout().register_shape == (4, 2) + + expected = a_host @ b_host + cp.testing.assert_allclose( + out, + cp.asarray(expected), + rtol=1e-5, + atol=1e-5, + )