RFC: AMD ROCm support with plugin

Documentation/Makefile

@@ -13,6 +13,7 @@ endif
 FOOTER		:= footer.txt
 SRC1		+= crit.txt
 SRC1		+= compel.txt
+SRC1		+= amdgpu_plugin.txt
 SRC8		+= criu.txt
 SRC		:= $(SRC1) $(SRC8)
 XMLS		:= $(patsubst %.txt,%.xml,$(SRC))

Documentation/amdgpu_plugin.txt

@@ -0,0 +1,99 @@
+ROCM Support(1)
+===============
+
+NAME
+----
+amdgpu_plugin - A plugin extention to CRIU to support checkpoint/restore in
+userspace for AMD GPUs.
+
+
+CURRENT SUPPORT
+---------------
+Single and Multi GPU systems (Gfx9)
+Checkpoint / Restore on same system
+Checkpoint / Restore inside a docker container
+Pytorch
+
+DESCRIPTION
+-----------
+Though *criu* is a great tool for checkpointing and restoring running
+applications, it has certain limitations such as it cannot handle
+applications that have device files open. In order to support *ROCm* based
+workloads with *criu* we need to augment criu's core functionality with a
+plugin based extention mechanism. *amdgpu_plugin* provides the necessary support
+to criu to allow Checkpoint / Restore with ROCm.
+
+
+Dependencies
+~~~~~~~~~~~~~~
+*amdkfd support*::
+    In order to snapshot the *VRAM* and other *GPU* device states, we require
+    an updated version of amdkfd(amdgpu) driver. The kernel patches are under
+    review currently.
+
+*criu 3.15*::
+    This work is rebased on latest criu release available at this time.
+
+
+OPTIONS
+-------
+Optional parameters can be passed in as environment variables before
+executing criu command.
+
+*KFD_FW_VER_CHECK*::
+    Enable or disable firmware version check.
+    If enabled, firmware version on restored gpu needs to be greater than or
+    equal firmware version on checkpointed GPU. Default:Enabled
+
+    E.g:
+    KFD_FW_VER_CHECK=0
+
+*KFD_SDMA_FW_VER_CHECK*::
+    Enable or disable SDMA firmware version check.
+    If enabled, SDMA firmware version on restored gpu needs to be greater than or
+    equal firmware version on checkpointed GPU. Default:Enabled
+
+    E.g:
+    KFD_SDMA_FW_VER_CHECK=0
+
+*KFD_CACHES_COUNT_CHECK*::
+    Enable or disable caches count check. If enabled, the caches count on
+    restored GPU needs to be greater than or equal caches count on checkpointed
+    GPU. Default:Enabled
+
+    E.g:
+    KFD_CACHES_COUNT_CHECK=0
+
+*KFD_NUM_GWS_CHECK*::
+    Enable or disable num_gws check. If enabled, the num_gws on
+    restored GPU needs to be greater than or equal num_gws on checkpointed
+    GPU. Default:Enabled
+
+    E.g:
+    KFD_NUM_GWS_CHECK=0
+
+*KFD_VRAM_SIZE_CHECK*::
+    Enable or disable VRAM size check. If enabled, the VRAM size on
+    restored GPU needs to be greater than or equal VRAM size on checkpointed
+    GPU. Default:Enabled
+
+    E.g:
+    KFD_VRAM_SIZE_CHECK=0
+
+*KFD_NUMA_CHECK*::
+    Enable or disable NUMA CPU region check. If enabled, the plugin will restore
+    GPUs that belong to one CPU NUMA region to the same CPU NUMA region.
+    Default:Enabled
+
+    E.g:
+    KFD_IGNORE_NUMA=1
+
+
+AUTHOR
+------
+The AMDKFD team.
+
+
+COPYRIGHT
+---------
+Copyright \(C) 2020-2021, Advanced Micro Devices, Inc. (AMD)

Makefile

@@ -147,7 +147,7 @@ HOSTCFLAGS		+= $(WARNINGS) $(DEFINES) -iquote include/
 export AFLAGS CFLAGS USERCLFAGS HOSTCFLAGS
 
 # Default target
-all: flog criu lib crit
+all: flog criu lib crit amdgpu_plugin
 .PHONY: all
 
 #
@@ -290,15 +290,19 @@ clean mrproper:
 	$(Q) $(MAKE) $(build)=crit $@
 .PHONY: clean mrproper
 
+clean-amdgpu_plugin:
+	$(Q) $(MAKE) -C plugins/amdgpu clean
+.PHONY: clean-amdgpu_plugin
+
 clean-top:
 	$(Q) $(MAKE) -C Documentation clean
 	$(Q) $(MAKE) $(build)=test/compel clean
 	$(Q) $(RM) .gitid
 .PHONY: clean-top
 
-clean: clean-top
+clean: clean-top clean-amdgpu_plugin
 
