}
 #endif
 
+#ifdef CONFIG_PREEMPT_TRACER
+void tracer_preempt_on(unsigned long a0, unsigned long a1);
+void tracer_preempt_off(unsigned long a0, unsigned long a1);
+#else
+static inline void tracer_preempt_on(unsigned long a0, unsigned long a1) { }
+static inline void tracer_preempt_off(unsigned long a0, unsigned long a1) { }
+#endif
+#ifdef CONFIG_IRQSOFF_TRACER
+void tracer_hardirqs_on(unsigned long a0, unsigned long a1);
+void tracer_hardirqs_off(unsigned long a0, unsigned long a1);
+#else
+static inline void tracer_hardirqs_on(unsigned long a0, unsigned long a1) { }
+static inline void tracer_hardirqs_off(unsigned long a0, unsigned long a1) { }
+#endif
+
 extern struct trace_iterator *tracepoint_print_iter;
 
 #endif /* _LINUX_KERNEL_TRACE_H */
 
 /*
  * We are only interested in hardirq on/off events:
  */
-static void tracer_hardirqs_on(void *none, unsigned long a0, unsigned long a1)
+void tracer_hardirqs_on(unsigned long a0, unsigned long a1)
 {
        unsigned int pc = preempt_count();
 
-       /*
-        * Tracepoint probes are expected to be called with preempt disabled,
-        * We don't care about being called with preempt disabled but we need
-        * to know in the future if that changes so we can remove the next
-        * preempt_enable.
-        */
-       WARN_ON_ONCE(pc < PREEMPT_DISABLE_OFFSET);
-
-       /* Use PREEMPT_DISABLE_OFFSET to handle !CONFIG_PREEMPT cases */
-       pc -= PREEMPT_DISABLE_OFFSET;
-
        if (!preempt_trace(pc) && irq_trace())
                stop_critical_timing(a0, a1, pc);
 }
 
-static void tracer_hardirqs_off(void *none, unsigned long a0, unsigned long a1)
+void tracer_hardirqs_off(unsigned long a0, unsigned long a1)
 {
        unsigned int pc = preempt_count();
 
-       /*
-        * Tracepoint probes are expected to be called with preempt disabled,
-        * We don't care about being called with preempt disabled but we need
-        * to know in the future if that changes so we can remove the next
-        * preempt_enable.
-        */
-       WARN_ON_ONCE(pc < PREEMPT_DISABLE_OFFSET);
-
-       /* Use PREEMPT_DISABLE_OFFSET to handle !CONFIG_PREEMPT cases */
-       pc -= PREEMPT_DISABLE_OFFSET;
-
        if (!preempt_trace(pc) && irq_trace())
                start_critical_timing(a0, a1, pc);
 }
 {
        trace_type = TRACER_IRQS_OFF;
 
-       register_trace_irq_disable(tracer_hardirqs_off, NULL);
-       register_trace_irq_enable(tracer_hardirqs_on, NULL);
        return __irqsoff_tracer_init(tr);
 }
 
 static void irqsoff_tracer_reset(struct trace_array *tr)
 {
-       unregister_trace_irq_disable(tracer_hardirqs_off, NULL);
-       unregister_trace_irq_enable(tracer_hardirqs_on, NULL);
        __irqsoff_tracer_reset(tr);
 }
 
 #endif /*  CONFIG_IRQSOFF_TRACER */
 
 #ifdef CONFIG_PREEMPT_TRACER
-static void tracer_preempt_on(void *none, unsigned long a0, unsigned long a1)
+void tracer_preempt_on(unsigned long a0, unsigned long a1)
 {
        int pc = preempt_count();
 
                stop_critical_timing(a0, a1, pc);
 }
 
