[PATCH v6 06/11] x86/virt/tdx: Optimize tdx_pamt_get/put()

[Rick Edgecombe <rick.p.edgecombe@intel.com>]

From: "Kirill A. Shutemov" <kirill.shutemov@linux.intel.com>

The Dynamic PAMT get/put helpers use a global spinlock to serialize all
refcount updates and SEAMCALL invocations. This gives correct behavior for
concurrent callers, but leads to contention. It is especially bad from the
KVM side, which is designed to allow faulting in EPT under a shared lock.
With the global spinlock, not only is the lock an exclusive one, but it is
for all TDs instead of just a single one.

But taking the global lock each time is actually unnecessary. Only the 0->1
and 1->0 refcount transitions actually need the lock (to pair with
SEAMCALLs that actually add and remove with the Dynamic PAMT pages). The
common case of incrementing or decrementing a non-zero refcount can be
done locklessly.

So create a fast and slow path. Check the refcount outside the lock and
only take it for the slowpath (0->1 and 1->0 transitions).

On the put side make the refcount adjustment and lock taking atomic so if
a 'get' happens between them, it doesn't cause the Dynamic PAMT to be
freed incorrectly. On the get side there is no technique for doing the
refcount adjustment and lock atomically, so check the refcount again
inside the lock.

Assisted-by: GitHub Copilot:claude-opus-4-6
Signed-off-by: Kirill A. Shutemov <kirill.shutemov@linux.intel.com>
Co-developed-by: Rick Edgecombe <rick.p.edgecombe@intel.com>
Signed-off-by: Rick Edgecombe <rick.p.edgecombe@intel.com>
---
v6:
 - Fix "tdx_pamt_add()" typo to "tdx_pamt_get()" in lost-race comment
 - Fix error path bug: set ret = -EIO and use WARN_ON_ONCE() instead of
   pr_err() for unexpected PAMT.ADD failures (Sean)
 - Use "set the refcount 0->1" wording to match atomic_set() usage
 - Wrap comments to 80 columns
 - Switch to atomic_dec_and_lock() and remove handling of races that are
   no longer needed as a result. Adjust comments as appropriate. (Dave)
 - Adjustments from dropping error helper patches
v4:
 - Use atomic_set() in the HPA_RANGE_NOT_FREE case (Kiryl)
 - Log, comment typos (Binbin)
 - Move PAMT page allocation after refcount check in tdx_pamt_get() to
   avoid an alloc/free in the common path.

v3:
 - Split out optimization from “x86/virt/tdx: Add tdx_alloc/free_page() helpers”
 - Remove edge case handling that I could not find a reason for
 - Write log
---
 arch/x86/virt/vmx/tdx/tdx.c | 102 +++++++++++++++++++++---------------
 1 file changed, 61 insertions(+), 41 deletions(-)

diff --git a/arch/x86/virt/vmx/tdx/tdx.c b/arch/x86/virt/vmx/tdx/tdx.c
index 50333eb96efa6..c41c632a4cdf2 100644
--- a/arch/x86/virt/vmx/tdx/tdx.c
+++ b/arch/x86/virt/vmx/tdx/tdx.c
@@ -2057,32 +2057,50 @@ static int tdx_pamt_get(kvm_pfn_t pfn)
 	if (!tdx_supports_dynamic_pamt(&tdx_sysinfo))
 		return 0;
 
+	pamt_refcount = tdx_find_pamt_refcount(pfn);
+
+	/*
+	 * If the pamt page is already added (i.e. refcount >= 1),
+	 * then just increment the refcount.
+	 */
+	if (atomic_inc_not_zero(pamt_refcount))
+		return 0;
+
 	ret = alloc_pamt_array(pamt_pages);
 	if (ret)
 		return ret;
 
-	pamt_refcount = tdx_find_pamt_refcount(pfn);
+	spin_lock(&pamt_lock);
 
