4 files changed, 156 insertions, 26 deletions
diff --git a/firmware/fw_base.S b/firmware/fw_base.S
index f596638..2c198cf 100644
--- a/firmware/fw_base.S
+++ b/firmware/fw_base.S
@@ -486,6 +486,16 @@ _trap_handler_all_mode:
 	REG_S	t0, SBI_TRAP_REGS_OFFSET(mepc)(sp)
 	csrr	t0, CSR_MSTATUS
 	REG_S	t0, SBI_TRAP_REGS_OFFSET(mstatus)(sp)
+	REG_S	zero, SBI_TRAP_REGS_OFFSET(mstatusH)(sp)
+#if __riscv_xlen == 32
+	csrr	t0, CSR_MISA
+	srli	t0, t0, ('H' - 'A')
+	andi	t0, t0, 0x1
+	beq	t0, zero, _skip_mstatush_save
+	csrr	t0, CSR_MSTATUSH
+	REG_S	t0, SBI_TRAP_REGS_OFFSET(mstatusH)(sp)
+_skip_mstatush_save:
+#endif
 
 	/* Save all general regisers except SP and T0 */
 	REG_S	zero, SBI_TRAP_REGS_OFFSET(zero)(sp)
@@ -560,6 +570,15 @@ _trap_handler_all_mode:
 	csrw	CSR_MEPC, t0
 	REG_L	t0, SBI_TRAP_REGS_OFFSET(mstatus)(sp)
 	csrw	CSR_MSTATUS, t0
+#if __riscv_xlen == 32
+	csrr	t0, CSR_MISA
+	srli	t0, t0, ('H' - 'A')
+	andi	t0, t0, 0x1
+	beq	t0, zero, _skip_mstatush_restore
+	REG_L	t0, SBI_TRAP_REGS_OFFSET(mstatusH)(sp)
+	csrw	CSR_MSTATUSH, t0
+_skip_mstatush_restore:
+#endif
 
 	/* Restore T0 */
 	REG_L	t0, SBI_TRAP_REGS_OFFSET(t0)(sp)
diff --git a/include/sbi/riscv_encoding.h b/include/sbi/riscv_encoding.h
index dace9e6..0e1ea59 100644
--- a/include/sbi/riscv_encoding.h
+++ b/include/sbi/riscv_encoding.h
@@ -58,8 +58,10 @@
 #define SSTATUS_UIE			0x00000001
 #define SSTATUS_SIE			0x00000002
 #define SSTATUS_UPIE			0x00000010
-#define SSTATUS_SPIE			0x00000020
-#define SSTATUS_SPP			0x00000100
+#define SSTATUS_SPIE_SHIFT		5
+#define SSTATUS_SPIE			(1UL << MSTATUS_SPIE_SHIFT)
+#define SSTATUS_SPP_SHIFT		8
+#define SSTATUS_SPP			(1UL << MSTATUS_SPP_SHIFT)
 #define SSTATUS_FS			0x00006000
 #define SSTATUS_XS			0x00018000
 #define SSTATUS_SUM			0x00040000
@@ -68,6 +70,14 @@
 #define SSTATUS_UXL			0x0000000300000000
 #define SSTATUS64_SD			0x8000000000000000
 
+#define HSTATUS_VTSR			0x00400000
+#define HSTATUS_VTVM			0x00100000
+#define HSTATUS_SP2V			0x00000200
+#define HSTATUS_SP2P			0x00000100
+#define HSTATUS_SPV			0x00000080
+#define HSTATUS_STL			0x00000040
+#define HSTATUS_SPRV			0x00000001
+
 #define DCSR_XDEBUGVER			(3U<<30)
 #define DCSR_NDRESET			(1<<29)
 #define DCSR_FULLRESET			(1<<28)
@@ -262,6 +272,23 @@
 #define CSR_STVAL			0x143
 #define CSR_SIP				0x144
 #define CSR_SATP			0x180
+
+#define CSR_HSTATUS			0x600
+#define CSR_HEDELEG			0x602
+#define CSR_HIDELEG			0x603
+#define CSR_HCOUNTERNEN		0x606
+#define CSR_HGATP			0x680
+
+#define CSR_VSSTATUS			0x200
+#define CSR_VSIE			0x204
+#define CSR_VSTVEC			0x205
+#define CSR_VSSCRATCH			0x240
+#define CSR_VSEPC			0x241
+#define CSR_VSCAUSE			0x242
+#define CSR_VSTVAL			0x243
+#define CSR_VSIP			0x244
+#define CSR_VSATP			0x280
+
 #define CSR_MSTATUS			0x300
 #define CSR_MISA			0x301
 #define CSR_MEDELEG			0x302
