xe_assert(xe, !xe_pm_runtime_suspended(xe));
 }
 
-bool xe_device_mem_access_get_if_ongoing(struct xe_device *xe)
-{
-       bool active;
-
-       if (xe_pm_read_callback_task(xe) == current)
-               return true;
-
-       active = xe_pm_runtime_get_if_active(xe);
-       if (active) {
-               int ref = atomic_inc_return(&xe->mem_access.ref);
-
-               xe_assert(xe, ref != S32_MAX);
-       }
-
-       return active;
-}
-
 void xe_device_mem_access_get(struct xe_device *xe)
 {
        int ref;
 
 }
 
 void xe_device_mem_access_get(struct xe_device *xe);
-bool xe_device_mem_access_get_if_ongoing(struct xe_device *xe);
 void xe_device_mem_access_put(struct xe_device *xe);
 
 void xe_device_assert_mem_access(struct xe_device *xe);
 
        bool ongoing;
        int len;
 
-       ongoing = xe_device_mem_access_get_if_ongoing(ct_to_xe(ct));
+       ongoing = xe_pm_runtime_get_if_active(ct_to_xe(ct));
        if (!ongoing && xe_pm_read_callback_task(ct_to_xe(ct)) == NULL)
                return;
 
        spin_unlock(&ct->fast_lock);
 
        if (ongoing)
-               xe_device_mem_access_put(xe);
+               xe_pm_runtime_put(xe);
 }
 
 /* Returns less than zero on error, 0 on done, 1 on more available */
         * responses, if the worker here is blocked on those callbacks
         * completing, creating a deadlock.
         */
-       ongoing = xe_device_mem_access_get_if_ongoing(ct_to_xe(ct));
+       ongoing = xe_pm_runtime_get_if_active(ct_to_xe(ct));
        if (!ongoing && xe_pm_read_callback_task(ct_to_xe(ct)) == NULL)
                return;
 
        } while (ret == 1);
 
        if (ongoing)
-               xe_device_mem_access_put(ct_to_xe(ct));
+               xe_pm_runtime_put(ct_to_xe(ct));
 }
 
 static void guc_ctb_snapshot_capture(struct xe_device *xe, struct guc_ctb *ctb,