From c7e5f1c823972c4a02bf031ae2f616cfab7bd363 Mon Sep 17 00:00:00 2001 From: Pavel Belevich Date: Sat, 22 Aug 2026 22:20:04 -0400 Subject: [PATCH] Version 14: Add shared-memory tiles --- README.md | 62 ++- src/mytriton/cuda_codegen.py | 250 ++++++++- src/mytriton/cuda_dot_staging.py | 426 +++++++++++++++ tests/test_shared_memory.py | 889 +++++++++++++++++++++++++++++++ 4 files changed, 1620 insertions(+), 7 deletions(-) create mode 100644 src/mytriton/cuda_dot_staging.py create mode 100644 tests/test_shared_memory.py diff --git a/README.md b/README.md index 0ca74f7..47bd625 100644 --- a/README.md +++ b/README.md @@ -54,6 +54,11 @@ MLIR's GPU/NVVM stack to a cubin. `[M, K] x [K, N] -> [M, N]` type inference, independent SSA verification, optimizer purity rules, and an explicit diagnostic for the not-yet-implemented CUDA lowering. +- [ver14](https://github.com/pbelevich/mytriton/tree/ver14): CUDA shared-memory + tile buffers, SSA pattern matching for canonical masked matrix loads, + cooperative staging of `tl.dot` operands, zero-filled boundary handling, + block synchronization, and an explicit diagnostic for the deferred + CUDA-core dot computation. ## AST frontend @@ -243,6 +248,52 @@ Version 13 defines the language and IR semantics only. CUDA lowering for cooperative shared-memory tiles and then lower `dot` to an ordinary CUDA-core multiply-accumulate loop. +## Shared-memory dot staging + +The CUDA backend recognizes canonical matrix tiles loaded for `tl.dot`: + +```python +a_values = tl.load( + a + a_rows * K + a_columns, + mask=(a_rows < M) & (a_columns < K), + other=0.0, +) +b_values = tl.load( + b + b_rows * N + b_columns, + mask=(b_rows < K) & (b_columns < N), + other=0.0, +) +result = tl.dot(a_values, b_values) +``` + +The supported pointer form is `base + rows * row_stride + columns`. Rows and +columns must be built from a scalar tile offset plus an expanded +`tl.arange(0, size)`. Each load must use a two-dimensional bounds mask and +`other=0.0`. + +The CUDA staging analysis follows the SSA use-def graph backwards from both +`dot` operands. Operations used only to describe the matrix tiles are removed +from ordinary per-thread scalar lowering, while scalar tile origins and values +shared with the output address remain available. + +Each CUDA thread copies linear shared-memory positions + +```text +threadIdx.x +threadIdx.x + threads_per_block +threadIdx.x + 2 * threads_per_block +... +``` + +until the complete A `[BM, BK]` or B `[BK, BN]` tile has been covered. Logical +row and column coordinates determine the global matrix address. Out-of-bounds +positions receive `0.0`, so edge tiles are safe without divergent barriers. +After both cooperative loads, the backend emits one `__syncthreads()`. + +Version 14 intentionally stops after staging and reports that CUDA computation +for `tl.dot` is not implemented. Version 15 will read the shared buffers, +perform the ordinary CUDA-core FMA loop over `BK`, and store the result tile. + ## Example ```python @@ -472,9 +523,11 @@ these rewrite passes because they are not region-aware yet. rank-1 and rank-2 logical blocks. Reduction lowering internally emits the CUDA shared-memory scratch buffers and synchronization needed for block-local reductions. Floating-point elementwise extrema propagate NaNs and choose the - right-hand operand when values compare equal. `tl.dot` is represented and - verified in SSA, but the CUDA backend intentionally rejects it until - shared-memory tile lowering is implemented. + right-hand operand when values compare equal. For canonical matrix-load + operands, `tl.dot` lowering emits shared-memory declarations, cooperative + masked loads with zero-filled boundaries, and a block barrier. It then + intentionally rejects the missing CUDA-core dot computation instead of + returning an incorrect kernel. - Reductions are currently single-block reductions over the SSA vector width. The vector width must be a power of two and must match the CUDA thread block size. Larger rows can be handled by statically unrolling multiple loads into @@ -487,7 +540,8 @@ these rewrite passes because they are not region-aware yet. CUDA loop. Cooperative tile layouts can describe A `[BM, BK]` and B `[BK, BN]` independently of C `[BM, BN]`. The `tl.dot` operation now has expression-tree, typed SSA, verification, and optimization semantics, but it - does not yet have CUDA lowering or shared-memory staging for its operands. + can only stage canonical load operands cooperatively in CUDA shared memory; + the multiply-accumulate computation remains deferred to Version 15. `tl.empty`, `tl.full`, and `tl.zeros` continue to represent logical per-thread values rather than shared-memory allocations. - MLIR lowering currently supports only `ptr` parameters as diff --git a/src/mytriton/cuda_codegen.py b/src/mytriton/cuda_codegen.py index 8e68011..ddd53d1 100644 --- a/src/mytriton/cuda_codegen.py +++ b/src/mytriton/cuda_codegen.py @@ -3,6 +3,17 @@ from typing import ClassVar from .block_shapes import CudaKernelLayout, cuda_kernel_layout +from .cuda_dot_staging import ( + CudaDotSharedBuffers, + CudaDotStagingAnalysis, + CudaDotStagingAnalyzer, + CudaDotStagingPlan, + CudaGlobalTile, + CudaGlobalTilePlan, + CudaSharedBuffer, + SSADefinitions, + cuda_scalar_nbytes, +) from .ssa import SSAForRange, SSAItem, SSAOp, SSAOperand, SSAValue from .trace import ( BOOL, @@ -34,6 +45,8 @@ def width(self) -> int: class SSACUDACodegen: + MAX_SHARED_MEMORY_BYTES: ClassVar[int] = 48 * 1024 + BINARY_OPS: ClassVar[dict[str, str]] = { "add": "+", "sub": "-", @@ -51,6 +64,12 @@ def __init__(self): thread_shape=(1,), ) self.shared_lines: list[str] = [] + self.shared_memory_bytes = 0 + self.definitions = SSADefinitions([]) + self.staging_analysis = CudaDotStagingAnalysis( + dot_plans={}, + staging_only_ids=frozenset(), + ) def cuda_type(self, ty: Type) -> str: if isinstance(ty, BlockType): @@ -133,6 +152,215 @@ def declare(self, result: SSAValue) -> None: self.lines.append(f" {self.cuda_type(result.ty)} {name};") self.values[result.id] = name + def reserve_shared_memory(self, additional_bytes: int) -> None: + required_bytes = self.shared_memory_bytes + additional_bytes + if required_bytes > self.MAX_SHARED_MEMORY_BYTES: + raise ValueError( + f"CUDA shared memory requires {required_bytes} bytes, " + f"exceeding the conservative {self.MAX_SHARED_MEMORY_BYTES}-byte limit" + ) + + self.shared_memory_bytes = required_bytes + + def append_shared_buffer_declaration(self, buffer: CudaSharedBuffer) -> None: + cuda_ty = self.cuda_type(buffer.element_ty) + self.shared_lines.append( + f" __shared__ {cuda_ty} {buffer.name}[{buffer.size}];" + ) + + def declare_shared_buffer( + self, + name: str, + logical_shape: tuple[int, ...], + element_ty: ScalarType, + ) -> CudaSharedBuffer: + buffer = CudaSharedBuffer( + name=name, + logical_shape=logical_shape, + element_ty=element_ty, + ) + + self.reserve_shared_memory(buffer.nbytes) + self.append_shared_buffer_declaration(buffer) + + return buffer + + def emit_cooperative_load( + self, + target: CudaSharedBuffer, + source: CudaGlobalTile, + *, + order: tuple[int, ...] = (1, 0), + ) -> None: + cooperative_layout = self.layout.cooperative_tile_layout( + target.logical_shape, + order=order, + ) + if cooperative_layout.order != (1, 0): + raise ValueError( + "CUDA cooperative dot loads require row-major order (1, 0), " + f"got {cooperative_layout.order}" + ) + + index = f"{target.name}_index" + row = f"{target.name}_row" + column = f"{target.name}_column" + global_row = f"{target.name}_global_row" + global_column = f"{target.name}_global_column" + source_index = f"{target.name}_source_index" + in_bounds = f"{target.name}_in_bounds" + + self.lines.extend( + [ + ( + f" for (int {index} = threadIdx.x; " + f"{index} < {cooperative_layout.size}; " + f"{index} += {cooperative_layout.threads_per_block}) {{" + ), + f" int {row} = {index} / {target.columns};", + f" int {column} = {index} % {target.columns};", + (f" int {global_row} = ({source.row_offset}) + {row};"), + (f" int {global_column} = ({source.column_offset}) + {column};"), + ( + f" int {source_index} = " + f"{global_row} * ({source.row_stride}) + {global_column};" + ), + ( + f" bool {in_bounds} = " + f"{global_row} < ({source.row_bound}) && " + f"{global_column} < ({source.column_bound});" + ), + ( + f" {target.element(row, column)} = " + f"{in_bounds} ? " + f"{source.base}[{source_index}] : {source.other};" + ), + " }", + ] + ) + + def emit_block_barrier(self) -> None: + self.lines.append(" __syncthreads();") + + def emit_dot_operand_staging( + self, + dot_result_id: int, + lhs_shape: tuple[int, ...], + rhs_shape: tuple[int, ...], + element_ty: ScalarType, + lhs_source: CudaGlobalTile, + rhs_source: CudaGlobalTile, + ) -> CudaDotSharedBuffers: + if len(lhs_shape) != 2 or len(rhs_shape) != 2 or lhs_shape[1] != rhs_shape[0]: + raise ValueError( + "dot staging expects compatible rank-2 operands, " + f"got {lhs_shape} and {rhs_shape}" + ) + + lhs = CudaSharedBuffer( + name=f"dot_lhs_{dot_result_id}", + logical_shape=lhs_shape, + element_ty=element_ty, + ) + rhs = CudaSharedBuffer( + name=f"dot_rhs_{dot_result_id}", + logical_shape=rhs_shape, + element_ty=element_ty, + ) + + # Reserve both operands before mutating the generated CUDA fragment. + self.reserve_shared_memory(lhs.nbytes + rhs.nbytes) + self.append_shared_buffer_declaration(lhs) + self.append_shared_buffer_declaration(rhs) + + self.emit_cooperative_load(lhs, lhs_source) + self.emit_cooperative_load(rhs, rhs_source) + self.emit_block_barrier() + + return CudaDotSharedBuffers( + lhs=lhs, + rhs=rhs, + ) + + def resolve_global_tile( + self, + plan: CudaGlobalTilePlan, + ) -> CudaGlobalTile: + base = self.pointer_operand(plan.base) + + if base.index != "0": + raise TypeError( + f"dot staging expects an unmodified global base pointer, got {base}" + ) + + return CudaGlobalTile( + base=base.base, + row_offset=self.expression_operand(plan.row_offset), + column_offset=self.expression_operand(plan.column_offset), + row_stride=self.expression_operand(plan.row_stride), + row_bound=self.expression_operand(plan.row_bound), + column_bound=self.expression_operand(plan.column_bound), + other=self.expression_operand(plan.other), + ) + + def emit_dot_operand_staging_from_ssa( + self, + op: SSAOp, + plan: CudaDotStagingPlan, + ) -> CudaDotSharedBuffers: + if op.opcode != "dot": + raise TypeError( + f"expected dot operation for shared staging, got {op.opcode}" + ) + + if op.result is None: + raise TypeError("dot operation requires a result") + + lhs, rhs = op.operands + + if not isinstance(lhs, SSAValue) or not isinstance( + lhs.ty, + BlockType, + ): + raise TypeError(f"dot lhs must be a block SSA value, got {lhs}") + + if not isinstance(rhs, SSAValue) or not isinstance( + rhs.ty, + BlockType, + ): + raise TypeError(f"dot rhs must be a block SSA value, got {rhs}") + + element_ty = lhs.ty.element + if not isinstance(element_ty, ScalarType): + raise TypeError( + f"dot shared-memory element must be scalar, got {element_ty}" + ) + + if rhs.ty.element != element_ty: + raise TypeError( + "dot shared-memory operands must have matching elements, " + f"got {lhs.ty.element} and {rhs.ty.element}" + ) + + # Resolve every operand before mutating shared_lines/lines. + lhs_source = self.resolve_global_tile(plan.lhs) + rhs_source = self.resolve_global_tile(plan.rhs) + + return self.emit_dot_operand_staging( + dot_result_id=op.result.id, + lhs_shape=lhs.ty.shape, + rhs_shape=rhs.ty.shape, + element_ty=element_ty, + lhs_source=lhs_source, + rhs_source=rhs_source, + ) + + def is_staging_only(self, op: SSAOp) -> bool: + return ( + op.result is not None + and op.result.id in self.staging_analysis.staging_only_ids + ) + def scalar_type(self, ty: Type) -> ScalarType | PointerType: return ty.element if isinstance(ty, BlockType) else ty @@ -215,6 +443,9 @@ def emit_reduction(self, op: SSAOp) -> None: value = self.expression_operand(operand) element_ty = input_ty.element + if not isinstance(element_ty, ScalarType): + raise TypeError(f"{op.opcode} expects scalar elements, got {element_ty}") + cuda_ty = self.cuda_type(element_ty) width = input_ty.size if width & (width - 1): @@ -223,6 +454,7 @@ def emit_reduction(self, op: SSAOp) -> None: shared = f"reduce_smem_{result.id}" stride = f"stride_{result.id}" + self.reserve_shared_memory(width * cuda_scalar_nbytes(element_ty)) self.shared_lines.append(f" __shared__ {cuda_ty} {shared}[{width}];") self.lines.extend( @@ -285,7 +517,7 @@ def emit_for_range(self, loop: SSAForRange) -> None: for body_op in loop.body: if isinstance(body_op, SSAForRange): self.emit_for_range(body_op) - else: + elif not self.is_staging_only(body_op): self.emit(body_op) for yielded, carried_name in zip(loop.yields, carried_names, strict=True): @@ -351,7 +583,16 @@ def emit(self, op: SSAOp) -> None: zero = False if element_ty == BOOL else 0.0 if element_ty == F32 else 0 self.assign(result, self.literal(zero)) elif op.opcode == "dot": - raise TypeError("CUDA lowering for tl.dot is not implemented") + if result.id not in self.staging_analysis.stageable_dot_ids: + raise TypeError("CUDA lowering for tl.dot is not implemented") + + plan = self.staging_analysis.plan_for(result.id) + self.emit_dot_operand_staging_from_ssa(op, plan) + + raise TypeError( + "CUDA shared-memory staging for tl.dot is implemented, " + "but CUDA computation for tl.dot is not implemented" + ) elif op.opcode in self.BINARY_OPS: lhs = self.expression_operand(op.operands[0]) rhs = self.expression_operand(op.operands[1]) @@ -481,7 +722,10 @@ def generate( ) -> str: self.lines = [] self.shared_lines = [] + self.shared_memory_bytes = 0 self.values = {} + self.definitions = SSADefinitions(ssa_ops) + self.staging_analysis = CudaDotStagingAnalyzer(self.definitions).analyze() self.layout = cuda_kernel_layout(ssa_ops) self.emit_rank2_prologue() @@ -493,7 +737,7 @@ def generate( for op in ssa_ops: if isinstance(op, SSAForRange): self.emit_for_range(op) - else: + elif not self.is_staging_only(op): self.emit(op) body = [ diff --git a/src/mytriton/cuda_dot_staging.py b/src/mytriton/cuda_dot_staging.py new file mode 100644 index 0000000..ed68f4e --- /dev/null +++ b/src/mytriton/cuda_dot_staging.py @@ -0,0 +1,426 @@ +from dataclasses import dataclass + +from .ssa import SSAForRange, SSAItem, SSAOp, SSAOperand, SSAValue +from .trace import F32, I32, BlockType, Const, Param, PointerType, ScalarType + + +def cuda_scalar_nbytes(ty: ScalarType) -> int: + if ty in (F32, I32): + return 4 + if ty.name == "bool": + return 1 + + raise TypeError(f"cannot determine CUDA storage size for {ty}") + + +@dataclass(frozen=True) +class CudaSharedBuffer: + name: str + logical_shape: tuple[int, ...] + element_ty: ScalarType + + def __post_init__(self) -> None: + if len(self.logical_shape) != 2 or any( + type(dim) is not int or dim <= 0 for dim in self.logical_shape + ): + raise ValueError( + "shared buffer must be a positive rank-2 tile, " + f"got {self.logical_shape}" + ) + + @property + def rows(self) -> int: + return self.logical_shape[0] + + @property + def columns(self) -> int: + return self.logical_shape[1] + + @property + def size(self) -> int: + return self.rows * self.columns + + @property + def nbytes(self) -> int: + return self.size * cuda_scalar_nbytes(self.element_ty) + + def element(self, row: str, column: str) -> str: + return f"{self.name}[({row}) * {self.columns} + ({column})]" + + +@dataclass(frozen=True) +class CudaGlobalTile: + base: str + row_offset: str + column_offset: str + row_stride: str + row_bound: str + column_bound: str + other: str = "0.0f" + + +@dataclass(frozen=True) +class CudaGlobalTilePlan: + base: SSAOperand + row_offset: SSAOperand + column_offset: SSAOperand + row_stride: SSAOperand + row_bound: SSAOperand + column_bound: SSAOperand + other: SSAOperand + + +@dataclass(frozen=True) +class CudaDotStagingPlan: + lhs: CudaGlobalTilePlan + rhs: CudaGlobalTilePlan + + +@dataclass(frozen=True) +class CudaDotSharedBuffers: + lhs: CudaSharedBuffer + rhs: CudaSharedBuffer + + @property + def reduction_size(self) -> int: + return self.lhs.columns + + +class SSADefinitions: + def __init__(self, ssa_ops: list[SSAItem]) -> None: + self.ops: dict[int, SSAOp] = {} + self.ordered_ops: list[SSAOp] = [] + self._collect(ssa_ops) + + def _collect(self, ssa_ops: list[SSAItem]) -> None: + for item in ssa_ops: + if isinstance(item, SSAForRange): + self._collect(item.body) + continue + + self.ordered_ops.append(item) + + if item.result is None: + continue + + result_id = item.result.id + if result_id in self.ops: + raise ValueError(f"duplicate SSA definition for %{result_id}") + + self.ops[result_id] = item + + def get(self, value: SSAValue) -> SSAOp | None: + return self.ops.get(value.id) + + def require( + self, + operand: SSAOperand, + opcode: str, + ) -> SSAOp: + if not isinstance(operand, SSAValue): + raise TypeError(f"expected SSA value defined by {opcode}, got {operand}") + + op = self.get(operand) + if op is None: + raise TypeError(f"SSA value {operand} has no operation definition") + + if op.opcode != opcode: + raise TypeError( + f"expected {operand} to be defined by {opcode}, got {op.opcode}" + ) + + return op + + def dependency_ids( + self, + *operands: SSAOperand, + ) -> set[int]: + result: set[int] = set() + + def visit(operand: SSAOperand) -> None: + if not isinstance(operand, SSAValue): + return + + if operand.id in result: + return + + op = self.get(operand) + if op is None: + return + + result.add(operand.id) + + for dependency in op.operands: + visit(dependency) + + for operand in operands: + visit(operand) + + return result + + +class CudaDotOperandMatcher: + def __init__(self, definitions: SSADefinitions) -> None: + self.definitions = definitions + + @staticmethod + def _is_i32_scalar(operand: SSAOperand) -> bool: + if isinstance(operand, SSAValue): + return operand.ty == I32 + + if isinstance(operand, Param): + return operand.ty == I32 + + return isinstance(operand, Const) and type(operand.value) is int + + def _split_block_and_scalar( + self, + lhs: SSAOperand, + rhs: SSAOperand, + opcode: str, + ) -> tuple[SSAValue, SSAOperand]: + for block, scalar in ((lhs, rhs), (rhs, lhs)): + if ( + isinstance(block, SSAValue) + and isinstance(block.ty, BlockType) + and block.ty.element == I32 + and self._is_i32_scalar(scalar) + ): + return block, scalar + + raise TypeError( + f"dot staging expects {opcode} of block and scalar, got {lhs} and {rhs}" + ) + + def _match_axis_offset( + self, + coordinates: SSAValue, + axis: int, + ) -> SSAOperand: + add = self.definitions.require(coordinates, "add") + lhs, rhs = add.operands + + for offset, expanded in ((lhs, rhs), (rhs, lhs)): + if not self._is_i32_scalar(offset): + continue + + if not isinstance(expanded, SSAValue): + continue + + expand = self.definitions.get(expanded) + if ( + expand is None + or expand.opcode != "expand_dims" + or expand.attrs.get("axis") != axis + ): + continue + + arange = self.definitions.require( + expand.operands[0], + "arange", + ) + if arange.attrs.get("start") != 0: + continue + + return offset + + raise TypeError( + "dot staging expects coordinates in the form " + f"scalar_offset + expand_dims(arange(0, size), axis={axis})" + ) + + def _match_mask( + self, + mask: SSAOperand, + rows: SSAValue, + columns: SSAValue, + ) -> tuple[SSAOperand, SSAOperand]: + conjunction = self.definitions.require(mask, "and") + + row_bound: SSAOperand = None + column_bound: SSAOperand = None + + for comparison_operand in conjunction.operands: + comparison = self.definitions.require( + comparison_operand, + "cmp_lt", + ) + coordinate, bound = comparison.operands + + if not self._is_i32_scalar(bound): + raise TypeError( + f"dot staging bounds must be scalar i32 values, got {bound}" + ) + + if coordinate == rows: + row_bound = bound + elif coordinate == columns: + column_bound = bound + + if row_bound is None or column_bound is None: + raise TypeError( + "dot staging expects mask (rows < row_bound) & (columns < column_bound)" + ) + + return row_bound, column_bound + + def match(self, load_value: SSAOperand) -> CudaGlobalTilePlan: + load = self.definitions.require(load_value, "load") + pointer, mask, other = load.operands + + outer_addptr = self.definitions.require(pointer, "addptr") + row_pointer, columns = outer_addptr.operands + + inner_addptr = self.definitions.require(row_pointer, "addptr") + base, row_offset_expression = inner_addptr.operands + + row_stride_mul = self.definitions.require( + row_offset_expression, + "mul", + ) + row_stride_lhs, row_stride_rhs = row_stride_mul.operands + rows, row_stride = self._split_block_and_scalar( + row_stride_lhs, + row_stride_rhs, + opcode="row-stride multiplication", + ) + + if not isinstance(columns, SSAValue): + raise TypeError( + f"dot staging expects block-shaped column coordinates, got {columns}" + ) + + if not isinstance(base, Param) or not isinstance( + base.ty, + PointerType, + ): + raise TypeError( + f"dot staging expects a global pointer parameter, got {base}" + ) + + if ( + not isinstance(other, Const) + or type(other.value) is not float + or other.value != 0.0 + ): + raise TypeError("dot staging requires masked loads with other=0.0") + + row_offset = self._match_axis_offset(rows, axis=1) + column_offset = self._match_axis_offset(columns, axis=0) + row_bound, column_bound = self._match_mask( + mask, + rows, + columns, + ) + + return CudaGlobalTilePlan( + base=base, + row_offset=row_offset, + column_offset=column_offset, + row_stride=row_stride, + row_bound=row_bound, + column_bound=column_bound, + other=other, + ) + + +@dataclass(frozen=True) +class CudaDotStagingAnalysis: + dot_plans: dict[int, CudaDotStagingPlan] + staging_only_ids: frozenset[int] + + @property + def stageable_dot_ids(self) -> frozenset[int]: + return frozenset(self.dot_plans) + + def plan_for(self, dot_result_id: int) -> CudaDotStagingPlan: + try: + return self.dot_plans[dot_result_id] + except KeyError as error: + raise KeyError( + f"dot result %{dot_result_id} has no staging plan" + ) from error + + +class CudaDotStagingAnalyzer: + def __init__(self, definitions: SSADefinitions) -> None: + self.definitions = definitions + + def analyze(self) -> CudaDotStagingAnalysis: + matcher = CudaDotOperandMatcher(self.definitions) + + dot_plans: dict[int, CudaDotStagingPlan] = {} + staging_dependency_ids: set[int] = set() + required_scalar_ids: set[int] = set() + + for op in self.definitions.ordered_ops: + if op.opcode != "dot" or op.result is None: + continue + + lhs, rhs = op.operands + if not isinstance(lhs, SSAValue) or not isinstance( + rhs, + SSAValue, + ): + continue + + lhs_definition = self.definitions.get(lhs) + rhs_definition = self.definitions.get(rhs) + + # Version 13 examples using zeros remain non-stageable. + if ( + lhs_definition is None + or lhs_definition.opcode != "load" + or rhs_definition is None + or rhs_definition.opcode != "load" + ): + continue + + plan = CudaDotStagingPlan( + lhs=matcher.match(lhs), + rhs=matcher.match(rhs), + ) + dot_plans[op.result.id] = plan + staging_dependency_ids.update(self.definitions.dependency_ids(lhs, rhs)) + + for operand_plan in (plan.lhs, plan.rhs): + required_scalar_ids.update( + self.definitions.dependency_ids( + operand_plan.base, + operand_plan.row_offset, + operand_plan.column_offset, + operand_plan.row_stride, + operand_plan.row_bound, + operand_plan.column_bound, + operand_plan.other, + ) + ) + + external_dependency_ids: set[int] = set() + + for op in self.definitions.ordered_ops: + if op.opcode == "dot": + continue + + result_id = op.result.id if op.result is not None else None + + if result_id in staging_dependency_ids: + continue + + for operand in op.operands: + if ( + isinstance(operand, SSAValue) + and operand.id in staging_dependency_ids + ): + external_dependency_ids.update( + self.definitions.dependency_ids(operand) + ) + + staging_only_ids = ( + staging_dependency_ids - required_scalar_ids - external_dependency_ids + ) + + return CudaDotStagingAnalysis( + dot_plans=dot_plans, + staging_only_ids=frozenset(staging_only_ids), + ) diff --git a/tests/test_shared_memory.py b/tests/test_shared_memory.py new file mode 100644 index 0000000..2f9fff6 --- /dev/null +++ b/tests/test_shared_memory.py @@ -0,0 +1,889 @@ +from textwrap import dedent + +import numpy as np +import pytest + +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.cuda_dot_staging import ( + CudaDotOperandMatcher, + CudaDotSharedBuffers, + CudaDotStagingAnalysis, + CudaDotStagingAnalyzer, + CudaDotStagingPlan, + CudaGlobalTile, + CudaGlobalTilePlan, + CudaSharedBuffer, + SSADefinitions, +) +from mytriton.ssa import SSAForRange, SSAItem, SSALowering, SSAOp, SSAValue +from mytriton.trace import ( + F32, + I32, + PTR_F32, + BlockType, + Const, + Param, + Ptr, + Value, + make_runtime_params, +) + + +@triton.jit +def shared_memory_staging_kernel( + a, + b, + out, + M, + N, + K, + k_base, + BM: tl.constexpr, + BK: tl.constexpr, + BN: tl.constexpr, +): + offsets_m = tl.program_id(0) * BM + tl.arange(0, BM)[:, None] + offsets_n = tl.program_id(1) * BN + tl.arange(0, BN)[None, :] + offsets_k = tl.arange(0, BK) + + a_rows = offsets_m + a_columns = k_base + offsets_k[None, :] + a_values = tl.load( + a + a_rows * K + a_columns, + mask=(a_rows < M) & (a_columns < K), + other=0.0, + ) + + b_rows = k_base + offsets_k[:, None] + b_columns = offsets_n + b_values = tl.load( + b + b_rows * N + b_columns, + mask=(b_rows < K) & (b_columns < N), + other=0.0, + ) + + result = tl.dot(a_values, b_values) + + output_pointers = out + offsets_m * N + offsets_n + output_mask = (offsets_m < M) & (offsets_n < N) + tl.store(output_pointers, result, mask=output_mask) + + +@triton.jit +def shared_memory_runtime_loop_staging_kernel( + a, + b, + out, + M, + N, + K, + BM: tl.constexpr, + BK: tl.constexpr, + BN: tl.constexpr, +): + offsets_m = tl.program_id(0) * BM + tl.arange(0, BM)[:, None] + offsets_n = tl.program_id(1) * BN + tl.arange(0, BN)[None, :] + offsets_k = tl.arange(0, BK) + acc = tl.zeros((BM, BN), tl.float32) + + for k_base in range(0, K, BK): + a_rows = offsets_m + a_columns = k_base + offsets_k[None, :] + a_values = tl.load( + a + a_rows * K + a_columns, + mask=(a_rows < M) & (a_columns < K), + other=0.0, + ) + + b_rows = k_base + offsets_k[:, None] + b_columns = offsets_n + b_values = tl.load( + b + b_rows * N + b_columns, + mask=(b_rows < K) & (b_columns < N), + other=0.0, + ) + + acc = acc + tl.dot(a_values, b_values) + + output_pointers = out + offsets_m * N + offsets_n + output_mask = (offsets_m < M) & (offsets_n < N) + tl.store(output_pointers, acc, mask=output_mask) + + +def test_cuda_shared_buffer_represents_flat_rank2_tile() -> None: + buffer = CudaSharedBuffer( + name="dot_lhs_7", + logical_shape=(4, 16), + element_ty=F32, + ) + + assert buffer.rows == 4 + assert buffer.columns == 16 + assert buffer.size == 64 + assert buffer.nbytes == 256 + assert buffer.element("row", "column") == "dot_lhs_7[(row) * 16 + (column)]" + + +@pytest.mark.parametrize( + "shape", + [ + (), + (8,), + (4, 0), + (4, -1), + (4, 8, 2), + ], +) +def test_cuda_shared_buffer_rejects_invalid_shape( + shape: tuple[int, ...], +) -> None: + with pytest.raises( + ValueError, + match="shared buffer must be a positive rank-2 tile", + ): + CudaSharedBuffer( + name="tile", + logical_shape=shape, + element_ty=F32, + ) + + +def test_cuda_codegen_declares_shared_buffer() -> None: + codegen = SSACUDACodegen() + + buffer = codegen.declare_shared_buffer( + name="dot_lhs_7", + logical_shape=(4, 16), + element_ty=F32, + ) + + assert buffer == CudaSharedBuffer( + name="dot_lhs_7", + logical_shape=(4, 16), + element_ty=F32, + ) + assert codegen.shared_lines == [ + " __shared__ float dot_lhs_7[64];", + ] + + +def test_cuda_codegen_emits_cooperative_masked_load() -> None: + codegen = SSACUDACodegen() + codegen.layout = CudaKernelLayout( + output_tile_shape=(4, 8), + thread_shape=(4, 8), + ) + + target = codegen.declare_shared_buffer( + name="dot_lhs_7", + logical_shape=(4, 16), + element_ty=F32, + ) + source = CudaGlobalTile( + base="a", + row_offset="blockIdx.x * 4", + column_offset="k_base", + row_stride="K", + row_bound="M", + column_bound="K", + ) + + codegen.emit_cooperative_load(target, source) + + assert codegen.lines == [ + ( + " for (int dot_lhs_7_index = threadIdx.x; " + "dot_lhs_7_index < 64; dot_lhs_7_index += 32) {" + ), + " int dot_lhs_7_row = dot_lhs_7_index / 16;", + " int dot_lhs_7_column = dot_lhs_7_index % 16;", + (" int dot_lhs_7_global_row = (blockIdx.x * 4) + dot_lhs_7_row;"), + (" int dot_lhs_7_global_column = (k_base) + dot_lhs_7_column;"), + ( + " int dot_lhs_7_source_index = " + "dot_lhs_7_global_row * (K) + dot_lhs_7_global_column;" + ), + ( + " bool dot_lhs_7_in_bounds = " + "dot_lhs_7_global_row < (M) && " + "dot_lhs_7_global_column < (K);" + ), + ( + " dot_lhs_7[(dot_lhs_7_row) * 16 + " + "(dot_lhs_7_column)] = dot_lhs_7_in_bounds ? " + "a[dot_lhs_7_source_index] : 0.0f;" + ), + " }", + ] + + +def test_cuda_codegen_rejects_non_row_major_cooperative_load() -> None: + codegen = SSACUDACodegen() + codegen.layout = CudaKernelLayout( + output_tile_shape=(4, 8), + thread_shape=(4, 8), + ) + target = CudaSharedBuffer( + name="tile", + logical_shape=(4, 8), + element_ty=F32, + ) + source = CudaGlobalTile( + base="a", + row_offset="0", + column_offset="0", + row_stride="8", + row_bound="4", + column_bound="8", + ) + + with pytest.raises( + ValueError, + match=r"cooperative dot loads require row-major order \(1, 0\)", + ): + codegen.emit_cooperative_load( + target, + source, + order=(0, 1), + ) + + assert codegen.lines == [] + + +def test_cuda_codegen_emits_block_barrier() -> None: + codegen = SSACUDACodegen() + + codegen.emit_block_barrier() + + assert codegen.lines == [ + " __syncthreads();", + ] + + +def test_cuda_codegen_stages_both_dot_operands() -> None: + codegen = SSACUDACodegen() + codegen.layout = CudaKernelLayout( + output_tile_shape=(4, 8), + thread_shape=(4, 8), + ) + + buffers = codegen.emit_dot_operand_staging( + dot_result_id=7, + lhs_shape=(4, 16), + rhs_shape=(16, 8), + element_ty=F32, + lhs_source=CudaGlobalTile( + base="a", + row_offset="blockIdx.x * 4", + column_offset="k_base", + row_stride="K", + row_bound="M", + column_bound="K", + ), + rhs_source=CudaGlobalTile( + base="b", + row_offset="k_base", + column_offset="blockIdx.y * 8", + row_stride="N", + row_bound="K", + column_bound="N", + ), + ) + + assert buffers == CudaDotSharedBuffers( + lhs=CudaSharedBuffer( + name="dot_lhs_7", + logical_shape=(4, 16), + element_ty=F32, + ), + rhs=CudaSharedBuffer( + name="dot_rhs_7", + logical_shape=(16, 8), + element_ty=F32, + ), + ) + assert buffers.reduction_size == 16 + + assert codegen.shared_lines == [ + " __shared__ float dot_lhs_7[64];", + " __shared__ float dot_rhs_7[128];", + ] + + assert len(codegen.lines) == 19 + + assert codegen.lines[0] == ( + " for (int dot_lhs_7_index = threadIdx.x; " + "dot_lhs_7_index < 64; dot_lhs_7_index += 32) {" + ) + assert codegen.lines[6] == ( + " bool dot_lhs_7_in_bounds = " + "dot_lhs_7_global_row < (M) && " + "dot_lhs_7_global_column < (K);" + ) + assert codegen.lines[7] == ( + " dot_lhs_7[(dot_lhs_7_row) * 16 + " + "(dot_lhs_7_column)] = dot_lhs_7_in_bounds ? " + "a[dot_lhs_7_source_index] : 0.0f;" + ) + + assert codegen.lines[9] == ( + " for (int dot_rhs_7_index = threadIdx.x; " + "dot_rhs_7_index < 128; dot_rhs_7_index += 32) {" + ) + assert codegen.lines[15] == ( + " bool dot_rhs_7_in_bounds = " + "dot_rhs_7_global_row < (K) && " + "dot_rhs_7_global_column < (N);" + ) + assert codegen.lines[16] == ( + " dot_rhs_7[(dot_rhs_7_row) * 8 + " + "(dot_rhs_7_column)] = dot_rhs_7_in_bounds ? " + "b[dot_rhs_7_source_index] : 0.0f;" + ) + + assert codegen.lines[-1] == " __syncthreads();" + + +def test_dot_operand_staging_rejects_incompatible_shapes() -> None: + codegen = SSACUDACodegen() + codegen.layout = CudaKernelLayout( + output_tile_shape=(4, 8), + thread_shape=(4, 8), + ) + + source = CudaGlobalTile( + base="x", + row_offset="0", + column_offset="0", + row_stride="16", + row_bound="16", + column_bound="16", + ) + + with pytest.raises( + ValueError, + match="dot staging expects compatible rank-2 operands", + ): + codegen.emit_dot_operand_staging( + dot_result_id=0, + lhs_shape=(4, 16), + rhs_shape=(8, 8), + element_ty=F32, + lhs_source=source, + rhs_source=source, + ) + + assert codegen.shared_lines == [] + assert codegen.lines == [] + + +def test_dot_operand_staging_rejects_shared_memory_over_budget() -> None: + codegen = SSACUDACodegen() + codegen.layout = CudaKernelLayout( + output_tile_shape=(1, 1), + thread_shape=(1, 1), + ) + source = CudaGlobalTile( + base="x", + row_offset="0", + column_offset="0", + row_stride="128", + row_bound="128", + column_bound="128", + ) + + with pytest.raises( + ValueError, + match="exceeding the conservative 49152-byte limit", + ): + codegen.emit_dot_operand_staging( + dot_result_id=0, + lhs_shape=(128, 128), + rhs_shape=(128, 1), + element_ty=F32, + lhs_source=source, + rhs_source=source, + ) + + assert codegen.shared_memory_bytes == 0 + assert codegen.shared_lines == [] + assert codegen.lines == [] + + +def make_tiled_dot_ssa() -> tuple[list[SSAItem], SSAValue]: + BM, BK, BN = 4, 16, 8 + + a = Ptr(Param("a", PTR_F32)) + b = Ptr(Param("b", PTR_F32)) + + M = Value(Param("M", I32)) + N = Value(Param("N", I32)) + K = Value(Param("K", I32)) + k_base = Value(Param("k_base", I32)) + + offsets_m = tl.program_id(0) * BM + tl.arange(0, BM)[:, None] + offsets_n = tl.program_id(1) * BN + tl.arange(0, BN)[None, :] + offsets_k = tl.arange(0, BK) + + a_rows = offsets_m + a_columns = k_base + offsets_k[None, :] + a_pointers = a + a_rows * K + a_columns + a_mask = (a_rows < M) & (a_columns < K) + lhs = tl.load(a_pointers, mask=a_mask, other=0.0) + + b_rows = k_base + offsets_k[:, None] + b_columns = offsets_n + b_pointers = b + b_rows * N + b_columns + b_mask = (b_rows < K) & (b_columns < N) + rhs = tl.load(b_pointers, mask=b_mask, other=0.0) + + dot = tl.dot(lhs, rhs) + + lowering = SSALowering() + result = lowering.lower_expr(dot.expr) + + assert isinstance(result, SSAValue) + + return lowering.ops, result + + +def test_ssa_definitions_find_dot_operand_loads() -> None: + ssa_ops, dot_result = make_tiled_dot_ssa() + definitions = SSADefinitions(ssa_ops) + + dot = definitions.require(dot_result, "dot") + lhs, rhs = dot.operands + + assert isinstance(lhs, SSAValue) + assert isinstance(rhs, SSAValue) + assert lhs.id == 14 + assert rhs.id == 28 + + lhs_load = definitions.require(lhs, "load") + rhs_load = definitions.require(rhs, "load") + + assert lhs_load.result is lhs + assert rhs_load.result is rhs + assert lhs_load.opcode == "load" + assert rhs_load.opcode == "load" + + assert definitions.get(SSAValue(id=1000, ty=I32)) is None + + +def test_dot_operand_matcher_recovers_lhs_global_tile() -> None: + ssa_ops, dot_result = make_tiled_dot_ssa() + definitions = SSADefinitions(ssa_ops) + dot = definitions.require(dot_result, "dot") + + lhs, _ = dot.operands + plan = CudaDotOperandMatcher(definitions).match(lhs) + + assert plan == CudaGlobalTilePlan( + base=Param("a", PTR_F32), + row_offset=SSAValue(id=1, ty=I32), + column_offset=Param("k_base", I32), + row_stride=Param("K", I32), + row_bound=Param("M", I32), + column_bound=Param("K", I32), + other=Const(0.0), + ) + + +def test_dot_operand_matcher_recovers_rhs_global_tile() -> None: + ssa_ops, dot_result = make_tiled_dot_ssa() + definitions = SSADefinitions(ssa_ops) + dot = definitions.require(dot_result, "dot") + + _, rhs = dot.operands + plan = CudaDotOperandMatcher(definitions).match(rhs) + + assert plan == CudaGlobalTilePlan( + base=Param("b", PTR_F32), + row_offset=Param("k_base", I32), + column_offset=SSAValue(id=20, ty=I32), + row_stride=Param("N", I32), + row_bound=Param("K", I32), + column_bound=Param("N", I32), + other=Const(0.0), + ) + + +def test_cuda_codegen_stages_dot_operands_from_ssa() -> None: + ssa_ops, dot_result = make_tiled_dot_ssa() + definitions = SSADefinitions(ssa_ops) + dot = definitions.require(dot_result, "dot") + analysis = CudaDotStagingAnalyzer(definitions).analyze() + + codegen = SSACUDACodegen() + codegen.layout = CudaKernelLayout( + output_tile_shape=(4, 8), + thread_shape=(4, 8), + ) + codegen.definitions = definitions + codegen.staging_analysis = analysis + + codegen.values[1] = "v1" + codegen.values[20] = "v20" + + buffers = codegen.emit_dot_operand_staging_from_ssa( + dot, + analysis.plan_for(dot_result.id), + ) + + assert buffers.reduction_size == 16 + assert codegen.shared_lines == [ + " __shared__ float dot_lhs_29[64];", + " __shared__ float dot_rhs_29[128];", + ] + + assert codegen.lines[3] == ( + " int dot_lhs_29_global_row = (v1) + dot_lhs_29_row;" + ) + assert codegen.lines[4] == ( + " int dot_lhs_29_global_column = (k_base) + dot_lhs_29_column;" + ) + assert codegen.lines[5] == ( + " int dot_lhs_29_source_index = " + "dot_lhs_29_global_row * (K) + " + "dot_lhs_29_global_column;" + ) + assert codegen.lines[6] == ( + " bool dot_lhs_29_in_bounds = " + "dot_lhs_29_global_row < (M) && " + "dot_lhs_29_global_column < (K);" + ) + + assert codegen.lines[12] == ( + " int dot_rhs_29_global_row = (k_base) + dot_rhs_29_row;" + ) + assert codegen.lines[13] == ( + " int dot_rhs_29_global_column = (v20) + dot_rhs_29_column;" + ) + assert codegen.lines[14] == ( + " int dot_rhs_29_source_index = " + "dot_rhs_29_global_row * (N) + " + "dot_rhs_29_global_column;" + ) + assert codegen.lines[15] == ( + " bool dot_rhs_29_in_bounds = " + "dot_rhs_29_global_row < (K) && " + "dot_rhs_29_global_column < (N);" + ) + + assert codegen.lines[-1] == " __syncthreads();" + + +def test_dot_staging_analysis_finds_staging_only_operations() -> None: + ssa_ops, dot_result = make_tiled_dot_ssa() + definitions = SSADefinitions(ssa_ops) + + analysis = CudaDotStagingAnalyzer(definitions).analyze() + dot = definitions.require(dot_result, "dot") + lhs, rhs = dot.operands + + assert analysis.dot_plans == { + dot_result.id: CudaDotStagingPlan( + lhs=CudaDotOperandMatcher(definitions).match(lhs), + rhs=CudaDotOperandMatcher(definitions).match(rhs), + ) + } + assert analysis.stageable_dot_ids == frozenset({dot_result.id}) + assert analysis.staging_only_ids == frozenset( + [ + *range(2, 19), + *range(21, 29), + ] + ) + + +def test_dot_staging_analysis_ignores_non_load_operands() -> None: + lhs = SSAValue( + id=0, + ty=BlockType((4, 16), F32), + ) + rhs = SSAValue( + id=1, + ty=BlockType((16, 8), F32), + ) + result = SSAValue( + id=2, + ty=BlockType((4, 8), F32), + ) + + ssa_ops: list[SSAItem] = [ + SSAOp(opcode="zeros", result=lhs), + SSAOp(opcode="zeros", result=rhs), + SSAOp( + opcode="dot", + operands=(lhs, rhs), + result=result, + ), + ] + + analysis = CudaDotStagingAnalyzer(SSADefinitions(ssa_ops)).analyze() + + assert analysis == CudaDotStagingAnalysis( + dot_plans={}, + staging_only_ids=frozenset(), + ) + + +def test_cuda_generate_reaches_dot_after_skipping_staging_operations() -> None: + ssa_ops, dot_result = make_tiled_dot_ssa() + + ssa_ops.append( + SSAOp( + opcode="store", + operands=( + Param("out", PTR_F32), + dot_result, + None, + ), + ) + ) + + codegen = SSACUDACodegen() + + with pytest.raises( + TypeError, + match=( + r"CUDA shared-memory staging for tl\.dot is implemented, " + r"but CUDA computation for tl\.dot is not implemented" + ), + ): + codegen.generate( + kernel_name="staged_dot_kernel", + ssa_ops=ssa_ops, + params=[], + ) + + assert codegen.staging_analysis.stageable_dot_ids == frozenset({dot_result.id}) + + expected_cuda_fragment = "\n".join( + f" {line}" if line else "" + for line in dedent( + """ + __shared__ float dot_lhs_29[64]; + __shared__ float dot_rhs_29[128]; + + int tile_i = threadIdx.x / 8; + int tile_j = threadIdx.x % 8; + int v0 = blockIdx.x; + int v1 = (v0 * 4); + int v19 = blockIdx.y; + int v20 = (v19 * 8); + for (int dot_lhs_29_index = threadIdx.x; dot_lhs_29_index < 64; dot_lhs_29_index += 32) { + int dot_lhs_29_row = dot_lhs_29_index / 16; + int dot_lhs_29_column = dot_lhs_29_index % 16; + int dot_lhs_29_global_row = (v1) + dot_lhs_29_row; + int dot_lhs_29_global_column = (k_base) + dot_lhs_29_column; + int dot_lhs_29_source_index = dot_lhs_29_global_row * (K) + dot_lhs_29_global_column; + bool dot_lhs_29_in_bounds = dot_lhs_29_global_row < (M) && dot_lhs_29_global_column < (K); + dot_lhs_29[(dot_lhs_29_row) * 16 + (dot_lhs_29_column)] = dot_lhs_29_in_bounds ? a[dot_lhs_29_source_index] : 0.0f; + } + for (int dot_rhs_29_index = threadIdx.x; dot_rhs_29_index < 128; dot_rhs_29_index += 32) { + int dot_rhs_29_row = dot_rhs_29_index / 8; + int dot_rhs_29_column = dot_rhs_29_index % 8; + int dot_rhs_29_global_row = (k_base) + dot_rhs_29_row; + int dot_rhs_29_global_column = (v20) + dot_rhs_29_column; + int dot_rhs_29_source_index = dot_rhs_29_global_row * (N) + dot_rhs_29_global_column; + bool dot_rhs_29_in_bounds = dot_rhs_29_global_row < (K) && dot_rhs_29_global_column < (N); + dot_rhs_29[(dot_rhs_29_row) * 8 + (dot_rhs_29_column)] = dot_rhs_29_in_bounds ? b[dot_rhs_29_source_index] : 0.0f; + } + __syncthreads(); + """ + ) + .strip() + .splitlines() + ) + actual_cuda_fragment = "\n".join( + [ + *codegen.shared_lines, + "", + *codegen.lines, + ] + ) + + assert actual_cuda_fragment == expected_cuda_fragment + + +def test_ast_frontend_reaches_shared_memory_dot_staging( + monkeypatch, +) -> None: + monkeypatch.setenv("MYTRITON_BACKEND", "cuda") + + M, N, K = 4, 8, 16 + BM, BK, BN = 4, 16, 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) + + shared_memory_staging_kernel.clear_cache() + + with pytest.raises( + TypeError, + match=( + r"CUDA shared-memory staging for tl\.dot is implemented, " + r"but CUDA computation for tl\.dot is not implemented" + ), + ): + shared_memory_staging_kernel[(1, 1)]( + a, + b, + out, + M, + N, + K, + 0, + BM=BM, + BK=BK, + BN=BN, + ) + + +def test_runtime_k_loop_reaches_shared_memory_dot_staging() -> None: + M, N, K = 4, 8, 32 + BM, BK, BN = 4, 16, 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) + + bound = shared_memory_runtime_loop_staging_kernel.signature.bind( + a, + b, + out, + M, + N, + K, + BM=BM, + BK=BK, + BN=BN, + ) + runtime_params = make_runtime_params( + shared_memory_runtime_loop_staging_kernel.signature, + bound.arguments, + ) + traced_ops, _ = trace_ast( + shared_memory_runtime_loop_staging_kernel.fn, + shared_memory_runtime_loop_staging_kernel.signature, + bound.arguments, + runtime_params=runtime_params, + ) + ssa_ops = SSALowering().lower(traced_ops) + + 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") + assert dot.result is not None + + definitions = SSADefinitions(ssa_ops) + analysis = CudaDotStagingAnalyzer(definitions).analyze() + plan = analysis.plan_for(dot.result.id) + + assert plan.lhs.column_offset == loop.index + assert plan.rhs.row_offset == loop.index + + codegen = SSACUDACodegen() + with pytest.raises( + TypeError, + match=( + r"CUDA shared-memory staging for tl\.dot is implemented, " + r"but CUDA computation for tl\.dot is not implemented" + ), + ): + codegen.generate( + kernel_name="runtime_loop_staged_dot_kernel", + ssa_ops=ssa_ops, + params=runtime_params, + ) + + loop_index = f"v{loop.index.id}" + assert any( + line.startswith(f" for (int {loop_index} = ") for line in codegen.lines + ) + assert any( + f"global_column = ({loop_index}) + dot_lhs_" in line for line in codegen.lines + ) + assert any( + f"global_row = ({loop_index}) + dot_rhs_" in line for line in codegen.lines + ) + assert codegen.shared_memory_bytes == (BM * BK + BK * BN) * 4 + + +def test_dot_operand_matcher_rejects_block_bound() -> None: + ssa_ops, dot_result = make_tiled_dot_ssa() + definitions = SSADefinitions(ssa_ops) + dot = definitions.require(dot_result, "dot") + + lhs, _ = dot.operands + lhs_load = definitions.require(lhs, "load") + mask = lhs_load.operands[1] + conjunction = definitions.require(mask, "and") + + row_comparison_operand = conjunction.operands[0] + row_comparison = definitions.require( + row_comparison_operand, + "cmp_lt", + ) + + coordinate, _ = row_comparison.operands + block_bound = SSAValue( + id=1000, + ty=BlockType((4, 1), I32), + ) + row_comparison.operands = (coordinate, block_bound) + + with pytest.raises( + TypeError, + match="dot staging bounds must be scalar i32", + ): + CudaDotOperandMatcher(definitions).match(lhs) + + +def test_dot_staging_analysis_preserves_external_dependencies() -> None: + ssa_ops, dot_result = make_tiled_dot_ssa() + initial_definitions = SSADefinitions(ssa_ops) + + rows = initial_definitions.ops[4].result + columns = initial_definitions.ops[23].result + + assert rows is not None + assert columns is not None + + output_offset = SSAValue( + id=30, + ty=BlockType((4, 8), I32), + ) + ssa_ops.append( + SSAOp( + opcode="add", + operands=(rows, columns), + result=output_offset, + ) + ) + ssa_ops.append( + SSAOp( + opcode="store", + operands=( + Param("out", PTR_F32), + output_offset, + None, + ), + ) + ) + + analysis = CudaDotStagingAnalyzer(SSADefinitions(ssa_ops)).analyze() + + # Shared with output address: must use ordinary lowering. + assert not {0, 1, 2, 3, 4} & analysis.staging_only_ids + assert not {19, 20, 21, 22, 23} & analysis.staging_only_ids + + # Private pointer/mask/load operations remain staging-only. + assert {5, 6, 10, 11, 12, 13, 14} <= analysis.staging_only_ids + assert {17, 18, 24, 25, 26, 27, 28} <= analysis.staging_only_ids + + assert dot_result.id in analysis.stageable_dot_ids