@@ -302,6 +329,7 @@
 #define CSR_DCSR			0x7b0
 #define CSR_DPC				0x7b1
 #define CSR_DSCRATCH			0x7b2
+
 #define CSR_MCYCLE			0xb00
 #define CSR_MINSTRET			0xb02
 #define CSR_MHPMCOUNTER3		0xb03
diff --git a/include/sbi/sbi_trap.h b/include/sbi/sbi_trap.h
index 19b95a7..bc26ad1 100644
--- a/include/sbi/sbi_trap.h
+++ b/include/sbi/sbi_trap.h
@@ -80,8 +80,10 @@
 #define SBI_TRAP_REGS_mepc			32
 /** Index of mstatus member in sbi_trap_regs */
 #define SBI_TRAP_REGS_mstatus			33
+/** Index of mstatusH member in sbi_trap_regs */
+#define SBI_TRAP_REGS_mstatusH			34
 /** Last member index in sbi_trap_regs */
-#define SBI_TRAP_REGS_last			34
+#define SBI_TRAP_REGS_last			35
 
 /* clang-format on */
 
@@ -164,6 +166,8 @@ struct sbi_trap_regs {
 	unsigned long mepc;
 	/** mstatus register state */
 	unsigned long mstatus;
+	/** mstatusH register state (only for 32-bit) */
+	unsigned long mstatusH;
 } __packed;
 
 struct sbi_scratch;
