From 772ad6bf264258eb0a6079841c25ffd1cda05942 Mon Sep 17 00:00:00 2001 From: Pavel Belevich Date: Sat, 22 Aug 2026 15:26:33 -0400 Subject: [PATCH] Version 13: Add tl.dot semantics --- README.md | 53 ++++- src/mytriton/cuda_codegen.py | 2 + src/mytriton/language.py | 2 + src/mytriton/optim.py | 1 + src/mytriton/ssa.py | 11 + src/mytriton/ssa_verification.py | 34 +++ src/mytriton/trace.py | 11 + src/mytriton/type_inference.py | 31 +++ tests/test_dot.py | 396 +++++++++++++++++++++++++++++++ 9 files changed, 534 insertions(+), 7 deletions(-) create mode 100644 tests/test_dot.py diff --git a/README.md b/README.md index 96a5ba0..0ca74f7 100644 --- a/README.md +++ b/README.md @@ -49,6 +49,11 @@ MLIR's GPU/NVVM stack to a cubin. organization, store-rooted output-layout inference, reduction-aware thread layouts, projected layouts for per-thread values, and cooperative layouts for distributing arbitrary rank-2 tiles across a CUDA thread block. +- [ver13](https://github.com/pbelevich/mytriton/tree/ver13): public `tl.dot` + semantics for rank-2 `f32` blocks, expression-tree and typed SSA operations, + `[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. ## AST frontend @@ -206,8 +211,37 @@ matrix multiplication operands such as A `[BM, BK]` and B `[BK, BN]` even when their shapes do not match the output tile or CUDA thread shape. Version 12 introduces the layout model and its validation. It does not yet emit -cooperative shared-memory loads or implement `tl.dot`; those are the next -lowering stages. +cooperative shared-memory loads; those are the next CUDA lowering stage. + +## `tl.dot` semantics + +Rank-2 `f32` blocks can be combined with the public `tl.dot` operation: + +```python +lhs = tl.zeros((BM, BK), tl.float32) +rhs = tl.zeros((BK, BN), tl.float32) +result = tl.dot(lhs, rhs) +``` + +The operands must have shapes `[M, K]` and `[K, N]`. Their inner dimensions +must match, and the result has shape `[M, N]`: + +```text +%0 = zeros {shape=(4, 16), dtype=f32} : block<4x16 x f32> +%1 = zeros {shape=(16, 8), dtype=f32} : block<16x8 x f32> +%2 = dot %0, %1 : block<4x8 x f32> +``` + +The expression-tree type inference and SSA verifier independently check operand +rank, `f32` element types, matching reduction dimensions, and the exact result +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. ## Example @@ -383,7 +417,9 @@ the positive constant step, definition order, and matching types and counts for carried inputs, region arguments, yielded values, and loop results. For block factory functions it checks that shapes are non-empty and positive, dtypes are supported, result block types match the declared shape/dtype, and `tl.full` has -a scalar fill value convertible to the requested dtype. +a scalar fill value convertible to the requested dtype. For `tl.dot`, it +requires two rank-2 `f32` operands, matching inner dimensions, and an exact +`[M, N]` rank-2 `f32` result. Straight-line verified SSA then runs through a small optimization pipeline: @@ -436,7 +472,9 @@ 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. + 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. - 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 @@ -446,9 +484,10 @@ these rewrite passes because they are not region-aware yet. 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 now describe A `[BM, BK]` and - B `[BK, BN]` independently of C `[BM, BN]`, but there is not yet a `tl.dot` - operation or CUDA lowering that stages these operands in shared memory. + 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. `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 e3a4cc0..8e68011 100644 --- a/src/mytriton/cuda_codegen.py +++ b/src/mytriton/cuda_codegen.py @@ -350,6 +350,8 @@ def emit(self, op: SSAOp) -> None: element_ty = self.scalar_type(result.ty) 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") elif op.opcode in self.BINARY_OPS: lhs = self.expression_operand(op.operands[0]) rhs = self.expression_operand(op.operands[1]) diff --git a/src/mytriton/language.py b/src/mytriton/language.py index b2739d4..6f32b36 100644 --- a/src/mytriton/language.py +++ b/src/mytriton/language.py @@ -1,6 +1,7 @@ from .trace import ( arange, constexpr, + dot, empty, exp, float32, @@ -23,6 +24,7 @@ __all__ = [ "arange", "constexpr", + "dot", "empty", "exp", "float32", diff --git a/src/mytriton/optim.py b/src/mytriton/optim.py index 17a5a39..4a77042 100644 --- a/src/mytriton/optim.py +++ b/src/mytriton/optim.py @@ -178,6 +178,7 @@ class CSEPass: "arange", "full", "zeros", + "dot", "add", "sub", "mul", diff --git a/src/mytriton/ssa.py b/src/mytriton/ssa.py index 96228d9..6a33bc3 100644 --- a/src/mytriton/ssa.py +++ b/src/mytriton/ssa.py @@ -7,6 +7,7 @@ Arange, BinOp, Const, + Dot, Empty, ExpandDims, ForRange, @@ -161,6 +162,16 @@ def lower_expr(self, expr): attrs={"shape": expr.shape, "dtype": expr.dtype}, ) + if isinstance(expr, Dot): + lhs = self.lower_expr(expr.lhs) + rhs = self.lower_expr(expr.rhs) + + return self.emit( + "dot", + expr, + operands=(lhs, rhs), + ) + if isinstance(expr, BinOp): if expr.op not in self.BINOPS: raise TypeError(f"Unsupported binary operator: {expr.op}") diff --git a/src/mytriton/ssa_verification.py b/src/mytriton/ssa_verification.py index 5917c0a..b157c70 100644 --- a/src/mytriton/ssa_verification.py +++ b/src/mytriton/ssa_verification.py @@ -26,6 +26,7 @@ class SSAVerifier: "empty": 0, "full": 1, "zeros": 0, + "dot": 2, "add": 2, "sub": 2, "mul": 2, @@ -147,6 +148,37 @@ def check_binary_numeric(self, index: int, op: SSAOp) -> None: expected_ty = self.with_shape(index, op, element, lhs_ty, rhs_ty) self.require_type(index, op, result_ty, expected_ty) + def check_dot(self, index: int, op: SSAOp) -> None: + lhs_ty = self.require_operand_type(index, op, op.operands[0], "lhs") + rhs_ty = self.require_operand_type(index, op, op.operands[1], "rhs") + result_ty = self.result_type(index, op) + + if not isinstance(lhs_ty, BlockType) or lhs_ty.rank != 2: + self.fail(index, op, f"dot lhs must be a rank-2 block, got {lhs_ty}") + + if not isinstance(rhs_ty, BlockType) or rhs_ty.rank != 2: + self.fail(index, op, f"dot rhs must be a rank-2 block, got {rhs_ty}") + + if lhs_ty.element != F32: + self.fail(index, op, f"dot lhs must have f32 elements, got {lhs_ty}") + + if rhs_ty.element != F32: + self.fail(index, op, f"dot rhs must have f32 elements, got {rhs_ty}") + + lhs_m, lhs_k = lhs_ty.shape + rhs_k, rhs_n = rhs_ty.shape + + if lhs_k != rhs_k: + self.fail( + index, + op, + "dot inner dimensions must match, " + f"got {lhs_ty.shape} and {rhs_ty.shape}", + ) + + expected_ty = BlockType((lhs_m, rhs_n), F32) + self.require_type(index, op, result_ty, expected_ty) + def check_unary(self, index: int, op: SSAOp) -> None: value_ty = self.require_operand_type(index, op, op.operands[0], "value") result_ty = self.result_type(index, op) @@ -565,6 +597,8 @@ def _verify_ops(self, ops: list[SSAItem], defined: set[int]) -> None: if op.opcode in {"add", "sub", "mul", "div", "cmp_lt"}: self.check_binary_numeric(index, op) + elif op.opcode == "dot": + self.check_dot(index, op) elif op.opcode == "and": self.check_binary_bool(index, op) elif op.opcode in {"neg", "exp"}: diff --git a/src/mytriton/trace.py b/src/mytriton/trace.py index 5dec038..dde9d93 100644 --- a/src/mytriton/trace.py +++ b/src/mytriton/trace.py @@ -82,6 +82,10 @@ def zeros( return Value(Zeros(_normalize_block_shape(shape), _require_block_dtype(dtype))) +def dot(lhs: Value, rhs: Value) -> Value: + return Value(Dot(unwrap(lhs), unwrap(rhs))) + + def load( ptr: Ptr, mask: Value | bool | None = None, @@ -277,6 +281,12 @@ class Zeros: dtype: ScalarType +@dataclass +class Dot: + lhs: Expression + rhs: Expression + + @dataclass class BinOp: op: str @@ -389,6 +399,7 @@ class ForRange: | Empty | Full | Zeros + | Dot | BinOp | AddPtr | Load diff --git a/src/mytriton/type_inference.py b/src/mytriton/type_inference.py index 5bad3f6..04c450e 100644 --- a/src/mytriton/type_inference.py +++ b/src/mytriton/type_inference.py @@ -10,6 +10,7 @@ BinOp, BlockType, Const, + Dot, Empty, ExpandDims, Full, @@ -89,6 +90,33 @@ def require_convertible( raise TypeError(f"{context} must be convertible to {destination}, got {source}") + def infer_dot(self, expr: Dot) -> BlockType: + lhs_ty = self.infer(expr.lhs) + rhs_ty = self.infer(expr.rhs) + + if not isinstance(lhs_ty, BlockType) or lhs_ty.rank != 2: + raise TypeError(f"dot lhs must be a rank-2 block, got {lhs_ty}") + + if not isinstance(rhs_ty, BlockType) or rhs_ty.rank != 2: + raise TypeError(f"dot rhs must be a rank-2 block, got {rhs_ty}") + + if lhs_ty.element != F32: + raise TypeError(f"dot lhs must have f32 elements, got {lhs_ty}") + + if rhs_ty.element != F32: + raise TypeError(f"dot rhs must have f32 elements, got {rhs_ty}") + + lhs_m, lhs_k = lhs_ty.shape + rhs_k, rhs_n = rhs_ty.shape + + if lhs_k != rhs_k: + raise TypeError( + "dot inner dimensions must match, " + f"got {lhs_ty.shape} and {rhs_ty.shape}" + ) + + return BlockType((lhs_m, rhs_n), F32) + def infer(self, expr) -> Type: key = id(expr) ty: Type @@ -145,6 +173,9 @@ def infer(self, expr) -> Type: ) ty = BlockType(expr.shape, expr.dtype) + elif isinstance(expr, Dot): + ty = self.infer_dot(expr) + elif isinstance(expr, BinOp): lhs = self.infer(expr.lhs) rhs = self.infer(expr.rhs) diff --git a/tests/test_dot.py b/tests/test_dot.py new file mode 100644 index 0000000..1f45f91 --- /dev/null +++ b/tests/test_dot.py @@ -0,0 +1,396 @@ +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 +from mytriton.block_shapes import cuda_threads_per_block +from mytriton.optim import CSEPass, DCEPass +from mytriton.ssa import SSAItem, SSALowering, SSAOp, SSAPrinter, SSAValue +from mytriton.ssa_verification import CompileError, SSAVerifier +from mytriton.trace import F32, I32, BlockType, Dot, Zeros +from mytriton.type_inference import TypeInference + + +@triton.jit +def dot_semantics_kernel( + out, + BM: tl.constexpr, + BK: tl.constexpr, + BN: tl.constexpr, +): + lhs = tl.zeros((BM, BK), tl.float32) + rhs = tl.zeros((BK, BN), tl.float32) + result = tl.dot(lhs, rhs) + + offsets_m = tl.arange(0, BM)[:, None] + offsets_n = tl.arange(0, BN)[None, :] + offsets = offsets_m * BN + offsets_n + + tl.store(out + offsets, result) + + +def infer_type(value) -> BlockType: + ty = TypeInference().infer(value.expr) + assert isinstance(ty, BlockType) + return ty + + +def test_dot_builds_expression_tree_node() -> None: + lhs = tl.zeros((4, 16), tl.float32) + rhs = tl.zeros((16, 8), tl.float32) + + result = tl.dot(lhs, rhs) + + assert isinstance(result.expr, Dot) + assert result.expr.lhs is lhs.expr + assert result.expr.rhs is rhs.expr + assert isinstance(result.expr.lhs, Zeros) + assert isinstance(result.expr.rhs, Zeros) + + +@pytest.mark.parametrize( + ("lhs_shape", "rhs_shape", "expected_shape"), + [ + ((4, 16), (16, 8), (4, 8)), + ((1, 16), (16, 8), (1, 8)), + ((4, 16), (16, 1), (4, 1)), + ((1, 1), (1, 1), (1, 1)), + ], +) +def test_dot_type_inference( + lhs_shape: tuple[int, int], + rhs_shape: tuple[int, int], + expected_shape: tuple[int, int], +) -> None: + lhs = tl.zeros(lhs_shape, tl.float32) + rhs = tl.zeros(rhs_shape, tl.float32) + + result = tl.dot(lhs, rhs) + + assert infer_type(result) == BlockType(expected_shape, F32) + + +def test_dot_rejects_non_rank2_lhs() -> None: + lhs = tl.zeros((16,), tl.float32) + rhs = tl.zeros((16, 8), tl.float32) + + with pytest.raises( + TypeError, + match="dot lhs must be a rank-2 block", + ): + TypeInference().infer(tl.dot(lhs, rhs).expr) + + +def test_dot_rejects_non_rank2_rhs() -> None: + lhs = tl.zeros((4, 16), tl.float32) + rhs = tl.zeros((16,), tl.float32) + + with pytest.raises( + TypeError, + match="dot rhs must be a rank-2 block", + ): + TypeInference().infer(tl.dot(lhs, rhs).expr) + + +def test_dot_rejects_mismatched_inner_dimensions() -> None: + lhs = tl.zeros((4, 16), tl.float32) + rhs = tl.zeros((8, 4), tl.float32) + + with pytest.raises( + TypeError, + match="dot inner dimensions must match", + ): + TypeInference().infer(tl.dot(lhs, rhs).expr) + + +def test_dot_rejects_non_f32_lhs() -> None: + lhs = tl.zeros((4, 16), tl.int32) + rhs = tl.zeros((16, 8), tl.float32) + + with pytest.raises( + TypeError, + match="dot lhs must have f32 elements", + ): + TypeInference().infer(tl.dot(lhs, rhs).expr) + + +def test_dot_rejects_non_f32_rhs() -> None: + lhs = tl.zeros((4, 16), tl.float32) + rhs = tl.zeros((16, 8), tl.int32) + + with pytest.raises( + TypeError, + match="dot rhs must have f32 elements", + ): + TypeInference().infer(tl.dot(lhs, rhs).expr) + + +def test_dot_lowers_to_ssa() -> None: + lhs = tl.zeros((4, 16), tl.float32) + rhs = tl.zeros((16, 8), tl.float32) + result = tl.dot(lhs, rhs) + + lowering = SSALowering() + ssa_result = lowering.lower_expr(result.expr) + + assert str(ssa_result) == "%2" + assert ssa_result.ty == BlockType((4, 8), F32) + + expected_ssa = dedent( + """\ + %0 = zeros {shape=(4, 16), dtype=f32} : block<4x16 x f32> + %1 = zeros {shape=(16, 8), dtype=f32} : block<16x8 x f32> + %2 = dot %0, %1 : block<4x8 x f32> + """ + ).rstrip("\n") + + assert SSAPrinter().print_ops(lowering.ops) == expected_ssa + + +def test_dot_ssa_reuses_shared_operand() -> None: + lhs = tl.zeros((4, 4), tl.float32) + result = tl.dot(lhs, lhs) + + lowering = SSALowering() + lowering.lower_expr(result.expr) + + expected_ssa = dedent( + """\ + %0 = zeros {shape=(4, 4), dtype=f32} : block<4x4 x f32> + %1 = dot %0, %0 : block<4x4 x f32> + """ + ).rstrip("\n") + + assert SSAPrinter().print_ops(lowering.ops) == expected_ssa + + +def test_dot_ssa_passes_verification() -> None: + lhs = tl.zeros((4, 16), tl.float32) + rhs = tl.zeros((16, 8), tl.float32) + result = tl.dot(lhs, rhs) + + lowering = SSALowering() + lowering.lower_expr(result.expr) + + verified = SSAVerifier(block_size=32).verify(lowering.ops) + + assert verified == lowering.ops + + +def make_dot_ssa( + lhs_ty: BlockType, + rhs_ty: BlockType, + result_ty: BlockType, +) -> list[SSAItem]: + lhs = SSAValue(id=0, ty=lhs_ty) + rhs = SSAValue(id=1, ty=rhs_ty) + result = SSAValue(id=2, ty=result_ty) + + return [ + SSAOp( + opcode="zeros", + result=lhs, + attrs={ + "shape": lhs_ty.shape, + "dtype": lhs_ty.element, + }, + ), + SSAOp( + opcode="zeros", + result=rhs, + attrs={ + "shape": rhs_ty.shape, + "dtype": rhs_ty.element, + }, + ), + SSAOp( + opcode="dot", + operands=(lhs, rhs), + result=result, + ), + ] + + +def test_dot_verifier_rejects_wrong_result_shape() -> None: + ops = make_dot_ssa( + lhs_ty=BlockType((4, 16), F32), + rhs_ty=BlockType((16, 8), F32), + result_ty=BlockType((4, 7), F32), + ) + + with pytest.raises( + CompileError, + match=r"expected block<4x8 x f32>, got block<4x7 x f32>", + ): + SSAVerifier(block_size=32).verify(ops) + + +def test_dot_verifier_rejects_mismatched_inner_dimensions() -> None: + ops = make_dot_ssa( + lhs_ty=BlockType((4, 16), F32), + rhs_ty=BlockType((8, 4), F32), + result_ty=BlockType((4, 4), F32), + ) + + with pytest.raises( + CompileError, + match="dot inner dimensions must match", + ): + SSAVerifier(block_size=32).verify(ops) + + +def test_dot_verifier_rejects_non_rank2_lhs() -> None: + ops = make_dot_ssa( + lhs_ty=BlockType((16,), F32), + rhs_ty=BlockType((16, 8), F32), + result_ty=BlockType((4, 8), F32), + ) + + with pytest.raises( + CompileError, + match="dot lhs must be a rank-2 block", + ): + SSAVerifier(block_size=32).verify(ops) + + +def test_dot_verifier_rejects_non_rank2_rhs() -> None: + ops = make_dot_ssa( + lhs_ty=BlockType((4, 16), F32), + rhs_ty=BlockType((16,), F32), + result_ty=BlockType((4, 8), F32), + ) + + with pytest.raises( + CompileError, + match="dot rhs must be a rank-2 block", + ): + SSAVerifier(block_size=32).verify(ops) + + +def test_dot_verifier_rejects_non_f32_lhs() -> None: + ops = make_dot_ssa( + lhs_ty=BlockType((4, 16), I32), + rhs_ty=BlockType((16, 8), F32), + result_ty=BlockType((4, 8), F32), + ) + + with pytest.raises( + CompileError, + match="dot lhs must have f32 elements", + ): + SSAVerifier(block_size=32).verify(ops) + + +def test_dot_verifier_rejects_non_f32_rhs() -> None: + ops = make_dot_ssa( + lhs_ty=BlockType((4, 16), F32), + rhs_ty=BlockType((16, 8), I32), + result_ty=BlockType((4, 8), F32), + ) + + with pytest.raises( + CompileError, + match="dot rhs must have f32 elements", + ): + SSAVerifier(block_size=32).verify(ops) + + +def test_ast_frontend_lowers_dot_to_verified_ssa() -> None: + out = np.empty(4 * 8, dtype=np.float32) + + bound = dot_semantics_kernel.signature.bind( + out, + BM=4, + BK=16, + BN=8, + ) + + ops, _ = trace( + dot_semantics_kernel.fn, + dot_semantics_kernel.signature, + bound.arguments, + ) + + ssa_ops = SSALowering().lower(ops) + threads_per_block = cuda_threads_per_block(ssa_ops) + + assert threads_per_block == 32 + SSAVerifier(threads_per_block).verify(ssa_ops) + + expected_ssa = dedent( + """\ + %0 = zeros {shape=(4, 16), dtype=f32} : block<4x16 x f32> + %1 = zeros {shape=(16, 8), dtype=f32} : block<16x8 x f32> + %2 = dot %0, %1 : block<4x8 x f32> + %3 = arange {start=0, end=4} : vector<4 x i32> + %4 = expand_dims %3 {axis=1} : block<4x1 x i32> + %5 = mul %4, 8 : block<4x1 x i32> + %6 = arange {start=0, end=8} : vector<8 x i32> + %7 = expand_dims %6 {axis=0} : block<1x8 x i32> + %8 = add %5, %7 : block<4x8 x i32> + %9 = addptr out, %8 : block<4x8 x ptr> + store %9, %2, none + """ + ).rstrip("\n") + + assert SSAPrinter().print_ops(ssa_ops) == expected_ssa + + +def test_cuda_backend_rejects_dot_before_lowering(monkeypatch) -> None: + monkeypatch.setenv("MYTRITON_BACKEND", "cuda") + out = np.empty(4 * 8, dtype=np.float32) + + dot_semantics_kernel.clear_cache() + + with pytest.raises( + TypeError, + match=r"CUDA lowering for tl\.dot is not implemented", + ): + dot_semantics_kernel[(1,)]( + out, + BM=4, + BK=16, + BN=8, + ) + + +def test_cse_reuses_duplicate_dot() -> None: + lhs = tl.zeros((4, 16), tl.float32) + rhs = tl.zeros((16, 8), tl.float32) + + first = tl.dot(lhs, rhs) + second = tl.dot(lhs, rhs) + result = first + second + + lowering = SSALowering() + lowering.lower_expr(result.expr) + + optimized = CSEPass().run(lowering.ops) + + expected_ssa = dedent( + """\ + %0 = zeros {shape=(4, 16), dtype=f32} : block<4x16 x f32> + %1 = zeros {shape=(16, 8), dtype=f32} : block<16x8 x f32> + %2 = dot %0, %1 : block<4x8 x f32> + %4 = add %2, %2 : block<4x8 x f32> + """ + ).rstrip("\n") + + assert SSAPrinter().print_ops(optimized) == expected_ssa + SSAVerifier(block_size=32).verify(optimized) + + +def test_dce_removes_unused_dot() -> None: + lhs = tl.zeros((4, 16), tl.float32) + rhs = tl.zeros((16, 8), tl.float32) + result = tl.dot(lhs, rhs) + + lowering = SSALowering() + lowering.lower_expr(result.expr) + + optimized = DCEPass().run(lowering.ops) + + assert optimized == []