2 files changed, 9 insertions, 3 deletions

diff --git a/virt/kvm/iommu.c b/virt/kvm/iommu.c
index a457d2138f49..fec1723de9b4 100644
--- a/virt/kvm/iommu.c
+++ b/virt/kvm/iommu.c
@@ -310,6 +310,11 @@ static void kvm_iommu_put_pages(struct kvm *kvm,
         }
 }
 
+void kvm_iommu_unmap_pages(struct kvm *kvm, struct kvm_memory_slot *slot)
+{
+        kvm_iommu_put_pages(kvm, slot->base_gfn, slot->npages);
+}
+
 static int kvm_iommu_unmap_memslots(struct kvm *kvm)
 {
         int idx;
@@ -320,7 +325,7 @@ static int kvm_iommu_unmap_memslots(struct kvm *kvm)
         slots = kvm_memslots(kvm);
 
         kvm_for_each_memslot(memslot, slots)
-                kvm_iommu_put_pages(kvm, memslot->base_gfn, memslot->npages);
+                kvm_iommu_unmap_pages(kvm, memslot);
 
         srcu_read_unlock(&kvm->srcu, idx);
 
diff --git a/virt/kvm/kvm_main.c b/virt/kvm/kvm_main.c
index 42b73930a6de..9739b533ca2e 100644
--- a/virt/kvm/kvm_main.c
+++ b/virt/kvm/kvm_main.c
@@ -808,12 +808,13 @@ int __kvm_set_memory_region(struct kvm *kvm,
         if (r)
                 goto out_free;
 
-        /* map the pages in iommu page table */
+        /* map/unmap the pages in iommu page table */
         if (npages) {
                 r = kvm_iommu_map_pages(kvm, &new);
                 if (r)
                         goto out_free;
-        }
+        } else
+                kvm_iommu_unmap_pages(kvm, &old);
 
         r = -ENOMEM;
         slots = kmemdup(kvm->memslots, sizeof(struct kvm_memslots),