TheRock/rocm-systems/projects/rocprof-trace-decoder/python/rocprof_trace_decoder/bindings.py

from __future__ import annotations

import ctypes
import importlib
import importlib.util
import os
from pathlib import Path
from typing import Callable, Iterable, Protocol, Union

from ._version import __version__ as _PACKAGE_VERSION
from .records import (
    DecoderStatus,
    Dispatch,
    Event,
    EventPayload,
    Instruction,
    Occupancy,
    OtherSimdInstruction,
    Pc,
    PerfEvent,
    Realtime,
    RecordType,
    ShaderData,
    TraceRecords,
    Wave,
    WaveState,
)

BytesLike = Union[bytes, bytearray, memoryview]


class DecoderError(RuntimeError):
    def __init__(self, status: int, message: str | None = None):
        self.status = int(status)
        super().__init__(message or f"rocprof-trace-decoder returned status {status}")


class IsaProvider(Protocol):
    def isa_for_pc(self, pc: Pc) -> tuple[str, int] | None:
        """Return (instruction text, instruction byte size) for *pc*."""


class _PcInfo(ctypes.Structure):
    _fields_ = [
        ("address", ctypes.c_uint64),
        ("code_object_id", ctypes.c_uint64),
    ]


class _PerfEvent(ctypes.Structure):
    _fields_ = [
        ("time", ctypes.c_int64),
        ("events0", ctypes.c_uint16),
        ("events1", ctypes.c_uint16),
        ("events2", ctypes.c_uint16),
        ("events3", ctypes.c_uint16),
        ("CU", ctypes.c_uint8),
        ("bank", ctypes.c_uint8),
    ]


class _Occupancy(ctypes.Structure):
    _fields_ = [
        ("pc", _PcInfo),
        ("time", ctypes.c_uint64),
        ("reserved", ctypes.c_uint8),
        ("cu", ctypes.c_uint8),
        ("simd", ctypes.c_uint8),
        ("wave_id", ctypes.c_uint8),
        ("_bits", ctypes.c_uint32),
    ]


class _WaveState(ctypes.Structure):
    _fields_ = [("type", ctypes.c_int32), ("duration", ctypes.c_int32)]


class _Instruction(ctypes.Structure):
    _fields_ = [
        ("_cat_stall", ctypes.c_uint32),
        ("duration", ctypes.c_int32),
        ("time", ctypes.c_int64),
        ("pc", _PcInfo),
    ]


class _Wave(ctypes.Structure):
    _fields_ = [
        ("cu", ctypes.c_uint8),
        ("simd", ctypes.c_uint8),
        ("wave_id", ctypes.c_uint8),
        ("contexts", ctypes.c_uint8),
        ("dispatcher", ctypes.c_uint8),
        ("workgroup_id", ctypes.c_uint8),
        ("cluster_id", ctypes.c_uint8),
        ("reserved", ctypes.c_uint8),
        ("size", ctypes.c_uint64),
        ("begin_time", ctypes.c_int64),
        ("end_time", ctypes.c_int64),
        ("timeline_size", ctypes.c_uint64),
        ("instructions_size", ctypes.c_uint64),
        ("timeline_array", ctypes.POINTER(_WaveState)),
        ("instructions_array", ctypes.POINTER(_Instruction)),
    ]


class _Realtime(ctypes.Structure):
    _fields_ = [
        ("shader_clock", ctypes.c_int64),
        ("realtime_clock", ctypes.c_uint64),
        ("reserved", ctypes.c_uint64),
    ]


class _ShaderData(ctypes.Structure):
    _fields_ = [
        ("time", ctypes.c_int64),
        ("value", ctypes.c_uint64),
        ("cu", ctypes.c_uint8),
        ("simd", ctypes.c_uint8),
        ("wave_id", ctypes.c_uint8),
        ("flags", ctypes.c_uint8),
        ("reserved", ctypes.c_uint32),
    ]


