loader: replace SSE4a EXTRQ with native trampolines on hosts without SSE4a (#839)

loader: add native EXTRQ patches with trap fallback

Co-authored-by: nmzik <Nmzik@mail.ru>
This commit is contained in:
Ewen trividic
2026-10-03 17:40:51 +02:00
committed by GitHub
parent b11aea803c
commit a4d26cf81a
4 changed files with 585 additions and 152 deletions
+151 -52
View File
@@ -132,12 +132,21 @@ struct DecodedFunction {
bool requires_conservative_red_zone_tracking {};
};
enum class InstructionReplacement { None, ReciprocalSquareRoot, ExtractQ };
struct InstructionRewrite {
bool protect_red_zone {};
bool protected_indirect_call {};
bool emulate_rsqrt {};
bool protect_red_zone {};
bool protected_indirect_call {};
InstructionReplacement replacement {};
};
InstructionPatchCounts& ReplacementCounts(GuestInstructionPatchResult& result,
InstructionReplacement replacement) {
ASSERT(replacement != InstructionReplacement::None);
return replacement == InstructionReplacement::ReciprocalSquareRoot ? result.reciprocal_sqrt
: result.extrq;
}
bool IsStackPointerRegister(ZydisRegister reg) {
return reg != ZYDIS_REGISTER_NONE &&
ZydisRegisterGetLargestEnclosing(ZYDIS_MACHINE_MODE_LONG_64, reg) == ZYDIS_REGISTER_RSP;
@@ -814,33 +823,108 @@ void GenerateReciprocalSquareRoot(const DecodedCodeInstruction& decoded,
generator.L(done);
}
void CollectReciprocalSquareRoots(const DecodedFunction& function,
std::map<uintptr_t, InstructionRewrite>& rewrite_sites,
GuestInstructionPatchResult& result) {
bool IsSupportedExtractQ(const DecodedCodeInstruction& decoded) {
if (decoded.instruction.mnemonic != ZYDIS_MNEMONIC_EXTRQ) {
return false;
}
const auto is_xmm = [](const ZydisDecodedOperand& operand) {
return operand.type == ZYDIS_OPERAND_TYPE_REGISTER &&
operand.reg.value >= ZYDIS_REGISTER_XMM0 &&
operand.reg.value <= ZYDIS_REGISTER_XMM15;
};
const bool immediate = decoded.instruction.operand_count_visible == 3 &&
decoded.operands[1].type == ZYDIS_OPERAND_TYPE_IMMEDIATE &&
decoded.operands[2].type == ZYDIS_OPERAND_TYPE_IMMEDIATE;
if (!is_xmm(decoded.operands[0]) ||
(!immediate &&
!(decoded.instruction.operand_count_visible == 2 && is_xmm(decoded.operands[1])))) {
return false;
}
// Match the trap handler's encoding support: 66 [REX] 0F 78/79 ModRM [imm8 imm8].
const auto* code = reinterpret_cast<const u8*>(decoded.address);
const size_t opcode = (code[1] & 0xf0u) == 0x40u ? 2 : 1;
return code[0] == 0x66 && code[opcode] == 0x0f &&
code[opcode + 1] == (immediate ? 0x78 : 0x79) &&
decoded.instruction.length == opcode + (immediate ? 5 : 3);
}
void GenerateExtractQ(const DecodedCodeInstruction& decoded, Xbyak::CodeGenerator& generator) {
const Xbyak::Xmm destination(decoded.operands[0].reg.value - ZYDIS_REGISTER_XMM0);
// Keep spills below the guest red zone on both hosts, and restore all modified flags/GPRs.
generator.lea(rsp, ptr[rsp - GuestRedZoneSize]);
generator.pushfq();
generator.push(rax);
generator.push(rcx);
generator.push(rdx);
if (decoded.operands[1].type == ZYDIS_OPERAND_TYPE_IMMEDIATE) {
const u32 length = decoded.operands[1].imm.value.u & 63u;
const u32 index = decoded.operands[2].imm.value.u & 63u;
const u64 mask = length == 0 ? UINT64_MAX : (u64 {1} << length) - 1;
generator.movq(rax, destination);
if (index != 0) {
generator.shr(rax, index);
}
generator.mov(rdx, mask);
} else {
const Xbyak::Xmm source(decoded.operands[1].reg.value - ZYDIS_REGISTER_XMM0);
generator.movq(rax, source);
generator.mov(ecx, eax);
generator.neg(ecx);
generator.mov(rdx, UINT64_MAX);
// CL shifts use six bits: -length gives 64-length, with zero retaining all 64 bits.
generator.shr(rdx, cl);
generator.shr(rax, 8);
generator.mov(ecx, eax);
generator.movq(rax, destination);
generator.shr(rax, cl);
}
generator.and_(rax, rdx);
// Match the emulator: clear the upper XMM half, preserving the upper YMM half.
generator.movq(destination, rax);
generator.pop(rdx);
generator.pop(rcx);
generator.pop(rax);
generator.popfq();
generator.lea(rsp, ptr[rsp + GuestRedZoneSize]);
}
void CollectAmdInstructions(const DecodedFunction& function,
std::map<uintptr_t, InstructionRewrite>& rewrite_sites,
GuestInstructionPatchResult& result) {
static const bool host_lacks_sse4a = !Xbyak::util::Cpu().has(Xbyak::util::Cpu::tSSE4a);
for (const auto& [address, decoded]: function.instructions) {
const auto& instruction = decoded.instruction;
if (instruction.mnemonic != ZYDIS_MNEMONIC_VRSQRTPS ||
instruction.encoding != ZYDIS_INSTRUCTION_ENCODING_VEX ||
instruction.raw.vex.offset != 0 || decoded.operands[0].size != 128 ||
decoded.operands[1].type != ZYDIS_OPERAND_TYPE_REGISTER) {
const auto& instruction = decoded.instruction;
InstructionReplacement replacement {};
if (instruction.mnemonic == ZYDIS_MNEMONIC_VRSQRTPS &&
instruction.encoding == ZYDIS_INSTRUCTION_ENCODING_VEX &&
instruction.raw.vex.offset == 0 && decoded.operands[0].size == 128 &&
decoded.operands[1].type == ZYDIS_OPERAND_TYPE_REGISTER) {
replacement = InstructionReplacement::ReciprocalSquareRoot;
} else if (host_lacks_sse4a && IsSupportedExtractQ(decoded)) {
replacement = InstructionReplacement::ExtractQ;
} else {
continue;
}
++result.reciprocal_sqrt_instruction_count;
rewrite_sites[address].emulate_rsqrt = true;
++ReplacementCounts(result, replacement).found;
rewrite_sites[address].replacement = replacement;
}
}
void MarkReciprocalSquareRootTrap(u8* code, const ZydisDecodedInstruction& instruction) {
// Clear a reserved VEX.vvvv bit, retaining both register operands for the handler.
code[instruction.raw.vex.size - 1] &= ~0x08u;
void MarkInstructionTrap(u8* code, const ZydisDecodedInstruction& instruction,
InstructionReplacement replacement) {
if (replacement == InstructionReplacement::ReciprocalSquareRoot) {
// Clear a reserved VEX.vvvv bit, retaining both register operands for the handler.
code[instruction.raw.vex.size - 1] &= ~0x08u;
}
// EXTRQ already raises #UD on hosts without SSE4a; keep its original encoding.
}
void GenerateReciprocalSquareRootTrap(const DecodedCodeInstruction& decoded,
Xbyak::CodeGenerator& generator) {
void GenerateInstructionTrap(const DecodedCodeInstruction& decoded,
InstructionReplacement replacement, Xbyak::CodeGenerator& generator) {
std::array<u8, ZYDIS_MAX_INSTRUCTION_LENGTH> code {};
std::memcpy(code.data(), reinterpret_cast<const void*>(decoded.address),
decoded.instruction.length);
MarkReciprocalSquareRootTrap(code.data(), decoded.instruction);
MarkInstructionTrap(code.data(), decoded.instruction, replacement);
generator.db(code.data(), decoded.instruction.length);
}
@@ -852,20 +936,28 @@ bool NeedsTrapRedZoneProtection(const DecodedCodeInstruction& decoded) {
#endif
}
void TrapUnrelocatedReciprocalSquareRoots(
const PatchModule& module, const DecodedFunction& function,
const std::map<uintptr_t, InstructionRewrite>& rewrite_sites,
GuestInstructionPatchResult& result) {
void TrapUnrelocatedInstructions(const PatchModule& module, const DecodedFunction& function,
const std::map<uintptr_t, InstructionRewrite>& rewrite_sites,
GuestInstructionPatchResult& result) {
for (const auto& [address, rewrite]: rewrite_sites) {
if (!rewrite.emulate_rsqrt || module.patched.contains(reinterpret_cast<u8*>(address))) {
if (rewrite.replacement == InstructionReplacement::None ||
module.patched.contains(reinterpret_cast<u8*>(address))) {
continue;
}
const auto& decoded = function.instructions.at(address);
if (NeedsTrapRedZoneProtection(decoded)) {
if (rewrite.replacement == InstructionReplacement::ExtractQ) {
// Unlike VRSQRTPS, leaving EXTRQ unchanged would still trap and corrupt the red
// zone.
EXIT("AMD CPU compatibility: cannot safely trap EXTRQ at %p (guest red zone is "
"live)\n",
reinterpret_cast<void*>(address));
}
continue;
}
MarkReciprocalSquareRootTrap(reinterpret_cast<u8*>(address), decoded.instruction);
++result.trapped_reciprocal_sqrt_instruction_count;
MarkInstructionTrap(reinterpret_cast<u8*>(address), decoded.instruction,
rewrite.replacement);
++ReplacementCounts(result, rewrite.replacement).trapped;
}
}
@@ -877,38 +969,44 @@ void RelocateGuestInstructions(PatchModule* module, const DecodedFunction& funct
uintptr_t patch_start {};
uintptr_t continuation {};
size_t patch_size {};
bool trap_rsqrt {};
bool trap_replacements {};
};
const auto emit_span = [&](RelocationSpan& span) -> std::optional<size_t> {
auto& generator = module->trampoline_gen;
const size_t trampoline_offset = generator.getSize();
const bool has_rsqrt = std::ranges::any_of(span.instructions, [&](const auto* decoded) {
const auto rewrite = rewrite_sites.find(decoded->address);
return rewrite != rewrite_sites.end() && rewrite->second.emulate_rsqrt;
});
for (bool trap_rsqrt: {false, true}) {
span.trap_rsqrt = trap_rsqrt;
const bool has_replacements =
std::ranges::any_of(span.instructions, [&](const auto* decoded) {
const auto rewrite = rewrite_sites.find(decoded->address);
return rewrite != rewrite_sites.end() &&
rewrite->second.replacement != InstructionReplacement::None;
});
for (bool trap_replacements: {false, true}) {
span.trap_replacements = trap_replacements;
Xbyak::ClearError();
bool encoded = true;
for (const auto* decoded: span.instructions) {
const auto rewrite = rewrite_sites.find(decoded->address);
const bool emulate_rsqrt =
rewrite != rewrite_sites.end() && rewrite->second.emulate_rsqrt;
const auto rewrite = rewrite_sites.find(decoded->address);
const auto replacement = rewrite != rewrite_sites.end()
? rewrite->second.replacement
: InstructionReplacement::None;
const bool protected_indirect_call =
rewrite != rewrite_sites.end() && rewrite->second.protected_indirect_call;
const bool protect_red_zone =
(rewrite != rewrite_sites.end() && rewrite->second.protect_red_zone &&
!protected_indirect_call) ||
(emulate_rsqrt && trap_rsqrt && NeedsTrapRedZoneProtection(*decoded));
(replacement != InstructionReplacement::None && trap_replacements &&
NeedsTrapRedZoneProtection(*decoded));
if (protect_red_zone) {
generator.lea(rsp, ptr[rsp - GuestRedZoneSize]);
}
if (emulate_rsqrt) {
if (trap_rsqrt) {
GenerateReciprocalSquareRootTrap(*decoded, generator);
} else {
if (replacement != InstructionReplacement::None) {
if (trap_replacements) {
GenerateInstructionTrap(*decoded, replacement, generator);
} else if (replacement == InstructionReplacement::ReciprocalSquareRoot) {
GenerateReciprocalSquareRoot(*decoded, generator);
} else {
GenerateExtractQ(*decoded, generator);
}
} else if (protected_indirect_call) {
encoded = GenerateProtectedIndirectCall(*decoded, generator);
@@ -936,7 +1034,7 @@ void RelocateGuestInstructions(PatchModule* module, const DecodedFunction& funct
EXIT("Guest instruction trampoline generation failed: %s\n",
Xbyak::ConvertErrorToString(error));
}
if (error != Xbyak::ERR_CODE_IS_TOO_BIG || !has_rsqrt) {
if (error != Xbyak::ERR_CODE_IS_TOO_BIG || !has_replacements) {
break;
}
}
@@ -953,11 +1051,12 @@ void RelocateGuestInstructions(PatchModule* module, const DecodedFunction& funct
if (rewrite->second.protect_red_zone && decoded->accesses_memory) {
++result.patched_memory_instruction_count;
}
if (rewrite->second.emulate_rsqrt) {
if (span.trap_rsqrt) {
++result.trapped_reciprocal_sqrt_instruction_count;
if (rewrite->second.replacement != InstructionReplacement::None) {
auto& counts = ReplacementCounts(result, rewrite->second.replacement);
if (span.trap_replacements) {
++counts.trapped;
} else {
++result.patched_reciprocal_sqrt_instruction_count;
++counts.native;
}
}
}
@@ -1328,9 +1427,9 @@ void RelocateGuestInstructions(PatchModule* module, const DecodedFunction& funct
GuestInstructionPatchResult PatchGuestInstructions(u64 segment_addr, u64 segment_size,
std::span<const uintptr_t> function_starts,
bool protect_memory, bool emulate_rsqrt) {
bool protect_memory, bool emulate_amd) {
GuestInstructionPatchResult result {};
if (!protect_memory && !emulate_rsqrt) {
if (!protect_memory && !emulate_amd) {
return result;
}
auto* module = GetContainingModule(reinterpret_cast<void*>(segment_addr));
@@ -1354,7 +1453,7 @@ GuestInstructionPatchResult PatchGuestInstructions(u64 segment_addr, u64 segment
const size_t trampoline_begin = module->trampoline_gen.getSize();
bool analyze_red_zone = protect_memory;
#if KYTY_PLATFORM == KYTY_PLATFORM_WINDOWS
analyze_red_zone |= emulate_rsqrt;
analyze_red_zone |= emulate_amd;
#endif
for (size_t function_index = 0; function_index < starts.size(); ++function_index) {
const uintptr_t function_start = starts[function_index];
@@ -1380,12 +1479,12 @@ GuestInstructionPatchResult PatchGuestInstructions(u64 segment_addr, u64 segment
if (protect_memory) {
CollectRedZoneMemoryInstructions(function, rewrite_sites, result);
}
if (emulate_rsqrt) {
CollectReciprocalSquareRoots(function, rewrite_sites, result);
if (emulate_amd) {
CollectAmdInstructions(function, rewrite_sites, result);
}
if (!rewrite_sites.empty()) {
RelocateGuestInstructions(module, function, rewrite_sites, result);
TrapUnrelocatedReciprocalSquareRoots(*module, function, rewrite_sites, result);
TrapUnrelocatedInstructions(*module, function, rewrite_sites, result);
}
}
const auto trampoline_addr =
+28 -14
View File
@@ -11,20 +11,34 @@
namespace Loader {
struct InstructionPatchCounts {
uint64_t found = 0;
uint64_t native = 0;
uint64_t trapped = 0;
uint64_t Skipped() const { return found - native - trapped; }
InstructionPatchCounts& operator+=(const InstructionPatchCounts& other) {
found += other.found;
native += other.native;
trapped += other.trapped;
return *this;
}
};
struct GuestInstructionPatchResult {
uint64_t function_count = 0;
uint64_t instruction_count = 0;
uint64_t red_zone_function_count = 0;
uint64_t memory_instruction_count = 0;
uint64_t short_memory_instruction_count = 0;
uint64_t patched_memory_instruction_count = 0;
uint64_t stack_dependent_memory_instruction_count = 0;
uint64_t control_flow_memory_instruction_count = 0;
uint64_t unrelocatable_memory_instruction_count = 0;
uint64_t indirect_red_zone_function_count = 0;
uint64_t reciprocal_sqrt_instruction_count = 0;
uint64_t patched_reciprocal_sqrt_instruction_count = 0;
uint64_t trapped_reciprocal_sqrt_instruction_count = 0;
uint64_t function_count = 0;
uint64_t instruction_count = 0;
uint64_t red_zone_function_count = 0;
uint64_t memory_instruction_count = 0;
uint64_t short_memory_instruction_count = 0;
uint64_t patched_memory_instruction_count = 0;
uint64_t stack_dependent_memory_instruction_count = 0;
uint64_t control_flow_memory_instruction_count = 0;
uint64_t unrelocatable_memory_instruction_count = 0;
uint64_t indirect_red_zone_function_count = 0;
InstructionPatchCounts reciprocal_sqrt;
InstructionPatchCounts extrq;
};
void RegisterGuestInstructionPatchModule(void* module_ptr, uint64_t module_size,
@@ -34,7 +48,7 @@ void UnregisterGuestInstructionPatchModule(void* module_ptr);
// Apply enabled instruction fixes using native trampolines or safe trap fallbacks.
GuestInstructionPatchResult PatchGuestInstructions(uint64_t segment_addr, uint64_t segment_size,
std::span<const uintptr_t> function_starts,
bool protect_memory, bool emulate_rsqrt);
bool protect_memory, bool emulate_amd);
bool DecodeEhFrameFunctionStarts(uint64_t eh_frame_header_addr, uint64_t eh_frame_header_size,
std::vector<uintptr_t>* function_starts);
+19 -22
View File
@@ -1813,16 +1813,16 @@ void RuntimeLinker::LoadProgramToMemory(Program* program) {
EXIT_IF(tls_handler_size > UINT64_MAX - program->base_size_aligned);
program->mapped_size = program->base_size_aligned + tls_handler_size;
#if !defined(__APPLE__)
const bool emulate_rsqrt = Config::AmdCpuEnabled();
const bool emulate_amd = Config::AmdCpuEnabled();
#else
const bool emulate_rsqrt = false;
const bool emulate_amd = false;
#endif
#if KYTY_PLATFORM == KYTY_PLATFORM_WINDOWS
const bool protect_memory_faults = Config::RedZoneProtectionEnabled();
#else
const bool protect_memory_faults = false;
#endif
const bool patch_guest_instructions = protect_memory_faults || emulate_rsqrt;
const bool patch_guest_instructions = protect_memory_faults || emulate_amd;
constexpr uint64_t INSTRUCTION_TRAMPOLINE_SIZE = 8u * 1024u * 1024u;
if (patch_guest_instructions) {
@@ -1940,18 +1940,15 @@ void RuntimeLinker::LoadProgramToMemory(Program* program) {
if (!have_function_starts) {
Log::WriteToConsoleAndLog(
fmt::format("{}: {} not patched (function boundaries unavailable)\n",
emulate_rsqrt ? "AMD CPU compatibility" : "Guest red-zone protection",
emulate_amd ? "AMD CPU compatibility" : "Guest red-zone protection",
module_name));
}
GuestInstructionPatchResult totals {};
for (const auto& [segment_addr, segment_size]: executable_segments) {
const auto result = PatchGuestInstructions(segment_addr, segment_size, function_starts,
protect_memory_faults, emulate_rsqrt);
totals.reciprocal_sqrt_instruction_count += result.reciprocal_sqrt_instruction_count;
totals.patched_reciprocal_sqrt_instruction_count +=
result.patched_reciprocal_sqrt_instruction_count;
totals.trapped_reciprocal_sqrt_instruction_count +=
result.trapped_reciprocal_sqrt_instruction_count;
protect_memory_faults, emulate_amd);
totals.reciprocal_sqrt += result.reciprocal_sqrt;
totals.extrq += result.extrq;
if (protect_memory_faults) {
LOGF("Windows guest red-zone patching: %s, functions=%" PRIu64 ", red_zone=%" PRIu64
", memory=%" PRIu64 ", patched=%" PRIu64 ", short=%" PRIu64 ", stack=%" PRIu64
@@ -1965,19 +1962,19 @@ void RuntimeLinker::LoadProgramToMemory(Program* program) {
result.unrelocatable_memory_instruction_count);
}
}
if (emulate_rsqrt && have_function_starts) {
const auto patched = totals.patched_reciprocal_sqrt_instruction_count +
totals.trapped_reciprocal_sqrt_instruction_count;
const auto skipped = totals.reciprocal_sqrt_instruction_count - patched;
const char* status = totals.reciprocal_sqrt_instruction_count == 0
? "no matching instructions"
: patched == 0 ? "not patched"
: skipped != 0 ? "partially patched"
: "patched";
if (emulate_amd && have_function_starts) {
const auto found = totals.reciprocal_sqrt.found + totals.extrq.found;
const auto skipped = totals.reciprocal_sqrt.Skipped() + totals.extrq.Skipped();
const char* status = found == 0 ? "no matching instructions"
: skipped == found ? "not patched"
: skipped != 0 ? "partially patched"
: "patched";
Log::WriteToConsoleAndLog(fmt::format(
"AMD CPU compatibility: {} {} (VRSQRTPS: native={}, trapped={}, skipped={})\n",
module_name, status, totals.patched_reciprocal_sqrt_instruction_count,
totals.trapped_reciprocal_sqrt_instruction_count, skipped));
"AMD CPU compatibility: {} {} (VRSQRTPS: native={}, trapped={}, skipped={}; "
"EXTRQ: native={}, trapped={}, skipped={})\n",
module_name, status, totals.reciprocal_sqrt.native, totals.reciprocal_sqrt.trapped,
totals.reciprocal_sqrt.Skipped(), totals.extrq.native, totals.extrq.trapped,
totals.extrq.Skipped()));
}
}