*/
 
 #include <linux/version.h>
+#include <linux/list.h>
 #include <linux/spinlock.h>
 #include <linux/hash.h>
 #include <linux/init.h>
 #include <linux/ptrace.h>
 #include <linux/preempt.h>
 #include <linux/percpu.h>
+#include <linux/kdebug.h>
 #include <asm/io.h>
 #include <asm/cacheflush.h>
 #include <asm/errno.h>
 #include <asm/tlbflush.h>
 #include <asm/pgtable.h>
 
-#include "kmmio.h"
+#include <linux/mmiotrace.h>
 
-#define KMMIO_HASH_BITS 6
-#define KMMIO_TABLE_SIZE (1 << KMMIO_HASH_BITS)
 #define KMMIO_PAGE_HASH_BITS 4
 #define KMMIO_PAGE_TABLE_SIZE (1 << KMMIO_PAGE_HASH_BITS)
 
+struct kmmio_fault_page {
+       struct list_head list;
+       struct kmmio_fault_page *release_next;
+       unsigned long page; /* location of the fault page */
+
+       /*
+        * Number of times this page has been registered as a part
+        * of a probe. If zero, page is disarmed and this may be freed.
+        * Used only by writers (RCU).
+        */
+       int count;
+};
+
+struct kmmio_delayed_release {
+       struct rcu_head rcu;
+       struct kmmio_fault_page *release_list;
+};
+
 struct kmmio_context {
        struct kmmio_fault_page *fpage;
        struct kmmio_probe *probe;
        unsigned long saved_flags;
+       unsigned long addr;
        int active;
 };
 
-static int kmmio_page_fault(struct pt_regs *regs, unsigned long error_code,
-                                               unsigned long address);
 static int kmmio_die_notifier(struct notifier_block *nb, unsigned long val,
                                                                void *args);
 
+static DECLARE_MUTEX(kmmio_init_mutex);
 static DEFINE_SPINLOCK(kmmio_lock);
 
 /* These are protected by kmmio_lock */
+static int kmmio_initialized;
 unsigned int kmmio_count;
-static unsigned int handler_registered;
+
+/* Read-protected by RCU, write-protected by kmmio_lock. */
 static struct list_head kmmio_page_table[KMMIO_PAGE_TABLE_SIZE];
 static LIST_HEAD(kmmio_probes);
 
+static struct list_head *kmmio_page_list(unsigned long page)
+{
+       return &kmmio_page_table[hash_long(page, KMMIO_PAGE_HASH_BITS)];
+}
+
 /* Accessed per-cpu */
 static DEFINE_PER_CPU(struct kmmio_context, kmmio_ctx);
 
+/* protected by kmmio_init_mutex */
 static struct notifier_block nb_die = {
        .notifier_call = kmmio_die_notifier
 };
 
-int init_kmmio(void)
+/**
+ * Makes sure kmmio is initialized and usable.
+ * This must be called before any other kmmio function defined here.
+ * May sleep.
+ */
+void reference_kmmio(void)
 {
-       int i;
-       for (i = 0; i < KMMIO_PAGE_TABLE_SIZE; i++)
-               INIT_LIST_HEAD(&kmmio_page_table[i]);
-
-       register_die_notifier(&nb_die);
-       return 0;
+       down(&kmmio_init_mutex);
+       spin_lock_irq(&kmmio_lock);
+       if (!kmmio_initialized) {
+               int i;
+               for (i = 0; i < KMMIO_PAGE_TABLE_SIZE; i++)
+                       INIT_LIST_HEAD(&kmmio_page_table[i]);
+               if (register_die_notifier(&nb_die))
+                       BUG();
+       }
+       kmmio_initialized++;
+       spin_unlock_irq(&kmmio_lock);
+       up(&kmmio_init_mutex);
 }
+EXPORT_SYMBOL_GPL(reference_kmmio);
 
