static inline void hugetlb_prefault_arch_hook(struct mm_struct *mm)
 {
-       hugetlb_setup(mm);
 }
 
 static inline int is_hugepage_only_range(struct mm_struct *mm,
 
 #ifndef __ASSEMBLY__
 
 #if defined(CONFIG_HUGETLB_PAGE) || defined(CONFIG_TRANSPARENT_HUGEPAGE)
-struct mm_struct;
-extern void hugetlb_setup(struct mm_struct *mm);
+struct pt_regs;
+extern void hugetlb_setup(struct pt_regs *regs);
 #endif
 
 #define WANT_PAGE_VIRTUAL
 
         nop
 
        /* It is a huge page, use huge page TSB entry address we
-        * calculated above.
+        * calculated above.  If the huge page TSB has not been
+        * allocated, setup a trap stack and call hugetlb_setup()
+        * to do so, then return from the trap to replay the TLB
+        * miss.
+        *
+        * This is necessary to handle the case of transparent huge
+        * pages where we don't really have a non-atomic context
+        * in which to allocate the hugepage TSB hash table.  When
+        * the 'mm' faults in the hugepage for the first time, we
+        * thus handle it here.  This also makes sure that we can
+        * allocate the TSB hash table on the correct NUMA node.
         */
        TRAP_LOAD_TRAP_BLOCK(%g7, %g2)
-       ldx             [%g7 + TRAP_PER_CPU_TSB_HUGE_TEMP], %g2
-       cmp             %g2, -1
-       movne           %xcc, %g2, %g1
+       ldx             [%g7 + TRAP_PER_CPU_TSB_HUGE_TEMP], %g1
+       cmp             %g1, -1
+       bne,pt          %xcc, 60f
+        nop
+
+661:   rdpr            %pstate, %g5
+       wrpr            %g5, PSTATE_AG | PSTATE_MG, %pstate
+       .section        .sun4v_2insn_patch, "ax"
+       .word           661b
+       SET_GL(1)
+       nop
+       .previous
+
+       rdpr    %tl, %g3
+       cmp     %g3, 1
+       bne,pn  %xcc, winfix_trampoline
+        nop
+       ba,pt   %xcc, etrap
+        rd     %pc, %g7
+       call    hugetlb_setup
+        add    %sp, PTREGS_OFF, %o0
+       ba,pt   %xcc, rtrap
+        nop
+
 60:
 #endif
 
 
 #if defined(CONFIG_HUGETLB_PAGE) || defined(CONFIG_TRANSPARENT_HUGEPAGE)
        mm_rss = mm->context.huge_pte_count;
        if (unlikely(mm_rss >
-                    mm->context.tsb_block[MM_TSB_HUGE].tsb_rss_limit))
-               tsb_grow(mm, MM_TSB_HUGE, mm_rss);
+                    mm->context.tsb_block[MM_TSB_HUGE].tsb_rss_limit)) {
+               if (mm->context.tsb_block[MM_TSB_HUGE].tsb)
+                       tsb_grow(mm, MM_TSB_HUGE, mm_rss);
+               else
+                       hugetlb_setup(regs);
+
+       }
 #endif
        return;
 
 
                load_secondary_context(mm);
 }
 
-void hugetlb_setup(struct mm_struct *mm)
+void hugetlb_setup(struct pt_regs *regs)
 {
-       struct tsb_config *tp = &mm->context.tsb_block[MM_TSB_HUGE];
+       struct mm_struct *mm = current->mm;
+       struct tsb_config *tp;
 
-       if (likely(tp->tsb != NULL))
-               return;
+       if (in_atomic() || !mm) {
+               const struct exception_table_entry *entry;
+
+               entry = search_exception_tables(regs->tpc);
+               if (entry) {
+                       regs->tpc = entry->fixup;
+                       regs->tnpc = regs->tpc + 4;
+                       return;
+               }
+               pr_alert("Unexpected HugeTLB setup in atomic context.\n");
+               die_if_kernel("HugeTSB in atomic", regs);
+       }
+
+       tp = &mm->context.tsb_block[MM_TSB_HUGE];
+       if (likely(tp->tsb == NULL))
+               tsb_grow(mm, MM_TSB_HUGE, 0);
 
-       tsb_grow(mm, MM_TSB_HUGE, 0);
        tsb_context_switch(mm);
        smp_tsb_sync(mm);
 
 
                        mm->context.huge_pte_count++;
                else
                        mm->context.huge_pte_count--;
-               if (mm->context.huge_pte_count == 1)
-                       hugetlb_setup(mm);
+
+               /* Do not try to allocate the TSB hash table if we
+                * don't have one already.  We have various locks held
+                * and thus we'll end up doing a GFP_KERNEL allocation
+                * in an atomic context.
+                *
+                * Instead, we let the first TLB miss on a hugepage
+                * take care of this.
+                */
        }
 
        if (!pmd_none(orig)) {