-mrproper-top: clean-top
+mrproper-top: clean-top clean-amdgpu_plugin
 	$(Q) $(RM) $(CONFIG_HEADER)
 	$(Q) $(RM) $(VERSION_HEADER)
 	$(Q) $(RM) $(COMPEL_VERSION_HEADER)
@@ -326,6 +330,10 @@ test: zdtm
 	$(Q) $(MAKE) -C test
 .PHONY: test
 
+amdgpu_plugin:
+	$(Q) $(MAKE) -C plugins/amdgpu all
+.PHONY: amdgpu_plugin
+
 #
 # Generating tar requires tag matched CRIU_VERSION.
 # If not found then simply use GIT's describe with
@@ -403,6 +411,7 @@ help:
 	@echo '      cscope          - Generate cscope database'
 	@echo '      test            - Run zdtm test-suite'
 	@echo '      gcov            - Make code coverage report'
+	@echo '      amdgpu_plugin   - Make AMD GPU plugin'
 .PHONY: help
 
 lint:

Makefile.install

@@ -7,6 +7,7 @@ MANDIR		?= $(PREFIX)/share/man
 INCLUDEDIR	?= $(PREFIX)/include
 LIBEXECDIR	?= $(PREFIX)/libexec
 RUNDIR		?= /run
+PLUGINDIR	?= /var/lib/criu
 
 #
 # For recent Debian/Ubuntu with multiarch support.
@@ -26,7 +27,7 @@ endif
 LIBDIR ?= $(PREFIX)/lib
 
 export PREFIX BINDIR SBINDIR MANDIR RUNDIR
-export LIBDIR INCLUDEDIR LIBEXECDIR
+export LIBDIR INCLUDEDIR LIBEXECDIR PLUGINDIR
 
 install-man:
 	$(Q) $(MAKE) -C Documentation install
@@ -40,12 +41,16 @@ install-criu: criu
 	$(Q) $(MAKE) $(build)=criu install
 .PHONY: install-criu
 
+install-amdgpu_plugin: amdgpu_plugin
+	$(Q) $(MAKE) -C plugins/amdgpu install
+.PHONY: install-amdgpu_plugin
+
 install-compel: $(compel-install-targets)
 	$(Q) $(MAKE) $(build)=compel install
 	$(Q) $(MAKE) $(build)=compel/plugins install
 .PHONY: install-compel
 
-install: install-man install-lib install-criu install-compel ;
+install: install-man install-lib install-criu install-compel install-amdgpu_plugin ;
 .PHONY: install
 
 uninstall:
@@ -54,4 +59,5 @@ uninstall:
 	$(Q) $(MAKE) $(build)=criu $@
 	$(Q) $(MAKE) $(build)=compel $@
 	$(Q) $(MAKE) $(build)=compel/plugins $@
+	$(Q) $(MAKE) -C plugins/amdgpu uninstall
 .PHONY: uninstall

criu/cr-restore.c

@@ -2451,6 +2451,21 @@ static int restore_root_task(struct pstree_item *init)
 		pr_err("Unable to flush breakpoints\n");
 
 	finalize_restore();
+	/*
+	 * Some external devices such as GPUs might need a very late
+	 * trigger to kick-off some events, memory notifiers and for
+	 * restarting the previously restored queues during criu restore
+	 * stage. This is needed since criu pie code may shuffle VMAs
+	 * around so things such as registering MMU notifiers (for GPU
+	 * mapped memory) could be done sanely once the pie code hands
+	 * over the control to master process.
+	 */
+	for_each_pstree_item(item) {
+		pr_info("Run late stage hook from criu master for external devices\n");
+		ret = run_plugins(RESUME_DEVICES_LATE, item->pid->real);
+		if (ret < 0)
+			pr_err("criu: restore late stage hook for external plugin failed\n");
+	}
 
 	ret = run_scripts(ACT_PRE_RESUME);
 	if (ret)

criu/file-ids.c

