[PATCH] riscv: vector: preserve state when scheduling at nonzero depth

[Karl Mehltretter <kmehltretter@gmail.com>]

The IN_SCHEDULE shortcut lets __switch_to_vector() discard Vector state at
a voluntary schedule point, where Vector registers are caller-saved.
Switch-in can then enable Vector without restoring state.

An interrupt or fault can also schedule at nonzero Vector nesting depth.
This triggers:

  WARNING: arch/riscv/include/asm/vector.h:376 at __schedule+0xfbc/0x10b4

The shortcut is then also taken on switch-in, so it skips NEED_RESTORE and
riscv_v_context_nesting_end() resumes with stale Vector registers. In the
vector usercopy loop, an interrupt between vsetvli and vle8.v/vse8.v can
therefore resume with another task's vl, vtype and vector registers. The
scalar loop state survives, so the copy can use the wrong vector length and
silently corrupt user data. A sleeping page fault in vectorized usercopy
can reach the same switch without CONFIG_PREEMPTION.

Use the shortcut only at depth zero. Nonzero-depth switches retain the
existing save and NEED_RESTORE protocol.

Fixes: d1049fc0de81 ("riscv: vector: Support calling schedule() for preemptible Vector")
Cc: stable@vger.kernel.org
Assisted-by: Codex:gpt-5.6-luna
Signed-off-by: Karl Mehltretter <kmehltretter@gmail.com>
---
Reproducer:

Originally caught by syzkaller fuzzing:

  [   75.067010] ------------[ cut here ]------------
  [   75.069361] WARNING: arch/riscv/include/asm/vector.h:376 at __schedule+0xfbc/0x10b4, CPU#1: syz.4.788/3533
  [   75.072830] CPU: 1 UID: 0 PID: 3533 Comm: syz.4.788 Not tainted 7.2.0-rc5-g11f985de3fef #3 PREEMPTLAZY
  [   75.075177] [<ffffffff810cb124>] __schedule+0xfbc/0x10b4
  [   75.075572] [<ffffffff810cb41c>] preempt_schedule_irq+0x2a/0x76
  [   75.075753] [<ffffffff810c6a06>] irqentry_exit+0x260/0xd48
  [   75.075911] [<ffffffff810c666a>] do_irq+0x34/0x48
  [   75.076076] [<ffffffff810d4f82>] handle_exception+0x146/0x174
  [   75.076380] [<ffffffff810c5654>] loop+0x4/0x26
  [   75.076505] [<ffffffff80749838>] copy_folio_from_iter_atomic+0x2d4/0xd64
  [   75.078599] ---[ end trace 0000000000000000 ]---

The following standalone workload exercises the same vectorized usercopy
path. Build vector-stress.c into the initramfs and mount debugfs before
running it. The kernel used CONFIG_PREEMPT_LAZY=y,
CONFIG_RISCV_ISA_V_PREEMPTIVE=y and CONFIG_KCOV=y. Run it under QEMU TCG
with:

  qemu-system-riscv64 -machine virt -cpu max -smp 1 -nographic \
    -kernel arch/riscv/boot/Image -initrd vector-diag-small-memcheck.cpio.gz \
    -append 'console=ttyS0 earlycon=sbi rdinit=/init'

The stress program checks every byte read back from the pipe. With this
patch, the workload completed with failures=0 and no Vector warning.

Testing:
  checkpatch.pl --strict, git diff --check, Image builds with
  CONFIG_RISCV_ISA_V_PREEMPTIVE=y and CONFIG_RISCV_ISA_V_PREEMPTIVE=n, and
  QEMU TCG one-vCPU stress runs. The patched workload completed with
  failures=0 and no warning.

No conflict with Andy Chiu's pending "riscv: optimize Vector context restore
on syscall" series; both apply independently.

vector-stress.c:

#define _GNU_SOURCE

#include <errno.h>
#include <fcntl.h>
#include <linux/kcov.h>
#include <pthread.h>
#include <sched.h>
#include <stdatomic.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/ioctl.h>
#include <sys/mman.h>
#include <sys/types.h>
#include <unistd.h>

#define WORKERS 4
#define ITERATIONS 2000
#define CHUNK (16 * 1024)
#define KCOV_ENTRIES (1 << 16)
#define KCOV_BYTES (KCOV_ENTRIES * sizeof(unsigned long))

static pthread_barrier_t start_barrier;
static atomic_int failures;

