diff --git a/compel/arch/aarch64/src/lib/include/uapi/asm/infect-types.h b/compel/arch/aarch64/src/lib/include/uapi/asm/infect-types.h
index 714881c57..eabf1f22f 100644
--- a/compel/arch/aarch64/src/lib/include/uapi/asm/infect-types.h
+++ b/compel/arch/aarch64/src/lib/include/uapi/asm/infect-types.h
@@ -5,7 +5,6 @@
 #include <signal.h>
 #include <sys/mman.h>
 #include <asm/ptrace.h>
-#include "common/page.h"
 
 #define SIGMAX			64
 #define SIGMAX_OLD		31
@@ -25,26 +24,6 @@ typedef struct user_fpsimd_state	user_fpregs_struct_t;
 
 #define user_regs_native(pregs)			true
 
-/*
- * Range for task size calculated from the following Linux kernel files:
- *   arch/arm64/include/asm/memory.h
- *   arch/arm64/Kconfig
- *
- * TODO: handle 32 bit tasks
- */
-#define TASK_SIZE_MIN (1UL << 39)
-#define TASK_SIZE_MAX (1UL << 48)
-
-static inline unsigned long task_size(void)
-{
-	unsigned long task_size;
-
-	for (task_size = TASK_SIZE_MIN; task_size < TASK_SIZE_MAX; task_size <<= 1)
-		if (munmap((void *)task_size, page_size()))
-			break;
-	return task_size;
-}
-
 #define AT_VECTOR_SIZE 40
 
 typedef uint64_t auxv_t;
diff --git a/compel/arch/aarch64/src/lib/infect.c b/compel/arch/aarch64/src/lib/infect.c
index 95d6fc592..0fc639a1e 100644
--- a/compel/arch/aarch64/src/lib/infect.c
+++ b/compel/arch/aarch64/src/lib/infect.c
@@ -3,6 +3,7 @@
 #include <sys/uio.h>
 #include <linux/elf.h>
 #include <compel/plugins/std/syscall-codes.h>
+#include "common/page.h"
 #include "uapi/compel/asm/infect-types.h"
 #include "log.h"
 #include "errno.h"
@@ -109,3 +110,24 @@ bool arch_can_dump_task(struct parasite_ctl *ctl)
 	 */
 	return true;
 }
+
+/*
+ * Range for task size calculated from the following Linux kernel files:
+ *   arch/arm64/include/asm/memory.h
+ *   arch/arm64/Kconfig
+ *
+ * TODO: handle 32 bit tasks
+ */
+#define TASK_SIZE_MIN (1UL << 39)
+#define TASK_SIZE_MAX (1UL << 48)
+
+unsigned long compel_task_size(void)
+{
+	unsigned long task_size;
+
+	for (task_size = TASK_SIZE_MIN; task_size < TASK_SIZE_MAX; task_size <<= 1)
+		if (munmap((void *)task_size, page_size()))
+			break;
+	return task_size;
+}
+
diff --git a/compel/arch/arm/src/lib/include/uapi/asm/infect-types.h b/compel/arch/arm/src/lib/include/uapi/asm/infect-types.h
index 9c2092e5d..b532c7dd9 100644
--- a/compel/arch/arm/src/lib/include/uapi/asm/infect-types.h
+++ b/compel/arch/arm/src/lib/include/uapi/asm/infect-types.h
@@ -3,7 +3,6 @@
 
 #include <stdint.h>
 #include <sys/mman.h>
-#include "common/page.h"
 
 #define SIGMAX			64
 #define SIGMAX_OLD		31
@@ -59,26 +58,6 @@ struct user_vfp_exc {
 
 #define user_regs_native(pregs)			true
 
-/*
- * Range for task size calculated from the following Linux kernel files:
- *   arch/arm/include/asm/memory.h
- *   arch/arm/Kconfig (PAGE_OFFSET values in Memory split section)
- */
-#define TASK_SIZE_MIN		0x3f000000
-#define TASK_SIZE_MAX		0xbf000000
-#define SZ_1G			0x40000000
-
-static inline unsigned long task_size(void)
-{
-	unsigned long task_size;
-
-	for (task_size = TASK_SIZE_MIN; task_size < TASK_SIZE_MAX; task_size += SZ_1G)
-		if (munmap((void *)task_size, page_size()))
-			break;
-
-	return task_size;
-}
-
 #define AT_VECTOR_SIZE		40
 
 typedef uint32_t auxv_t;
diff --git a/compel/arch/arm/src/lib/infect.c b/compel/arch/arm/src/lib/infect.c
index b45239211..7d9a23024 100644
--- a/compel/arch/arm/src/lib/infect.c
+++ b/compel/arch/arm/src/lib/infect.c
@@ -2,6 +2,7 @@
 #include <sys/types.h>
 #include <compel/plugins/std/syscall-codes.h>
 #include <compel/asm/processor-flags.h>
+#include "common/page.h"
 #include "uapi/compel/asm/infect-types.h"
 #include "log.h"
 #include "errno.h"
@@ -120,3 +121,24 @@ bool arch_can_dump_task(struct parasite_ctl *ctl)
 	 */
 	return true;
 }
