Last active
February 17, 2024 20:24
-
-
Save glandium/01d54cefdb70561b5f6675e08f2990f2 to your computer and use it in GitHub Desktop.
Linux kernel module for Zen workaround for rr
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| obj-m = zen_workaround.o | |
| all: | |
| make -C /lib/modules/$(shell uname -r)/build M=$(PWD) modules | |
| clean: | |
| make -C /lib/modules/$(shell uname -r)/build M=$(PWD) clean |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| #include <linux/module.h> | |
| #include <linux/tracepoint.h> | |
| #include <linux/suspend.h> | |
| #define MODULE_NAME "zen_workaround" | |
| #define SPECLOCKMAP_DISABLE BIT_64(54) | |
| u64 set_speclockmap_disable(u64 msr) { | |
| return msr | SPECLOCKMAP_DISABLE; | |
| } | |
| u64 unset_speclockmap_disable(u64 msr) { | |
| return msr & ~SPECLOCKMAP_DISABLE; | |
| } | |
| typedef u64 (*edit_msr_func_t)(u64); | |
| static void edit_ls_cfg_on_cpu(void *info) | |
| { | |
| int cpu = get_cpu(); | |
| u64 value = 0; | |
| if (!rdmsrl_safe(MSR_AMD64_LS_CFG, &value)) { | |
| edit_msr_func_t edit_msr = (edit_msr_func_t) info; | |
| u64 new_value = edit_msr(value); | |
| if (!wrmsrl_safe(MSR_AMD64_LS_CFG, new_value)) { | |
| pr_info("MSR_AMD64_LS_CFG for cpu %d was 0x%llx, setting to 0x%llx\n", | |
| cpu, value, new_value); | |
| } else { | |
| pr_err("MSR_AMD64_LS_CFG for cpu %d was 0x%llx, setting to 0x%llx failed\n", | |
| cpu, value, new_value); | |
| } | |
| } | |
| put_cpu(); | |
| } | |
| static void do_zen_workaround(edit_msr_func_t edit_msr) | |
| { | |
| smp_call_function(edit_ls_cfg_on_cpu, edit_msr, 1); | |
| edit_ls_cfg_on_cpu(edit_msr); | |
| } | |
| void on_write_msr(void *data, unsigned int msr, u64 val, int failed) | |
| { | |
| if (msr == MSR_AMD64_LS_CFG && !(val & SPECLOCKMAP_DISABLE)) { | |
| native_wrmsrl(MSR_AMD64_LS_CFG, set_speclockmap_disable(val)); | |
| } | |
| } | |
| static int install_probe(void) | |
| { | |
| return !boot_cpu_has(X86_FEATURE_AMD_SSBD) && !boot_cpu_has(X86_FEATURE_VIRT_SSBD); | |
| } | |
| static int enable_zen_workaround(void) | |
| { | |
| if (install_probe()) { | |
| int ret = tracepoint_probe_register(&__tracepoint_write_msr, on_write_msr, NULL); | |
| if (ret) { | |
| pr_err("Registering tracepoint probe failed\n"); | |
| return ret; | |
| } | |
| } | |
| do_zen_workaround(set_speclockmap_disable); | |
| return 0; | |
| } | |
| static int pm_notification(struct notifier_block *this, unsigned long event, void *ptr) | |
| { | |
| switch (event) { | |
| case PM_POST_SUSPEND: | |
| case PM_POST_HIBERNATION: | |
| case PM_POST_RESTORE: | |
| enable_zen_workaround(); | |
| break; | |
| case PM_HIBERNATION_PREPARE: | |
| case PM_SUSPEND_PREPARE: | |
| if (install_probe()) { | |
| tracepoint_probe_unregister(&__tracepoint_write_msr, on_write_msr, NULL); | |
| } | |
| break; | |
| } | |
| return NOTIFY_DONE; | |
| } | |
| static struct notifier_block pm_notifier = { | |
| .notifier_call = pm_notification, | |
| }; | |
| static int __init zen_workaround_init(void) | |
| { | |
| if (!boot_cpu_has(X86_FEATURE_ZEN)) { | |
| pr_err("Cannot use the Zen workaround on a non-Zen CPU\n"); | |
| return -EINVAL; | |
| } | |
| enable_zen_workaround(); | |
| register_pm_notifier(&pm_notifier); | |
| return 0; | |
| } | |
| module_init(zen_workaround_init); | |
| static void __exit zen_workaround_exit(void) | |
| { | |
| unregister_pm_notifier(&pm_notifier); | |
| if (install_probe()) { | |
| tracepoint_probe_unregister(&__tracepoint_write_msr, on_write_msr, NULL); | |
| } | |
| do_zen_workaround(unset_speclockmap_disable); | |
| } | |
| module_exit(zen_workaround_exit); | |
| MODULE_LICENSE("GPL"); |
On NixOS, as per https://nixos.wiki/wiki/Linux_kernel#Developing_kernel_modules, one should be able to just:
make -C $(nix-build -E '(import <nixpkgs> {}).linux.dev' --no-out-link)/lib/modules/*/build M=$(pwd) modulesThe Makefile in the gist is fine to use because the obj-m = zen_workaround.o line will be used and the rest ignored.
In case you have SecureBoot enabled (tested on Ubuntu):
sign: $(obj-m)
/lib/modules/$(shell uname -r)/build/scripts/sign-file sha512 /var/lib/shim-signed/mok/MOK.priv /var/lib/shim-signed/mok/MOK.der zen_workaround.ko
I have packaged this into a NixOS module.
Kernel 6.5.8 on Ubuntu 23.10:
[ 9058.474358] BUG: scheduling while atomic: swapper/15/0/0x00000002
[ 9058.474360] Modules linked in:
[ 9058.474360] BUG: scheduling while atomic: swapper/5/0/0x00000002
[ 9058.474361] zen_workaround(O+)
[ 9058.474363] Modules linked in: zen_workaround(O+)
Fix:
--- zen_workaround.c.orig 2023-11-06 16:37:23.000582213 +0100
+++ zen_workaround.c 2023-11-06 16:37:24.628603272 +0100
@@ -32,6 +32,8 @@
cpu, value, new_value);
}
}
+
+ put_cpu();
}
static void do_zen_workaround(edit_msr_func_t edit_msr)
Fix:
--- zen_workaround.c.orig 2023-11-06 16:37:23.000582213 +0100 +++ zen_workaround.c 2023-11-06 16:37:24.628603272 +0100 @@ -32,6 +32,8 @@ cpu, value, new_value); } } + + put_cpu(); } static void do_zen_workaround(edit_msr_func_t edit_msr)
Applied. Thanks.
Could this module be moved to a proper repository? this would make it easier to track updates.
Maybe it should just be in the rr repo. Open an issue?
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Latest version avoids requiring internal kernel symbols altogether.