return 0;
        }
 
-       xe_device_mem_access_get(xe);
+       xe_pm_runtime_get(xe);
 
        for_each_tile(tile, xe, id) {
                /* For igfx run only for primary tile */
                ccs_test_run_tile(xe, tile, test);
        }
 
-       xe_device_mem_access_put(xe);
+       xe_pm_runtime_put(xe);
 
        return 0;
 }
                return 0;
        }
 
-       xe_device_mem_access_get(xe);
+       xe_pm_runtime_get(xe);
 
        for_each_tile(tile, xe, id)
                evict_test_run_tile(xe, tile, test);
 
-       xe_device_mem_access_put(xe);
+       xe_pm_runtime_put(xe);
 
        return 0;
 }
 
 #include "tests/xe_pci_test.h"
 
 #include "xe_pci.h"
+#include "xe_pm.h"
 
 static bool sanity_fence_failed(struct xe_device *xe, struct dma_fence *fence,
                                const char *str, struct kunit *test)
        struct xe_tile *tile;
        int id;
 
+       xe_pm_runtime_get(xe);
+
        for_each_tile(tile, xe, id) {
                struct xe_migrate *m = tile->migrate;
 
                kunit_info(test, "Testing tile id %d.\n", id);
                xe_vm_lock(m->q->vm, true);
-               xe_device_mem_access_get(xe);
                xe_migrate_sanity_test(m, test);
-               xe_device_mem_access_put(xe);
                xe_vm_unlock(m->q->vm);
        }
 
+       xe_pm_runtime_put(xe);
+
        return 0;
 }
 
 
 #include "tests/xe_pci_test.h"
 #include "tests/xe_test.h"
 
-#include "xe_pci.h"
+#include "xe_device.h"
 #include "xe_gt.h"
 #include "xe_mocs.h"
-#include "xe_device.h"
+#include "xe_pci.h"
+#include "xe_pm.h"
 
 struct live_mocs {
        struct xe_mocs_info table;
 
        struct kunit *test = xe_cur_kunit();
 
-       xe_device_mem_access_get(gt_to_xe(gt));
        ret = xe_force_wake_get(gt_to_fw(gt), XE_FW_GT);
        KUNIT_ASSERT_EQ_MSG(test, ret, 0, "Forcewake Failed.\n");
        mocs_dbg(>_to_xe(gt)->drm, "L3CC entries:%d\n", info->n_entries);
                                   XELP_LNCFCMOCS(i).addr);
        }
        xe_force_wake_put(gt_to_fw(gt), XE_FW_GT);
-       xe_device_mem_access_put(gt_to_xe(gt));
 }
 
 static void read_mocs_table(struct xe_gt *gt,
 
        struct kunit *test = xe_cur_kunit();
 
-       xe_device_mem_access_get(gt_to_xe(gt));
        ret = xe_force_wake_get(gt_to_fw(gt), XE_FW_GT);
        KUNIT_ASSERT_EQ_MSG(test, ret, 0, "Forcewake Failed.\n");
        mocs_dbg(>_to_xe(gt)->drm, "Global MOCS entries:%d\n", info->n_entries);
                                   XELP_GLOBAL_MOCS(i).addr);
        }
        xe_force_wake_put(gt_to_fw(gt), XE_FW_GT);
-       xe_device_mem_access_put(gt_to_xe(gt));
 }
 
 static int mocs_kernel_test_run_device(struct xe_device *xe)
        unsigned int flags;
        int id;
 
+       xe_pm_runtime_get(xe);
+
        for_each_gt(gt, xe, id) {
                flags = live_mocs_init(&mocs, gt);
                if (flags & HAS_GLOBAL_MOCS)
                if (flags & HAS_LNCF_MOCS)
                        read_l3cc_table(gt, &mocs.table);
        }
+
+       xe_pm_runtime_put(xe);
+
        return 0;
 }
 
        int id;
        struct kunit *test = xe_cur_kunit();
 
+       xe_pm_runtime_get(xe);
+
        for_each_gt(gt, xe, id) {
                flags = live_mocs_init(&mocs, gt);
                kunit_info(test, "mocs_reset_test before reset\n");
                if (flags & HAS_LNCF_MOCS)
                        read_l3cc_table(gt, &mocs.table);
        }
+
+       xe_pm_runtime_put(xe);
+
        return 0;
 }