*/
 
 #include <linux/bitops.h>
+#include <linux/entry-kvm.h>
 #include <linux/errno.h>
 #include <linux/err.h>
 #include <linux/kdebug.h>
        run->exit_reason = KVM_EXIT_UNKNOWN;
        while (ret > 0) {
                /* Check conditions before entering the guest */
-               cond_resched();
+               ret = xfer_to_guest_mode_handle_work(vcpu);
+               if (!ret)
+                       ret = 1;
 
                kvm_riscv_gstage_vmid_update(vcpu);
 
 
                local_irq_disable();
 
-               /*
-                * Exit if we have a signal pending so that we can deliver
-                * the signal to user space.
-                */
-               if (signal_pending(current)) {
-                       ret = -EINTR;
-                       run->exit_reason = KVM_EXIT_INTR;
-                       ++vcpu->stat.signal_exits;
-               }
-
                /*
                 * Ensure we set mode to IN_GUEST_MODE after we disable
                 * interrupts and before the final VCPU requests check.
 
                if (ret <= 0 ||
                    kvm_riscv_gstage_vmid_ver_changed(&vcpu->kvm->arch.vmid) ||
-                   kvm_request_pending(vcpu)) {
+                   kvm_request_pending(vcpu) ||
+                   xfer_to_guest_mode_work_pending()) {
                        vcpu->mode = OUTSIDE_GUEST_MODE;
                        local_irq_enable();
                        kvm_vcpu_srcu_read_lock(vcpu);