#endif
 
 #include <asm/ptrace.h>
+#include <asm/thread_info.h>
 
 /*
  * Stack pushing/popping (register pairs only). Equivalent to store decrement
        msr     daifclr, #8
        .endm
 
-       .macro  disable_step, tmp
+       .macro  disable_step_tsk, flgs, tmp
+       tbz     \flgs, #TIF_SINGLESTEP, 9990f
        mrs     \tmp, mdscr_el1
        bic     \tmp, \tmp, #1
        msr     mdscr_el1, \tmp
+       isb     // Synchronise with enable_dbg
+9990:
        .endm
 
-       .macro  enable_step, tmp
+       .macro  enable_step_tsk, flgs, tmp
+       tbz     \flgs, #TIF_SINGLESTEP, 9990f
+       disable_dbg
        mrs     \tmp, mdscr_el1
        orr     \tmp, \tmp, #1
        msr     mdscr_el1, \tmp
+9990:
        .endm
 
-       .macro  enable_dbg_if_not_stepping, tmp
-       mrs     \tmp, mdscr_el1
-       tbnz    \tmp, #0, 9990f
-       enable_dbg
-9990:
+/*
+ * Enable both debug exceptions and interrupts. This is likely to be
+ * faster than two daifclr operations, since writes to this register
+ * are self-synchronising.
+ */
+       .macro  enable_dbg_and_irq
+       msr     daifclr, #(8 | 2)
        .endm
 
 /*
 
        push    x0, x1
        .if     \el == 0
        mrs     x21, sp_el0
+       get_thread_info tsk                     // Ensure MDSCR_EL1.SS is clear,
+       ldr     x19, [tsk, #TI_FLAGS]           // since we can unmask debug
+       disable_step_tsk x19, x20               // exceptions when scheduling.
        .else
        add     x21, sp, #S_FRAME_SIZE
        .endif
         * Data abort handling
         */
        mrs     x0, far_el1
-       enable_dbg_if_not_stepping x2
+       enable_dbg
        // re-enable interrupts if they were enabled in the aborted context
        tbnz    x23, #7, 1f                     // PSR_I_BIT
        enable_irq
         * Stack or PC alignment exception handling
         */
        mrs     x0, far_el1
+       enable_dbg
        mov     x1, x25
        mov     x2, sp
        b       do_sp_pc_abort
        /*
         * Undefined instruction
         */
+       enable_dbg
        mov     x0, sp
        b       do_undefinstr
 el1_dbg:
        mrs     x0, far_el1
        mov     x2, sp                          // struct pt_regs
        bl      do_debug_exception
-
+       enable_dbg
        kernel_exit 1
 el1_inv:
        // TODO: add support for undefined instructions in kernel mode
+       enable_dbg
        mov     x0, sp
        mov     x1, #BAD_SYNC
        mrs     x2, esr_el1
        .align  6
 el1_irq:
        kernel_entry 1
-       enable_dbg_if_not_stepping x0
+       enable_dbg
 #ifdef CONFIG_TRACE_IRQFLAGS
        bl      trace_hardirqs_off
 #endif
 #ifdef CONFIG_PREEMPT
 el1_preempt:
        mov     x24, lr
