2 files changed, 15 insertions, 4 deletions

diff --git a/include/linux/kvm_host.h b/include/linux/kvm_host.h
index ab428552af8e..9272db03a3e5 100644
--- a/include/linux/kvm_host.h
+++ b/include/linux/kvm_host.h
@@ -513,6 +513,7 @@ struct kvm_assigned_dev_kernel {
         struct kvm *kvm;
         spinlock_t intx_lock;
         char irq_name[32];
+        struct pci_saved_state *pci_saved_state;
 };
 
 struct kvm_irq_mask_notifier {
diff --git a/virt/kvm/assigned-dev.c b/virt/kvm/assigned-dev.c
index ae72ae604c89..6cc4b97ec458 100644
--- a/virt/kvm/assigned-dev.c
+++ b/virt/kvm/assigned-dev.c
@@ -197,8 +197,13 @@ static void kvm_free_assigned_device(struct kvm *kvm,
 {
         kvm_free_assigned_irq(kvm, assigned_dev);
 
-        __pci_reset_function(assigned_dev->dev);
-        pci_restore_state(assigned_dev->dev);
+        pci_reset_function(assigned_dev->dev);
+        if (pci_load_and_free_saved_state(assigned_dev->dev,
+                                          &assigned_dev->pci_saved_state))
+                printk(KERN_INFO "%s: Couldn't reload %s saved state\n",
+                       __func__, dev_name(&assigned_dev->dev->dev));
+        else
+                pci_restore_state(assigned_dev->dev);
 
         pci_release_regions(assigned_dev->dev);
         pci_disable_device(assigned_dev->dev);
@@ -516,7 +521,10 @@ static int kvm_vm_ioctl_assign_device(struct kvm *kvm,
 
         pci_reset_function(dev);
         pci_save_state(dev);
-
+        match->pci_saved_state = pci_store_saved_state(dev);
+        if (!match->pci_saved_state)
+                printk(KERN_DEBUG "%s: Couldn't store %s saved state\n",
+                       __func__, dev_name(&dev->dev));
         match->assigned_dev_id = assigned_dev->assigned_dev_id;
         match->host_segnr = assigned_dev->segnr;
         match->host_busnr = assigned_dev->busnr;
@@ -546,7 +554,9 @@ out:
         mutex_unlock(&kvm->lock);
         return r;
 out_list_del:
-        pci_restore_state(dev);
+        if (pci_load_and_free_saved_state(dev, &match->pci_saved_state))
+                printk(KERN_INFO "%s: Couldn't reload %s saved state\n",
+                       __func__, dev_name(&dev->dev));
         list_del(&match->list);
         pci_release_regions(dev);
 out_disable: