[edk2-platforms][PATCH 27/30] RISC-V/PlatformPkg: Updates for the latest OpenSBI

["Abner Chang" <abner.chang@hpe.com>]

Code changes to incorporate with OpenSBI commit ID:
a731c7e36988c3308e1978ecde491f2f6182d490

Cc: Sunil V L <sunilvl@ventanamicro.com>
Cc: Daniel Schaefer <daniel.schaefer@hpe.com>

Signed-off-by: Abner Chang <abner.chang@hpe.com>
---
 .../OpensbiPlatformLib/OpensbiPlatformLib.inf |  10 +-
 .../PlatformPkg/Universal/Sec/SecMain.inf     |   4 +
 .../Library/OpensbiPlatformLib/Platform.c     |  57 ----
 .../Universal/Sec/Edk2OpenSbiPlatform.c       | 149 ---------
 .../PlatformPkg/Universal/Sec/SecMain.c       |  48 ++-
 .../Universal/Sec/Riscv64/SecEntry.S          | 300 ++++++++++--------
 6 files changed, 212 insertions(+), 356 deletions(-)

diff --git a/Platform/RISC-V/PlatformPkg/Library/OpensbiPlatformLib/OpensbiPlatformLib.inf b/Platform/RISC-V/PlatformPkg/Library/OpensbiPlatformLib/OpensbiPlatformLib.inf
index 909fbffa8d..2e1227733a 100644
--- a/Platform/RISC-V/PlatformPkg/Library/OpensbiPlatformLib/OpensbiPlatformLib.inf
+++ b/Platform/RISC-V/PlatformPkg/Library/OpensbiPlatformLib/OpensbiPlatformLib.inf
@@ -51,12 +51,4 @@
   gUefiRiscVPlatformPkgTokenSpaceGuid.PcdBootableHartNumber
   gUefiRiscVPlatformPkgTokenSpaceGuid.PcdOpenSbiStackSize
 
-  gSiFiveU5SeriesPlatformsPkgTokenSpaceGuid.PcdU5UartBase
-  gSiFiveU5SeriesPlatformsPkgTokenSpaceGuid.PcdU5PlatformSystemClock
-
-  gUefiRiscVPlatformPkgTokenSpaceGuid.PcdRootFirmwareDomainBaseAddress
-  gUefiRiscVPlatformPkgTokenSpaceGuid.PcdRootFirmwareDomainSize
-  gUefiRiscVPlatformPkgTokenSpaceGuid.PcdFirmwareDomainBaseAddress
-  gUefiRiscVPlatformPkgTokenSpaceGuid.PcdFirmwareDomainSize
-  gUefiRiscVPlatformPkgTokenSpaceGuid.PcdVariableFirmwareRegionBaseAddress
-  gUefiRiscVPlatformPkgTokenSpaceGuid.PcdVariableFirmwareRegionSize
+
diff --git a/Platform/RISC-V/PlatformPkg/Universal/Sec/SecMain.inf b/Platform/RISC-V/PlatformPkg/Universal/Sec/SecMain.inf
index 1cfbef961f..dd5f01ab4d 100644
--- a/Platform/RISC-V/PlatformPkg/Universal/Sec/SecMain.inf
+++ b/Platform/RISC-V/PlatformPkg/Universal/Sec/SecMain.inf
@@ -66,6 +66,10 @@
   gUefiRiscVPlatformPkgTokenSpaceGuid.PcdBootableHartIndexToId
   gUefiRiscVPlatformPkgTokenSpaceGuid.PcdRootFirmwareDomainBaseAddress
   gUefiRiscVPlatformPkgTokenSpaceGuid.PcdRootFirmwareDomainSize
+  gUefiRiscVPlatformPkgTokenSpaceGuid.PcdFirmwareDomainBaseAddress
+  gUefiRiscVPlatformPkgTokenSpaceGuid.PcdFirmwareDomainSize
+  gUefiRiscVPlatformPkgTokenSpaceGuid.PcdVariableFirmwareRegionBaseAddress
+  gUefiRiscVPlatformPkgTokenSpaceGuid.PcdVariableFirmwareRegionSize
   gUefiRiscVPlatformPkgTokenSpaceGuid.PcdOpenSbiStackSize
   gUefiRiscVPlatformPkgTokenSpaceGuid.PcdScratchRamBase
   gUefiRiscVPlatformPkgTokenSpaceGuid.PcdScratchRamSize
diff --git a/Platform/RISC-V/PlatformPkg/Library/OpensbiPlatformLib/Platform.c b/Platform/RISC-V/PlatformPkg/Library/OpensbiPlatformLib/Platform.c
index b477b81d74..c62d235333 100644
--- a/Platform/RISC-V/PlatformPkg/Library/OpensbiPlatformLib/Platform.c
+++ b/Platform/RISC-V/PlatformPkg/Library/OpensbiPlatformLib/Platform.c
@@ -197,68 +197,11 @@ static u64 generic_tlbr_flush_limit(void)
   return SBI_PLATFORM_TLB_RANGE_FLUSH_LIMIT_DEFAULT;
 }
 
