asm volatile("lgdt %0"::"m" (*dtr));
 }
 
-static inline void native_load_idt(const struct desc_ptr *dtr)
+static __always_inline void native_load_idt(const struct desc_ptr *dtr)
 {
        asm volatile("lidt %0"::"m" (*dtr));
 }
 
 #ifdef CONFIG_X86_64
 DECLARE_PER_CPU(u32, debug_idt_ctr);
-static inline bool is_debug_idt_enabled(void)
+static __always_inline bool is_debug_idt_enabled(void)
 {
        if (this_cpu_read(debug_idt_ctr))
                return true;
        return false;
 }
 
-static inline void load_debug_idt(void)
+static __always_inline void load_debug_idt(void)
 {
        load_idt((const struct desc_ptr *)&debug_idt_descr);
 }
  * that doesn't need to disable interrupts, as nothing should be
  * bothering the CPU then.
  */
-static inline void load_current_idt(void)
+static __always_inline void load_current_idt(void)
 {
        if (is_debug_idt_enabled())
                load_debug_idt();
 
 DEFINE_PER_CPU(int, debug_stack_usage);
 DEFINE_PER_CPU(u32, debug_idt_ctr);
 
-void debug_stack_set_zero(void)
+noinstr void debug_stack_set_zero(void)
 {
        this_cpu_inc(debug_idt_ctr);
        load_current_idt();
 }
-NOKPROBE_SYMBOL(debug_stack_set_zero);
 
-void debug_stack_reset(void)
+noinstr void debug_stack_reset(void)
 {
        if (WARN_ON(!this_cpu_read(debug_idt_ctr)))
                return;
        if (this_cpu_dec_return(debug_idt_ctr) == 0)
                load_current_idt();
 }
-NOKPROBE_SYMBOL(debug_stack_reset);
 
 #else  /* CONFIG_X86_64 */
 
 
 static DEFINE_PER_CPU(bool, swallow_nmi);
 static DEFINE_PER_CPU(unsigned long, last_nmi_rip);
 
-static void default_do_nmi(struct pt_regs *regs)
+static noinstr void default_do_nmi(struct pt_regs *regs)
 {
        unsigned char reason = 0;
        int handled;
 
        __this_cpu_write(last_nmi_rip, regs->ip);
 
+       instrumentation_begin();
+
        handled = nmi_handle(NMI_LOCAL, regs);
        __this_cpu_add(nmi_stats.normal, handled);
        if (handled) {
                 */
                if (handled > 1)
                        __this_cpu_write(swallow_nmi, true);
-               return;
+               goto out;
        }
 
        /*
 #endif
                __this_cpu_add(nmi_stats.external, 1);
                raw_spin_unlock(&nmi_reason_lock);
-               return;
+               goto out;
        }
        raw_spin_unlock(&nmi_reason_lock);
 
                __this_cpu_add(nmi_stats.swallow, 1);
        else
                unknown_nmi_error(reason, regs);
+
+out:
+       instrumentation_end();
 }
-NOKPROBE_SYMBOL(default_do_nmi);
 
 /*
  * NMIs can page fault or hit breakpoints which will cause it to lose
  */
 static DEFINE_PER_CPU(int, update_debug_stack);
 
-static bool notrace is_debug_stack(unsigned long addr)
+static noinstr bool is_debug_stack(unsigned long addr)
 {
        struct cea_exception_stacks *cs = __this_cpu_read(cea_exception_stacks);
        unsigned long top = CEA_ESTACK_TOP(cs, DB);
         */
        return addr >= bot && addr < top;
 }
-NOKPROBE_SYMBOL(is_debug_stack);
 #endif
 
 DEFINE_IDTENTRY_NMI(exc_nmi)