Cleanup and improve kvm_load/save_registers usage