[RFC 21/26] plugins/amdgpu: Enable restore with only the drm node open

It is currently not possible to restore a process which did not have a KFD
device open.

We can work around that by assuming the restore is being done on the same
system with the same GPU id. This can fail over reboots when the GPU id
changes, but that also could be worked around with perhaps something like
a "--ignore-gpu-id" command line switch.

In the long term we probably need to think of adding the relevant GPU
discovery APIs to the amdgpu kernel driver, which would enable this to
work more smoothly.

Signed-off-by: Tvrtko Ursulin <[email protected]>
---
 plugins/amdgpu/amdgpu_plugin.c          | 62 +++++++++++++++++++++----
 plugins/amdgpu/amdgpu_plugin_drm.c      |  5 +-
 plugins/amdgpu/amdgpu_plugin_topology.c |  2 +
 plugins/amdgpu/amdgpu_plugin_topology.h |  2 +
 4 files changed, 62 insertions(+), 9 deletions(-)

diff --git a/plugins/amdgpu/amdgpu_plugin.c b/plugins/amdgpu/amdgpu_plugin.c
index 5ab3420c1371..693df0ea7e39 100644
--- a/plugins/amdgpu/amdgpu_plugin.c
+++ b/plugins/amdgpu/amdgpu_plugin.c
@@ -1844,11 +1844,11 @@ static int amdgpu_plugin_restore_drm_file(int id, bool *retry_needed)
 	size_t img_size;
 	int fd, ret;
 
-	/* This is restorer plugin for renderD nodes. Criu doesn't guarantee that they will
-	 * be called before the plugin is called for kfd file descriptor.
-	 * TODO: Currently, this code will only work if this function is called for /dev/kfd
-	 * first as we assume restore_maps is already filled. Need to fix this later.
+	/* This is restorer plugin for renderD nodes. Criu doesn't guarantee
+	 * that they will be called before the plugin is called for kfd file
+	 * descriptor.
 	 */
+
 	snprintf(img_path, sizeof(img_path), IMG_DRM_FILE, id);
 	ret = load_img(img_path, &buf, &img_size);
 	if (ret < 0) {
@@ -1876,14 +1876,60 @@ static int amdgpu_plugin_restore_drm_file(int id, bool *retry_needed)
 
 	pr_info("render node gpu_id = 0x%04x\n", rd->gpu_id);
 
-	target_gpu_id = maps_get_dest_gpu(&restore_maps, rd->gpu_id);
-	if (!target_gpu_id) {
-		fd = -ENODEV;
-		goto fail;
+	if (fd_next == -1) {
+		ret = find_unused_fd_pid(getpid());
+		if (ret < 0) {
+			pr_err("Failed to find unused fd (fd:%d)\n", ret);
+			goto fail;
+		}
+		fd_next = ret;
+	}
+
+	if (!dest_topology.parsed) {
+		pr_info("Parsing local topology for render node restore\n");
+		ret = topology_parse(&dest_topology, "Local");
+		if (ret) {
+			pr_err("Failed to parse local system topology %d\n",
+			       ret);
+			goto fail;
+		}
+	}
+
+	if (restore_maps.mapped_cnt) {
+		target_gpu_id = maps_get_dest_gpu(&restore_maps, rd->gpu_id);
+		if (!target_gpu_id) {
+			pr_perror("Unable to map gpu_id 0x%04x\n", rd->gpu_id);
+			fd = -ENODEV;
+			goto fail;
+		}
+	} else {
+		unsigned int num_gpus = 0;
+
+		pr_info("Assuming same system with a single gpu_id 0x%04x\n",
+			rd->gpu_id);
+		list_for_each_entry(tp_node, &dest_topology.nodes,
+				    listm_system) {
+			if (NODE_IS_GPU(tp_node)) {
+				num_gpus++;
+				target_gpu_id = tp_node->gpu_id;
+			}
+		}
+		if (num_gpus != 1) {
+			pr_perror("Unexpectedly found %u GPUs\n", num_gpus);
+			fd = -EINVAL;
+			goto fail;
+		} else if (target_gpu_id != rd->gpu_id) {
+			pr_perror("Unexpectedly found gpu_id 0x%04x (expected 0x%04x)\n",
+				  target_gpu_id, rd->gpu_id);
+			fd = -EINVAL;
+			goto fail;
+		}
 	}
 
 	tp_node = sys_get_node_by_gpu_id(&dest_topology, target_gpu_id);
 	if (!tp_node) {
+		pr_perror("Unable to find target gpu_id=0x%04x\n",
+			  target_gpu_id);
 		fd = -ENODEV;
 		goto fail;
 	}
diff --git a/plugins/amdgpu/amdgpu_plugin_drm.c b/plugins/amdgpu/amdgpu_plugin_drm.c
index 99a6804f51eb..30a0626b7b5f 100644
--- a/plugins/amdgpu/amdgpu_plugin_drm.c
+++ b/plugins/amdgpu/amdgpu_plugin_drm.c
@@ -428,7 +428,10 @@ int amdgpu_plugin_drm_dump_file(int fd, int id, struct stat *drm)
 	}
 
 	/* Get the GPU_ID of the DRM device */
-	rd->gpu_id = maps_get_dest_gpu(&checkpoint_maps, tp_node->gpu_id);
+	if (checkpoint_maps.mapped_cnt)
+		rd->gpu_id = maps_get_dest_gpu(&checkpoint_maps, tp_node->gpu_id);
+	else
+		rd->gpu_id = tp_node->gpu_id; /* Render node only */
 	if (!rd->gpu_id) {
 		pr_err("Failed to find valid gpu_id for the device = %d\n", tp_node->gpu_id);
 		return -ENODEV;
diff --git a/plugins/amdgpu/amdgpu_plugin_topology.c b/plugins/amdgpu/amdgpu_plugin_topology.c
index 730f2e028430..b26c12e4c7d4 100644
--- a/plugins/amdgpu/amdgpu_plugin_topology.c
+++ b/plugins/amdgpu/amdgpu_plugin_topology.c
@@ -307,6 +307,7 @@ void maps_init(struct device_maps *maps)
 	INIT_LIST_HEAD(&maps->gpu_maps);
 	maps->tail_cpu = 0;
 	maps->tail_gpu = 0;
+	maps->mapped_cnt = 0;
 }
 
 void maps_free(struct device_maps *maps)
@@ -1220,6 +1221,7 @@ static bool map_devices(struct tp_system *src_sys, struct tp_system *dest_sys, s
 
 			if (map_devices(src_sys, dest_sys, src_nodes, dest_nodes, maps)) {
 				pr_debug("Matched nodes 0x%04X and after\n", dest_node->gpu_id);
+				maps->mapped_cnt++;
 				return true;
 			} else {
 				/* We could not map remaining nodes in the list. Add dest node back
diff --git a/plugins/amdgpu/amdgpu_plugin_topology.h b/plugins/amdgpu/amdgpu_plugin_topology.h
index e19f8e7ce9af..ab850d6bdce6 100644
--- a/plugins/amdgpu/amdgpu_plugin_topology.h
+++ b/plugins/amdgpu/amdgpu_plugin_topology.h
@@ -92,6 +92,8 @@ struct id_map {
 };
 
 struct device_maps {
+	unsigned int mapped_cnt;
+
 	struct list_head cpu_maps; /* CPUs are mapped using node_id */
 	struct list_head gpu_maps;
 
-- 
2.52.0