-static int generic_system_reset_check(u32 reset_type, u32 reset_reason)
-{
-  if (generic_plat && generic_plat->system_reset_check)
-    return generic_plat->system_reset_check(reset_type,
-              reset_reason,
-              generic_plat_match);
-  return fdt_system_reset_check(reset_type, reset_reason);
-}
-
-static void generic_system_reset(u32 reset_type, u32 reset_reason)
-{
-  if (generic_plat && generic_plat->system_reset) {
-    generic_plat->system_reset(reset_type, reset_reason,
-             generic_plat_match);
-    return;
-  }
-
-  fdt_system_reset(reset_type, reset_reason);
-}
-
-#define EDK2_ROOT_FW_REGION       0
-#define EDK2_FW_REGION            1
-#define EDK2_VARIABLE_REGION    2
-#define EDK2_ALL_REGION            3
-#define EDK2_END_REGION            4
-static struct sbi_domain_memregion root_memregs[EDK2_END_REGION + 1] = { 0 };
-
-struct sbi_domain_memregion *get_mem_regions(void) {
-  /* EDK2 root firmware domain memory region */
-  root_memregs[EDK2_ROOT_FW_REGION].order = log2roundup(FixedPcdGet32(PcdRootFirmwareDomainSize));
-  root_memregs[EDK2_ROOT_FW_REGION].base = FixedPcdGet32(PcdRootFirmwareDomainBaseAddress);
-  root_memregs[EDK2_ROOT_FW_REGION].flags = 0;
-
-    /*EDK2 firmware domain memory region */
-  root_memregs[EDK2_FW_REGION].order = log2roundup(FixedPcdGet32(PcdFirmwareDomainSize));
-  root_memregs[EDK2_FW_REGION].base = FixedPcdGet32(PcdFirmwareDomainBaseAddress);
-  root_memregs[EDK2_FW_REGION].flags = SBI_DOMAIN_MEMREGION_EXECUTABLE | SBI_DOMAIN_MEMREGION_READABLE;
-
-    /*EDK2 firmware domain memory region */
-  root_memregs[EDK2_VARIABLE_REGION].order = log2roundup(FixedPcdGet32(PcdVariableFirmwareRegionSize));
-  root_memregs[EDK2_VARIABLE_REGION].base = FixedPcdGet32(PcdVariableFirmwareRegionBaseAddress);
-  root_memregs[EDK2_VARIABLE_REGION].flags = SBI_DOMAIN_MEMREGION_READABLE | SBI_DOMAIN_MEMREGION_WRITEABLE;
-
-  /* EDK2 domain allow everything memory region */
-  root_memregs[EDK2_ALL_REGION].order = __riscv_xlen;
-  root_memregs[EDK2_ALL_REGION].base = 0;
-  root_memregs[EDK2_ALL_REGION].flags = (SBI_DOMAIN_MEMREGION_READABLE |
-            SBI_DOMAIN_MEMREGION_WRITEABLE |
-            SBI_DOMAIN_MEMREGION_EXECUTABLE);
-
-  /* EDK2 domain memory region end */
-  root_memregs[EDK2_END_REGION].order = 0;
-
-  return root_memregs;
-}
-
 const struct sbi_platform_operations platform_ops = {
   .early_init            = generic_early_init,
   .final_init            = generic_final_init,
   .early_exit            = generic_early_exit,
   .final_exit            = generic_final_exit,
-  .domains_root_regions  = get_mem_regions,
   .domains_init          = generic_domains_init,
   .console_init          = fdt_serial_init,
   .irqchip_init          = fdt_irqchip_init,
diff --git a/Platform/RISC-V/PlatformPkg/Universal/Sec/Edk2OpenSbiPlatform.c b/Platform/RISC-V/PlatformPkg/Universal/Sec/Edk2OpenSbiPlatform.c
index 79b2f33675..779705489c 100644
--- a/Platform/RISC-V/PlatformPkg/Universal/Sec/Edk2OpenSbiPlatform.c
+++ b/Platform/RISC-V/PlatformPkg/Universal/Sec/Edk2OpenSbiPlatform.c
@@ -117,18 +117,6 @@ int Edk2OpensbiPlatforMMISAGetXLEN (VOID)
     return 0;
 }
 
-/** Get platform specific root domain memory regions */
-struct sbi_domain_memregion *
-Edk2OpensbiPlatformGetMemRegions (VOID)
-{
-    DEBUG((DEBUG_INFO, "%a: Entry\n", __FUNCTION__));
-
-    if (platform_ops.domains_root_regions) {
-        return platform_ops.domains_root_regions ();
-    }
-    return 0;
-}
-
 /** Initialize (or populate) domains for the platform */
 int Edk2OpensbiPlatformDomainsInit (VOID)
 {
@@ -140,25 +128,6 @@ int Edk2OpensbiPlatformDomainsInit (VOID)
     return 0;
 }
 
-/** Write a character to the platform console output */
-VOID Edk2OpensbiPlatformSerialPutc (
-    CHAR8 Ch
-    )
-{
-    if (platform_ops.console_putc) {
-        return platform_ops.console_putc (Ch);
-    }
-}
-
-/** Read a character from the platform console input */
-int Edk2OpensbiPlatformSerialGetc (VOID)
-{
-    if (platform_ops.console_getc) {
-        return platform_ops.console_getc ();
-    }
-    return 0;
-}
-
 /** Initialize the platform console */
 int Edk2OpensbiPlatformSerialInit (VOID)
 {
@@ -193,30 +162,6 @@ VOID Edk2OpensbiPlatformIrqchipExit (VOID)
     }
 }
 