+
+/*
+ * Range for task size calculated from the following Linux kernel files:
+ *   arch/arm/include/asm/memory.h
+ *   arch/arm/Kconfig (PAGE_OFFSET values in Memory split section)
+ */
+#define TASK_SIZE_MIN		0x3f000000
+#define TASK_SIZE_MAX		0xbf000000
+#define SZ_1G			0x40000000
+
+unsigned long compel_task_size(void)
+{
+	unsigned long task_size;
+
+	for (task_size = TASK_SIZE_MIN; task_size < TASK_SIZE_MAX; task_size += SZ_1G)
+		if (munmap((void *)task_size, page_size()))
+			break;
+
+	return task_size;
+}
+
diff --git a/compel/arch/ppc64/src/lib/include/uapi/asm/infect-types.h b/compel/arch/ppc64/src/lib/include/uapi/asm/infect-types.h
index f243def73..ecf643d68 100644
--- a/compel/arch/ppc64/src/lib/include/uapi/asm/infect-types.h
+++ b/compel/arch/ppc64/src/lib/include/uapi/asm/infect-types.h
@@ -93,16 +93,6 @@ typedef uint64_t auxv_t;
 /* Not used but the structure parasite_dump_thread needs a tls_t field */
 typedef uint64_t tls_t;
 
-/*
- * Copied for the Linux kernel arch/powerpc/include/asm/processor.h
- *
- * NOTE: 32bit tasks are not supported.
- */
-#define TASK_SIZE_USER64	(0x0000400000000000UL)
-#define TASK_SIZE		TASK_SIZE_USER64
-
-static inline unsigned long task_size(void) { return TASK_SIZE; }
-
 #define ARCH_SI_TRAP TRAP_BRKPT
 
 #define __NR(syscall, compat)	__NR_##syscall
diff --git a/compel/arch/ppc64/src/lib/infect.c b/compel/arch/ppc64/src/lib/infect.c
index b208f3dbb..1f1437414 100644
--- a/compel/arch/ppc64/src/lib/infect.c
+++ b/compel/arch/ppc64/src/lib/infect.c
@@ -316,3 +316,13 @@ bool arch_can_dump_task(struct parasite_ctl *ctl)
 	 */
 	return true;
 }
+
+/*
+ * Copied for the Linux kernel arch/powerpc/include/asm/processor.h
+ *
+ * NOTE: 32bit tasks are not supported.
+ */
+#define TASK_SIZE_USER64	(0x0000400000000000UL)
+#define TASK_SIZE		TASK_SIZE_USER64
+
+unsigned long compel_task_size(void) { return TASK_SIZE; }
diff --git a/compel/arch/x86/src/lib/include/uapi/asm/infect-types.h b/compel/arch/x86/src/lib/include/uapi/asm/infect-types.h
index bbc6bcf22..cb12ff00f 100644
--- a/compel/arch/x86/src/lib/include/uapi/asm/infect-types.h
+++ b/compel/arch/x86/src/lib/include/uapi/asm/infect-types.h
@@ -4,9 +4,6 @@
 #include <stdint.h>
 #include <stdbool.h>
 #include <signal.h>
-#include "log.h"
-#include "common/bug.h"
-#include "common/page.h"
 #include <compel/plugins/std/asm/syscall-types.h>
 
 #define SIGMAX			64
@@ -123,17 +120,9 @@ typedef struct {
 
 typedef struct xsave_struct user_fpregs_struct_t;
 
-#ifdef CONFIG_X86_64
-# define TASK_SIZE	((1UL << 47) - PAGE_SIZE)
-#else
-/*
- * Task size may be limited to 3G but we need a
- * higher limit, because it's backward compatible.
- */
-# define TASK_SIZE	(0xffffe000)
-#endif
-
-static inline unsigned long task_size(void) { return TASK_SIZE; }
+#define REG_RES(regs)		get_user_reg(&regs, ax)
+#define REG_IP(regs)		get_user_reg(&regs, ip)
+#define REG_SYSCALL_NR(regs)	get_user_reg(&regs, orig_ax)
 
 typedef uint64_t auxv_t;
 
diff --git a/compel/arch/x86/src/lib/infect.c b/compel/arch/x86/src/lib/infect.c
index af46c6526..a9a0556aa 100644
--- a/compel/arch/x86/src/lib/infect.c
+++ b/compel/arch/x86/src/lib/infect.c
@@ -349,3 +349,15 @@ int ptrace_set_regs(pid_t pid, user_regs_struct_t *regs)
 	}
 	return ptrace(PTRACE_SETREGSET, pid, NT_PRSTATUS, &iov);
 }
+
+#ifdef CONFIG_X86_64
+# define TASK_SIZE	((1UL << 47) - PAGE_SIZE)
+#else
+/*
+ * Task size may be limited to 3G but we need a
+ * higher limit, because it's backward compatible.
+ */
+# define TASK_SIZE	(0xffffe000)
+#endif
+
+unsigned long compel_task_size(void) { return TASK_SIZE; }
diff --git a/compel/include/uapi/infect.h b/compel/include/uapi/infect.h
index adc40f2d9..c8a6d9f6e 100644
--- a/compel/include/uapi/infect.h
+++ b/compel/include/uapi/infect.h
@@ -141,4 +141,6 @@ extern int compel_get_thread_regs(struct parasite_thread_ctl *, save_regs_t, voi
 
 extern void compel_relocs_apply(void *mem, void *vbase, size_t size, compel_reloc_t *elf_relocs, size_t nr_relocs);
 
+extern unsigned long compel_task_size(void);
+
 #endif
diff --git a/criu/kerndat.c b/criu/kerndat.c
index 9f7eb0fd7..414e9c414 100644
--- a/criu/kerndat.c
+++ b/criu/kerndat.c
@@ -28,6 +28,7 @@
 #include "config.h"
 #include "sk-inet.h"
 #include <compel/plugins/std/syscall-codes.h>
+#include <compel/compel.h>
 
 struct kerndat_s kdat = {
 };
@@ -386,7 +387,7 @@ static bool kerndat_has_memfd_create(void)
 
 static int get_task_size(void)
 {
-	kdat.task_size = task_size();
+	kdat.task_size = compel_task_size();
 	pr_debug("Found task size of %lx\n", kdat.task_size);
 	return 0;
 }