3 files changed, 8 insertions, 10 deletions

diff --git a/include/linux/kvm_host.h b/include/linux/kvm_host.h
index f1f78deece10..9af240387fe6 100644
--- a/include/linux/kvm_host.h
+++ b/include/linux/kvm_host.h
@@ -440,8 +440,7 @@ void kvm_free_irq_source_id(struct kvm *kvm, int irq_source_id);
 #define KVM_IOMMU_CACHE_COHERENCY       0x1
 
 #ifdef CONFIG_IOMMU_API
-int kvm_iommu_map_pages(struct kvm *kvm, gfn_t base_gfn,
-                        unsigned long npages);
+int kvm_iommu_map_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,
diff --git a/virt/kvm/iommu.c b/virt/kvm/iommu.c
index bc697a66a883..cf567d8033db 100644
--- a/virt/kvm/iommu.c
+++ b/virt/kvm/iommu.c
@@ -32,10 +32,10 @@ static int kvm_iommu_unmap_memslots(struct kvm *kvm);
 static void kvm_iommu_put_pages(struct kvm *kvm,
                                 gfn_t base_gfn, unsigned long npages);
 
-int kvm_iommu_map_pages(struct kvm *kvm,
-                        gfn_t base_gfn, unsigned long npages)
+int kvm_iommu_map_pages(struct kvm *kvm, struct kvm_memory_slot *slot)
 {
-        gfn_t gfn = base_gfn;
+        gfn_t gfn = slot->base_gfn;
+        unsigned long npages = slot->npages;
         pfn_t pfn;
         int i, r = 0;
         struct iommu_domain *domain = kvm->arch.iommu_domain;
@@ -54,7 +54,7 @@ int kvm_iommu_map_pages(struct kvm *kvm,
                 if (iommu_iova_to_phys(domain, gfn_to_gpa(gfn)))
                         continue;
 
-                pfn = gfn_to_pfn(kvm, gfn);
+                pfn = gfn_to_pfn_memslot(kvm, slot, gfn);
                 r = iommu_map_range(domain,
                                     gfn_to_gpa(gfn),
                                     pfn_to_hpa(pfn),
@@ -69,7 +69,7 @@ int kvm_iommu_map_pages(struct kvm *kvm,
         return 0;
 
 unmap_pages:
-        kvm_iommu_put_pages(kvm, base_gfn, i);
+        kvm_iommu_put_pages(kvm, slot->base_gfn, i);
         return r;
 }
 
@@ -81,8 +81,7 @@ static int kvm_iommu_map_memslots(struct kvm *kvm)
         slots = kvm->memslots;
 
         for (i = 0; i < slots->nmemslots; i++) {
-                r = kvm_iommu_map_pages(kvm, slots->memslots[i].base_gfn,
-                                        slots->memslots[i].npages);
+                r = kvm_iommu_map_pages(kvm, &slots->memslots[i]);
                 if (r)
                         break;
         }
diff --git a/virt/kvm/kvm_main.c b/virt/kvm/kvm_main.c
index 4e2321c733f7..87d296d8b270 100644
--- a/virt/kvm/kvm_main.c
+++ b/virt/kvm/kvm_main.c
@@ -684,7 +684,7 @@ skip_lpage:
         spin_unlock(&kvm->mmu_lock);
 #ifdef CONFIG_DMAR
         /* map the pages in iommu page table */
-        r = kvm_iommu_map_pages(kvm, base_gfn, npages);
+        r = kvm_iommu_map_pages(kvm, memslot);
         if (r)
                 goto out;
 #endif