#include <linux/kdebug.h>
 #include <linux/debugfs.h>
 #include <linux/ratelimit.h>
+#include <linux/context_tracking.h>
 
 #include <asm/emulated_ops.h>
 #include <asm/pgtable.h>
 
 void machine_check_exception(struct pt_regs *regs)
 {
+       enum ctx_state prev_state = exception_enter();
        int recover = 0;
 
        __get_cpu_var(irq_stat).mce_exceptions++;
                recover = cur_cpu_spec->machine_check(regs);
 
        if (recover > 0)
-               return;
+               goto bail;
 
 #if defined(CONFIG_8xx) && defined(CONFIG_PCI)
        /* the qspan pci read routines can cause machine checks -- Cort
         * -- BenH
         */
        bad_page_fault(regs, regs->dar, SIGBUS);
-       return;
+       goto bail;
 #endif
 
        if (debugger_fault_handler(regs))
-               return;
+               goto bail;
 
        if (check_io_access(regs))
-               return;
+               goto bail;
 
        die("Machine check", regs, SIGBUS);
 
        /* Must die if the interrupt is not recoverable */
        if (!(regs->msr & MSR_RI))
                panic("Unrecoverable Machine check");
+
+bail:
+       exception_exit(prev_state);
 }
 
 void SMIException(struct pt_regs *regs)
 
 void unknown_exception(struct pt_regs *regs)
 {
+       enum ctx_state prev_state = exception_enter();
+
        printk("Bad trap at PC: %lx, SR: %lx, vector=%lx\n",
               regs->nip, regs->msr, regs->trap);
 
        _exception(SIGTRAP, regs, 0, 0);
+
+       exception_exit(prev_state);
 }
 
 void instruction_breakpoint_exception(struct pt_regs *regs)
 {
+       enum ctx_state prev_state = exception_enter();
+
        if (notify_die(DIE_IABR_MATCH, "iabr_match", regs, 5,
                                        5, SIGTRAP) == NOTIFY_STOP)
-               return;
+               goto bail;
        if (debugger_iabr_match(regs))
-               return;
+               goto bail;
        _exception(SIGTRAP, regs, TRAP_BRKPT, regs->nip);
+
+bail:
+       exception_exit(prev_state);
 }
 
 void RunModeException(struct pt_regs *regs)
 
 void __kprobes single_step_exception(struct pt_regs *regs)
 {
+       enum ctx_state prev_state = exception_enter();
+
        clear_single_step(regs);
 
        if (notify_die(DIE_SSTEP, "single_step", regs, 5,
                                        5, SIGTRAP) == NOTIFY_STOP)
-               return;
+               goto bail;
        if (debugger_sstep(regs))
-               return;
+               goto bail;
 
        _exception(SIGTRAP, regs, TRAP_TRACE, regs->nip);
+
+bail:
+       exception_exit(prev_state);
 }
 
 /*
 
 void __kprobes program_check_exception(struct pt_regs *regs)
 {
+       enum ctx_state prev_state = exception_enter();
        unsigned int reason = get_reason(regs);
        extern int do_mathemu(struct pt_regs *regs);
 
        if (reason & REASON_FP) {
                /* IEEE FP exception */
                parse_fpe(regs);
-               return;
+               goto bail;
        }
        if (reason & REASON_TRAP) {
                /* Debugger is first in line to stop recursive faults in
                 * rcu_lock, notify_die, or atomic_notifier_call_chain */
                if (debugger_bpt(regs))
-                       return;
+                       goto bail;
 
                /* trap exception */
                if (notify_die(DIE_BPT, "breakpoint", regs, 5, 5, SIGTRAP)
                                == NOTIFY_STOP)
-                       return;
+                       goto bail;
 
                if (!(regs->msr & MSR_PR) &&  /* not user-mode */
                    report_bug(regs->nip, regs) == BUG_TRAP_TYPE_WARN) {
                        regs->nip += 4;
-                       return;
+                       goto bail;
                }
                _exception(SIGTRAP, regs, TRAP_BRKPT, regs->nip);
-               return;
+               goto bail;
        }
 #ifdef CONFIG_PPC_TRANSACTIONAL_MEM
        if (reason & REASON_TM) {
                if (!user_mode(regs) &&
                    report_bug(regs->nip, regs) == BUG_TRAP_TYPE_WARN) {
                        regs->nip += 4;
-                       return;
+                       goto bail;
                }
                /* If usermode caused this, it's done something illegal and
                 * gets a SIGILL slap on the wrist.  We call it an illegal
                 */
                if (user_mode(regs)) {
                        _exception(SIGILL, regs, ILL_ILLOPN, regs->nip);
-                       return;
+                       goto bail;
                } else {
                        printk(KERN_EMERG "Unexpected TM Bad Thing exception "
                               "at %lx (msr 0x%x)\n", regs->nip, reason);
        switch (do_mathemu(regs)) {
        case 0:
                emulate_single_step(regs);
-               return;
+               goto bail;
        case 1: {
                        int code = 0;
                        code = __parse_fpscr(current->thread.fpscr.val);
                        _exception(SIGFPE, regs, code, regs->nip);
-                       return;
+                       goto bail;
                }
        case -EFAULT:
                _exception(SIGSEGV, regs, SEGV_MAPERR, regs->nip);
-               return;
+               goto bail;
        }
        /* fall through on any other errors */
 #endif /* CONFIG_MATH_EMULATION */
                case 0:
                        regs->nip += 4;
                        emulate_single_step(regs);
-                       return;
+                       goto bail;
                case -EFAULT:
                        _exception(SIGSEGV, regs, SEGV_MAPERR, regs->nip);
-                       return;
+                       goto bail;
                }
        }
 
                _exception(SIGILL, regs, ILL_PRVOPC, regs->nip);
        else
                _exception(SIGILL, regs, ILL_ILLOPC, regs->nip);
