Skip to content

Commit

Permalink
Various fixes and lock TSC freq on AMD if possible
Browse files Browse the repository at this point in the history
  • Loading branch information
VisualEhrmanntraut committed May 27, 2024
1 parent 7777cad commit 3165865
Show file tree
Hide file tree
Showing 2 changed files with 81 additions and 24 deletions.
96 changes: 74 additions & 22 deletions ForgedInvariant/ForgedInvariant.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -8,68 +8,108 @@
#include <i386/proc_reg.h>

ForgedInvariantMain *ForgedInvariantMain::callback = nullptr;
_Atomic(bool) ForgedInvariantMain::syncCompleted = ATOMIC_VAR_INIT(false);
_Atomic(bool) ForgedInvariantMain::systemAwake = ATOMIC_VAR_INIT(true);
_Atomic(bool) ForgedInvariantMain::synchronised = ATOMIC_VAR_INIT(false);
_Atomic(int) ForgedInvariantMain::threadsEngaged = ATOMIC_VAR_INIT(0);
_Atomic(UInt64) ForgedInvariantMain::targetTSC = ATOMIC_VAR_INIT(0);

constexpr UInt32 MSR_IA32_TSC = 0x10;
constexpr UInt32 MSR_IA32_TSC_ADJUST = 0x3B;
constexpr UInt32 CPUID_LEAF7_TSC_ADJUST = getBit<UInt32>(1);
constexpr UInt32 MSR_TSC = 0x10;
constexpr UInt32 MSR_TSC_ADJUST = 0x3B;
constexpr UInt32 MSR_HWCR = 0xC0010015;
constexpr UInt64 MSR_HWCR_LOCK_TSC_TO_CURR_P0 = getBit<UInt64>(21);

void ForgedInvariantMain::resetTscAdjust(void *) { wrmsr64(MSR_IA32_TSC_ADJUST, 0); }
void ForgedInvariantMain::resetTscAdjust(void *) { wrmsr64(MSR_TSC_ADJUST, 0); }

void ForgedInvariantMain::lockTscFreqIfPossible() {
// On AMD Family 17h and newer, we can take advantage of the LockTscToCurrentP0 bit
// which allows us to lock the frequency of the TSC to the current P0 frequency
// and prevent it from changing regardless of future changes to it.
if (callback->constants.lockTSCFreqUsingHWCR) {
wrmsr64(MSR_HWCR, rdmsr64(MSR_HWCR) | MSR_HWCR_LOCK_TSC_TO_CURR_P0);
}
}

void ForgedInvariantMain::setTscValue(void *) {
// Barrier: Wait until all threads have reached this point.
lockTscFreqIfPossible();

// Thread: Hey, I'm here! What did I miss?
atomic_fetch_add_explicit(&threadsEngaged, 1, memory_order_relaxed);
while (atomic_load_explicit(&threadsEngaged, memory_order_relaxed) != callback->constants.threadCount) {}

// If we are the target thread, store the value for the other threads.
// Otherwise, wait until the TSC value is set.
if (cpu_number() == callback->constants.targetThread) {
atomic_store_explicit(&targetTSC, rdtsc64(), memory_order_relaxed);
} else {
// Otherwise, wait until the TSC value is set, then set the TSC MSR with the new value.
while (atomic_load_explicit(&targetTSC, memory_order_relaxed) == 0) {}
wrmsr64(MSR_IA32_TSC, atomic_load_explicit(&targetTSC, memory_order_relaxed));
}

// Barrier: Wait until all threads have reached this point.
while (atomic_load_explicit(&threadsEngaged, memory_order_relaxed) != callback->constants.threadCount) {}

// Set the TSC value of all threads to the same exact one.
wrmsr64(MSR_TSC, atomic_load_explicit(&targetTSC, memory_order_relaxed));
}