-	scoped_guard(spinlock, &pamt_lock) {
-		/*
-		 * If the pamt page is already added (i.e. refcount >= 1),
-		 * then just increment the refcount.
-		 */
-		if (atomic_read(pamt_refcount)) {
-			atomic_inc(pamt_refcount);
-			goto out_free;
-		}
-
-		/* Try to add the pamt page and take the refcount 0->1. */
-		tdx_status = tdh_phymem_pamt_add(pfn, pamt_pages);
-		if (WARN_ON_ONCE(tdx_status != TDX_SUCCESS)) {
-			ret = -EIO;
-			goto out_free;
-		}
-
-		atomic_set(pamt_refcount, 1);
+	/*
+	 * Unlike tdx_pamt_put() which uses atomic_dec_and_lock() to
+	 * atomically handle the 1->0 transition, the get side has no
+	 * equivalent combined primitive for 0->1. Recheck under the
+	 * lock since another get may have already done the 0->1
+	 * transition after both saw atomic_inc_not_zero() fail.
+	 */
+	if (atomic_read(pamt_refcount)) {
+		atomic_inc(pamt_refcount);
+		spin_unlock(&pamt_lock);
+		goto out_free;
 	}
 
+	tdx_status = tdh_phymem_pamt_add(pfn, pamt_pages);
+	if (tdx_status == TDX_SUCCESS) {
+		/*
+		 * The refcount is zero, and this locked path is the
+		 * only way to increase it from 0->1.
+		 */
+		atomic_set(pamt_refcount, 1);
+	} else {
+		WARN_ON_ONCE(1);
+		ret = -EIO;
+		spin_unlock(&pamt_lock);
+		goto out_free;
+	}
+
+	spin_unlock(&pamt_lock);
+
 	return 0;
 out_free:
 	free_pamt_array(pamt_pages);
@@ -2104,32 +2122,34 @@ static void tdx_pamt_put(kvm_pfn_t pfn)
 
 	pamt_refcount = tdx_find_pamt_refcount(pfn);
 
-	scoped_guard(spinlock, &pamt_lock) {
+	/*
+	 * If there is more than 1 reference on the pamt page, don't
+	 * remove it yet. Just decrement the refcount.
+	 */
+	if (!atomic_dec_and_lock(pamt_refcount, &pamt_lock))
+		return;
+
+	tdx_status = tdh_phymem_pamt_remove(pfn, pamt_pages);
+
+	/*
+	 * Don't free pamt_pages as it could hold garbage when
+	 * tdh_phymem_pamt_remove() fails.  Don't panic/BUG_ON(), as
+	 * there is no risk of data corruption, but do yell loudly as
+	 * failure indicates a kernel bug, memory is being leaked, and
+	 * the dangling PAMT entry may cause future operations to fail.
+	 */
+	if (WARN_ON_ONCE(tdx_status != TDX_SUCCESS)) {
 		/*
-		 * If the there are more than 1 references on the pamt page,
-		 * don't remove it yet. Just decrement the refcount.
+		 * atomic_dec_and_lock() already decremented it to 0,
+		 * but the PAMT entry still exists since REMOVE failed.
 		 */
-		if (atomic_read(pamt_refcount) > 1) {
-			atomic_dec(pamt_refcount);
-			return;
-		}
-
-		/* Try to remove the pamt page and take the refcount 1->0. */
-		tdx_status = tdh_phymem_pamt_remove(pfn, pamt_pages);
-
-		/*
-		 * Don't free pamt_pages as it could hold garbage when
-		 * tdh_phymem_pamt_remove() fails.  Don't panic/BUG_ON(), as
-		 * there is no risk of data corruption, but do yell loudly as
-		 * failure indicates a kernel bug, memory is being leaked, and
-		 * the dangling PAMT entry may cause future operations to fail.
-		 */
-		if (WARN_ON_ONCE(tdx_status != TDX_SUCCESS))
-			return;
-
-		atomic_set(pamt_refcount, 0);
+		atomic_set(pamt_refcount, 1);
+		spin_unlock(&pamt_lock);
+		return;
 	}
 
+	spin_unlock(&pamt_lock);
+
 	free_pamt_array(pamt_pages);
 }
 
-- 
2.54.0