*/
 # if THREAD_SIZE >= PAGE_SIZE || defined(CONFIG_VMAP_STACK)
 
-#ifdef CONFIG_VMAP_STACK
+#  ifdef CONFIG_VMAP_STACK
 /*
  * vmalloc() is a bit slow, and calling vfree() enough times will force a TLB
  * flush.  Try to minimize the number of calls by caching stacks.
 
        return 0;
 }
-#endif
 
 static unsigned long *alloc_thread_stack_node(struct task_struct *tsk, int node)
 {
-#ifdef CONFIG_VMAP_STACK
        void *stack;
        int i;
 
                tsk->stack = stack;
        }
        return stack;
-#else
-       struct page *page = alloc_pages_node(node, THREADINFO_GFP,
-                                            THREAD_SIZE_ORDER);
-
-       if (likely(page)) {
-               tsk->stack = kasan_reset_tag(page_address(page));
-               return tsk->stack;
-       }
-       return NULL;
-#endif
 }
 
-static inline void free_thread_stack(struct task_struct *tsk)
+static void free_thread_stack(struct task_struct *tsk)
 {
-#ifdef CONFIG_VMAP_STACK
        struct vm_struct *vm = task_stack_vm_area(tsk);
+       int i;
 
-       if (vm) {
-               int i;
+       for (i = 0; i < THREAD_SIZE / PAGE_SIZE; i++)
+               memcg_kmem_uncharge_page(vm->pages[i], 0);
 
-               for (i = 0; i < THREAD_SIZE / PAGE_SIZE; i++)
-                       memcg_kmem_uncharge_page(vm->pages[i], 0);
+       for (i = 0; i < NR_CACHED_STACKS; i++) {
+               if (this_cpu_cmpxchg(cached_stacks[i], NULL,
+                                    tsk->stack_vm_area) != NULL)
+                       continue;
 
-               for (i = 0; i < NR_CACHED_STACKS; i++) {
-                       if (this_cpu_cmpxchg(cached_stacks[i],
-                                       NULL, tsk->stack_vm_area) != NULL)
-                               continue;
+               tsk->stack = NULL;
+               tsk->stack_vm_area = NULL;
+               return;
+       }
+       vfree_atomic(tsk->stack);
+       tsk->stack = NULL;
+       tsk->stack_vm_area = NULL;
+}
 
-                       return;
-               }
+#  else /* !CONFIG_VMAP_STACK */
 
-               vfree_atomic(tsk->stack);
-               return;
+static unsigned long *alloc_thread_stack_node(struct task_struct *tsk, int node)
+{
+       struct page *page = alloc_pages_node(node, THREADINFO_GFP,
+                                            THREAD_SIZE_ORDER);
+
+       if (likely(page)) {
+               tsk->stack = kasan_reset_tag(page_address(page));
+               return tsk->stack;
        }
-#endif
+       return NULL;
+}
 
+static void free_thread_stack(struct task_struct *tsk)
+{
        __free_pages(virt_to_page(tsk->stack), THREAD_SIZE_ORDER);
+       tsk->stack = NULL;
 }
-# else
+
+#  endif /* CONFIG_VMAP_STACK */
+# else /* !(THREAD_SIZE >= PAGE_SIZE || defined(CONFIG_VMAP_STACK)) */
+
 static struct kmem_cache *thread_stack_cache;
 
 static unsigned long *alloc_thread_stack_node(struct task_struct *tsk,
 static void free_thread_stack(struct task_struct *tsk)
 {
        kmem_cache_free(thread_stack_cache, tsk->stack);
+       tsk->stack = NULL;
 }
 
 void thread_stack_cache_init(void)
                                        THREAD_SIZE, NULL);
        BUG_ON(thread_stack_cache == NULL);
 }
-# endif
-#endif
+
+# endif /* THREAD_SIZE >= PAGE_SIZE || defined(CONFIG_VMAP_STACK) */
+#endif /* !CONFIG_ARCH_THREAD_STACK_ALLOCATOR */
 
 /* SLAB cache for signal_struct structures (tsk->signal) */
 static struct kmem_cache *signal_cachep;
 
        account_kernel_stack(tsk, -1);
        free_thread_stack(tsk);
-       tsk->stack = NULL;
-#ifdef CONFIG_VMAP_STACK
-       tsk->stack_vm_area = NULL;
-#endif
 }
 
 #ifdef CONFIG_THREAD_INFO_IN_TASK