-1:     enable_dbg
-       bl      preempt_schedule_irq            // irq en/disable is done inside
+1:     bl      preempt_schedule_irq            // irq en/disable is done inside
        ldr     x0, [tsk, #TI_FLAGS]            // get new tasks TI_FLAGS
        tbnz    x0, #TIF_NEED_RESCHED, 1b       // needs rescheduling?
        ret     x24
        lsr     x24, x25, #ESR_EL1_EC_SHIFT     // exception class
        cmp     x24, #ESR_EL1_EC_SVC64          // SVC in 64-bit state
        b.eq    el0_svc
-       adr     lr, ret_from_exception
+       adr     lr, ret_to_user
        cmp     x24, #ESR_EL1_EC_DABT_EL0       // data abort in EL0
        b.eq    el0_da
        cmp     x24, #ESR_EL1_EC_IABT_EL0       // instruction abort in EL0
        lsr     x24, x25, #ESR_EL1_EC_SHIFT     // exception class
        cmp     x24, #ESR_EL1_EC_SVC32          // SVC in 32-bit state
        b.eq    el0_svc_compat
-       adr     lr, ret_from_exception
+       adr     lr, ret_to_user
        cmp     x24, #ESR_EL1_EC_DABT_EL0       // data abort in EL0
        b.eq    el0_da
        cmp     x24, #ESR_EL1_EC_IABT_EL0       // instruction abort in EL0
         */
        mrs     x0, far_el1
        bic     x0, x0, #(0xff << 56)
-       disable_step x1
-       isb
-       enable_dbg
        // enable interrupts before calling the main handler
-       enable_irq
+       enable_dbg_and_irq
        mov     x1, x25
        mov     x2, sp
        b       do_mem_abort
         * Instruction abort handling
         */
        mrs     x0, far_el1
-       disable_step x1
-       isb
-       enable_dbg
        // enable interrupts before calling the main handler
-       enable_irq
+       enable_dbg_and_irq
        orr     x1, x25, #1 << 24               // use reserved ISS bit for instruction aborts
        mov     x2, sp
        b       do_mem_abort
        /*
         * Floating Point or Advanced SIMD access
         */
+       enable_dbg
        mov     x0, x25
        mov     x1, sp
        b       do_fpsimd_acc
        /*
         * Floating Point or Advanced SIMD exception
         */
+       enable_dbg
        mov     x0, x25
        mov     x1, sp
        b       do_fpsimd_exc
         * Stack or PC alignment exception handling
         */
        mrs     x0, far_el1
-       disable_step x1
-       isb
-       enable_dbg
        // enable interrupts before calling the main handler
-       enable_irq
+       enable_dbg_and_irq
        mov     x1, x25
        mov     x2, sp
        b       do_sp_pc_abort
        /*
         * Undefined instruction
         */
-       mov     x0, sp
        // enable interrupts before calling the main handler
-       enable_irq
+       enable_dbg_and_irq
+       mov     x0, sp
        b       do_undefinstr
 el0_dbg:
        /*
         */
        tbnz    x24, #0, el0_inv                // EL0 only
        mrs     x0, far_el1
-       disable_step x1
        mov     x1, x25
        mov     x2, sp
-       b       do_debug_exception
+       bl      do_debug_exception
+       enable_dbg
+       b       ret_to_user
 el0_inv:
+       enable_dbg
        mov     x0, sp
        mov     x1, #BAD_SYNC
        mrs     x2, esr_el1
 el0_irq:
        kernel_entry 0
 el0_irq_naked:
-       disable_step x1
-       isb
        enable_dbg
 #ifdef CONFIG_TRACE_IRQFLAGS
        bl      trace_hardirqs_off
 #endif
 
        irq_handler
-       get_thread_info tsk
 
 #ifdef CONFIG_TRACE_IRQFLAGS
        bl      trace_hardirqs_on
        b       ret_to_user
 ENDPROC(el0_irq)
 
-/*
- * This is the return code to user mode for abort handlers
- */
-ret_from_exception:
-       get_thread_info tsk
-       b       ret_to_user
-ENDPROC(ret_from_exception)
-
 /*
  * Register switch for AArch64. The callee-saved registers need to be saved
  * and restored. On entry:
        ldr     x1, [tsk, #TI_FLAGS]
        and     x2, x1, #_TIF_WORK_MASK
        cbnz    x2, fast_work_pending
-       tbz     x1, #TIF_SINGLESTEP, fast_exit
-       disable_dbg
-       enable_step x2
-fast_exit:
+       enable_step_tsk x1, x2
        kernel_exit 0, ret = 1
 
 /*
        bl      do_notify_resume
        b       ret_to_user
 work_resched:
-       enable_dbg
        bl      schedule
 
 /*
        ldr     x1, [tsk, #TI_FLAGS]
        and     x2, x1, #_TIF_WORK_MASK
        cbnz    x2, work_pending
-       tbz     x1, #TIF_SINGLESTEP, no_work_pending
-       disable_dbg
-       enable_step x2
+       enable_step_tsk x1, x2
 no_work_pending:
        kernel_exit 0, ret = 0
 ENDPROC(ret_to_user)
        mov     sc_nr, #__NR_syscalls
 el0_svc_naked:                                 // compat entry point
        stp     x0, scno, [sp, #S_ORIG_X0]      // save the original x0 and syscall number
-       disable_step x16
-       isb
-       enable_dbg
-       enable_irq
+       enable_dbg_and_irq
 
-       get_thread_info tsk
        ldr     x16, [tsk, #TI_FLAGS]           // check for syscall tracing
        tbnz    x16, #TIF_SYSCALL_TRACE, __sys_trace // are we tracing syscalls?
        adr     lr, ret_fast_syscall            // return address