ptesync
 1:     blr
 
+/*
+ * Flush the TLB in hash mode. Hash must flush with RIC=2 once for process
+ * and one for partition scope to clear process and partition table entries.
+ */
 __init_tlb_power9:
-       li      r6,POWER9_TLB_SETS_HASH
+       li      r6,POWER9_TLB_SETS_HASH - 1
        mtctr   r6
        li      r7,0xc00        /* IS field = 0b11 */
+       li      r8,0
        ptesync
-2:     tlbiel  r7
-       addi    r7,r7,0x1000
+       PPC_TLBIEL(7, 8, 2, 1, 0)
+       PPC_TLBIEL(7, 8, 2, 0, 0)
+2:     addi    r7,r7,0x1000
+       PPC_TLBIEL(7, 8, 0, 0, 0)
        bdnz    2b
        ptesync
 1:     blr
 
 
 static void cpufeatures_flush_tlb(void)
 {
-       unsigned long rb;
-       unsigned int i, num_sets;
-
        /*
         * This is a temporary measure to keep equivalent TLB flush as the
         * cputable based setup code.
        case PVR_POWER8:
        case PVR_POWER8E:
        case PVR_POWER8NVL:
-               num_sets = POWER8_TLB_SETS;
+               __flush_tlb_power8(POWER8_TLB_SETS);
                break;
        case PVR_POWER9:
-               num_sets = POWER9_TLB_SETS_HASH;
+               __flush_tlb_power9(POWER9_TLB_SETS_HASH);
                break;
        default:
-               num_sets = 1;
                pr_err("unknown CPU version for boot TLB flush\n");
                break;
        }
-
-       asm volatile("ptesync" : : : "memory");
-       rb = TLBIEL_INVAL_SET;
-       for (i = 0; i < num_sets; i++) {
-               asm volatile("tlbiel %0" : : "r" (rb));
-               rb += 1 << TLBIEL_INVAL_SET_SHIFT;
-       }
-       asm volatile("ptesync" : : : "memory");
 }
 
 static void __restore_cpu_cpufeatures(void)
 
        asm volatile("ptesync" : : : "memory");
 }
 
+static void flush_tlb_300(unsigned int num_sets, unsigned int action)
+{
+       unsigned long rb;
+       unsigned int i;
+       unsigned int r;
+
+       switch (action) {
+       case TLB_INVAL_SCOPE_GLOBAL:
+               rb = TLBIEL_INVAL_SET;
+               break;
+       case TLB_INVAL_SCOPE_LPID:
+               rb = TLBIEL_INVAL_SET_LPID;
+               break;
+       default:
+               BUG();
+               break;
+       }
+
+       asm volatile("ptesync" : : : "memory");
+
+       if (early_radix_enabled())
+               r = 1;
+       else
+               r = 0;
+
+       /*
+        * First flush table/PWC caches with set 0, then flush the
+        * rest of the sets, partition scope. Radix must then do it
+        * all again with process scope. Hash just has to flush
+        * process table.
+        */
+       asm volatile(PPC_TLBIEL(%0, %1, %2, %3, %4) : :
+                       "r"(rb), "r"(0), "i"(2), "i"(0), "r"(r));
+       for (i = 1; i < num_sets; i++) {
+               unsigned long set = i * (1<<TLBIEL_INVAL_SET_SHIFT);
+
+               asm volatile(PPC_TLBIEL(%0, %1, %2, %3, %4) : :
+                               "r"(rb+set), "r"(0), "i"(2), "i"(0), "r"(r));
+       }
+
+       asm volatile(PPC_TLBIEL(%0, %1, %2, %3, %4) : :
+                       "r"(rb), "r"(0), "i"(2), "i"(1), "r"(r));
+       if (early_radix_enabled()) {
+               for (i = 1; i < num_sets; i++) {
+                       unsigned long set = i * (1<<TLBIEL_INVAL_SET_SHIFT);
+
+                       asm volatile(PPC_TLBIEL(%0, %1, %2, %3, %4) : :
+                               "r"(rb+set), "r"(0), "i"(2), "i"(1), "r"(r));
+               }
+       }
+
+       asm volatile("ptesync" : : : "memory");
+}
+
 /*
  * Generic routines to flush TLB on POWER processors. These routines
  * are used as flush_tlb hook in the cpu_spec.
        else
                num_sets = POWER9_TLB_SETS_HASH;
 
-       flush_tlb_206(num_sets, action);
+       flush_tlb_300(num_sets, action);
 }