[PATCH v2] drm/amdkfd: save/restore MQD across hibernation when MQD is in VRAM

On gfx942/944 and gfx9.5.0, KFD compute-queue MQDs live in a pinned VRAM
BO (mqd_on_vram()). Pinned VRAM is skipped by TTM eviction and not saved
across S4 hibernation, so the MQD is garbage on resume and the first
submission faults. Save it to a system-RAM shadow at suspend and restore
it on resume, reusing the CRIU checkpoint_mqd/restore_mqd primitives.

v2:
 - Also snapshot the MQD when a process is already runtime-evicted at
   hibernation start. The qpd->evicted refcount early-return previously
   skipped the save loop, losing the MQD across the VRAM wipe (both the
   cpsch and nocpsch evict paths).

Signed-off-by: Shikang Fan <[email protected]>
---
 .../drm/amd/amdkfd/kfd_device_queue_manager.c | 120 +++++++++++++++++-
 drivers/gpu/drm/amd/amdkfd/kfd_priv.h         |   9 ++
 drivers/gpu/drm/amd/amdkfd/kfd_queue.c        |   2 +
 3 files changed, 129 insertions(+), 2 deletions(-)

diff --git a/drivers/gpu/drm/amd/amdkfd/kfd_device_queue_manager.c b/drivers/gpu/drm/amd/amdkfd/kfd_device_queue_manager.c
index 51ee9c39104b..d168686ee324 100644
--- a/drivers/gpu/drm/amd/amdkfd/kfd_device_queue_manager.c
+++ b/drivers/gpu/drm/amd/amdkfd/kfd_device_queue_manager.c
@@ -65,6 +65,9 @@ static int map_queues_cpsch(struct device_queue_manager *dqm);
 static void deallocate_sdma_queue(struct device_queue_manager *dqm,
 				struct queue *q);
 
+static int dqm_alloc_mqd_backup(struct device_queue_manager *dqm, struct queue *q);
+static void dqm_save_mqd_backup(struct device_queue_manager *dqm, struct queue *q);
+
 static inline void deallocate_hqd(struct device_queue_manager *dqm,
 				struct queue *q);
 static int allocate_hqd(struct device_queue_manager *dqm, struct queue *q);
