#include <linux/init.h>
 #include <linux/interrupt.h>
+#include <linux/uaccess.h>
 #include <asm/kmap_types.h>
 
 /* undef for production */
 
 static inline void *kmap_atomic(struct page *page, enum km_type type)
 {
+       pagefault_disable();
        return page_address(page);
 }
 
-static inline void kunmap_atomic(void *kvaddr, enum km_type type) { }
-#define kmap_atomic_pfn(pfn, idx)      page_address(pfn_to_page(pfn))
+static inline void kunmap_atomic(void *kvaddr, enum km_type type)
+{
+       pagefault_enable();
+}
+
+#define kmap_atomic_pfn(pfn, idx) kmap_atomic(pfn_to_page(pfn), (idx))
 
 #define kmap_atomic_to_page(ptr) virt_to_page(ptr)
 
 
 
 #include <linux/fs.h>
 #include <linux/mm.h>
+#include <linux/uaccess.h>
 
 #include <asm/cacheflush.h>
 
 
 #define kunmap(page) do { (void) (page); } while (0)
 
-#define kmap_atomic(page, idx)         page_address(page)
-#define kunmap_atomic(addr, idx)       do { } while (0)
-#define kmap_atomic_pfn(pfn, idx)      page_address(pfn_to_page(pfn))
+#define kmap_atomic(page, idx) \
+       ({ pagefault_disable(); page_address(page); })
+#define kunmap_atomic(addr, idx)       do { pagefault_enable(); } while (0)
+#define kmap_atomic_pfn(pfn, idx)      kmap_atomic(pfn_to_page(pfn), (idx))
 #define kmap_atomic_to_page(ptr)       virt_to_page(ptr)
 #endif