diff --git a/lib/sbi/sbi_trap.c b/lib/sbi/sbi_trap.c
index 82f7b65..93a0404 100644
--- a/lib/sbi/sbi_trap.c
+++ b/lib/sbi/sbi_trap.c
@@ -79,41 +79,119 @@ static void __noreturn sbi_trap_error(const char *msg, int rc, u32 hartid,
 int sbi_trap_redirect(struct sbi_trap_regs *regs, struct sbi_scratch *scratch,
 		      ulong epc, ulong cause, ulong tval)
 {
-	ulong new_mstatus, prev_mode;
+	ulong hstatus, vsstatus, prev_mode;
+#if __riscv_xlen == 32
+	bool prev_virt = (regs->mstatusH & MSTATUSH_MPV) ? TRUE : FALSE;
+	bool prev_stage2 = (regs->mstatusH & MSTATUSH_MTL) ? TRUE : FALSE;
+#else
+	bool prev_virt = (regs->mstatus & MSTATUS_MPV) ? TRUE : FALSE;
+	bool prev_stage2 = (regs->mstatus & MSTATUS_MTL) ? TRUE : FALSE;
+#endif
+	/* By default, we redirect to HS-mode */
+	bool next_virt = FALSE;
 
 	/* Sanity check on previous mode */
 	prev_mode = (regs->mstatus & MSTATUS_MPP) >> MSTATUS_MPP_SHIFT;
 	if (prev_mode != PRV_S && prev_mode != PRV_U)
 		return SBI_ENOTSUPP;
 
-	/* Update S-mode exception info */
-	csr_write(CSR_STVAL, tval);
-	csr_write(CSR_SEPC, epc);
-	csr_write(CSR_SCAUSE, cause);
+	/* For certain exceptions from VS/VU-mode we redirect to VS-mode */
+	if (misa_extension('H') && prev_virt && !prev_stage2) {
+		switch (cause) {
+		case CAUSE_FETCH_PAGE_FAULT:
+		case CAUSE_LOAD_PAGE_FAULT:
+		case CAUSE_STORE_PAGE_FAULT:
+			next_virt = TRUE;
+			break;
+		default:
+			break;
+		};
+	}
+
+	/* Update MSTATUS MPV and MTL bits */
+#if __riscv_xlen == 32
+	regs->mstatusH &= ~MSTATUSH_MPV;
+	regs->mstatusH |= (next_virt) ? MSTATUSH_MPV : 0UL;
+	regs->mstatusH &= ~MSTATUSH_MTL;
+#else
+	regs->mstatus &= ~MSTATUS_MPV;
+	regs->mstatus |= (next_virt) ? MSTATUS_MPV : 0UL;
+	regs->mstatus &= ~MSTATUS_MTL;
+#endif
+
+	/* Update HSTATUS for VS/VU-mode to HS-mode transition */
+	if (misa_extension('H') && prev_virt && !next_virt) {
+		/* Update HSTATUS SP2P, SP2V, SPV, and STL bits */
+		hstatus = csr_read(CSR_HSTATUS);
+		hstatus &= ~HSTATUS_SP2P;
+		hstatus |= (regs->mstatus & MSTATUS_SPP) ? HSTATUS_SP2P : 0;
+		hstatus &= ~HSTATUS_SP2V;
+		hstatus |= (hstatus & HSTATUS_SPV) ? HSTATUS_SP2V : 0;
+		hstatus &= ~HSTATUS_SPV;
+		hstatus |= (prev_virt) ? HSTATUS_SPV : 0;
+		hstatus &= ~HSTATUS_STL;
+		hstatus |= (prev_stage2) ? HSTATUS_STL : 0;
+		csr_write(CSR_HSTATUS, hstatus);
+	}
 
-	/* Set MEPC to S-mode exception vector base */
-	regs->mepc = csr_read(CSR_STVEC);
+	/* Update exception related CSRs */
+	if (next_virt) {
+		/* Update VS-mode exception info */
+		csr_write(CSR_VSTVAL, tval);
+		csr_write(CSR_VSEPC, epc);
+		csr_write(CSR_VSCAUSE, cause);
 
-	/* Initial value of new MSTATUS */
-	new_mstatus = regs->mstatus;
+		/* Set MEPC to VS-mode exception vector base */
+		regs->mepc = csr_read(CSR_VSTVEC);
 
-	/* Clear MPP, SPP, SPIE, and SIE */
-	new_mstatus &=
-		~(MSTATUS_MPP | MSTATUS_SPP | MSTATUS_SPIE | MSTATUS_SIE);
+		/* Set MPP to VS-mode */
+		regs->mstatus &= ~MSTATUS_MPP;
+		regs->mstatus |= (PRV_S << MSTATUS_MPP_SHIFT);
 
-	/* Set SPP */
-	if (prev_mode == PRV_S)
-		new_mstatus |= (1UL << MSTATUS_SPP_SHIFT);
+		/* Get VS-mode SSTATUS CSR */
+		vsstatus = csr_read(CSR_VSSTATUS);
 
-	/* Set SPIE */
-	if (regs->mstatus & MSTATUS_SIE)
-		new_mstatus |= (1UL << MSTATUS_SPIE_SHIFT);
+		/* Set SPP for VS-mode */
+		vsstatus &= ~SSTATUS_SPP;
+		if (prev_mode == PRV_S)
+			vsstatus |= (1UL << SSTATUS_SPP_SHIFT);
 
-	/* Set MPP */
-	new_mstatus |= (PRV_S << MSTATUS_MPP_SHIFT);
+		/* Set SPIE for VS-mode */
+		vsstatus &= ~SSTATUS_SPIE;
+		if (vsstatus & SSTATUS_SIE)
+			vsstatus |= (1UL << SSTATUS_SPIE_SHIFT);
 
-	/* Set new value in MSTATUS */
-	regs->mstatus = new_mstatus;
+		/* Clear SIE for VS-mode */
+		vsstatus &= ~SSTATUS_SIE;
+
+		/* Update VS-mode SSTATUS CSR */
+		csr_write(CSR_VSSTATUS, vsstatus);
+	} else {
+		/* Update S-mode exception info */
+		csr_write(CSR_STVAL, tval);
+		csr_write(CSR_SEPC, epc);
+		csr_write(CSR_SCAUSE, cause);
+
+		/* Set MEPC to S-mode exception vector base */
+		regs->mepc = csr_read(CSR_STVEC);
+
+		/* Set MPP to S-mode */
+		regs->mstatus &= ~MSTATUS_MPP;
+		regs->mstatus |= (PRV_S << MSTATUS_MPP_SHIFT);
+
+		/* Set SPP for S-mode*/
+		regs->mstatus &= ~MSTATUS_SPP;
+		if (prev_mode == PRV_S)
+			regs->mstatus |= (1UL << MSTATUS_SPP_SHIFT);
+
+		/* Set SPIE for S-mode */
+		regs->mstatus &= ~MSTATUS_SPIE;
+		if (regs->mstatus & MSTATUS_SIE)
+			regs->mstatus |= (1UL << MSTATUS_SPIE_SHIFT);
+
+		/* Clear SIE for S-mode */
+		regs->mstatus &= ~MSTATUS_SIE;
+	}
 
 	return 0;
 }
@@ -195,7 +273,8 @@ void sbi_trap_handler(struct sbi_trap_regs *regs, struct sbi_scratch *scratch)
 		break;
 	default:
 		/* If the trap came from S or U mode, redirect it there */
-		rc = sbi_trap_redirect(regs, scratch, regs->mepc, mcause, mtval);
+		rc = sbi_trap_redirect(regs, scratch, regs->mepc,
+				       mcause, mtval);
 		break;
 	};