void ForgedInvariantMain::syncTsc() {
atomic_store_explicit(&synchronised, false, memory_order_relaxed);

// If we are on macOS 12 and newer and TSC_ADJUST is supported, just reset it.
atomic_store_explicit(&syncCompleted, false, memory_order_relaxed);
if (this->constants.newSyncMethod && this->constants.supportsTscAdjust) {
// Otherwise, we have to synchronise the TSC value itself.
if (this->constants.resetTscAdjust) {
mp_rendezvous_no_intrs(resetTscAdjust, nullptr);
} else {
// Otherwise, we have to sync the TSC value itself.
atomic_store_explicit(&threadsEngaged, 0, memory_order_relaxed);
atomic_store_explicit(&targetTSC, 0, memory_order_relaxed);
mp_rendezvous_no_intrs(setTscValue, nullptr);
}
atomic_store_explicit(&syncCompleted, true, memory_order_relaxed);

atomic_store_explicit(&synchronised, true, memory_order_relaxed);
}

void ForgedInvariantMain::wrapXcpmUrgency(int urgency, UInt64 rtPeriod, UInt64 rtDeadline) {
if (!atomic_load_explicit(&syncCompleted, memory_order_relaxed)) { return; }
if (!atomic_load_explicit(&synchronised, memory_order_relaxed)) { return; }
FunctionCast(wrapXcpmUrgency, callback->orgXcpmUrgency)(urgency, rtPeriod, rtDeadline);
}

constexpr UInt8 kIOPMTracePointWakeCPUs = 0x23;
constexpr UInt8 kIOPMTracePointSleepCPUs = 0x18;
constexpr UInt8 kIOPMTracePointWakePlatformActions = 0x22;

void ForgedInvariantMain::wrapTracePoint(void *that, UInt8 point) {
if (point == kIOPMTracePointWakeCPUs) { callback->syncTsc(); }
switch (point) {
case kIOPMTracePointSleepCPUs: {
atomic_store_explicit(&systemAwake, false, memory_order_relaxed);
atomic_store_explicit(&synchronised, false, memory_order_relaxed);
break;
}
case kIOPMTracePointWakePlatformActions: {
atomic_store_explicit(&systemAwake, true, memory_order_relaxed);
callback->syncTsc();
break;
}
default: {
break;
}
}
FunctionCast(wrapTracePoint, callback->orgTracePoint)(that, point);
}

void ForgedInvariantMain::wrapClockGetCalendarMicrotime(clock_sec_t *secs, clock_usec_t *microsecs) {
if (!atomic_load_explicit(&syncCompleted, memory_order_relaxed)) { callback->syncTsc(); }
// Don't need to synchronise the TSC if it's already synchronised.
// Also shouldn't synchronise it if the OS is supposed to be sleeping.
if (!atomic_load_explicit(&synchronised, memory_order_relaxed) &&
atomic_load_explicit(&systemAwake, memory_order_relaxed)) {
callback->syncTsc();
}
FunctionCast(wrapClockGetCalendarMicrotime, callback->orgClockGetCalendarMicrotime)(secs, microsecs);
}

void ForgedInvariantMain::processPatcher(KernelPatcher &patcher) {
KernelPatcher::RouteRequest requests[] {
KernelPatcher::RouteRequest requests[] = {
{"_xcpm_urgency", wrapXcpmUrgency, this->orgXcpmUrgency},
{"__ZN14IOPMrootDomain10tracePointEh", wrapTracePoint, this->orgTracePoint},
{"_clock_get_calendar_microtime", wrapClockGetCalendarMicrotime, this->orgClockGetCalendarMicrotime},
};

PANIC_COND(!patcher.routeMultiple(KernelPatcher::KernelID, requests), "Main", "Failed to route symbols");
}

Expand All @@ -78,11 +118,12 @@ void ForgedInvariantMain::init() {
callback = this;

// Monterey got a change in the task scheduling requiring the target TSC value
// to be in sync with the first thread and not the last one.
// to be in synchronise with the first thread and not the last one.
this->constants.newSyncMethod = getKernelVersion() >= KernelVersion::Monterey;
UInt32 reg = 0;
// CPUID Leaf 7 Count 0 Bit 1 defines whether a CPU supports TSC_ADJUST, according to the Intel SDM.
this->constants.supportsTscAdjust = CPUInfo::getCpuid(7, 0, nullptr, &reg) && (reg & getBit<UInt32>(1)) != 0;
this->constants.resetTscAdjust =
this->constants.newSyncMethod && CPUInfo::getCpuid(7, 0, nullptr, &reg) && (reg & CPUID_LEAF7_TSC_ADJUST) != 0;

switch (BaseDeviceInfo::get().cpuVendor) {
case CPUInfo::CpuVendor::Unknown: {
Expand All @@ -99,6 +140,14 @@ void ForgedInvariantMain::init() {
"1 thread...");
this->constants.threadCount = 1;
}
reg = 0;
if (CPUInfo::getCpuid(1, 0, &reg)) {
CPUInfo::CpuVersion *ver = reinterpret_cast<CPUInfo::CpuVersion *>(&reg);
UInt32 family = ver->family == 0xF ? ver->family + ver->extendedFamily : ver->family;
this->constants.lockTSCFreqUsingHWCR = family >= 0x17;
} else {
this->constants.lockTSCFreqUsingHWCR = false;
}
break;
}
case CPUInfo::CpuVendor::Intel: {
Expand All @@ -110,11 +159,14 @@ void ForgedInvariantMain::init() {

this->constants.targetThread = this->constants.newSyncMethod ? 0 : (this->constants.threadCount - 1);

DBGLOG("Main", "New sync method: %s.", this->constants.newSyncMethod ? "yes" : "no");
DBGLOG("Main", "Supports TSC_ADJUST: %s.", this->constants.supportsTscAdjust ? "yes" : "no");
DBGLOG("Main", "New synchronise method: %s.", this->constants.newSyncMethod ? "yes" : "no");
DBGLOG("Main", "Supports TSC_ADJUST: %s.", this->constants.resetTscAdjust ? "yes" : "no");
DBGLOG("Main", "Supports LockTscToCurrentP0: %s.", this->constants.lockTSCFreqUsingHWCR ? "yes" : "no");
DBGLOG("Main", "Thread count: %d.", this->constants.threadCount);
DBGLOG("Main", "Target thread: %d.", this->constants.targetThread);

lockTscFreqIfPossible();

lilu.onPatcherLoadForce(
[](void *user, KernelPatcher &patcher) { static_cast<ForgedInvariantMain *>(user)->processPatcher(patcher); },
this);
Expand Down
9 changes: 7 additions & 2 deletions ForgedInvariant/ForgedInvariant.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -7,14 +7,16 @@

struct ForgedInvariantConstants {
bool newSyncMethod {false};
bool supportsTscAdjust {false};
bool resetTscAdjust {false};
bool lockTSCFreqUsingHWCR {false};
int threadCount {0};
int targetThread {0};
};

class ForgedInvariantMain {
static ForgedInvariantMain *callback;
static _Atomic(bool) syncCompleted;
static _Atomic(bool) systemAwake;
static _Atomic(bool) synchronised;
static _Atomic(int) threadsEngaged;
static _Atomic(UInt64) targetTSC;

Expand All @@ -23,8 +25,11 @@ class ForgedInvariantMain {
mach_vm_address_t orgTracePoint {0};
mach_vm_address_t orgClockGetCalendarMicrotime {0};
mach_vm_address_t orgRtcNanotimeRead {0};
mach_vm_address_t orgIOHibernateSystemHasSlept {0};
mach_vm_address_t orgIOHibernateSystemWake {0};

static void resetTscAdjust(void *);
static void lockTscFreqIfPossible();
static void setTscValue(void *);

void syncTsc();
Expand Down

0 comments on commit 3165865

Please sign in to comment.