-/** Send IPI to a target HART */
-VOID Edk2OpensbiPlatformIpiSend (
-    UINT32 TargetHart
-    )
-{
-    DEBUG((DEBUG_INFO, "%a: Entry\n", __FUNCTION__));
-
-    if (platform_ops.ipi_send) {
-        return platform_ops.ipi_send (TargetHart);
-    }
-}
-
-/** Clear IPI for a target HART */
-VOID Edk2OpensbiPlatformIpiClear (
-    UINT32 TargetHart
-    )
-{
-    DEBUG((DEBUG_INFO, "%a: Entry\n", __FUNCTION__));
-
-    if (platform_ops.ipi_clear) {
-        return platform_ops.ipi_clear (TargetHart);
-    }
-}
-
 /** Initialize IPI for current HART */
 int Edk2OpensbiPlatformIpiInit (
     BOOLEAN ColdBoot
@@ -251,33 +196,6 @@ UINT64 Edk2OpensbiPlatformTlbrFlushLimit (VOID)
     return 0;
 }
 
-/** Get platform timer value */
-UINT64 Edk2OpensbiPlatformTimerValue (VOID)
-{
-    if (platform_ops.timer_value) {
-        return platform_ops.timer_value ();
-    }
-    return 0;
-}
-
-/** Start platform timer event for current HART */
-VOID Edk2OpensbiPlatformTimerEventStart (
-    UINT64 NextEvent
-    )
-{
-    if (platform_ops.timer_event_start) {
-        return platform_ops.timer_event_start (NextEvent);
-    }
-}
-
-/** Stop platform timer event for current HART */
-VOID Edk2OpensbiPlatformTimerEventStop (VOID)
-{
-    if (platform_ops.timer_event_stop) {
-        return platform_ops.timer_event_stop ();
-    }
-}
-
 /** Initialize platform timer for current HART */
 int Edk2OpensbiPlatformTimerInit (
     BOOLEAN ColdBoot
@@ -301,61 +219,6 @@ VOID Edk2OpensbiPlatformTimerExit (VOID)
     }
 }
 