class _OtherSimdInstruction(ctypes.Structure):
    _fields_ = [
        ("size", ctypes.c_uint64),
        ("time", ctypes.c_int64),
        ("cycles", ctypes.c_uint16),
        ("wgp", ctypes.c_uint8),
        ("category", ctypes.c_uint8),
    ]


class _ClusterBarrier(ctypes.Structure):
    _fields_ = [("cluster_id", ctypes.c_int32), ("barrier_id", ctypes.c_int32)]


class _EventPayload(ctypes.Union):
    _fields_ = [
        ("raw", ctypes.c_uint64),
        ("code_object_id", ctypes.c_uint64),
        ("cluster_barrier", _ClusterBarrier),
    ]


class _Event(ctypes.Structure):
    _fields_ = [
        ("size", ctypes.c_uint64),
        ("time", ctypes.c_int64),
        ("type", ctypes.c_int),
        ("me_id", ctypes.c_uint8),
        ("pipe_id", ctypes.c_uint8),
        ("flags", ctypes.c_uint16),
        ("payload", _EventPayload),
        ("byte_offset", ctypes.c_uint64),
    ]


class _Dispatch(ctypes.Structure):
    _fields_ = [
        ("size", ctypes.c_uint64),
        ("time", ctypes.c_int64),
        ("me_id", ctypes.c_uint8),
        ("pipe_id", ctypes.c_uint8),
        ("user_sgprs", ctypes.c_uint16),
        ("flags", ctypes.c_int),
        ("vgprs", ctypes.c_uint32),
        ("sgprs", ctypes.c_uint32),
        ("lds_size", ctypes.c_uint32),
        ("thread_dim_x", ctypes.c_uint32),
        ("thread_dim_y", ctypes.c_uint32),
        ("thread_dim_z", ctypes.c_uint32),
        ("dispatch_pkt_addr", ctypes.c_uint64),
        ("byte_offset", ctypes.c_uint64),
        ("entry_point", _PcInfo),
    ]


class _Handle(ctypes.Structure):
    _fields_ = [("handle", ctypes.c_uint64)]


_TRACE_CB = ctypes.CFUNCTYPE(
    ctypes.c_int,
    ctypes.c_int,
    ctypes.c_void_p,
    ctypes.c_uint64,
    ctypes.c_void_p,
)

_ISA_CB = ctypes.CFUNCTYPE(
    ctypes.c_int,
    ctypes.c_void_p,
    ctypes.POINTER(ctypes.c_uint64),
    ctypes.POINTER(ctypes.c_uint64),
    _PcInfo,
    ctypes.c_void_p,
)

_SE_DATA_CB = ctypes.CFUNCTYPE(
    ctypes.c_uint64,
    ctypes.POINTER(ctypes.POINTER(ctypes.c_uint8)),
    ctypes.POINTER(ctypes.c_uint64),
    ctypes.c_void_p,
)


def _pc(c_pc: _PcInfo) -> Pc:
    return Pc(address=int(c_pc.address), code_object_id=int(c_pc.code_object_id))


def _find_library(explicit: str | os.PathLike[str] | None = None) -> str:
    if explicit:
        return str(explicit)

    env = os.environ.get("ROCPROF_TRACE_DECODER_LIB")
    if env:
        return env

    names = _library_names()
    for root in _library_roots():
        for name in names:
            candidate = root / name
            if candidate.exists():
                return str(candidate)
        for candidate in sorted(root.glob("librocprof-trace-decoder.so.*"), reverse=True):
            if candidate.is_file():
                return str(candidate)

    # Let the platform loader try its normal search path.
    return "librocprof-trace-decoder.so"


def _library_names() -> list[str]:
    names = ["librocprof-trace-decoder.so", "librocprof-trace-decoder.dylib"]
    version_parts = _PACKAGE_VERSION.split(".")
    if len(version_parts) >= 2:
        names.append(f"librocprof-trace-decoder.so.{version_parts[0]}.{version_parts[1]}")
    return names


