diff --git a/init/main.c b/init/main.c
index c4bc168fb11f..f241ce702dbb 100644
--- a/init/main.c
+++ b/init/main.c
@@ -899,11 +899,29 @@ static void __init early_numa_node_init(void)
 #endif
 }
 
-#ifdef CONFIG_ARCH_WANTS_DYNAMIC_TASK_STRUCT
+#ifdef CONFIG_GKI_DYNAMIC_TASK_STRUCT_SIZE
 static void __init setup_arch_task_struct_size(void)
 {
 	arch_task_struct_size = sizeof(struct task_struct);
 }
+
+static int __init task_struct_vendor_size_setup(char *str)
+{
+	u64 size;
+
+	if (!str)
+		return -EINVAL;
+
+	size = memparse(str, &str);
+
+	if (size < 0 || size > CONFIG_GKI_TASK_STRUCT_VENDOR_SIZE_MAX)
+		return -EINVAL;
+
+	arch_task_struct_size = sizeof(struct task_struct) + size;
+
+	return 0;
+}
+early_param("android_arch_task_struct_size", task_struct_vendor_size_setup);
 #endif
 
 asmlinkage __visible __init __no_sanitize_address __noreturn __no_stack_protector
@@ -929,8 +947,8 @@ void start_kernel(void)
 	boot_cpu_init();
 	page_address_init();
 	pr_notice("%s", linux_banner);
-#ifdef CONFIG_ARCH_WANTS_DYNAMIC_TASK_STRUCT
-        setup_arch_task_struct_size();
+#ifdef CONFIG_GKI_DYNAMIC_TASK_STRUCT_SIZE
+	setup_arch_task_struct_size();
 #endif
 	setup_arch(&command_line);
 	/* Static keys and static calls are needed by LSMs */
diff --git a/kernel/fork.c b/kernel/fork.c
index 34f2ebdd8969..786eb6f8cf47 100644
--- a/kernel/fork.c
+++ b/kernel/fork.c
@@ -1034,24 +1034,6 @@ static void __init set_max_threads(unsigned int max_threads_suggested)
 #ifdef CONFIG_ARCH_WANTS_DYNAMIC_TASK_STRUCT
 /* Initialized by the architecture: */
 int arch_task_struct_size __read_mostly;
-
-static int __init arch_task_struct_size_setup(char *str)
-{
-	u64 size;
-
-	if (!str)
-		return -EINVAL;
-
-	size = memparse(str, &str);
-
-	if (size < 0 || size > CONFIG_GKI_TASK_STRUCT_VENDOR_SIZE_MAX)
-		return -EINVAL;
-
-	arch_task_struct_size = sizeof(struct task_struct) + size;
-
-	return 0;
-}
-early_param("android_task_struct_vendor_size", arch_task_struct_size_setup);
 #endif
 
 static void __init task_struct_whitelist(unsigned long *offset, unsigned long *size)