--- /dev/null
+From 7c7b962938ddda6a9cd095de557ee5250706ea88 Mon Sep 17 00:00:00 2001
+From: Alistair Popple <apopple@nvidia.com>
+Date: Thu, 30 Mar 2023 12:25:19 +1100
+Subject: mm: take a page reference when removing device exclusive entries
+
+From: Alistair Popple <apopple@nvidia.com>
+
+commit 7c7b962938ddda6a9cd095de557ee5250706ea88 upstream.
+
+Device exclusive page table entries are used to prevent CPU access to a
+page whilst it is being accessed from a device. Typically this is used to
+implement atomic operations when the underlying bus does not support
+atomic access. When a CPU thread encounters a device exclusive entry it
+locks the page and restores the original entry after calling mmu notifiers
+to signal drivers that exclusive access is no longer available.
+
+The device exclusive entry holds a reference to the page making it safe to
+access the struct page whilst the entry is present. However the fault
+handling code does not hold the PTL when taking the page lock. This means
+if there are multiple threads faulting concurrently on the device
+exclusive entry one will remove the entry whilst others will wait on the
+page lock without holding a reference.
+
+This can lead to threads locking or waiting on a folio with a zero
+refcount. Whilst mmap_lock prevents the pages getting freed via munmap()
+they may still be freed by a migration. This leads to warnings such as
+PAGE_FLAGS_CHECK_AT_FREE due to the page being locked when the refcount
+drops to zero.
+
+Fix this by trying to take a reference on the folio before locking it.
+The code already checks the PTE under the PTL and aborts if the entry is
+no longer there. It is also possible the folio has been unmapped, freed
+and re-allocated allowing a reference to be taken on an unrelated folio.
+This case is also detected by the PTE check and the folio is unlocked
+without further changes.
+
+Link: https://lkml.kernel.org/r/20230330012519.804116-1-apopple@nvidia.com
+Fixes: b756a3b5e7ea ("mm: device exclusive memory access")
+Signed-off-by: Alistair Popple <apopple@nvidia.com>
+Reviewed-by: Ralph Campbell <rcampbell@nvidia.com>
+Reviewed-by: John Hubbard <jhubbard@nvidia.com>
+Acked-by: David Hildenbrand <david@redhat.com>
+Cc: Matthew Wilcox (Oracle) <willy@infradead.org>
+Cc: Christoph Hellwig <hch@infradead.org>
+Cc: <stable@vger.kernel.org>
+Signed-off-by: Andrew Morton <akpm@linux-foundation.org>
+Signed-off-by: Greg Kroah-Hartman <gregkh@linuxfoundation.org>
+---
+ mm/memory.c | 16 +++++++++++++++-
+ 1 file changed, 15 insertions(+), 1 deletion(-)
+
+--- a/mm/memory.c
++++ b/mm/memory.c
+@@ -3619,8 +3619,21 @@ static vm_fault_t remove_device_exclusiv
+ struct vm_area_struct *vma = vmf->vma;
+ struct mmu_notifier_range range;
+
+- if (!folio_lock_or_retry(folio, vma->vm_mm, vmf->flags))
++ /*
++ * We need a reference to lock the folio because we don't hold
++ * the PTL so a racing thread can remove the device-exclusive
++ * entry and unmap it. If the folio is free the entry must
++ * have been removed already. If it happens to have already
++ * been re-allocated after being freed all we do is lock and
++ * unlock it.
++ */
++ if (!folio_try_get(folio))
++ return 0;
++
++ if (!folio_lock_or_retry(folio, vma->vm_mm, vmf->flags)) {
++ folio_put(folio);
+ return VM_FAULT_RETRY;
++ }
+ mmu_notifier_range_init_owner(&range, MMU_NOTIFY_EXCLUSIVE, 0, vma,
+ vma->vm_mm, vmf->address & PAGE_MASK,
+ (vmf->address & PAGE_MASK) + PAGE_SIZE, NULL);
+@@ -3633,6 +3646,7 @@ static vm_fault_t remove_device_exclusiv
+
+ pte_unmap_unlock(vmf->pte, vmf->ptl);
+ folio_unlock(folio);
++ folio_put(folio);
+
+ mmu_notifier_invalidate_range_end(&range);
+ return 0;