@@ -77,11 +77,18 @@ int fd_id_generate_special(struct fd_parms *p, u32 *id)
 {
 	if (p) {
 		struct fd_id *fi;
+		struct stat st_kfd;
 
 		fi = fd_id_cache_lookup(p);
 		if (fi) {
-			*id = fi->id;
-			return 0;
+			if (stat("/dev/kfd", &st_kfd) == -1) {
+				*id = fi->id;
+				return 0;
+			} else {
+				/* Don't cache the id */
+				*id = fd_tree.subid++;
+				return 1;
+			}
 		}
 	}
 

criu/files-reg.c

@@ -2314,6 +2314,23 @@ static int open_filemap(int pid, struct vma_area *vma)
 	BUG_ON((vma->vmfd == NULL) || !vma->e->has_fdflags);
 	flags = vma->e->fdflags;
 
+	/* update the new device file page offsets and file paths set during restore */
+	if (vma->e->status & VMA_UNSUPP) {
+		uint64_t new_pgoff;
+		char new_path[PATH_MAX];
+		int ret;
+
+		struct reg_file_info *rfi = container_of(vma->vmfd, struct reg_file_info, d);
+		ret = run_plugins(UPDATE_VMA_MAP, rfi->rfe->name, new_path, vma->e->start,
+					vma->e->pgoff, &new_pgoff);
+		if (ret == 1) {
+			pr_info("New mmap %#016"PRIx64"->%#016"PRIx64" path %s\n", vma->e->pgoff,
+					new_pgoff, new_path);
+			vma->e->pgoff = new_pgoff;
+			strcpy(rfi->path, new_path);
+		}
+	}
+
 	if (ctx.flags != flags || ctx.desc != vma->vmfd) {
 		if (vma->e->status & VMA_AREA_MEMFD)
 			ret = memfd_open(vma->vmfd, &flags);
@@ -2331,6 +2348,7 @@ static int open_filemap(int pid, struct vma_area *vma)
 
 	ctx.vma = vma;
 	vma->e->fd = ctx.fd;
+
 	return 0;
 }
 

criu/include/criu-plugin.h

@@ -22,6 +22,7 @@
 
 #include <limits.h>
 #include <stdbool.h>
+#include <stdint.h>
 
 #define CRIU_PLUGIN_GEN_VERSION(a,b,c)	(((a) << 16) + ((b) << 8) + (c))
 #define CRIU_PLUGIN_VERSION_MAJOR	0
@@ -50,6 +51,10 @@ enum {
 
 	CR_PLUGIN_HOOK__DUMP_EXT_LINK		= 6,
 
+	CR_PLUGIN_HOOK__UPDATE_VMA_MAP		= 7,
+
+	CR_PLUGIN_HOOK__RESUME_DEVICES_LATE	= 8,
+
 	CR_PLUGIN_HOOK__MAX
 };
 
@@ -63,6 +68,9 @@ DECLARE_PLUGIN_HOOK_ARGS(CR_PLUGIN_HOOK__RESTORE_EXT_FILE, int id);
 DECLARE_PLUGIN_HOOK_ARGS(CR_PLUGIN_HOOK__DUMP_EXT_MOUNT, char *mountpoint, int id);
 DECLARE_PLUGIN_HOOK_ARGS(CR_PLUGIN_HOOK__RESTORE_EXT_MOUNT, int id, char *mountpoint, char *old_root, int *is_file);
 DECLARE_PLUGIN_HOOK_ARGS(CR_PLUGIN_HOOK__DUMP_EXT_LINK, int index, int type, char *kind);
+DECLARE_PLUGIN_HOOK_ARGS(CR_PLUGIN_HOOK__UPDATE_VMA_MAP, const char* old_path, char *new_path,
+                               const uint64_t addr, const uint64_t old_pgoff, uint64_t *new_pgoff);
+DECLARE_PLUGIN_HOOK_ARGS(CR_PLUGIN_HOOK__RESUME_DEVICES_LATE, int pid);
 
 enum {
 	CR_PLUGIN_STAGE__DUMP,
@@ -128,5 +136,9 @@ typedef int (cr_plugin_restore_file_t)(int id);
 typedef int (cr_plugin_dump_ext_mount_t)(char *mountpoint, int id);
 typedef int (cr_plugin_restore_ext_mount_t)(int id, char *mountpoint, char *old_root, int *is_file);
 typedef int (cr_plugin_dump_ext_link_t)(int index, int type, char *kind);
+typedef int (cr_plugin_update_vma_offset_t)(const char* old_path, char *new_path,
+					    const uint64_t addr, const uint64_t old_pgoff,
+					    uint64_t *new_pgoff);
+typedef int (cr_plugin_resume_devices_late_t)(int pid);
 
 #endif /* __CRIU_PLUGIN_H__ */

criu/include/proc_parse.h

@@ -8,6 +8,9 @@
 #define PROC_TASK_COMM_LEN	32
 #define PROC_TASK_COMM_LEN_FMT	"(%31s"
 
+#define DRM_FIRST_RENDER_NODE 128
+#define DRM_LAST_RENDER_NODE 255
+
 struct proc_pid_stat {
 	int			pid;
 	char			comm[PROC_TASK_COMM_LEN];

criu/plugin.c

@@ -53,6 +53,8 @@ static cr_plugin_desc_t *cr_gen_plugin_desc(void *h, char *path)
 	__assign_hook(DUMP_EXT_MOUNT,		"cr_plugin_dump_ext_mount");
 	__assign_hook(RESTORE_EXT_MOUNT,	"cr_plugin_restore_ext_mount");
 	__assign_hook(DUMP_EXT_LINK,		"cr_plugin_dump_ext_link");
+	__assign_hook(UPDATE_VMA_MAP,		"cr_plugin_update_vma_map");
+	__assign_hook(RESUME_DEVICES_LATE,	"cr_plugin_resume_devices_late");
 
 #undef __assign_hook