mempref_cli.py

import ctypes
import sys, os
import datetime, time
from enum import Enum


class NewCPUModels(Enum):
    TGL = (6, 140)  # Tiger Lake
    ADL = (6, 151)  # Alder Lake
    RPL = (6, 183)  # Raptor Lake
    MTL = (6, 170)  # Meteor Lake
    LNL = (6, 189)  # Lunar Lake
    ARL = (6, 197)  # Arrow Lake
    PTL = (6, 204)  # Phoenix Lake


def is_admin():
    try:
        return ctypes.windll.shell32.IsUserAnAdmin()
    except:
        return False


def bytes_to_human_readable(num_bytes):
    for unit in ["B", "KB", "MB", "GB", "TB"]:
        if num_bytes < 1024.0:
            return f"{num_bytes:.3f} {unit}"
        num_bytes /= 1024.0
    return f"{num_bytes:.6f} PB"


class OLSStatus(Enum):
    OLS_DLL_NO_ERROR = 0
    OLS_DLL_UNSUPPORTED_PLATFORM = 1
    OLS_DLL_DRIVER_NOT_LOADED = 2
    OLS_DLL_DRIVER_NOT_FOUND = 3
    OLS_DLL_DRIVER_UNLOADED = 4
    OLS_DLL_DRIVER_NOT_LOADED_ON_NETWORK = 5
    OLS_DLL_UNKNOWN_ERROR = 9
    OLS_DLL_DRIVER_INVALID_PARAM = 10
    OLS_DLL_DRIVER_SC_MANAGER_NOT_OPENED = 11
    OLS_DLL_DRIVER_SC_DRIVER_NOT_INSTALLED = 12
    OLS_DLL_DRIVER_SC_DRIVER_NOT_STARTED = 13
    OLS_DLL_DRIVER_SC_DRIVER_NOT_REMOVED = 14


class WinRing0:
    def __init__(self):
        driver_path = os.path.dirname(os.path.abspath(__file__)) + "\\"
        self.lib = ctypes.WinDLL(driver_path + "WinRing0x64.dll")
        driver_path += "WinRing0x64.sys"
        print(f"Driver path: {driver_path}")
        self.lib.SetDriverPath(ctypes.c_char_p(driver_path.encode("utf-8")))

        res = self.lib.InitializeOls()
        if res != 0:
            raise Exception(f"Failed to initialize WinRing0: {res}")

        res = self.lib.GetDllStatus()
        if res != 0:
            raise Exception(f"Failed to initialize WinRing0: {res}")

    def __del__(self):
        res = self.lib.DeinitializeOls()
        # // TRUE: success, FALSE: failure
        if res == 0:
            print(f"Failed to deinitialize WinRing0: {res}")

    def cpuid(self, index):
        # BOOL WINAPI Cpuid(DWORD index, PDWORD eax, PDWORD ebx, PDWORD ecx, PDWORD edx)
        # TRUE: success, FALSE: failure
        eax = ctypes.c_uint32()
        ebx = ctypes.c_uint32()
        ecx = ctypes.c_uint32()
        edx = ctypes.c_uint32()
        res = self.lib.Cpuid(
            ctypes.c_uint32(index),
            ctypes.byref(eax),
            ctypes.byref(ebx),
            ctypes.byref(ecx),
            ctypes.byref(edx),
        )
        if res == 0:
            raise Exception(f"Failed to execute CPUID {index}: {res}")
        return eax.value, ebx.value, ecx.value, edx.value

    def getCPUFamilyModelFromCPUID(self, cpuid_array):
        eax = cpuid_array[0]
        family_id = (eax >> 8) & 0xF
        extended_family_id = (eax >> 20) & 0xFF
        model_id = (eax >> 4) & 0xF
        extended_model_id = (eax >> 16) & 0xF
        cpu_family = (
            family_id if family_id != 0x0F else (extended_family_id + family_id)
        )
        cpu_model = (
            (model_id + (extended_model_id << 4))
            if (family_id == 0x06 or family_id == 0x0F)
            else model_id
        )
        return cpu_family, cpu_model

    def read_msr(self, msr):

        # BOOL // TRUE: success, FALSE: failure
        # WINAPI Rdmsr(
        #     DWORD index,	// MSR index
        #     PDWORD eax,   // bit  0-31
        #     PDWORD edx	// bit 32-63
        # );

        eax = ctypes.c_uint32()
        edx = ctypes.c_uint32()
        res = self.lib.Rdmsr(ctypes.c_uint32(msr), ctypes.byref(eax), ctypes.byref(edx))
        if res == 0:
            raise Exception(f"Failed to read MSR {msr}: {res}")
        return (eax.value, edx.value)

    def ReadPciConfigDword(self, bus, device, function, offset):
        # PciBusDevFunc(Bus, Dev, Func)	((Bus&0xFF)<<8) | ((Dev&0x1F)<<3) | (Func&7)
        #
        # DWORD WINAPI ReadPciConfigDword(
        #     DWORD pciAddress,	// PCI Device Address
        #     BYTE regAddress	// Configuration Address 0-255
        # );

        pci_address = ((bus & 0xFF) << 8) | ((device & 0x1F) << 3) | (function & 7)
        pci_address = ctypes.c_uint32(pci_address)
        reg_address = ctypes.c_uint8(offset)
        self.lib.ReadPciConfigDword.restype = ctypes.c_uint32
        value = self.lib.ReadPciConfigDword(pci_address, reg_address)
        return value


