1 files changed, 22 insertions, 10 deletions

diff --git a/virt/kvm/kvm_main.c b/virt/kvm/kvm_main.c
index 7e140683ff14..1148c96a4817 100644
--- a/virt/kvm/kvm_main.c
+++ b/virt/kvm/kvm_main.c
@@ -516,16 +516,32 @@ out_err_nodisable:
         return ERR_PTR(r);
 }
 
+/*
+ * Avoid using vmalloc for a small buffer.
+ * Should not be used when the size is statically known.
+ */
+static void *kvm_kvzalloc(unsigned long size)
+{
+        if (size > PAGE_SIZE)
+                return vzalloc(size);
+        else
+                return kzalloc(size, GFP_KERNEL);
+}
+
+static void kvm_kvfree(const void *addr)
+{
+        if (is_vmalloc_addr(addr))
+                vfree(addr);
+        else
+                kfree(addr);
+}
+
 static void kvm_destroy_dirty_bitmap(struct kvm_memory_slot *memslot)
 {
         if (!memslot->dirty_bitmap)
                 return;
 
-        if (2 * kvm_dirty_bitmap_bytes(memslot) > PAGE_SIZE)
-                vfree(memslot->dirty_bitmap);
-        else
-                kfree(memslot->dirty_bitmap);
-
+        kvm_kvfree(memslot->dirty_bitmap);
         memslot->dirty_bitmap = NULL;
 }
 
@@ -617,11 +633,7 @@ static int kvm_create_dirty_bitmap(struct kvm_memory_slot *memslot)
 #ifndef CONFIG_S390
         unsigned long dirty_bytes = 2 * kvm_dirty_bitmap_bytes(memslot);
 
-        if (dirty_bytes > PAGE_SIZE)
-                memslot->dirty_bitmap = vzalloc(dirty_bytes);
-        else
-                memslot->dirty_bitmap = kzalloc(dirty_bytes, GFP_KERNEL);
-
+        memslot->dirty_bitmap = kvm_kvzalloc(dirty_bytes);
         if (!memslot->dirty_bitmap)
                 return -ENOMEM;