291 lines
14 KiB
Python
291 lines
14 KiB
Python
#!/usr/bin/env python3
|
|
# SPDX-License-Identifier: AGPL-3.0-or-later
|
|
# Copyright (C) 2026 David (Corey) Herrera
|
|
"""
|
|
emit_tensix_cb_formats.py -- compile-time circular-buffer (CB) descriptor generator.
|
|
|
|
Foundry compiles each compute kernel once per TRISC thread (unpack, math, pack). Every
|
|
compile includes a generated header (ft_tensix_cb_formats.h) that bakes in, per CB index,
|
|
the data formats and tile geometry the three threads need. This tool writes that header
|
|
from a .cbtable file (or a uniform Float16_b default). Output is plain C, no includes.
|
|
|
|
Format codes are the hardware data-format codes from the public ISA documentation
|
|
(tt-isa-documentation, data-format and tile pages); everything else is Foundry-defined.
|
|
|
|
Status: Complete for the formats modelled (Float32, Float16, Tf32, Float16_b, Bfp8_b, Bfp4_b).
|
|
Issues: None
|
|
Todo: None
|
|
"""
|
|
|
|
import argparse
|
|
import re
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
TOOL_NAME = "emit_tensix_cb_formats"
|
|
MAILBOX_HEADER_RELATIVE_PATH = "../ft_tt_mailbox.h"
|
|
CB_COUNT_PATTERN = re.compile(r"^\s*#\s*define\s+FT_TT_NUM_CIRCULAR_BUFFERS\s+(\d+)u?\b", re.M)
|
|
|
|
FORMAT_CODES = {"Float32": 0, "Float16": 1, "Tf32": 4, "Float16_b": 5, "Bfp8_b": 6, "Bfp4_b": 7}
|
|
INVALID_FORMAT_CODE = 0xFF
|
|
BYTES_PER_ELEMENT = {"Float32": 4, "Tf32": 4, "Float16": 2, "Float16_b": 2}
|
|
MANTISSA_BITS = {"Bfp8_b": 8, "Bfp4_b": 4}
|
|
B_EXPONENT_FAMILY = {"Float16_b", "Bfp8_b", "Bfp4_b", "Bfp2_b"}
|
|
THIRTY_TWO_BIT_FORMATS = {"Float32", "Int32"}
|
|
|
|
SETTING_DEFAULTS = {"math_fidelity": 4, "approx": 0, "dst_accum_mode": 0}
|
|
GEOMETRY_DEFAULTS = (32, 32, 16, 4, 0, 0)
|
|
GEOMETRY_FIELD_NAMES = ("tile_rows", "tile_cols", "face_rows", "num_faces",
|
|
"partial_face", "narrow_tile")
|
|
SHARED_EXPONENT_ROW_ROUNDING = 16
|
|
FACE_COLUMNS = 16
|
|
UINT8_MAX = 255
|
|
UINT16_MAX = 65535
|
|
|
|
# Array suffix -> C type and default for an index that describes no CB.
|
|
GEOMETRY_ARRAY_SUFFIXES = ("num_faces", "partial_face", "face_rows", "narrow_tile",
|
|
"tile_rows", "tile_cols")
|
|
|
|
# Column key -> the C array suffix kernels use. The names keep one fixed sort order across
|
|
# the arrays: GCC's whole-program optimiser can order spill slots by symbol name, so renaming
|
|
# arrays out of that order moved two stack slots in one image (measured 2026-10-08, gdn).
|
|
PUBLIC_ARRAY_SUFFIX = {"num_faces": "tile_num_faces", "face_rows": "tile_face_rows",
|
|
"tile_bytes": "tile_size", "face_grid_rows": "num_face_rows",
|
|
"face_grid_cols": "num_face_cols"}
|
|
|
|
|
|
class CbTableError(Exception):
|
|
"""A malformed or unsupported input; reported as `<tool>: <message>`, exit 2."""
|
|
|
|
|
|
def read_circular_buffer_count(tool_directory):
|
|
"""Return N from `#define FT_TT_NUM_CIRCULAR_BUFFERS <digits>[u]` in ft_tt_mailbox.h."""
|
|
header_path = tool_directory / MAILBOX_HEADER_RELATIVE_PATH
|
|
try:
|
|
header_text = header_path.read_text()
|
|
except OSError as error:
|
|
raise CbTableError(f"cannot read {header_path}: {error}")
|
|
match = CB_COUNT_PATTERN.search(header_text)
|
|
if not match:
|
|
raise CbTableError(f"FT_TT_NUM_CIRCULAR_BUFFERS not defined in {header_path}")
|
|
return int(match.group(1))
|
|
|
|
|
|
def parse_integer(token, context):
|
|
try:
|
|
return int(token, 0)
|
|
except ValueError:
|
|
raise CbTableError(f"{context}: '{token}' is not an integer")
|
|
|
|
|
|
def make_cb_description(index, format_name, geometry):
|
|
"""Validate one CB line and derive its face grid and tile byte count."""
|
|
if format_name not in FORMAT_CODES:
|
|
raise CbTableError(f"cb {index}: format '{format_name}' is not modelled")
|
|
fields = dict(zip(GEOMETRY_FIELD_NAMES, geometry))
|
|
if fields["tile_cols"] % FACE_COLUMNS != 0 or fields["tile_cols"] <= 0:
|
|
raise CbTableError(f"cb {index}: tile_cols {fields['tile_cols']} is not a multiple of 16")
|
|
face_cols = min(fields["tile_cols"] // FACE_COLUMNS, fields["num_faces"])
|
|
if face_cols == 0 or fields["num_faces"] % face_cols != 0:
|
|
raise CbTableError(f"cb {index}: num_faces {fields['num_faces']} does not tile "
|
|
f"{face_cols} face column(s)")
|
|
fields.update(index=index, format=format_name, face_grid_cols=face_cols,
|
|
face_grid_rows=fields["num_faces"] // face_cols,
|
|
tile_bytes=compute_tile_bytes(format_name, fields))
|
|
for name in GEOMETRY_FIELD_NAMES + ("face_grid_cols", "face_grid_rows"):
|
|
if not 0 <= fields[name] <= UINT8_MAX:
|
|
raise CbTableError(f"cb {index}: {name} {fields[name]} does not fit uint8_t")
|
|
if not 0 <= fields["tile_bytes"] <= UINT16_MAX:
|
|
raise CbTableError(f"cb {index}: tile bytes {fields['tile_bytes']} do not fit uint16_t")
|
|
return fields
|
|
|
|
|
|
def compute_tile_bytes(format_name, fields):
|
|
"""Block floats: packed mantissas plus one shared exponent per face row (rounded to 16)."""
|
|
elements = fields["tile_rows"] * fields["tile_cols"]
|
|
if format_name in MANTISSA_BITS:
|
|
exponent_rows = fields["face_rows"] * fields["num_faces"]
|
|
rounded_rows = -(-exponent_rows // SHARED_EXPONENT_ROW_ROUNDING) * SHARED_EXPONENT_ROW_ROUNDING
|
|
return elements * MANTISSA_BITS[format_name] // 8 + rounded_rows
|
|
return elements * BYTES_PER_ELEMENT[format_name]
|
|
|
|
|
|
def parse_cb_line(arguments, descriptions, circular_buffer_count):
|
|
if len(arguments) < 2:
|
|
raise CbTableError("cb needs at least <index> <format>")
|
|
index = parse_integer(arguments[0], "cb index")
|
|
if not 0 <= index < circular_buffer_count:
|
|
raise CbTableError(f"cb index {index} out of range (0..{circular_buffer_count - 1})")
|
|
if index in descriptions:
|
|
raise CbTableError(f"cb index {index} given twice")
|
|
if len(arguments) - 2 > len(GEOMETRY_DEFAULTS):
|
|
raise CbTableError(f"cb {index}: too many fields")
|
|
given = [parse_integer(token, f"cb {index}") for token in arguments[2:]]
|
|
geometry = tuple(given) + GEOMETRY_DEFAULTS[len(given):]
|
|
descriptions[index] = make_cb_description(index, arguments[1], geometry)
|
|
|
|
|
|
def parse_cb_table(text, circular_buffer_count):
|
|
"""Return (settings, {index: description}) from cbtable text."""
|
|
settings = dict(SETTING_DEFAULTS)
|
|
descriptions = {}
|
|
for line_number, raw_line in enumerate(text.splitlines(), 1):
|
|
tokens = raw_line.split("#", 1)[0].split()
|
|
if not tokens:
|
|
continue
|
|
directive, arguments = tokens[0], tokens[1:]
|
|
try:
|
|
if directive in settings:
|
|
if len(arguments) != 1:
|
|
raise CbTableError(f"{directive} needs exactly one value")
|
|
settings[directive] = parse_integer(arguments[0], directive)
|
|
elif directive == "cb":
|
|
parse_cb_line(arguments, descriptions, circular_buffer_count)
|
|
else:
|
|
raise CbTableError(f"unknown directive '{directive}'")
|
|
except CbTableError as error:
|
|
raise CbTableError(f"line {line_number}: {error}")
|
|
if not descriptions:
|
|
raise CbTableError("table has no cb line")
|
|
return settings, descriptions
|
|
|
|
|
|
def make_uniform_default(circular_buffer_count):
|
|
descriptions = {index: make_cb_description(index, "Float16_b", GEOMETRY_DEFAULTS)
|
|
for index in range(circular_buffer_count)}
|
|
return dict(SETTING_DEFAULTS), descriptions
|
|
|
|
|
|
def choose_intermediate_format(settings, descriptions):
|
|
"""The 16-bit format that Float32 buffers pass through between unpack and pack."""
|
|
ordered = [descriptions[index] for index in sorted(descriptions)]
|
|
if settings["dst_accum_mode"] and all(d["format"] == "Float32" for d in ordered):
|
|
return "Tf32"
|
|
for description in ordered:
|
|
if description["format"] not in ("Float32", "Tf32"):
|
|
return "Float16_b" if description["format"] in B_EXPONENT_FAMILY else "Float16"
|
|
return "Float16_b"
|
|
|
|
|
|
def derive_role_formats(format_name, intermediate_format, dst_accum_mode):
|
|
"""(unpack source, unpack destination, pack source, pack destination)."""
|
|
if format_name in MANTISSA_BITS:
|
|
return (format_name,) * 4
|
|
if format_name == "Float32":
|
|
pack_source = "Float32" if dst_accum_mode else intermediate_format
|
|
return ("Float32", intermediate_format, pack_source, "Float32")
|
|
return (format_name,) * 4
|
|
|
|
|
|
def format_array(c_type, name, values, circular_buffer_count):
|
|
joined = ",".join(str(value) for value in values)
|
|
assert len(values) == circular_buffer_count
|
|
return f"FT_CB_TABLE_CONST {c_type} {name}[{circular_buffer_count}] = {{\n {joined}\n}};\n"
|
|
|
|
|
|
def build_columns(descriptions, roles, circular_buffer_count):
|
|
"""Per-index column lists, keyed by field name, with unused-index defaults filled in."""
|
|
def column(getter, default):
|
|
return [getter(i) if i in descriptions else default for i in range(circular_buffer_count)]
|
|
|
|
columns = {}
|
|
for name in GEOMETRY_ARRAY_SUFFIXES + ("tile_bytes", "face_grid_rows", "face_grid_cols"):
|
|
columns[name] = column(lambda i, n=name: descriptions[i][n], 0)
|
|
for position, label in enumerate(("unpack_src", "unpack_dst", "pack_src", "pack_dst")):
|
|
columns[label] = column(lambda i, p=position: FORMAT_CODES[roles[i][p]], INVALID_FORMAT_CODE)
|
|
return columns
|
|
|
|
|
|
def emit_stage_arrays(prefix, destination_key, columns, circular_buffer_count):
|
|
"""Arrays for one stage ('unpack' or 'pack'); the source-format array is emitted by the caller."""
|
|
text = format_array("uint8_t", f"ft_cb_{prefix}_dst_format", columns[destination_key],
|
|
circular_buffer_count)
|
|
for suffix in GEOMETRY_ARRAY_SUFFIXES:
|
|
text += format_array("uint8_t", f"ft_cb_{prefix}_{PUBLIC_ARRAY_SUFFIX.get(suffix, suffix)}", columns[suffix],
|
|
circular_buffer_count)
|
|
text += format_array("uint16_t", f"ft_cb_{prefix}_tile_size", columns["tile_bytes"],
|
|
circular_buffer_count)
|
|
for suffix in ("face_grid_rows", "face_grid_cols"):
|
|
text += format_array("uint8_t", f"ft_cb_{prefix}_{PUBLIC_ARRAY_SUFFIX.get(suffix, suffix)}", columns[suffix],
|
|
circular_buffer_count)
|
|
return text
|
|
|
|
|
|
def has_any_32bit_unpack(descriptions, roles):
|
|
return any(roles[i][0] in THIRTY_TWO_BIT_FORMATS and roles[i][1] in THIRTY_TWO_BIT_FORMATS
|
|
for i in descriptions)
|
|
|
|
|
|
def render_header(settings, descriptions, circular_buffer_count):
|
|
intermediate = choose_intermediate_format(settings, descriptions)
|
|
roles = {i: derive_role_formats(d["format"], intermediate, settings["dst_accum_mode"])
|
|
for i, d in descriptions.items()}
|
|
columns = build_columns(descriptions, roles, circular_buffer_count)
|
|
count = circular_buffer_count
|
|
source_array = format_array("uint8_t", "ft_cb_unpack_src_format", columns["unpack_src"], count)
|
|
lines = ["#pragma once", "", "#define FT_CB_TABLE_CONST static const", ""]
|
|
lines.append("#if defined(FT_TRISC_MATH) || defined(FT_TRISC_PACK)")
|
|
lines.append(f"#define FT_CB_MATH_FIDELITY {settings['math_fidelity']}")
|
|
lines.append(f"#define FT_CB_APPROX {1 if settings['approx'] else 0}")
|
|
lines.append("#endif\n")
|
|
lines.append("#if !defined(FT_TRISC_PACK)")
|
|
lines.append(source_array + emit_stage_arrays("unpack", "unpack_dst", columns, count)
|
|
+ "#endif\n")
|
|
lines.append("#if !defined(FT_TRISC_MATH) && !defined(FT_TRISC_UNPACK)")
|
|
lines.append(source_array)
|
|
lines.append(format_array("uint8_t", "ft_cb_pack_src_format", columns["pack_src"], count)
|
|
+ emit_stage_arrays("pack", "pack_dst", columns, count) + "#endif\n")
|
|
lines.append("#if defined(FT_TRISC_MATH) || defined(FT_TRISC_PACK) || defined(FT_TRISC_UNPACK)")
|
|
lines.append(f"#define FT_CB_DST_ACCUM_FP32 {1 if settings['dst_accum_mode'] else 0}")
|
|
lines.append("#endif\n")
|
|
lines.append(f"#define FT_CB_ANY_32BIT_UNPACK {1 if has_any_32bit_unpack(descriptions, roles) else 0}")
|
|
return "\n".join(lines) + "\n"
|
|
|
|
|
|
def summarize(descriptions, uniform, circular_buffer_count):
|
|
if uniform:
|
|
return f"uniform Float16_b 32x32 4-face 2048 B on all {circular_buffer_count} CBs"
|
|
return " ".join(f"{i}:{descriptions[i]['format']}/{descriptions[i]['tile_bytes']}B"
|
|
for i in sorted(descriptions))
|
|
|
|
|
|
def parse_arguments(argv):
|
|
parser = argparse.ArgumentParser(prog=TOOL_NAME, description=__doc__.split("\n")[1])
|
|
parser.add_argument("-o", "--output", required=True, help="header file to write")
|
|
parser.add_argument("--cb-table", help="input .cbtable file")
|
|
parser.add_argument("--uniform-default", action="store_true",
|
|
help="describe every CB as Float16_b 32x32")
|
|
arguments = parser.parse_args(argv)
|
|
if bool(arguments.cb_table) == arguments.uniform_default:
|
|
raise CbTableError("give exactly one of --cb-table or --uniform-default")
|
|
return arguments
|
|
|
|
|
|
def main(argv):
|
|
try:
|
|
arguments = parse_arguments(argv)
|
|
count = read_circular_buffer_count(Path(__file__).resolve().parent)
|
|
if arguments.uniform_default:
|
|
settings, descriptions = make_uniform_default(count)
|
|
else:
|
|
try:
|
|
text = Path(arguments.cb_table).read_text()
|
|
except OSError as error:
|
|
raise CbTableError(f"cannot read {arguments.cb_table}: {error}")
|
|
settings, descriptions = parse_cb_table(text, count)
|
|
header_text = render_header(settings, descriptions, count)
|
|
Path(arguments.output).write_text(header_text)
|
|
except CbTableError as error:
|
|
print(f"{TOOL_NAME}: {error}", file=sys.stderr)
|
|
return 2
|
|
except OSError as error:
|
|
print(f"{TOOL_NAME}: {error}", file=sys.stderr)
|
|
return 2
|
|
print(f"{TOOL_NAME}: {summarize(descriptions, arguments.uniform_default, count)}",
|
|
file=sys.stderr)
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main(sys.argv[1:]))
|