@@ -802,6 +805,10 @@ static int create_queue_nocpsch(struct device_queue_manager *dqm,
 		mqd_mgr->init_mqd(mqd_mgr, &q->mqd, q->mqd_mem_obj,
 					&q->gart_mqd_addr, &q->properties);
 
+	retval = dqm_alloc_mqd_backup(dqm, q);
+	if (retval)
+		goto out_free_mqd;
+
 	if (q->properties.is_active) {
 		if (!dqm->sched_running) {
 			WARN_ONCE(1, "Load non-HWS mqd while stopped\n");
@@ -1263,8 +1270,16 @@ static int evict_process_queues_nocpsch(struct device_queue_manager *dqm,
 	int retval, ret = 0;
 
 	dqm_lock(dqm);
-	if (qpd->evicted++ > 0) /* already evicted, do nothing */
+	if (qpd->evicted++ > 0) { /* already evicted, do nothing */
+		/* A process already runtime-evicted when hibernation starts
+		 * skips the save loop below, so snapshot its VRAM MQDs here
+		 * before the S4 image is taken. dqm_save_mqd_backup() is a
+		 * no-op outside S4.
+		 */
+		list_for_each_entry(q, &qpd->queues_list, list)
+			dqm_save_mqd_backup(dqm, q);
 		goto out;
+	}
 
 	pdd = qpd_to_pdd(qpd);
 	pr_debug_ratelimited("Evicting process pid %d queues\n",
@@ -1297,6 +1312,8 @@ static int evict_process_queues_nocpsch(struct device_queue_manager *dqm,
 			 * maintain a consistent eviction state
 			 */
 			ret = retval;
+
+		dqm_save_mqd_backup(dqm, q);
 	}
 
 out:
@@ -1304,6 +1321,88 @@ static int evict_process_queues_nocpsch(struct device_queue_manager *dqm,
 	return ret;
 }
 
+/* MQD software-shadow save/restore across S4 hibernation, reusing the CRIU
+ * checkpoint_mqd/restore_mqd primitives.
+ */
+static int dqm_alloc_mqd_backup(struct device_queue_manager *dqm, struct queue *q)
+{
+	struct mqd_manager *mqd_mgr;
+	uint32_t mqd_size, ctl_stack_size = 0;
+
+	if (!mqd_on_vram(dqm->dev->adev))
+		return 0;
+	if (q->properties.type != KFD_QUEUE_TYPE_COMPUTE)
+		return 0;
+	if (!q->mqd)
+		return 0;
+
+	mqd_mgr = dqm->mqd_mgrs[get_mqd_type_from_queue_type(q->properties.type)];
+	if (!mqd_mgr->checkpoint_mqd || !mqd_mgr->restore_mqd)
+		return 0;
+
+	mqd_size = AMDGPU_MQD_SIZE_ALIGN(mqd_mgr->mqd_size) *
+		   NUM_XCC(dqm->dev->xcc_mask);
+	if (mqd_mgr->get_checkpoint_info)
+		mqd_mgr->get_checkpoint_info(mqd_mgr, q->mqd, &ctl_stack_size);
+
+	if (!q->mqd_backup) {
+		q->mqd_backup = kzalloc(mqd_size, GFP_KERNEL);
+		if (!q->mqd_backup)
+			return -ENOMEM;
+		q->mqd_backup_size = mqd_size;
+	}
+	if (ctl_stack_size && !q->ctl_stack_backup) {
+		q->ctl_stack_backup = kzalloc(ctl_stack_size, GFP_KERNEL);
+		if (!q->ctl_stack_backup) {
+			kfree(q->mqd_backup);
+			q->mqd_backup = NULL;
+			q->mqd_backup_size = 0;
+			return -ENOMEM;
+		}
+		q->ctl_stack_backup_size = ctl_stack_size;
+	}
+
+	return 0;
+}
+
+static void dqm_save_mqd_backup(struct device_queue_manager *dqm, struct queue *q)
+{
+	struct mqd_manager *mqd_mgr;
+
+	/* in_s4 is reliably set at evict/suspend time; skip normal runtime eviction
+	 * (pinned VRAM MQD stays intact then).
+	 */
+	if (!dqm->dev->adev->in_s4)
+		return;
+	if (!q->mqd_backup)
+		return;
+	if (!q->mqd || !q->mqd_mem_obj || !q->mqd_mem_obj->cpu_ptr)
+		return;
+
+	mqd_mgr = dqm->mqd_mgrs[get_mqd_type_from_queue_type(q->properties.type)];
+	mqd_mgr->checkpoint_mqd(mqd_mgr, q->mqd, q->mqd_backup, q->ctl_stack_backup);
+	q->mqd_backup_valid = true;
+}
+
+static void dqm_restore_mqd_backup(struct device_queue_manager *dqm, struct queue *q)
+{
+	struct mqd_manager *mqd_mgr;
+
+	if (!q->mqd_backup_valid || !q->mqd_backup)
+		return;
+	if (!q->mqd_mem_obj)
+		return;
+
+	mqd_mgr = dqm->mqd_mgrs[get_mqd_type_from_queue_type(q->properties.type)];
+	if (!mqd_mgr->restore_mqd)
+		return;
+
+	mqd_mgr->restore_mqd(mqd_mgr, &q->mqd, q->mqd_mem_obj, &q->gart_mqd_addr,
+			     &q->properties, q->mqd_backup, q->ctl_stack_backup,
+			     q->ctl_stack_backup_size);
+	q->mqd_backup_valid = false;
+}
+
 static int evict_process_queues_cpsch(struct device_queue_manager *dqm,
 				      struct qcm_process_device *qpd)
 {
@@ -1313,8 +1412,16 @@ static int evict_process_queues_cpsch(struct device_queue_manager *dqm,
 	int retval = 0;
 
 	dqm_lock(dqm);
-	if (qpd->evicted++ > 0) /* already evicted, do nothing */
+	if (qpd->evicted++ > 0) { /* already evicted, do nothing */
+		/* A process already runtime-evicted when hibernation starts
+		 * skips the save loop below, so snapshot its VRAM MQDs here
+		 * before the S4 image is taken. dqm_save_mqd_backup() is a
+		 * no-op outside S4.
+		 */
+		list_for_each_entry(q, &qpd->queues_list, list)
+			dqm_save_mqd_backup(dqm, q);
 		goto out;
+	}
 
 	pdd = qpd_to_pdd(qpd);
 
@@ -1350,6 +1457,8 @@ static int evict_process_queues_cpsch(struct device_queue_manager *dqm,
 				goto out;
 			}
 		}
+
+		dqm_save_mqd_backup(dqm, q);
 	}
 
 	if (!dqm->dev->kfd->shared_resources.enable_mes) {
@@ -1423,6 +1532,7 @@ static int restore_process_queues_nocpsch(struct device_queue_manager *dqm,
 
 		mqd_mgr = dqm->mqd_mgrs[get_mqd_type_from_queue_type(
 				q->properties.type)];
+		dqm_restore_mqd_backup(dqm, q);
 		q->properties.is_active = true;
 		increment_queue_count(dqm, qpd, q);
 
@@ -1486,6 +1596,7 @@ static int restore_process_queues_cpsch(struct device_queue_manager *dqm,
 		if (!QUEUE_IS_ACTIVE(q->properties))
 			continue;
 
+		dqm_restore_mqd_backup(dqm, q);
 		q->properties.is_active = true;
 		increment_queue_count(dqm, &pdd->qpd, q);
 
@@ -2161,6 +2272,10 @@ static int create_queue_cpsch(struct device_queue_manager *dqm, struct queue *q,
 		mqd_mgr->init_mqd(mqd_mgr, &q->mqd, q->mqd_mem_obj,
 					&q->gart_mqd_addr, &q->properties);
 
+	retval = dqm_alloc_mqd_backup(dqm, q);
+	if (retval)
+		goto out_free_mqd;
+
 	list_add(&q->list, &qpd->queues_list);
 	qpd->queue_count++;
 
@@ -2193,6 +2308,7 @@ static int create_queue_cpsch(struct device_queue_manager *dqm, struct queue *q,
 	list_del(&q->list);
 	if (q->properties.is_active)
 		decrement_queue_count(dqm, qpd, q);
+out_free_mqd:
 	mqd_mgr->free_mqd(mqd_mgr, q->mqd, q->mqd_mem_obj);
 	dqm_unlock(dqm);
 out_deallocate_doorbell:
diff --git a/drivers/gpu/drm/amd/amdkfd/kfd_priv.h b/drivers/gpu/drm/amd/amdkfd/kfd_priv.h
index 88191a4c1657..1f728f13813e 100644
--- a/drivers/gpu/drm/amd/amdkfd/kfd_priv.h
+++ b/drivers/gpu/drm/amd/amdkfd/kfd_priv.h
@@ -637,6 +637,15 @@ struct queue {
 	void *gang_ctx_cpu_ptr;
 
 	struct amdgpu_bo *wptr_bo_gart;
+
+	/* system-RAM shadow of a VRAM-resident MQD (+ control stack)
+	 * for hibernation snapshot/restore
+	 */
+	void *mqd_backup;
+	void *ctl_stack_backup;
+	uint32_t mqd_backup_size;
+	uint32_t ctl_stack_backup_size;
+	bool mqd_backup_valid;
 };
 
 enum KFD_MQD_TYPE {
diff --git a/drivers/gpu/drm/amd/amdkfd/kfd_queue.c b/drivers/gpu/drm/amd/amdkfd/kfd_queue.c
index 98a5512b701b..cbc90bf87792 100644
--- a/drivers/gpu/drm/amd/amdkfd/kfd_queue.c
+++ b/drivers/gpu/drm/amd/amdkfd/kfd_queue.c
@@ -83,6 +83,8 @@ int init_queue(struct queue **q, const struct queue_properties *properties)
 
 void uninit_queue(struct queue *q)
 {
+	kfree(q->mqd_backup);
+	kfree(q->ctl_stack_backup);
 	kfree(q);
 }
 
-- 
2.34.1