[PATCH v2] arm64: tlb: Fix TLBI RANGE operand

[Gavin Shan <gshan@redhat.com>]

KVM/arm64 relies on TLBI RANGE feature to flush TLBs when the dirty
bitmap is collected by VMM and the corresponding PTEs need to be
write-protected during live migration. Unfortunately, the operand
passed to the TLBI RANGE instruction isn't correctly sorted out by
commit d1d3aa98b1d4 ("arm64: tlb: Use the TLBI RANGE feature in arm64").
It leads to crash on the destination VM after live migration because
TLBs aren't flushed completely and some of the dirty pages are missed.

For example, I have a VM where 8GB memory is assigned, starting from
0x40000000 (1GB). Note that the host has 4KB as the base page size.
All TLBs for VM can be covered by one TLBI RANGE operation. However,
the operand 0xffff708000040000 is set for scale -9, and  -1 is returned
from __TLBI_RANGE_NUM() for scale 3/2/1/0 and rejected by the loop in
__flush_tlb_range_op(). __TLBI_RANGE_NUM() isn't expected to work
like this because all the pages should be covered by scale 3/2/1/0,
plus an additional page if needed.

Fix the macro __TLBI_RANGE_NUM() so that the correct NUM and TLBI RANGE
operand are provided for each scale level. With the changes, [-1 31]
instead of [-1 30] can be returned from the macro, meaning the TLBs for
0x200000 pages (8GB memory) can be flushed in one shoot at scale 3. The
macro TLBI_RANGE_MASK is dropped since no one uses it any more.

Fixes: d1d3aa98b1d4 ("arm64: tlb: Use the TLBI RANGE feature in arm64")
Cc: stable@kernel.org # v5.10+
Reported-by: Yihuang Yu <yihyu@redhat.com>
Suggested-by: Marc Zyngier <maz@kernel.org>
Signed-off-by: Gavin Shan <gshan@redhat.com>
---
v2: Improve __TLBI_RANGE_NUM() as Marc suggested
---
 arch/arm64/include/asm/tlbflush.h | 13 +++++++++----
 1 file changed, 9 insertions(+), 4 deletions(-)

diff --git a/arch/arm64/include/asm/tlbflush.h b/arch/arm64/include/asm/tlbflush.h
index 3b0e8248e1a4..cd9b71c30366 100644
--- a/arch/arm64/include/asm/tlbflush.h
+++ b/arch/arm64/include/asm/tlbflush.h
@@ -161,12 +161,17 @@ static inline unsigned long get_trans_granule(void)
 #define MAX_TLBI_RANGE_PAGES		__TLBI_RANGE_PAGES(31, 3)
 
 /*
- * Generate 'num' values from -1 to 30 with -1 rejected by the
+ * Generate 'num' values from -1 to 31 with -1 rejected by the
  * __flush_tlb_range() loop below.
  */
-#define TLBI_RANGE_MASK			GENMASK_ULL(4, 0)
-#define __TLBI_RANGE_NUM(pages, scale)	\
-	((((pages) >> (5 * (scale) + 1)) & TLBI_RANGE_MASK) - 1)
+#define __TLBI_RANGE_NUM(pages, scale)					\
+	({								\
+		int __pages = min((pages),				\
+				  __TLBI_RANGE_PAGES(31, (scale)));	\
+		int __numplus1 = __pages >> (5 * (scale) + 1);		\
+									\
+		(__numplus1 - 1);					\
+	})
 
 /*
  *	TLB Invalidation
-- 
2.44.0