#include <asm/regdef.h>
 #include <asm/stackframe.h>
 
+#define INVTLB_ADDR_GFALSE_AND_ASID    5
+
 #define PTRS_PER_PGD_BITS      (PAGE_SHIFT - 3)
 #define PTRS_PER_PUD_BITS      (PAGE_SHIFT - 3)
 #define PTRS_PER_PMD_BITS      (PAGE_SHIFT - 3)
        ori             t0, ra, _PAGE_VALID
        st.d            t0, t1, 0
 #endif
-       tlbsrch
-       addu16i.d       t1, zero, -(CSR_TLBIDX_EHINV >> 16)
-       addi.d          ra, t1, 0
-       csrxchg         ra, t1, LOONGARCH_CSR_TLBIDX
-       tlbwr
-
-       csrxchg         zero, t1, LOONGARCH_CSR_TLBIDX
+       csrrd           ra, LOONGARCH_CSR_ASID
+       csrrd           t1, LOONGARCH_CSR_BADV
+       andi            ra, ra, CSR_ASID_ASID
+       invtlb          INVTLB_ADDR_GFALSE_AND_ASID, ra, t1
 
        /*
         * A huge PTE describes an area the size of the
        ori             t0, ra, (_PAGE_VALID | _PAGE_DIRTY | _PAGE_MODIFIED)
        st.d            t0, t1, 0
 #endif
-       tlbsrch
-       addu16i.d       t1, zero, -(CSR_TLBIDX_EHINV >> 16)
-       addi.d          ra, t1, 0
-       csrxchg         ra, t1, LOONGARCH_CSR_TLBIDX
-       tlbwr
+       csrrd           ra, LOONGARCH_CSR_ASID
+       csrrd           t1, LOONGARCH_CSR_BADV
+       andi            ra, ra, CSR_ASID_ASID
+       invtlb          INVTLB_ADDR_GFALSE_AND_ASID, ra, t1
 
-       csrxchg         zero, t1, LOONGARCH_CSR_TLBIDX
        /*
         * A huge PTE describes an area the size of the
         * configured huge page size. This is twice the
        ori             t0, ra, (_PAGE_VALID | _PAGE_DIRTY | _PAGE_MODIFIED)
        st.d            t0, t1, 0
 #endif
+       csrrd           ra, LOONGARCH_CSR_ASID
+       csrrd           t1, LOONGARCH_CSR_BADV
+       andi            ra, ra, CSR_ASID_ASID
+       invtlb          INVTLB_ADDR_GFALSE_AND_ASID, ra, t1
+
        /*
         * A huge PTE describes an area the size of the
         * configured huge page size. This is twice the
        addu16i.d       t1, zero, (PS_HUGE_SIZE << (CSR_TLBIDX_PS_SHIFT - 16))
        csrxchg         t1, t0, LOONGARCH_CSR_TLBIDX
 
-       tlbwr
+       tlbfill
 
        /* Reset default page size */
        addu16i.d       t0, zero, (CSR_TLBIDX_PS >> 16)