#include <linux/hardirq.h>
 #include <linux/thread_info.h>
 #include <linux/mm.h>
+#include <linux/oom.h>
+#include <linux/sched.h>
 #include <linux/uaccess.h>
 #include <linux/module.h>
 #include <linux/kallsyms.h>
        struct task_struct *p;
        struct mm_struct *mm;
        unsigned long flags, offset;
-       unsigned char in_atomic = (bfin_read_IPEND() & 0x10) || in_atomic();
        struct rb_node *n;
 
 #ifdef CONFIG_KALLSYMS
         */
        write_lock_irqsave(&tasklist_lock, flags);
        for_each_process(p) {
-               mm = (in_atomic ? p->mm : get_task_mm(p));
-               if (!mm)
-                       continue;
+               struct task_struct *t;
 
-               if (!down_read_trylock(&mm->mmap_sem)) {
-                       if (!in_atomic)
-                               mmput(mm);
+               t = find_lock_task_mm(p);
+               if (!t)
                        continue;
-               }
+
+               mm = t->mm;
+               if (!down_read_trylock(&mm->mmap_sem))
+                       goto __continue;
 
                for (n = rb_first(&mm->mm_rb); n; n = rb_next(n)) {
                        struct vm_area_struct *vma;
 
                        if (address >= vma->vm_start && address < vma->vm_end) {
                                char _tmpbuf[256];
-                               char *name = p->comm;
+                               char *name = t->comm;
                                struct file *file = vma->vm_file;
 
                                if (file) {
                                                name, vma->vm_start, vma->vm_end);
 
                                up_read(&mm->mmap_sem);
-                               if (!in_atomic)
-                                       mmput(mm);
+                               task_unlock(t);
 
                                if (buf[0] == '\0')
                                        sprintf(buf, "[ %s ] dynamic memory", name);
                }
 
                up_read(&mm->mmap_sem);
-               if (!in_atomic)
-                       mmput(mm);
+__continue:
+               task_unlock(t);
        }
 
        /*