diff --git a/xen/common/page_alloc.c b/xen/common/page_alloc.c
index e866e0d864..ea59cd1a4a 100644
--- a/xen/common/page_alloc.c
+++ b/xen/common/page_alloc.c
@@ -162,13 +162,6 @@
 static char __initdata opt_badpage[100] = "";
 string_param("badpage", opt_badpage);
 
-/*
- * Heap allocations may need TLB flushes which may require IRQs to be
- * enabled (except when only 1 PCPU is online).
- */
-#define ASSERT_ALLOC_CONTEXT() \
-    ASSERT(!in_irq() && (local_irq_is_enabled() || num_online_cpus() <= 1))
-
 /*
  * no-bootscrub -> Free pages are not zeroed during boot.
  */
diff --git a/xen/common/xmalloc_tlsf.c b/xen/common/xmalloc_tlsf.c
index d2ad909502..75bdf18c4e 100644
--- a/xen/common/xmalloc_tlsf.c
+++ b/xen/common/xmalloc_tlsf.c
@@ -378,6 +378,8 @@ void *xmem_pool_alloc(unsigned long size, struct xmem_pool *pool)
     int fl, sl;
     unsigned long tmp_size;
 
+    ASSERT_ALLOC_CONTEXT();
+
     if ( size < MIN_BLOCK_SIZE )
         size = MIN_BLOCK_SIZE;
     else
@@ -456,6 +458,8 @@ void xmem_pool_free(void *ptr, struct xmem_pool *pool)
     struct bhdr *b, *tmp_b;
     int fl = 0, sl = 0;
 
+    ASSERT_ALLOC_CONTEXT();
+
     if ( unlikely(ptr == NULL) )
         return;
 
@@ -594,7 +598,7 @@ void *_xmalloc(unsigned long size, unsigned long align)
 {
     void *p = NULL;
 
-    ASSERT(!in_irq());
+    ASSERT_ALLOC_CONTEXT();
 
     if ( !size )
         return ZERO_BLOCK_PTR;
@@ -697,11 +701,11 @@ void *_xrealloc(void *ptr, unsigned long size, unsigned long align)
 
 void xfree(void *p)
 {
+    ASSERT_ALLOC_CONTEXT();
+
     if ( p == NULL || p == ZERO_BLOCK_PTR )
         return;
 
-    ASSERT(!in_irq());
-
     if ( !((unsigned long)p & (PAGE_SIZE - 1)) )
     {
         unsigned long size = PFN_ORDER(virt_to_page(p));
diff --git a/xen/include/xen/irq.h b/xen/include/xen/irq.h
index d8beadd16b..300625e56d 100644
--- a/xen/include/xen/irq.h
+++ b/xen/include/xen/irq.h
@@ -10,6 +10,13 @@
 #include <asm/hardirq.h>
 #include <public/event_channel.h>
 
+/*
+ * Heap allocations may need TLB flushes which may require IRQs to be
+ * enabled (except when only 1 PCPU is online).
+ */
+#define ASSERT_ALLOC_CONTEXT() \
+    ASSERT(!in_irq() && (local_irq_is_enabled() || num_online_cpus() <= 1))
+
 struct irqaction {
     void (*handler)(int, void *, struct cpu_user_regs *);
     const char *name;