+
+bail:
+       exception_exit(prev_state);
 }
 
 void alignment_exception(struct pt_regs *regs)
 {
+       enum ctx_state prev_state = exception_enter();
        int sig, code, fixed = 0;
 
        /* We restore the interrupt state now */
        if (fixed == 1) {
                regs->nip += 4; /* skip over emulated instruction */
                emulate_single_step(regs);
-               return;
+               goto bail;
        }
 
        /* Operand address was bad */
                _exception(sig, regs, code, regs->dar);
        else
                bad_page_fault(regs, regs->dar, sig);
+
+bail:
+       exception_exit(prev_state);
 }
 
 void StackOverflow(struct pt_regs *regs)
 
 void kernel_fp_unavailable_exception(struct pt_regs *regs)
 {
+       enum ctx_state prev_state = exception_enter();
+
        printk(KERN_EMERG "Unrecoverable FP Unavailable Exception "
                          "%lx at %lx\n", regs->trap, regs->nip);
        die("Unrecoverable FP Unavailable Exception", regs, SIGABRT);
+
+       exception_exit(prev_state);
 }
 
 void altivec_unavailable_exception(struct pt_regs *regs)
 {
+       enum ctx_state prev_state = exception_enter();
+
        if (user_mode(regs)) {
                /* A user program has executed an altivec instruction,
                   but this kernel doesn't support altivec. */
                _exception(SIGILL, regs, ILL_ILLOPC, regs->nip);
-               return;
+               goto bail;
        }
 
        printk(KERN_EMERG "Unrecoverable VMX/Altivec Unavailable Exception "
                        "%lx at %lx\n", regs->trap, regs->nip);
        die("Unrecoverable VMX/Altivec Unavailable Exception", regs, SIGABRT);
+
+bail:
+       exception_exit(prev_state);
 }
 
 void vsx_unavailable_exception(struct pt_regs *regs)
 
 #include <linux/perf_event.h>
 #include <linux/magic.h>
 #include <linux/ratelimit.h>
+#include <linux/context_tracking.h>
 
 #include <asm/firmware.h>
 #include <asm/page.h>
 int __kprobes do_page_fault(struct pt_regs *regs, unsigned long address,
                            unsigned long error_code)
 {
+       enum ctx_state prev_state = exception_enter();
        struct vm_area_struct * vma;
        struct mm_struct *mm = current->mm;
        unsigned int flags = FAULT_FLAG_ALLOW_RETRY | FAULT_FLAG_KILLABLE;
        int trap = TRAP(regs);
        int is_exec = trap == 0x400;
        int fault;
+       int rc = 0;
 
 #if !(defined(CONFIG_4xx) || defined(CONFIG_BOOKE))
        /*
         * look at it
         */
        if (error_code & ICSWX_DSI_UCT) {
-               int rc = acop_handle_fault(regs, address, error_code);
+               rc = acop_handle_fault(regs, address, error_code);
                if (rc)
-                       return rc;
+                       goto bail;
        }
 #endif /* CONFIG_PPC_ICSWX */
 
        if (notify_page_fault(regs))
-               return 0;
+               goto bail;
 
        if (unlikely(debugger_fault_handler(regs)))
-               return 0;
+               goto bail;
 
        /* On a kernel SLB miss we can only check for a valid exception entry */
-       if (!user_mode(regs) && (address >= TASK_SIZE))
-               return SIGSEGV;
+       if (!user_mode(regs) && (address >= TASK_SIZE)) {
+               rc = SIGSEGV;
+               goto bail;
+       }
 
 #if !(defined(CONFIG_4xx) || defined(CONFIG_BOOKE) || \
                             defined(CONFIG_PPC_BOOK3S_64))
        if (error_code & DSISR_DABRMATCH) {
                /* breakpoint match */
                do_break(regs, address, error_code);
-               return 0;
+               goto bail;
        }
 #endif
 
                local_irq_enable();
 
        if (in_atomic() || mm == NULL) {
-               if (!user_mode(regs))
-                       return SIGSEGV;
+               if (!user_mode(regs)) {
+                       rc = SIGSEGV;
+                       goto bail;
+               }
                /* in_atomic() in user mode is really bad,
                   as is current->mm == NULL. */
                printk(KERN_EMERG "Page fault in user mode with "
         */
        fault = handle_mm_fault(mm, vma, address, flags);
        if (unlikely(fault & (VM_FAULT_RETRY|VM_FAULT_ERROR))) {
-               int rc = mm_fault_error(regs, address, fault);
+               rc = mm_fault_error(regs, address, fault);
                if (rc >= MM_FAULT_RETURN)
-                       return rc;
+                       goto bail;
+               else
+                       rc = 0;
        }
 
        /*
        }
 
        up_read(&mm->mmap_sem);
-       return 0;
+       goto bail;
 
 bad_area:
        up_read(&mm->mmap_sem);
        /* User mode accesses cause a SIGSEGV */
        if (user_mode(regs)) {
                _exception(SIGSEGV, regs, code, address);
-               return 0;
+               goto bail;
        }
 
        if (is_exec && (error_code & DSISR_PROTFAULT))
                                   " page (%lx) - exploit attempt? (uid: %d)\n",
                                   address, from_kuid(&init_user_ns, current_uid()));
 
-       return SIGSEGV;
+       rc = SIGSEGV;
+
+bail:
+       exception_exit(prev_state);
+       return rc;
 
 }
 
 
 #include <linux/init.h>
 #include <linux/signal.h>
 #include <linux/memblock.h>
