int track_pfn_vma_new(struct vm_area_struct *vma, pgprot_t *prot,
                        unsigned long pfn, unsigned long size)
 {
+       resource_size_t paddr = (resource_size_t)pfn << PAGE_SHIFT;
        unsigned long flags;
-       resource_size_t paddr;
-       unsigned long vma_size = vma->vm_end - vma->vm_start;
 
-       if (is_linear_pfn_mapping(vma)) {
-               /* reserve the whole chunk starting from vm_pgoff */
-               paddr = (resource_size_t)vma->vm_pgoff << PAGE_SHIFT;
-               return reserve_pfn_range(paddr, vma_size, prot, 0);
-       }
+       /* reserve the whole chunk starting from paddr */
+       if (is_linear_pfn_mapping(vma))
+               return reserve_pfn_range(paddr, size, prot, 0);
 
        if (!pat_enabled)
                return 0;
 
        /* for vm_insert_pfn and friends, we set prot based on lookup */
-       flags = lookup_memtype(pfn << PAGE_SHIFT);
+       flags = lookup_memtype(paddr);
        *prot = __pgprot((pgprot_val(vma->vm_page_prot) & (~_PAGE_CACHE_MASK)) |
                         flags);
 
 /*
  * untrack_pfn_vma is called while unmapping a pfnmap for a region.
  * untrack can be called for a specific region indicated by pfn and size or
- * can be for the entire vma (in which case size can be zero).
+ * can be for the entire vma (in which case pfn, size are zero).
  */
 void untrack_pfn_vma(struct vm_area_struct *vma, unsigned long pfn,
                        unsigned long size)
 {
        resource_size_t paddr;
-       unsigned long vma_size = vma->vm_end - vma->vm_start;
+       unsigned long prot;
 
-       if (is_linear_pfn_mapping(vma)) {
-               /* free the whole chunk starting from vm_pgoff */
-               paddr = (resource_size_t)vma->vm_pgoff << PAGE_SHIFT;
-               free_pfn_range(paddr, vma_size);
+       if (!is_linear_pfn_mapping(vma))
                return;
+
+       /* free the chunk starting from pfn or the whole chunk */
+       paddr = (resource_size_t)pfn << PAGE_SHIFT;
+       if (!paddr && !size) {
+               if (follow_phys(vma, vma->vm_start, 0, &prot, &paddr)) {
+                       WARN_ON_ONCE(1);
+                       return;
+               }
+
+               size = vma->vm_end - vma->vm_start;
        }
+       free_pfn_range(paddr, size);
 }
 
 pgprot_t pgprot_writecombine(pgprot_t prot)