#include <asm/mmu_context.h>
 #include <asm/msa.h>
 #include <asm/setup.h>
+#include <asm/tlbex.h>
 #include <asm/uasm.h>
 
 /* Register names */
         */
        unsigned int kscratch_mask = cpu_data[0].kscratch_mask;
 
+       if (pgd_reg != -1)
+               kscratch_mask &= ~BIT(pgd_reg);
+
        /* Pick a scratch register for storing VCPU */
        if (kscratch_mask) {
                scratch_vcpu[0] = c0_kscratch();
        return p;
 }
 
+/**
+ * kvm_mips_build_tlb_refill_exception() - Assemble TLB refill handler.
+ * @addr:      Address to start writing code.
+ * @handler:   Address of common handler (within range of @addr).
+ *
+ * Assemble TLB refill exception fast path handler for guest execution.
+ *
+ * Returns:    Next address after end of written function.
+ */
+void *kvm_mips_build_tlb_refill_exception(void *addr, void *handler)
+{
+       u32 *p = addr;
+       struct uasm_label labels[2];
+       struct uasm_reloc relocs[2];
+       struct uasm_label *l = labels;
+       struct uasm_reloc *r = relocs;
+
+       memset(labels, 0, sizeof(labels));
+       memset(relocs, 0, sizeof(relocs));
+
+       /* Save guest k1 into scratch register */
+       UASM_i_MTC0(&p, K1, scratch_tmp[0], scratch_tmp[1]);
+
+       /* Get the VCPU pointer from the VCPU scratch register */
+       UASM_i_MFC0(&p, K1, scratch_vcpu[0], scratch_vcpu[1]);
+
+       /* Save guest k0 into VCPU structure */
+       UASM_i_SW(&p, K0, offsetof(struct kvm_vcpu, arch.gprs[K0]), K1);
+
+       /*
+        * Some of the common tlbex code uses current_cpu_type(). For KVM we
+        * assume symmetry and just disable preemption to silence the warning.
+        */
+       preempt_disable();
+
+       /*
+        * Now for the actual refill bit. A lot of this can be common with the
+        * Linux TLB refill handler, however we don't need to handle so many
+        * cases. We only need to handle user mode refills, and user mode runs
+        * with 32-bit addressing.
+        *
+        * Therefore the branch to label_vmalloc generated by build_get_pmde64()
+        * that isn't resolved should never actually get taken and is harmless
+        * to leave in place for now.
+        */
+
+#ifdef CONFIG_64BIT
+       build_get_pmde64(&p, &l, &r, K0, K1); /* get pmd in K1 */
+#else
+       build_get_pgde32(&p, K0, K1); /* get pgd in K1 */
+#endif
+
+       /* we don't support huge pages yet */
+
+       build_get_ptep(&p, K0, K1);
+       build_update_entries(&p, K0, K1);
+       build_tlb_write_entry(&p, &l, &r, tlb_random);
+
+       preempt_enable();
+
+       /* Get the VCPU pointer from the VCPU scratch register again */
+       UASM_i_MFC0(&p, K1, scratch_vcpu[0], scratch_vcpu[1]);
+
+       /* Restore the guest's k0/k1 registers */
+       UASM_i_LW(&p, K0, offsetof(struct kvm_vcpu, arch.gprs[K0]), K1);
+       uasm_i_ehb(&p);
+       UASM_i_MFC0(&p, K1, scratch_tmp[0], scratch_tmp[1]);
+
+       /* Jump to guest */
+       uasm_i_eret(&p);
+
+       return p;
+}
+
 /**
  * kvm_mips_build_exception() - Assemble first level guest exception handler.
  * @addr:      Address to start writing code.
 
 struct kvm_vcpu *kvm_arch_vcpu_create(struct kvm *kvm, unsigned int id)
 {
        int err, size;
-       void *gebase, *p, *handler;
+       void *gebase, *p, *handler, *refill_start, *refill_end;
        int i;
 
        struct kvm_vcpu *vcpu = kzalloc(sizeof(struct kvm_vcpu), GFP_KERNEL);
        /* Build guest exception vectors dynamically in unmapped memory */
        handler = gebase + 0x2000;
 
-       /* TLB Refill, EXL = 0 */
-       kvm_mips_build_exception(gebase, handler);
+       /* TLB refill */
+       refill_start = gebase;
+       refill_end = kvm_mips_build_tlb_refill_exception(refill_start, handler);
 
        /* General Exception Entry point */
        kvm_mips_build_exception(gebase + 0x180, handler);
        pr_debug("#include <asm/regdef.h>\n");
        pr_debug("\n");
        dump_handler("kvm_vcpu_run", vcpu->arch.vcpu_run, p);
+       dump_handler("kvm_tlb_refill", refill_start, refill_end);
        dump_handler("kvm_gen_exc", gebase + 0x180, gebase + 0x200);
        dump_handler("kvm_exit", gebase + 0x2000, vcpu->arch.vcpu_run);