class ASMMAP:
    def __init__(self):
        driver_path = os.path.dirname(os.path.abspath(__file__)) + "\\"
        self.lib = ctypes.WinDLL(driver_path + "asmmap64.dll")
        driver_path += "asmmap64.sys"
        print(f"Driver path: {driver_path}")
        res = self.lib.asmmap64_init(ctypes.c_char_p(driver_path.encode("utf-8")))
        if res != 0:
            raise Exception(f"Failed to initialize ASMMAP: {res}")
        self.mmap_dict = {}

    def __del__(self):
        to_unmmap_dict = self.mmap_dict.copy()
        for virt_addr, (phys_addr, size) in to_unmmap_dict.items():
            try:
                self.unmmap(virt_addr, size)
            except Exception as e:
                print(f"Failed to unmmap {virt_addr:#010x} with size {size:#010x}: {e}")

        res = self.lib.asmmap64_close()
        if res != 0:
            print(f"Failed to deinitialize ASMMAP: {res}")

    def mmap(self, phys_addr, size):
        # void* asmmap64_mmap(uint64_t phys_addr, uint32_t length)
        self.lib.asmmap64_mmap.restype = ctypes.c_void_p
        ptr = self.lib.asmmap64_mmap(ctypes.c_uint64(phys_addr), ctypes.c_uint32(size))
        if not ptr:
            raise Exception(f"Failed to mmap {phys_addr:#010x} with size {size:#010x}")
        self.mmap_dict[ptr] = (phys_addr, size)
        return ptr

    def unmmap(self, virt_addr, size):
        # int asmmap64_unmmap(void* virt_addr, uint32_t length)
        res = self.lib.asmmap64_unmmap(
            ctypes.c_void_p(virt_addr), ctypes.c_uint32(size)
        )
        if res != 0:
            raise Exception(
                f"Failed to unmmap {virt_addr:#010x} with size {size:#010x}: {res}"
            )
        self.mmap_dict.pop(virt_addr, None)


class IMCEvent:
    # 0xD800：Alder Lake/Raptor Lake IMC
    # 0x5000：Taiger Lake IMC
    IMC_COUNT = 2
    IMC_STRIDE = 0x10000
    IMC_EVENT_BASE_OFFSET = [0x5000, 0xD800]
    IMC_DATA_TOTAL_CNT_OFFSET = 0x40
    IMC_DATA_READ_CNT_OFFSET = 0x58
    IMC_DATA_WRITE_CNT_OFFSET = 0xA0
    IMC_ACCESS_SIZE = 64

    PCM_CLIENT_IMC_DRAM_BASE = 0x5000
    PCM_CLIENT_IMC_DRAM_DATA_READS = 0x5050
    PCM_CLIENT_IMC_DRAM_DATA_WRITES = 0x5054

    def __init__(self, cpu_model: tuple, mchbar_phys: int, asmmap: ASMMAP):
        self.cpu_model = cpu_model
        self.mchbar_phys = mchbar_phys
        self.asmmap = asmmap
        self.last_read_cnt = 0
        self.last_write_cnt = 0
        self._cached_read_val = None
        self._cached_write_val = None
        if cpu_model in (m.value for m in NewCPUModels):
            self.is_new = True
            self.IMC_COUNT = 2
            self.mmap_sz = self.IMC_COUNT * self.IMC_STRIDE
            self.mchbar = asmmap.mmap(self.mchbar_phys, self.mmap_sz)
        else:
            self.is_new = False
            self.IMC_COUNT = 1
            self.mmap_sz = 0x10000
            self.mchbar = asmmap.mmap(self.mchbar_phys, self.mmap_sz)

    def __del__(self):
        if hasattr(self, "mchbar") and self.mchbar:
            del self.mchbar

    def get_imc_read_cnt(self):
        total_read_cnt = 0
        if self.is_new:
            for j in self.IMC_EVENT_BASE_OFFSET:
                for i in range(self.IMC_COUNT):
                    offset = j + i * self.IMC_STRIDE
                    read_val = ctypes.c_uint64.from_address(
                        self.mchbar + offset + self.IMC_DATA_READ_CNT_OFFSET
                    ).value
                    if read_val == 0xFFFFFFFFFFFFFFFF:
                        read_val = 0
                    total_read_cnt += read_val
        else:
            # 32bit for old cpuse.g. Skylake, Kaby Lake, Coffee Lake
            if self._cached_read_val is not None:
                total_read_cnt = self._cached_read_val
                self._cached_read_val = None
            else:
                read_val = ctypes.c_uint32.from_address(
                    self.mchbar + self.PCM_CLIENT_IMC_DRAM_DATA_READS
                ).value
                if self.last_read_cnt != 0 and read_val < self.last_read_cnt:
                    # 处理计数器回绕
                    read_val += 0x100000000
                self.last_read_cnt = read_val
                total_read_cnt = read_val
        return total_read_cnt

    def get_imc_write_cnt(self):
        total_write_cnt = 0
        if self.is_new:
            for j in self.IMC_EVENT_BASE_OFFSET:
                for i in range(self.IMC_COUNT):
                    offset = j + i * self.IMC_STRIDE
                    write_val = ctypes.c_uint64.from_address(
                        self.mchbar + offset + self.IMC_DATA_WRITE_CNT_OFFSET
                    ).value
                    if write_val == 0xFFFFFFFFFFFFFFFF:
                        write_val = 0
                    total_write_cnt += write_val
        else:
            # 32bit for old cpuse.g. Skylake, Kaby Lake, Coffee Lake
            if self._cached_write_val is not None:
                total_write_cnt = self._cached_write_val
                self._cached_write_val = None
            else:
                write_val = ctypes.c_uint32.from_address(
                    self.mchbar + self.PCM_CLIENT_IMC_DRAM_DATA_WRITES
                ).value
                if self.last_write_cnt != 0 and write_val < self.last_write_cnt:
                    # 处理计数器回绕
                    write_val += 0x100000000
                self.last_write_cnt = write_val
                total_write_cnt = write_val
        return total_write_cnt

    def get_imc_total_cnt(self):
        if self.is_new:
            total_cnt = 0
            for j in self.IMC_EVENT_BASE_OFFSET:
                for i in range(self.IMC_COUNT):
                    offset = j + i * self.IMC_STRIDE
                    val = ctypes.c_uint64.from_address(
                        self.mchbar + offset + self.IMC_DATA_TOTAL_CNT_OFFSET
                    ).value
                    if val == 0xFFFFFFFFFFFFFFFF:
                        val = 0
                    total_cnt += val
            return total_cnt
        else:
            read_cnt = self.get_imc_read_cnt()
            write_cnt = self.get_imc_write_cnt()
            self._cached_read_val = read_cnt
            self._cached_write_val = write_cnt
            return read_cnt + write_cnt

    def get_metrics(self):
        total_cnt = self.get_imc_total_cnt()
        read_cnt = self.get_imc_read_cnt()
        write_cnt = self.get_imc_write_cnt()
        return total_cnt, read_cnt, write_cnt