+#include <linux/context_tracking.h>
 
 #include <asm/processor.h>
 #include <asm/pgtable.h>
  */
 int hash_page(unsigned long ea, unsigned long access, unsigned long trap)
 {
+       enum ctx_state prev_state = exception_enter();
        pgd_t *pgdir;
        unsigned long vsid;
        struct mm_struct *mm;
                mm = current->mm;
                if (! mm) {
                        DBG_LOW(" user region with no mm !\n");
-                       return 1;
+                       rc = 1;
+                       goto bail;
                }
                psize = get_slice_psize(mm, ea);
                ssize = user_segment_size(ea);
                /* Not a valid range
                 * Send the problem up to do_page_fault 
                 */
-               return 1;
+               rc = 1;
+               goto bail;
        }
        DBG_LOW(" mm=%p, mm->pgdir=%p, vsid=%016lx\n", mm, mm->pgd, vsid);
 
        /* Bad address. */
        if (!vsid) {
                DBG_LOW("Bad address!\n");
-               return 1;
+               rc = 1;
+               goto bail;
        }
        /* Get pgdir */
        pgdir = mm->pgd;
-       if (pgdir == NULL)
-               return 1;
+       if (pgdir == NULL) {
+               rc = 1;
+               goto bail;
+       }
 
        /* Check CPU locality */
        tmp = cpumask_of(smp_processor_id());
        ptep = find_linux_pte_or_hugepte(pgdir, ea, &hugeshift);
        if (ptep == NULL || !pte_present(*ptep)) {
                DBG_LOW(" no PTE !\n");
-               return 1;
+               rc = 1;
+               goto bail;
        }
 
        /* Add _PAGE_PRESENT to the required access perm */
         */
        if (access & ~pte_val(*ptep)) {
                DBG_LOW(" no access !\n");
-               return 1;
+               rc = 1;
+               goto bail;
        }
 
 #ifdef CONFIG_HUGETLB_PAGE
-       if (hugeshift)
-               return __hash_page_huge(ea, access, vsid, ptep, trap, local,
+       if (hugeshift) {
+               rc = __hash_page_huge(ea, access, vsid, ptep, trap, local,
                                        ssize, hugeshift, psize);
+               goto bail;
+       }
 #endif /* CONFIG_HUGETLB_PAGE */
 
 #ifndef CONFIG_PPC_64K_PAGES
                pte_val(*(ptep + PTRS_PER_PTE)));
 #endif
        DBG_LOW(" -> rc=%d\n", rc);
+
+bail:
+       exception_exit(prev_state);
        return rc;
 }
 EXPORT_SYMBOL_GPL(hash_page);
  */
 void low_hash_fault(struct pt_regs *regs, unsigned long address, int rc)
 {
+       enum ctx_state prev_state = exception_enter();
+
        if (user_mode(regs)) {
 #ifdef CONFIG_PPC_SUBPAGE_PROT
                if (rc == -2)
                        _exception(SIGBUS, regs, BUS_ADRERR, address);
        } else
                bad_page_fault(regs, address, SIGBUS);
+
+       exception_exit(prev_state);
 }
 
 long hpte_insert_repeating(unsigned long hash, unsigned long vpn,