#define ISPSSPM0_IUNIT_POWER_ON                0x0
 #define ISPSSPM0_IUNIT_POWER_OFF       0x3
 
-static int isp_probe(struct pci_dev *dev, const struct pci_device_id *id)
+static int isp_set_power(struct pci_dev *dev, bool enable)
 {
        unsigned long timeout;
-       u32 val;
-
-       pci_write_config_dword(dev, PCI_INTERRUPT_CTRL, 0);
-
-       /*
-        * MRFLD IUNIT DPHY is located in an always-power-on island
-        * MRFLD HW design need all CSI ports are disabled before
-        * powering down the IUNIT.
-        */
-       pci_read_config_dword(dev, PCI_CSI_CONTROL, &val);
-       val |= PCI_CSI_CONTROL_PORTS_OFF_MASK;
-       pci_write_config_dword(dev, PCI_CSI_CONTROL, val);
+       u32 val = enable ? ISPSSPM0_IUNIT_POWER_ON :
+               ISPSSPM0_IUNIT_POWER_OFF;
 
-       /* Write 0x3 to ISPSSPM0 bit[1:0] to power off the IUNIT */
+       /* Write to ISPSSPM0 bit[1:0] to power on/off the IUNIT */
        iosf_mbi_modify(BT_MBI_UNIT_PMC, MBI_REG_READ, ISPSSPM0,
-                       ISPSSPM0_IUNIT_POWER_OFF, ISPSSPM0_ISPSSC_MASK);
+                       val, ISPSSPM0_ISPSSC_MASK);
 
        /*
         * There should be no IUNIT access while power-down is
         * in progress HW sighting: 4567865
         * Wait up to 50 ms for the IUNIT to shut down.
+        * And we do the same for power on.
         */
        timeout = jiffies + msecs_to_jiffies(50);
        while (1) {
-               /* Wait until ISPSSPM0 bit[25:24] shows 0x3 */
-               iosf_mbi_read(BT_MBI_UNIT_PMC, MBI_REG_READ, ISPSSPM0, &val);
-               val = (val & ISPSSPM0_ISPSSS_MASK) >> ISPSSPM0_ISPSSS_OFFSET;
-               if (val == ISPSSPM0_IUNIT_POWER_OFF)
+               u32 tmp;
+
+               /* Wait until ISPSSPM0 bit[25:24] shows the right value */
+               iosf_mbi_read(BT_MBI_UNIT_PMC, MBI_REG_READ, ISPSSPM0, &tmp);
+               tmp = (tmp & ISPSSPM0_ISPSSS_MASK) >> ISPSSPM0_ISPSSS_OFFSET;
+               if (tmp == val)
                        break;
 
                if (time_after(jiffies, timeout)) {
-                       dev_err(&dev->dev, "IUNIT power-off timeout.\n");
+                       dev_err(&dev->dev, "IUNIT power-%s timeout.\n",
+                               enable ? "on" : "off");
                        return -EBUSY;
                }
                usleep_range(1000, 2000);
        }
 
+       return 0;
+}
+
+static int isp_probe(struct pci_dev *dev, const struct pci_device_id *id)
+{
        pm_runtime_allow(&dev->dev);
        pm_runtime_put_sync_suspend(&dev->dev);
 
 
 static int isp_pci_suspend(struct device *dev)
 {
+       struct pci_dev *pdev = to_pci_dev(dev);
+       u32 val;
+
+       pci_write_config_dword(pdev, PCI_INTERRUPT_CTRL, 0);
+
+       /*
+        * MRFLD IUNIT DPHY is located in an always-power-on island
+        * MRFLD HW design need all CSI ports are disabled before
+        * powering down the IUNIT.
+        */
+       pci_read_config_dword(pdev, PCI_CSI_CONTROL, &val);
+       val |= PCI_CSI_CONTROL_PORTS_OFF_MASK;
+       pci_write_config_dword(pdev, PCI_CSI_CONTROL, val);
+
+       /*
+        * We lose config space access when punit power gates
+        * the ISP. Can't use pci_set_power_state() because
+        * pmcsr won't actually change when we write to it.
+        */
+       pci_save_state(pdev);
+       pdev->current_state = PCI_D3cold;
+       isp_set_power(pdev, false);
+
        return 0;
 }
 
 static int isp_pci_resume(struct device *dev)
 {
+       struct pci_dev *pdev = to_pci_dev(dev);
+
+       isp_set_power(pdev, true);
+       pdev->current_state = PCI_D0;
+       pci_restore_state(pdev);
+
        return 0;
 }