2 files changed, 7 insertions, 4 deletions

diff --git a/drivers/xen/gntdev.c b/drivers/xen/gntdev.c
index 1ffd03bf8e10..7f1241608489 100644
--- a/drivers/xen/gntdev.c
+++ b/drivers/xen/gntdev.c
@@ -314,8 +314,9 @@ static int __unmap_grant_pages(struct grant_map *map, int offset, int pages)
                 }
         }
 
-        err = gnttab_unmap_refs(map->unmap_ops + offset, map->pages + offset,
-                                pages, true);
+        err = gnttab_unmap_refs(map->unmap_ops + offset,
+                        use_ptemod ? map->kmap_ops + offset : NULL, map->pages + offset,
+                        pages);
         if (err)
                 return err;
 
diff --git a/drivers/xen/grant-table.c b/drivers/xen/grant-table.c
index 0bfc1ef11259..006726688baf 100644
--- a/drivers/xen/grant-table.c
+++ b/drivers/xen/grant-table.c
@@ -870,7 +870,8 @@ int gnttab_map_refs(struct gnttab_map_grant_ref *map_ops,
 EXPORT_SYMBOL_GPL(gnttab_map_refs);
 
 int gnttab_unmap_refs(struct gnttab_unmap_grant_ref *unmap_ops,
-                      struct page **pages, unsigned int count, bool clear_pte)
+                      struct gnttab_map_grant_ref *kmap_ops,
+                      struct page **pages, unsigned int count)
 {
         int i, ret;
         bool lazy = false;
@@ -888,7 +889,8 @@ int gnttab_unmap_refs(struct gnttab_unmap_grant_ref *unmap_ops,
         }
 
         for (i = 0; i < count; i++) {
-                ret = m2p_remove_override(pages[i], clear_pte);
+                ret = m2p_remove_override(pages[i], kmap_ops ?
+                                       &kmap_ops[i] : NULL);
                 if (ret)
                         return ret;
         }