static int kcov_start(unsigned long **area, int *fd, int id)
{
	unsigned long *map;
	int kfd, saved_errno;

	kfd = open("/sys/kernel/debug/kcov", O_RDWR);
	if (kfd < 0) {
		dprintf(STDERR_FILENO, "worker %d: KCOV open: %s\n", id,
			strerror(errno));
		return -1;
	}
	if (ioctl(kfd, KCOV_INIT_TRACE, KCOV_ENTRIES) < 0) {
		dprintf(STDERR_FILENO, "worker %d: KCOV init: %s\n", id,
			strerror(errno));
		goto fail_close;
	}
	map = mmap(NULL, KCOV_BYTES,
			PROT_READ | PROT_WRITE, MAP_SHARED, kfd, 0);
	if (map == MAP_FAILED) {
		dprintf(STDERR_FILENO, "worker %d: KCOV mmap: %s\n", id,
			strerror(errno));
		goto fail_close;
	}
	if (ioctl(kfd, KCOV_ENABLE, KCOV_TRACE_PC) < 0) {
		dprintf(STDERR_FILENO, "worker %d: KCOV enable: %s\n", id,
			strerror(errno));
		saved_errno = errno;
		munmap(map, KCOV_BYTES);
		errno = saved_errno;
		goto fail_close;
	}
	*area = map;
	*fd = kfd;
	return 0;

fail_close:
	saved_errno = errno;
	close(kfd);
	errno = saved_errno;
	return -1;
}

static void *worker(void *arg)
{
	unsigned long *area;
	unsigned char *buffer;
	int pipefd[2], kfd, id = (int)(uintptr_t)arg;

	if (kcov_start(&area, &kfd, id) < 0) {
		atomic_fetch_add(&failures, 1);
		return NULL;
	}
	if (pipe2(pipefd, O_CLOEXEC) < 0) {
		dprintf(STDERR_FILENO, "worker %d: pipe failed: %s\n", id,
			strerror(errno));
		atomic_fetch_add(&failures, 1);
		goto out_kcov;
	}
	buffer = aligned_alloc(64, CHUNK);
	if (!buffer) {
		dprintf(STDERR_FILENO, "worker %d: allocation failed: %s\n", id,
			strerror(errno));
		atomic_fetch_add(&failures, 1);
		goto out_pipe;
	}
	memset(buffer, 0x30 + id, CHUNK);
	pthread_barrier_wait(&start_barrier);

	for (int i = 0; i < ITERATIONS; i++) {
		size_t done = 0;

		while (done < CHUNK) {
			ssize_t n = write(pipefd[1], buffer + done, CHUNK - done);
			if (n < 0 && errno == EINTR)
				continue;
			if (n <= 0) {
				atomic_fetch_add(&failures, 1);
				goto out_buffer;
			}
			done += n;
		}
		done = 0;
		while (done < CHUNK) {
			ssize_t n = read(pipefd[0], buffer + done, CHUNK - done);
			if (n < 0 && errno == EINTR)
				continue;
			if (n <= 0) {
				atomic_fetch_add(&failures, 1);
				goto out_buffer;
			}
			done += n;
		}
		for (size_t j = 0; j < CHUNK; j++) {
			if (buffer[j] != (unsigned char)(0x30 + id)) {
				dprintf(STDERR_FILENO,
					"worker %d: data mismatch at %zu\n", id, j);
				atomic_fetch_add(&failures, 1);
				goto out_buffer;
			}
		}
		if ((i & 7) == 0)
			sched_yield();
	}

out_buffer:
	free(buffer);
out_pipe:
	close(pipefd[0]);
	close(pipefd[1]);
out_kcov:
	ioctl(kfd, KCOV_DISABLE, 0);
	munmap(area, KCOV_BYTES);
	close(kfd);
	return NULL;
}

int main(void)
{
	pthread_t threads[WORKERS];

	pthread_barrier_init(&start_barrier, NULL, WORKERS);
	for (int i = 0; i < WORKERS; i++)
		if (pthread_create(&threads[i], NULL, worker, (void *)(uintptr_t)i))
			atomic_fetch_add(&failures, 1);
	for (int i = 0; i < WORKERS; i++)
		pthread_join(threads[i], NULL);
	pthread_barrier_destroy(&start_barrier);

	printf("vector-stress complete failures=%d\n", atomic_load(&failures));
	return atomic_load(&failures) ? 1 : 0;
}

 arch/riscv/include/asm/vector.h | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/arch/riscv/include/asm/vector.h b/arch/riscv/include/asm/vector.h
index 00cb9c0982b1a..1766bb7494d3b 100644
--- a/arch/riscv/include/asm/vector.h
+++ b/arch/riscv/include/asm/vector.h
@@ -372,8 +372,8 @@ static inline void __switch_to_vector(struct task_struct *prev,
 	struct pt_regs *regs;
 
 	if (riscv_preempt_v_started(prev)) {
-		if (riscv_v_is_on()) {
-			WARN_ON(prev->thread.riscv_v_flags & RISCV_V_CTX_DEPTH_MASK);
+		if (riscv_v_is_on() &&
+		    !(prev->thread.riscv_v_flags & RISCV_V_CTX_DEPTH_MASK)) {
 			riscv_v_disable();
 			prev->thread.riscv_v_flags |= RISCV_PREEMPT_V_IN_SCHEDULE;
 		}
-- 
2.53.0