def _library_roots() -> list[Path]:
    roots: list[Path] = []
    seen: set[Path] = set()

    def add(path: Path) -> None:
        if path not in seen:
            roots.append(path)
            seen.add(path)

    for raw in os.environ.get("LD_LIBRARY_PATH", "").split(os.pathsep):
        if raw:
            add(Path(raw).expanduser())

    for var in ("ROCM_HOME", "ROCM_PATH"):
        value = os.environ.get(var)
        if value:
            add(Path(value).expanduser() / "lib")

    for root in _therock_package_roots():
        add(root / "lib")

    add(Path("/opt/rocm/lib"))
    return roots


def _therock_package_roots() -> list[Path]:
    roots: list[Path] = []
    seen: set[Path] = set()

    def add_package_root(package_name: str) -> None:
        spec = importlib.util.find_spec(package_name)
        if spec is None or spec.origin is None:
            return
        root = Path(spec.origin).resolve().parent
        if root not in seen:
            roots.append(root)
            seen.add(root)

    for dist_info_name in (
        "rocm_sdk_core._dist_info",
        "rocm_profiler._dist_info",
        "rocm_sdk._dist_info",
    ):
        try:
            dist_info = importlib.import_module(dist_info_name)
        except Exception:
            continue
        for logical_name in ("core", "devel", "profiler"):
            entry = getattr(dist_info, "ALL_PACKAGES", {}).get(logical_name)
            if entry is None or getattr(entry, "is_target_specific", False):
                continue
            try:
                add_package_root(entry.get_py_package_name())
            except Exception:
                continue

    for package_name in ("_rocm_sdk_core", "_rocm_sdk_devel", "_rocm_profiler"):
        add_package_root(package_name)

    return roots


