diff --git a/include/linux/kvm_host.h b/include/linux/kvm_host.h
index 7e8a29f..6136821 100644
--- a/include/linux/kvm_host.h
+++ b/include/linux/kvm_host.h
@@ -562,6 +562,7 @@ void kvm_free_irq_source_id(struct kvm *kvm, int
irq_source_id);

 #ifdef CONFIG_IOMMU_API
 int kvm_iommu_map_pages(struct kvm *kvm, struct kvm_memory_slot *slot);
+void kvm_iommu_unmap_pages(struct kvm *kvm, struct kvm_memory_slot *slot);
 int kvm_iommu_map_guest(struct kvm *kvm);
 int kvm_iommu_unmap_guest(struct kvm *kvm);
 int kvm_assign_device(struct kvm *kvm,
@@ -575,6 +576,11 @@ static inline int kvm_iommu_map_pages(struct kvm *kvm,
 	return 0;
 }

+static inline void kvm_iommu_unmap_pages(struct kvm *kvm,
+					 struct kvm_memory_slot *slot)
+{
+}
+
 static inline int kvm_iommu_map_guest(struct kvm *kvm)
 {
 	return -ENODEV;
diff --git a/virt/kvm/iommu.c b/virt/kvm/iommu.c
index a195c07..981f55a 100644
--- a/virt/kvm/iommu.c
+++ b/virt/kvm/iommu.c
@@ -309,6 +309,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 i, idx;
@@ -318,8 +323,7 @@ static int kvm_iommu_unmap_memslots(struct kvm *kvm)
 	slots = kvm_memslots(kvm);

 	for (i = 0; i < slots->nmemslots; i++) {
-		kvm_iommu_put_pages(kvm, slots->memslots[i].base_gfn,
-				    slots->memslots[i].npages);
+		kvm_iommu_unmap_pages(kvm, &slots->memslots[i]);
 	}
 	srcu_read_unlock(&kvm->srcu, idx);

diff --git a/virt/kvm/kvm_main.c b/virt/kvm/kvm_main.c
index 5007b7d..63e01de 100644
--- a/virt/kvm/kvm_main.c
+++ b/virt/kvm/kvm_main.c
@@ -802,12 +802,13 @@ skip_lpage:
 	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 = kzalloc(sizeof(struct kvm_memslots), GFP_KERNEL);