tt-flash 3.4.1__py3-none-any.whl

tt_flash/__init__.py

@@ -0,0 +1,59 @@
+# SPDX-FileCopyrightText: © 2024 Tenstorrent AI ULC
+# SPDX-License-Identifier: Apache-2.0
+
+# Adapted from https://github.com/python-poetry/poetry/issues/273#issuecomment-1877789967
+# This will get the semantic version from the current pyproject package definition.
+
+from typing import Any
+
+try:
+    import importlib.metadata as importlib_metadata
+except ModuleNotFoundError:
+    import importlib_metadata
+from pathlib import Path
+
+__package_version = "unknown"
+
+
+def __get_package_version() -> str:
+    """Find the version of this package."""
+    global __package_version
+
+    if __package_version != "unknown":
+        # We already set it at some point in the past,
+        # so return that previous value without any
+        # extra work.
+        return __package_version
+
+    try:
+        # Try to get the version of the current package if
+        # it is running from a distribution.
+        __package_version = importlib_metadata.version("tt-flash")
+    except importlib.metadata.PackageNotFoundError:
+        # Fall back on getting it from a local pyproject.toml.
+        # This works in a development environment where the
+        # package has not been installed from a distribution.
+        try:
+            # This gets added to the standard library as tomllib in Python3.11
+            # therefore we expect to hit a ModuleNotFoundError.
+            import tomli as toml
+        except ModuleNotFoundError:
+            import tomllib as toml
+
+        pyproject_toml_file = Path(__file__).parent.parent / "pyproject.toml"
+        if pyproject_toml_file.exists() and pyproject_toml_file.is_file():
+            __package_version = toml.loads(pyproject_toml_file.read_text())["project"][
+                "version"
+            ]
+            # Indicate it might be locally modified or unreleased.
+            __package_version = __package_version + "+"
+
+    return __package_version
+
+
+def __getattr__(name: str) -> Any:
+    """Get package attributes."""
+    if name == "__version__":
+        return __get_package_version()
+    else:
+        raise AttributeError(f"No attribute {name} in module {__name__}.")

tt_flash/blackhole.py

@@ -0,0 +1,70 @@
+# SPDX-FileCopyrightText: © 2024 Tenstorrent AI ULC
+# SPDX-License-Identifier: Apache-2.0
+
+import ctypes
+
+from tt_flash.boot_fs import tt_boot_fs_fd
+from tt_flash.error import TTError
+from . import boot_fs
+
+from tt_flash.chip import BhChip
+
+
+def writeback_boardcfg(chip: BhChip, write: bytearray) -> bytearray:
+    # Find boardcfg on chip
+    fd_in_spi = boot_fs.read_tag(
+        lambda addr, size: chip.spi_read(addr, size), "boardcfg"
+    )
+    if fd_in_spi is None:
+        raise TTError("Couldn't find boardcfg on chip")
+
+    # Find boardcfg in current fd
+    fd_to_flash = boot_fs.read_tag(
+        lambda addr, size: write[addr : addr + size], "boardcfg"
+    )
+    if fd_to_flash is None:
+        raise TTError("Couldn't find boardcfg in flash package")
+    fd_as_data = bytes(fd_in_spi[1])
+    write[fd_to_flash[0] : fd_to_flash[0] + len(fd_as_data)] = fd_as_data
+
+    flashed_fd = boot_fs.read_tag(
+        lambda addr, size: write[addr : addr + size], "boardcfg"
+    )
+    assert flashed_fd[1] == fd_in_spi[1], f"{flashed_fd[1]} != {fd_in_spi[1]}"
+
+    return write
+
+
+TAG_HANDLERS = {"write-boardcfg": writeback_boardcfg}
+
+
+def boot_fs_write(
+    chip: BhChip, boardname_to_display: str, mask: dict, write: bytearray
+) -> bytearray:
+    param_handlers = []
+    for v in mask:
+        tag = v.get("tag", None)
+
+        if tag is None or not isinstance(tag, str):
+            raise TTError(
+                f"Invalid mask format for {boardname_to_display}; expected to see a list of dicts with keys 'tag'"
+            )
+
+        if tag in TAG_HANDLERS:
+            param_handlers.append(TAG_HANDLERS[tag])
+        else:
+            if len(TAG_HANDLERS) > 0:
+                pretty_tags = [f"'{x}'" for x in TAG_HANDLERS.keys()]
+                pretty_tags[-1] = f"or {pretty_tags[-1]}"
+                raise TTError(
+                    f"Invalid tag {tag} for {boardname_to_display}; expected to see one of {pretty_tags}"
+                )
+            else:
+                raise TTError(
+                    f"Invalid tag {tag} for {boardname_to_display}; there aren't any tags defined!"
+                )
+
+    for handler in param_handlers:
+        write = handler(chip, write)
+
+    return write

