regs->ip = current->utask->xol_vaddr;
        pre_xol_rip_insn(auprobe, regs, autask);
 
+       autask->saved_tf = !!(regs->flags & X86_EFLAGS_TF);
+       regs->flags |= X86_EFLAGS_TF;
+       if (test_tsk_thread_flag(current, TIF_BLOCKSTEP))
+               set_task_blockstep(current, false);
+
        return 0;
 }
 
        if (auprobe->fixups & UPROBE_FIX_CALL)
                result = adjust_ret_addr(regs->sp, correction);
 
+       /*
+        * arch_uprobe_pre_xol() doesn't save the state of TIF_BLOCKSTEP
+        * so we can get an extra SIGTRAP if we do not clear TF. We need
+        * to examine the opcode to make it right.
+        */
+       if (utask->autask.saved_tf)
+               send_sig(SIGTRAP, current, 0);
+       else if (!(auprobe->fixups & UPROBE_FIX_SETF))
+               regs->flags &= ~X86_EFLAGS_TF;
+
        return result;
 }
 
        current->thread.trap_nr = utask->autask.saved_trap_nr;
        handle_riprel_post_xol(auprobe, regs, NULL);
        instruction_pointer_set(regs, utask->vaddr);
+
+       /* clear TF if it was set by us in arch_uprobe_pre_xol() */
+       if (!utask->autask.saved_tf)
+               regs->flags &= ~X86_EFLAGS_TF;
 }
 
 /*
                send_sig(SIGTRAP, current, 0);
        return ret;
 }
-
-void arch_uprobe_enable_step(struct arch_uprobe *auprobe)
-{
-       struct task_struct *task = current;
-       struct arch_uprobe_task *autask = &task->utask->autask;
-       struct pt_regs *regs = task_pt_regs(task);
-
-       autask->saved_tf = !!(regs->flags & X86_EFLAGS_TF);
-
-       regs->flags |= X86_EFLAGS_TF;
-       if (test_tsk_thread_flag(task, TIF_BLOCKSTEP))
-               set_task_blockstep(task, false);
-}
-
-void arch_uprobe_disable_step(struct arch_uprobe *auprobe)
-{
-       struct task_struct *task = current;
-       struct arch_uprobe_task *autask = &task->utask->autask;
-       bool trapped = (task->utask->state == UTASK_SSTEP_TRAPPED);
-       struct pt_regs *regs = task_pt_regs(task);
-       /*
-        * The state of TIF_BLOCKSTEP was not saved so we can get an extra
-        * SIGTRAP if we do not clear TF. We need to examine the opcode to
-        * make it right.
-        */
-       if (unlikely(trapped)) {
-               if (!autask->saved_tf)
-                       regs->flags &= ~X86_EFLAGS_TF;
-       } else {
-               if (autask->saved_tf)
-                       send_sig(SIGTRAP, task, 0);
-               else if (!(auprobe->fixups & UPROBE_FIX_SETF))
-                       regs->flags &= ~X86_EFLAGS_TF;
-       }
-}