throttled/lenovo_fix.py

#!/usr/bin/env python3

import configparser
import dbus
import glob
import os
import struct
import subprocess

from collections import defaultdict
from dbus.mainloop.glib import DBusGMainLoop
from periphery import MMIO
from threading import Event, Thread

try:
    from gi.repository import GObject
except ImportError:
    import gobject as GObject

SYSFS_POWER_PATH = '/sys/class/power_supply/AC/online'
CONFIG_PATH = '/etc/lenovo_fix.conf'

VOLTAGE_PLANES = {
    'CORE': 0,
    'GPU': 1,
    'CACHE': 2,
    'UNCORE': 3,
    'ANALOGIO': 4,
}


def writemsr(msr, val):
    n = glob.glob('/dev/cpu/[0-9]*/msr')
    for c in n:
        f = os.open(c, os.O_WRONLY)
        os.lseek(f, msr, os.SEEK_SET)
        os.write(f, struct.pack('Q', val))
        os.close(f)
    if not n:
        try:
            subprocess.check_call(('modprobe', 'msr'))
        except subprocess.CalledProcessError:
            raise OSError("Unable to load msr module.")


def is_on_battery():
    with open(SYSFS_POWER_PATH) as f:
        return not bool(int(f.read()))


def calc_time_window_vars(t):
    for Y in range(2**5):
        for Z in range(2**2):
            if t <= (2**Y) * (1. + Z / 4.) * 0.000977:
                return (Y, Z)
    raise Exception('Unable to find a good combination!')


def undervolt(config):
    for plane in VOLTAGE_PLANES:
        writemsr(0x150, calc_undervolt_msr(plane, config.getfloat('UNDERVOLT', plane)))


def calc_undervolt_msr(plane, offset):
    assert offset <= 0
    assert plane in VOLTAGE_PLANES
    offset = int(round(offset * 1.024))
    offset = 0xFFE00000 & ((offset & 0xFFF) << 21)
    return 0x8000001100000000 | (VOLTAGE_PLANES[plane] << 40) | offset


def load_config():
    config = configparser.ConfigParser()
    config.read(CONFIG_PATH)

    for power_source in ('AC', 'BATTERY'):
        assert 0 < config.getfloat(power_source, 'Update_Rate_s')
        assert 0 < config.getfloat(power_source, 'PL1_Tdp_W')
        assert 0 < config.getfloat(power_source, 'PL1_Duration_s')
        assert 0 < config.getfloat(power_source, 'PL2_Tdp_W')
        assert 0 < config.getfloat(power_source, 'PL2_Duration_S')
        assert 40 < config.getfloat(power_source, 'Trip_Temp_C') < 98

    for plane in VOLTAGE_PLANES:
        assert config.getfloat('UNDERVOLT', plane) <= 0

    return config


def calc_reg_values(config):
    regs = defaultdict(dict)
    for power_source in ('AC', 'BATTERY'):
        # the critical temperature for this CPU is 100 C
        trip_offset = int(round(100 - config.getfloat(power_source, 'Trip_Temp_C')))
        regs[power_source]['MSR_TEMPERATURE_TARGET'] = trip_offset << 24

        # 0.125 is the power unit of this CPU
        PL1 = int(round(config.getfloat(power_source, 'PL1_Tdp_W') / 0.125))
        Y, Z = calc_time_window_vars(config.getfloat(power_source, 'PL1_Duration_s'))
        TW1 = Y | (Z << 5)

        PL2 = int(round(config.getfloat(power_source, 'PL2_Tdp_W') / 0.125))
        Y, Z = calc_time_window_vars(config.getfloat(power_source, 'PL2_Duration_s'))
        TW2 = Y | (Z << 5)

        regs[power_source]['MSR_PKG_POWER_LIMIT'] = PL1 | (1 << 15) | (TW1 << 17) | (PL2 << 32) | (1 << 47) | (
            TW2 << 49)

    return regs


def power_thread(config, regs, exit_event):
    mchbar_mmio = MMIO(0xfed159a0, 8)

    while not exit_event.is_set():
        power_source = 'BATTERY' if is_on_battery() else 'AC'

        # set temperature trip point
        writemsr(0x1a2, regs[power_source]['MSR_TEMPERATURE_TARGET'])

        # set PL1/2 on MSR
        writemsr(0x610, regs[power_source]['MSR_PKG_POWER_LIMIT'])
        # set MCHBAR register to the same PL1/2 values
        mchbar_mmio.write32(0, regs[power_source]['MSR_PKG_POWER_LIMIT'] & 0xffffffff)
        mchbar_mmio.write32(4, regs[power_source]['MSR_PKG_POWER_LIMIT'] >> 32)

        exit_event.wait(config.getfloat(power_source, 'Update_Rate_s'))


def main():
    config = load_config()
    regs = calc_reg_values(config)

    if not config.getboolean('GENERAL', 'Enabled'):
        return

    exit_event = Event()
    t = Thread(target=power_thread, args=(config, regs, exit_event))
    t.start()

    undervolt(config)

    # handle dbus events for applying undervolt on resume from sleep/hybernate
    def handle_sleep_callback(sleeping):
        if not sleeping:
            undervolt(config)

    DBusGMainLoop(set_as_default=True)
    bus = dbus.SystemBus()

    # add dbus receiver only if undervolt is enabled in config
    if any(config.getfloat('UNDERVOLT', plane) != 0 for plane in VOLTAGE_PLANES):
        bus.add_signal_receiver(handle_sleep_callback, 'PrepareForSleep', 'org.freedesktop.login1.Manager',
                                'org.freedesktop.login1')

    try:
        GObject.threads_init()
        loop = GObject.MainLoop()
        loop.run()
    except (KeyboardInterrupt, SystemExit):
        pass

    exit_event.set()
    loop.quit()
    t.join()


if __name__ == '__main__':
    main()