tt_flash/boot_fs.py

@@ -0,0 +1,148 @@
+# SPDX-FileCopyrightText: © 2024 Tenstorrent AI ULC
+# SPDX-License-Identifier: Apache-2.0
+
+from typing import Callable, Optional, Tuple
+import ctypes
+
+# Define constants
+TT_BOOT_FS_FD_HEAD_ADDR = 0x0
+TT_BOOT_FS_SECURITY_BINARY_FD_ADDR = 0x3FE0
+TT_BOOT_FS_FAILOVER_HEAD_ADDR = 0x4000
+IMAGE_TAG_SIZE = 8
+
+
+class ExtendedStructure(ctypes.Structure):
+    def __eq__(self, other):
+        if not isinstance(other, self.__class__):
+            return False
+        for field in self._fields_:
+            field_name = field[0]
+
+            self_value = getattr(self, field_name)
+            other_value = getattr(other, field_name)
+
+            # Handle comparison for ctypes.Array fields
+            if isinstance(self_value, ctypes.Array):
+                if len(self_value) != len(other_value):
+                    return False
+                for i in range(len(self_value)):
+                    if self_value[i] != other_value[i]:
+                        return False
+            else:
+                if self_value != other_value:
+                    return False
+        return True
+
+    def __ne__(self, other):
+        return not self.__eq__(other)
+
+    def __repr__(self):
+        field_strings = []
+        for field in self._fields_:
+            field_name = field[0]
+
+            field_value = getattr(self, field_name)
+
+            # Handle string representation for ctypes.Array fields
+            if isinstance(field_value, ctypes.Array):
+                array_str = ", ".join(str(x) for x in field_value)
+                field_strings.append(f"{field_name}=[{array_str}]")
+            else:
+                field_strings.append(f"{field_name}={field_value}")
+
+        fields_repr = ", ".join(field_strings)
+        return f"{self.__class__.__name__}({fields_repr})"
+
+
+class ExtendedUnion(ctypes.Union):
+    def __eq__(self, other):
+        for fld in self._fields_:
+            if getattr(self, fld[0]) != getattr(other, fld[0]):
+                return False
+        return True
+
+    def __ne__(self, other):
+        for fld in self._fields_:
+            if getattr(self, fld[0]) != getattr(other, fld[0]):
+                return True
+        return False
+
+    def __repr__(self):
+        field_strings = []
+        for field in self._fields_:
+            field_name = field[0]
+
+            field_value = getattr(self, field_name)
+            field_strings.append(f"{field_name}={field_value}")
+        fields_repr = ", ".join(field_strings)
+        return f"{self.__class__.__name__}({fields_repr})"
+
+
+# Define fd_flags structure
+class fd_flags(ExtendedStructure):
+    _fields_ = [
+        ("image_size", ctypes.c_uint32, 24),
+        ("invalid", ctypes.c_uint32, 1),
+        ("executable", ctypes.c_uint32, 1),
+        ("fd_flags_rsvd", ctypes.c_uint32, 6),
+    ]
+
+
+# Define fd_flags union
+class fd_flags_u(ExtendedUnion):
+    _fields_ = [("val", ctypes.c_uint32), ("f", fd_flags)]
+
+
+# Define security_fd_flags structure
+class security_fd_flags(ExtendedStructure):
+    _fields_ = [
+        ("signature_size", ctypes.c_uint32, 12),
+        ("sb_phase", ctypes.c_uint32, 8),  # 0 - Phase0A, 1 - Phase0B
+    ]
+
+
+# Define security_fd_flags union
+class security_fd_flags_u(ExtendedUnion):
+    _fields_ = [("val", ctypes.c_uint32), ("f", security_fd_flags)]
+
+
+# Define tt_boot_fs_fd structure (File descriptor)
+class tt_boot_fs_fd(ExtendedStructure):
+    _fields_ = [
+        ("spi_addr", ctypes.c_uint32),
+        ("copy_dest", ctypes.c_uint32),
+        ("flags", fd_flags_u),
+        ("data_crc", ctypes.c_uint32),
+        ("security_flags", security_fd_flags_u),
+        ("image_tag", ctypes.c_uint8 * IMAGE_TAG_SIZE),
+        ("fd_crc", ctypes.c_uint32),
+    ]
+
+    def image_tag_str(self):
+        output = ""
+        for c in self.image_tag:
+            if c == "\0":
+                break
+            output += chr(c)
+        return output
+
+
+def read_fd(reader, addr: int) -> tt_boot_fs_fd:
+    fd = reader(addr, ctypes.sizeof(tt_boot_fs_fd))
+    return tt_boot_fs_fd.from_buffer_copy(fd)
+
+
+def read_tag(
+    reader: Callable[[int, int], bytes], tag: str
+) -> Optional[Tuple[int, tt_boot_fs_fd]]:
+    curr_addr = 0
+    while True:
+        fd = read_fd(reader, curr_addr)
+
+        if fd.flags.f.invalid != 0:
+            return None
+
+        if fd.image_tag_str() == tag:
+            return curr_addr, fd
+
+        curr_addr += ctypes.sizeof(tt_boot_fs_fd)

