diff --git a/gpu/mse.py b/gpu/mse.py index 7bd6f13..8520616 100644 --- a/gpu/mse.py +++ b/gpu/mse.py @@ -21,11 +21,13 @@ # DEALINGS IN THE SOFTWARE. # +import atexit +from logging import debug + from utils import NiceStruct + from .mnoc import GpuMnoc -from logging import debug -import atexit class MseHeader(NiceStruct): _fields_ = [ @@ -107,7 +109,7 @@ def send_cmd(self, cmd_class, cmd_opcode, cmd_data, reset=False): mse_header.ssid = 9 mse_header.dsid = 4 - data_to_send = mse_header.to_int_array() + cmd_data + data_to_send = mse_header.to_int_list() + cmd_data self.mnoc.send_data(data_to_send) @@ -118,7 +120,7 @@ def process_incoming(self): while True: resp_data = self.mnoc.receive_data() - mse_header.from_int_array(resp_data[:4]) + mse_header.from_ints(resp_data[:4]) if mse_header.is_response: return resp_data[4:] else: @@ -166,6 +168,5 @@ def goodbye(self): def get_platform_info(self): platform_info_raw = self.send_cmd(2, 0x30, []) platform_info = GetPlatformInfoRsp() - platform_info.from_int_array(platform_info_raw) + platform_info.from_ints(platform_info_raw) return platform_info - diff --git a/nvidia_gpu_tools.py b/nvidia_gpu_tools.py index 43fdcfd..c59694a 100755 --- a/nvidia_gpu_tools.py +++ b/nvidia_gpu_tools.py @@ -24,6 +24,7 @@ # from __future__ import print_function +import array import collections import time import sys @@ -32,7 +33,7 @@ from pathlib import Path from utils import platform_config -from utils import data_from_int, ints_from_bytearray, read_ints_from_path +from utils import data_from_int, array_view_from_bytearray, read_ints_from_path from utils import formatted_tuple_from_data from gpu.defines import * from pci.defines import * @@ -3040,14 +3041,13 @@ def is_driver_loaded(self): return False def dump_bar0(self): - bar0_data = bytearray() + bar0_data_array = array.array('I') for offset in range(0, self.bar0_size, 4): if offset % (128 * 1024) == 0: debug("Dumped %d bytes so far", offset) data = self.bar0.read32(offset) - bar0_data.extend(data_from_int(data, 4)) - - return bar0_data + bar0_data_array.append(data) + return memoryview(bar0_data_array.tobytes()).toreadonly() def flr_resettable_scratch(self): return 0xdfe0 diff --git a/pytest.ini b/pytest.ini new file mode 100644 index 0000000..35c4f10 --- /dev/null +++ b/pytest.ini @@ -0,0 +1,3 @@ +[pytest] +markers = + performance: Lightweight throughput checks for core utils diff --git a/tests/test_ints_helpers.py b/tests/test_ints_helpers.py new file mode 100644 index 0000000..ef074d5 --- /dev/null +++ b/tests/test_ints_helpers.py @@ -0,0 +1,102 @@ +import array +import struct +import time +from collections import namedtuple + +import pytest + +from utils.formatted_tuple import FormattedTuple +from utils.ints_to_bytes import ( + array_view_from_bytearray, + bytearray_view_from_ints, + data_from_int, + int_from_data, + ints_from_data, +) +from utils.nice_struct import NiceStruct + + +class DemoStruct(NiceStruct): + name = "DemoStruct" + _fields_ = [ + ("a", "I"), + ("b", "I"), + ] + + +class DemoFormattedTuple(FormattedTuple): + namedtuple = namedtuple("DemoTuple", ["value"]) + struct = struct.Struct(" None: + if size not in _INT_SIZE_TO_FORMAT: + raise AssertionError(f"Unhandled size {size}") + + +def _byte_view(data: Any) -> memoryview: + if isinstance(data, memoryview): + view = data + else: + try: + view = memoryview(data) + except TypeError: + view = memoryview(bytes(data)) + + if not view.contiguous: + view = memoryview(bytes(view)) + + if view.format != "B": + try: + view = view.cast("B") + except TypeError: + view = memoryview(bytes(view)) + + return view + + +def ints_from_data(data: Any, size: int) -> list[int]: + _require_int_size(size) + view = _byte_view(data) + + if view.nbytes % size != 0: + raise struct.error(f"unpack requires a buffer of {size} bytes") + + if size == 1: + return list(view) + + return view.cast(_INT_SIZE_TO_FORMAT[size]).tolist() + + +def int_from_data(data: Any, size: int) -> int: + _require_int_size(size) + + if isinstance(data, (bytes, bytearray)): + if len(data) != size: + raise struct.error(f"unpack requires a buffer of {size} bytes") + if size == 1: + return data[0] + return _STRUCT_BY_SIZE[size].unpack(data)[0] + + view = _byte_view(data) + + if view.nbytes != size: + raise struct.error(f"unpack requires a buffer of {size} bytes") + + if size == 1: + return view[0] + + return _STRUCT_BY_SIZE[size].unpack_from(view)[0] + + +def data_from_int(integer: int, size: int = 4) -> bytes: + _require_int_size(size) + return integer.to_bytes(size, byteorder=sys.byteorder) + + +def bytearray_view_from_ints(int_array: Iterable[int], type_code: str = "I") -> memoryview: + arr = array.array(type_code, int_array) + return memoryview(arr).cast("B").toreadonly() + + +def array_view_from_bytearray(ba: Any, type_code: str = "I") -> memoryview: + arr = array.array(type_code) + arr.frombytes(_byte_view(ba)) + return memoryview(arr).toreadonly() + + +def read_ints_from_path(path: str, offset: int, int_size: int, int_num: int = -1) -> list[int]: + with open(path, "rb") as f: f.seek(offset, 0) if int_num == -1: size = -1 diff --git a/utils/nice_struct.py b/utils/nice_struct.py index 82359ba..2bdd8b7 100644 --- a/utils/nice_struct.py +++ b/utils/nice_struct.py @@ -22,8 +22,11 @@ # import struct +from typing import Iterable + from .ints_to_bytes import * + class NiceStructMeta(type): def __init__(cls, name, bases, attrs): super().__init__(name, bases, attrs) @@ -187,21 +190,20 @@ def to_bytes(self): return struct.pack(self.fmt_string, *packed_values) - def from_int_array(self, ints): - self.from_bytes(bytearray_from_ints(ints)) + def from_ints(self, ints: Iterable[int]) -> None: + self.from_bytes(bytearray_view_from_ints(ints)) - def to_int_array(self, int_size=4): - as_bytes = self.to_bytes() - return ints_from_bytearray(as_bytes, int_size) + def to_int_list(self, int_size: int = 4) -> list[int]: + return ints_from_data(self.to_bytes(), int_size) - def to_int(self, int_size=4): - int_array = self.to_int_array(int_size) - if len(int_array) != 1: + def to_int(self, int_size: int = 4) -> int: + int_list = self.to_int_list(int_size) + if len(int_list) != 1: raise ValueError(f"{self.name} doesn't fit in a single int with {int_size} bytes") - return int_array[0] + return int_list[0] - def from_int(self, integer): - self.from_int_array([integer]) + def from_int(self, integer: int) -> None: + self.from_ints([integer]) class NiceStructArray: def __init__(self, struct_class, count):