class Decoder:
    def __init__(self, lib_path: str | os.PathLike[str] | None = None):
        self.lib_path = _find_library(lib_path)
        self._lib = ctypes.CDLL(self.lib_path)
        self._configure_library()
        self._handle = _Handle()
        self._closed = False
        self._isa_cb: _ISA_CB | None = None
        self._trace_cb: _TRACE_CB | None = None

        status = self._lib.rocprof_trace_decoder_create_handle(ctypes.byref(self._handle))
        self._check(status)

    def __enter__(self) -> "Decoder":
        return self

    def __exit__(self, *_exc: object) -> None:
        self.close()

    def close(self) -> None:
        if self._closed:
            return
        status = self._lib.rocprof_trace_decoder_destroy_handle(self._handle)
        self._check(status)
        self._closed = True
        self._isa_cb = None
        self._trace_cb = None

    def _ensure_open(self) -> None:
        if self._closed:
            raise DecoderError(DecoderStatus.ERROR_INVALID_ARGUMENT, "Decoder handle is closed")

    def _configure_library(self) -> None:
        lib = self._lib
        lib.rocprof_trace_decoder_create_handle.argtypes = [ctypes.POINTER(_Handle)]
        lib.rocprof_trace_decoder_create_handle.restype = ctypes.c_int

        lib.rocprof_trace_decoder_destroy_handle.argtypes = [_Handle]
        lib.rocprof_trace_decoder_destroy_handle.restype = ctypes.c_int

        lib.rocprof_trace_decoder_codeobj_load.argtypes = [
            _Handle,
            ctypes.c_uint64,
            ctypes.c_uint64,
            ctypes.c_uint64,
            ctypes.c_void_p,
            ctypes.c_uint64,
        ]
        lib.rocprof_trace_decoder_codeobj_load.restype = ctypes.c_int

        lib.rocprof_trace_decoder_codeobj_unload.argtypes = [_Handle, ctypes.c_uint64]
        lib.rocprof_trace_decoder_codeobj_unload.restype = ctypes.c_int

        lib.rocprof_trace_decoder_set_isa_callback.argtypes = [_Handle, _ISA_CB, ctypes.c_void_p]
        lib.rocprof_trace_decoder_set_isa_callback.restype = ctypes.c_int

        lib.rocprof_trace_decoder_set_se_data_callback.argtypes = [
            _Handle,
            _SE_DATA_CB,
            ctypes.c_void_p,
        ]
        lib.rocprof_trace_decoder_set_se_data_callback.restype = ctypes.c_int

        lib.rocprof_trace_decoder_parse.argtypes = [
            _Handle,
            ctypes.c_void_p,
            ctypes.c_uint64,
            _TRACE_CB,
            ctypes.c_void_p,
        ]
        lib.rocprof_trace_decoder_parse.restype = ctypes.c_int

        lib.rocprof_trace_decoder_get_info_string.argtypes = [ctypes.c_int]
        lib.rocprof_trace_decoder_get_info_string.restype = ctypes.c_char_p

        lib.rocprof_trace_decoder_get_status_string.argtypes = [ctypes.c_int]
        lib.rocprof_trace_decoder_get_status_string.restype = ctypes.c_char_p

    def _check(self, status: int) -> None:
        if int(status) == int(DecoderStatus.SUCCESS):
            return
        message = None
        try:
            raw = self._lib.rocprof_trace_decoder_get_status_string(int(status))
            if raw:
                message = raw.decode("utf-8", errors="replace")
        except Exception:
            pass
        raise DecoderError(status, message)

    def info_string(self, info: int) -> str:
        raw = self._lib.rocprof_trace_decoder_get_info_string(int(info))
        return raw.decode("utf-8", errors="replace") if raw else ""

    def status_string(self, status: int) -> str:
        raw = self._lib.rocprof_trace_decoder_get_status_string(int(status))
        return raw.decode("utf-8", errors="replace") if raw else ""

    def load_code_object_data(
        self,
        load_id: int,
        load_addr: int,
        load_size: int,
        data: BytesLike,
    ) -> None:
        """Compatibility wrapper for the original data-loading argument order.

        *load_size* is the loaded memory range size, not the code-object byte
        length.
        """
        self._load_code_object_bytes(
            data,
            load_id=load_id,
            load_addr=load_addr,
            load_size=load_size,
        )

    def load_code_object(
        self,
        data: BytesLike | str | os.PathLike[str],
        load_id: int,
        load_addr: int,
        load_size: int,
    ) -> None:
        """Load a code object into this decoder handle.

        *data* may be a bytes-like object or a filesystem path. New code that
        wants to be explicit about paths can use `load_code_object_file()`.
        *load_size* is the size of the loaded memory range and is intentionally
        separate from the code-object byte length passed to the native API.
        """
        if _is_bytes_like(data):
            self._load_code_object_bytes(
                data,
                load_id=load_id,
                load_addr=load_addr,
                load_size=load_size,
            )
        else:
            self.load_code_object_file(
                data,
                load_id=load_id,
                load_addr=load_addr,
                load_size=load_size,
            )

    def load_code_object_file(
        self,
        path: str | os.PathLike[str],
        load_id: int,
        load_addr: int,
        load_size: int,
    ) -> None:
        """Load a code object from a file path.

        *load_size* is the loaded memory range size, not the file length.
        """
        self._load_code_object_bytes(
            Path(path).read_bytes(),
            load_id=load_id,
            load_addr=load_addr,
            load_size=load_size,
        )

    def _load_code_object_bytes(
        self,
        data: BytesLike,
        *,
        load_id: int,
        load_addr: int,
        load_size: int,
    ) -> None:
        self._ensure_open()
        blob = bytes(data)
        buf = ctypes.create_string_buffer(blob)
        status = self._lib.rocprof_trace_decoder_codeobj_load(
            self._handle, load_id, load_addr, load_size, buf, len(blob)
        )
        self._check(status)

    def unload_code_object(self, load_id: int) -> None:
        self._ensure_open()
        status = self._lib.rocprof_trace_decoder_codeobj_unload(self._handle, load_id)
        self._check(status)

    def parse_file(
        self,
        path: str | os.PathLike[str],
        isa: IsaProvider | None = None,
        on_batch: Callable[[RecordType, list[object] | int], None] | None = None,
    ) -> TraceRecords:
        return self.parse(Path(path).read_bytes(), isa=isa, on_batch=on_batch)

    def parse_chunks(
        self,
        chunks: Iterable[BytesLike],
        isa: IsaProvider | None = None,
        on_batch: Callable[[RecordType, list[object] | int], None] | None = None,
    ) -> TraceRecords:
        out = TraceRecords()
        for chunk in chunks:
            partial = self.parse(chunk, isa=isa, on_batch=on_batch)
            merge_records(out, partial)
        return out

    def parse(
        self,
        data: BytesLike,
        isa: IsaProvider | None = None,
        on_batch: Callable[[RecordType, list[object] | int], None] | None = None,
    ) -> TraceRecords:
        """Decode a binary ATT payload into Python records.

        *isa* is called by the native decoder when it needs instruction text
        for a PC. *on_batch*, when provided, receives each converted native
        record batch before this method returns the accumulated TraceRecords.
        """
        self._ensure_open()

        raw = bytes(data)
        records = TraceRecords()
        callback_error: list[BaseException] = []

        if isa is not None:
            self._isa_cb = self._make_isa_callback(isa)
        else:
            self._isa_cb = _ISA_CB()
        self._check(
            self._lib.rocprof_trace_decoder_set_isa_callback(self._handle, self._isa_cb, None)
        )

        def _trace(record_type: int, events: int | None, size: int, _userdata: int | None) -> int:
            try:
                batch_type = RecordType(record_type)
                batch = self._convert_batch(batch_type, events, int(size))
                _append_batch(records, batch_type, batch)
                if on_batch is not None:
                    on_batch(batch_type, batch)
                return int(DecoderStatus.SUCCESS)
            except BaseException as exc:
                callback_error.append(exc)
                return int(DecoderStatus.ERROR)

        trace_cb = _TRACE_CB(_trace)
        self._trace_cb = trace_cb

        if raw:
            data_buf = (ctypes.c_uint8 * len(raw)).from_buffer_copy(raw)
            data_ptr = ctypes.cast(data_buf, ctypes.c_void_p)
            data_size = len(raw)
        else:
            data_buf = None
            data_ptr = None
            data_size = 0

        status = self._lib.rocprof_trace_decoder_parse(
            self._handle, data_ptr, data_size, trace_cb, None
        )
        if callback_error:
            raise callback_error[0]
        self._check(status)
        return records

    def _make_isa_callback(self, isa: IsaProvider) -> _ISA_CB:
        def _isa(
            instr_buf: int,
            mem_size_p: ctypes.POINTER(ctypes.c_uint64),
            size_p: ctypes.POINTER(ctypes.c_uint64),
            c_pc: _PcInfo,
            _userdata: int | None,
        ) -> int:
            pc = _pc(c_pc)
            resolved = isa.isa_for_pc(pc)
            if resolved is None:
                return int(DecoderStatus.ERROR_INVALID_ARGUMENT)

            text, memory_size = resolved
            encoded = text.encode("utf-8")
            available = int(size_p[0])
            size_p[0] = len(encoded)
            mem_size_p[0] = max(int(memory_size), 4)
            if len(encoded) > available:
                return int(DecoderStatus.ERROR_OUT_OF_RESOURCES)
            if encoded:
                ctypes.memmove(instr_buf, encoded, len(encoded))
            return int(DecoderStatus.SUCCESS)

        return _ISA_CB(_isa)

    def _convert_batch(
        self, record_type: RecordType, events: int | None, size: int
    ) -> list[object] | int:
        if record_type == RecordType.GFXIP:
            return int(events or 0)
        if size <= 0:
            return []
        if not events:
            return []

        if record_type == RecordType.INFO:
            ptr = ctypes.cast(events, ctypes.POINTER(ctypes.c_int))
            return [int(ptr[i]) for i in range(size)]
        if record_type == RecordType.OCCUPANCY:
            ptr = ctypes.cast(events, ctypes.POINTER(_Occupancy))
            return [_convert_occupancy(ptr[i]) for i in range(size)]
        if record_type == RecordType.PERFEVENT:
            ptr = ctypes.cast(events, ctypes.POINTER(_PerfEvent))
            return [_convert_perf(ptr[i]) for i in range(size)]
        if record_type == RecordType.WAVE:
            ptr = ctypes.cast(events, ctypes.POINTER(_Wave))
            return [_convert_wave(ptr[i]) for i in range(size)]
        if record_type == RecordType.EVENT:
            ptr = ctypes.cast(events, ctypes.POINTER(_Event))
            return [_convert_event(ptr[i]) for i in range(size)]
        if record_type == RecordType.SHADERDATA:
            ptr = ctypes.cast(events, ctypes.POINTER(_ShaderData))
            return [_convert_shaderdata(ptr[i]) for i in range(size)]
        if record_type == RecordType.REALTIME:
            ptr = ctypes.cast(events, ctypes.POINTER(_Realtime))
            return [_convert_realtime(ptr[i]) for i in range(size)]
        if record_type == RecordType.RT_FREQUENCY:
            ptr = ctypes.cast(events, ctypes.POINTER(ctypes.c_uint64))
            return int(ptr[0])
        if record_type == RecordType.INST_OTHER_SIMD:
            ptr = ctypes.cast(events, ctypes.POINTER(_OtherSimdInstruction))
            return [_convert_other_simd(ptr[i]) for i in range(size)]
        if record_type == RecordType.DISPATCH:
            ptr = ctypes.cast(events, ctypes.POINTER(_Dispatch))
            return [_convert_dispatch(ptr[i]) for i in range(size)]
        return []