tt_flash/chip.py

@@ -0,0 +1,376 @@
+# SPDX-FileCopyrightText: © 2024 Tenstorrent AI ULC
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+import time
+from typing import Union, Optional
+import sys
+import yaml
+
+from pyluwen import PciChip, Telemetry
+from pyluwen import detect_chips as luwen_detect_chips
+from pyluwen import detect_chips_fallible as luwen_detect_chips_fallible
+
+from tt_flash import utility
+from tt_flash.error import TTError
+
+
+@dataclass
+class FwVersion:
+    allow_exception: bool
+    exception: Optional[Exception]
+    running: Optional[tuple[int, int, int, int]]
+    spi: Optional[tuple[int, int, int, int]]
+
+
+def get_bundle_version_v1(chip: TTChip) -> FwVersion:
+    """
+    Get the currently running bundle version for gs and wh, using a legacy method
+
+    @param chip
+
+    @return the detected fw bundle version.
+    """
+    running_bundle_version = None
+    spi_bundle_version = None
+    exception = None
+
+    try:
+        fw_version = chip.arc_msg(
+            chip.fw_defines["MSG_TYPE_FW_VERSION"], wait_for_done=True, arg0=0, arg1=0
+        )[0]
+
+        # Pre fw version 5 we don't have bundle support
+        # this version of tt-flash only works with bundled fw
+        # so it's safe to assume that we need to update
+        if fw_version >= chip.min_fw_version():
+            temp = chip.arc_msg(
+                chip.fw_defines["MSG_TYPE_FW_VERSION"],
+                wait_for_done=True,
+                arg0=1,
+                arg1=0,
+            )[0]
+
+            if temp not in [0xFFFFFFFF, 0xDEAD]:
+                patch = temp & 0xFF
+                minor = (temp >> 8) & 0xFF
+                major = (temp >> 16) & 0xFF
+                component = (temp >> 24) & 0xFF
+                running_bundle_version = (component, major, minor, patch)
+
+            # There is a version of the firmware that doesn't correctly return an error when setting arg0 to an unknown option.
+            # The running_bundle_version and fw_version can never be the same (as mandated by the version formatting) so I can safely check to see if they are the same when checking for this older FW.
+            if (
+                running_bundle_version != 0xDEAD
+                and fw_version != running_bundle_version
+            ):
+                temp = chip.arc_msg(
+                    chip.fw_defines["MSG_TYPE_FW_VERSION"],
+                    wait_for_done=True,
+                    arg0=2,
+                    arg1=0,
+                )[0]
+
+                if temp not in [0xFFFFFFFF, 0xDEAD]:
+                    patch = temp & 0xFF
+                    minor = (temp >> 8) & 0xFF
+                    major = (temp >> 16) & 0xFF
+                    component = (temp >> 24) & 0xFF
+                    spi_bundle_version = (component, major, minor, patch)
+    except Exception as e:
+        exception = e
+
+    return FwVersion(
+        allow_exception=True,
+        exception=exception,
+        running=running_bundle_version,
+        spi=spi_bundle_version,
+    )
+
+
+def get_chip_data(chip, file, internal: bool):
+    with utility.package_root_path() as path:
+        if isinstance(chip, WhChip):
+            prefix = "wormhole"
+        elif isinstance(chip, GsChip):
+            prefix = "grayskull"
+        elif isinstance(chip, BhChip):
+            prefix = "blackhole"
+        else:
+            raise TTError("Only support flashing Wh or GS chips")
+        if internal:
+            prefix = f".ignored/{prefix}"
+        else:
+            prefix = f"data/{prefix}"
+        return open(str(path.joinpath(f"{prefix}/{file}")))
+
+
+def init_fw_defines(chip):
+    return yaml.safe_load(get_chip_data(chip, "fw_defines.yaml", False))
+
+
+class TTChip:
+    def __init__(self, chip: PciChip):
+        self.luwen_chip = chip
+        self.interface_id = chip.pci_interface_id()
+
+        self.fw_defines = init_fw_defines(self)
+
+        self.telmetry_cache = None
+
+    def reinit(self, callback=None):
+        self.luwen_chip = PciChip(self.interface_id)
+        self.telmetry_cache = None
+
+        chip_count = 0
+        block_count = 0
+        last_draw = time.time()
+
+        def chip_detect_callback(status):
+            nonlocal chip_count, last_draw, block_count
+
+            if status.new_chip():
+                chip_count += 1
+            elif status.correct_down():
+                chip_count -= 1
+            chip_count = max(chip_count, 0)
+
+            if sys.stdout.isatty():
+                current_time = time.time()
+                if current_time - last_draw > 0.1:
+                    last_draw = current_time
+
+                    if block_count > 0:
+                        print(f"\033[{block_count}A", end="", flush=True)
+                        print(f"\033[J", end="", flush=True)
+
+                    print(f"\rDetected Chips: {chip_count}\n", end="", flush=True)
+                    block_count = 1
+
+                    status_string = status.status_string()
+                    if status_string is not None:
+                        for line in status_string.splitlines():
+                            block_count += 1
+                            print(f"\r{line}", flush=True)
+            else:
+                time.sleep(0.01)
+
+        self.luwen_chip.init(
+            callback=chip_detect_callback if callback is None else callback
+        )
+
+    def get_telemetry(self) -> Telemetry:
+        self.telmetry_cache = self.luwen_chip.get_telemetry()
+        return self.telmetry_cache
+
+    def get_telemetry_unchanged(self) -> Telemetry:
+        if self.telmetry_cache is None:
+            self.telmetry_cache = self.luwen_chip.get_telemetry()
+
+        return self.telmetry_cache
+
+    def __vnum_to_version(self, version: int) -> tuple[int, int, int, int]:
+        return (
+            (version >> 24) & 0xFF,
+            (version >> 16) & 0xFF,
+            (version >> 8) & 0xFF,
+            version & 0xFF,
+        )
+
+    def m3_fw_app_version(self):
+        telem = self.get_telemetry_unchanged()
+        return self.__vnum_to_version(telem.m3_app_fw_version)
+
+    def smbus_fw_version(self):
+        telem = self.get_telemetry_unchanged()
+        return self.__vnum_to_version(telem.arc1_fw_version)
+
+    def arc_l2_fw_version(self):
+        telem = self.get_telemetry_unchanged()
+        return self.__vnum_to_version(telem.arc0_fw_version)
+
+    def board_type(self):
+        return self.luwen_chip.pci_board_type()
+
+    def axi_write32(self, addr: int, value: int):
+        self.luwen_chip.axi_write32(addr, value)
+
+    def axi_write(self, addr: int, data: bytes):
+        self.luwen_chip.axi_write(addr, data)
+
+    def axi_read32(self, addr: int) -> int:
+        return self.luwen_chip.axi_read32(addr)
+
+    def axi_read(self, addr: int, size: int) -> bytes:
+        data = bytearray(size)
+        self.luwen_chip.axi_read(addr, data)
+
+        return bytes(data)
+
+    def spi_write(self, addr: int, data: bytes):
+        self.luwen_chip.spi_write(addr, data)
+
+    def spi_read(self, addr: int, size: int) -> bytes:
+        data = bytearray(size)
+        self.luwen_chip.spi_read(addr, data)
+
+        return bytes(data)
+
+    def arc_msg(self, *args, **kwargs):
+        return self.luwen_chip.arc_msg(*args, **kwargs)
+
+    @abstractmethod
+    def min_fw_version(self):
+        pass
+
+    @abstractmethod
+    def get_bundle_version(self) -> FwVersion:
+        pass
+
+
+class BhChip(TTChip):
+    def min_fw_version(self):
+        return 0x0
+
+    def __repr__(self):
+        return f"Blackhole[{self.interface_id}]"
+
+    def get_bundle_version(self) -> FwVersion:
+        running = None
+        spi = None
+        exception = None
+        try:
+            # Read running FW bundle version from telemetry
+            telem = self.get_telemetry_unchanged()
+            temp = telem.fw_bundle_version
+            patch = temp & 0xFF
+            minor = (temp >> 8) & 0xFF
+            major = (temp >> 16) & 0xFF
+            component = (temp >> 24) & 0xFF
+            running = (component, major, minor, patch)
+
+            # Read SPI FW bundle version
+            cmfwcfg = self.luwen_chip.decode_boot_fs_table("cmfwcfg")
+            temp = cmfwcfg["fw_bundle_version"]
+            patch = temp & 0xFF
+            minor = (temp >> 8) & 0xFF
+            major = (temp >> 16) & 0xFF
+            component = (temp >> 24) & 0xFF
+            spi = (component, major, minor, patch)
+        except Exception as e:
+            exception = e
+
+        return FwVersion(
+            allow_exception=True, exception=exception, running=running, spi=spi
+        )
+
+
+class WhChip(TTChip):
+    def min_fw_version(self):
+        return 0x2170000
+
+    def __repr__(self):
+        return f"Wormhole[{self.interface_id}]"
+
+    def get_bundle_version(self) -> FwVersion:
+        return get_bundle_version_v1(self)
+
+
+class GsChip(TTChip):
+    def min_fw_version(self):
+        return 0x1050000
+
+    def __repr__(self):
+        return f"Grayskull[{self.interface_id}]"
+
+    def get_bundle_version(self) -> FwVersion:
+        return get_bundle_version_v1(self)
+
+
+def detect_local_chips(
+    ignore_ethernet: bool = False,
+) -> list[Union[GsChip, WhChip, BhChip]]:
+    """
+    This will create a chip which only gaurentees that you have communication with the chip.
+    """
+
+    chip_count = 0
+    block_count = 0
+    last_draw = time.time()
+    did_draw = False
+
+    def chip_detect_callback(status):
+        nonlocal chip_count, last_draw, block_count, did_draw
+
+        if status.new_chip():
+            chip_count += 1
+        elif status.correct_down():
+            chip_count -= 1
+        chip_count = max(chip_count, 0)
+
+        if sys.stdout.isatty():
+            did_draw = True
+            current_time = time.time()
+            if current_time - last_draw > 0.1:
+                last_draw = current_time
+
+                if block_count > 0:
+                    print(f"\033[{block_count}A", end="", flush=True)
+                    print(f"\033[J", end="", flush=True)
+
+                print(f"\rDetected Chips: {chip_count}\n", end="", flush=True)
+                block_count = 1
+
+                status_string = status.status_string()
+                if status_string is not None:
+                    for line in status_string.splitlines():
+                        block_count += 1
+                        print(f"\r{line}", flush=True)
+        else:
+            time.sleep(0.01)
+
+    output = []
+    for device in luwen_detect_chips_fallible(
+        local_only=True,
+        continue_on_failure=False,
+        callback=chip_detect_callback,
+        noc_safe=ignore_ethernet,
+    ):
+        if not device.have_comms():
+            raise Exception(
+                f"Do not have communication with {device}, you should reset or remove this device from your system before continuing."
+            )
+
+        device = device.force_upgrade()
+
+        if device.as_gs() is not None:
+            output.append(GsChip(device.as_gs()))
+        elif device.as_wh() is not None:
+            output.append(WhChip(device.as_wh()))
+        elif device.as_bh() is not None:
+            output.append(BhChip(device.as_bh()))
+        else:
+            raise ValueError("Did not recognize board")
+
+    if not did_draw:
+        print(f"\tDetected Chips: {chip_count}")
+
+    return output
+
+
+def detect_chips(local_only: bool = False) -> list[Union[GsChip, WhChip, BhChip]]:
+    output = []
+    for device in luwen_detect_chips(local_only=local_only):
+        if device.as_gs() is not None:
+            output.append(GsChip(device.as_gs()))
+        elif device.as_wh() is not None:
+            output.append(WhChip(device.as_wh()))
+        elif device.as_bh() is not None:
+            output.append(BhChip(device.as_bh()))
+        else:
+            raise ValueError("Did not recognize board")
+
+    return output