-void cleanup_kmmio(void)
+/**
+ * Clean up kmmio after use. This must be called for every call to
+ * reference_kmmio(). All probes registered after the corresponding
+ * reference_kmmio() must have been unregistered when calling this.
+ * May sleep.
+ */
+void unreference_kmmio(void)
 {
-       /*
-        * Assume the following have been already cleaned by calling
-        * unregister_kmmio_probe() appropriately:
-        * kmmio_page_table, kmmio_probes
-        */
-       if (handler_registered) {
-               if (mmiotrace_unregister_pf(&kmmio_page_fault))
-                       BUG();
-               synchronize_rcu();
+       bool unreg = false;
+
+       down(&kmmio_init_mutex);
+       spin_lock_irq(&kmmio_lock);
+
+       if (kmmio_initialized == 1) {
+               BUG_ON(is_kmmio_active());
+               unreg = true;
        }
-       unregister_die_notifier(&nb_die);
+       kmmio_initialized--;
+       BUG_ON(kmmio_initialized < 0);
+       spin_unlock_irq(&kmmio_lock);
+
+       if (unreg)
+               unregister_die_notifier(&nb_die); /* calls sync_rcu() */
+       up(&kmmio_init_mutex);
 }
+EXPORT_SYMBOL(unreference_kmmio);
 
 /*
  * this is basically a dynamic stabbing problem:
  * Overlap a Point (might be simple)
  * Space Efficient Dynamic Stabbing with Fast Queries - Mikkel Thorup
  */
-/* Get the kmmio at this addr (if any). You must be holding kmmio_lock. */
+/* Get the kmmio at this addr (if any). You must be holding RCU read lock. */
 static struct kmmio_probe *get_kmmio_probe(unsigned long addr)
 {
        struct kmmio_probe *p;
-       list_for_each_entry(p, &kmmio_probes, list) {
+       list_for_each_entry_rcu(p, &kmmio_probes, list) {
                if (addr >= p->addr && addr <= (p->addr + p->len))
                        return p;
        }
        return NULL;
 }
 
+/* You must be holding RCU read lock. */
 static struct kmmio_fault_page *get_kmmio_fault_page(unsigned long page)
 {
-       struct list_head *head, *tmp;
+       struct list_head *head;
+       struct kmmio_fault_page *p;
 
        page &= PAGE_MASK;
-       head = &kmmio_page_table[hash_long(page, KMMIO_PAGE_HASH_BITS)];
-       list_for_each(tmp, head) {
-               struct kmmio_fault_page *p
-                       = list_entry(tmp, struct kmmio_fault_page, list);
+       head = kmmio_page_list(page);
+       list_for_each_entry_rcu(p, head, list) {
                if (p->page == page)
                        return p;
        }
-
        return NULL;
 }
 
+/** Mark the given page as not present. Access to it will trigger a fault. */
 static void arm_kmmio_fault_page(unsigned long page, int *page_level)
 {
        unsigned long address = page & PAGE_MASK;
        pte_t *pte = lookup_address(address, &level);
 
        if (!pte) {
-               printk(KERN_ERR "Error in %s: no pte for page 0x%08lx\n",
-                                               __FUNCTION__, page);
+               pr_err("kmmio: Error in %s: no pte for page 0x%08lx\n",
+                                                       __func__, page);
                return;
        }
 
        __flush_tlb_one(page);
 }
 
+/** Mark the given page as present. */
 static void disarm_kmmio_fault_page(unsigned long page, int *page_level)
 {
        unsigned long address = page & PAGE_MASK;
        pte_t *pte = lookup_address(address, &level);
 
        if (!pte) {
-               printk(KERN_ERR "Error in %s: no pte for page 0x%08lx\n",
-                                               __FUNCTION__, page);
+               pr_err("kmmio: Error in %s: no pte for page 0x%08lx\n",
+                                                       __func__, page);
                return;
        }
 
        __flush_tlb_one(page);
 }
 
+/*
+ * This is being called from do_page_fault().
+ *
+ * We may be in an interrupt or a critical section. Also prefecthing may
+ * trigger a page fault. We may be in the middle of process switch.
+ * We cannot take any locks, because we could be executing especially
+ * within a kmmio critical section.
+ *
+ * Local interrupts are disabled, so preemption cannot happen.
+ * Do not enable interrupts, do not sleep, and watch out for other CPUs.
+ */
 /*
  * Interrupts are disabled on entry as trap3 is an interrupt gate
  * and they remain disabled thorough out this function.
  */
-static int kmmio_handler(struct pt_regs *regs, unsigned long addr)
+int kmmio_handler(struct pt_regs *regs, unsigned long addr)
 {
-       struct kmmio_context *ctx = &get_cpu_var(kmmio_ctx);
+       struct kmmio_context *ctx;
+       struct kmmio_fault_page *faultpage;
 
        /*
         * Preemption is now disabled to prevent process switch during
         * XXX what if an interrupt occurs between returning from
         * do_page_fault() and entering the single-step exception handler?
         * And that interrupt triggers a kmmio trap?
+        * XXX If we tracing an interrupt service routine or whatever, is
+        * this enough to keep it on the current cpu?
         */
        preempt_disable();
 
-       /* interrupts disabled and CPU-local data => atomicity guaranteed. */
+       rcu_read_lock();
+       faultpage = get_kmmio_fault_page(addr);
+       if (!faultpage) {
+               /*
+                * Either this page fault is not caused by kmmio, or
+                * another CPU just pulled the kmmio probe from under
+                * our feet. In the latter case all hell breaks loose.
+                */
+               goto no_kmmio;
+       }
+
+       ctx = &get_cpu_var(kmmio_ctx);
        if (ctx->active) {
                /*
-                * This avoids a deadlock with kmmio_lock.
+                * Prevent overwriting already in-flight context.
                 * If this page fault really was due to kmmio trap,
                 * all hell breaks loose.
                 */
-               printk(KERN_EMERG "mmiotrace: recursive probe hit on CPU %d, "
-                                       "for address %lu. Ignoring.\n",
+               pr_emerg("kmmio: recursive probe hit on CPU %d, "
+                                       "for address 0x%08lx. Ignoring.\n",
                                        smp_processor_id(), addr);
-               goto no_kmmio;
+               goto no_kmmio_ctx;
        }
        ctx->active++;
 
-       /*
-        * Acquire the kmmio lock to prevent changes affecting
-        * get_kmmio_fault_page() and get_kmmio_probe(), since we save their
-        * returned pointers.
-        * The lock is released in post_kmmio_handler().
-        * XXX: could/should get_kmmio_*() be using RCU instead of spinlock?
-        */
-       spin_lock(&kmmio_lock);
-
-       ctx->fpage = get_kmmio_fault_page(addr);
-       if (!ctx->fpage) {
-               /* this page fault is not caused by kmmio */
-               goto no_kmmio_locked;
-       }
-
+       ctx->fpage = faultpage;
        ctx->probe = get_kmmio_probe(addr);
        ctx->saved_flags = (regs->flags & (TF_MASK|IF_MASK));
+       ctx->addr = addr;
 
        if (ctx->probe && ctx->probe->pre_handler)
                ctx->probe->pre_handler(ctx->probe, regs, addr);
        regs->flags |= TF_MASK;
        regs->flags &= ~IF_MASK;
 
-       /* We hold lock, now we set present bit in PTE and single step. */
+       /* Now we set present bit in PTE and single step. */
        disarm_kmmio_fault_page(ctx->fpage->page, NULL);
 
        put_cpu_var(kmmio_ctx);
+       rcu_read_unlock();
        return 1;
 
-no_kmmio_locked:
-       spin_unlock(&kmmio_lock);
-       ctx->active--;
+no_kmmio_ctx:
+       put_cpu_var(kmmio_ctx);
 no_kmmio:
+       rcu_read_unlock();
        preempt_enable_no_resched();
-       put_cpu_var(kmmio_ctx);
-       /* page fault not handled by kmmio */
-       return 0;
+       return 0; /* page fault not handled by kmmio */
 }
 
 /*
  * Interrupts are disabled on entry as trap1 is an interrupt gate
  * and they remain disabled thorough out this function.
- * And we hold kmmio lock.
+ * This must always get called as the pair to kmmio_handler().
  */
 static int post_kmmio_handler(unsigned long condition, struct pt_regs *regs)
 {
        int ret = 0;
+       struct kmmio_probe *probe;
+       struct kmmio_fault_page *faultpage;
        struct kmmio_context *ctx = &get_cpu_var(kmmio_ctx);
 
        if (!ctx->active)
                goto out;
 
+       rcu_read_lock();
+
+       faultpage = get_kmmio_fault_page(ctx->addr);
+       probe = get_kmmio_probe(ctx->addr);
+       if (faultpage != ctx->fpage || probe != ctx->probe) {
+               /*
+                * The trace setup changed after kmmio_handler() and before
+                * running this respective post handler. User does not want
+                * the result anymore.
+                */
+               ctx->probe = NULL;
+               ctx->fpage = NULL;
+       }
+
        if (ctx->probe && ctx->probe->post_handler)
                ctx->probe->post_handler(ctx->probe, condition, regs);
 
-       arm_kmmio_fault_page(ctx->fpage->page, NULL);
+       if (ctx->fpage)
+               arm_kmmio_fault_page(ctx->fpage->page, NULL);
 
        regs->flags &= ~TF_MASK;
        regs->flags |= ctx->saved_flags;
 
        /* These were acquired in kmmio_handler(). */
        ctx->active--;
-       spin_unlock(&kmmio_lock);
+       BUG_ON(ctx->active);
        preempt_enable_no_resched();
 
        /*
        if (!(regs->flags & TF_MASK))
                ret = 1;
 
+       rcu_read_unlock();
 out:
        put_cpu_var(kmmio_ctx);
        return ret;
 }
 
+/* You must be holding kmmio_lock. */
 static int add_kmmio_fault_page(unsigned long page)
 {
        struct kmmio_fault_page *f;
        page &= PAGE_MASK;
        f = get_kmmio_fault_page(page);
        if (f) {
+               if (!f->count)
+                       arm_kmmio_fault_page(f->page, NULL);
                f->count++;
                return 0;
        }
 
        f->count = 1;
        f->page = page;
-       list_add(&f->list,
-                &kmmio_page_table[hash_long(f->page, KMMIO_PAGE_HASH_BITS)]);
+       list_add_rcu(&f->list, kmmio_page_list(f->page));
 
        arm_kmmio_fault_page(f->page, NULL);
 
        return 0;
 }
 
-static void release_kmmio_fault_page(unsigned long page)
+/* You must be holding kmmio_lock. */
+static void release_kmmio_fault_page(unsigned long page,
+                               struct kmmio_fault_page **release_list)
 {
        struct kmmio_fault_page *f;
 
                return;
 
        f->count--;
+       BUG_ON(f->count < 0);
        if (!f->count) {
                disarm_kmmio_fault_page(f->page, NULL);
-               list_del(&f->list);
+               f->release_next = *release_list;
+               *release_list = f;
        }
 }
 
                ret = -EEXIST;
                goto out;
        }
-       list_add(&p->list, &kmmio_probes);
-       /*printk("adding fault pages...\n");*/
+       list_add_rcu(&p->list, &kmmio_probes);
        while (size < p->len) {
                if (add_kmmio_fault_page(p->addr + size))
-                       printk(KERN_ERR "mmio: Unable to set page fault.\n");
+                       pr_err("kmmio: Unable to set page fault.\n");
                size += PAGE_SIZE;
        }
-
-       if (!handler_registered) {
-               if (mmiotrace_register_pf(&kmmio_page_fault))
-                       printk(KERN_ERR "mmiotrace: Cannot register page "
-                                       "fault handler.\n");
-               else
-                       handler_registered++;
-       }
-
 out:
        spin_unlock_irq(&kmmio_lock);
        /*
         * XXX: What should I do here?
         * Here was a call to global_flush_tlb(), but it does not exist
-        * anymore.
+        * anymore. It seems it's not needed after all.
         */
        return ret;
 }
+EXPORT_SYMBOL(register_kmmio_probe);
 
+static void rcu_free_kmmio_fault_pages(struct rcu_head *head)
+{
+       struct kmmio_delayed_release *dr = container_of(
+                                               head,
+                                               struct kmmio_delayed_release,
+                                               rcu);
+       struct kmmio_fault_page *p = dr->release_list;
+       while (p) {
+               struct kmmio_fault_page *next = p->release_next;
+               BUG_ON(p->count);
+               kfree(p);
+               p = next;
+       }
+       kfree(dr);
+}
+
+static void remove_kmmio_fault_pages(struct rcu_head *head)
+{
+       struct kmmio_delayed_release *dr = container_of(
+                                               head,
+                                               struct kmmio_delayed_release,
+                                               rcu);
+       struct kmmio_fault_page *p = dr->release_list;
+       struct kmmio_fault_page **prevp = &dr->release_list;
+       unsigned long flags;
+       spin_lock_irqsave(&kmmio_lock, flags);
+       while (p) {
+               if (!p->count)
+                       list_del_rcu(&p->list);
+               else
+                       *prevp = p->release_next;
+               prevp = &p->release_next;
+               p = p->release_next;
+       }
+       spin_unlock_irqrestore(&kmmio_lock, flags);
+       /* This is the real RCU destroy call. */
+       call_rcu(&dr->rcu, rcu_free_kmmio_fault_pages);
+}
+
+/*
+ * Remove a kmmio probe. You have to synchronize_rcu() before you can be
+ * sure that the callbacks will not be called anymore.
+ *
+ * Unregistering a kmmio fault page has three steps:
+ * 1. release_kmmio_fault_page()
+ *    Disarm the page, wait a grace period to let all faults finish.
+ * 2. remove_kmmio_fault_pages()
+ *    Remove the pages from kmmio_page_table.
+ * 3. rcu_free_kmmio_fault_pages()
+ *    Actally free the kmmio_fault_page structs as with RCU.
+ */
 void unregister_kmmio_probe(struct kmmio_probe *p)
 {
        unsigned long size = 0;
+       struct kmmio_fault_page *release_list = NULL;
+       struct kmmio_delayed_release *drelease;
 
        spin_lock_irq(&kmmio_lock);
        while (size < p->len) {
-               release_kmmio_fault_page(p->addr + size);
+               release_kmmio_fault_page(p->addr + size, &release_list);
                size += PAGE_SIZE;
        }
-       list_del(&p->list);
+       list_del_rcu(&p->list);
        kmmio_count--;
        spin_unlock_irq(&kmmio_lock);
-}
 
-/*
- * According to 2.6.20, mainly x86_64 arch:
- * This is being called from do_page_fault(), via the page fault notifier
- * chain. The chain is called for both user space faults and kernel space
- * faults (address >= TASK_SIZE64), except not on faults serviced by
- * vmalloc_fault().
- *
- * We may be in an interrupt or a critical section. Also prefecthing may
- * trigger a page fault. We may be in the middle of process switch.
- * The page fault hook functionality has put us inside RCU read lock.
- *
- * Local interrupts are disabled, so preemption cannot happen.
- * Do not enable interrupts, do not sleep, and watch out for other CPUs.
- */
-static int kmmio_page_fault(struct pt_regs *regs, unsigned long error_code,
-                                               unsigned long address)
-{
-       if (is_kmmio_active())
-               if (kmmio_handler(regs, address) == 1)
-                       return -1;
-       return 0;
+       drelease = kmalloc(sizeof(*drelease), GFP_ATOMIC);
+       if (!drelease) {
+               pr_crit("kmmio: leaking kmmio_fault_page objects.\n");
+               return;
+       }
+       drelease->release_list = release_list;
+
+       /*
+        * This is not really RCU here. We have just disarmed a set of
+        * pages so that they cannot trigger page faults anymore. However,
+        * we cannot remove the pages from kmmio_page_table,
+        * because a probe hit might be in flight on another CPU. The
+        * pages are collected into a list, and they will be removed from
+        * kmmio_page_table when it is certain that no probe hit related to
+        * these pages can be in flight. RCU grace period sounds like a
+        * good choice.
+        *
+        * If we removed the pages too early, kmmio page fault handler might
+        * not find the respective kmmio_fault_page and determine it's not
+        * a kmmio fault, when it actually is. This would lead to madness.
+        */
+       call_rcu(&drelease->rcu, remove_kmmio_fault_pages);
 }
+EXPORT_SYMBOL(unregister_kmmio_probe);
 
 static int kmmio_die_notifier(struct notifier_block *nb, unsigned long val,
                                                                void *args)
 
 #include <asm/atomic.h>
 #include <linux/percpu.h>
 
-#include "kmmio.h"
 #include "pf_in.h"
 
 /* This app's relay channel files will appear in /debug/mmio-trace */
        pte_t *pte = lookup_address(address, &level);
 
        if (!pte) {
-               printk(KERN_ERR "Error in %s: no pte for page 0x%08lx\n",
-                                               __FUNCTION__, address);
+               pr_err(MODULE_NAME ": Error in %s: no pte for page 0x%08lx\n",
+                                                       __func__, address);
                return;
        }
 
        if (level == PG_LEVEL_2M) {
-               printk(KERN_EMERG MODULE_NAME ": 4MB pages are not "
-                                               "currently supported: %lx\n",
-                                               address);
+               pr_emerg(MODULE_NAME ": 4MB pages are not currently "
+                                               "supported: %lx\n", address);
                BUG();
        }
-       printk(KERN_DEBUG MODULE_NAME ": pte for 0x%lx: 0x%lx 0x%lx\n",
+       pr_info(MODULE_NAME ": pte for 0x%lx: 0x%lx 0x%lx\n",
                                        address, pte_val(*pte),
                                        pte_val(*pte) & _PAGE_PRESENT);
 }
 static void die_kmmio_nesting_error(struct pt_regs *regs, unsigned long addr)
 {
        const struct trap_reason *my_reason = &get_cpu_var(pf_reason);
-       printk(KERN_EMERG MODULE_NAME ": unexpected fault for address: %lx, "
+       pr_emerg(MODULE_NAME ": unexpected fault for address: %lx, "
                                        "last fault for address: %lx\n",
                                        addr, my_reason->addr);
        print_pte(addr);
        print_symbol(KERN_EMERG "faulting EIP is at %s\n", regs->ip);
        print_symbol(KERN_EMERG "last faulting EIP was at %s\n",
                                                        my_reason->ip);
-       printk(KERN_EMERG
-                       "eax: %08lx   ebx: %08lx   ecx: %08lx   edx: %08lx\n",
+       pr_emerg("eax: %08lx   ebx: %08lx   ecx: %08lx   edx: %08lx\n",
                        regs->ax, regs->bx, regs->cx, regs->dx);
-       printk(KERN_EMERG
-                       "esi: %08lx   edi: %08lx   ebp: %08lx   esp: %08lx\n",
+       pr_emerg("esi: %08lx   edi: %08lx   ebp: %08lx   esp: %08lx\n",
                        regs->si, regs->di, regs->bp, regs->sp);
 #else
        print_symbol(KERN_EMERG "faulting RIP is at %s\n", regs->ip);
        print_symbol(KERN_EMERG "last faulting RIP was at %s\n",
                                                        my_reason->ip);
-       printk(KERN_EMERG "rax: %016lx   rcx: %016lx   rdx: %016lx\n",
+       pr_emerg("rax: %016lx   rcx: %016lx   rdx: %016lx\n",
                                        regs->ax, regs->cx, regs->dx);
-       printk(KERN_EMERG "rsi: %016lx   rdi: %016lx   "
-                               "rbp: %016lx   rsp: %016lx\n",
+       pr_emerg("rsi: %016lx   rdi: %016lx   rbp: %016lx   rsp: %016lx\n",
                                regs->si, regs->di, regs->bp, regs->sp);
 #endif
        put_cpu_var(pf_reason);
        struct trap_reason *my_reason = &get_cpu_var(pf_reason);
        struct mm_io_header_rw *my_trace = &get_cpu_var(cpu_trace);
 
+       /*
+        * XXX: This might not get called, if the probe is removed while
+        * trace hit is on flight.
+        */
+
        /* this should always return the active_trace count to 0 */
        my_reason->active_traces--;
        if (my_reason->active_traces) {
-               printk(KERN_EMERG MODULE_NAME ": unexpected post handler");
+               pr_emerg(MODULE_NAME ": unexpected post handler");
                BUG();
        }
 
        atomic_t *drop = &per_cpu(dropped, cpu);
        int count;
        if (relay_buf_full(buf)) {
-               if (atomic_inc_return(drop) == 1) {
-                       printk(KERN_ERR MODULE_NAME ": cpu %d buffer full!\n",
-                                                                       cpu);
-               }
+               if (atomic_inc_return(drop) == 1)
+                       pr_err(MODULE_NAME ": cpu %d buffer full!\n", cpu);
                return 0;
-       } else if ((count = atomic_read(drop))) {
-               printk(KERN_ERR MODULE_NAME
-                                       ": cpu %d buffer no longer full, "
-                                       "missed %d events.\n",
-                                       cpu, count);
+       }
+       count = atomic_read(drop);
+       if (count) {
+               pr_err(MODULE_NAME ": cpu %d buffer no longer full, "
+                                               "missed %d events.\n",
+                                               cpu, count);
                atomic_sub(count, drop);
        }
 
        /* Don't trace the low PCI/ISA area, it's always mapped.. */
        if (!ISA_trace && (offset < ISA_END_ADDRESS) &&
                                        (offset + size > ISA_START_ADDRESS)) {
-               printk(KERN_NOTICE MODULE_NAME ": Ignoring map of low "
-                                               "PCI/ISA area (0x%lx-0x%lx)\n",
+               pr_notice(MODULE_NAME ": Ignoring map of low PCI/ISA area "
+                                               "(0x%lx-0x%lx)\n",
                                                offset, offset + size);
                return;
        }
 void __iomem *ioremap_cache_trace(unsigned long offset, unsigned long size)
 {
        void __iomem *p = ioremap_cache(offset, size);
-       printk(KERN_DEBUG MODULE_NAME ": ioremap_cache(0x%lx, 0x%lx) = %p\n",
+       pr_debug(MODULE_NAME ": ioremap_cache(0x%lx, 0x%lx) = %p\n",
                                                        offset, size, p);
        ioremap_trace_core(offset, size, p);
        return p;
 void __iomem *ioremap_nocache_trace(unsigned long offset, unsigned long size)
 {
        void __iomem *p = ioremap_nocache(offset, size);
-       printk(KERN_DEBUG MODULE_NAME ": ioremap_nocache(0x%lx, 0x%lx) = %p\n",
+       pr_debug(MODULE_NAME ": ioremap_nocache(0x%lx, 0x%lx) = %p\n",
                                                        offset, size, p);
        ioremap_trace_core(offset, size, p);
        return p;
        };
        struct remap_trace *trace;
        struct remap_trace *tmp;
-       printk(KERN_DEBUG MODULE_NAME ": Unmapping %p.\n", addr);
+       pr_debug(MODULE_NAME ": Unmapping %p.\n", addr);
        record_timestamp(&event.header);
 
        spin_lock(&trace_list_lock);
 
        spin_lock(&trace_list_lock);
        list_for_each_entry_safe(trace, tmp, &trace_list, list) {
-               printk(KERN_WARNING MODULE_NAME ": purging non-iounmapped "
+               pr_warning(MODULE_NAME ": purging non-iounmapped "
                                        "trace @0x%08lx, size 0x%lx.\n",
                                        trace->probe.addr, trace->probe.len);
                if (!nommiotrace)
 
        dir = debugfs_create_dir(APP_DIR, NULL);
        if (!dir) {
-               printk(KERN_ERR MODULE_NAME
-                               ": Couldn't create relay app directory.\n");
+               pr_err(MODULE_NAME ": Couldn't create relay app directory.\n");
                return -ENOMEM;
        }
 
        chan = create_channel(subbuf_size, n_subbufs);
        if (!chan) {
                debugfs_remove(dir);
-               printk(KERN_ERR MODULE_NAME
-                               ": relay app channel creation failed\n");
+               pr_err(MODULE_NAME ": relay app channel creation failed\n");
                return -ENOMEM;
        }
 
-       init_kmmio();
+       reference_kmmio();
 
        proc_marker_file = create_proc_entry(MARKER_FILE, 0, NULL);
        if (proc_marker_file)
                proc_marker_file->write_proc = write_marker;
 
-       printk(KERN_DEBUG MODULE_NAME ": loaded.\n");
+       pr_debug(MODULE_NAME ": loaded.\n");
        if (nommiotrace)
-               printk(KERN_DEBUG MODULE_NAME ": MMIO tracing disabled.\n");
+               pr_info(MODULE_NAME ": MMIO tracing disabled.\n");
        if (ISA_trace)
-               printk(KERN_WARNING MODULE_NAME
-                               ": Warning! low ISA range will be traced.\n");
+               pr_warning(MODULE_NAME ": Warning! low ISA range will be "
+                                                               "traced.\n");
        return 0;
 }
 
 static void __exit cleanup(void)
 {
-       printk(KERN_DEBUG MODULE_NAME ": unload...\n");
+       pr_debug(MODULE_NAME ": unload...\n");
        clear_trace_list();
-       cleanup_kmmio();
+       unreference_kmmio();
        remove_proc_entry(MARKER_FILE, NULL);
        destroy_channel();
        if (dir)