Version 13: Add tl.dot semantics

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<f32>` parameters as

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])

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",

src/mytriton/optim.py

@@ -178,6 +178,7 @@ class CSEPass:
         "arange",
         "full",
         "zeros",
+        "dot",
         "add",
         "sub",
         "mul",

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}")

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"}:

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

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)