def _convert_perf(c: _PerfEvent) -> PerfEvent:
    return PerfEvent(
        time=int(c.time),
        events0=int(c.events0),
        events1=int(c.events1),
        events2=int(c.events2),
        events3=int(c.events3),
        cu=int(c.CU),
        bank=int(c.bank),
    )


def _convert_occupancy(c: _Occupancy) -> Occupancy:
    bits = int(c._bits)
    return Occupancy(
        pc=_pc(c.pc),
        time=int(c.time),
        reserved=int(c.reserved),
        cu=int(c.cu),
        simd=int(c.simd),
        wave_id=int(c.wave_id),
        start=bits & 0x1,
        me_id=(bits >> 1) & 0x7,
        pipe_id=(bits >> 4) & 0xF,
        is_ext=(bits >> 8) & 0x1,
        workgroup_id=(bits >> 9) & 0x7F,
        cluster_id=(bits >> 16) & 0x1F,
    )


def _convert_instruction(c: _Instruction) -> Instruction:
    raw = int(c._cat_stall)
    return Instruction(
        category=raw & 0xFF,
        stall=(raw >> 8) & 0xFFFFFF,
        duration=int(c.duration),
        time=int(c.time),
        pc=_pc(c.pc),
    )


def _convert_wave(c: _Wave) -> Wave:
    timeline = []
    if c.timeline_array:
        timeline = [
            WaveState(
                type=int(c.timeline_array[i].type),
                duration=int(c.timeline_array[i].duration),
            )
            for i in range(int(c.timeline_size))
        ]

    instructions = []
    if c.instructions_array:
        instructions = [
            _convert_instruction(c.instructions_array[i]) for i in range(int(c.instructions_size))
        ]

    return Wave(
        cu=int(c.cu),
        simd=int(c.simd),
        wave_id=int(c.wave_id),
        contexts=int(c.contexts),
        dispatcher=int(c.dispatcher),
        workgroup_id=int(c.workgroup_id),
        cluster_id=int(c.cluster_id),
        reserved=int(c.reserved),
        size=int(c.size),
        begin_time=int(c.begin_time),
        end_time=int(c.end_time),
        timeline=timeline,
        instructions=instructions,
    )