tt_flash/data/blackhole/fw_defines.yaml

@@ -0,0 +1,5 @@
+# SPDX-FileCopyrightText: © 2024 Tenstorrent AI ULC
+# SPDX-License-Identifier: Apache-2.0
+
+MSG_TYPE_ARC_STATE0:                      0xa0
+MSG_TYPE_ARC_STATE3:                      0xa3

tt_flash/data/grayskull/fw_defines.yaml

@@ -0,0 +1,6 @@
+# SPDX-FileCopyrightText: © 2024 Tenstorrent AI ULC
+# SPDX-License-Identifier: Apache-2.0
+
+MSG_TYPE_ARC_STATE0:                      0xa0
+MSG_TYPE_ARC_STATE3:                      0xa3
+MSG_TYPE_FW_VERSION:                      0xb9

tt_flash/data/wormhole/fw_defines.yaml

@@ -0,0 +1,7 @@
+# SPDX-FileCopyrightText: © 2024 Tenstorrent AI ULC
+# SPDX-License-Identifier: Apache-2.0
+
+MSG_TYPE_ARC_STATE0:                      0xa0
+MSG_TYPE_ARC_STATE3:                      0xa3
+MSG_TYPE_FW_VERSION:                      0xb9
+MSG_TRIGGER_SPI_COPY_LtoR:                0x50

tt_flash/error.py

@@ -0,0 +1,8 @@
+# SPDX-FileCopyrightText: © 2024 Tenstorrent AI ULC
+# SPDX-License-Identifier: Apache-2.0
+
+
+class TTError(Exception):
+    """Base class for exceptions in this module."""
+
+    pass