From 7909dd8a43725e1a74c71a8e9cbae470a0758949 Mon Sep 17 00:00:00 2001 From: Pavel Belevich Date: Sat, 22 Aug 2026 22:57:35 -0400 Subject: [PATCH] Version 15: Lower tl.dot to CUDA cores --- README.md | 72 +++-- src/mytriton/cuda_codegen.py | 83 +++++- tests/test_shared_memory.py | 530 +++++++++++++++++++++++++++++------ 3 files changed, 581 insertions(+), 104 deletions(-) diff --git a/README.md b/README.md index 47bd625..c5a9291 100644 --- a/README.md +++ b/README.md @@ -59,6 +59,11 @@ MLIR's GPU/NVVM stack to a cubin. cooperative staging of `tl.dot` operands, zero-filled boundary handling, block synchronization, and an explicit diagnostic for the deferred CUDA-core dot computation. +- [ver15](https://github.com/pbelevich/mytriton/tree/ver15): working + CUDA-core lowering for canonical `tl.dot` matrix tiles, one register + accumulator per output thread, an FMA loop over shared-memory operands, + synchronization before tile reuse, runtime traversal of multiple K-tiles, + and CUDA correctness tests for masked edge tiles. ## AST frontend @@ -243,12 +248,11 @@ type. `dot` is a pure SSA operation, so duplicate operations are eligible for common subexpression elimination and unused operations can be removed by dead-code elimination. -Version 13 defines the language and IR semantics only. CUDA lowering for -`tl.dot` is intentionally not implemented yet. The next versions will introduce -cooperative shared-memory tiles and then lower `dot` to an ordinary CUDA-core -multiply-accumulate loop. +Version 13 defines the language and IR semantics only. Version 14 adds +cooperative shared-memory staging for canonical matrix loads, and Version 15 +lowers the staged operands to an ordinary CUDA-core multiply-accumulate loop. -## Shared-memory dot staging +## Shared-memory CUDA-core dot The CUDA backend recognizes canonical matrix tiles loaded for `tl.dot`: @@ -288,11 +292,39 @@ 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()`. +After both cooperative loads, the backend emits `__syncthreads()` so no thread +starts reading a tile before all writes have completed. -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. +Each CUDA thread owns one output coordinate `(tile_i, tile_j)` and one `f32` +register accumulator. It performs the ordinary CUDA-core reduction: + +```text +accumulator = 0.0 + +for k in range(BK): + accumulator += shared_a[tile_i, k] * shared_b[k, tile_j] +``` + +A second `__syncthreads()` ensures every thread has finished reading the current +shared buffers before a runtime K-loop iteration overwrites them with the next +tiles. Partial K-tiles are zero-filled by the existing load masks. + +A complete tiled matmul can therefore accumulate several `tl.dot` results: + +```python +acc = tl.zeros((BM, BN), tl.float32) + +for k_base in range(0, K, BK): + # Build and load A [BM, BK] and B [BK, BN] tiles. + acc = acc + tl.dot(a_values, b_values) + +tl.store(output_pointers, acc, mask=output_mask) +``` + +Version 14 intentionally stops after shared-memory staging. Version 15 adds the +CUDA-core FMA loop, safe shared-buffer reuse, and execution across multiple +K-tiles. It still uses one CUDA thread per output element; register tiles with +multiple results per thread are deferred to Version 16. ## Example @@ -525,23 +557,21 @@ these rewrite passes because they are not region-aware yet. reductions. Floating-point elementwise extrema propagate NaNs and choose the 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. + masked loads with zero-filled boundaries, a CUDA-core FMA loop, and the + barriers required before reading and reusing the shared tiles. Runtime + `range` loops can accumulate multiple K-tiles into one result. - 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 one block-local partial vector, as in the long-row sum test, but there is no multi-block reduction yet. -- Matrix multiplication support is intentionally naive so far. The current - rank-2 matmul kernel computes one output tile with one CUDA thread per output - element and repeatedly reads from global memory. It can traverse the - reduction dimension with either an unrolled `tl.static_range` or a runtime - 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 - can only stage canonical load operands cooperatively in CUDA shared memory; - the multiply-accumulate computation remains deferred to Version 15. +- Matrix multiplication supports a correct tiled CUDA-core implementation for + canonical `tl.dot` operands. A `[BM, BK]` and B `[BK, BN]` are loaded + cooperatively into shared memory, each thread computes one C element, and a + runtime CUDA loop can traverse the complete K dimension. The implementation + prioritizes correctness over performance: it has no per-thread register + tiles, vectorized loads, shared-memory padding or swizzling, double buffering, + asynchronous copies, tensor-core instructions, or autotuning yet. `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 ddd53d1..d98b37d 100644 --- a/src/mytriton/cuda_codegen.py +++ b/src/mytriton/cuda_codegen.py @@ -282,6 +282,77 @@ def emit_dot_operand_staging( rhs=rhs, ) + def emit_dot_from_shared_memory( + self, + result: SSAValue, + buffers: CudaDotSharedBuffers, + ) -> None: + if not isinstance(result.ty, BlockType) or result.ty.rank != 2: + raise TypeError(f"CUDA-core dot requires a rank-2 result, got {result.ty}") + + if buffers.lhs.columns != buffers.rhs.rows: + raise TypeError( + "CUDA-core dot requires matching reduction dimensions, " + f"got {buffers.lhs.logical_shape} and " + f"{buffers.rhs.logical_shape}" + ) + + expected_shape = ( + buffers.lhs.rows, + buffers.rhs.columns, + ) + if result.ty.shape != expected_shape: + raise TypeError( + f"CUDA-core dot expected result shape {expected_shape}, " + f"got {result.ty.shape}" + ) + + if result.ty.shape != self.layout.thread_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}" + ) + + if ( + result.ty.element != F32 + or buffers.lhs.element_ty != F32 + or buffers.rhs.element_ty != F32 + ): + raise TypeError("CUDA-core dot currently supports only f32") + + result_name = f"v{result.id}" + 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, + ) + + 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};"), + " }", + ] + ) + self.values[result.id] = result_name + + # All threads must finish reading the current tiles before another + # runtime K-loop iteration overwrites the shared buffers. + self.emit_block_barrier() + def resolve_global_tile( self, plan: CudaGlobalTilePlan, @@ -587,11 +658,13 @@ def emit(self, op: SSAOp) -> None: 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" + buffers = self.emit_dot_operand_staging_from_ssa( + op, + plan, + ) + self.emit_dot_from_shared_memory( + result, + buffers, ) elif op.opcode in self.BINARY_OPS: lhs = self.expression_operand(op.operands[0]) diff --git a/tests/test_shared_memory.py b/tests/test_shared_memory.py index 2f9fff6..1cf86f2 100644 --- a/tests/test_shared_memory.py +++ b/tests/test_shared_memory.py @@ -19,7 +19,14 @@ CudaSharedBuffer, SSADefinitions, ) -from mytriton.ssa import SSAForRange, SSAItem, SSALowering, SSAOp, SSAValue +from mytriton.ssa import ( + SSAForRange, + SSAItem, + SSALowering, + SSAOp, + SSAPrinter, + SSAValue, +) from mytriton.trace import ( F32, I32, @@ -34,7 +41,7 @@ @triton.jit -def shared_memory_staging_kernel( +def single_tile_dot_kernel( a, b, out, @@ -74,7 +81,7 @@ def shared_memory_staging_kernel( @triton.jit -def shared_memory_runtime_loop_staging_kernel( +def tiled_matmul_kernel( a, b, out, @@ -632,7 +639,7 @@ def test_dot_staging_analysis_ignores_non_load_operands() -> None: ) -def test_cuda_generate_reaches_dot_after_skipping_staging_operations() -> None: +def test_cuda_generate_lowers_dot_to_cuda_cores() -> None: ssa_ops, dot_result = make_tiled_dot_ssa() ssa_ops.append( @@ -646,27 +653,29 @@ def test_cuda_generate_reaches_dot_after_skipping_staging_operations() -> None: ) ) - codegen = SSACUDACodegen() + params = [ + Param("a", PTR_F32), + Param("b", PTR_F32), + Param("out", PTR_F32), + Param("M", I32), + Param("N", I32), + Param("K", I32), + Param("k_base", I32), + ] - 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=[], - ) + codegen = SSACUDACodegen() + cuda_src = codegen.generate( + kernel_name="cuda_core_dot_kernel", + ssa_ops=ssa_ops, + params=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( - """ + expected_cuda_src = dedent( + """\ + extern "C" __global__ + void cuda_core_dot_kernel(float* a, float* b, float* out, int M, int N, int K, int k_base) { __shared__ float dot_lhs_29[64]; __shared__ float dot_rhs_29[128]; @@ -695,25 +704,23 @@ def test_cuda_generate_reaches_dot_after_skipping_staging_operations() -> None: 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, - ] - ) + float v29 = 0.0f; + for (int dot_k_29 = 0; dot_k_29 < 16; ++dot_k_29) { + v29 += dot_lhs_29[(tile_i) * 16 + (dot_k_29)] * dot_rhs_29[(dot_k_29) * 8 + (tile_j)]; + } + __syncthreads(); + out[0] = v29; + } + """ + ).rstrip("\n") - assert actual_cuda_fragment == expected_cuda_fragment + assert cuda_src == expected_cuda_src -def test_ast_frontend_reaches_shared_memory_dot_staging( +def test_ast_frontend_lowers_shared_memory_dot_to_cuda( monkeypatch, ) -> None: + monkeypatch.setenv("MYTRITON_FRONTEND", "ast") monkeypatch.setenv("MYTRITON_BACKEND", "cuda") M, N, K = 4, 8, 16 @@ -723,37 +730,169 @@ def test_ast_frontend_reaches_shared_memory_dot_staging( b = np.zeros((K, N), dtype=np.float32) out = np.zeros((M, N), dtype=np.float32) - shared_memory_staging_kernel.clear_cache() + single_tile_dot_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, - ) + _, ssa_ops, cuda_src = single_tile_dot_kernel[(1, 1)]( + a, + b, + out, + M, + N, + K, + 0, + BM=BM, + BK=BK, + BN=BN, + ) + expected_ssa = dedent( + """\ + %0 = program_id {axis=0} : i32 + %1 = mul %0, 4 : i32 + %2 = arange {start=0, end=4} : vector<4 x i32> + %3 = expand_dims %2 {axis=1} : block<4x1 x i32> + %4 = add %1, %3 : block<4x1 x i32> + %5 = mul %4, K : block<4x1 x i32> + %6 = addptr a, %5 : block<4x1 x ptr> + %7 = arange {start=0, end=16} : vector<16 x i32> + %8 = expand_dims %7 {axis=0} : block<1x16 x i32> + %9 = add k_base, %8 : block<1x16 x i32> + %10 = addptr %6, %9 : block<4x16 x ptr> + %11 = cmp_lt %4, M : block<4x1 x bool> + %12 = cmp_lt %9, K : block<1x16 x bool> + %13 = and %11, %12 : block<4x16 x bool> + %14 = load %10, %13, 0.0 : block<4x16 x f32> + %15 = expand_dims %7 {axis=1} : block<16x1 x i32> + %16 = add k_base, %15 : block<16x1 x i32> + %17 = mul %16, N : block<16x1 x i32> + %18 = addptr b, %17 : block<16x1 x ptr> + %19 = program_id {axis=1} : i32 + %20 = mul %19, 8 : i32 + %21 = arange {start=0, end=8} : vector<8 x i32> + %22 = expand_dims %21 {axis=0} : block<1x8 x i32> + %23 = add %20, %22 : block<1x8 x i32> + %24 = addptr %18, %23 : block<16x8 x ptr> + %25 = cmp_lt %16, K : block<16x1 x bool> + %26 = cmp_lt %23, N : block<1x8 x bool> + %27 = and %25, %26 : block<16x8 x bool> + %28 = load %24, %27, 0.0 : block<16x8 x f32> + %29 = dot %14, %28 : block<4x8 x f32> + %30 = mul %4, N : block<4x1 x i32> + %31 = addptr out, %30 : block<4x1 x ptr> + %32 = addptr %31, %23 : block<4x8 x ptr> + %35 = and %11, %26 : block<4x8 x bool> + store %32, %29, %35 + """ + ).rstrip("\n") + + expected_cuda_src = dedent( + """\ + extern "C" __global__ + void single_tile_dot_kernel(float* a, float* b, float* out, int M, int N, int K, int k_base) { + __shared__ float dot_lhs_29[64]; + __shared__ float dot_rhs_29[128]; -def test_runtime_k_loop_reaches_shared_memory_dot_staging() -> None: + int tile_i = threadIdx.x / 8; + int tile_j = threadIdx.x % 8; + int v0 = blockIdx.x; + int v1 = (v0 * 4); + int v3 = tile_i; + int v4 = (v1 + v3); + bool v11 = (v4 < M); + int v19 = blockIdx.y; + int v20 = (v19 * 8); + int v22 = tile_j; + int v23 = (v20 + v22); + bool v26 = (v23 < N); + 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(); + float v29 = 0.0f; + for (int dot_k_29 = 0; dot_k_29 < 16; ++dot_k_29) { + v29 += dot_lhs_29[(tile_i) * 16 + (dot_k_29)] * dot_rhs_29[(dot_k_29) * 8 + (tile_j)]; + } + __syncthreads(); + int v30 = (v4 * N); + bool v35 = (v11 && v26); + if (v35) { + out[(v30 + v23)] = v29; + } + } + """ + ).rstrip("\n") + + assert SSAPrinter().print_ops(ssa_ops) == expected_ssa + assert cuda_src == expected_cuda_src + + +@pytest.mark.execution +def test_shared_memory_dot_executes_single_k_tile( + cp, + monkeypatch, +) -> None: + monkeypatch.setenv("MYTRITON_FRONTEND", "ast") + monkeypatch.setenv("MYTRITON_BACKEND", "cuda") + + M, N, K = 7, 13, 7 + BM, BK, BN = 4, 8, 8 + + a = cp.arange(M * K, dtype=cp.float32).reshape(M, K) / K + b = cp.arange(K * N, dtype=cp.float32).reshape(K, N) / N + 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 @ b + cp.testing.assert_allclose( + out, + expected, + rtol=1e-5, + atol=1e-5, + ) + + +def test_runtime_k_loop_lowers_dot_with_reusable_shared_memory() -> 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( + bound = tiled_matmul_kernel.signature.bind( a, b, out, @@ -765,12 +904,12 @@ def test_runtime_k_loop_reaches_shared_memory_dot_staging() -> None: BN=BN, ) runtime_params = make_runtime_params( - shared_memory_runtime_loop_staging_kernel.signature, + tiled_matmul_kernel.signature, bound.arguments, ) traced_ops, _ = trace_ast( - shared_memory_runtime_loop_staging_kernel.fn, - shared_memory_runtime_loop_staging_kernel.signature, + tiled_matmul_kernel.fn, + tiled_matmul_kernel.signature, bound.arguments, runtime_params=runtime_params, ) @@ -787,33 +926,117 @@ def test_runtime_k_loop_reaches_shared_memory_dot_staging() -> None: assert plan.lhs.column_offset == loop.index assert plan.rhs.row_offset == loop.index + accumulation = next( + op + for op in loop.body + if (isinstance(op, SSAOp) and op.opcode == "add" and dot.result in op.operands) + ) + assert accumulation.result is not None + assert len(loop.results) == 1 + 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, - ) + cuda_src = codegen.generate( + kernel_name="tiled_matmul_kernel", + ssa_ops=ssa_ops, + params=runtime_params, + ) + dot_id = dot.result.id loop_index = f"v{loop.index.id}" - assert any( - line.startswith(f" for (int {loop_index} = ") for line in codegen.lines + accumulator = f"v{loop.results[0].id}" + accumulation_result = f"v{accumulation.result.id}" + + outer_loop = ( + f" for (int {loop_index} = 0; {loop_index} < K; {loop_index} += {BK}) {{" + ) + lhs_load = f" for (int dot_lhs_{dot_id}_index = threadIdx.x; " + rhs_load = f" for (int dot_rhs_{dot_id}_index = threadIdx.x; " + fma_loop = ( + f" for (int dot_k_{dot_id} = 0; " + f"dot_k_{dot_id} < {BK}; ++dot_k_{dot_id}) {{" ) - assert any( - f"global_column = ({loop_index}) + dot_lhs_" in line for line in codegen.lines + barrier = " __syncthreads();" + + outer_loop_position = cuda_src.index(outer_loop) + lhs_load_position = cuda_src.index(lhs_load, outer_loop_position) + rhs_load_position = cuda_src.index(rhs_load, lhs_load_position) + load_barrier_position = cuda_src.index(barrier, rhs_load_position) + fma_position = cuda_src.index(fma_loop, load_barrier_position) + reuse_barrier_position = cuda_src.index(barrier, fma_position) + + expected_accumulation = ( + f" float {accumulation_result} = " + f"({accumulator} + v{dot_id});\n" + f" {accumulator} = {accumulation_result};" ) - assert any( - f"global_row = ({loop_index}) + dot_rhs_" in line for line in codegen.lines + accumulation_position = cuda_src.index( + expected_accumulation, + reuse_barrier_position, ) + + assert ( + outer_loop_position + < lhs_load_position + < rhs_load_position + < load_barrier_position + < fma_position + < reuse_barrier_position + < accumulation_position + ) + + assert f"global_column = ({loop_index}) + dot_lhs_{dot_id}_column;" in cuda_src + assert f"global_row = ({loop_index}) + dot_rhs_{dot_id}_row;" in cuda_src + + assert cuda_src.count("__syncthreads();") == 2 assert codegen.shared_memory_bytes == (BM * BK + BK * BN) * 4 +@pytest.mark.execution +def test_shared_memory_dot_executes_multiple_k_tiles( + cp, + monkeypatch, +) -> None: + monkeypatch.setenv("MYTRITON_FRONTEND", "ast") + monkeypatch.setenv("MYTRITON_BACKEND", "cuda") + + M, N, K = 7, 13, 19 + BM, BK, BN = 4, 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, + ) + + def test_dot_operand_matcher_rejects_block_bound() -> None: ssa_ops, dot_result = make_tiled_dot_ssa() definitions = SSADefinitions(ssa_ops) @@ -887,3 +1110,154 @@ def test_dot_staging_analysis_preserves_external_dependencies() -> None: assert {17, 18, 24, 25, 26, 27, 28} <= analysis.staging_only_ids assert dot_result.id in analysis.stageable_dot_ids + + +def test_cuda_codegen_computes_dot_from_shared_memory() -> None: + codegen = SSACUDACodegen() + codegen.layout = CudaKernelLayout( + output_tile_shape=(4, 8), + thread_shape=(4, 8), + ) + + result = SSAValue( + id=7, + ty=BlockType((4, 8), F32), + ) + 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, + ), + ) + + codegen.emit_dot_from_shared_memory(result, buffers) + + assert codegen.lines == [ + " float v7 = 0.0f;", + " for (int dot_k_7 = 0; dot_k_7 < 16; ++dot_k_7) {", + ( + " v7 += " + "dot_lhs_7[(tile_i) * 16 + (dot_k_7)] * " + "dot_rhs_7[(dot_k_7) * 8 + (tile_j)];" + ), + " }", + " __syncthreads();", + ] + assert codegen.values[result.id] == "v7" + + +def test_cuda_core_dot_rejects_mismatched_reduction_dimensions() -> None: + codegen = SSACUDACodegen() + codegen.layout = CudaKernelLayout( + output_tile_shape=(4, 8), + thread_shape=(4, 8), + ) + + result = SSAValue( + id=7, + ty=BlockType((4, 8), F32), + ) + buffers = CudaDotSharedBuffers( + lhs=CudaSharedBuffer( + name="dot_lhs_7", + logical_shape=(4, 16), + element_ty=F32, + ), + rhs=CudaSharedBuffer( + name="dot_rhs_7", + logical_shape=(8, 8), + element_ty=F32, + ), + ) + + with pytest.raises( + TypeError, + match="matching reduction dimensions", + ): + codegen.emit_dot_from_shared_memory( + result, + buffers, + ) + + assert codegen.lines == [] + assert codegen.values == {} + + +def test_cuda_core_dot_rejects_wrong_result_shape() -> None: + codegen = SSACUDACodegen() + codegen.layout = CudaKernelLayout( + output_tile_shape=(4, 7), + thread_shape=(4, 7), + ) + + result = SSAValue( + id=7, + ty=BlockType((4, 7), F32), + ) + 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, + ), + ) + + with pytest.raises( + TypeError, + match=r"expected result shape \(4, 8\), got \(4, 7\)", + ): + codegen.emit_dot_from_shared_memory( + result, + buffers, + ) + + assert codegen.lines == [] + assert codegen.values == {} + + +def test_cuda_core_dot_requires_one_thread_per_result_element() -> None: + codegen = SSACUDACodegen() + codegen.layout = CudaKernelLayout( + output_tile_shape=(4, 8), + thread_shape=(2, 8), + ) + + result = SSAValue( + id=7, + ty=BlockType((4, 8), F32), + ) + 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, + ), + ) + + with pytest.raises( + TypeError, + match="requires one CUDA thread per result element", + ): + codegen.emit_dot_from_shared_memory( + result, + buffers, + ) + + assert codegen.lines == [] + assert codegen.values == {}