if __name__ == "__main__":
    if not is_admin():
        ctypes.windll.shell32.ShellExecuteW(
            None, "runas", sys.executable, " ".join(sys.argv), None, 1
        )
        sys.exit()
    winring0 = WinRing0()
    cpuid = winring0.cpuid(1)
    cpu_model = winring0.getCPUFamilyModelFromCPUID(cpuid)
    if cpu_model[0] != 6:
        print(f"Unsupported CPU family: {cpu_model[0]}")
        sys.exit()

    asmmap = ASMMAP()
    MCHBAR = winring0.ReadPciConfigDword(0, 0, 0, 0x48)
    MCHBAR = MCHBAR & 0xFFFFF000
    if MCHBAR == 0:
        print("unsupported platform.")
        sys.exit()
    print(f"MCHBAR: {MCHBAR:#010x}")
    imc_controller = IMCEvent(cpu_model, MCHBAR, asmmap)

    try:
        prev_total, prev_read, prev_write = imc_controller.get_metrics()
        start_time = datetime.datetime.now()
        while True:
            time.sleep(1)
            cur_total, cur_read, cur_write = imc_controller.get_metrics()
            end_time = datetime.datetime.now()
            delta_t = (end_time - start_time).total_seconds()

            mem_access_total = cur_total * imc_controller.IMC_ACCESS_SIZE
            mem_read_total = cur_read * imc_controller.IMC_ACCESS_SIZE
            mem_write_total = cur_write * imc_controller.IMC_ACCESS_SIZE

            print(
                f"Total: {bytes_to_human_readable(mem_access_total)}, Read: {bytes_to_human_readable(mem_read_total)}, Write: {bytes_to_human_readable(mem_write_total)}",
                end="; ",
            )

            access_speed = (
                (cur_total - prev_total) * imc_controller.IMC_ACCESS_SIZE / delta_t
            )
            read_speed = (
                (cur_read - prev_read) * imc_controller.IMC_ACCESS_SIZE / delta_t
            )
            write_speed = (
                (cur_write - prev_write) * imc_controller.IMC_ACCESS_SIZE / delta_t
            )
            print(
                f"Access Speed: {bytes_to_human_readable(access_speed)}/s, Read Speed: {bytes_to_human_readable(read_speed)}/s, Write Speed: {bytes_to_human_readable(write_speed)}/s"
            )

            prev_total, prev_read, prev_write = cur_total, cur_read, cur_write
            start_time = end_time
    except KeyboardInterrupt:
        print("Exiting...")
    finally:
        del asmmap
        del winring0