PCI: Enable ATS at the device state restore

During S3 or S4 resume or PCI reset, ATS regs aren't restored correctly.
This patch enables ATS at the device state restore if PCI device has ATS
capability.

Signed-off-by: Xudong Hao <xudong.hao@intel.com>
Signed-off-by: Xiantao Zhang <xiantao.zhang@intel.com>
Signed-off-by: Jesse Barnes <jbarnes@virtuousgeek.org>
diff --git a/drivers/pci/ats.c b/drivers/pci/ats.c
index e11ebaf..a4a1b36 100644
--- a/drivers/pci/ats.c
+++ b/drivers/pci/ats.c
@@ -128,6 +128,23 @@
 }
 EXPORT_SYMBOL_GPL(pci_disable_ats);
 
+void pci_restore_ats_state(struct pci_dev *dev)
+{
+	u16 ctrl;
+
+	if (!pci_ats_enabled(dev))
+		return;
+	if (!pci_find_ext_capability(dev, PCI_EXT_CAP_ID_ATS))
+		BUG();
+
+	ctrl = PCI_ATS_CTRL_ENABLE;
+	if (!dev->is_virtfn)
+		ctrl |= PCI_ATS_CTRL_STU(dev->ats->stu - PCI_ATS_MIN_STU);
+
+	pci_write_config_word(dev, dev->ats->pos + PCI_ATS_CTRL, ctrl);
+}
+EXPORT_SYMBOL_GPL(pci_restore_ats_state);
+
 /**
  * pci_ats_queue_depth - query the ATS Invalidate Queue Depth
  * @dev: the PCI device
diff --git a/drivers/pci/pci.c b/drivers/pci/pci.c
index 54343aa..97fff78 100644
--- a/drivers/pci/pci.c
+++ b/drivers/pci/pci.c
@@ -965,6 +965,7 @@
 
 	/* PCI Express register must be restored first */
 	pci_restore_pcie_state(dev);
+	pci_restore_ats_state(dev);
 
 	/*
 	 * The Base Address register should be programmed before the command
diff --git a/drivers/pci/pci.h b/drivers/pci/pci.h
index 3b6e4ed..1009a5e 100644
--- a/drivers/pci/pci.h
+++ b/drivers/pci/pci.h
@@ -251,6 +251,14 @@
 	u8 __iomem *mstate;	/* VF Migration State Array */
 };
 
+#ifdef CONFIG_PCI_ATS
+extern void pci_restore_ats_state(struct pci_dev *dev);
+#else
+static inline void pci_restore_ats_state(struct pci_dev *dev)
+{
+}
+#endif /* CONFIG_PCI_ATS */
+
 #ifdef CONFIG_PCI_IOV
 extern int pci_iov_init(struct pci_dev *dev);
 extern void pci_iov_release(struct pci_dev *dev);