1 files changed, 71 insertions, 16 deletions

diff --git a/drivers/pci/ats.c b/drivers/pci/ats.c
index eeb9fb2b47aa..ad8ddbbbf245 100644
--- a/drivers/pci/ats.c
+++ b/drivers/pci/ats.c
@@ -153,23 +153,27 @@ int pci_enable_pri(struct pci_dev *pdev, u32 reqs)
         u32 max_requests;
         int pos;
 
+        if (WARN_ON(pdev->pri_enabled))
+                return -EBUSY;
+
         pos = pci_find_ext_capability(pdev, PCI_EXT_CAP_ID_PRI);
         if (!pos)
                 return -EINVAL;
 
-        pci_read_config_word(pdev, pos + PCI_PRI_CTRL, &control);
         pci_read_config_word(pdev, pos + PCI_PRI_STATUS, &status);
-        if ((control & PCI_PRI_CTRL_ENABLE) ||
-            !(status & PCI_PRI_STATUS_STOPPED))
+        if (!(status & PCI_PRI_STATUS_STOPPED))
                 return -EBUSY;
 
         pci_read_config_dword(pdev, pos + PCI_PRI_MAX_REQ, &max_requests);
         reqs = min(max_requests, reqs);
+        pdev->pri_reqs_alloc = reqs;
         pci_write_config_dword(pdev, pos + PCI_PRI_ALLOC_REQ, reqs);
 
-        control |= PCI_PRI_CTRL_ENABLE;
+        control = PCI_PRI_CTRL_ENABLE;
         pci_write_config_word(pdev, pos + PCI_PRI_CTRL, control);
 
+        pdev->pri_enabled = 1;
+
         return 0;
 }
 EXPORT_SYMBOL_GPL(pci_enable_pri);
@@ -185,6 +189,9 @@ void pci_disable_pri(struct pci_dev *pdev)
         u16 control;
         int pos;
 
+        if (WARN_ON(!pdev->pri_enabled))
+                return;
+
         pos = pci_find_ext_capability(pdev, PCI_EXT_CAP_ID_PRI);
         if (!pos)
                 return;
@@ -192,10 +199,34 @@ void pci_disable_pri(struct pci_dev *pdev)
         pci_read_config_word(pdev, pos + PCI_PRI_CTRL, &control);
         control &= ~PCI_PRI_CTRL_ENABLE;
         pci_write_config_word(pdev, pos + PCI_PRI_CTRL, control);
+
+        pdev->pri_enabled = 0;
 }
 EXPORT_SYMBOL_GPL(pci_disable_pri);
 
 /**
+ * pci_restore_pri_state - Restore PRI
+ * @pdev: PCI device structure
+ */
+void pci_restore_pri_state(struct pci_dev *pdev)
+{
+        u16 control = PCI_PRI_CTRL_ENABLE;
+        u32 reqs = pdev->pri_reqs_alloc;
+        int pos;
+
+        if (!pdev->pri_enabled)
+                return;
+
+        pos = pci_find_ext_capability(pdev, PCI_EXT_CAP_ID_PRI);
+        if (!pos)
+                return;
+
+        pci_write_config_dword(pdev, pos + PCI_PRI_ALLOC_REQ, reqs);
+        pci_write_config_word(pdev, pos + PCI_PRI_CTRL, control);
+}
+EXPORT_SYMBOL_GPL(pci_restore_pri_state);
+
+/**
  * pci_reset_pri - Resets device's PRI state
  * @pdev: PCI device structure
  *
@@ -207,16 +238,14 @@ int pci_reset_pri(struct pci_dev *pdev)
         u16 control;
         int pos;
 
+        if (WARN_ON(pdev->pri_enabled))
+                return -EBUSY;
+
         pos = pci_find_ext_capability(pdev, PCI_EXT_CAP_ID_PRI);
         if (!pos)
                 return -EINVAL;
 
-        pci_read_config_word(pdev, pos + PCI_PRI_CTRL, &control);
-        if (control & PCI_PRI_CTRL_ENABLE)
-                return -EBUSY;
-
-        control |= PCI_PRI_CTRL_RESET;
-
+        control = PCI_PRI_CTRL_RESET;
         pci_write_config_word(pdev, pos + PCI_PRI_CTRL, control);
 
         return 0;
@@ -239,16 +268,14 @@ int pci_enable_pasid(struct pci_dev *pdev, int features)
         u16 control, supported;
         int pos;
 
+        if (WARN_ON(pdev->pasid_enabled))
+                return -EBUSY;
+
         pos = pci_find_ext_capability(pdev, PCI_EXT_CAP_ID_PASID);
         if (!pos)
                 return -EINVAL;
 
-        pci_read_config_word(pdev, pos + PCI_PASID_CTRL, &control);
         pci_read_config_word(pdev, pos + PCI_PASID_CAP, &supported);
-
-        if (control & PCI_PASID_CTRL_ENABLE)
-                return -EINVAL;
-
         supported &= PCI_PASID_CAP_EXEC | PCI_PASID_CAP_PRIV;
 
         /* User wants to enable anything unsupported? */
@@ -256,9 +283,12 @@ int pci_enable_pasid(struct pci_dev *pdev, int features)
                 return -EINVAL;
 
         control = PCI_PASID_CTRL_ENABLE | features;
+        pdev->pasid_features = features;
 
         pci_write_config_word(pdev, pos + PCI_PASID_CTRL, control);
 
+        pdev->pasid_enabled = 1;
+
         return 0;
 }
 EXPORT_SYMBOL_GPL(pci_enable_pasid);
@@ -266,22 +296,47 @@ EXPORT_SYMBOL_GPL(pci_enable_pasid);
 /**
  * pci_disable_pasid - Disable the PASID capability
  * @pdev: PCI device structure
- *
  */
 void pci_disable_pasid(struct pci_dev *pdev)
 {
         u16 control = 0;
         int pos;
 
+        if (WARN_ON(!pdev->pasid_enabled))
+                return;
+
         pos = pci_find_ext_capability(pdev, PCI_EXT_CAP_ID_PASID);
         if (!pos)
                 return;
 
         pci_write_config_word(pdev, pos + PCI_PASID_CTRL, control);
+
+        pdev->pasid_enabled = 0;
 }
 EXPORT_SYMBOL_GPL(pci_disable_pasid);
 
 /**
+ * pci_restore_pasid_state - Restore PASID capabilities
+ * @pdev: PCI device structure
+ */
+void pci_restore_pasid_state(struct pci_dev *pdev)
+{
+        u16 control;
+        int pos;
+
+        if (!pdev->pasid_enabled)
+                return;
+
+        pos = pci_find_ext_capability(pdev, PCI_EXT_CAP_ID_PASID);
+        if (!pos)
+                return;
+
+        control = PCI_PASID_CTRL_ENABLE | pdev->pasid_features;
+        pci_write_config_word(pdev, pos + PCI_PASID_CTRL, control);
+}
+EXPORT_SYMBOL_GPL(pci_restore_pasid_state);
+
+/**
  * pci_pasid_features - Check which PASID features are supported
  * @pdev: PCI device structure
  *