-/** Bringup the given hart */
-int Edk2OpensbiPlatformHartStart (
-    UINT32 HartId,
-    ulong Saddr
-    )
-{
-    DEBUG((DEBUG_INFO, "%a: Entry\n", __FUNCTION__));
-
-    if (platform_ops.hart_start) {
-        return platform_ops.hart_start (HartId, Saddr);
-    }
-    return 0;
-}
-/**
-  Stop the current hart from running. This call doesn't expect to
-  return if success.
-**/
-int Edk2OpensbiPlatformHartStop (VOID)
-{
-    DEBUG((DEBUG_INFO, "%a: Entry\n", __FUNCTION__));
-
-    if (platform_ops.hart_stop) {
-        return platform_ops.hart_stop ();
-    }
-    return 0;
-}
-
-/**
- Check whether reset type and reason supported by the platform*
-
-**/
-int Edk2OpensbiPlatformSystemResetCheck (
-    UINT32 ResetType,
-    UINT32 ResetReason
-    )
-{
-    if (platform_ops.system_reset_check) {
-        return platform_ops.system_reset_check (ResetType, ResetReason);
-    }
-    return 0;
-}
-
-/** Reset the platform */
-VOID Edk2OpensbiPlatformSystemReset (
-    UINT32 ResetType,
-    UINT32 ResetReason
-    )
-{
-    DEBUG((DEBUG_INFO, "%a: Entry\n", __FUNCTION__));
-
-    if (platform_ops.system_reset) {
-        return platform_ops.system_reset (ResetType, ResetReason);
-    }
-}
-
 /** platform specific SBI extension implementation probe function */
 int Edk2OpensbiPlatformVendorExtCheck (
     long ExtId
@@ -400,27 +263,15 @@ const struct sbi_platform_operations Edk2OpensbiPlatformOps = {
     .final_exit             = Edk2OpensbiPlatformFinalExit,
     .misa_check_extension   = Edk2OpensbiPlatforMMISACheckExtension,
     .misa_get_xlen          = Edk2OpensbiPlatforMMISAGetXLEN,
-    .domains_root_regions   = Edk2OpensbiPlatformGetMemRegions,
     .domains_init           = Edk2OpensbiPlatformDomainsInit,
-    .console_putc           = Edk2OpensbiPlatformSerialPutc,
-    .console_getc           = Edk2OpensbiPlatformSerialGetc,
     .console_init           = Edk2OpensbiPlatformSerialInit,
     .irqchip_init           = Edk2OpensbiPlatformIrqchipInit,
     .irqchip_exit           = Edk2OpensbiPlatformIrqchipExit,
-    .ipi_send               = Edk2OpensbiPlatformIpiSend,
-    .ipi_clear              = Edk2OpensbiPlatformIpiClear,
     .ipi_init               = Edk2OpensbiPlatformIpiInit,
     .ipi_exit               = Edk2OpensbiPlatformIpiExit,
     .get_tlbr_flush_limit   = Edk2OpensbiPlatformTlbrFlushLimit,
-    .timer_value            = Edk2OpensbiPlatformTimerValue,
-    .timer_event_stop       = Edk2OpensbiPlatformTimerEventStop,
-    .timer_event_start      = Edk2OpensbiPlatformTimerEventStart,
     .timer_init             = Edk2OpensbiPlatformTimerInit,
     .timer_exit             = Edk2OpensbiPlatformTimerExit,
-    .hart_start             = Edk2OpensbiPlatformHartStart,
-    .hart_stop              = Edk2OpensbiPlatformHartStop,
-    .system_reset_check     = Edk2OpensbiPlatformSystemResetCheck,
-    .system_reset           = Edk2OpensbiPlatformSystemReset,
     .vendor_ext_check       = Edk2OpensbiPlatformVendorExtCheck,
     .vendor_ext_provider    = Edk2OpensbiPlatformVendorExtProvider,
 };
diff --git a/Platform/RISC-V/PlatformPkg/Universal/Sec/SecMain.c b/Platform/RISC-V/PlatformPkg/Universal/Sec/SecMain.c
index 93ff8a598d..3bc3690047 100644
--- a/Platform/RISC-V/PlatformPkg/Universal/Sec/SecMain.c
+++ b/Platform/RISC-V/PlatformPkg/Universal/Sec/SecMain.c
@@ -15,10 +15,12 @@
 #include <sbi/riscv_asm.h>
 #include <sbi/riscv_atomic.h>
 #include <sbi/sbi_console.h>  // Reference to header file in opensbi
+#include <sbi/sbi_domain.h>
 #include <sbi/sbi_hart.h>     // Reference to header file in opensbi
-#include <sbi/sbi_hartmask.h>    // Reference to header file in opensbi
+#include <sbi/sbi_hartmask.h> // Reference to header file in opensbi
 #include <sbi/sbi_scratch.h>  // Reference to header file in opensbi
 #include <sbi/sbi_platform.h> // Reference to header file in opensbi
+#include <sbi/sbi_math.h>     // Reference to header file in opensbi
 #include <sbi/sbi_init.h>     // Reference to header file in opensbi
 #include <sbi/sbi_ecall.h>    // Reference to header file in opensbi
 #include <sbi/sbi_trap.h>     // Reference to header file in opensbi
@@ -31,8 +33,41 @@ extern struct sbi_platform_operations Edk2OpensbiPlatformOps;
 atomic_t BootHartDone = ATOMIC_INITIALIZER(0);
 atomic_t NonBootHartMessageLock = ATOMIC_INITIALIZER(0);
 
+int sbi_domain_root_add_memregion(const struct sbi_domain_memregion *reg);
+
 typedef struct sbi_scratch *(*hartid2scratch)(ulong hartid, ulong hartindex);
 
+struct sbi_domain_memregion fw_memregs;
+
+int SecSetEdk2FwMemoryRegions (VOID) {
+  int Ret;
+
+  Ret = 0;
+
+  //
+  // EDK2 PEI domain memory region
+  //
+  fw_memregs.order = log2roundup(FixedPcdGet32(PcdFirmwareDomainSize));
+  fw_memregs.base = FixedPcdGet32(PcdFirmwareDomainBaseAddress);
+  fw_memregs.flags = SBI_DOMAIN_MEMREGION_EXECUTABLE | SBI_DOMAIN_MEMREGION_READABLE;
+  Ret = sbi_domain_root_add_memregion ((const struct sbi_domain_memregion *)&fw_memregs);
+  if (Ret != 0) {
+    DEBUG ((DEBUG_ERROR, "%a: Add firmware regiosn of FW Domain fail\n", __FUNCTION__));
+  }
+
+  //
+  // EDK2 EFI Variable domain memory region
+  //
+  fw_memregs.order = log2roundup(FixedPcdGet32(PcdVariableFirmwareRegionSize));
+  fw_memregs.base = FixedPcdGet32(PcdVariableFirmwareRegionBaseAddress);
+  fw_memregs.flags = SBI_DOMAIN_MEMREGION_READABLE | SBI_DOMAIN_MEMREGION_WRITEABLE;
+  Ret = sbi_domain_root_add_memregion ((const struct sbi_domain_memregion *)&fw_memregs);
+  if (Ret != 0) {
+    DEBUG ((DEBUG_ERROR, "%a: Add firmware regiosn of variable FW Domain fail\n", __FUNCTION__));
+  }
+  return Ret;
+}
+
 /**
   Locates a section within a series of sections
   with the specified section type.
@@ -405,6 +440,13 @@ SecPostOpenSbiPlatformEarlylInit(
     DEBUG ((DEBUG_INFO, "%a: Non boot hart %d.\n", __FUNCTION__, HartId));
     return 0;
   }
+  //
+  // Setup firmware memory region.
+  //
+  if (SecSetEdk2FwMemoryRegions () != 0) {
+    ASSERT (FALSE);
+  }
+
   //
   // Boot HART is already in the process of OpenSBI initialization.
   // We can let other HART to keep booting.
@@ -477,7 +519,7 @@ SecPostOpenSbiPlatformFinalInit (
     }
   }
 
-  DEBUG((DEBUG_INFO, "%a: Jump to PEI Core with \n", __FUNCTION__));
+  DEBUG((DEBUG_INFO, "%a: Will jump to PEI Core in OpenSBI with \n", __FUNCTION__));
   DEBUG((DEBUG_INFO, "  sbi_scratch = %x\n", SbiScratch));
   DEBUG((DEBUG_INFO, "  sbi_platform = %x\n", SbiPlatform));
   DEBUG((DEBUG_INFO, "  FirmwareContext = %x\n", FirmwareContext));
@@ -793,7 +835,7 @@ VOID EFIAPI SecCoreStartUpWithStack(
   sbi_init(Scratch);
 }
 
-void OpensbiDebugPrint (char *debugstr, ...)
+VOID OpensbiDebugPrint (CHAR8 *debugstr, ...)
 {
   VA_LIST  Marker;
 
diff --git a/Platform/RISC-V/PlatformPkg/Universal/Sec/Riscv64/SecEntry.S b/Platform/RISC-V/PlatformPkg/Universal/Sec/Riscv64/SecEntry.S
index dc410703e0..96087738a3 100644
--- a/Platform/RISC-V/PlatformPkg/Universal/Sec/Riscv64/SecEntry.S
+++ b/Platform/RISC-V/PlatformPkg/Universal/Sec/Riscv64/SecEntry.S
@@ -18,6 +18,12 @@
 
 #include <SecMain.h>
 
+.macro  MOV_3R __d0, __s0, __d1, __s1, __d2, __s2
+        add     \__d0, \__s0, zero
+        add     \__d1, \__s1, zero
+        add     \__d2, \__s2, zero
+.endm
+
 .text
 .align 3
 
@@ -90,7 +96,11 @@ _scratch_init:
   la    a4, _hartid_to_scratch
   sd    a4, SBI_SCRATCH_HARTID_TO_SCRATCH_OFFSET(tp) /* Save _hartid_to_scratch function in scratch buffer*/
   sd    zero, SBI_SCRATCH_TMP0_OFFSET(tp)
-
+  /* Store trap-exit function address in scratch space */
+  lla   a4, _trap_exit
+  sd    a4, SBI_SCRATCH_TRAP_EXIT_OFFSET(tp)
+  /* Clear tmp0 in scratch space */
+  sd    zero, SBI_SCRATCH_TMP0_OFFSET(tp)
 #ifdef FW_OPTIONS
   li    a4, FW_OPTIONS
   sd    a4, SBI_SCRATCH_OPTIONS_OFFSET(tp)
@@ -322,160 +332,174 @@ _uninitialized_hart_wait:
   wfi
   j     _uninitialized_hart_wait
 
-  .align 3
-  .section .entry, "ax", %progbits
-  .align 3
-  .globl _trap_handler
-_trap_handler:
-
+.macro  TRAP_SAVE_AND_SETUP_SP_T0
   /* Swap TP and MSCRATCH */
-  csrrw tp, CSR_MSCRATCH, tp
+  csrrw   tp, CSR_MSCRATCH, tp
 
   /* Save T0 in scratch space */
-  REG_S t0, SBI_SCRATCH_TMP0_OFFSET(tp)
+  REG_S   t0, SBI_SCRATCH_TMP0_OFFSET(tp)
 
-  /* Check which mode we came from */
-  csrr  t0, CSR_MSTATUS
-  srl   t0, t0, MSTATUS_MPP_SHIFT
-  and   t0, t0, PRV_M
-  xori  t0, t0, PRV_M
-  beq   t0, zero, _trap_handler_m_mode
-
-  /* We came from S-mode or U-mode */
-_trap_handler_s_mode:
-  /* Set T0 to original SP */
-  add   t0, sp, zero
-
-  /* Setup exception stack */
-  add   sp, tp, -(SBI_TRAP_REGS_SIZE)
-
-  /* Jump to code common for all modes */
-  j     _trap_handler_all_mode
-
-  /* We came from M-mode */
-_trap_handler_m_mode:
-  /* Set T0 to original SP */
-  add   t0, sp, zero
-
-  /* Re-use current SP as exception stack */
-  add   sp, sp, -(SBI_TRAP_REGS_SIZE)
-
-_trap_handler_all_mode:
-  /* Save original SP (from T0) on stack */
-  REG_S t0, SBI_TRAP_REGS_OFFSET(sp)(sp)
+  /*
+   * Set T0 to appropriate exception stack
+   *
+   * Came_From_M_Mode = ((MSTATUS.MPP < PRV_M) ? 1 : 0) - 1;
+   * Exception_Stack = TP ^ (Came_From_M_Mode & (SP ^ TP))
+   *
+   * Came_From_M_Mode = 0    ==>    Exception_Stack = TP
+   * Came_From_M_Mode = -1   ==>    Exception_Stack = SP
+   */
+  csrr    t0, CSR_MSTATUS
+  srl     t0, t0, MSTATUS_MPP_SHIFT
+  and     t0, t0, PRV_M
+  slti    t0, t0, PRV_M
+  add     t0, t0, -1
+  xor     sp, sp, tp
+  and     t0, t0, sp
+  xor     sp, sp, tp
+  xor     t0, tp, t0
+
+  /* Save original SP on exception stack */
+  REG_S   sp, (SBI_TRAP_REGS_OFFSET(sp) - SBI_TRAP_REGS_SIZE)(t0)
+
+  /* Set SP to exception stack and make room for trap registers */
+  add     sp, t0, -(SBI_TRAP_REGS_SIZE)
 
   /* Restore T0 from scratch space */
-  REG_L t0, SBI_SCRATCH_TMP0_OFFSET(tp)
+  REG_L   t0, SBI_SCRATCH_TMP0_OFFSET(tp)
 
   /* Save T0 on stack */
-  REG_S t0, SBI_TRAP_REGS_OFFSET(t0)(sp)
+  REG_S   t0, SBI_TRAP_REGS_OFFSET(t0)(sp)
 
   /* Swap TP and MSCRATCH */
-  csrrw tp, CSR_MSCRATCH, tp
+  csrrw   tp, CSR_MSCRATCH, tp
+.endm
 
+.macro  TRAP_SAVE_MEPC_MSTATUS have_mstatush
   /* Save MEPC and MSTATUS CSRs */
-  csrr  t0, CSR_MEPC
-  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
+  csrr    t0, CSR_MEPC
+  REG_S   t0, SBI_TRAP_REGS_OFFSET(mepc)(sp)
+  csrr    t0, CSR_MSTATUS
+  REG_S   t0, SBI_TRAP_REGS_OFFSET(mstatus)(sp)
+.if \have_mstatush
+  csrr    t0, CSR_MSTATUSH
+  REG_S   t0, SBI_TRAP_REGS_OFFSET(mstatusH)(sp)
+.else
+  REG_S   zero, SBI_TRAP_REGS_OFFSET(mstatusH)(sp)
+.endif
+.endm
+
+.macro  TRAP_SAVE_GENERAL_REGS_EXCEPT_SP_T0
+  /* Save all general regisers except SP and T0 */
+  REG_S   zero, SBI_TRAP_REGS_OFFSET(zero)(sp)
+  REG_S   ra, SBI_TRAP_REGS_OFFSET(ra)(sp)
+  REG_S   gp, SBI_TRAP_REGS_OFFSET(gp)(sp)
+  REG_S   tp, SBI_TRAP_REGS_OFFSET(tp)(sp)
+  REG_S   t1, SBI_TRAP_REGS_OFFSET(t1)(sp)
+  REG_S   t2, SBI_TRAP_REGS_OFFSET(t2)(sp)
+  REG_S   s0, SBI_TRAP_REGS_OFFSET(s0)(sp)
+  REG_S   s1, SBI_TRAP_REGS_OFFSET(s1)(sp)
+  REG_S   a0, SBI_TRAP_REGS_OFFSET(a0)(sp)
+  REG_S   a1, SBI_TRAP_REGS_OFFSET(a1)(sp)
+  REG_S   a2, SBI_TRAP_REGS_OFFSET(a2)(sp)
+  REG_S   a3, SBI_TRAP_REGS_OFFSET(a3)(sp)
+  REG_S   a4, SBI_TRAP_REGS_OFFSET(a4)(sp)
+  REG_S   a5, SBI_TRAP_REGS_OFFSET(a5)(sp)
+  REG_S   a6, SBI_TRAP_REGS_OFFSET(a6)(sp)
+  REG_S   a7, SBI_TRAP_REGS_OFFSET(a7)(sp)
+  REG_S   s2, SBI_TRAP_REGS_OFFSET(s2)(sp)
+  REG_S   s3, SBI_TRAP_REGS_OFFSET(s3)(sp)
+  REG_S   s4, SBI_TRAP_REGS_OFFSET(s4)(sp)
+  REG_S   s5, SBI_TRAP_REGS_OFFSET(s5)(sp)
+  REG_S   s6, SBI_TRAP_REGS_OFFSET(s6)(sp)
+  REG_S   s7, SBI_TRAP_REGS_OFFSET(s7)(sp)
+  REG_S   s8, SBI_TRAP_REGS_OFFSET(s8)(sp)
+  REG_S   s9, SBI_TRAP_REGS_OFFSET(s9)(sp)
+  REG_S   s10, SBI_TRAP_REGS_OFFSET(s10)(sp)
+  REG_S   s11, SBI_TRAP_REGS_OFFSET(s11)(sp)
+  REG_S   t3, SBI_TRAP_REGS_OFFSET(t3)(sp)
+  REG_S   t4, SBI_TRAP_REGS_OFFSET(t4)(sp)
+  REG_S   t5, SBI_TRAP_REGS_OFFSET(t5)(sp)
+  REG_S   t6, SBI_TRAP_REGS_OFFSET(t6)(sp)
+.endm
+
+.macro  TRAP_CALL_C_ROUTINE
+  /* Call C routine */
+  add     a0, sp, zero
+  call    sbi_trap_handler
+.endm
+
+.macro  TRAP_RESTORE_GENERAL_REGS_EXCEPT_A0_T0
+  /* Restore all general regisers except A0 and T0 */
+  REG_L   ra, SBI_TRAP_REGS_OFFSET(ra)(a0)
+  REG_L   sp, SBI_TRAP_REGS_OFFSET(sp)(a0)
+  REG_L   gp, SBI_TRAP_REGS_OFFSET(gp)(a0)
+  REG_L   tp, SBI_TRAP_REGS_OFFSET(tp)(a0)
+  REG_L   t1, SBI_TRAP_REGS_OFFSET(t1)(a0)
+  REG_L   t2, SBI_TRAP_REGS_OFFSET(t2)(a0)
+  REG_L   s0, SBI_TRAP_REGS_OFFSET(s0)(a0)
+  REG_L   s1, SBI_TRAP_REGS_OFFSET(s1)(a0)
+  REG_L   a1, SBI_TRAP_REGS_OFFSET(a1)(a0)
+  REG_L   a2, SBI_TRAP_REGS_OFFSET(a2)(a0)
+  REG_L   a3, SBI_TRAP_REGS_OFFSET(a3)(a0)
+  REG_L   a4, SBI_TRAP_REGS_OFFSET(a4)(a0)
+  REG_L   a5, SBI_TRAP_REGS_OFFSET(a5)(a0)
+  REG_L   a6, SBI_TRAP_REGS_OFFSET(a6)(a0)
+  REG_L   a7, SBI_TRAP_REGS_OFFSET(a7)(a0)
+  REG_L   s2, SBI_TRAP_REGS_OFFSET(s2)(a0)
+  REG_L   s3, SBI_TRAP_REGS_OFFSET(s3)(a0)
+  REG_L   s4, SBI_TRAP_REGS_OFFSET(s4)(a0)
+  REG_L   s5, SBI_TRAP_REGS_OFFSET(s5)(a0)
+  REG_L   s6, SBI_TRAP_REGS_OFFSET(s6)(a0)
+  REG_L   s7, SBI_TRAP_REGS_OFFSET(s7)(a0)
+  REG_L   s8, SBI_TRAP_REGS_OFFSET(s8)(a0)
+  REG_L   s9, SBI_TRAP_REGS_OFFSET(s9)(a0)
+  REG_L   s10, SBI_TRAP_REGS_OFFSET(s10)(a0)
+  REG_L   s11, SBI_TRAP_REGS_OFFSET(s11)(a0)
+  REG_L   t3, SBI_TRAP_REGS_OFFSET(t3)(a0)
+  REG_L   t4, SBI_TRAP_REGS_OFFSET(t4)(a0)
+  REG_L   t5, SBI_TRAP_REGS_OFFSET(t5)(a0)
+  REG_L   t6, SBI_TRAP_REGS_OFFSET(t6)(a0)
+.endm
+
+.macro  TRAP_RESTORE_MEPC_MSTATUS have_mstatush
+  /* Restore MEPC and MSTATUS CSRs */
+  REG_L   t0, SBI_TRAP_REGS_OFFSET(mepc)(a0)
+  csrw    CSR_MEPC, t0
+  REG_L   t0, SBI_TRAP_REGS_OFFSET(mstatus)(a0)
+  csrw    CSR_MSTATUS, t0
+.if \have_mstatush
+  REG_L   t0, SBI_TRAP_REGS_OFFSET(mstatusH)(a0)
+  csrw    CSR_MSTATUSH, t0
+.endif
+.endm
+
+.macro TRAP_RESTORE_A0_T0
+  /* Restore T0 */
+  REG_L   t0, SBI_TRAP_REGS_OFFSET(t0)(a0)
 
-  /* Save all general registers except SP and T0 */
-  REG_S zero, SBI_TRAP_REGS_OFFSET(zero)(sp)
-  REG_S ra, SBI_TRAP_REGS_OFFSET(ra)(sp)
-  REG_S gp, SBI_TRAP_REGS_OFFSET(gp)(sp)
-  REG_S tp, SBI_TRAP_REGS_OFFSET(tp)(sp)
-  REG_S t1, SBI_TRAP_REGS_OFFSET(t1)(sp)
-  REG_S t2, SBI_TRAP_REGS_OFFSET(t2)(sp)
-  REG_S s0, SBI_TRAP_REGS_OFFSET(s0)(sp)
-  REG_S s1, SBI_TRAP_REGS_OFFSET(s1)(sp)
-  REG_S a0, SBI_TRAP_REGS_OFFSET(a0)(sp)
-  REG_S a1, SBI_TRAP_REGS_OFFSET(a1)(sp)
-  REG_S a2, SBI_TRAP_REGS_OFFSET(a2)(sp)
-  REG_S a3, SBI_TRAP_REGS_OFFSET(a3)(sp)
-  REG_S a4, SBI_TRAP_REGS_OFFSET(a4)(sp)
-  REG_S a5, SBI_TRAP_REGS_OFFSET(a5)(sp)
-  REG_S a6, SBI_TRAP_REGS_OFFSET(a6)(sp)
-  REG_S a7, SBI_TRAP_REGS_OFFSET(a7)(sp)
-  REG_S s2, SBI_TRAP_REGS_OFFSET(s2)(sp)
-  REG_S s3, SBI_TRAP_REGS_OFFSET(s3)(sp)
-  REG_S s4, SBI_TRAP_REGS_OFFSET(s4)(sp)
-  REG_S s5, SBI_TRAP_REGS_OFFSET(s5)(sp)
-  REG_S s6, SBI_TRAP_REGS_OFFSET(s6)(sp)
-  REG_S s7, SBI_TRAP_REGS_OFFSET(s7)(sp)
-  REG_S s8, SBI_TRAP_REGS_OFFSET(s8)(sp)
-  REG_S s9, SBI_TRAP_REGS_OFFSET(s9)(sp)
-  REG_S s10, SBI_TRAP_REGS_OFFSET(s10)(sp)
-  REG_S s11, SBI_TRAP_REGS_OFFSET(s11)(sp)
-  REG_S t3, SBI_TRAP_REGS_OFFSET(t3)(sp)
-  REG_S t4, SBI_TRAP_REGS_OFFSET(t4)(sp)
-  REG_S t5, SBI_TRAP_REGS_OFFSET(t5)(sp)
-  REG_S t6, SBI_TRAP_REGS_OFFSET(t6)(sp)
+  /* Restore A0 */
+  REG_L   a0, SBI_TRAP_REGS_OFFSET(a0)(a0)
+.endm
 
-  /* Call C routine */
-  add   a0, sp, zero
-  call  sbi_trap_handler
-
-  /* Restore all general registers except SP and T0 */
-  REG_L ra, SBI_TRAP_REGS_OFFSET(ra)(sp)
-  REG_L gp, SBI_TRAP_REGS_OFFSET(gp)(sp)
-  REG_L tp, SBI_TRAP_REGS_OFFSET(tp)(sp)
-  REG_L t1, SBI_TRAP_REGS_OFFSET(t1)(sp)
-  REG_L t2, SBI_TRAP_REGS_OFFSET(t2)(sp)
-  REG_L s0, SBI_TRAP_REGS_OFFSET(s0)(sp)
-  REG_L s1, SBI_TRAP_REGS_OFFSET(s1)(sp)
-  REG_L a0, SBI_TRAP_REGS_OFFSET(a0)(sp)
-  REG_L a1, SBI_TRAP_REGS_OFFSET(a1)(sp)
-  REG_L a2, SBI_TRAP_REGS_OFFSET(a2)(sp)
-  REG_L a3, SBI_TRAP_REGS_OFFSET(a3)(sp)
-  REG_L a4, SBI_TRAP_REGS_OFFSET(a4)(sp)
-  REG_L a5, SBI_TRAP_REGS_OFFSET(a5)(sp)
-  REG_L a6, SBI_TRAP_REGS_OFFSET(a6)(sp)
-  REG_L a7, SBI_TRAP_REGS_OFFSET(a7)(sp)
-  REG_L s2, SBI_TRAP_REGS_OFFSET(s2)(sp)
-  REG_L s3, SBI_TRAP_REGS_OFFSET(s3)(sp)
-  REG_L s4, SBI_TRAP_REGS_OFFSET(s4)(sp)
-  REG_L s5, SBI_TRAP_REGS_OFFSET(s5)(sp)
-  REG_L s6, SBI_TRAP_REGS_OFFSET(s6)(sp)
-  REG_L s7, SBI_TRAP_REGS_OFFSET(s7)(sp)
-  REG_L s8, SBI_TRAP_REGS_OFFSET(s8)(sp)
-  REG_L s9, SBI_TRAP_REGS_OFFSET(s9)(sp)
-  REG_L s10, SBI_TRAP_REGS_OFFSET(s10)(sp)
-  REG_L s11, SBI_TRAP_REGS_OFFSET(s11)(sp)
-  REG_L t3, SBI_TRAP_REGS_OFFSET(t3)(sp)
-  REG_L t4, SBI_TRAP_REGS_OFFSET(t4)(sp)
-  REG_L t5, SBI_TRAP_REGS_OFFSET(t5)(sp)
-  REG_L t6, SBI_TRAP_REGS_OFFSET(t6)(sp)
+  .section .entry, "ax", %progbits
+  .align 3
+  .globl _trap_handler
+  .globl _trap_exit
+_trap_handler:
+  TRAP_SAVE_AND_SETUP_SP_T0
 
-  /* Restore MEPC and MSTATUS CSRs */
-  REG_L t0, SBI_TRAP_REGS_OFFSET(mepc)(sp)
-  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
+  TRAP_SAVE_MEPC_MSTATUS 0
 
-  /* Restore T0 */
-  REG_L   t0, SBI_TRAP_REGS_OFFSET(t0)(sp)
+  TRAP_SAVE_GENERAL_REGS_EXCEPT_SP_T0
+
+  TRAP_CALL_C_ROUTINE
+
+_trap_exit:
+  TRAP_RESTORE_GENERAL_REGS_EXCEPT_A0_T0
+
+  TRAP_RESTORE_MEPC_MSTATUS 0
 
-  /* Restore SP */
-  REG_L   sp, SBI_TRAP_REGS_OFFSET(sp)(sp)
+  TRAP_RESTORE_A0_T0
 
   mret
 
-- 
2.31.1