[PATCH RFC v4 02/12] riscv/mm: add untagged_addr_remote_unlocked()

[Rik van Riel <riel@surriel.com>]

__access_remote_vm() reads another task's memory under the mmap lock. That
lock is contended on large multi-threaded workloads, where an mmap() or
munmap() stalls readers like /proc/PID/cmdline or process_vm_readv() even
though the target memory is resident and reachable under the per-VMA lock.

Reaching it under the per-VMA lock means looking up the VMA first, which
requires untagging the remote address without the mmap lock. riscv's
untagged_addr_remote() asserts that lock and has no unlocked variant; the
generic untagged_addr_remote_unlocked() falls back to untagged_addr(),
which masks against current->mm, not the target mm.

Add untagged_addr_remote_unlocked(), which masks @addr against the target
mm without asserting the lock, and redefine untagged_addr_remote() on top
of it so it keeps the assertion.

The mask width is mm->context.pmlen, written while a program is single
threaded, and frozen with MM_CONTEXT_LOCK_PMLEN once it spawns a second
thread.

That read need not be stable: pmlen is already read locklessly by
untagged_addr() and mm_untag_mask(), and a remote untag is best-effort.
Annotate the lockless reads with READ_ONCE() and the writes with
WRITE_ONCE(), so a reader sees one whole value and the accesses are
KCSAN-clean.

Assisted-by: Claude:claude-opus-4.8
Signed-off-by: Rik van Riel <riel@surriel.com>
---
 arch/riscv/include/asm/mmu_context.h |  4 ++--
 arch/riscv/include/asm/uaccess.h     | 10 +++++++---
 arch/riscv/kernel/process.c          | 12 +++++++-----
 3 files changed, 16 insertions(+), 10 deletions(-)

diff --git a/arch/riscv/include/asm/mmu_context.h b/arch/riscv/include/asm/mmu_context.h
index dbf27a78df6c..3ce16796e5a2 100644
--- a/arch/riscv/include/asm/mmu_context.h
+++ b/arch/riscv/include/asm/mmu_context.h
@@ -21,7 +21,7 @@ static inline void activate_mm(struct mm_struct *prev,
 			       struct mm_struct *next)
 {
 #ifdef CONFIG_RISCV_ISA_SUPM
-	next->context.pmlen = 0;
+	WRITE_ONCE(next->context.pmlen, 0);
 #endif
 	switch_mm(prev, next, NULL);
 }
@@ -44,7 +44,7 @@ DECLARE_STATIC_KEY_FALSE(use_asid_allocator);
 #define mm_untag_mask mm_untag_mask
 static inline unsigned long mm_untag_mask(struct mm_struct *mm)
 {
-	return -1UL >> mm->context.pmlen;
+	return -1UL >> READ_ONCE(mm->context.pmlen);
 }
 #endif
 
diff --git a/arch/riscv/include/asm/uaccess.h b/arch/riscv/include/asm/uaccess.h
index 5d4ec15584cf..53806e0f7dcf 100644
--- a/arch/riscv/include/asm/uaccess.h
+++ b/arch/riscv/include/asm/uaccess.h
@@ -16,7 +16,7 @@
 static inline unsigned long __untagged_addr_remote(struct mm_struct *mm, unsigned long addr)
 {
 	if (riscv_has_extension_unlikely(RISCV_ISA_EXT_SUPM)) {
-		u8 pmlen = mm->context.pmlen;
+		u8 pmlen = READ_ONCE(mm->context.pmlen);
 
 		/* Virtual addresses are sign-extended; physical addresses are zero-extended. */
 		if (IS_ENABLED(CONFIG_MMU))
@@ -33,12 +33,16 @@ static inline unsigned long __untagged_addr_remote(struct mm_struct *mm, unsigne
 	(__force __typeof__(addr))__untagged_addr_remote(current->mm, __addr);	\
 })
 
-#define untagged_addr_remote(mm, addr) ({					\
+#define untagged_addr_remote_unlocked(mm, addr) ({				\
 	unsigned long __addr = (__force unsigned long)(addr);			\
-	mmap_assert_locked(mm);							\
 	(__force __typeof__(addr))__untagged_addr_remote(mm, __addr);		\
 })
 
+#define untagged_addr_remote(mm, addr) ({					\
+	mmap_assert_locked(mm);							\
+	untagged_addr_remote_unlocked(mm, addr);				\
+})
+
 #define access_ok(addr, size) likely(__access_ok(untagged_addr(addr), size))
 #else
 #define untagged_addr(addr) (addr)
diff --git a/arch/riscv/kernel/process.c b/arch/riscv/kernel/process.c
index b2df7f72241a..6ae7552fed09 100644
--- a/arch/riscv/kernel/process.c
+++ b/arch/riscv/kernel/process.c
@@ -357,13 +357,15 @@ long set_tagged_addr_ctrl(struct task_struct *task, unsigned long arg)
 	if (mmap_write_lock_killable(mm))
 		return -EINTR;
 
-	if (test_bit(MM_CONTEXT_LOCK_PMLEN, &mm->context.flags) && mm->context.pmlen != pmlen) {
-		mmap_write_unlock(mm);
-		return -EBUSY;
+	if (test_bit(MM_CONTEXT_LOCK_PMLEN, &mm->context.flags)) {
+		if (READ_ONCE(mm->context.pmlen) != pmlen) {
+			mmap_write_unlock(mm);
+			return -EBUSY;
+		}
 	}
 
 	envcfg_update_bits(task, ENVCFG_PMM, pmm);
-	mm->context.pmlen = pmlen;
+	WRITE_ONCE(mm->context.pmlen, pmlen);
 
 	mmap_write_unlock(mm);
 
@@ -394,7 +396,7 @@ long get_tagged_addr_ctrl(struct task_struct *task)
 		break;
 	}
 
-	if (task->mm->context.pmlen)
+	if (READ_ONCE(task->mm->context.pmlen))
 		ret |= PR_TAGGED_ADDR_ENABLE;
 
 	return ret;
-- 
2.53.0-Meta