#include "xe_pm.h"
 #include "xe_step.h"
 
+enum toggle_d3cold {
+       D3COLD_DISABLE,
+       D3COLD_ENABLE,
+};
+
 struct xe_subplatform_desc {
        enum xe_subplatform subplatform;
        const char *name;
        return 0;
 }
 
+static void d3cold_toggle(struct pci_dev *pdev, enum toggle_d3cold toggle)
+{
+       struct xe_device *xe = pdev_to_xe_device(pdev);
+       struct pci_dev *root_pdev;
+
+       if (!xe->d3cold.capable)
+               return;
+
+       root_pdev = pcie_find_root_port(pdev);
+       if (!root_pdev)
+               return;
+
+       switch (toggle) {
+       case D3COLD_DISABLE:
+               pci_d3cold_disable(root_pdev);
+               break;
+       case D3COLD_ENABLE:
+               pci_d3cold_enable(root_pdev);
+               break;
+       }
+}
+
 static int xe_pci_runtime_suspend(struct device *dev)
 {
        struct pci_dev *pdev = to_pci_dev(dev);
                pci_ignore_hotplug(pdev);
                pci_set_power_state(pdev, PCI_D3cold);
        } else {
+               d3cold_toggle(pdev, D3COLD_DISABLE);
                pci_set_power_state(pdev, PCI_D3hot);
        }
 
                        return err;
 
                pci_set_master(pdev);
+       } else {
+               d3cold_toggle(pdev, D3COLD_ENABLE);
        }
 
        return xe_pm_runtime_resume(xe);
        if (!xe->d3cold.capable) {
                xe->d3cold.allowed = false;
        } else {
+               xe_pm_d3cold_allowed_toggle(xe);
+
                /*
                 * TODO: d3cold should be allowed (true) if
                 * (IS_DGFX(xe) && !xe_device_mem_access_ongoing(xe))
                 * but maybe include some other conditions. So, before
                 * we can re-enable the D3cold, we need to:
                 * 1. rewrite the VRAM save / restore to avoid buffer object locks
-                * 2. block D3cold if we have a big amount of device memory in use
-                *    in order to reduce the latency.
-                * 3. at resume, detect if we really lost power and avoid memory
+                * 2. at resume, detect if we really lost power and avoid memory
                 *    restoration if we were only up to d3cold
                 */
                xe->d3cold.allowed = false;
 
 
        return 0;
 }
+
+void xe_pm_d3cold_allowed_toggle(struct xe_device *xe)
+{
+       struct ttm_resource_manager *man;
+       u32 total_vram_used_mb = 0;
+       u64 vram_used;
+       int i;
+
+       for (i = XE_PL_VRAM0; i <= XE_PL_VRAM1; ++i) {
+               man = ttm_manager_type(&xe->ttm, i);
+               if (man) {
+                       vram_used = ttm_resource_manager_usage(man);
+                       total_vram_used_mb += DIV_ROUND_UP_ULL(vram_used, 1024 * 1024);
+               }
+       }
+
+       mutex_lock(&xe->d3cold.lock);
+
+       if (total_vram_used_mb < xe->d3cold.vram_threshold)
+               xe->d3cold.allowed = true;
+       else
+               xe->d3cold.allowed = false;
+
+       mutex_unlock(&xe->d3cold.lock);
+}