From 9637858ad9b64ac8c02adcc0ec0b1d749ae6a6fa Mon Sep 17 00:00:00 2001 From: Michael C Li <63594253+Happymic@users.noreply.github.com> Date: Sat, 5 Sep 2026 21:11:15 +0100 Subject: [PATCH] Add Matrix SRAM views and prepared L_TILE lowering --- .github/workflows/ci.yml | 17 + assembler/assembly_to_binary.py | 133 ++- assembler/tests/test_matrix_view.py | 166 +++ aten/plena/isa_matrix.py | 35 +- aten/plena/isa_matrix_view.py | 661 ++++++++++++ aten/plena/memory.py | 72 ++ aten/plena/program_lcompute.py | 1521 +++++++++++++++++++++++++++ aten/plena/program_matrix_ops.py | 442 +++++++- aten/tests/test_lcompute.py | 690 ++++++++++++ aten/tests/test_matrix_view.py | 399 +++++++ doc/operation.svh | 7 +- doc/plena_isa_spec.md | 64 ++ 12 files changed, 4170 insertions(+), 37 deletions(-) create mode 100644 assembler/tests/test_matrix_view.py create mode 100644 aten/plena/isa_matrix_view.py create mode 100644 aten/plena/program_lcompute.py create mode 100644 aten/tests/test_lcompute.py create mode 100644 aten/tests/test_matrix_view.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 3732abc..3cc9dfb 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -135,3 +135,20 @@ jobs: asm.generate_binary('/tmp/smolvlm2.asm', '/tmp/smolvlm2.mem') print('SmolVLM2 assembled OK') " + + matrix-view-tests: + name: Matrix SRAM view and L_TILE tests + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: '3.10' + - name: Install test dependencies + run: pip install pytest pyyaml + - name: Test encoding, layouts, ownership and lowering + run: | + PYTHONPATH=. python3 -m pytest -q \ + assembler/tests/test_matrix_view.py \ + aten/tests/test_matrix_view.py \ + aten/tests/test_lcompute.py diff --git a/assembler/assembly_to_binary.py b/assembler/assembly_to_binary.py index 9cbeb35..24ff4b6 100644 --- a/assembler/assembly_to_binary.py +++ b/assembler/assembly_to_binary.py @@ -1,5 +1,14 @@ from utils.load_config import load_svh_settings +from compiler.aten.plena.isa_matrix_view import ( + LTilePrimitive, + MatrixViewAxis, + encode_matrix_view_dma_word, + encode_l_tile_exec, + encode_l_tile_cfg, + validate_matrix_view_dominance, +) + from .parser import load_isa_definitions, parse_asm_file @@ -23,6 +32,16 @@ "V_TOPK", } ) +_PSEUDO_OPCODE_ALIASES = { + "L_TILE_CFG": "L_TILE", + "L_TILE_EXEC": "L_TILE", + "V_ADD_VV.MV": "V_ADD_VV", + "V_SUB_VV.MV": "V_SUB_VV", + "V_MUL_VV.MV": "V_MUL_VV", + "H_PREFETCH_V.MV": "H_PREFETCH_V", + "H_STORE_V.MV": "H_STORE_V", +} + _IMM_RS1_RD_OPS = frozenset( { "S_ADDI_INT", @@ -89,7 +108,8 @@ def _convert_to_binary(self, instruction): :return: Binary representation of the instruction """ # Example conversion logic (to be replaced with actual logic) - opcode = self.isa_definitions[instruction.opcode] + mnemonic = instruction.opcode + opcode = self.isa_definitions[_PSEUDO_OPCODE_ALIASES.get(mnemonic, mnemonic)] rd = instruction.rd rs1 = instruction.rs1 rs2 = instruction.rs2 @@ -105,7 +125,114 @@ def _convert_to_binary(self, instruction): # Treat omitted rmask deterministically as "mask disabled" instead of crashing on None << ... rmask = 0 - if instruction.opcode in _IMM_RS1_RD_OPS: + if mnemonic == "L_TILE_CFG": + # Text: L_TILE_CFG slot, gp_shape, gp_map. + slot = rd + shape_register = rs1 + map_register = rs2 + if slot is None or shape_register is None or map_register is None: + raise ValueError("L_TILE_CFG requires slot, shape register, and map register") + if not 0 <= slot < 4: + raise ValueError(f"L_TILE_CFG slot must be in [0, 4), got {slot}") + binary_instruction = encode_l_tile_cfg( + slot=slot, + shape_register=shape_register, + map_register=map_register, + ) + elif mnemonic == "L_TILE_EXEC": + # Text: L_TILE_EXEC gp_dst, gp_src1, gp_scale, primitive[, axis_mask]. + # axis_mask[0] selects source columns and axis_mask[1] selects scale + # columns. Omission is the byte-compatible row/row form. + if rd is None or rs1 is None or rs2 is None or rstride is None: + raise ValueError( + "L_TILE_EXEC requires destination/source base registers and primitive" + ) + try: + primitive = LTilePrimitive(rstride) + except ValueError as error: + raise ValueError(f"reserved L_TILE primitive {rstride}") from error + axis_mask = 0 if funct1 is None else funct1 + if not isinstance(axis_mask, int) or not 0 <= axis_mask <= 0b11: + raise ValueError("L_TILE_EXEC axis mask must be in [0, 3]") + binary_instruction = encode_l_tile_exec( + dst_register=rd, + src1_register=rs1, + src2_register=rs2, + primitive=primitive, + source_axis=MatrixViewAxis((axis_mask >> 0) & 1), + scale_axis=MatrixViewAxis((axis_mask >> 1) & 1), + ) + elif mnemonic in {"H_PREFETCH_V.MV", "H_STORE_V.MV"}: + # Existing vector DMA with an explicit Matrix-view destination or + # source: rd, rs1, rs2, rstride, precision, view_slot. + if None in (rd, rs1, rs2, rstride, funct1, instruction.funct2): + raise ValueError(f"{mnemonic} requires all six operands") + if not isinstance(funct1, int) or not 0 <= funct1 <= 2: + raise ValueError(f"{mnemonic}: precision must be in [0, 2]") + slot = instruction.funct2 + if not isinstance(slot, int) or not 0 <= slot < 4: + raise ValueError(f"{mnemonic}: Matrix view slot must be in [0, 4)") + legacy_word = ( + (funct1 << (opw + 4 * ow)) + + (rstride << (opw + 3 * ow)) + + (rs2 << (opw + 2 * ow)) + + (rs1 << (opw + ow)) + + (rd << opw) + + opcode + ) + binary_instruction = encode_matrix_view_dma_word(legacy_word, slot=slot) + elif mnemonic == "M_MM_WO" and rstride is not None: + # View-qualified existing writeback. The 18-bit immediate uses its + # top bit as an explicit view marker and the next two bits as the + # slot. Legacy three-operand words retain marker=0 byte-for-byte. + if rd is None or rs1 is None or imm is None: + raise ValueError("M_MM_WO requires base, stride register, and offset") + if not 0 <= rstride < 4: + raise ValueError(f"M_MM_WO Matrix view slot must be in [0, 4), got {rstride}") + if not 0 <= imm < (1 << 15): + raise ValueError( + f"view-qualified M_MM_WO offset must fit 15 bits, got {imm}" + ) + encoded_imm = (1 << 17) | (rstride << 15) | imm + binary_instruction = ( + (encoded_imm << (opw + 2 * ow)) + + (rs1 << (opw + ow)) + + (rd << opw) + + opcode + ) + elif mnemonic in _RS2_RS1_RD_OPS and mnemonic.startswith("M_") and rstride is not None: + # A fourth Matrix operand is an explicit view slot. funct1=0 keeps + # the legacy no-view word; codes 1..4 select slots 0..3. + if not 0 <= rstride < 4: + raise ValueError(f"{mnemonic}: Matrix view slot must be in [0, 4), got {rstride}") + binary_instruction = ( + ((rstride + 1) << (opw + 4 * ow)) + + (rs2 << (opw + 2 * ow)) + + (rs1 << (opw + ow)) + + (rd << opw) + + opcode + ) + elif mnemonic in {"V_ADD_VV.MV", "V_SUB_VV.MV", "V_MUL_VV.MV"}: + if not isinstance(funct1, int) or isinstance(funct1, bool): + raise TypeError(f"{mnemonic}: Matrix-view operand mask must be an int") + if funct1 == 0: + raise ValueError(f"{mnemonic}: Matrix-view operand mask cannot be zero") + if not 1 <= funct1 <= 7: + raise ValueError(f"{mnemonic}: Matrix-view operand mask must be in [1, 7]") + binary_instruction = ( + ((funct1 | 0x8) << (opw + 4 * ow)) + + ((rmask or 0) << (opw + 3 * ow)) + + ((rs2 or 0) << (opw + 2 * ow)) + + ((rs1 or 0) << (opw + ow)) + + ((rd or 0) << opw) + + opcode + ) + elif instruction.opcode in _IMM_RS1_RD_OPS: + if mnemonic == "M_MM_WO" and not 0 <= imm < (1 << 17): + raise ValueError( + "legacy M_MM_WO offset must fit 17 bits; bit 17 is the " + f"Matrix-view marker, got {imm}" + ) binary_instruction = (imm << (opw + 2 * ow)) + (rs1 << (opw + ow)) + (rd << opw) + opcode elif instruction.opcode in _IMM_RD_OPS: binary_instruction = (imm << (opw + ow)) + (rd << opw) + opcode @@ -153,6 +280,8 @@ def generate_binary(self, asm_file: str, output_file: str): """ Generate binary instructions from the assembled instructions. """ + with open(asm_file) as source: + validate_matrix_view_dominance(source.read()) instructions = parse_asm_file(asm_file) binary_instructions = [] for instruction in instructions: diff --git a/assembler/tests/test_matrix_view.py b/assembler/tests/test_matrix_view.py new file mode 100644 index 0000000..9f86510 --- /dev/null +++ b/assembler/tests/test_matrix_view.py @@ -0,0 +1,166 @@ +from pathlib import Path + +import pytest + +from assembler.assembly_to_binary import AssemblyToBinary +from assembler.parser import Instruction +from compiler.aten.plena.isa_matrix_view import LTilePrimitive + + +ROOT = Path(__file__).resolve().parents[2] + + +def _assembler() -> AssemblyToBinary: + return AssemblyToBinary( + str(ROOT / "doc" / "operation.svh"), + str(ROOT / "doc" / "configuration.svh"), + ) + + +def test_config_uses_0x3f_and_the_config_funct() -> None: + word = _assembler()._convert_to_binary( + Instruction("L_TILE_CFG", 2, 7, 9, None, None, None) + ) + assert word & 0x3F == 0x3F + assert (word >> 22) & 0xF == 1 + + +@pytest.mark.parametrize("primitive", list(LTilePrimitive)) +def test_exec_uses_0x3f_and_a_distinct_funct( + primitive: LTilePrimitive, +) -> None: + word = _assembler()._convert_to_binary( + Instruction("L_TILE_EXEC", 4, 5, 6, int(primitive), None, None) + ) + assert word & 0x3F == 0x3F + assert (word >> 22) & 0xF == 3 + assert (word >> 18) & 0xF == int(primitive) + + +def test_exec_rejects_a_reserved_primitive() -> None: + with pytest.raises(ValueError, match="reserved L_TILE primitive"): + _assembler()._convert_to_binary( + Instruction("L_TILE_EXEC", 4, 5, 6, 3, None, None) + ) + + +def test_exec_encodes_optional_operand_axes_without_spending_an_opcode() -> None: + word = _assembler()._convert_to_binary( + Instruction("L_TILE_EXEC", 4, 5, 6, 1, 0b10, None) + ) + assert word & 0x3F == 0x3F + assert (word >> 22) & 0xF == 3 + assert (word >> 26) & 0b11 == 0b10 + + +def test_exec_rejects_reserved_operand_axes() -> None: + with pytest.raises(ValueError, match="axis mask"): + _assembler()._convert_to_binary( + Instruction("L_TILE_EXEC", 4, 5, 6, 1, 4, None) + ) + + +def test_matrix_fourth_operand_is_an_explicit_view_slot() -> None: + assembler = _assembler() + # M_BMV architecturally consumes rd as its matrix offset, so a non-zero rd + # locks every register field without violating M_MV's rd=0 convention. + plain = assembler._convert_to_binary(Instruction("M_BMV", 5, 1, 2, None, None, None)) + viewed = assembler._convert_to_binary(Instruction("M_BMV", 5, 1, 2, 3, None, None)) + assert plain == 0x09 | (5 << 6) | (1 << 10) | (2 << 14) + assert viewed & ((1 << 22) - 1) == plain + assert (viewed >> 22) & 0xF == 4 + + +def test_view_slot_out_of_range_is_rejected() -> None: + with pytest.raises(ValueError, match="view slot"): + _assembler()._convert_to_binary(Instruction("M_TMM", 0, 1, 2, 4, None, None)) + + +def test_matrix_view_vector_alias_keeps_arithmetic_opcode_and_marks_operands() -> None: + assembler = _assembler() + plain = assembler._convert_to_binary( + Instruction("V_ADD_VV", 4, 5, 6, 0, 0, None) + ) + viewed = assembler._convert_to_binary( + Instruction("V_ADD_VV.MV", 4, 5, 6, 0, 0b110, None) + ) + assert plain & 0x3F == viewed & 0x3F == 0x0D + assert plain >> 22 == 0 + assert (viewed >> 22) & 0xF == 0x8 | 0b110 + + +def test_matrix_view_vector_alias_rejects_an_empty_operand_mask() -> None: + with pytest.raises(ValueError, match="cannot be zero"): + _assembler()._convert_to_binary( + Instruction("V_MUL_VV.MV", 1, 1, 2, 0, 0, None) + ) + + +def test_accumulator_writeback_uses_existing_m_mm_wo_with_explicit_view() -> None: + assembler = _assembler() + legacy = assembler._convert_to_binary( + Instruction("M_MM_WO", 4, 0, None, None, None, None, imm=5) + ) + viewed = assembler._convert_to_binary( + Instruction("M_MM_WO", 4, 0, None, 2, None, None, imm=5) + ) + assert legacy & 0x3F == viewed & 0x3F == 0x06 + legacy_imm = legacy >> 14 + viewed_imm = viewed >> 14 + assert legacy_imm == 5 + assert viewed_imm == (1 << 17) | (2 << 15) | 5 + + +@pytest.mark.parametrize( + ("mnemonic", "opcode"), + [("H_PREFETCH_V.MV", 0x29), ("H_STORE_V.MV", 0x2A)], +) +def test_vector_dma_matrix_view_form_reuses_opcode_and_names_slot( + mnemonic: str, opcode: int +) -> None: + word = _assembler()._convert_to_binary( + Instruction(mnemonic, 1, 2, 3, 1, 2, 3) + ) + assert word & 0x3F == opcode + assert word >> 31 == 1 + assert (word >> 29) & 0b11 == 3 + assert (word >> 26) & 0b111 == 0 + assert (word >> 22) & 0xF == 2 + assert word & ((1 << 26) - 1) == ( + opcode | (1 << 6) | (2 << 10) | (3 << 14) | (1 << 18) | (2 << 22) + ) + + +@pytest.mark.parametrize("precision", [0, 1, 2]) +def test_vector_dma_matrix_view_accepts_all_canonical_precisions( + precision: int, +) -> None: + word = _assembler()._convert_to_binary( + Instruction("H_PREFETCH_V.MV", 1, 2, 3, 1, precision, 0) + ) + assert (word >> 22) & 0xF == precision + + +def test_vector_dma_matrix_view_form_rejects_reserved_slot() -> None: + with pytest.raises(ValueError, match="view slot"): + _assembler()._convert_to_binary( + Instruction("H_PREFETCH_V.MV", 1, 2, 3, 1, 2, 4) + ) + + +def test_legacy_m_mm_wo_cannot_alias_the_view_marker() -> None: + with pytest.raises(ValueError, match="bit 17"): + _assembler()._convert_to_binary( + Instruction("M_MM_WO", 4, 0, None, None, None, None, imm=1 << 17) + ) + + +def test_generate_binary_enforces_matrix_view_dominance(tmp_path: Path) -> None: + source = tmp_path / "bad.asm" + output = tmp_path / "bad.mem" + source.write_text("M_TMM gp0, gp1, gp2, 2\n") + + with pytest.raises(ValueError, match="before a dominating configuration"): + _assembler().generate_binary(str(source), str(output)) + + assert not output.exists() diff --git a/aten/plena/isa_matrix.py b/aten/plena/isa_matrix.py index 0f26b68..dad9d55 100644 --- a/aten/plena/isa_matrix.py +++ b/aten/plena/isa_matrix.py @@ -373,6 +373,9 @@ def vram_sub_projection_microtile_accumulate_asm( k_block_count: int, write_out: bool, gp_regs: list[int] | None = None, + matrix_view_base: int | None = None, + matrix_view_logical_offset: int | None = None, + matrix_view_slot: int = 0, ) -> str: """Emit M_MM for one 4x4 microtile, optionally flushing with M_MM_WO. @@ -437,8 +440,23 @@ def vram_sub_projection_microtile_accumulate_asm( lines.append(f"M_MM 0, gp{gp_mat}, gp{gp_act}") if write_out: - lines.extend(load_large_int(gp_result, result_addr)) - lines.append(f"M_MM_WO gp{gp_result}, gp0, 0") + if matrix_view_base is not None: + if matrix_view_logical_offset is None: + raise ValueError( + "Matrix-view writeback requires a logical element offset" + ) + lines.extend(load_large_int(gp_result, matrix_view_base)) + lines.append( + f"M_MM_WO gp{gp_result}, gp0, " + f"{matrix_view_logical_offset}, {matrix_view_slot}" + ) + else: + if matrix_view_logical_offset is not None: + raise ValueError( + "Matrix-view logical offset was provided without a Matrix base" + ) + lines.extend(load_large_int(gp_result, result_addr)) + lines.append(f"M_MM_WO gp{gp_result}, gp0, 0") return "\n".join(lines) + "\n" @@ -457,11 +475,16 @@ def vram_sub_projection_microtile_accumulate_to( k_block_start: int, k_block_count: int, write_out: bool, + gp_regs: list[int] | None = None, + matrix_view_base: int | None = None, + matrix_view_logical_offset: int | None = None, + matrix_view_slot: int = 0, ) -> str: result_vram_addr, _target_base_addr, _target_rows = self._target_tile_addr( target_matrix, target_row_idx, target_col_idx ) - gp_regs = self.register_allocator.allocate_gp(3) + owns_gp_regs = gp_regs is None + gp_regs = self.register_allocator.allocate_gp(3) if gp_regs is None else gp_regs try: asm = self.vram_sub_projection_microtile_accumulate_asm( vram_mat_name=vram_mat_name, @@ -475,9 +498,13 @@ def vram_sub_projection_microtile_accumulate_to( k_block_count=k_block_count, write_out=write_out, gp_regs=gp_regs, + matrix_view_base=matrix_view_base, + matrix_view_logical_offset=matrix_view_logical_offset, + matrix_view_slot=matrix_view_slot, ) finally: - self.register_allocator.free_gp(gp_regs) + if owns_gp_regs: + self.register_allocator.free_gp(gp_regs) return self._emit(asm) def vram_sub_projection_packed_skinny_microtile_accumulate_asm( diff --git a/aten/plena/isa_matrix_view.py b/aten/plena/isa_matrix_view.py new file mode 100644 index 0000000..4b09e2e --- /dev/null +++ b/aten/plena/isa_matrix_view.py @@ -0,0 +1,661 @@ +"""Packed ISA contract for compiler-programmable Matrix-SRAM views. + +The view describes physical placement only. It does not encode an arithmetic +operation, a model, or a runtime traversal. Existing Matrix/Vector opcodes +select a configured slot explicitly and retain their original mathematics. + +Two GP values are sufficient for the hot form: + +``shape`` + ``rows_minus_one[11:0] | cols_minus_one[23:12] | + tile_count_minus_one[31:24]`` + +``mapping`` + ``tile_pitch_rows[15:0] | reserved[21:16] | + tile_phase_stride[27:22] | flags[31:28]`` + +``tile_pitch_rows`` is measured in full physical Matrix-SRAM rows (one bank +word from every bank), not elements. A zero pitch is legal only when the +inter-tile phase keeps every logical bank word distinct. That compact form +is what places several logical head rows in the same physical row but in +different banks. The injectivity check below rejects an unsafe zero pitch. +""" + +from __future__ import annotations + +from collections import deque +from dataclasses import dataclass +from enum import IntEnum, IntFlag + + +L_MVIEW_OPCODE = 0x3F +L_MVIEW_MAX_SLOTS = 4 +# Version 3 removes the programmable row coefficient and optional storage +# precision. Bits [21:16] and flag bits [2:0] are now reserved and must trap. +L_MVIEW_CONTRACT_VERSION = 3 +MATRIX_VIEW_DMA_MARKER = 1 << 31 +MATRIX_VIEW_DMA_SLOT_SHIFT = 29 +MATRIX_VIEW_DMA_HIGH_RESERVED_MASK = 0b111 << 26 + +_DIM_BITS = 12 +_DIM_MASK = (1 << _DIM_BITS) - 1 +_TILE_COUNT_BITS = 8 +_TILE_COUNT_MASK = (1 << _TILE_COUNT_BITS) - 1 +_PITCH_BITS = 16 +_PITCH_MASK = (1 << _PITCH_BITS) - 1 +_PHASE_BITS = 6 +_PHASE_MASK = (1 << _PHASE_BITS) - 1 + + +class MatrixViewForm(IntEnum): + """Function selector carried by the shared ``L_TILE`` opcode.""" + + CONFIG = 1 + EXEC = 3 + + +class LTilePrimitive(IntEnum): + """Model-independent tensor primitives executed over Matrix-SRAM views. + + Shapes, broadcasting and reduction axes come from the three configured + views. No primitive names a model or owns persistent state. + """ + + # dst[row, :] = scale[row, 0] * dst[row, :] + # + scale[row, 1] * src[row_or_broadcast, :] + # The decoder walks logical rows; the scale view is segment-broadcast. + SCALE_ACCUM = 0 + + # dst[0, col] += sum_row(src[row, col] * scale[row, 0]). The source is + # consumed as a group of logical columns, so this form exercises the + # transposable/diagonal Matrix-SRAM path rather than materialising a + # transpose or reducing one row at a time. + DOT_REDUCE = 1 + + # dst[row, :] += vector[row_or_broadcast, :] * scale[row, 0]. This is the + # generic outer/rank-1 update used by linear recurrences. + OUTER_UPDATE = 2 + + +class MatrixViewAxis(IntEnum): + """Logical line direction selected by an ``L_TILE.EXEC`` operand. + + Placement remains a property of the configured view. Axis is a use-site + property: the same physical cells may be consumed as logical rows during + decode or logical columns at a transpose boundary without creating a + second copy. + """ + + ROW = 0 + COLUMN = 1 + + +class MatrixViewFlags(IntFlag): + """Placement and element properties for one explicit scratchpad view.""" + + # Bits 0..2 are reserved. Bounds are always strict and storage is BF16. + BROADCAST_MINOR = 1 << 3 + + +@dataclass(frozen=True) +class MatrixViewShape: + rows: int + cols: int + tile_count: int = 1 + + def validate(self) -> None: + for name, value, maximum in ( + ("rows", self.rows, 1 << _DIM_BITS), + ("cols", self.cols, 1 << _DIM_BITS), + ("tile_count", self.tile_count, 1 << _TILE_COUNT_BITS), + ): + if not 1 <= value <= maximum: + raise ValueError(f"{name} must be in [1, {maximum}], got {value}") + + def pack(self) -> int: + self.validate() + return ( + (self.rows - 1) + | (self.cols - 1) << _DIM_BITS + | (self.tile_count - 1) << (2 * _DIM_BITS) + ) + + @classmethod + def unpack(cls, word: int) -> "MatrixViewShape": + _require_u32(word, "shape word") + return cls( + rows=(word & _DIM_MASK) + 1, + cols=((word >> _DIM_BITS) & _DIM_MASK) + 1, + tile_count=((word >> (2 * _DIM_BITS)) & _TILE_COUNT_MASK) + 1, + ) + + +@dataclass(frozen=True) +class MatrixViewMap: + """Compiler-selected physical placement for one Matrix tensor view. + + PLENA's published diagonal row mapping remains fixed. The only programmable + map term is ``tile_phase_stride``: a six-bit phase applied between logical + tiles. A tensor's allocation base supplies the constant bank phase, so no + model-specific field or group identifier is needed. + + """ + + tile_pitch_rows: int + tile_phase_stride: int = 0 + flags: MatrixViewFlags = MatrixViewFlags(0) + + @property + def phased_enabled(self) -> bool: + """Whether the descriptor applies an inter-tile bank phase.""" + + return self.tile_phase_stride != 0 + + def validate(self) -> None: + if not 0 <= self.tile_pitch_rows <= _PITCH_MASK: + raise ValueError( + f"tile_pitch_rows must be in [0, {_PITCH_MASK}], got {self.tile_pitch_rows}" + ) + if not 0 <= self.tile_phase_stride <= _PHASE_MASK: + raise ValueError( + "tile_phase_stride must fit " + f"{_PHASE_BITS} bits, got {self.tile_phase_stride}" + ) + unknown = int(self.flags) & ~int(MatrixViewFlags.BROADCAST_MINOR) + if unknown: + raise ValueError(f"unknown Matrix-view flags 0x{unknown:x}") + + def pack(self) -> int: + self.validate() + return ( + self.tile_pitch_rows + | self.tile_phase_stride << 22 + | int(self.flags) << 28 + ) + + @classmethod + def unpack(cls, word: int) -> "MatrixViewMap": + _require_u32(word, "mapping word") + if (word >> 16) & _PHASE_MASK: + raise ValueError("Matrix-view mapping bits [21:16] are reserved") + mapping = cls( + tile_pitch_rows=word & _PITCH_MASK, + tile_phase_stride=(word >> 22) & _PHASE_MASK, + flags=MatrixViewFlags((word >> 28) & 0xF), + ) + mapping.validate() + return mapping + + +@dataclass(frozen=True) +class MatrixViewDescriptor: + shape: MatrixViewShape + mapping: MatrixViewMap + + def validate_for_machine(self, *, banks: int, bank_width: int) -> None: + self.shape.validate() + self.mapping.validate() + if not 1 <= banks <= 64 or banks & (banks - 1): + raise ValueError(f"banks must be a positive power of two no larger than 64, got {banks}") + if bank_width < 1: + raise ValueError(f"bank_width must be positive, got {bank_width}") + if self.shape.cols % bank_width: + raise ValueError( + f"cols {self.shape.cols} must contain whole {bank_width}-element bank words" + ) + words_per_row = self.shape.cols // bank_width + row_groups = (words_per_row + banks - 1) // banks + alpha = 1 + tile_phase_stride = self.mapping.tile_phase_stride + occupied: dict[tuple[int, int], tuple[int, int, int]] = {} + for tile in range(self.shape.tile_count): + for row in range(self.shape.rows): + for word in range(words_per_row): + bank_row = ( + tile * self.mapping.tile_pitch_rows + + row * row_groups + + word // banks + ) + bank = ( + alpha * bank_row + tile_phase_stride * tile + word + ) % banks + key = (bank, bank_row) + previous = occupied.setdefault(key, (tile, row, word)) + if previous != (tile, row, word): + raise ValueError( + "Matrix view aliases logical bank words: " + f"{previous} and {(tile, row, word)} both map to " + f"bank={bank}, row={bank_row}" + ) + + +@dataclass(frozen=True) +class MatrixViewAllocation: + """One compiler-owned tensor view placed at an explicit Matrix-SRAM base.""" + + name: str + base: int + descriptor: MatrixViewDescriptor + + +def validate_disjoint_matrix_views( + allocations: list[MatrixViewAllocation] | tuple[MatrixViewAllocation, ...], + *, + mlen: int, + banks: int, + bank_width: int, + depth_rows: int, + fixed_alpha: int = 1, + fixed_gamma: int = 0, +) -> dict[str, int]: + """Prove that all simultaneously-live views occupy distinct bank words. + + This is a static scratchpad allocation check, not a cache simulation. It + uses the same coordinate equation as the Rust Matrix SRAM and returns a few + capacity facts for reports. A bad placement fails during compilation. + """ + + if mlen != banks * bank_width: + raise ValueError("mlen must equal banks * bank_width") + if depth_rows <= 0: + raise ValueError("depth_rows must be positive") + if not 0 <= fixed_alpha < banks or not 0 <= fixed_gamma < banks: + raise ValueError("fixed Matrix-SRAM coefficients must fit the bank count") + + occupied: dict[tuple[int, int], tuple[str, int, int, int]] = {} + max_row = 0 + for allocation in allocations: + descriptor = allocation.descriptor + descriptor.validate_for_machine(banks=banks, bank_width=bank_width) + if allocation.base < 0 or allocation.base % bank_width: + raise ValueError( + f"{allocation.name}: base {allocation.base} is not bank-word aligned" + ) + base_row, base_offset = divmod(allocation.base, mlen) + base_bank = base_offset // bank_width + words_per_row = descriptor.shape.cols // bank_width + row_groups = (words_per_row + banks - 1) // banks + alpha = fixed_alpha + tile_phase_stride = descriptor.mapping.tile_phase_stride + for tile in range(descriptor.shape.tile_count): + for row in range(descriptor.shape.rows): + for word in range(words_per_row): + bank_row = ( + base_row + + tile * descriptor.mapping.tile_pitch_rows + + row * row_groups + + word // banks + ) + if bank_row >= depth_rows: + raise ValueError( + f"{allocation.name}: bank row {bank_row} exceeds " + f"Matrix-SRAM depth {depth_rows}" + ) + bank = ( + base_bank + + alpha * bank_row + + tile_phase_stride * tile + + fixed_gamma * (bank_row // banks) + + word + ) % banks + key = (bank, bank_row) + current = (allocation.name, tile, row, word) + previous = occupied.setdefault(key, current) + if previous != current: + raise ValueError( + "Matrix views alias physical bank word " + f"bank={bank}, row={bank_row}: {previous} and {current}" + ) + max_row = max(max_row, bank_row + 1) + return { + "bank_words": len(occupied), + "capacity_bank_words": banks * depth_rows, + "max_bank_row": max_row, + } + + +def encode_l_tile_cfg(*, slot: int, shape_register: int, map_register: int) -> int: + """Encode ``L_TILE_CFG slot, shape_reg, map_reg`` canonically.""" + + _require_slot(slot) + _require_register(shape_register, "shape_register") + _require_register(map_register, "map_register") + return ( + L_MVIEW_OPCODE + | shape_register << 6 + | map_register << 10 + | slot << 14 + | int(MatrixViewForm.CONFIG) << 22 + ) + + +def encode_l_tile_exec( + *, + dst_register: int, + src1_register: int, + src2_register: int, + primitive: LTilePrimitive | int, + source_axis: MatrixViewAxis | int = MatrixViewAxis.ROW, + scale_axis: MatrixViewAxis | int = MatrixViewAxis.ROW, +) -> int: + """Encode ``L_TILE_EXEC dst, src1, src2, primitive`` canonically. + + Slots 0/1/2 are the destination/source-1/source-2 descriptors. Keeping the + slot assignment positional leaves all operand bases explicit in the + instruction and avoids an implicit SELECT state. + """ + + for name, register in ( + ("dst_register", dst_register), + ("src1_register", src1_register), + ("src2_register", src2_register), + ): + _require_register(register, name) + try: + selected = LTilePrimitive(int(primitive)) + except ValueError as error: + raise ValueError(f"reserved L_TILE primitive {primitive}") from error + try: + selected_source_axis = MatrixViewAxis(int(source_axis)) + selected_scale_axis = MatrixViewAxis(int(scale_axis)) + except ValueError as error: + raise ValueError("reserved L_TILE operand axis") from error + return ( + L_MVIEW_OPCODE + | dst_register << 6 + | src1_register << 10 + | src2_register << 14 + | int(selected) << 18 + | int(MatrixViewForm.EXEC) << 22 + | int(selected_source_axis) << 26 + | int(selected_scale_axis) << 27 + ) + + +def decode_l_tile_exec_word( + word: int, +) -> tuple[int, int, int, LTilePrimitive, MatrixViewAxis, MatrixViewAxis]: + """Decode one canonical ``L_TILE_EXEC`` word.""" + + _require_u32(word, "instruction word") + if word & 0x3F != L_MVIEW_OPCODE or ((word >> 22) & 0xF) != MatrixViewForm.EXEC: + raise ValueError("word is not an L_TILE_EXEC encoding") + if word >> 28: + raise ValueError("reserved L_TILE_EXEC bits [31:28] must be zero") + try: + primitive = LTilePrimitive((word >> 18) & 0xF) + except ValueError as error: + raise ValueError("reserved L_TILE primitive") from error + return ( + (word >> 6) & 0xF, + (word >> 10) & 0xF, + (word >> 14) & 0xF, + primitive, + MatrixViewAxis((word >> 26) & 0x1), + MatrixViewAxis((word >> 27) & 0x1), + ) + + +def encode_matrix_view_dma_word(legacy_word: int, *, slot: int) -> int: + """Qualify an existing vector DMA word with an explicit Matrix view. + + The physical opcode and every legacy operand field remain unchanged. Bit + 31 marks the addressing form, bits 30:29 name a configured view, and bits + 28:26 remain canonical zero. This is not a new transfer operation. + """ + + _require_u32(legacy_word, "legacy DMA word") + _require_slot(slot) + if legacy_word >> 26: + raise ValueError("legacy DMA word uses reserved bits [31:26]") + if legacy_word & 0x3F not in {0x29, 0x2A}: + raise ValueError("Matrix-view DMA requires H_PREFETCH_V or H_STORE_V") + return legacy_word | MATRIX_VIEW_DMA_MARKER | (slot << MATRIX_VIEW_DMA_SLOT_SHIFT) + + +def decode_matrix_view_dma_slot(word: int) -> int | None: + """Return the explicit Matrix-view slot, or ``None`` for a legacy DMA.""" + + _require_u32(word, "DMA word") + if word >> 26 == 0: + return None + if not word & MATRIX_VIEW_DMA_MARKER or word & MATRIX_VIEW_DMA_HIGH_RESERVED_MASK: + raise ValueError("non-canonical Matrix-view DMA bits [31:26]") + slot = (word >> MATRIX_VIEW_DMA_SLOT_SHIFT) & 0b11 + _require_slot(slot) + return slot + + +def decode_l_mview_word(word: int) -> tuple[MatrixViewForm, int, int, int]: + """Return ``(form, slot, operand_a, operand_b)`` for a canonical word.""" + + _require_u32(word, "instruction word") + if word & 0x3F != L_MVIEW_OPCODE or word >> 26: + raise ValueError("word is not a canonical L_MVIEW encoding") + try: + form = MatrixViewForm((word >> 22) & 0xF) + except ValueError as error: + raise ValueError("reserved L_MVIEW form") from error + slot = (word >> 14) & 0xF + _require_slot(slot) + operand_a = (word >> 6) & 0xF + operand_b = (word >> 10) & 0xF + if (word >> 18) & 0xF: + raise ValueError("reserved L_MVIEW bits [21:18] must be zero") + if form is MatrixViewForm.EXEC: + raise ValueError("use decode_l_tile_exec_word for EXEC form") + return form, slot, operand_a, operand_b + + +def validate_matrix_view_dominance(assembly: str) -> None: + """Reject a Matrix consumer whose explicit view is not configured first. + + This is a syntactic must-dataflow property over the static loop control-flow + graph. There is no implicit SELECT state. Intersections at loop back-edges + prevent a first-iteration-only configuration from being + accepted as dominating every dynamic use. Inside a loop, C_BREAK has both a + fallthrough edge and an edge to the instruction after the matching loop end. + Intersecting those paths is conservative for both the public debug-exception + wording and the transactional emulator's loop-break behavior. + """ + + consumers = { + "M_MM", + "M_TMM", + "M_BMM", + "M_BTMM", + "M_MV", + "M_TMV", + "M_BMV", + "M_BTMV", + } + instructions: list[tuple[int, str, list[str]]] = [] + for line_number, raw in enumerate(assembly.splitlines(), start=1): + line = raw.split(";", 1)[0].split("//", 1)[0].strip() + if not line: + continue + parts = line.split(maxsplit=1) + opcode = parts[0] + operands = [] if len(parts) == 1 else [part.strip() for part in parts[1].split(",")] + instructions.append((line_number, opcode, operands)) + + loop_stack: list[int] = [] + end_to_start: dict[int, int] = {} + start_to_end: dict[int, int] = {} + break_to_start: dict[int, int] = {} + for index, (line_number, opcode, operands) in enumerate(instructions): + if opcode == "L_TILE_CFG": + if len(operands) != 3: + raise ValueError(f"line {line_number}: malformed {opcode}") + _require_slot(int(operands[0], 0)) + if opcode == "C_LOOP_START": + loop_stack.append(index) + elif opcode == "C_BREAK" and loop_stack: + break_to_start[index] = loop_stack[-1] + elif opcode == "C_LOOP_END": + if not loop_stack: + raise ValueError(f"line {line_number}: C_LOOP_END without C_LOOP_START") + start = loop_stack.pop() + end_to_start[index] = start + start_to_end[start] = index + if loop_stack: + line_number = instructions[loop_stack[-1]][0] + raise ValueError(f"line {line_number}: C_LOOP_START without C_LOOP_END") + + def successors(index: int) -> tuple[int, ...]: + _, opcode, _ = instructions[index] + fallthrough = index + 1 + if opcode == "C_LOOP_END": + start = end_to_start[index] + result = [start + 1] + if fallthrough < len(instructions): + result.append(fallthrough) + return tuple(result) + if opcode == "C_BREAK" and index in break_to_start: + result = [] + if fallthrough < len(instructions): + result.append(fallthrough) + after_loop = start_to_end[break_to_start[index]] + 1 + if after_loop < len(instructions): + result.append(after_loop) + return tuple(dict.fromkeys(result)) + return (fallthrough,) if fallthrough < len(instructions) else () + + State = frozenset[tuple[str, int]] + + def transfer(state: State, opcode: str, operands: list[str]) -> State: + result = set(state) + if opcode == "L_TILE_CFG": + slot = int(operands[0], 0) + result.difference_update({token for token in result if token[1] == slot}) + result.update({("shape", slot), ("mapping", slot), ("configured", slot)}) + return frozenset(result) + + if not instructions: + return + in_states: list[State | None] = [None] * len(instructions) + in_states[0] = frozenset() + work = deque([0]) + while work: + index = work.popleft() + state = in_states[index] + assert state is not None + _, opcode, operands = instructions[index] + output = transfer(state, opcode, operands) + for target in successors(index): + previous = in_states[target] + merged = output if previous is None else previous & output + if merged != previous: + in_states[target] = merged + work.append(target) + + for (line_number, opcode, operands), state in zip(instructions, in_states, strict=True): + if state is None: + continue + if opcode == "L_TILE_CFG": + if len(operands) != 3: + raise ValueError(f"line {line_number}: malformed L_TILE_CFG") + slot = int(operands[0], 0) + _require_slot(slot) + continue + if opcode == "L_TILE_EXEC": + if len(operands) not in {4, 5}: + raise ValueError(f"line {line_number}: malformed L_TILE_EXEC") + LTilePrimitive(int(operands[3], 0)) + if len(operands) == 5: + axis_mask = int(operands[4], 0) + if not 0 <= axis_mask <= 0b11: + raise ValueError( + f"line {line_number}: L_TILE_EXEC axis mask must be in [0, 3]" + ) + for slot in range(3): + if ("configured", slot) not in state: + raise ValueError( + f"line {line_number}: L_TILE_EXEC consumes Matrix view {slot} " + "before a dominating configuration" + ) + continue + if opcode == "M_MM_WO" and len(operands) == 4: + slot = int(operands[3], 0) + _require_slot(slot) + if ("configured", slot) not in state: + raise ValueError( + f"line {line_number}: M_MM_WO consumes Matrix view {slot} " + "before a dominating configuration" + ) + continue + if opcode in {"H_PREFETCH_V.MV", "H_STORE_V.MV"}: + if len(operands) != 6: + raise ValueError(f"line {line_number}: malformed {opcode}") + slot = int(operands[5], 0) + _require_slot(slot) + if ("configured", slot) not in state: + raise ValueError( + f"line {line_number}: {opcode} consumes Matrix view {slot} " + "before a dominating configuration" + ) + continue + if opcode.endswith(".MV"): + if len(operands) != 5: + raise ValueError(f"line {line_number}: malformed {opcode}") + mask = int(operands[4], 0) + if not 1 <= mask <= 0b111: + raise ValueError( + f"line {line_number}: Matrix-view operand mask must be in [1, 7]" + ) + for slot in range(3): + if mask & (1 << slot) and ("configured", slot) not in state: + raise ValueError( + f"line {line_number}: {opcode} consumes Matrix view {slot} " + "before a dominating configuration" + ) + continue + if opcode in consumers and len(operands) == 4: + slot = int(operands[3], 0) + _require_slot(slot) + if ("configured", slot) not in state: + raise ValueError( + f"line {line_number}: {opcode} consumes Matrix view {slot} " + "before a dominating configuration" + ) + + +def _require_slot(slot: int) -> None: + if not 0 <= slot < L_MVIEW_MAX_SLOTS: + raise ValueError(f"slot must be in [0, {L_MVIEW_MAX_SLOTS}), got {slot}") + + +def _require_register(register: int, name: str) -> None: + if not 0 <= register < 16: + raise ValueError(f"{name} must fit four bits, got {register}") + + +def _require_u32(word: int, name: str) -> None: + if not 0 <= word <= 0xFFFF_FFFF: + raise ValueError(f"{name} must be an unsigned 32-bit integer") + + +__all__ = [ + "LTilePrimitive", + "L_MVIEW_CONTRACT_VERSION", + "L_MVIEW_MAX_SLOTS", + "L_MVIEW_OPCODE", + "MATRIX_VIEW_DMA_HIGH_RESERVED_MASK", + "MATRIX_VIEW_DMA_MARKER", + "MATRIX_VIEW_DMA_SLOT_SHIFT", + "MatrixViewDescriptor", + "MatrixViewAllocation", + "MatrixViewAxis", + "MatrixViewFlags", + "MatrixViewForm", + "MatrixViewMap", + "MatrixViewShape", + "decode_l_tile_exec_word", + "decode_matrix_view_dma_slot", + "decode_l_mview_word", + "encode_l_tile_exec", + "encode_matrix_view_dma_word", + "encode_l_tile_cfg", + "validate_matrix_view_dominance", + "validate_disjoint_matrix_views", +] diff --git a/aten/plena/memory.py b/aten/plena/memory.py index 7dff198..266bbbd 100644 --- a/aten/plena/memory.py +++ b/aten/plena/memory.py @@ -396,10 +396,82 @@ def __init__(self, total_size: int | None = None, *, mlen: int = MLEN, tile_capa if total_size is None: total_size = self.tile_elems * tile_capacity super().__init__(total_size=total_size, alignment=self.tile_elems, mem_name="MRAM") + self._reserved: dict[str, MemoryBlock] = {} def allocate(self, name: str, size: int) -> int: return self._vmm.allocate(name, size) + @property + def reserved_blocks(self) -> tuple[MemoryBlock, ...]: + """Static allocations that remain occupied after a weight reset.""" + + return tuple(self._reserved.values()) + + @property + def capacity_after_reset(self) -> int: + """Contiguous elements the reset allocator can use for weight tiles. + + Reset preserves the original addresses and starts its bump pointer + after the last reservation; holes before that point are not reused. + """ + + end = max((block.addr + block.size for block in self._reserved.values()), default=0) + return self.total_size - end + + def free(self, name: str, strict: bool = True) -> MemoryBlock | None: + if name in self._reserved: + raise ValueError(f"MRAM reservation {name!r} cannot be freed as a transient weight") + return super().free(name, strict=strict) + + def reserve(self, name: str, size: int, *, required_tail: int = 0) -> int: + """Reserve compiler-managed Matrix scratch across allocator resets. + + This is an explicit static allocation, not a cache. Its address is + fixed in the generated program and it has no tag, lookup, replacement, + or runtime ownership state. ``required_tail`` checks weight capacity + after the resulting reset before changing any allocation. + """ + + if required_tail < 0: + raise ValueError("required_tail must be non-negative") + if name in self._reserved: + block = self._reserved[name] + if self._vmm._align(size) != block.size: + raise ValueError( + f"MRAM reservation {name!r} already has size {block.size}, " + f"not {self._vmm._align(size)}" + ) + if self.capacity_after_reset < required_tail: + raise ValueError("MRAM reservation leaves insufficient room for weight tiles") + return block.addr + if any(block.name == name for block in self._vmm.used_stack): + raise ValueError(f"MRAM reservation {name!r} conflicts with a live allocation") + aligned_size = self._vmm._align(size) + best = min( + ((block.size - aligned_size, i) for i, block in enumerate(self._vmm.free_stack) + if block.size >= aligned_size), + default=None, + ) + address = ( + self._vmm.free_stack[best[1]].addr + if best is not None else self._vmm._align(self.next_free) + ) + reserved_end = self.total_size - self.capacity_after_reset + if required_tail and self.total_size - max(reserved_end, address + aligned_size) < required_tail: + raise ValueError("MRAM reservation leaves insufficient room for weight tiles") + addr = self._vmm.allocate(name, size) + block = next(block for block in self._vmm.used_stack if block.name == name) + self._reserved[name] = block + return addr + + def reset(self): + """Release transient weight tiles while retaining static reservations.""" + + reserved = tuple(self._reserved.values()) + self._vmm.reset() + for block in sorted(reserved, key=lambda item: item.addr): + self._vmm.mark_used(block.addr, block.size, block.name) + class VRAMAllocator(MemoryAllocatorBase): """VRAM address allocator with MLEN-aligned best-fit reuse + bump allocation.""" diff --git a/aten/plena/program_lcompute.py b/aten/plena/program_lcompute.py new file mode 100644 index 0000000..643afef --- /dev/null +++ b/aten/plena/program_lcompute.py @@ -0,0 +1,1521 @@ +"""Emit prepared BF16 Mamba-2 and KDA recurrences through Matrix SRAM views. + +Inputs are post-projection/post-convolution fields described by the explicit HBM +manifest. This builder emits state and field DMA, L_TILE primitives, and output +stores; it does not compile the surrounding model. ``fixed`` uses one diagonal +base per chunk. ``affine`` retains the existing API spelling for the phased +layout, whose row coefficient remains fixed at one. +""" + +from __future__ import annotations + +from dataclasses import asdict, dataclass + +from enum import Enum + +from compiler.asm_templates._imm import load_large_int +from compiler.aten.plena.isa_matrix_view import ( + LTilePrimitive, + MatrixViewAllocation, + MatrixViewAxis, + MatrixViewDescriptor, + MatrixViewFlags, + MatrixViewMap, + MatrixViewShape, + validate_disjoint_matrix_views, +) + + +PAPER_MLEN = 2048 +PAPER_BANKS = 64 +PAPER_BANK_WIDTH = 32 +BF16_BYTES = 2 +ONE_MIB = 1024 * 1024 +STATE_DMA_VIEW = 3 +STATE_PRECISION_SELECTOR = 2 +GP_ADDRESS_SPACE_BYTES = 1 << 32 + + +class _StringEnum(str, Enum): + __str__ = str.__str__ + __format__ = str.__format__ + + +class RecurrenceLayout(_StringEnum): + FIXED = "fixed" + AFFINE = "affine" + + +class RecurrenceKind(_StringEnum): + """Algorithm-independent field contract selected by one recurrence.""" + + MAMBA = "mamba" + KDA = "kda" + + +@dataclass(frozen=True) +class LoweringRegisters: + """Five caller-owned GP registers used while emitting ``L_TILE``.""" + + destination: int = 9 + source: int = 10 + scale: int = 11 + shape: int = 12 + mapping: int = 13 + + def validate(self) -> None: + values = ( + self.destination, + self.source, + self.scale, + self.shape, + self.mapping, + ) + if len(set(values)) != len(values): + raise ValueError("L_TILE lowering registers must be distinct") + if any(register < 1 or register > 15 for register in values): + raise ValueError("L_TILE lowering registers must be GP1..GP15") + + +@dataclass(frozen=True) +class MatrixSramPoint: + """Physical Matrix-SRAM capacity used by one recurrence tile.""" + + mlen: int = PAPER_MLEN + banks: int = PAPER_BANKS + bank_width: int = PAPER_BANK_WIDTH + capacity_bytes: int = ONE_MIB + element_bytes: int = BF16_BYTES + + def validate(self) -> None: + if self.mlen != self.banks * self.bank_width: + raise ValueError("MLEN must equal banks * bank_width") + if not 1 <= self.banks <= 64 or self.banks & (self.banks - 1): + raise ValueError("Matrix-SRAM bank count must be a power of two no larger than 64") + row_bytes = self.mlen * self.element_bytes + if self.capacity_bytes <= 0 or self.capacity_bytes % row_bytes: + raise ValueError("Matrix-SRAM capacity must contain whole physical rows") + + @property + def depth_rows(self) -> int: + self.validate() + return self.capacity_bytes // (self.mlen * self.element_bytes) + + @property + def capacity_bank_words(self) -> int: + return self.depth_rows * self.banks + + +@dataclass(frozen=True) +class MatrixRecurrenceSpec: + name: str + kind: RecurrenceKind + heads: int + row_elements: int + recurrence_rows: int + primitives: tuple[LTilePrimitive, ...] + + @property + def state_bytes_per_head(self) -> int: + return self.recurrence_rows * self.row_elements * BF16_BYTES + + @property + def state_bytes_per_layer(self) -> int: + return self.heads * self.state_bytes_per_head + + def group_heads(self, point: MatrixSramPoint) -> int: + """Choose the largest group that fits capacity and one MLEN packet.""" + + capacity_heads = point.capacity_bytes // (2 * self.state_bytes_per_head) + packet_heads = point.mlen // self.row_elements + heads = min(capacity_heads, packet_heads) + if heads < 1: + raise ValueError(f"{self.name}: one state head does not fit Matrix SRAM") + return min(self.heads, heads) + + def validate(self, point: MatrixSramPoint) -> None: + point.validate() + if self.heads % self.group_heads(point): + raise ValueError( + f"{self.name}: {self.heads} heads do not divide into " + f"groups of {self.group_heads(point)}" + ) + if self.row_elements % point.bank_width: + raise ValueError("one recurrent row must contain whole bank words") + + +NEMOTRON_MAMBA = MatrixRecurrenceSpec( + name="nemotron3_mamba2", + kind=RecurrenceKind.MAMBA, + heads=64, + row_elements=64, + recurrence_rows=128, + primitives=( + LTilePrimitive.SCALE_ACCUM, + LTilePrimitive.SCALE_ACCUM, + LTilePrimitive.DOT_REDUCE, + LTilePrimitive.SCALE_ACCUM, + ), +) + +KIMI_KDA = MatrixRecurrenceSpec( + name="kimi_k3_kda", + kind=RecurrenceKind.KDA, + heads=96, + row_elements=128, + recurrence_rows=128, + primitives=( + LTilePrimitive.SCALE_ACCUM, + LTilePrimitive.DOT_REDUCE, + LTilePrimitive.SCALE_ACCUM, + LTilePrimitive.OUTER_UPDATE, + LTilePrimitive.DOT_REDUCE, + ), +) + + +@dataclass(frozen=True) +class RecurrenceWorkingSet: + spec: MatrixRecurrenceSpec + point: MatrixSramPoint + layout: RecurrenceLayout + group_heads: int + state_rows_per_chunk: int + allocations: tuple[MatrixViewAllocation, ...] + capacity_facts: dict[str, int] + + @property + def chunks(self) -> int: + return self.spec.recurrence_rows // self.state_rows_per_chunk + + @property + def groups(self) -> int: + return self.spec.heads // self.group_heads + + @property + def packet_values(self) -> int: + return self.group_heads * self.spec.row_elements + + def allocation(self, name: str) -> MatrixViewAllocation: + for allocation in self.allocations: + if allocation.name == name: + return allocation + raise KeyError(name) + + def to_dict(self) -> dict[str, object]: + return { + "layout": self.layout, + "group_heads": self.group_heads, + "groups": self.groups, + "state_rows_per_chunk": self.state_rows_per_chunk, + "chunks": self.chunks, + "packet_values": self.packet_values, + "vector_lane_utilization": self.packet_values / self.point.mlen, + "capacity_bytes": self.point.capacity_bytes, + "depth_rows": self.point.depth_rows, + "capacity_facts": self.capacity_facts, + "allocations": [ + { + "name": allocation.name, + "base": allocation.base, + "shape": asdict(allocation.descriptor.shape), + "mapping": { + "tile_pitch_rows": allocation.descriptor.mapping.tile_pitch_rows, + "tile_phase_stride": ( + allocation.descriptor.mapping.tile_phase_stride + ), + "flags": int(allocation.descriptor.mapping.flags), + }, + } + for allocation in self.allocations + ], + } + + +@dataclass(frozen=True) +class RecurrenceFieldPacket: + """One compiler-visible prepared recurrence operand in HBM. + + This is a functional interface for the standalone Matrix-SRAM recurrence + program. It is deliberately not a hidden descriptor fetch: the compiler + emits an ordinary viewed DMA for every packet. ``logical_values`` is the + tensor payload consumed by the view; ``transfer_values`` includes the VLEN + tail that the existing HBM interface reads for a short packet. + """ + + field: str + target: str + group: int + chunk: int | None + hbm_byte_offset: int + logical_values: int + transfer_values: int + + @property + def storage_bytes(self) -> int: + return self.transfer_values * BF16_BYTES + + @property + def key(self) -> tuple[str, int, int | None]: + return (self.field, self.group, self.chunk) + + def to_dict(self) -> dict[str, object]: + return asdict(self) | {"storage_bytes": self.storage_bytes} + + +@dataclass(frozen=True) +class RecurrenceFieldManifest: + """Static HBM placement for all non-state operands of one recurrence.""" + + base: int + end: int + packets: tuple[RecurrenceFieldPacket, ...] + + def packet( + self, + field: str, + *, + group: int, + chunk: int | None = None, + ) -> RecurrenceFieldPacket: + key = (field, group, chunk) + matches = [packet for packet in self.packets if packet.key == key] + if len(matches) != 1: + raise KeyError(f"expected one recurrence field packet {key}, found {len(matches)}") + return matches[0] + + def to_dict(self) -> dict[str, object]: + return { + "base": self.base, + "end": self.end, + "storage_bytes": self.end - self.base, + "packets": [packet.to_dict() for packet in self.packets], + } + + +def _addr(point: MatrixSramPoint, row: int, bank_phase: int) -> int: + if not 0 <= bank_phase < point.banks: + raise ValueError(f"bank phase {bank_phase} is outside the Matrix SRAM") + return row * point.mlen + bank_phase * point.bank_width + + +def _descriptor( + *, + rows: int, + cols: int, + tiles: int, + pitch: int, + affine: bool, + tile_phase_stride: int = 0, + broadcast: bool = False, +) -> MatrixViewDescriptor: + flags = MatrixViewFlags(0) + if broadcast: + flags |= MatrixViewFlags.BROADCAST_MINOR + return MatrixViewDescriptor( + MatrixViewShape(rows=rows, cols=cols, tile_count=tiles), + MatrixViewMap( + tile_pitch_rows=pitch, + tile_phase_stride=tile_phase_stride if affine else 0, + flags=flags, + ), + ) + + +def build_recurrence_working_set( + spec: MatrixRecurrenceSpec, + *, + layout: RecurrenceLayout | str, + point: MatrixSramPoint | None = None, +) -> RecurrenceWorkingSet: + """Build and statically prove one same-capacity recurrence allocation.""" + + point = point or MatrixSramPoint() + layout = RecurrenceLayout(layout) + spec.validate(point) + group_heads = spec.group_heads(point) + # Both controls start from the same maximal head group. Fixed wiring may + # need shorter state-row chunks, but silently shrinking only its head group + # would give the affine treatment an extra scheduling freedom and make the + # comparison unfair. + words_per_row = spec.row_elements // point.bank_width + # One scale packet carries either one scalar per head (DOT/OUTER) or one + # [a,b] pair per head (SCALE_ACCUM). It is read one cycle before the state + # packet and segment-broadcast across each logical state row. Replicating + # one scalar into an entire bank word per head would waste 16-32x capacity + # and model traffic that the architecture never needs. + scalar_cols = max( + point.bank_width, + ((2 * group_heads + point.bank_width - 1) // point.bank_width) + * point.bank_width, + ) + affine = layout is RecurrenceLayout.AFFINE + if affine: + state_rows = spec.recurrence_rows + state = _descriptor( + rows=state_rows, + cols=spec.row_elements, + tiles=group_heads, + pitch=0, + affine=True, + tile_phase_stride=words_per_row, + ) + scalar = _descriptor( + rows=state_rows, + cols=scalar_cols, + tiles=1, + pitch=state_rows, + affine=True, + tile_phase_stride=1, + broadcast=True, + ) + vector = _descriptor( + rows=1, + cols=spec.row_elements, + tiles=group_heads, + # A one-row temporary does not need state-style row sharing. The + # existing diagonal map reaches the packet floor when consecutive + # head tiles are spaced by one logical row width. Pitch 1 would + # overlap adjacent heads by all but one bank word (2--4x service). + pitch=words_per_row, + affine=True, + ) + occupied_state_banks = group_heads * words_per_row + # At the 1 MiB point the head-group state occupies half of every + # physical row, so the live fields use the remaining bank phases. At + # the 2 MiB point the larger group occupies all 64 banks; place fields + # in the second 128-row half instead of inventing a non-existent bank + # phase 64. Both are compiler-owned scratchpad allocations. + fields_base_row = state_rows if occupied_state_banks == point.banks else 0 + fields_first_bank = 0 if fields_base_row else occupied_state_banks + if spec.kind is RecurrenceKind.KDA: + head_pair = _descriptor( + rows=2, + cols=state_rows, + tiles=group_heads, + # Decay stores two logical rows (a/b) per head. Physical rows + # remain eight rows apart so all 2x4 words are disjoint. A + # four-bank tile phase makes both the complete viewed DMA + # (128 words, two per bank) and each 32-word column packet + # conflict-free for the BF16 16-head group. + pitch=2 * words_per_row, + affine=True, + tile_phase_stride=words_per_row, + broadcast=True, + ) + head_scalar = _descriptor( + rows=1, + cols=state_rows, + tiles=group_heads, + # One 128-value row occupies four bank words. The matching + # four-row tile pitch makes adjacent heads occupy disjoint + # four-bank slices in a full viewed DMA. + pitch=words_per_row, + affine=True, + broadcast=True, + ) + phases = { + "state": 0, + # Decay occupies both logical scalar rows and therefore uses + # the opposite half-bank phase at both supported capacities. + # Combined with head_pair's affine coefficients this reaches + # the theoretical one-port service floor without aliasing the + # resident state or the other live fields. + "decay": point.banks // 2, + "key": fields_first_bank + words_per_row, + "query": fields_first_bank + 2 * words_per_row, + "beta": fields_first_bank + 3 * words_per_row, + "error": fields_first_bank + 3 * words_per_row + 1, + "prediction": fields_first_bank + 4 * words_per_row + 1, + "output": fields_first_bank + 5 * words_per_row + 1, + } + base_rows = dict.fromkeys(phases, fields_base_row) + base_rows["state"] = 0 + descriptors = { + "state": state, + "decay": head_pair, + "key": head_scalar, + "query": head_scalar, + "beta": _shape_rows(scalar, 1), + "error": vector, + "prediction": vector, + "output": vector, + } + else: + scalar_words = scalar_cols // point.bank_width + phases = { + "state": 0, + "dt": fields_first_bank, + "update": fields_first_bank + scalar_words, + "c": fields_first_bank + 2 * scalar_words, + "d": fields_first_bank + 3 * scalar_words, + "x": fields_first_bank + 4 * scalar_words, + "scratch": fields_first_bank + 4 * scalar_words + words_per_row, + "output": fields_first_bank + 4 * scalar_words + 2 * words_per_row, + } + base_rows = dict.fromkeys(phases, fields_base_row) + base_rows["state"] = 0 + descriptors = { + "state": state, + "dt": _shape_rows(scalar, 1), + "update": scalar, + "c": scalar, + "d": _shape_rows(scalar, 1), + "x": vector, + "scratch": vector, + "output": vector, + } + else: + # One fixed base and descriptor cover each chunk. Keep its row count + # within the bank-word width so the head tiles remain disjoint. + state_rows = words_per_row + if state_rows < 1 or spec.recurrence_rows % state_rows: + raise ValueError(f"{spec.name}: fixed state chunks do not divide recurrence rows") + state = _descriptor( + rows=state_rows, + cols=spec.row_elements, + tiles=group_heads, + pitch=state_rows, + affine=False, + ) + scalar = _descriptor( + rows=state_rows, + cols=scalar_cols, + tiles=1, + pitch=state_rows, + affine=False, + broadcast=True, + ) + vector = _descriptor( + rows=1, + cols=spec.row_elements, + tiles=group_heads, + pitch=state_rows, + affine=False, + ) + if spec.kind is RecurrenceKind.KDA: + padded_key_cols = max(point.bank_width, state_rows) + head_pair = _descriptor( + rows=2, + cols=padded_key_cols, + tiles=group_heads, + pitch=2, + affine=False, + broadcast=True, + ) + head_scalar = _descriptor( + rows=1, + cols=padded_key_cols, + tiles=group_heads, + pitch=1, + affine=False, + broadcast=True, + ) + phases = { + "state": 0, + "decay": words_per_row, + "key": 2 * words_per_row, + "query": 3 * words_per_row, + "beta": 4 * words_per_row, + "error": 4 * words_per_row + 1, + "prediction": 5 * words_per_row + 1, + "output": 6 * words_per_row + 1, + } + descriptors = { + "state": state, + "decay": head_pair, + "key": head_scalar, + "query": head_scalar, + "beta": _shape_rows(scalar, 1), + "error": vector, + "prediction": vector, + "output": vector, + } + else: + scalar_words = scalar_cols // point.bank_width + phases = { + "state": 0, + "dt": words_per_row, + "update": words_per_row + scalar_words, + "c": words_per_row + 2 * scalar_words, + "d": words_per_row + 3 * scalar_words, + "x": words_per_row + 4 * scalar_words, + "scratch": 2 * words_per_row + 4 * scalar_words, + "output": 3 * words_per_row + 4 * scalar_words, + } + descriptors = { + "state": state, + "dt": _shape_rows(scalar, 1), + "update": scalar, + "c": scalar, + "d": _shape_rows(scalar, 1), + "x": vector, + "scratch": vector, + "output": vector, + } + base_rows = dict.fromkeys(descriptors, 0) + + allocations = tuple( + MatrixViewAllocation( + name, + _addr(point, base_rows[name], phases[name]), + descriptor, + ) + for name, descriptor in descriptors.items() + ) + facts = validate_disjoint_matrix_views( + allocations, + mlen=point.mlen, + banks=point.banks, + bank_width=point.bank_width, + depth_rows=point.depth_rows, + ) + return RecurrenceWorkingSet( + spec=spec, + point=point, + layout=layout, + group_heads=group_heads, + state_rows_per_chunk=state_rows, + allocations=allocations, + capacity_facts=facts, + ) + + +def _shape_rows(descriptor: MatrixViewDescriptor, rows: int) -> MatrixViewDescriptor: + return MatrixViewDescriptor( + MatrixViewShape( + rows=rows, + cols=descriptor.shape.cols, + tile_count=descriptor.shape.tile_count, + ), + descriptor.mapping, + ) + + +def _round_up(value: int, multiple: int) -> int: + if multiple <= 0: + raise ValueError("rounding multiple must be positive") + return ((value + multiple - 1) // multiple) * multiple + + +def _validate_hbm_offset_range(base: int, end: int, name: str) -> None: + """Keep a DMA arena inside its unsigned 32-bit GP offset window. + + The selected architectural aN register can supply a higher HBM base; the + addresses loaded into GP registers by this lowering are offsets from it. + An exclusive end equal to 2**32 is valid: its last byte is 2**32 - 1. + """ + if base < 0 or base >= GP_ADDRESS_SPACE_BYTES or end > GP_ADDRESS_SPACE_BYTES: + raise ValueError(f"{name} must fit the unsigned 32-bit GP DMA offset window") + + +def build_recurrence_field_manifest( + working_set: RecurrenceWorkingSet, + *, + field_hbm_base: int, +) -> RecurrenceFieldManifest: + """Place every prepared non-state operand in a disjoint BF16 HBM packet. + + The standalone recurrence consumes post-projection/post-convolution values, + so decay/exponential and compact scale pairs are already prepared. The + integrated model lowering may replace these DMAs with direct Matrix/Vector + producer writeback, but it must preserve this exact logical field contract. + """ + + if field_hbm_base < 0 or field_hbm_base % 64: + raise ValueError("field_hbm_base must be a non-negative 64-byte address") + _validate_hbm_offset_range(field_hbm_base, field_hbm_base, "prepared fields") + packets: list[RecurrenceFieldPacket] = [] + cursor = field_hbm_base + + def add(field: str, target: str, group: int, chunk: int | None = None) -> None: + nonlocal cursor + descriptor = working_set.allocation(target).descriptor + logical_values = ( + descriptor.shape.rows + * descriptor.shape.cols + * descriptor.shape.tile_count + ) + transfer_values = _round_up(logical_values, working_set.point.mlen) + packets.append( + RecurrenceFieldPacket( + field=field, + target=target, + group=group, + chunk=chunk, + hbm_byte_offset=cursor, + logical_values=logical_values, + transfer_values=transfer_values, + ) + ) + cursor += transfer_values * BF16_BYTES + + for group in range(working_set.groups): + if working_set.spec.kind is RecurrenceKind.MAMBA: + for field, target in ( + ("x", "x"), + ("scratch_zero", "scratch"), + ("dt", "dt"), + ("output_zero", "output"), + ("d", "d"), + ): + add(field, target, group) + for chunk in range(working_set.chunks): + add("update", "update", group, chunk) + add("c", "c", group, chunk) + add("output_result", "output", group) + elif working_set.spec.kind is RecurrenceKind.KDA: + for field, target in ( + ("prediction_zero", "prediction"), + ("value", "error"), + ("beta", "beta"), + ("output_zero", "output"), + ): + add(field, target, group) + for chunk in range(working_set.chunks): + add("decay", "decay", group, chunk) + add("key", "key", group, chunk) + add("query", "query", group, chunk) + add("output_result", "output", group) + else: # pragma: no cover - all public specs are enumerated above + raise ValueError(f"no recurrence-field ABI for {working_set.spec.name}") + + _validate_hbm_offset_range(field_hbm_base, cursor, "prepared fields") + return RecurrenceFieldManifest( + base=field_hbm_base, + end=cursor, + packets=tuple(packets), + ) + + +def _configure_view( + *, + slot: int, + descriptor: MatrixViewDescriptor, + shape_register: int, + map_register: int, +) -> list[str]: + return [ + *load_large_int(shape_register, descriptor.shape.pack()), + *load_large_int(map_register, descriptor.mapping.pack()), + f"L_TILE_CFG {slot}, gp{shape_register}, gp{map_register}", + ] + + +def _emit_exec( + lines: list[str], + *, + dst: MatrixViewAllocation, + src: MatrixViewAllocation, + scales: MatrixViewAllocation, + primitive: LTilePrimitive, + rows: int | None = None, + source_axis: MatrixViewAxis = MatrixViewAxis.ROW, + scale_axis: MatrixViewAxis = MatrixViewAxis.ROW, + label: str, + registers: LoweringRegisters, +) -> None: + dst_descriptor = dst.descriptor if rows is None else _shape_rows(dst.descriptor, rows) + src_rows = rows if source_axis is MatrixViewAxis.ROW else None + if source_axis is MatrixViewAxis.ROW and src.descriptor.shape.rows == 1: + src_rows = 1 + src_descriptor = ( + src.descriptor if src_rows is None else _shape_rows(src.descriptor, src_rows) + ) + scale_rows = rows if scale_axis is MatrixViewAxis.ROW else None + scale_descriptor = ( + scales.descriptor + if scale_rows is None + else _shape_rows(scales.descriptor, scale_rows) + ) + lines.append(f"; @l_tile_step={label}") + lines.extend( + _configure_view( + slot=0, + descriptor=dst_descriptor, + shape_register=registers.shape, + map_register=registers.mapping, + ) + ) + lines.extend( + _configure_view( + slot=1, + descriptor=src_descriptor, + shape_register=registers.shape, + map_register=registers.mapping, + ) + ) + lines.extend( + _configure_view( + slot=2, + descriptor=scale_descriptor, + shape_register=registers.shape, + map_register=registers.mapping, + ) + ) + lines.extend(load_large_int(registers.destination, dst.base)) + lines.extend(load_large_int(registers.source, src.base)) + lines.extend(load_large_int(registers.scale, scales.base)) + axis_mask = int(source_axis) | int(scale_axis) << 1 + suffix = "" if axis_mask == 0 else f", {axis_mask}" + lines.append( + f"L_TILE_EXEC gp{registers.destination}, gp{registers.source}, " + f"gp{registers.scale}, {int(primitive)}{suffix}" + ) + + +def _emit_state_transfer( + lines: list[str], + *, + working_set: RecurrenceWorkingSet, + state: MatrixViewAllocation, + direction: str, + group: int, + chunk: int, + rows: int, + values: int, + state_hbm_base: int, + hbm_address_register: int, + registers: LoweringRegisters, + snapshot_hbm_base: int | None = None, +) -> None: + if direction not in {"load", "reload_intermediate", "store", "store_intermediate"}: + raise ValueError(f"unsupported Matrix state transfer direction {direction!r}") + if not 0 <= hbm_address_register <= 15: + raise ValueError("HBM address register must be a0..a15") + expected_values = working_set.group_heads * rows * working_set.spec.row_elements + if values != expected_values: + raise ValueError( + f"state packet has {values} values, expected {expected_values}" + ) + + # Persistent state is laid out packet-major in HBM: + # [head-group][state-row chunk][head-in-group][row][lane]. This makes one + # explicit DMA packet contiguous without a hidden gather engine. Fixed and + # affine variants move the same number of BF16 values; only their on-chip + # bank placement differs. + packet = group * working_set.chunks + chunk + hbm_byte_offset = state_hbm_base + packet * values * BF16_BYTES + descriptor = _shape_rows(state.descriptor, rows) + lines.append( + f"; @matrix_state_{direction} group={group} chunk={chunk} " + f"rows={rows} values={values} precision=bf16 hbm_byte_offset={hbm_byte_offset}" + ) + lines.extend( + _configure_view( + slot=STATE_DMA_VIEW, + descriptor=descriptor, + shape_register=registers.shape, + map_register=registers.mapping, + ) + ) + lines.extend(load_large_int(registers.destination, state.base)) + lines.extend(load_large_int(registers.source, hbm_byte_offset)) + opcode = ( + "H_PREFETCH_V.MV" + if direction in {"load", "reload_intermediate"} + else "H_STORE_V.MV" + ) + lines.append( + f"{opcode} gp{registers.destination}, gp{registers.source}, " + f"a{hbm_address_register}, 0, {STATE_PRECISION_SELECTOR}, {STATE_DMA_VIEW}" + ) + + if direction == "store" and snapshot_hbm_base is not None: + snapshot_offset = snapshot_hbm_base + packet * values * BF16_BYTES + lines.append(f"; @state_snapshot group={group} chunk={chunk} hbm_byte_offset={snapshot_offset}") + lines.extend(load_large_int(registers.source, snapshot_offset)) + lines.append( + f"H_STORE_V.MV gp{registers.destination}, gp{registers.source}, " + f"a{hbm_address_register}, 0, {STATE_PRECISION_SELECTOR}, {STATE_DMA_VIEW}" + ) + + +def _emit_field_load( + lines: list[str], + *, + working_set: RecurrenceWorkingSet, + manifest: RecurrenceFieldManifest, + field: str, + group: int, + chunk: int | None, + hbm_address_register: int, + registers: LoweringRegisters, +) -> None: + """Emit a real viewed DMA for one prepared recurrence operand.""" + + packet = manifest.packet(field, group=group, chunk=chunk) + target = working_set.allocation(packet.target) + descriptor = target.descriptor + logical_values = ( + descriptor.shape.rows + * descriptor.shape.cols + * descriptor.shape.tile_count + ) + if logical_values != packet.logical_values: + raise AssertionError( + f"field {packet.key} changed shape after its HBM manifest was built" + ) + chunk_text = "none" if chunk is None else str(chunk) + lines.append( + f"; @matrix_field_load field={field} target={packet.target} " + f"group={group} chunk={chunk_text} logical_values={packet.logical_values} " + f"transfer_values={packet.transfer_values} precision=bf16 " + f"hbm_byte_offset={packet.hbm_byte_offset}" + ) + lines.extend( + _configure_view( + slot=STATE_DMA_VIEW, + descriptor=descriptor, + shape_register=registers.shape, + map_register=registers.mapping, + ) + ) + lines.extend(load_large_int(registers.destination, target.base)) + lines.extend(load_large_int(registers.source, packet.hbm_byte_offset)) + lines.append( + f"H_PREFETCH_V.MV gp{registers.destination}, gp{registers.source}, " + f"a{hbm_address_register}, 0, {STATE_PRECISION_SELECTOR}, {STATE_DMA_VIEW}" + ) + + +def _emit_field_store( + lines: list[str], + *, + working_set: RecurrenceWorkingSet, + manifest: RecurrenceFieldManifest, + field: str, + group: int, + chunk: int | None, + hbm_address_register: int, + registers: LoweringRegisters, +) -> None: + """Store one produced Matrix view before its scratch allocation is reused.""" + + packet = manifest.packet(field, group=group, chunk=chunk) + source = working_set.allocation(packet.target) + descriptor = source.descriptor + chunk_text = "none" if chunk is None else str(chunk) + lines.append( + f"; @matrix_field_store field={field} source={packet.target} " + f"group={group} chunk={chunk_text} logical_values={packet.logical_values} " + f"transfer_values={packet.transfer_values} precision=bf16 " + f"hbm_byte_offset={packet.hbm_byte_offset}" + ) + lines.extend( + _configure_view( + slot=STATE_DMA_VIEW, + descriptor=descriptor, + shape_register=registers.shape, + map_register=registers.mapping, + ) + ) + lines.extend(load_large_int(registers.destination, source.base)) + lines.extend(load_large_int(registers.source, packet.hbm_byte_offset)) + lines.append( + f"H_STORE_V.MV gp{registers.destination}, gp{registers.source}, " + f"a{hbm_address_register}, 0, {STATE_PRECISION_SELECTOR}, {STATE_DMA_VIEW}" + ) + + +def _lower_mamba( + working_set: RecurrenceWorkingSet, + field_manifest: RecurrenceFieldManifest, + registers: LoweringRegisters, + *, + state_hbm_base: int, + hbm_address_register: int, + snapshot_hbm_base: int | None = None, +) -> list[str]: + lines: list[str] = [] + state = working_set.allocation("state") + dt = working_set.allocation("dt") + update = working_set.allocation("update") + c = working_set.allocation("c") + d = working_set.allocation("d") + x = working_set.allocation("x") + scratch = working_set.allocation("scratch") + output = working_set.allocation("output") + + for group in range(working_set.groups): + lines.append(f"; @head_group={group}/{working_set.groups}") + for field in ("x", "scratch_zero", "dt", "output_zero"): + _emit_field_load( + lines, + working_set=working_set, + manifest=field_manifest, + field=field, + group=group, + chunk=None, + hbm_address_register=hbm_address_register, + registers=registers, + ) + _emit_exec( + lines, + dst=scratch, + src=x, + scales=dt, + primitive=LTilePrimitive.SCALE_ACCUM, + rows=1, + label="mamba_dt_times_x", + registers=registers, + ) + for chunk in range(working_set.chunks): + values = ( + working_set.group_heads + * working_set.state_rows_per_chunk + * working_set.spec.row_elements + ) + _emit_state_transfer( + lines, + working_set=working_set, + state=state, + direction="load", + group=group, + chunk=chunk, + rows=working_set.state_rows_per_chunk, + values=values, + state_hbm_base=state_hbm_base, + snapshot_hbm_base=snapshot_hbm_base, + hbm_address_register=hbm_address_register, + registers=registers, + ) + for field in ("update", "c"): + _emit_field_load( + lines, + working_set=working_set, + manifest=field_manifest, + field=field, + group=group, + chunk=chunk, + hbm_address_register=hbm_address_register, + registers=registers, + ) + _emit_exec( + lines, + dst=state, + src=scratch, + scales=update, + primitive=LTilePrimitive.SCALE_ACCUM, + label="mamba_state_decay_rank1_update", + registers=registers, + ) + _emit_exec( + lines, + dst=output, + src=state, + scales=c, + primitive=LTilePrimitive.DOT_REDUCE, + label="mamba_c_readout", + registers=registers, + ) + _emit_state_transfer( + lines, + working_set=working_set, + state=state, + direction="store", + group=group, + chunk=chunk, + rows=working_set.state_rows_per_chunk, + values=values, + state_hbm_base=state_hbm_base, + snapshot_hbm_base=snapshot_hbm_base, + hbm_address_register=hbm_address_register, + registers=registers, + ) + _emit_field_load( + lines, + working_set=working_set, + manifest=field_manifest, + field="d", + group=group, + chunk=None, + hbm_address_register=hbm_address_register, + registers=registers, + ) + _emit_exec( + lines, + dst=output, + src=x, + scales=d, + primitive=LTilePrimitive.SCALE_ACCUM, + rows=1, + label="mamba_skip", + registers=registers, + ) + _emit_field_store( + lines, + working_set=working_set, + manifest=field_manifest, + field="output_result", + group=group, + chunk=None, + hbm_address_register=hbm_address_register, + registers=registers, + ) + return lines + + +def _lower_kda( + working_set: RecurrenceWorkingSet, + field_manifest: RecurrenceFieldManifest, + registers: LoweringRegisters, + *, + state_hbm_base: int, + hbm_address_register: int, + snapshot_hbm_base: int | None = None, +) -> list[str]: + lines: list[str] = [] + state = working_set.allocation("state") + decay = working_set.allocation("decay") + key = working_set.allocation("key") + query = working_set.allocation("query") + beta = working_set.allocation("beta") + error = working_set.allocation("error") + prediction = working_set.allocation("prediction") + output = working_set.allocation("output") + + for group in range(working_set.groups): + lines.append(f"; @head_group={group}/{working_set.groups}") + _emit_field_load( + lines, + working_set=working_set, + manifest=field_manifest, + field="prediction_zero", + group=group, + chunk=None, + hbm_address_register=hbm_address_register, + registers=registers, + ) + # Fixed wiring cannot hold a complete head-group state and its fields; + # retain decayed chunks explicitly between the predict and update passes. + for chunk in range(working_set.chunks): + values = ( + working_set.group_heads + * working_set.state_rows_per_chunk + * working_set.spec.row_elements + ) + _emit_state_transfer( + lines, + working_set=working_set, + state=state, + direction="load", + group=group, + chunk=chunk, + rows=working_set.state_rows_per_chunk, + values=values, + state_hbm_base=state_hbm_base, + snapshot_hbm_base=snapshot_hbm_base, + hbm_address_register=hbm_address_register, + registers=registers, + ) + for field in ("decay", "key"): + _emit_field_load( + lines, + working_set=working_set, + manifest=field_manifest, + field=field, + group=group, + chunk=chunk, + hbm_address_register=hbm_address_register, + registers=registers, + ) + _emit_exec( + lines, + dst=state, + src=prediction, + scales=decay, + primitive=LTilePrimitive.SCALE_ACCUM, + scale_axis=MatrixViewAxis.COLUMN, + label="kda_decay", + registers=registers, + ) + _emit_exec( + lines, + dst=prediction, + src=state, + scales=key, + primitive=LTilePrimitive.DOT_REDUCE, + scale_axis=MatrixViewAxis.COLUMN, + label="kda_prediction", + registers=registers, + ) + if working_set.chunks > 1: + _emit_state_transfer( + lines, + working_set=working_set, + state=state, + direction="store_intermediate", + group=group, + chunk=chunk, + rows=working_set.state_rows_per_chunk, + values=values, + state_hbm_base=state_hbm_base, + snapshot_hbm_base=snapshot_hbm_base, + hbm_address_register=hbm_address_register, + registers=registers, + ) + for field in ("value", "beta", "output_zero"): + _emit_field_load( + lines, + working_set=working_set, + manifest=field_manifest, + field=field, + group=group, + chunk=None, + hbm_address_register=hbm_address_register, + registers=registers, + ) + _emit_exec( + lines, + dst=error, + src=prediction, + scales=beta, + primitive=LTilePrimitive.SCALE_ACCUM, + rows=1, + label="kda_beta_error", + registers=registers, + ) + for chunk in range(working_set.chunks): + values = ( + working_set.group_heads + * working_set.state_rows_per_chunk + * working_set.spec.row_elements + ) + if working_set.chunks > 1: + _emit_state_transfer( + lines, + working_set=working_set, + state=state, + direction="reload_intermediate", + group=group, + chunk=chunk, + rows=working_set.state_rows_per_chunk, + values=values, + state_hbm_base=state_hbm_base, + snapshot_hbm_base=snapshot_hbm_base, + hbm_address_register=hbm_address_register, + registers=registers, + ) + if working_set.chunks > 1: + _emit_field_load( + lines, + working_set=working_set, + manifest=field_manifest, + field="key", + group=group, + chunk=chunk, + hbm_address_register=hbm_address_register, + registers=registers, + ) + _emit_field_load( + lines, + working_set=working_set, + manifest=field_manifest, + field="query", + group=group, + chunk=chunk, + hbm_address_register=hbm_address_register, + registers=registers, + ) + _emit_exec( + lines, + dst=state, + src=error, + scales=key, + primitive=LTilePrimitive.OUTER_UPDATE, + scale_axis=MatrixViewAxis.COLUMN, + label="kda_rank1_update", + registers=registers, + ) + _emit_exec( + lines, + dst=output, + src=state, + scales=query, + primitive=LTilePrimitive.DOT_REDUCE, + scale_axis=MatrixViewAxis.COLUMN, + label="kda_readout", + registers=registers, + ) + _emit_state_transfer( + lines, + working_set=working_set, + state=state, + direction="store", + group=group, + chunk=chunk, + rows=working_set.state_rows_per_chunk, + values=values, + state_hbm_base=state_hbm_base, + snapshot_hbm_base=snapshot_hbm_base, + hbm_address_register=hbm_address_register, + registers=registers, + ) + _emit_field_store( + lines, + working_set=working_set, + manifest=field_manifest, + field="output_result", + group=group, + chunk=None, + hbm_address_register=hbm_address_register, + registers=registers, + ) + return lines + + +def lower_matrix_recurrence( + spec: MatrixRecurrenceSpec, + *, + co_layout: bool | None = None, + layout: RecurrenceLayout | str | None = None, + point: MatrixSramPoint | None = None, + mlen: int | None = None, + blen: int | None = None, + registers: LoweringRegisters | None = None, + state_hbm_base: int = 0, + field_hbm_base: int | None = None, + hbm_address_register: int = 0, + snapshot_hbm_base: int | None = None, +) -> str: + """Emit one official-shape layer's complete recurrent Matrix-SRAM path. + + ``co_layout`` is retained as a compatibility spelling for existing callers; + new code should pass ``layout``. ``mlen``/``blen`` may only restate the + physical point and cannot silently change its capacity. + """ + + if layout is None: + layout = RecurrenceLayout.AFFINE if co_layout else RecurrenceLayout.FIXED + elif co_layout is not None and (RecurrenceLayout(layout) is RecurrenceLayout.AFFINE) != co_layout: + raise ValueError("co_layout and layout select different mechanisms") + point = point or MatrixSramPoint() + if mlen is not None and mlen != point.mlen: + raise ValueError("mlen must match the explicit Matrix-SRAM point") + if blen is not None and blen != point.bank_width: + raise ValueError("blen must match the explicit Matrix-SRAM bank width") + registers = registers or LoweringRegisters() + registers.validate() + if state_hbm_base < 0: + raise ValueError("state_hbm_base must be a non-negative byte address") + if state_hbm_base % 64: + raise ValueError("state_hbm_base must be 64-byte aligned") + state_hbm_end = state_hbm_base + spec.state_bytes_per_layer + _validate_hbm_offset_range(state_hbm_base, state_hbm_end, "state") + if not 0 <= hbm_address_register <= 15: + raise ValueError("hbm_address_register must be in [0, 15]") + working_set = build_recurrence_working_set(spec, layout=layout, point=point) + if field_hbm_base is None: + field_hbm_base = _round_up( + state_hbm_base + spec.state_bytes_per_layer, + 64, + ) + field_manifest = build_recurrence_field_manifest( + working_set, + field_hbm_base=field_hbm_base, + ) + if state_hbm_base < field_manifest.end and field_manifest.base < state_hbm_end: + raise ValueError("prepared fields overlap live recurrent state") + if snapshot_hbm_base is not None: + if snapshot_hbm_base < 0 or snapshot_hbm_base % 64: + raise ValueError("snapshot_hbm_base must be a non-negative aligned byte address") + snapshot_end = snapshot_hbm_base + spec.state_bytes_per_layer + _validate_hbm_offset_range(snapshot_hbm_base, snapshot_end, "snapshot") + for start, end in ((state_hbm_base, state_hbm_end), + (field_manifest.base, field_manifest.end)): + if snapshot_hbm_base < end and start < snapshot_end: + raise ValueError("snapshot overlaps live state or prepared fields") + lines = [ + f"; @stage={spec.name}_matrix_recurrence", + f"; @layout={working_set.layout}", + f"; @matrix_sram_bytes={point.capacity_bytes}", + f"; @group_heads={working_set.group_heads}", + f"; @state_rows_per_chunk={working_set.state_rows_per_chunk}", + "; @state_precision=bf16", + "; @state_storage=compiler_managed_matrix_sram_no_cache", + "; @field_contract=prepared_bf16_post_projection_post_conv", + f"; @field_hbm_base={field_manifest.base}", + f"; @field_hbm_end={field_manifest.end}", + ] + lines.extend( + _lower_mamba( + working_set, + field_manifest, + registers, + state_hbm_base=state_hbm_base, + snapshot_hbm_base=snapshot_hbm_base, + hbm_address_register=hbm_address_register, + ) + if spec.kind is RecurrenceKind.MAMBA + else _lower_kda( + working_set, + field_manifest, + registers, + state_hbm_base=state_hbm_base, + snapshot_hbm_base=snapshot_hbm_base, + hbm_address_register=hbm_address_register, + ) + ) + return "\n".join(lines) + "\n" + + + + +def validate_recurrence_field_loads( + assembly: str, + *, + expected: set[tuple[str, int, int | None]] | None = None, +) -> set[tuple[str, int, int | None]]: + """Require every field marker to own one executable viewed prefetch. + + This intentionally rejects the old ``@matrix_field_write`` comments. A + comment is documentation, not evidence that a producer populated Matrix + SRAM. + """ + + if "@matrix_field_write=" in assembly: + raise ValueError("legacy comment-only Matrix field write is not executable") + lines = assembly.splitlines() + observed: set[tuple[str, int, int | None]] = set() + for index, raw in enumerate(lines): + line = raw.strip() + if not line.startswith("; @matrix_field_load"): + continue + tokens = { + token.split("=", 1)[0]: token.split("=", 1)[1] + for token in line.split() + if "=" in token + } + required = { + "field", + "target", + "group", + "chunk", + "logical_values", + "transfer_values", + "precision", + "hbm_byte_offset", + } + missing = required - tokens.keys() + if missing: + raise ValueError(f"malformed Matrix field marker missing {sorted(missing)}") + chunk = None if tokens["chunk"] == "none" else int(tokens["chunk"]) + key = (tokens["field"], int(tokens["group"]), chunk) + if key in observed and tokens["field"] not in {"key"}: + raise ValueError(f"duplicate Matrix field load marker {key}") + observed.add(key) + end = next( + ( + cursor + for cursor in range(index + 1, len(lines)) + if lines[cursor].strip().startswith("; @matrix_field_load") + or lines[cursor].strip().startswith("; @matrix_field_store") + or lines[cursor].strip().startswith("; @matrix_state_") + or lines[cursor].strip().startswith("; @l_tile_step=") + ), + len(lines), + ) + block = "\n".join(lines[index:end]) + if "L_TILE_CFG 3" not in block or "H_PREFETCH_V.MV" not in block: + raise ValueError(f"Matrix field marker {key} has no executable viewed DMA") + if tokens["precision"] != "bf16" or ", 0, 2, 3" not in block: + raise ValueError(f"Matrix field marker {key} does not use BF16 state traffic") + logical = int(tokens["logical_values"]) + transferred = int(tokens["transfer_values"]) + if transferred < logical or transferred % PAPER_MLEN: + raise ValueError(f"Matrix field marker {key} has an invalid DMA extent") + if expected is not None and observed != expected: + raise ValueError( + f"Matrix field coverage differs: missing={sorted(expected - observed)!r}, " + f"extra={sorted(observed - expected)!r}" + ) + return observed + + +def validate_recurrence_output_stores( + assembly: str, + *, + expected_groups: int, +) -> dict[int, int]: + """Prove that every head-group output leaves Matrix SRAM exactly once.""" + + lines = assembly.splitlines() + stores: dict[int, int] = {} + store_indices: dict[int, int] = {} + current_group: int | None = None + last_exec_by_group: dict[int, int] = {} + store_offsets: set[int] = set() + for index, raw in enumerate(lines): + line = raw.strip() + if line.startswith("; @head_group="): + current_group = int(line.split("=", 1)[1].split("/", 1)[0]) + elif line.startswith("L_TILE_EXEC"): + if current_group is None: + raise ValueError("L_TILE_EXEC appears before a head-group marker") + last_exec_by_group[current_group] = index + elif line.startswith("; @matrix_field_store"): + tokens = { + token.split("=", 1)[0]: token.split("=", 1)[1] + for token in line.split() + if "=" in token + } + if tokens.get("field") != "output_result": + raise ValueError(f"unexpected recurrence output field {tokens.get('field')}") + group = int(tokens["group"]) + if group in stores: + raise ValueError(f"head group {group} stores its output more than once") + if group != current_group or group not in last_exec_by_group: + raise ValueError(f"head group {group} output is stored before any EXEC") + offset = int(tokens["hbm_byte_offset"]) + if offset in store_offsets: + raise ValueError(f"head groups alias output HBM offset {offset}") + store_offsets.add(offset) + end = next( + ( + cursor + for cursor in range(index + 1, len(lines)) + if lines[cursor].strip().startswith("; @head_group=") + or lines[cursor].strip().startswith("; @matrix_field_") + or lines[cursor].strip().startswith("; @matrix_state_") + or lines[cursor].strip().startswith("; @l_tile_step=") + ), + len(lines), + ) + block = "\n".join(lines[index:end]) + if "L_TILE_CFG 3" not in block or "H_STORE_V.MV" not in block: + raise ValueError(f"head group {group} output marker has no viewed store") + stores[group] = offset + store_indices[group] = index + expected = set(range(expected_groups)) + if set(stores) != expected: + raise ValueError( + "recurrence output coverage differs: " + f"missing={sorted(expected - set(stores))}, " + f"extra={sorted(set(stores) - expected)}" + ) + for group, store_index in store_indices.items(): + if store_index <= last_exec_by_group[group]: + raise ValueError(f"head group {group} output is not stored after its final EXEC") + return stores + + +__all__ = [ + "BF16_BYTES", + "KIMI_KDA", + "LoweringRegisters", + "NEMOTRON_MAMBA", + "ONE_MIB", + "MatrixRecurrenceSpec", + "MatrixSramPoint", + "RecurrenceLayout", + "RecurrenceKind", + "RecurrenceFieldManifest", + "RecurrenceFieldPacket", + "RecurrenceWorkingSet", + "build_recurrence_field_manifest", + "build_recurrence_working_set", + "lower_matrix_recurrence", + "validate_recurrence_field_loads", + "validate_recurrence_output_stores", +] diff --git a/aten/plena/program_matrix_ops.py b/aten/plena/program_matrix_ops.py index 5a5f3bc..046fe96 100644 --- a/aten/plena/program_matrix_ops.py +++ b/aten/plena/program_matrix_ops.py @@ -4,6 +4,8 @@ import math +from compiler.asm_templates._imm import load_large_int +from compiler.aten.plena.isa_matrix_view import LTilePrimitive, MatrixViewDescriptor from compiler.aten.plena.vars import InputVar, TensorVar, VRAMMatrixVar @@ -31,6 +33,194 @@ def _matrix_precision_code(matrix_precision: str | int) -> int: class ProgramMatrixOpsMixin: + def reserve_matrix_view_scratch_v0( + self, + name: str = "__matrix_view_scratch", + *, + descriptor: MatrixViewDescriptor | None = None, + ) -> int: + """Reserve one existing Matrix-SRAM tile for explicit view traffic.""" + + if descriptor is not None: + self._validate_matrix_view_scratch_descriptor(descriptor) + return self.mram_allocator.reserve(name, self.mlen * self.mlen) + + def _validate_matrix_view_descriptor( + self, descriptor: MatrixViewDescriptor + ) -> None: + if not isinstance(descriptor, MatrixViewDescriptor): + raise TypeError( + "descriptor must be MatrixViewDescriptor, got " + f"{type(descriptor).__name__}" + ) + if self.mlen % self.blen: + raise ValueError("MLEN must contain a whole number of Matrix bank words") + descriptor.validate_for_machine( + banks=self.mlen // self.blen, + bank_width=self.blen, + ) + final_physical_row = self._matrix_view_physical_rows(descriptor) + available_physical_rows = self.mram_capacity_elems // self.mlen + if final_physical_row > available_physical_rows: + raise ValueError( + "Matrix view does not fit Matrix SRAM: " + f"needs {final_physical_row} physical rows, has " + f"{available_physical_rows}" + ) + + def _matrix_view_physical_rows(self, descriptor: MatrixViewDescriptor) -> int: + words_per_row = descriptor.shape.cols // self.blen + row_groups = (words_per_row + self.mlen // self.blen - 1) // ( + self.mlen // self.blen + ) + return ( + (descriptor.shape.tile_count - 1) + * descriptor.mapping.tile_pitch_rows + + descriptor.shape.rows * row_groups + ) + + def _validate_matrix_view_scratch_descriptor(self, descriptor: MatrixViewDescriptor) -> None: + self._validate_matrix_view_descriptor(descriptor) + rows = self._matrix_view_physical_rows(descriptor) + if rows > self.mlen: + raise ValueError( + "direct Matrix-view projection must fit its one reserved scratch tile: " + f"needs {rows} physical rows, has {self.mlen}" + ) + + def _validate_matrix_view_projection( + self, descriptor: MatrixViewDescriptor, *, logical_rows: int, slot: int + ) -> None: + self._validate_matrix_view_scratch_descriptor(descriptor) + if not 0 <= slot < 4: + raise ValueError(f"Matrix-view slot must be in [0, 4), got {slot}") + if not 1 <= logical_rows <= self.blen: + raise ValueError( + "direct Matrix-view projection currently requires one BLEN-sized " + f"decode row block, got {logical_rows} logical rows" + ) + if descriptor.shape.rows != logical_rows: + raise ValueError( + "Matrix-view projection rows must match the live decode rows: " + f"expected {logical_rows}, got {descriptor.shape.rows}" + ) + if descriptor.shape.cols * descriptor.shape.tile_count != self.mlen: + raise ValueError( + "Matrix-view projection must describe one complete MLEN-wide " + "consumer packet: " + f"cols * tile_count = {descriptor.shape.cols * descriptor.shape.tile_count}, " + f"MLEN = {self.mlen}" + ) + + def _matrix_view_projection_storage( + self, descriptor: MatrixViewDescriptor, base: int | None + ) -> tuple[int, int]: + """Keep every viewed word inside a persistent, tile-aligned owner. + + A whole-row footprint within one reserved tile is conservative but + sufficient: weights allocated after reset cannot share any bank row. + """ + + tile_elems = self.mlen * self.mlen + if base is None: + base = self.mram_allocator.reserve( + "__matrix_view_scratch", tile_elems, required_tail=tile_elems + ) + if base < 0 or base % tile_elems: + raise ValueError("Matrix-view base must be a non-negative tile-aligned reservation base") + end = base + self._matrix_view_physical_rows(descriptor) * self.mlen + if end > self.mram_capacity_elems: + raise ValueError("Matrix-view base plus footprint exceeds Matrix SRAM") + if not any( + block.addr == base and end <= block.addr + block.size + for block in self.mram_allocator.reserved_blocks + ): + raise ValueError("Matrix-view base must belong to a persistent MRAM reservation") + max_k_tiles = self.mram_allocator.capacity_after_reset // tile_elems + if max_k_tiles <= 0: + raise ValueError( + "Matrix-view projection needs reserved scratch and at least one " + "weight tile in Matrix SRAM" + ) + return base, max_k_tiles + + def configure_matrix_view_v0( + self, + descriptor: MatrixViewDescriptor, + *, + slot: int = 0, + ) -> str: + """Emit the packed hot-form Matrix-view configuration.""" + + self._validate_matrix_view_descriptor(descriptor) + if not 0 <= slot < 4: + raise ValueError(f"Matrix-view slot must be in [0, 4), got {slot}") + registers = self.register_allocator.allocate_gp(2) + try: + shape_register, map_register = registers + lines = [ + "; Configure compiler-managed Matrix SRAM affine view", + *load_large_int(shape_register, descriptor.shape.pack()), + *load_large_int(map_register, descriptor.mapping.pack()), + f"L_TILE_CFG {slot}, gp{shape_register}, gp{map_register}", + ] + return self._emit("\n".join(lines) + "\n") + finally: + self.register_allocator.free_gp(registers) + + def matrix_view_accumulator_store_v0( + self, + *, + matrix_base: int, + logical_offset: int, + slot: int = 0, + ) -> str: + """Flush the current Matrix accumulator directly through a view.""" + + registers = self.register_allocator.allocate_gp(1) + try: + (matrix_register,) = registers + lines = [ + *load_large_int(matrix_register, matrix_base), + f"M_MM_WO gp{matrix_register}, gp0, {logical_offset}, {slot}", + ] + return self._emit("\n".join(lines) + "\n") + finally: + self.register_allocator.free_gp(registers) + + def matrix_tile_execute_v0( + self, + *, + destination_base: int, + source_base: int, + scale_base: int, + primitive: LTilePrimitive, + ) -> str: + """Execute one deterministic recurrence primitive over views 0/1/2. + + The three views must already dominate this call. All tensor bases are + explicit architectural operands; the instruction has no hidden model + state, cache lookup, queue, or runtime-selected loop bounds. + """ + + if not isinstance(primitive, LTilePrimitive): + raise TypeError("primitive must be an LTilePrimitive") + registers = self.register_allocator.allocate_gp(3) + try: + destination, source, scale = registers + lines = [ + *load_large_int(destination, destination_base), + *load_large_int(source, source_base), + *load_large_int(scale, scale_base), + ( + f"L_TILE_EXEC gp{destination}, gp{source}, gp{scale}, " + f"{int(primitive)}" + ), + ] + return self._emit("\n".join(lines) + "\n") + finally: + self.register_allocator.free_gp(registers) + # ======================================================================== # Matrix Projection and VRAM Operations # ======================================================================== @@ -166,6 +356,10 @@ def vram_sub_projection_stream_k_accum_to( matrix_precision: str | int = "keyvalue", set_scale: bool = False, hbm_element_bytes: int = 2, + matrix_view_descriptor: MatrixViewDescriptor | None = None, + matrix_view_base: int | None = None, + matrix_view_slot: int = 0, + configure_matrix_view: bool = True, ): """Project one output tile while keeping K chunks in the FP32 accumulator. @@ -175,12 +369,25 @@ def vram_sub_projection_stream_k_accum_to( narrow: it reloads MRAM per 4x4 output microtile and only writes once, preserving the matrix-machine accumulator across K chunks. """ - vram_matrix, mram_input, target = self._prepare_projection( - vram_matrix, mram_input, target, auto_reset_mram=True - ) if max_k_tiles <= 0: raise ValueError(f"max_k_tiles must be > 0, got {max_k_tiles}") + if matrix_view_descriptor is not None: + logical_rows = min( + self.mlen, max(0, target.shape[0] - target_row_idx * self.mlen) + ) + self._validate_matrix_view_projection( + matrix_view_descriptor, logical_rows=logical_rows, slot=matrix_view_slot + ) + matrix_view_base, available_k_tiles = self._matrix_view_projection_storage( + matrix_view_descriptor, matrix_view_base + ) + max_k_tiles = min(max_k_tiles, available_k_tiles) + elif matrix_view_base is not None: + raise ValueError("matrix_view_base requires matrix_view_descriptor") + vram_matrix, mram_input, target = self._prepare_projection( + vram_matrix, mram_input, target, auto_reset_mram=True + ) vram_layout = self.vram_matrices[vram_matrix.name] vram_row_blocks = vram_layout.get_row_blocks(vram_row_idx) physical_k = max(vram_matrix.physical_shape[1], mram_input.physical_shape[0]) @@ -188,35 +395,163 @@ def vram_sub_projection_stream_k_accum_to( tiles_per_mlen = self.mlen // self.blen valid_rows = vram_row_blocks[0].valid_shape[0] if vram_row_blocks[0].valid_shape else self.mlen row_loop_count = min(tiles_per_mlen, max(1, math.ceil(valid_rows / self.blen))) + if matrix_view_descriptor is not None: + row_loop_count = 1 + if configure_matrix_view: + self.configure_matrix_view_v0( + matrix_view_descriptor, + slot=matrix_view_slot, + ) chunks = list(_iter_k_chunks(num_k_tiles, max_k_tiles)) - - for micro_col_idx in range(tiles_per_mlen): - for micro_row_idx in range(row_loop_count): - for chunk_idx, (k_block_start, k_block_count) in enumerate(chunks): - super().reset_mram() - super().load_sub_matrix_col( - name=mram_input.name, - col_idx=mram_col_idx, - k_block_start=k_block_start, - k_block_count=k_block_count, - precision=_matrix_precision_code(matrix_precision), - set_scale=set_scale, - hbm_element_bytes=hbm_element_bytes, + compact_view_loop = ( + matrix_view_descriptor is not None + and len(chunks) == 1 + and row_loop_count == 1 + ) + gp_regs = self.register_allocator.allocate_gp(7 if compact_view_loop else 3) + try: + # If the complete K span fits beside the statically reserved view, + # load it once and reuse it for every output micro-column. The old + # ordering reloaded the same weight tiles once per BLEN column, + # making direct affine writeback up to MLEN/BLEN times more + # expensive than the ordinary projection despite identical GEMM + # work. Multi-chunk K still uses the conservative reload path + # because this machine has only one Matrix accumulator. + if len(chunks) == 1: + k_block_start, k_block_count = chunks[0] + super().reset_mram() + super().load_sub_matrix_col( + name=mram_input.name, + col_idx=mram_col_idx, + k_block_start=k_block_start, + k_block_count=k_block_count, + precision=_matrix_precision_code(matrix_precision), + set_scale=set_scale, + hbm_element_bytes=hbm_element_bytes, + ) + if compact_view_loop: + assert matrix_view_descriptor is not None + assert matrix_view_base is not None + ( + gp_act_base, + gp_act_work, + gp_mat_base, + gp_mat_work, + gp_result, + gp_logical_offset, + gp_loop, + ) = gp_regs + mram_layout = self.hbm_matrices[mram_input.name] + mram_col_blocks = mram_layout.get_col_blocks(mram_col_idx)[ + k_block_start : k_block_start + k_block_count + ] + mram_col_start_addr = self._loaded_mram_start( + mram_col_blocks, + lambda block: ( + f"{mram_input.name}[{block.row_idx}][{mram_col_idx}]" + ), ) - super().vram_sub_projection_microtile_accumulate_to( - vram_mat_name=vram_matrix.name, - vram_row_idx=vram_row_idx, - mram_mat_name=mram_input.name, - mram_col_idx=mram_col_idx, - target_matrix=target.name, - target_row_idx=target_row_idx, - target_col_idx=target_col_idx, - micro_row_idx=micro_row_idx, - micro_col_idx=micro_col_idx, - k_block_start=k_block_start, - k_block_count=k_block_count, - write_out=(chunk_idx == len(chunks) - 1), + full_batch = ( + vram_layout.physical_shape or vram_layout.full_shape + )[0] + vram_row_start_addr = vram_row_blocks[k_block_start].vram_addr + lines = [ + "; Compact Matrix-view projection loop: preserve the ordinary " + "weight reuse while advancing the logical affine destination", + *load_large_int(gp_act_base, vram_row_start_addr), + *load_large_int(gp_mat_base, mram_col_start_addr), + *load_large_int(gp_result, matrix_view_base), + *load_large_int(gp_logical_offset, 0), + f"C_LOOP_START gp{gp_loop}, {tiles_per_mlen}", + f"S_ADD_INT gp{gp_act_work}, gp{gp_act_base}, gp0", + f"S_ADD_INT gp{gp_mat_work}, gp{gp_mat_base}, gp0", + ] + for hidden_tile in range(k_block_count): + lines.append(f"M_MM 0, gp{gp_mat_work}, gp{gp_act_work}") + if hidden_tile + 1 < k_block_count: + lines.append( + f"S_ADDI_INT gp{gp_mat_work}, gp{gp_mat_work}, " + f"{self.mlen * self.mlen}" + ) + lines.append( + f"S_ADDI_INT gp{gp_act_work}, gp{gp_act_work}, " + f"{full_batch * self.mlen}" + ) + lines.extend( + [ + f"M_MM_WO gp{gp_result}, gp{gp_logical_offset}, 0, " + f"{matrix_view_slot}", + f"S_ADDI_INT gp{gp_mat_base}, gp{gp_mat_base}, " + f"{self.blen * self.mlen}", + f"S_ADDI_INT gp{gp_logical_offset}, gp{gp_logical_offset}, " + f"{self.blen}", + f"C_LOOP_END gp{gp_loop}", + ] ) + self._emit("\n".join(lines) + "\n") + else: + for micro_col_idx in range(tiles_per_mlen): + for micro_row_idx in range(row_loop_count): + super().vram_sub_projection_microtile_accumulate_to( + vram_mat_name=vram_matrix.name, + vram_row_idx=vram_row_idx, + mram_mat_name=mram_input.name, + mram_col_idx=mram_col_idx, + target_matrix=target.name, + target_row_idx=target_row_idx, + target_col_idx=target_col_idx, + micro_row_idx=micro_row_idx, + micro_col_idx=micro_col_idx, + k_block_start=k_block_start, + k_block_count=k_block_count, + write_out=True, + gp_regs=gp_regs, + matrix_view_base=matrix_view_base, + matrix_view_logical_offset=( + micro_col_idx * self.blen + if matrix_view_descriptor is not None + else None + ), + matrix_view_slot=matrix_view_slot, + ) + else: + for micro_col_idx in range(tiles_per_mlen): + for micro_row_idx in range(row_loop_count): + for chunk_idx, (k_block_start, k_block_count) in enumerate(chunks): + super().reset_mram() + super().load_sub_matrix_col( + name=mram_input.name, + col_idx=mram_col_idx, + k_block_start=k_block_start, + k_block_count=k_block_count, + precision=_matrix_precision_code(matrix_precision), + set_scale=set_scale, + hbm_element_bytes=hbm_element_bytes, + ) + super().vram_sub_projection_microtile_accumulate_to( + vram_mat_name=vram_matrix.name, + vram_row_idx=vram_row_idx, + mram_mat_name=mram_input.name, + mram_col_idx=mram_col_idx, + target_matrix=target.name, + target_row_idx=target_row_idx, + target_col_idx=target_col_idx, + micro_row_idx=micro_row_idx, + micro_col_idx=micro_col_idx, + k_block_start=k_block_start, + k_block_count=k_block_count, + write_out=(chunk_idx == len(chunks) - 1), + gp_regs=gp_regs, + matrix_view_base=matrix_view_base, + matrix_view_logical_offset=( + micro_col_idx * self.blen + if matrix_view_descriptor is not None + else None + ), + matrix_view_slot=matrix_view_slot, + ) + finally: + self.register_allocator.free_gp(gp_regs) def vram_sub_projection_packed_skinny_stream_k_accum_to( self, @@ -318,6 +653,8 @@ def linear_projection( matrix_precision: str | int = "weights", set_scale: bool = True, hbm_element_bytes: int = 1, + matrix_view_descriptor: MatrixViewDescriptor | None = None, + matrix_view_slot: int = 0, ): """Emit tiled PLENA linear projection, including K-split accumulation.""" mlen = self.mlen @@ -338,7 +675,25 @@ def linear_projection( num_row_blocks = math.ceil(physical_rows / mlen) num_col_blocks = math.ceil(physical_out_features / mlen) num_k_tiles = math.ceil(physical_k / mlen) - max_k_tiles = self.mram_tile_capacity + matrix_view_base = None + if matrix_view_descriptor is not None: + self._validate_matrix_view_projection( + matrix_view_descriptor, logical_rows=rows, slot=matrix_view_slot + ) + if num_col_blocks != 1 or num_row_blocks != 1: + raise ValueError( + "direct Matrix-view projection supports exactly one output tile; " + "multiple output blocks need separate views or an explicit consumer/store" + ) + matrix_view_base, max_k_tiles = self._matrix_view_projection_storage( + matrix_view_descriptor, None + ) + self.configure_matrix_view_v0( + matrix_view_descriptor, + slot=matrix_view_slot, + ) + else: + max_k_tiles = self.mram_tile_capacity # When rows is not a multiple of mlen the hardware still operates on # full tiles; only the first `rows` rows contain valid output. @@ -365,6 +720,29 @@ def emit_projection(row_idx, col_idx, target, target_row_idx, target_col_idx, ** **k_split, ) + if matrix_view_descriptor is not None: + assert matrix_view_base is not None + for col_idx in range(num_col_blocks): + for row_idx in range(num_row_blocks): + self.vram_sub_projection_stream_k_accum_to( + input_var, + row_idx, + weight_var, + col_idx, + output, + row_idx, + col_idx, + max_k_tiles=max_k_tiles, + matrix_precision=matrix_precision, + set_scale=set_scale, + hbm_element_bytes=hbm_element_bytes, + matrix_view_descriptor=matrix_view_descriptor, + matrix_view_base=matrix_view_base, + matrix_view_slot=matrix_view_slot, + configure_matrix_view=False, + ) + return output + if num_k_tiles <= max_k_tiles: for col_idx in range(num_col_blocks): for row_idx in range(num_row_blocks): @@ -476,6 +854,8 @@ def linear_projection_bf16( weight_var: InputVar, name: str = "linear_out_bf16", physical_shape: tuple[int, int] | None = None, + matrix_view_descriptor: MatrixViewDescriptor | None = None, + matrix_view_slot: int = 0, ): """Emit a high-precision BF16 matrix projection through HBM_M_KV_TYPE. @@ -491,6 +871,8 @@ def linear_projection_bf16( matrix_precision="keyvalue", set_scale=False, hbm_element_bytes=2, + matrix_view_descriptor=matrix_view_descriptor, + matrix_view_slot=matrix_view_slot, ) def linear_projection_bias_bf16( diff --git a/aten/tests/test_lcompute.py b/aten/tests/test_lcompute.py new file mode 100644 index 0000000..86bdc43 --- /dev/null +++ b/aten/tests/test_lcompute.py @@ -0,0 +1,690 @@ +from pathlib import Path +from copy import deepcopy + +from compiler.aten.plena.isa_matrix_view import ( + MatrixViewAllocation, MatrixViewDescriptor, MatrixViewMap, MatrixViewShape, + validate_disjoint_matrix_views, +) + +import pytest + +from assembler.assembly_to_binary import AssemblyToBinary +from compiler.aten.plena import PlenaCompiler +from compiler.aten.plena.program_lcompute import ( + BF16_BYTES, + KIMI_KDA, + NEMOTRON_MAMBA, + ONE_MIB, + MatrixRecurrenceSpec, + MatrixSramPoint, + RecurrenceKind, + RecurrenceLayout, + build_recurrence_field_manifest, + build_recurrence_working_set, + lower_matrix_recurrence, + validate_recurrence_field_loads, + validate_recurrence_output_stores, +) +from compiler.aten.plena.isa_matrix_view import MatrixViewFlags, validate_matrix_view_dominance + + +ROOT = Path(__file__).resolve().parents[2] + + +def test_recurrence_point_enforces_the_simulator_bank_limit() -> None: + MatrixSramPoint(mlen=64, banks=64, bank_width=1).validate() + with pytest.raises(ValueError, match="no larger than 64"): + MatrixSramPoint(mlen=128, banks=128, bank_width=1).validate() + + +def test_official_shapes_and_plena_bf16_state_sizes_are_explicit() -> None: + assert (NEMOTRON_MAMBA.heads, NEMOTRON_MAMBA.row_elements) == (64, 64) + assert (KIMI_KDA.heads, KIMI_KDA.row_elements) == (96, 128) + assert NEMOTRON_MAMBA.recurrence_rows == KIMI_KDA.recurrence_rows == 128 + assert NEMOTRON_MAMBA.state_bytes_per_layer == ONE_MIB + assert KIMI_KDA.state_bytes_per_layer == 3 * ONE_MIB + + +def test_mamba_field_contract_is_reusable_by_a_real_checkpoint_shape() -> None: + checkpoint = MatrixRecurrenceSpec( + name="mamba2_130m", + kind=RecurrenceKind.MAMBA, + heads=24, + row_elements=64, + recurrence_rows=128, + primitives=NEMOTRON_MAMBA.primitives, + ) + point = MatrixSramPoint() + working_set = build_recurrence_working_set( + checkpoint, + layout=RecurrenceLayout.AFFINE, + point=point, + ) + assembly = lower_matrix_recurrence( + checkpoint, + layout=RecurrenceLayout.AFFINE, + point=point, + ) + + assert working_set.group_heads == 24 + assert working_set.groups == 1 + assert checkpoint.state_bytes_per_layer == 24 * 128 * 64 * BF16_BYTES + assert "@stage=mamba2_130m_matrix_recurrence" in assembly + assert assembly.count("L_TILE_EXEC") == 4 + validate_matrix_view_dominance(assembly) + + +@pytest.mark.parametrize( + ("capacity_bytes", "spec", "fixed_heads", "affine_heads"), + [ + (ONE_MIB, NEMOTRON_MAMBA, 32, 32), + (ONE_MIB, KIMI_KDA, 16, 16), + (2 * ONE_MIB, NEMOTRON_MAMBA, 32, 32), + (2 * ONE_MIB, KIMI_KDA, 16, 16), + ], +) +def test_capacity_point_sets_an_honest_head_group( + capacity_bytes: int, + spec, + fixed_heads: int, + affine_heads: int, +) -> None: + point = MatrixSramPoint(capacity_bytes=capacity_bytes) + for layout, group_heads in ( + (RecurrenceLayout.FIXED, fixed_heads), + (RecurrenceLayout.AFFINE, affine_heads), + ): + working_set = build_recurrence_working_set(spec, layout=layout, point=point) + assert working_set.group_heads == group_heads + assert working_set.packet_values == group_heads * spec.row_elements + assert working_set.capacity_facts["bank_words"] <= point.capacity_bank_words + assert working_set.capacity_facts["max_bank_row"] <= point.depth_rows + + +@pytest.mark.parametrize( + ("capacity_bytes", "spec", "expected_scalar_cols"), + [ + (ONE_MIB, NEMOTRON_MAMBA, 64), + (ONE_MIB, KIMI_KDA, 32), + (2 * ONE_MIB, NEMOTRON_MAMBA, 64), + (2 * ONE_MIB, KIMI_KDA, 32), + ], +) +def test_recurrence_scalar_and_head_major_fields_match_their_consumers( + capacity_bytes: int, + spec, + expected_scalar_cols: int, +) -> None: + point = MatrixSramPoint(capacity_bytes=capacity_bytes) + for layout in RecurrenceLayout: + working_set = build_recurrence_working_set(spec, layout=layout, point=point) + compact = {"beta"} if spec is KIMI_KDA else {"dt", "update", "c", "d"} + for name in compact: + descriptor = working_set.allocation(name).descriptor + assert descriptor.shape.tile_count == 1 + assert descriptor.shape.cols == expected_scalar_cols + assert descriptor.mapping.flags & MatrixViewFlags.BROADCAST_MINOR + # One row stores one [a,b] pair per head at most; it must not + # replicate each scalar across a complete recurrent row. + assert descriptor.shape.cols < (working_set.group_heads * spec.row_elements) + if spec is KIMI_KDA: + for name, rows in (("decay", 2), ("key", 1), ("query", 1)): + descriptor = working_set.allocation(name).descriptor + assert descriptor.shape.tile_count == working_set.group_heads + assert descriptor.shape.rows == rows + assert descriptor.shape.cols >= working_set.state_rows_per_chunk + assert descriptor.mapping.flags & MatrixViewFlags.BROADCAST_MINOR + + +def test_every_active_recurrence_view_uses_uniform_bf16() -> None: + for spec in (NEMOTRON_MAMBA, KIMI_KDA): + assert spec.state_bytes_per_layer == ( + spec.heads + * spec.recurrence_rows + * spec.row_elements + * BF16_BYTES + ) + for layout in RecurrenceLayout: + working_set = build_recurrence_working_set(spec, layout=layout) + assert working_set.point.element_bytes == BF16_BYTES + assert all( + not allocation.descriptor.mapping.flags + & ~MatrixViewFlags.BROADCAST_MINOR + for allocation in working_set.allocations + ) + + +def test_state_transfers_are_real_viewed_dma_not_timing_comments() -> None: + for spec in (NEMOTRON_MAMBA, KIMI_KDA): + assembly = lower_matrix_recurrence( + spec, + layout="affine", + state_hbm_base=0x20000, + hbm_address_register=2, + ) + lines = assembly.splitlines() + markers = [ + index + for index, line in enumerate(lines) + if line.startswith("; @matrix_state_") + ] + assert markers + assert assembly.count("H_PREFETCH_V.MV") > 0 + assert assembly.count("H_STORE_V.MV") > 0 + for index in markers: + marker = lines[index] + assert "precision=bf16" in marker + assert "hbm_byte_offset=" in marker + next_marker = next( + ( + candidate + for candidate in range(index + 1, len(lines)) + if lines[candidate].startswith("; @matrix_state_") + or lines[candidate].startswith("; @matrix_field_load") + or lines[candidate].startswith("; @l_tile_step=") + ), + len(lines), + ) + transfer = "\n".join(lines[index:next_marker]) + assert "L_TILE_CFG 3" in transfer + assert ("H_PREFETCH_V.MV" in transfer) ^ ("H_STORE_V.MV" in transfer) + assert ", a2, 0, 2, 3" in transfer + + +def test_every_recurrence_field_is_loaded_into_a_consumer_view() -> None: + for spec in (NEMOTRON_MAMBA, KIMI_KDA): + for layout in RecurrenceLayout: + working_set = build_recurrence_working_set(spec, layout=layout) + manifest = build_recurrence_field_manifest( + working_set, + field_hbm_base=spec.state_bytes_per_layer, + ) + assembly = lower_matrix_recurrence(spec, layout=layout) + assert validate_recurrence_field_loads( + assembly, + expected={ + packet.key + for packet in manifest.packets + if packet.field != "output_result" + }, + ) + assert assembly.count("H_PREFETCH_V.MV") > working_set.groups + + mamba = lower_matrix_recurrence(NEMOTRON_MAMBA, layout="affine") + for field in ("x", "scratch_zero", "dt", "update", "c", "output_zero", "d"): + assert f"@matrix_field_load field={field} " in mamba + for step in ( + "mamba_dt_times_x", + "mamba_state_decay_rank1_update", + "mamba_c_readout", + "mamba_skip", + ): + assert f"@l_tile_step={step}" in mamba + + kimi = lower_matrix_recurrence(KIMI_KDA, layout="affine") + for field in ( + "prediction_zero", + "decay", + "key", + "value", + "beta", + "output_zero", + "query", + ): + assert f"@matrix_field_load field={field} " in kimi + for step in ( + "kda_decay", + "kda_prediction", + "kda_beta_error", + "kda_rank1_update", + "kda_readout", + ): + assert f"@l_tile_step={step}" in kimi + + +def test_comment_only_field_claim_is_rejected() -> None: + with pytest.raises(ValueError, match="comment-only"): + validate_recurrence_field_loads( + "; @matrix_field_write=x,dt layout=consumer_view\nL_TILE_EXEC gp1, gp2, gp3, 0\n" + ) + + +@pytest.mark.parametrize("spec", [NEMOTRON_MAMBA, KIMI_KDA]) +@pytest.mark.parametrize("layout", list(RecurrenceLayout)) +def test_every_head_group_output_is_stored_once_without_aliasing(spec, layout) -> None: + working_set = build_recurrence_working_set(spec, layout=layout) + assembly = lower_matrix_recurrence(spec, layout=layout) + stores = validate_recurrence_output_stores( + assembly, + expected_groups=working_set.groups, + ) + assert len(stores) == working_set.groups + assert len(set(stores.values())) == working_set.groups + + +def test_missing_head_group_output_store_is_rejected() -> None: + working_set = build_recurrence_working_set( + NEMOTRON_MAMBA, + layout=RecurrenceLayout.AFFINE, + ) + assembly = lower_matrix_recurrence( + NEMOTRON_MAMBA, + layout=RecurrenceLayout.AFFINE, + ) + marker = "; @matrix_field_store field=output_result" + first = assembly.index(marker) + next_group = assembly.index("; @head_group=1/", first) + broken = assembly[:first] + assembly[next_group:] + with pytest.raises(ValueError, match="output coverage differs"): + validate_recurrence_output_stores( + broken, + expected_groups=working_set.groups, + ) + + +@pytest.mark.parametrize("spec", [NEMOTRON_MAMBA, KIMI_KDA]) +def test_fixed_chunks_have_distinct_field_packets(spec) -> None: + working_set = build_recurrence_working_set(spec, layout="fixed") + assert working_set.chunks > 1 + manifest = build_recurrence_field_manifest( + working_set, + field_hbm_base=spec.state_bytes_per_layer, + ) + chunk_fields = ( + ("update", "c") if spec is NEMOTRON_MAMBA else ("decay", "key", "query") + ) + for group in range(working_set.groups): + for field in chunk_fields: + packets = [ + manifest.packet(field, group=group, chunk=chunk) + for chunk in range(working_set.chunks) + ] + assert ( + len({packet.hbm_byte_offset for packet in packets}) + == working_set.chunks + ) + assert all(packet.logical_values > 0 for packet in packets) + + +def test_complete_lowerings_assemble_and_views_dominate_every_exec( + tmp_path: Path, +) -> None: + assembler = AssemblyToBinary( + str(ROOT / "doc" / "operation.svh"), + str(ROOT / "doc" / "configuration.svh"), + ) + for capacity_bytes in (ONE_MIB, 2 * ONE_MIB): + point = MatrixSramPoint(capacity_bytes=capacity_bytes) + for spec in (NEMOTRON_MAMBA, KIMI_KDA): + for layout in RecurrenceLayout: + assembly = lower_matrix_recurrence(spec, layout=layout, point=point) + validate_matrix_view_dominance(assembly) + stem = f"{spec.name}_{capacity_bytes}_{layout}" + asm_path = tmp_path / f"{stem}.asm" + mem_path = tmp_path / f"{stem}.mem" + asm_path.write_text(assembly) + assembler.generate_binary(str(asm_path), str(mem_path)) + words = [line for line in mem_path.read_text().splitlines() if line] + assert words + assert all(0 <= int(word, 16) < 2**32 for word in words) + + +@pytest.mark.parametrize("spec", [NEMOTRON_MAMBA, KIMI_KDA]) +@pytest.mark.parametrize("layout", [RecurrenceLayout.FIXED, RecurrenceLayout.AFFINE]) +def test_diagnostic_snapshots_are_explicit_and_do_not_replace_persistent_state(spec, layout): + working = build_recurrence_working_set(spec, layout=layout) + program = lower_matrix_recurrence(spec, layout=layout, snapshot_hbm_base=64 * 1024 * 1024) + assert program.count("@state_snapshot ") == working.groups * working.chunks + assert program.count("@matrix_state_store ") == working.groups * working.chunks + for bad_base in (0, 3, spec.state_bytes_per_layer): + with pytest.raises(ValueError, match="snapshot"): + lower_matrix_recurrence(spec, layout=layout, snapshot_hbm_base=bad_base) + + +@pytest.mark.parametrize("field_base", [0, 64, KIMI_KDA.state_bytes_per_layer - 64]) +def test_explicit_prepared_fields_cannot_overlap_live_state(field_base): + with pytest.raises(ValueError, match="fields overlap live recurrent state"): + lower_matrix_recurrence( + KIMI_KDA, layout="affine", state_hbm_base=0, field_hbm_base=field_base + ) + + +def test_adjacent_state_and_field_arenas_are_legal_in_both_orders(): + working = build_recurrence_working_set(KIMI_KDA, layout="affine") + fields = build_recurrence_field_manifest(working, field_hbm_base=0) + for state_base, field_base in ((0, KIMI_KDA.state_bytes_per_layer), (fields.end, 0)): + program = lower_matrix_recurrence( + KIMI_KDA, layout="affine", + state_hbm_base=state_base, field_hbm_base=field_base, + ) + assert "L_TILE_EXEC" in program + + +@pytest.mark.parametrize("base", [2**32, 2**32 - 64]) +@pytest.mark.parametrize("arena", ["state", "field", "snapshot"]) +def test_dma_arenas_cannot_wrap_the_gp_address_window(arena, base): + with pytest.raises(ValueError, match="32-bit GP DMA offset window"): + lower_matrix_recurrence( + KIMI_KDA, layout="affine", **{f"{arena}_hbm_base": base} + ) + + +def test_exact_top_of_gp_window_is_legal_with_nonzero_hbm_base_register(tmp_path): + working = build_recurrence_working_set(KIMI_KDA, layout="affine") + fields = build_recurrence_field_manifest(working, field_hbm_base=0) + field_base = 2**32 - fields.end + state_base = field_base - KIMI_KDA.state_bytes_per_layer + program = lower_matrix_recurrence( + KIMI_KDA, layout="affine", state_hbm_base=state_base, + field_hbm_base=field_base, hbm_address_register=7, + ) + assert f"@field_hbm_end={2**32}" in program + assert ", a7, 0, 2, 3" in program + asm_path, mem_path = tmp_path / "top.asm", tmp_path / "top.mem" + asm_path.write_text(program) + assembler = AssemblyToBinary(ROOT / "doc/operation.svh", ROOT / "doc/configuration.svh") + assembler.generate_binary(str(asm_path), str(mem_path)) + assert all(0 <= int(line, 16) < 2**32 for line in mem_path.read_text().splitlines()) + + +def test_projection_accumulator_directly_writes_a_configured_matrix_view(tmp_path): + program = PlenaCompiler(mlen=64, blen=4, mram_tile_capacity=64) + x_hbm = program.input( + "view_x", + shape=(1, 64), + physical_shape=(4, 64), + real_data_ratio=1.0, + ) + x = program.load_batch(x_hbm, name="view_x_vram") + weight = program.input( + "view_weight", + shape=(64, 64), + physical_shape=(64, 64), + real_data_ratio=1.0, + ) + output = program.alloc( + "view_output", + 1, + 64, + strict=False, + physical_shape=(4, 64), + ) + descriptor = MatrixViewDescriptor( + shape=MatrixViewShape(rows=1, cols=8, tile_count=8), + mapping=MatrixViewMap(tile_pitch_rows=2), + ) + + program.vram_sub_projection_stream_k_accum_to( + x, + 0, + weight, + 0, + output, + 0, + 0, + max_k_tiles=1, + matrix_view_descriptor=descriptor, + matrix_view_slot=2, + ) + assembly = program.get_code() + + validate_matrix_view_dominance(assembly) + assert assembly.count("L_TILE_CFG") == 1 + assert assembly.count("M_MM_WO") == 1 + assert "C_LOOP_START" in assembly + assert "S_ADDI_INT gp" in assembly + assert ", 4" in assembly + # All sixteen output micro-columns reuse the one resident weight tile. + assert assembly.count("H_PREFETCH_M") == 1 + assert "L_MVIEW_LOAD" not in assembly + assert "L_MVIEW_WO" not in assembly + + asm_path = tmp_path / "matrix_view_projection.asm" + mem_path = tmp_path / "matrix_view_projection.mem" + asm_path.write_text(assembly) + assembler = AssemblyToBinary( + str(ROOT / "doc/operation.svh"), + str(ROOT / "doc/configuration.svh"), + ) + words = assembler.generate_binary(str(asm_path), str(mem_path)) + assert sum(word & 0x3F == 0x3F for word in words) == 1 + viewed_writebacks = [word for word in words if word & 0x3F == 0x06] + assert len(viewed_writebacks) == 1 + assert all((word >> 31) & 1 for word in viewed_writebacks) + + +def test_wide_projection_reserves_existing_matrix_scratch_and_streams_weights(): + program = PlenaCompiler(mlen=64, blen=4, mram_tile_capacity=2) + x_hbm = program.input( + "wide_view_x", + shape=(1, 128), + physical_shape=(4, 128), + real_data_ratio=1.0, + ) + x = program.load_batch(x_hbm, name="wide_view_x_vram") + weight = program.input( + "wide_view_weight", + shape=(128, 64), + physical_shape=(128, 64), + real_data_ratio=1.0, + ) + descriptor = MatrixViewDescriptor( + shape=MatrixViewShape(rows=1, cols=8, tile_count=8), + mapping=MatrixViewMap(tile_pitch_rows=2), + ) + + program.linear_projection_bf16( + x, + weight, + name="wide_view_output", + matrix_view_descriptor=descriptor, + matrix_view_slot=1, + ) + assembly = program.get_code() + + validate_matrix_view_dominance(assembly) + assert assembly.count("L_TILE_CFG") == 1 + assert assembly.count("M_MM_WO") == 16 + assert "L_MVIEW_LOAD" not in assembly + # One Matrix tile is statically reserved for the view. The remaining tile + # streams the two K chunks without aliasing that scratch address. + assert assembly.count("H_PREFETCH_M") == 32 + + +def test_direct_view_projection_reuses_a_fitting_wide_k_weight_set(): + program = PlenaCompiler(mlen=64, blen=4, mram_tile_capacity=64) + x_hbm = program.input( + "fit_view_x", + shape=(1, 128), + physical_shape=(4, 128), + real_data_ratio=1.0, + ) + x = program.load_batch(x_hbm, name="fit_view_x_vram") + weight = program.input( + "fit_view_weight", + shape=(128, 64), + physical_shape=(128, 64), + real_data_ratio=1.0, + ) + descriptor = MatrixViewDescriptor( + shape=MatrixViewShape(rows=1, cols=8, tile_count=8), + mapping=MatrixViewMap(tile_pitch_rows=2), + ) + + program.linear_projection_bf16( + x, + weight, + name="fit_view_output", + matrix_view_descriptor=descriptor, + ) + assembly = program.get_code() + + assert assembly.count("M_MM_WO") == 1 + assert "C_LOOP_START" in assembly + assert assembly.count("H_PREFETCH_M") == 2 + + +def test_direct_view_projection_rejects_a_partial_consumer_packet(): + program = PlenaCompiler(mlen=64, blen=4, mram_tile_capacity=64) + x_hbm = program.input( + "bad_view_x", + shape=(1, 64), + physical_shape=(4, 64), + real_data_ratio=1.0, + ) + x = program.load_batch(x_hbm, name="bad_view_x_vram") + weight = program.input( + "bad_view_weight", + shape=(64, 64), + physical_shape=(64, 64), + real_data_ratio=1.0, + ) + incomplete = MatrixViewDescriptor( + shape=MatrixViewShape(rows=1, cols=8, tile_count=4), + mapping=MatrixViewMap(tile_pitch_rows=2), + ) + + try: + program.linear_projection_bf16( + x, + weight, + name="bad_view_output", + matrix_view_descriptor=incomplete, + ) + except ValueError as error: + assert "complete MLEN-wide consumer packet" in str(error) + else: + raise AssertionError("partial Matrix consumer packet was accepted") + + +def _view_projection_inputs(program, *, k=64, cols=64): + x_hbm = program.input("guard_x", shape=(1, k), physical_shape=(4, k), real_data_ratio=1.0) + x = program.load_batch(x_hbm, name="guard_x_vram") + weight = program.input("guard_weight", shape=(k, cols), physical_shape=(k, cols), real_data_ratio=1.0) + return x, weight + + +def _projection_state(program): + return deepcopy(( + program.get_code(), + vars(program.mram_allocator._vmm), + program.mram_allocator.reserved_blocks, + vars(program.vram_allocator._vmm), + vars(program.register_allocator), + )) + + +@pytest.mark.parametrize("entry", ["linear", "tile"]) +def test_direct_view_rejects_live_weight_alias_before_emitting_or_allocating(entry): + # tile 4 of this descriptor would overwrite the weight at 4096. + program = PlenaCompiler(mlen=64, blen=4, mram_tile_capacity=2) + x, weight = _view_projection_inputs(program) + descriptor = MatrixViewDescriptor(MatrixViewShape(1, 8, 8), MatrixViewMap(16)) + output = program.alloc("guard_output", 1, 64, strict=False, physical_shape=(4, 64)) + before = _projection_state(program) + with pytest.raises(ValueError, match="one reserved scratch tile"): + if entry == "linear": + program.linear_projection_bf16(x, weight, matrix_view_descriptor=descriptor) + else: + program.vram_sub_projection_stream_k_accum_to( + x, 0, weight, 0, output, 0, 0, max_k_tiles=1, + matrix_view_descriptor=descriptor, + ) + assert _projection_state(program) == before + + +@pytest.mark.parametrize("cols, physical_shape", [(128, None), (64, (4, 128)), (64, (128, 64))]) +def test_direct_view_rejects_multiple_output_blocks_without_losing_first_block(cols, physical_shape): + # independent output blocks previously reset the same view offset. + program = PlenaCompiler(mlen=64, blen=4, mram_tile_capacity=2) + x, weight = _view_projection_inputs(program, cols=cols) + descriptor = MatrixViewDescriptor(MatrixViewShape(1, 8, 8), MatrixViewMap(2)) + before = _projection_state(program) + with pytest.raises(ValueError, match="exactly one output tile"): + program.linear_projection_bf16( + x, weight, matrix_view_descriptor=descriptor, physical_shape=physical_shape + ) + assert _projection_state(program) == before + + +@pytest.mark.parametrize("base, error", [ + (0, "persistent MRAM reservation"), + (4, "tile-aligned reservation base"), + (-4096, "tile-aligned reservation base"), + (3 * 4096, "footprint exceeds Matrix SRAM"), +]) +def test_direct_view_requires_a_reserved_in_bounds_owner(base, error): + program = PlenaCompiler(mlen=64, blen=4, mram_tile_capacity=3) + x, weight = _view_projection_inputs(program) + output = program.alloc("guard_output", 1, 64, strict=False, physical_shape=(4, 64)) + # A transient allocation is not sufficient: _prepare_projection resets it. + assert program.mram_allocator.allocate("old_weight", 4096) == 0 + descriptor = MatrixViewDescriptor(MatrixViewShape(1, 8, 8), MatrixViewMap(2)) + before = _projection_state(program) + with pytest.raises(ValueError, match=error): + program.vram_sub_projection_stream_k_accum_to( + x, 0, weight, 0, output, 0, 0, max_k_tiles=1, + matrix_view_descriptor=descriptor, matrix_view_base=base, + ) + assert _projection_state(program) == before + + +def test_direct_view_rejects_missing_weight_capacity_before_reserving_scratch(): + program = PlenaCompiler(mlen=64, blen=4, mram_tile_capacity=1) + x, weight = _view_projection_inputs(program) + descriptor = MatrixViewDescriptor(MatrixViewShape(1, 8, 8), MatrixViewMap(2)) + before = _projection_state(program) + with pytest.raises(ValueError, match="insufficient room for weight tiles"): + program.linear_projection_bf16(x, weight, matrix_view_descriptor=descriptor) + assert _projection_state(program) == before + + +@pytest.mark.parametrize("entry", ["linear", "tile"]) +def test_nonzero_reserved_view_keeps_weights_disjoint_when_streaming_k(entry, tmp_path): + program = PlenaCompiler(mlen=64, blen=4, mram_tile_capacity=3) + # A pre-existing allocation puts scratch at a nonzero base. Reset does not + # reuse the hole below scratch, leaving only ONE tile for two K chunks. + program.mram_allocator.allocate("old_weight", 4096) + base = program.reserve_matrix_view_scratch_v0() + assert base == 4096 + x, weight = _view_projection_inputs(program, k=128) + descriptor = MatrixViewDescriptor(MatrixViewShape(1, 8, 8), MatrixViewMap(2)) + if entry == "linear": + program.linear_projection_bf16(x, weight, matrix_view_descriptor=descriptor) + else: + output = program.alloc("guard_output", 1, 64, strict=False, physical_shape=(4, 64)) + program.vram_sub_projection_stream_k_accum_to( + x, 0, weight, 0, output, 0, 0, max_k_tiles=2, + matrix_view_descriptor=descriptor, matrix_view_base=base, + ) + assembly = program.get_code() + assert assembly.count("H_PREFETCH_M") == 32 + assert assembly.count("M_MM_WO") == 16 + resident = [block for block in program.mram_allocator._vmm.used_stack + if block.name != "__matrix_view_scratch"] + assert [block.addr for block in resident] == [8192] + facts = validate_disjoint_matrix_views( + [MatrixViewAllocation("scratch", base, descriptor), + MatrixViewAllocation("weight", resident[0].addr, + MatrixViewDescriptor(MatrixViewShape(64, 64), MatrixViewMap(64)))], + mlen=64, banks=16, bank_width=4, depth_rows=192, + ) + assert facts["max_bank_row"] == 192 + asm_path, mem_path = tmp_path / "nonzero.asm", tmp_path / "nonzero.mem" + asm_path.write_text(assembly) + words = AssemblyToBinary(str(ROOT / "doc/operation.svh"), str(ROOT / "doc/configuration.svh")).generate_binary( + str(asm_path), str(mem_path) + ) + assert words + + +def test_reserved_view_cannot_be_freed_or_confused_with_a_transient_weight(): + program = PlenaCompiler(mlen=64, blen=4, mram_tile_capacity=3) + program.reserve_matrix_view_scratch_v0() + with pytest.raises(ValueError, match="cannot be freed"): + program.mram_allocator.free("__matrix_view_scratch") + program.mram_allocator.allocate("live_weight", 4096) + before = _projection_state(program) + with pytest.raises(ValueError, match="conflicts with a live allocation"): + program.reserve_matrix_view_scratch_v0("live_weight") + assert _projection_state(program) == before diff --git a/aten/tests/test_matrix_view.py b/aten/tests/test_matrix_view.py new file mode 100644 index 0000000..07af12b --- /dev/null +++ b/aten/tests/test_matrix_view.py @@ -0,0 +1,399 @@ +from __future__ import annotations + +import pytest + +from compiler.aten.plena import PlenaCompiler +from compiler.aten.plena.isa_matrix_view import ( + L_MVIEW_CONTRACT_VERSION, + LTilePrimitive, + L_MVIEW_OPCODE, + MatrixViewAllocation, + MatrixViewAxis, + MatrixViewDescriptor, + MatrixViewFlags, + MatrixViewForm, + MatrixViewMap, + MatrixViewShape, + decode_l_tile_exec_word, + decode_matrix_view_dma_slot, + decode_l_mview_word, + encode_l_tile_exec, + encode_l_tile_cfg, + encode_matrix_view_dma_word, + validate_matrix_view_dominance, + validate_disjoint_matrix_views, +) + + +def test_matrix_view_enforces_the_simulator_bank_limit() -> None: + descriptor = MatrixViewDescriptor(MatrixViewShape(1, 1), MatrixViewMap(1)) + descriptor.validate_for_machine(banks=64, bank_width=1) + with pytest.raises(ValueError, match="no larger than 64"): + descriptor.validate_for_machine(banks=128, bank_width=1) + + +def test_matrix_view_v3_golden_mapping_words() -> None: + assert L_MVIEW_CONTRACT_VERSION == 3 + assert MatrixViewMap(tile_pitch_rows=64).pack() == 0x00000040 + assert ( + MatrixViewMap(tile_pitch_rows=0, tile_phase_stride=4).pack() + == 0x01000000 + ) + assert ( + MatrixViewMap( + tile_pitch_rows=0, + tile_phase_stride=4, + flags=MatrixViewFlags.BROADCAST_MINOR, + ).pack() + == 0x81000000 + ) + + +def test_real_shapes_and_fixed_wiring_mapping_roundtrip() -> None: + shape = MatrixViewShape(rows=128, cols=128, tile_count=96) + mapping = MatrixViewMap(tile_pitch_rows=128) + + assert MatrixViewShape.unpack(shape.pack()) == shape + assert MatrixViewMap.unpack(mapping.pack()) == mapping + + +def test_phased_mapping_roundtrips_without_model_specific_fields() -> None: + mapping = MatrixViewMap(tile_pitch_rows=128, tile_phase_stride=11) + restored = MatrixViewMap.unpack(mapping.pack()) + assert restored.tile_pitch_rows == 128 + assert restored.tile_phase_stride == 11 + + +def test_programmable_row_coefficient_is_rejected_before_rtl() -> None: + word = 128 | (3 << 16) | (11 << 22) + with pytest.raises(ValueError, match=r"bits \[21:16\] are reserved"): + MatrixViewMap.unpack(word) + + +def test_zero_pitch_compacts_tiles_into_distinct_banks() -> None: + descriptor = MatrixViewDescriptor( + MatrixViewShape(rows=128, cols=128, tile_count=8), + MatrixViewMap( + tile_pitch_rows=0, + tile_phase_stride=4, + ), + ) + descriptor.validate_for_machine(banks=64, bank_width=32) + assert MatrixViewMap.unpack(descriptor.mapping.pack()).tile_pitch_rows == 0 + + +def test_zero_pitch_without_a_tile_phase_is_rejected_as_aliasing() -> None: + descriptor = MatrixViewDescriptor( + MatrixViewShape(rows=128, cols=128, tile_count=8), + MatrixViewMap(tile_pitch_rows=0), + ) + with pytest.raises(ValueError, match="aliases logical bank words"): + descriptor.validate_for_machine(banks=64, bank_width=32) + + +@pytest.mark.parametrize("reserved_flag", [1 << 0, 1 << 1, 1 << 2]) +def test_unused_mapping_flag_bits_are_reserved(reserved_flag: int) -> None: + word = 4 | (reserved_flag << 28) + with pytest.raises(ValueError, match="unknown Matrix-view flags"): + MatrixViewMap.unpack(word) + + + + + + + + +def test_l_tile_exec_rejects_reserved_primitives_and_high_bits() -> None: + with pytest.raises(ValueError, match="reserved L_TILE primitive"): + encode_l_tile_exec( + dst_register=1, src1_register=2, src2_register=3, primitive=3 + ) + valid = encode_l_tile_exec( + dst_register=1, + src1_register=2, + src2_register=3, + primitive=LTilePrimitive.DOT_REDUCE, + ) + with pytest.raises(ValueError, match="reserved L_TILE_EXEC bits"): + decode_l_tile_exec_word(valid | (1 << 31)) + + + + +def test_descriptor_rejects_a_pitch_that_aliases_tiles() -> None: + descriptor = MatrixViewDescriptor( + MatrixViewShape(rows=128, cols=128, tile_count=16), + MatrixViewMap(tile_pitch_rows=127), + ) + with pytest.raises(ValueError, match="aliases logical bank words"): + descriptor.validate_for_machine(banks=64, bank_width=2) + + +@pytest.mark.parametrize( + ("banks", "bank_width"), + [ + (8, 4), + (16, 4), + (32, 4), + (64, 32), + ], +) +def test_descriptor_does_not_encode_machine_bank_geometry( + banks: int, bank_width: int +) -> None: + descriptor = MatrixViewDescriptor( + MatrixViewShape(rows=64, cols=2048, tile_count=32), + MatrixViewMap( + tile_pitch_rows=4096, + flags=MatrixViewFlags(0), + ), + ) + descriptor.validate_for_machine(banks=banks, bank_width=bank_width) + + +@pytest.mark.parametrize( + "word", + [ + L_MVIEW_OPCODE, + L_MVIEW_OPCODE | (2 << 22), + L_MVIEW_OPCODE | (6 << 22), + encode_l_tile_cfg(slot=0, shape_register=1, map_register=2) | (1 << 18), + encode_l_tile_cfg(slot=0, shape_register=1, map_register=2) | (1 << 31), + ], +) +def test_noncanonical_or_reserved_words_fail(word: int) -> None: + with pytest.raises(ValueError): + decode_l_mview_word(word) + + +def test_matrix_consumer_must_be_dominated_by_an_explicit_view() -> None: + validate_matrix_view_dominance( + """ +L_TILE_CFG 2, gp7, gp9 +M_MM 0, gp1, gp2, 2 +""" + ) + with pytest.raises(ValueError, match="before a dominating configuration"): + validate_matrix_view_dominance("M_TMM 0, gp1, gp2, 2") + + +def test_matrix_view_dma_requires_a_dominating_explicit_view() -> None: + validate_matrix_view_dominance( + "L_TILE_CFG 2, gp7, gp9\n" + "H_PREFETCH_V.MV gp1, gp2, a0, 0, 2, 2" + ) + with pytest.raises(ValueError, match="before a dominating configuration"): + validate_matrix_view_dominance( + "H_STORE_V.MV gp1, gp2, a0, 0, 2, 2" + ) + + +def test_l_tile_exec_requires_all_three_explicit_views() -> None: + configured = "\n".join( + f"L_TILE_CFG {slot}, gp7, gp9" for slot in range(3) + ) + validate_matrix_view_dominance( + configured + "\nL_TILE_EXEC gp1, gp2, gp3, 0, 2" + ) + with pytest.raises(ValueError, match="Matrix view 2"): + validate_matrix_view_dominance( + "L_TILE_CFG 0, gp7, gp9\n" + "L_TILE_CFG 1, gp7, gp9\n" + "L_TILE_EXEC gp1, gp2, gp3, 0" + ) + + with pytest.raises(ValueError, match="axis mask"): + validate_matrix_view_dominance( + configured + "\nL_TILE_EXEC gp1, gp2, gp3, 0, 4" + ) + + +def test_loop_backedge_requires_the_view_on_every_iteration() -> None: + with pytest.raises(ValueError, match="before a dominating configuration"): + validate_matrix_view_dominance( + """ +C_LOOP_START gp10, 2 +M_MV 0, gp1, gp2, 1 +C_LOOP_END gp10 +""" + ) + + validate_matrix_view_dominance( + """ +C_LOOP_START gp10, 2 +L_TILE_CFG 1, gp7, gp9 +M_MV 0, gp1, gp2, 1 +C_LOOP_END gp10 +""" + ) + + +def test_break_cannot_hide_an_unconfigured_reachable_consumer() -> None: + with pytest.raises(ValueError, match="before a dominating configuration"): + validate_matrix_view_dominance( + """ +C_BREAK +M_MV 0, gp1, gp2, 1 +""" + ) + + with pytest.raises(ValueError, match="before a dominating configuration"): + validate_matrix_view_dominance( + """ +C_LOOP_START gp10, 2 +C_BREAK +L_TILE_CFG 1, gp7, gp9 +C_LOOP_END gp10 +M_MV 0, gp1, gp2, 1 +""" + ) + + +def test_matrix_view_vector_operands_require_dominating_views() -> None: + validate_matrix_view_dominance( + "L_TILE_CFG 1, gp2, gp3\n" + "L_TILE_CFG 2, gp2, gp3\n" + "V_MUL_VV.MV gp4, gp5, gp6, 0, 6" + ) + with pytest.raises(ValueError, match="before a dominating configuration"): + validate_matrix_view_dominance("V_MUL_VV.MV gp4, gp5, gp6, 0, 6") + + +def test_direct_accumulator_writeback_is_canonical_and_dominated() -> None: + validate_matrix_view_dominance( + "L_TILE_CFG 1, gp2, gp3\nM_MM_WO gp4, gp0, 5, 1" + ) + with pytest.raises(ValueError, match="before a dominating configuration"): + validate_matrix_view_dominance("M_MM_WO gp4, gp0, 5, 1") + + +def test_compiler_emits_one_generic_matrix_view_data_path() -> None: + program = PlenaCompiler(mlen=64, blen=4, mram_tile_capacity=4) + descriptor = MatrixViewDescriptor( + MatrixViewShape(rows=1, cols=4, tile_count=16), + MatrixViewMap(tile_pitch_rows=1), + ) + matrix_base = program.reserve_matrix_view_scratch_v0() + + program.configure_matrix_view_v0(descriptor, slot=2) + program.matrix_view_accumulator_store_v0( + matrix_base=matrix_base, + logical_offset=12, + slot=2, + ) + assembly = program.compile() + + validate_matrix_view_dominance(assembly) + assert assembly.count("L_TILE_CFG") == 1 + assert assembly.count("M_MM_WO") == 1 + assert "L_MVIEW_LOAD" not in assembly + assert "L_MVIEW_STORE" not in assembly + + +def test_matrix_view_reservation_survives_weight_reset() -> None: + program = PlenaCompiler(mlen=64, blen=4, mram_tile_capacity=4) + scratch = program.reserve_matrix_view_scratch_v0() + first_weight = program.mram_allocator.allocate("weight_before_reset", 64 * 64) + program.reset_mram() + same_scratch = program.reserve_matrix_view_scratch_v0() + second_weight = program.mram_allocator.allocate("weight_after_reset", 64 * 64) + + assert scratch == same_scratch == 0 + assert first_weight == second_weight == 64 * 64 + + +def _paper_addr(row: int, bank_phase: int) -> int: + return row * 2048 + bank_phase * 32 + + +def test_compiler_proves_official_kda_chunk_and_fields_are_disjoint() -> None: + state = MatrixViewDescriptor( + MatrixViewShape(rows=16, cols=128, tile_count=16), + MatrixViewMap( + tile_pitch_rows=8, + tile_phase_stride=4, + ), + ) + scalar = MatrixViewDescriptor( + MatrixViewShape(rows=16, cols=32, tile_count=16), + MatrixViewMap( + tile_pitch_rows=1, + tile_phase_stride=3, + ), + ) + vector = MatrixViewDescriptor( + MatrixViewShape(rows=1, cols=128, tile_count=16), + MatrixViewMap( + tile_pitch_rows=1, + tile_phase_stride=3, + ), + ) + allocations = [ + MatrixViewAllocation("state", _paper_addr(0, 0), state), + MatrixViewAllocation("decay", _paper_addr(136, 0), scalar), + MatrixViewAllocation("key", _paper_addr(136, 1), scalar), + MatrixViewAllocation("query", _paper_addr(136, 2), scalar), + MatrixViewAllocation("value_or_error", _paper_addr(168, 0), vector), + MatrixViewAllocation("prediction_or_output", _paper_addr(168, 4), vector), + ] + facts = validate_disjoint_matrix_views( + allocations, mlen=2048, banks=64, bank_width=32, depth_rows=256 + ) + assert facts["max_bank_row"] == 184 + assert facts["bank_words"] < facts["capacity_bank_words"] + + +def test_compiler_proves_official_mamba_chunk_and_fields_are_disjoint() -> None: + state = MatrixViewDescriptor( + MatrixViewShape(rows=16, cols=64, tile_count=32), + MatrixViewMap( + tile_pitch_rows=4, + tile_phase_stride=2, + ), + ) + scalar = MatrixViewDescriptor( + MatrixViewShape(rows=16, cols=32, tile_count=32), + MatrixViewMap( + tile_pitch_rows=1, + tile_phase_stride=1, + ), + ) + vector = MatrixViewDescriptor( + MatrixViewShape(rows=1, cols=64, tile_count=32), + MatrixViewMap( + tile_pitch_rows=1, + tile_phase_stride=1, + ), + ) + allocations = [ + MatrixViewAllocation("state", _paper_addr(0, 0), state), + MatrixViewAllocation("decay_and_b", _paper_addr(140, 0), scalar), + MatrixViewAllocation("c", _paper_addr(140, 16), scalar), + MatrixViewAllocation("dt_and_skip", _paper_addr(140, 32), scalar), + MatrixViewAllocation("x", _paper_addr(188, 0), vector), + MatrixViewAllocation("scratch", _paper_addr(188, 2), vector), + MatrixViewAllocation("output", _paper_addr(188, 4), vector), + ] + facts = validate_disjoint_matrix_views( + allocations, mlen=2048, banks=64, bank_width=32, depth_rows=256 + ) + assert facts["max_bank_row"] == 220 + + +def test_compiler_rejects_cross_view_aliasing() -> None: + descriptor = MatrixViewDescriptor( + MatrixViewShape(rows=1, cols=128, tile_count=16), + MatrixViewMap(tile_pitch_rows=16), + ) + with pytest.raises(ValueError, match="first.*second"): + validate_disjoint_matrix_views( + [ + MatrixViewAllocation("first", 0, descriptor), + MatrixViewAllocation("second", 0, descriptor), + ], + mlen=2048, + banks=64, + bank_width=32, + depth_rows=256, + ) diff --git a/doc/operation.svh b/doc/operation.svh index 79b3651..f282318 100644 --- a/doc/operation.svh +++ b/doc/operation.svh @@ -180,7 +180,12 @@ typedef enum logic [instruction_pkg::OPCODE_WIDTH - 1:0] { V_MAX_VF = 6'h35, V_MIN_VF = 6'h36, V_TOPK = 6'h37, - C_SET_TOPK_REG = 6'h38 + C_SET_TOPK_REG = 6'h38, + + // Packed Matrix-SRAM views and deterministic tile primitives. + // CFG funct1=1; EXEC funct1=3. Other funct1 values are reserved. + // Full descriptor, explicit-view DMA and consumer encoding: plena_isa_spec.md. + L_TILE = 6'h3F } CUSTOM_ISA_OPCODE; typedef enum logic [2:0] { diff --git a/doc/plena_isa_spec.md b/doc/plena_isa_spec.md index 89d0bdd..4829e59 100644 --- a/doc/plena_isa_spec.md +++ b/doc/plena_isa_spec.md @@ -804,3 +804,67 @@ End of a hardware loop. If the loop counter (in register `rd`) is greater than 0 **Description:** Shift the vector elements left by the amount specified by `gp_reg`. Elements are shifted left, and zeros are filled in from the right. For example, `[a0, a1, a2, a3]` with shift=2 becomes `[a2, a3, 0, 0]`. + +### Matrix SRAM views and L_TILE + +`L_TILE` uses opcode `0x3F`. Four explicit view slots describe BF16 values in +Matrix SRAM; the existing Matrix, Vector and DMA opcodes select a slot only +through the qualified forms below. Ordinary DMA precision selectors retain +their existing meaning. This is a Compiler/Simulator interface proposal; RTL +resource reuse, timing and PPA have not been established. + +| Assembly form | Encoding and behavior | +| --- | --- | +| `L_TILE_CFG slot, gp_shape, gp_map` | `funct1=1`; atomically binds one of four slots. | +| `L_TILE_EXEC gp_dst, gp_src, gp_scale, primitive[, axes]` | `funct1=3`; uses slots 0/1/2 and explicit Matrix SRAM bases. | +| `H_PREFETCH_V.MV rd, rs1, rs2, stride, precision, slot` | Existing opcode `0x29`; bit 31 marks Matrix-view DMA, bits 30:29 select its slot. | +| `H_STORE_V.MV rd, rs1, rs2, stride, precision, slot` | Same qualifier on existing opcode `0x2A`. | +| `M_MM_WO gp_base, gp_offset, immediate, slot` | Existing opcode `0x06`; writes the Matrix accumulator directly into the selected view. | +| `M_… rd, rs1, rs2, slot` | Supported Matrix consumers encode `slot + 1` in `funct1`; zero remains ordinary addressing. | +| `V_ADD_VV.MV`, `V_SUB_VV.MV`, `V_MUL_VV.MV` | Existing arithmetic opcodes; `funct1[3]=1` marks viewed operands and bits 2:0 select destination/source1/source2 slots. | + +The packed shape word is `rows-1[11:0]`, `cols-1[23:12]`, +`tiles-1[31:24]`. The mapping word is `tile_pitch_rows[15:0]`, +reserved-zero bits 21:16, `tile_phase_stride[27:22]`, and flags 31:28. +Only flag bit 3 (`BROADCAST_MINOR`) is defined. The diagonal row coefficient +remains one; the descriptor adds a phase between logical tiles. A zero tile +pitch is accepted only when logical bank words remain disjoint. Compiler +validation checks descriptor bounds, live allocation overlap and configuration +dominance across loops and reachable branches. +The supported bank count is a power of two from 1 through 64, matching the Simulator. + +For `CFG`, instruction fields 9:6/13:10/17:14 carry shape GP/map GP/slot, +bits 21:18 and 31:26 are zero. For `EXEC`, those fields carry destination, +source and scale GP registers, bits 21:18 carry the primitive, and bits 26/27 +choose source/scale row (0) or column (1) axes. Bits 31:28 are zero. +Primitive 0 performs scale-and-accumulate, 1 adds a dot reduction to the existing +destination, and 2 performs an outer update. +Other function selectors and primitives are reserved. + +Explicit `.MV` DMA reserves bits 28:26. Its precision selectors are +0=Activation, 1=KeyValue and 2=BF16 state; the prepared recurrence path uses 2. +This leaves weight, activation and KV formats independent. Viewed `M_MM_WO` +uses immediate bit 17 as the marker, bits 16:15 as the slot and bits 14:0 as the +offset. **Compatibility:** previously assembled writebacks with immediate bit +17 set are reinterpreted as viewed writes (for example `0x80000046` selects +view 0). The assembler now rejects unqualified offsets that require that bit; +existing binaries need review before use with this extension. + +`isa_matrix_view.py` owns packing and layout validation; the existing Matrix +emitters and `program_matrix_ops.py` provide direct projection writeback. +Scratch is an explicit persistent MRAM reservation so weight allocator resets +cannot overwrite it. The direct projection currently supports one complete +MLEN-wide output packet with at most BLEN live rows, and retains accumulator +values across K chunks. + +`program_lcompute.py` emits complete Mamba-2/KDA recurrent cores from prepared +BF16 coefficients, including state/field DMA and per-head-group output stores. +Persistent state and operands are BF16. The Rust L_TILE primitives use FP32 +intermediates and accumulation, then round results to BF16 on writeback; this +numeric contract does not establish additional physical storage or its mapping. +The builder validates disjoint state, field and optional snapshot HBM ranges +within the 32-bit GP offset window. It exposes fixed and phased layouts +(`affine` is the existing API spelling). Projection, convolution, gates and whole-model +checkpoint ingestion are outside that prepared-input contract. Simulator tests +execute four-token official-shape recurrences and an actual projection consumer; +these do not establish whole-model numerical execution or hardware speedup.