#include "kvm_util.h"
 #include <linux/stringify.h>
 
-static inline uint64_t __kvm_reg_id(uint64_t type, uint64_t idx,
-                                   uint64_t  size)
+static inline uint64_t __kvm_reg_id(uint64_t type, uint64_t subtype,
+                                   uint64_t idx, uint64_t size)
 {
-       return KVM_REG_RISCV | type | idx | size;
+       return KVM_REG_RISCV | type | subtype | idx | size;
 }
 
 #if __riscv_xlen == 64
 #define KVM_REG_SIZE_ULONG     KVM_REG_SIZE_U32
 #endif
 
-#define RISCV_CONFIG_REG(name) __kvm_reg_id(KVM_REG_RISCV_CONFIG, \
-                                            KVM_REG_RISCV_CONFIG_REG(name), \
-                                            KVM_REG_SIZE_ULONG)
+#define RISCV_CONFIG_REG(name)         __kvm_reg_id(KVM_REG_RISCV_CONFIG, 0,           \
+                                                    KVM_REG_RISCV_CONFIG_REG(name),    \
+                                                    KVM_REG_SIZE_ULONG)
 
-#define RISCV_CORE_REG(name)   __kvm_reg_id(KVM_REG_RISCV_CORE, \
-                                            KVM_REG_RISCV_CORE_REG(name), \
-                                            KVM_REG_SIZE_ULONG)
+#define RISCV_CORE_REG(name)           __kvm_reg_id(KVM_REG_RISCV_CORE, 0,             \
+                                                    KVM_REG_RISCV_CORE_REG(name),      \
+                                                    KVM_REG_SIZE_ULONG)
 
-#define RISCV_CSR_REG(name)    __kvm_reg_id(KVM_REG_RISCV_CSR, \
-                                            KVM_REG_RISCV_CSR_REG(name), \
-                                            KVM_REG_SIZE_ULONG)
+#define RISCV_GENERAL_CSR_REG(name)    __kvm_reg_id(KVM_REG_RISCV_CSR,                 \
+                                                    KVM_REG_RISCV_CSR_GENERAL,         \
+                                                    KVM_REG_RISCV_CSR_REG(name),       \
+                                                    KVM_REG_SIZE_ULONG)
 
-#define RISCV_TIMER_REG(name)  __kvm_reg_id(KVM_REG_RISCV_TIMER, \
-                                            KVM_REG_RISCV_TIMER_REG(name), \
-                                            KVM_REG_SIZE_U64)
+#define RISCV_TIMER_REG(name)          __kvm_reg_id(KVM_REG_RISCV_TIMER, 0,            \
+                                                    KVM_REG_RISCV_TIMER_REG(name),     \
+                                                    KVM_REG_SIZE_U64)
 
-#define RISCV_ISA_EXT_REG(idx) __kvm_reg_id(KVM_REG_RISCV_ISA_EXT, \
-                                            idx, KVM_REG_SIZE_ULONG)
+#define RISCV_ISA_EXT_REG(idx)         __kvm_reg_id(KVM_REG_RISCV_ISA_EXT,             \
+                                                    KVM_REG_RISCV_ISA_SINGLE,          \
+                                                    idx, KVM_REG_SIZE_ULONG)
+
+#define RISCV_SBI_EXT_REG(idx)         __kvm_reg_id(KVM_REG_RISCV_SBI_EXT,             \
+                                                    KVM_REG_RISCV_SBI_SINGLE,          \
+                                                    idx, KVM_REG_SIZE_ULONG)
 
 /* L3 index Bit[47:39] */
 #define PGTBL_L3_INDEX_MASK                    0x0000FF8000000000ULL
 
        satp = (vm->pgd >> PGTBL_PAGE_SIZE_SHIFT) & SATP_PPN;
        satp |= SATP_MODE_48;
 
-       vcpu_set_reg(vcpu, RISCV_CSR_REG(satp), satp);
+       vcpu_set_reg(vcpu, RISCV_GENERAL_CSR_REG(satp), satp);
 }
 
 void vcpu_arch_dump(FILE *stream, struct kvm_vcpu *vcpu, uint8_t indent)
        vcpu_set_reg(vcpu, RISCV_CORE_REG(regs.pc), (unsigned long)guest_code);
 
        /* Setup default exception vector of guest */
-       vcpu_set_reg(vcpu, RISCV_CSR_REG(stvec), (unsigned long)guest_unexp_trap);
+       vcpu_set_reg(vcpu, RISCV_GENERAL_CSR_REG(stvec), (unsigned long)guest_unexp_trap);
 
        return vcpu;
 }