def _convert_realtime(c: _Realtime) -> Realtime:
    return Realtime(
        shader_clock=int(c.shader_clock),
        realtime_clock=int(c.realtime_clock),
        reserved=int(c.reserved),
    )


def _convert_shaderdata(c: _ShaderData) -> ShaderData:
    return ShaderData(
        time=int(c.time),
        value=int(c.value),
        cu=int(c.cu),
        simd=int(c.simd),
        wave_id=int(c.wave_id),
        flags=int(c.flags),
        reserved=int(c.reserved),
    )


def _convert_other_simd(c: _OtherSimdInstruction) -> OtherSimdInstruction:
    return OtherSimdInstruction(
        size=int(c.size),
        time=int(c.time),
        cycles=int(c.cycles),
        wgp=int(c.wgp),
        category=int(c.category),
    )


def _convert_event(c: _Event) -> Event:
    payload = EventPayload(
        raw=int(c.payload.raw),
        code_object_id=int(c.payload.code_object_id),
        cluster_id=int(c.payload.cluster_barrier.cluster_id),
        barrier_id=int(c.payload.cluster_barrier.barrier_id),
    )
    return Event(
        size=int(c.size),
        time=int(c.time),
        type=int(c.type),
        me_id=int(c.me_id),
        pipe_id=int(c.pipe_id),
        flags=int(c.flags),
        payload=payload,
        byte_offset=int(c.byte_offset),
    )