-static void tracer_preempt_off(void *none, unsigned long a0, unsigned long a1)
+void tracer_preempt_off(unsigned long a0, unsigned long a1)
 {
        int pc = preempt_count();
 
 {
        trace_type = TRACER_PREEMPT_OFF;
 
-       register_trace_preempt_disable(tracer_preempt_off, NULL);
-       register_trace_preempt_enable(tracer_preempt_on, NULL);
        return __irqsoff_tracer_init(tr);
 }
 
 static void preemptoff_tracer_reset(struct trace_array *tr)
 {
-       unregister_trace_preempt_disable(tracer_preempt_off, NULL);
-       unregister_trace_preempt_enable(tracer_preempt_on, NULL);
        __irqsoff_tracer_reset(tr);
 }
 
 {
        trace_type = TRACER_IRQS_OFF | TRACER_PREEMPT_OFF;
 
-       register_trace_irq_disable(tracer_hardirqs_off, NULL);
-       register_trace_irq_enable(tracer_hardirqs_on, NULL);
-       register_trace_preempt_disable(tracer_preempt_off, NULL);
-       register_trace_preempt_enable(tracer_preempt_on, NULL);
-
        return __irqsoff_tracer_init(tr);
 }
 
 static void preemptirqsoff_tracer_reset(struct trace_array *tr)
 {
-       unregister_trace_irq_disable(tracer_hardirqs_off, NULL);
-       unregister_trace_irq_enable(tracer_hardirqs_on, NULL);
-       unregister_trace_preempt_disable(tracer_preempt_off, NULL);
-       unregister_trace_preempt_enable(tracer_preempt_on, NULL);
-
        __irqsoff_tracer_reset(tr);
 }
 
 
 #include <linux/uaccess.h>
 #include <linux/module.h>
 #include <linux/ftrace.h>
+#include "trace.h"
 
 #define CREATE_TRACE_POINTS
 #include <trace/events/preemptirq.h>
 void trace_hardirqs_on(void)
 {
        if (this_cpu_read(tracing_irq_cpu)) {
-               trace_irq_enable_rcuidle(CALLER_ADDR0, CALLER_ADDR1);
+               if (!in_nmi())
+                       trace_irq_enable_rcuidle(CALLER_ADDR0, CALLER_ADDR1);
+               tracer_hardirqs_on(CALLER_ADDR0, CALLER_ADDR1);
                this_cpu_write(tracing_irq_cpu, 0);
        }
 
 {
        if (!this_cpu_read(tracing_irq_cpu)) {
                this_cpu_write(tracing_irq_cpu, 1);
-               trace_irq_disable_rcuidle(CALLER_ADDR0, CALLER_ADDR1);
+               tracer_hardirqs_off(CALLER_ADDR0, CALLER_ADDR1);
+               if (!in_nmi())
+                       trace_irq_disable_rcuidle(CALLER_ADDR0, CALLER_ADDR1);
        }
 
        lockdep_hardirqs_off(CALLER_ADDR0);
 __visible void trace_hardirqs_on_caller(unsigned long caller_addr)
 {
        if (this_cpu_read(tracing_irq_cpu)) {
-               trace_irq_enable_rcuidle(CALLER_ADDR0, caller_addr);
+               if (!in_nmi())
+                       trace_irq_enable_rcuidle(CALLER_ADDR0, caller_addr);
+               tracer_hardirqs_on(CALLER_ADDR0, caller_addr);
                this_cpu_write(tracing_irq_cpu, 0);
        }
 
 {
        if (!this_cpu_read(tracing_irq_cpu)) {
                this_cpu_write(tracing_irq_cpu, 1);
-               trace_irq_disable_rcuidle(CALLER_ADDR0, caller_addr);
+               tracer_hardirqs_off(CALLER_ADDR0, caller_addr);
+               if (!in_nmi())
+                       trace_irq_disable_rcuidle(CALLER_ADDR0, caller_addr);
        }
 
        lockdep_hardirqs_off(CALLER_ADDR0);
 
 void trace_preempt_on(unsigned long a0, unsigned long a1)
 {
-       trace_preempt_enable_rcuidle(a0, a1);
+       if (!in_nmi())
+               trace_preempt_enable_rcuidle(a0, a1);
+       tracer_preempt_on(a0, a1);
 }
 
 void trace_preempt_off(unsigned long a0, unsigned long a1)
 {
-       trace_preempt_disable_rcuidle(a0, a1);
+       if (!in_nmi())
+               trace_preempt_disable_rcuidle(a0, a1);
+       tracer_preempt_off(a0, a1);
 }
 #endif