diff --git a/am/include/arch/riscv.h b/am/include/arch/riscv.h
index 9c9de20..9709050 100644
--- a/am/include/arch/riscv.h
+++ b/am/include/arch/riscv.h
@@ -1,9 +1,15 @@
 #ifndef ARCH_H__
 #define ARCH_H__
 
+#ifdef __riscv_e
+#define NR_REGS 16
+#else
+#define NR_REGS 32
+#endif
+
 struct Context {
   // TODO: fix the order of these members to match trap.S
-  uintptr_t mepc, mcause, gpr[32], mstatus;
+  uintptr_t mepc, mcause, gpr[NR_REGS], mstatus;
   void *pdir;
 };