def _convert_dispatch(c: _Dispatch) -> Dispatch:
    return Dispatch(
        size=int(c.size),
        time=int(c.time),
        me_id=int(c.me_id),
        pipe_id=int(c.pipe_id),
        user_sgprs=int(c.user_sgprs),
        flags=int(c.flags),
        vgprs=int(c.vgprs),
        sgprs=int(c.sgprs),
        lds_size=int(c.lds_size),
        thread_dim_x=int(c.thread_dim_x),
        thread_dim_y=int(c.thread_dim_y),
        thread_dim_z=int(c.thread_dim_z),
        dispatch_pkt_addr=int(c.dispatch_pkt_addr),
        byte_offset=int(c.byte_offset),
        entry_point=_pc(c.entry_point),
    )


def _append_batch(
    records: TraceRecords,
    record_type: RecordType,
    batch: list[object] | int,
) -> None:
    records.batches.append((record_type, batch))
    if record_type == RecordType.GFXIP:
        records.gfxip = int(batch)
    elif record_type == RecordType.INFO:
        records.info.extend(batch)  # type: ignore[arg-type]
    elif record_type == RecordType.OCCUPANCY:
        records.occupancy.extend(batch)  # type: ignore[arg-type]
    elif record_type == RecordType.PERFEVENT:
        records.perf_events.extend(batch)  # type: ignore[arg-type]
    elif record_type == RecordType.WAVE:
        records.waves.extend(batch)  # type: ignore[arg-type]
    elif record_type == RecordType.EVENT:
        records.events.extend(batch)  # type: ignore[arg-type]
    elif record_type == RecordType.SHADERDATA:
        records.shaderdata.extend(batch)  # type: ignore[arg-type]
    elif record_type == RecordType.REALTIME:
        records.realtime.extend(batch)  # type: ignore[arg-type]
    elif record_type == RecordType.RT_FREQUENCY:
        records.realtime_frequency = int(batch)
    elif record_type == RecordType.INST_OTHER_SIMD:
        records.other_simd.extend(batch)  # type: ignore[arg-type]
    elif record_type == RecordType.DISPATCH:
        records.dispatches.extend(batch)  # type: ignore[arg-type]


def merge_records(dst: TraceRecords, src: TraceRecords) -> TraceRecords:
    if src.gfxip is not None:
        dst.gfxip = src.gfxip
    if src.realtime_frequency is not None:
        dst.realtime_frequency = src.realtime_frequency
    dst.info.extend(src.info)
    dst.occupancy.extend(src.occupancy)
    dst.perf_events.extend(src.perf_events)
    dst.waves.extend(src.waves)
    dst.events.extend(src.events)
    dst.shaderdata.extend(src.shaderdata)
    dst.realtime.extend(src.realtime)
    dst.other_simd.extend(src.other_simd)
    dst.dispatches.extend(src.dispatches)
    dst.batches.extend(src.batches)
    return dst


def _is_bytes_like(value: object) -> bool:
    return isinstance(value, (bytes, bytearray, memoryview))


__all__ = ["BytesLike", "Decoder", "DecoderError", "IsaProvider", "merge_records"]