0001-patch-v2.patch

From 02f6467c5da02a0d16805ad6fe506a2a3d68fbd7 Mon Sep 17 00:00:00 2001
From: fengeh <2748250768@qq.com>
Date: Thu, 19 Aug 2021 17:25:21 +0800
Subject: [PATCH] patch-v2

---
 arch/riscv/Kconfig                  |  14 ++
 arch/riscv/configs/defconfig        |   4 +-
 arch/riscv/include/asm/pgalloc.h    |  32 ++++
 arch/riscv/include/asm/pgtable-64.h |  21 +++
 arch/riscv/include/asm/pgtable.h    |  81 +++++++-
 arch/riscv/include/asm/sbi.h        |  10 +
 arch/riscv/kernel/head.S            |   7 +-
 arch/riscv/kernel/sbi.c             |  61 ++++++
 arch/riscv/kernel/setup.c           |   5 +-
 arch/riscv/kernel/time.c            |  11 +-
 arch/riscv/mm/fault.c               |   1 +
 arch/riscv/mm/init.c                |   5 +-
 fs/stat.c                           |   1 +
 include/asm-generic/pgalloc.h       | 150 +++++++++++++++
 include/linux/mm.h                  |   4 +
 include/linux/mmzone.h              |   2 +-
 include/linux/pgtable.h             |  42 +++++
 include/linux/pt_area.h             |  25 +++
 init/main.c                         |  86 +++++++++
 mm/Makefile                         |   4 +
 mm/filemap.c                        |  21 ++-
 mm/memory.c                         | 254 ++++++++++++++++++++++++-
 mm/oom_kill.c                       |  16 +-
 mm/pt_area.c                        | 283 ++++++++++++++++++++++++++++
 mm/swap_state.c                     |  17 +-
 mm/vmalloc.c                        |  95 +++++++++-
 26 files changed, 1227 insertions(+), 25 deletions(-)
 create mode 100644 include/linux/pt_area.h
 create mode 100644 mm/pt_area.c

diff --git a/arch/riscv/Kconfig b/arch/riscv/Kconfig
index 44377fd78..b0f2d66e5 100644
--- a/arch/riscv/Kconfig
+++ b/arch/riscv/Kconfig
@@ -111,6 +111,20 @@ config RISCV_SBI
 	depends on !RISCV_M_MODE
 	default y
 
+# set if we enable PT_AREA for penglai enclave
+config PT_AREA
+  	bool 
+	default y
+
+config PT_AREA_BATCH
+  	bool
+	depends on PT_AREA
+	default y
+
+config EARLY_PRINTK
+  	bool 
+	default y
+
 config MMU
 	bool "MMU-based Paged Memory Management Support"
 	default y
diff --git a/arch/riscv/configs/defconfig b/arch/riscv/configs/defconfig
index d222d353d..5c641aeb2 100644
--- a/arch/riscv/configs/defconfig
+++ b/arch/riscv/configs/defconfig
@@ -118,7 +118,7 @@ CONFIG_DEBUG_RT_MUTEXES=y
 CONFIG_DEBUG_SPINLOCK=y
 CONFIG_DEBUG_MUTEXES=y
 CONFIG_DEBUG_RWSEMS=y
-CONFIG_DEBUG_ATOMIC_SLEEP=y
+CONFIG_DEBUG_ATOMIC_SLEEP=n
 CONFIG_STACKTRACE=y
 CONFIG_DEBUG_LIST=y
 CONFIG_DEBUG_PLIST=y
@@ -131,3 +131,5 @@ CONFIG_DEBUG_BLOCK_EXT_DEVT=y
 CONFIG_MEMTEST=y
 # CONFIG_SYSFS_SYSCALL is not set
 CONFIG_EFI=y
+CONFIG_CMDLINE="earlycon=ttyS0 console=hvc"
+CONFIG_CMDLINE_FORCE=y
diff --git a/arch/riscv/include/asm/pgalloc.h b/arch/riscv/include/asm/pgalloc.h
index 23b1544e0..4653ad5e8 100644
--- a/arch/riscv/include/asm/pgalloc.h
+++ b/arch/riscv/include/asm/pgalloc.h
@@ -10,6 +10,11 @@
 #include <linux/mm.h>
 #include <asm/tlb.h>
 
+#ifdef CONFIG_PT_AREA
+#include <linux/pt_area.h>
+#include <asm/sbi.h>
+#endif
+
 #ifdef CONFIG_MMU
 #include <asm-generic/pgalloc.h>
 
@@ -39,7 +44,33 @@ static inline void pud_populate(struct mm_struct *mm, pud_t *pud, pmd_t *pmd)
 #endif /* __PAGETABLE_PMD_FOLDED */
 
 #define pmd_pgtable(pmd)	pmd_page(pmd)
+#ifdef CONFIG_PT_AREA
+static inline pgd_t *pgd_alloc(struct mm_struct *mm)
+{
+	pgd_t *pgd;
 
+	pgd = (pgd_t *)alloc_pt_pgd_page();
+	if(likely(pgd != NULL))
+	{
+		if(enclave_module_installed)
+		{
+		SBI_PENGLAI_ECALL_4(SBI_SM_SET_PTE, SBI_PTE_MEMSET, __pa(pgd), 0, USER_PTRS_PER_PGD * sizeof(pgd_t));
+		SBI_PENGLAI_ECALL_4(SBI_SM_SET_PTE, SBI_PTE_MEMCPY, __pa(pgd + USER_PTRS_PER_PGD),
+			__pa(init_mm.pgd + USER_PTRS_PER_PGD),
+			(PTRS_PER_PGD - USER_PTRS_PER_PGD) * sizeof(pgd_t));
+		}
+		else
+		{
+		memset(pgd, 0, USER_PTRS_PER_PGD * sizeof(pgd_t));
+		/* Copy kernel mappings */
+		memcpy(pgd + USER_PTRS_PER_PGD,
+			init_mm.pgd + USER_PTRS_PER_PGD,
+			(PTRS_PER_PGD - USER_PTRS_PER_PGD) * sizeof(pgd_t));
+		}
+	}
+	return pgd;
+}
+#else /* CONFIG_PT_AREA */
 static inline pgd_t *pgd_alloc(struct mm_struct *mm)
 {
 	pgd_t *pgd;
@@ -54,6 +85,7 @@ static inline pgd_t *pgd_alloc(struct mm_struct *mm)
 	}
 	return pgd;
 }
+#endif /* CONFIG_PT_AREA */
 
 #ifndef __PAGETABLE_PMD_FOLDED
 
diff --git a/arch/riscv/include/asm/pgtable-64.h b/arch/riscv/include/asm/pgtable-64.h
index f3b0da64c..e7129dd69 100644
--- a/arch/riscv/include/asm/pgtable-64.h
+++ b/arch/riscv/include/asm/pgtable-64.h
@@ -7,6 +7,7 @@
 #define _ASM_RISCV_PGTABLE_64_H
 
 #include <linux/const.h>
+#include <asm/sbi.h>
 
 #define PGDIR_SHIFT     30
 /* Size of region mapped by a page global directory */
@@ -50,10 +51,30 @@ static inline int pud_leaf(pud_t pud)
 	       (pud_val(pud) & (_PAGE_READ | _PAGE_WRITE | _PAGE_EXEC));
 }
 
+#ifdef CONFIG_PT_AREA
+extern int enclave_module_installed;
+#define SBI_SM_SET_PTE 101
+#define SBI_SET_PTE_ONE 1
+#define SBI_PTE_MEMSET 2
+#define SBI_PTE_MEMCPY 3
+
+static void set_pud(pud_t *pudp, pud_t pud)
+{
+  if(enclave_module_installed)
+  {
+    SBI_PENGLAI_ECALL_4(SBI_SM_SET_PTE, SBI_SET_PTE_ONE, __pa(pudp), pud.p4d.pgd.pgd, 0);
+  }
+  else
+  {
+    *pudp = pud;
+  }
+}
+#else
 static inline void set_pud(pud_t *pudp, pud_t pud)
 {
 	*pudp = pud;
 }
+#endif /*CONFIG_PT_AREA*/
 
 static inline void pud_clear(pud_t *pudp)
 {
diff --git a/arch/riscv/include/asm/pgtable.h b/arch/riscv/include/asm/pgtable.h
index 183f1f4b2..de5aed65a 100644
--- a/arch/riscv/include/asm/pgtable.h
+++ b/arch/riscv/include/asm/pgtable.h
@@ -18,6 +18,8 @@
 #include <asm/page.h>
 #include <asm/tlbflush.h>
 #include <linux/mm_types.h>
+#include <asm/sbi.h>
+#include <linux/kernel.h>
 
 #ifdef CONFIG_MMU
 
@@ -114,6 +116,10 @@
 #define _PAGE_IOREMAP _PAGE_KERNEL
 
 extern pgd_t swapper_pg_dir[];
+//penglai extension
+#ifdef CONFIG_PT_AREA
+extern pgd_t *new_swapper_pg_dir;
+#endif /*CONFIG_PT_AREA*/
 
 /* MAP_PRIVATE permissions: xwr (copy-on-write) */
 #define __P000	PAGE_NONE
@@ -159,7 +165,16 @@ static inline int pmd_leaf(pmd_t pmd)
 
 static inline void set_pmd(pmd_t *pmdp, pmd_t pmd)
 {
+	#ifdef CONFIG_PT_AREA
+	if(enclave_module_installed)
+	{
+		SBI_PENGLAI_ECALL_4(SBI_SM_SET_PTE, SBI_SET_PTE_ONE, __pa(pmdp), pmd.pmd, 0);
+	}
+	else
+		*pmdp=pmd;
+	#else
 	*pmdp = pmd;
+	#endif /*CONFIG_PT_AREA*/
 }
 
 static inline void pmd_clear(pmd_t *pmdp)
@@ -320,6 +335,7 @@ static inline int pte_same(pte_t pte_a, pte_t pte_b)
 	return pte_val(pte_a) == pte_val(pte_b);
 }
 
+
 /*
  * Certain architectures need to do special things when PTEs within
  * a page table are directly modified.  Thus, the following hook is
@@ -327,7 +343,19 @@ static inline int pte_same(pte_t pte_a, pte_t pte_b)
  */
 static inline void set_pte(pte_t *ptep, pte_t pteval)
 {
-	*ptep = pteval;
+	#ifdef CONFIG_PT_AREA
+
+	if(enclave_module_installed)
+	{
+		// printk("set_pte\n");
+		// dump_stack();
+		SBI_PENGLAI_ECALL_4(SBI_SM_SET_PTE, SBI_SET_PTE_ONE, __pa(ptep), pteval.pte, 0);
+	}
+	else
+		*ptep = pteval;
+	#else
+		*ptep = pteval;
+	#endif /*CONFIG_PT_AREA*/
 }
 
 void flush_icache_pte(pte_t pte);
@@ -341,12 +369,33 @@ static inline void set_pte_at(struct mm_struct *mm,
 	set_pte(ptep, pteval);
 }
 
+#ifdef CONFIG_PT_AREA_BATCH
+static inline void delay_set_pte_at(struct mm_struct *mm,
+	unsigned long addr, pte_t *ptep, pte_t pteval)
+{
+	if (pte_present(pteval) && pte_exec(pteval))
+		flush_icache_pte(pteval);
+
+	// set_pte(ptep, pteval);
+}
+#endif
+
 static inline void pte_clear(struct mm_struct *mm,
 	unsigned long addr, pte_t *ptep)
 {
 	set_pte_at(mm, addr, ptep, __pte(0));
 }
 
+// delay pte set operation
+#ifdef  CONFIG_PT_AREA_BATCH
+static inline void delay_pte_clear(struct mm_struct *mm,
+	unsigned long addr, pte_t *ptep)
+{
+	if (pte_present(__pte(0)) && pte_exec(__pte(0)))
+		flush_icache_pte(__pte(0));
+}
+#endif
+
 #define __HAVE_ARCH_PTEP_SET_ACCESS_FLAGS
 static inline int ptep_set_access_flags(struct vm_area_struct *vma,
 					unsigned long address, pte_t *ptep,
@@ -365,7 +414,18 @@ static inline int ptep_set_access_flags(struct vm_area_struct *vma,
 static inline pte_t ptep_get_and_clear(struct mm_struct *mm,
 				       unsigned long address, pte_t *ptep)
 {
-	return __pte(atomic_long_xchg((atomic_long_t *)ptep, 0));
+	#ifdef CONFIG_PT_AREA
+	if(enclave_module_installed)
+	{
+		pte_t pte = *ptep;
+		pte_clear(mm, address, ptep);
+		return pte;
+	}
+	else
+		return __pte(atomic_long_xchg((atomic_long_t *)ptep, 0));
+	#else
+		return __pte(atomic_long_xchg((atomic_long_t *)ptep, 0));
+	#endif
 }
 
 #define __HAVE_ARCH_PTEP_TEST_AND_CLEAR_YOUNG
@@ -382,7 +442,22 @@ static inline int ptep_test_and_clear_young(struct vm_area_struct *vma,
 static inline void ptep_set_wrprotect(struct mm_struct *mm,
 				      unsigned long address, pte_t *ptep)
 {
-	atomic_long_and(~(unsigned long)_PAGE_WRITE, (atomic_long_t *)ptep);
+	#ifdef CONFIG_PT_AREA
+	if(enclave_module_installed)
+	{
+		pte_t pteval;
+		pteval.pte = (~(unsigned long)_PAGE_WRITE) & (ptep->pte);
+		// printk("ptep_set_wrprotect\n");
+		// dump_stack();
+		SBI_PENGLAI_ECALL_4(SBI_SM_SET_PTE, SBI_SET_PTE_ONE, __pa(ptep), pteval.pte, 0);
+	}
+	else
+	{
+		atomic_long_and(~(unsigned long)_PAGE_WRITE, (atomic_long_t *)ptep);
+	}
+	#else
+		atomic_long_and(~(unsigned long)_PAGE_WRITE, (atomic_long_t *)ptep);
+	#endif
 }
 
 #define __HAVE_ARCH_PTEP_CLEAR_YOUNG_FLUSH
diff --git a/arch/riscv/include/asm/sbi.h b/arch/riscv/include/asm/sbi.h
index 653edb25d..56ac0c7db 100644
--- a/arch/riscv/include/asm/sbi.h
+++ b/arch/riscv/include/asm/sbi.h
@@ -27,6 +27,9 @@ enum sbi_ext_id {
 	SBI_EXT_IPI = 0x735049,
 	SBI_EXT_RFENCE = 0x52464E43,
 	SBI_EXT_HSM = 0x48534D,
+	//TODO:
+	//Why this magic number
+	SBI_EXT_PENGLAI = 0x100100,
 };
 
 enum sbi_ext_base_fid {
@@ -95,6 +98,13 @@ struct sbiret sbi_ecall(int ext, int fid, unsigned long arg0,
 			unsigned long arg3, unsigned long arg4,
 			unsigned long arg5);
 
+long SBI_PENGLAI_ECALL_0(int fid);
+long SBI_PENGLAI_ECALL_1(int fid, unsigned long arg0);
+long SBI_PENGLAI_ECALL_2(int fid, unsigned long arg0, unsigned long arg1);
+long SBI_PENGLAI_ECALL_3(int fid, unsigned long arg0, unsigned long arg1, unsigned long arg2);
+long SBI_PENGLAI_ECALL_4(int fid, unsigned long arg0, unsigned long arg1, unsigned long arg2, unsigned long arg3);
+long SBI_PENGLAI_ECALL_5(int fid, unsigned long arg0, unsigned long arg1, unsigned long arg2, unsigned long arg3, unsigned long arg4);
+
 void sbi_console_putchar(int ch);
 int sbi_console_getchar(void);
 void sbi_set_timer(uint64_t stime_value);
diff --git a/arch/riscv/kernel/head.S b/arch/riscv/kernel/head.S
index 7e849797c..dab0f3c3a 100644
--- a/arch/riscv/kernel/head.S
+++ b/arch/riscv/kernel/head.S
@@ -155,7 +155,12 @@ secondary_start_common:
 
 #ifdef CONFIG_MMU
 	/* Enable virtual memory and relocate to virtual address */
-	la a0, swapper_pg_dir
+	#ifdef CONFIG_PT_AREA
+		la a0, new_swapper_pg_dir
+		ld a0, (a0)
+	#else
+		la a0, swapper_pg_dir
+	#endif
 	call relocate
 #endif
 	call setup_trap_vector
diff --git a/arch/riscv/kernel/sbi.c b/arch/riscv/kernel/sbi.c
index 226ccce0f..98ce3c655 100644
--- a/arch/riscv/kernel/sbi.c
+++ b/arch/riscv/kernel/sbi.c
@@ -46,6 +46,67 @@ struct sbiret sbi_ecall(int ext, int fid, unsigned long arg0,
 }
 EXPORT_SYMBOL(sbi_ecall);
 
+long SBI_PENGLAI_ECALL_0(int fid)
+{
+	struct sbiret ret;
+	ret = sbi_ecall(SBI_EXT_PENGLAI, fid, 0, 0, 0, 0, 0, 0);
+	if (ret.error == 0)
+		return ret.value;
+	return ret.error;
+}
+
+long SBI_PENGLAI_ECALL_1(int fid, unsigned long arg0)
+{
+	struct sbiret ret;
+	ret = sbi_ecall(SBI_EXT_PENGLAI, fid, arg0, 0, 0, 0, 0, 0);
+	if (ret.error == 0)
+		return ret.value;
+	return ret.error;
+}
+
+long SBI_PENGLAI_ECALL_2(int fid, unsigned long arg0, unsigned long arg1)
+{
+	struct sbiret ret;
+	ret = sbi_ecall(SBI_EXT_PENGLAI, fid, arg0, arg1, 0, 0, 0, 0);
+	if (ret.error == 0)
+		return ret.value;
+	return ret.error;
+}
+
+long SBI_PENGLAI_ECALL_3(int fid, unsigned long arg0, unsigned long arg1, unsigned long arg2)
+{
+	struct sbiret ret;
+	ret = sbi_ecall(SBI_EXT_PENGLAI, fid, arg0, arg1, arg2, 0, 0, 0);
+	if (ret.error == 0)
+		return ret.value;
+	return ret.error;
+}
+
+long SBI_PENGLAI_ECALL_4(int fid, unsigned long arg0, unsigned long arg1, unsigned long arg2, unsigned long arg3)
+{
+	struct sbiret ret;
+	ret = sbi_ecall(SBI_EXT_PENGLAI, fid, arg0, arg1, arg2, arg3, 0, 0);
+	if (ret.error == 0)
+		return ret.value;
+	return ret.error;
+}
+
+long SBI_PENGLAI_ECALL_5(int fid, unsigned long arg0, unsigned long arg1, unsigned long arg2, unsigned long arg3, unsigned long arg4)
+{
+	struct sbiret ret;
+	ret = sbi_ecall(SBI_EXT_PENGLAI, fid, arg0, arg1, arg2, arg3, arg4, 0);
+	if (ret.error == 0)
+		return ret.value;
+	return ret.error;
+}
+
+EXPORT_SYMBOL(SBI_PENGLAI_ECALL_0);
+EXPORT_SYMBOL(SBI_PENGLAI_ECALL_1);
+EXPORT_SYMBOL(SBI_PENGLAI_ECALL_2);
+EXPORT_SYMBOL(SBI_PENGLAI_ECALL_3);
+EXPORT_SYMBOL(SBI_PENGLAI_ECALL_4);
+EXPORT_SYMBOL(SBI_PENGLAI_ECALL_5);
+
 int sbi_err_map_linux_errno(int err)
 {
 	switch (err) {
diff --git a/arch/riscv/kernel/setup.c b/arch/riscv/kernel/setup.c
index 117f3212a..e7ef5ca0b 100644
--- a/arch/riscv/kernel/setup.c
+++ b/arch/riscv/kernel/setup.c
@@ -54,9 +54,12 @@ static DEFINE_PER_CPU(struct cpu, cpu_devices);
 static void __init parse_dtb(void)
 {
 	/* Early scan of device tree from init memory */
+	strlcpy(boot_command_line, CONFIG_CMDLINE, COMMAND_LINE_SIZE);
 	if (early_init_dt_scan(dtb_early_va))
+	{
+		strlcpy(boot_command_line, CONFIG_CMDLINE, COMMAND_LINE_SIZE);
 		return;
-
+	}
 	pr_err("No DTB passed to the kernel\n");
 #ifdef CONFIG_CMDLINE_FORCE
 	strlcpy(boot_command_line, CONFIG_CMDLINE, COMMAND_LINE_SIZE);
diff --git a/arch/riscv/kernel/time.c b/arch/riscv/kernel/time.c
index 4d3a1048a..0a76e5bf8 100644
--- a/arch/riscv/kernel/time.c
+++ b/arch/riscv/kernel/time.c
@@ -17,11 +17,12 @@ void __init time_init(void)
 	struct device_node *cpu;
 	u32 prop;
 
-	cpu = of_find_node_by_path("/cpus");
-	if (!cpu || of_property_read_u32(cpu, "timebase-frequency", &prop))
-		panic(KERN_WARNING "RISC-V system with no 'timebase-frequency' in DTS\n");
-	of_node_put(cpu);
-	riscv_timebase = prop;
+	// cpu = of_find_node_by_path("/cpus");
+	// if (!cpu || of_property_read_u32(cpu, "timebase-frequency", &prop))
+	// 	panic(KERN_WARNING "RISC-V system with no 'timebase-frequency' in DTS\n");
+	// of_node_put(cpu);
+	// riscv_timebase = prop;
+	riscv_timebase = 1000000;
 
 	lpj_fine = riscv_timebase / HZ;
 	timer_probe();
diff --git a/arch/riscv/mm/fault.c b/arch/riscv/mm/fault.c
index 3c8b9e433..e5278cf2a 100644
--- a/arch/riscv/mm/fault.c
+++ b/arch/riscv/mm/fault.c
@@ -33,6 +33,7 @@ static inline void no_context(struct pt_regs *regs, unsigned long addr)
 	pr_alert("Unable to handle kernel %s at virtual address " REG_FMT "\n",
 		(addr < PAGE_SIZE) ? "NULL pointer dereference" :
 		"paging request", addr);
+	dump_stack();
 	die(regs, "Oops");
 	do_exit(SIGKILL);
 }
diff --git a/arch/riscv/mm/init.c b/arch/riscv/mm/init.c
index 8e577f14f..43cc17989 100644
--- a/arch/riscv/mm/init.c
+++ b/arch/riscv/mm/init.c
@@ -208,9 +208,12 @@ EXPORT_SYMBOL(pfn_base);
 
 pgd_t swapper_pg_dir[PTRS_PER_PGD] __page_aligned_bss;
 pgd_t trampoline_pg_dir[PTRS_PER_PGD] __page_aligned_bss;
+#ifdef CONFIG_PT_AREA
+pgd_t *new_swapper_pg_dir;
+#endif
 pte_t fixmap_pte[PTRS_PER_PTE] __page_aligned_bss;
 
-#define MAX_EARLY_MAPPING_SIZE	SZ_128M
+#define MAX_EARLY_MAPPING_SIZE	SZ_256M
 
 pgd_t early_pg_dir[PTRS_PER_PGD] __initdata __aligned(PAGE_SIZE);
 
diff --git a/fs/stat.c b/fs/stat.c
index dacecdda2..8fdd475b1 100644
--- a/fs/stat.c
+++ b/fs/stat.c
@@ -207,6 +207,7 @@ int vfs_fstatat(int dfd, const char __user *filename,
 	return vfs_statx(dfd, filename, flags | AT_NO_AUTOMOUNT,
 			 stat, STATX_BASIC_STATS);
 }
+EXPORT_SYMBOL(vfs_fstatat);
 
 #ifdef __ARCH_WANT_OLD_STAT
 
diff --git a/include/asm-generic/pgalloc.h b/include/asm-generic/pgalloc.h
index 02932efad..f53fded62 100644
--- a/include/asm-generic/pgalloc.h
+++ b/include/asm-generic/pgalloc.h
@@ -2,11 +2,159 @@
 #ifndef __ASM_GENERIC_PGALLOC_H
 #define __ASM_GENERIC_PGALLOC_H
 
+#ifdef CONFIG_PT_AREA
+#include <linux/pt_area.h>
+#include <asm/atomic.h>
+#endif
+
 #ifdef CONFIG_MMU
 
 #define GFP_PGTABLE_KERNEL	(GFP_KERNEL | __GFP_ZERO)
 #define GFP_PGTABLE_USER	(GFP_PGTABLE_KERNEL | __GFP_ACCOUNT)
 
+#ifdef CONFIG_PT_AREA
+
+static inline pte_t *__pte_alloc_one_kernel(struct mm_struct *mm)
+{
+	pte_t* pte;
+	pte = (pte_t*) alloc_pt_pte_page();
+	return pte;
+}
+
+#ifndef __HAVE_ARCH_PTE_ALLOC_ONE_KERNEL
+static inline pte_t *pte_alloc_one_kernel(struct mm_struct *mm)
+{
+	return __pte_alloc_one_kernel(mm);
+}
+#endif
+
+
+
+
+static inline void pte_free_kernel(struct mm_struct *mm, pte_t *pte)
+{
+	free_pt_pte_page((unsigned long)pte);
+}
+
+static inline pgtable_t __pte_alloc_one(struct mm_struct *mm, gfp_t gfp)
+{
+	struct page *pte;
+	char *pt_pte_addr = alloc_pt_pte_page();
+	if (pt_pte_addr == NULL)
+		return NULL;
+
+	pte = virt_to_page((void*)(pt_pte_addr));
+	if (!pte)
+		return NULL;
+	if (!pgtable_pte_page_ctor(pte)) {
+		free_pt_pte_page((unsigned long)(page_address(pte)));
+		return NULL;
+	}
+	return pte;
+}
+
+#ifndef __HAVE_ARCH_PTE_ALLOC_ONE
+static inline pgtable_t pte_alloc_one(struct mm_struct *mm)
+{
+	return __pte_alloc_one(mm, GFP_PGTABLE_USER);
+}
+#endif
+
+static inline void pte_free(struct mm_struct *mm, struct page *pte_page)
+{
+	pgtable_pte_page_dtor(pte_page);
+	*(unsigned long *)&pte_page->ptl = 0;
+	free_pt_pte_page((unsigned long)(page_address(pte_page)));
+}
+
+static inline void pte_free_no_dtor(struct mm_struct *mm, struct page *pte_page)
+{
+	*(unsigned long *)&pte_page->ptl = 0;
+	free_pt_pte_page((unsigned long)(page_address(pte_page)));
+}
+
+static inline int check_pte(struct mm_struct *mm, struct page *pte_page)
+{
+	return check_pt_pte_page((unsigned long)(page_address(pte_page)));
+}
+
+#if CONFIG_PGTABLE_LEVELS > 2
+
+#ifndef __HAVE_ARCH_PMD_ALLOC_ONE
+static inline pmd_t *pmd_alloc_one(struct mm_struct *mm, unsigned long addr)
+{
+	struct page *page;
+	gfp_t gfp = GFP_PGTABLE_USER;
+
+	if (mm == &init_mm)
+		gfp = GFP_PGTABLE_KERNEL;
+
+	char *pt_pmd_addr = alloc_pt_pmd_page();
+	if (pt_pmd_addr == NULL)
+		return NULL;
+
+	page = virt_to_page((void *)pt_pmd_addr);
+	if (!page)
+		return NULL;
+	if (!pgtable_pmd_page_ctor(page)) {
+		free_pt_pmd_page((unsigned long)(page_address(page)));
+		return NULL;
+	}
+	return (pmd_t *)page_address(page);
+}
+#endif
+
+#ifndef __HAVE_ARCH_PMD_FREE
+static inline void pmd_free(struct mm_struct *mm, pmd_t *pmd)
+{
+	struct page *page = virt_to_page(pmd);
+	pgtable_pmd_page_dtor(page);
+	*(unsigned long *)&page->ptl = 0;
+	free_pt_pmd_page((unsigned long)pmd);
+}
+#endif
+
+#ifndef __HAVE_ARCH_PMD_FREE
+static inline void pmd_free_no_dtor(struct mm_struct *mm, pmd_t *pmd)
+{
+	free_pt_pmd_page((unsigned long)pmd);
+}
+#endif
+
+#endif /* CONFIG_PGTABLE_LEVELS > 2 */
+
+#if CONFIG_PGTABLE_LEVELS > 3
+
+#ifndef __HAVE_ARCH_PUD_ALLOC_ONE
+static inline pud_t *pud_alloc_one(struct mm_struct *mm, unsigned long addr)
+{
+	gfp_t gfp = GFP_PGTABLE_USER;
+
+	if (mm == &init_mm)
+		gfp = GFP_PGTABLE_KERNEL;
+	return (pud_t *)get_zeroed_page(gfp);
+}
+#endif
+
+static inline void pud_free(struct mm_struct *mm, pud_t *pud)
+{
+	BUG_ON((unsigned long)pud & (PAGE_SIZE-1));
+	free_page((unsigned long)pud);
+}
+
+#endif /* CONFIG_PGTABLE_LEVELS > 3 */
+
+#ifndef __HAVE_ARCH_PGD_FREE
+static inline void pgd_free(struct mm_struct *mm, pgd_t *pgd)
+{
+	free_pt_pgd_page((unsigned long)pgd);
+}
+#endif
+
+
+
+#else
+
 /**
  * __pte_alloc_one_kernel - allocate a page for PTE-level kernel page table
  * @mm: the mm_struct of the current context
@@ -182,6 +330,8 @@ static inline void pgd_free(struct mm_struct *mm, pgd_t *pgd)
 }
 #endif
 
+#endif /* CONFIG_PT_AREA */
+
 #endif /* CONFIG_MMU */
 
 #endif /* __ASM_GENERIC_PGALLOC_H */
diff --git a/include/linux/mm.h b/include/linux/mm.h
index db6ae4d3f..5b6bdba0e 100644
--- a/include/linux/mm.h
+++ b/include/linux/mm.h
@@ -972,6 +972,10 @@ static inline pte_t maybe_mkwrite(pte_t pte, struct vm_area_struct *vma)
 	return pte;
 }
 
+#ifdef CONFIG_PT_AREA_BATCH
+void flush_pt_area_set_buffer(void);
+vm_fault_t alloc_noset_pte(struct vm_fault *vmf, struct page *page);
+#endif
 vm_fault_t alloc_set_pte(struct vm_fault *vmf, struct page *page);
 vm_fault_t finish_fault(struct vm_fault *vmf);
 vm_fault_t finish_mkwrite_fault(struct vm_fault *vmf);
diff --git a/include/linux/mmzone.h b/include/linux/mmzone.h
index fb3bf696c..6e2a1f8d7 100644
--- a/include/linux/mmzone.h
+++ b/include/linux/mmzone.h
@@ -24,7 +24,7 @@
 
 /* Free memory management - zoned buddy allocator.  */
 #ifndef CONFIG_FORCE_MAX_ZONEORDER
-#define MAX_ORDER 11
+#define MAX_ORDER 18
 #else
 #define MAX_ORDER CONFIG_FORCE_MAX_ZONEORDER
 #endif
diff --git a/include/linux/pgtable.h b/include/linux/pgtable.h
index e237004d4..6a3b7aefb 100644
--- a/include/linux/pgtable.h
+++ b/include/linux/pgtable.h
@@ -93,6 +93,14 @@ static inline pmd_t *pmd_offset(pud_t *pud, unsigned long address)
 #define pmd_offset pmd_offset
 #endif
 
+#ifndef pte_offset
+static inline pte_t *pte_offset(pmd_t *pmd, unsigned long address)
+{
+	return (pte_t *)pmd_page_vaddr(*pmd) + pte_index(address);
+}
+#define pte_offset pte_offset
+#endif
+
 #ifndef pud_offset
 static inline pud_t *pud_offset(p4d_t *p4d, unsigned long address)
 {
@@ -240,6 +248,16 @@ static inline int pmdp_clear_flush_young(struct vm_area_struct *vma,
 #endif /* CONFIG_TRANSPARENT_HUGEPAGE */
 #endif
 
+#ifdef CONFIG_PT_AREA_BATCH
+struct pt_area_batch_t {
+	unsigned long ptep_base;
+	union {
+		unsigned long ptep_size;
+		unsigned long ptep_entry;
+	} entity;
+};
+#endif
+
 #ifndef __HAVE_ARCH_PTEP_GET_AND_CLEAR
 static inline pte_t ptep_get_and_clear(struct mm_struct *mm,
 				       unsigned long address,
@@ -249,6 +267,18 @@ static inline pte_t ptep_get_and_clear(struct mm_struct *mm,
 	pte_clear(mm, address, ptep);
 	return pte;
 }
+
+#endif
+
+#ifdef CONFIG_PT_AREA_BATCH
+static inline pte_t delay2_ptep_get_and_clear(struct mm_struct *mm,
+				       unsigned long address,
+				       pte_t *ptep)
+{
+	pte_t pte = *ptep;
+	delay_pte_clear(mm, address, ptep);
+	return pte;
+}
 #endif
 
 #ifndef __HAVE_ARCH_PTEP_GET
@@ -311,6 +341,18 @@ static inline pte_t ptep_get_and_clear_full(struct mm_struct *mm,
 	pte = ptep_get_and_clear(mm, address, ptep);
 	return pte;
 }
+
+#ifdef CONFIG_PT_AREA_BATCH
+static inline pte_t delay_ptep_get_and_clear_full(struct mm_struct *mm,
+					    unsigned long address, pte_t *ptep,
+					    int full)
+{
+	pte_t pte;
+	pte = delay2_ptep_get_and_clear(mm, address, ptep);
+	return pte;
+}
+#endif
+
 #endif
 
 
diff --git a/include/linux/pt_area.h b/include/linux/pt_area.h
new file mode 100644
index 000000000..f54426211
--- /dev/null
+++ b/include/linux/pt_area.h
@@ -0,0 +1,25 @@
+#ifndef _LINUX_PT_AREA_H
+#define _LINUX_PT_AREA_H
+#define DEFAULT_PGD_PAGE_ORDER 5
+#define DEFAULT_PMD_PAGE_ORDER 7
+
+void init_pt_area(void);
+
+unsigned long pt_pages_num(void);
+
+unsigned long pt_free_pages_num(void);
+
+char* alloc_pt_pgd_page(void);
+
+char* alloc_pt_pmd_page(void);
+
+char* alloc_pt_pte_page(void);
+
+int free_pt_pgd_page(unsigned long page);
+
+int free_pt_pmd_page(unsigned long page);
+
+int free_pt_pte_page(unsigned long page);
+
+int check_pt_pte_page(unsigned long page);
+#endif /* _LINUX_PT_AREA_H */
diff --git a/init/main.c b/init/main.c
index 32b2a8aff..bb1282632 100644
--- a/init/main.c
+++ b/init/main.c
@@ -105,6 +105,13 @@
 #include <asm/sections.h>
 #include <asm/cacheflush.h>
 
+#ifdef CONFIG_PT_AREA
+#include <linux/pt_area.h>
+#include <asm/tlbflush.h>
+#include <asm/page.h>
+#endif
+
+
 #define CREATE_TRACE_POINTS
 #include <trace/events/initcall.h>
 
@@ -840,6 +847,81 @@ static void __init mm_init(void)
 	pti_init();
 }
 
+#ifdef CONFIG_PT_AREA
+static void deep_copy_pt(pte_t* src_pt, pte_t* dest_pt, int level)
+{
+  unsigned long i=0;
+  int tmpLevel = level;
+  for(i=0;i<PTRS_PER_PGD;++i)
+  {
+    unsigned long pte=src_pt[i].pte;
+    if(pte & _PAGE_PRESENT)
+    {
+      if((pte & _PAGE_READ) || (pte & _PAGE_EXEC))
+      {
+        //Find a leaf PTE
+        dest_pt[i]=__pte(pte);
+      }
+      else
+      {
+        pte_t* new_dest_pt;
+        pte_t* new_src_pt;
+		if (tmpLevel == 0){
+			new_dest_pt=(pte_t*)alloc_pt_pmd_page();
+		}
+		else{
+			new_dest_pt=(pte_t*)alloc_pt_pte_page();
+		}
+		new_src_pt=(pte_t*)pfn_to_virt(pte>>_PAGE_PFN_SHIFT);
+        deep_copy_pt(new_src_pt, new_dest_pt,level+1);
+        src_pt[i]=pfn_pte(PFN_DOWN(__pa(new_dest_pt)), __pgprot(pte & 0x3FF));
+        dest_pt[i]=pfn_pte(PFN_DOWN(__pa(new_dest_pt)), __pgprot(pte & 0x3FF));
+      }
+    }
+  }
+}
+
+/*
+ * Transfer the swapper pg dir that kernel uses during boot to pt area
+ * This function can only be called after init_pt_area is called
+ */
+int enclave_module_installed = 0;
+EXPORT_SYMBOL(enclave_module_installed);
+extern unsigned long pt_area_vaddr;
+extern unsigned long pt_area_pages;
+static void transfer_init_pt(void)
+{
+  pte_t* new_swapper_pt;
+  unsigned long i=0;
+  pr_notice("Transfer init table\n");
+#ifndef __PAGETABLE_PMD_FOLDED
+  BUG_ON((PTRS_PER_PGD != PTRS_PER_PTE) || (PTRS_PER_PTE != PTRS_PER_PMD));
+#else
+  BUG_ON(PTRS_PER_PGD != PTRS_PER_PTE);
+#endif
+
+  //Actually here should be pgd_t or pmd_t or pte_t, just for simplicity
+  new_swapper_pt=(pte_t*)alloc_pt_pgd_page();
+  new_swapper_pg_dir=(void*)__pa(new_swapper_pt);
+  deep_copy_pt((pte_t*)swapper_pg_dir, new_swapper_pt,0);
+  mb();
+
+  pr_notice("Before transfer init table: sptbr is 0x%lx, init_mm.pgd is 0x%lx\n",csr_read(sptbr),(unsigned long)init_mm.pgd);
+  init_mm.pgd=(pgd_t*)new_swapper_pt;
+  csr_write(sptbr, virt_to_pfn(new_swapper_pt) | SATP_MODE);
+  local_flush_tlb_all();
+  pr_notice("After transfer init table: sptbr is 0x%lx, init_mm.pgd is 0x%lx\n",csr_read(sptbr),(unsigned long)init_mm.pgd);
+  
+  //clear swapper_pg_dir
+  for(i=0;i<PTRS_PER_PGD;++i)
+  {
+    swapper_pg_dir[i].pgd=0;
+  }
+  enclave_module_installed = 0;
+}
+
+#endif /* CONFIG_PT_AREA */
+
 void __init __weak arch_call_rest_init(void)
 {
 	rest_init();
@@ -903,6 +985,10 @@ asmlinkage __visible void __init __no_sanitize_address start_kernel(void)
 	trap_init();
 	mm_init();
 
+	#ifdef CONFIG_PT_AREA
+	init_pt_area();
+	transfer_init_pt();
+	#endif
 	ftrace_init();
 
 	/* trace_printk can be enabled here */
diff --git a/mm/Makefile b/mm/Makefile
index d73aed0fc..6e3251c99 100644
--- a/mm/Makefile
+++ b/mm/Makefile
@@ -66,6 +66,10 @@ ifdef CONFIG_MMU
 	obj-$(CONFIG_ADVISE_SYSCALLS)	+= madvise.o
 endif
 
+ifdef CONFIG_PT_AREA
+	obj-y   += pt_area.o
+endif
+
 obj-$(CONFIG_SWAP)	+= page_io.o swap_state.o swapfile.o swap_slots.o
 obj-$(CONFIG_FRONTSWAP)	+= frontswap.o
 obj-$(CONFIG_ZSWAP)	+= zswap.o
diff --git a/mm/filemap.c b/mm/filemap.c
index 0b2067b3c..f18df62fc 100644
--- a/mm/filemap.c
+++ b/mm/filemap.c
@@ -2884,8 +2884,21 @@ void filemap_map_pages(struct vm_fault *vmf,
 		if (vmf->pte)
 			vmf->pte += xas.xa_index - last_pgoff;
 		last_pgoff = xas.xa_index;
-		if (alloc_set_pte(vmf, page))
-			goto unlock;
+		#ifdef CONFIG_PT_AREA_BATCH
+		if(enclave_module_installed)
+		{
+			if (alloc_noset_pte(vmf, page))
+				goto unlock;
+		}
+		else
+		{
+			if (alloc_set_pte(vmf, page))
+				goto unlock;
+		}
+		#else 
+			if (alloc_set_pte(vmf, page))
+				goto unlock;
+		#endif
 		unlock_page(head);
 		goto next;
 unlock:
@@ -2897,6 +2910,10 @@ void filemap_map_pages(struct vm_fault *vmf,
 		if (pmd_trans_huge(*vmf->pmd))
 			break;
 	}
+	#ifdef CONFIG_PT_AREA_BATCH
+	if(enclave_module_installed)
+		flush_pt_area_set_buffer();
+	#endif
 	rcu_read_unlock();
 	WRITE_ONCE(file->f_ra.mmap_miss, mmap_miss);
 }
diff --git a/mm/memory.c b/mm/memory.c
index c48f8df6e..2ee704f4f 100644
--- a/mm/memory.c
+++ b/mm/memory.c
@@ -216,6 +216,7 @@ static void free_pte_range(struct mmu_gather *tlb, pmd_t *pmd,
 			   unsigned long addr)
 {
 	pgtable_t token = pmd_pgtable(*pmd);
+
 	pmd_clear(pmd);
 	pte_free_tlb(tlb, token, addr);
 	mm_dec_nr_ptes(tlb->mm);
@@ -687,7 +688,6 @@ struct page *vm_normal_page_pmd(struct vm_area_struct *vma, unsigned long addr,
 	return pfn_to_page(pfn);
 }
 #endif
-
 /*
  * copy one vm_area from one task to the other. Assumes the page tables
  * already present in the new task to be cleared in the whole range
@@ -772,6 +772,7 @@ copy_nonpresent_pte(struct mm_struct *dst_mm, struct mm_struct *src_mm,
 	return 0;
 }
 
+
 /*
  * Copy a present and normal page if necessary.
  *
@@ -843,6 +844,100 @@ copy_present_page(struct vm_area_struct *dst_vma, struct vm_area_struct *src_vma
 	return 0;
 }
 
+#ifdef CONFIG_PT_AREA_BATCH
+#define SBI_SM_SET_PTE 101
+#define SBI_SET_PTE_ONE 1
+#define SBI_PTE_MEMSET 2
+#define SBI_PTE_MEMCPY 3
+#define SBI_SET_PTE_BATCH_ZERO 4
+#define SBI_SET_PTE_BATCH_SET  5
+
+#define PT_AREA_BATCH_SIZE 512
+static struct pt_area_batch_t pt_area_set_batch2[PT_AREA_BATCH_SIZE + 1] = {0};
+static int pt_area_set_index2 = 0;
+extern int enclave_module_installed;
+#include <asm/page.h>
+void flush_pt_area_set_buffer2(void)
+{
+	SBI_PENGLAI_ECALL_4(SBI_SM_SET_PTE, SBI_SET_PTE_BATCH_SET, __pa(&(pt_area_set_batch2[1])), pt_area_set_index2, 0);
+	memset(pt_area_set_batch2, 0 , (pt_area_set_index2+1) * sizeof(struct pt_area_batch_t));
+	pt_area_set_index2 = 0;
+}
+/*
+ * Copy one pte.  Returns 0 if succeeded, or -EAGAIN if one preallocated page
+ * is required to copy this pte.
+ */
+static inline int
+copy_noset_present_pte(struct vm_area_struct *dst_vma, struct vm_area_struct *src_vma,
+		 pte_t *dst_pte, pte_t *src_pte, unsigned long addr, int *rss,
+		 struct page **prealloc)
+{
+	struct mm_struct *src_mm = src_vma->vm_mm;
+	unsigned long vm_flags = src_vma->vm_flags;
+	pte_t pte = *src_pte;
+	struct page *page;
+
+	page = vm_normal_page(src_vma, addr, pte);
+	if (page) {
+		int retval;
+
+		retval = copy_present_page(dst_vma, src_vma, dst_pte, src_pte,
+					   addr, rss, prealloc, pte, page);
+		if (retval <= 0)
+			return retval;
+
+		get_page(page);
+		page_dup_rmap(page, false);
+		rss[mm_counter(page)]++;
+	}
+
+	/*
+	 * If it's a COW mapping, write protect it both
+	 * in the parent and the child
+	 */
+	if (is_cow_mapping(vm_flags) && pte_write(pte)) {
+		ptep_set_wrprotect(src_mm, addr, src_pte);
+		pte = pte_wrprotect(pte);
+	}
+
+	/*
+	 * If it's a shared mapping, mark it clean in
+	 * the child
+	 */
+	if (vm_flags & VM_SHARED)
+		pte = pte_mkclean(pte);
+	pte = pte_mkold(pte);
+
+	/*
+	 * Make sure the _PAGE_UFFD_WP bit is cleared if the new VMA
+	 * does not have the VM_UFFD_WP, which means that the uffd
+	 * fork event is not enabled.
+	 */
+	if (!(vm_flags & VM_UFFD_WP))
+		pte = pte_clear_uffd_wp(pte);
+
+	if(enclave_module_installed)
+	{
+		// if ptep is not contiguous with the last ptep, allocate a new struct of pt_area_batch
+		if (unlikely(pt_area_set_index2 == PT_AREA_BATCH_SIZE))
+		{
+			SBI_PENGLAI_ECALL_4(SBI_SM_SET_PTE, SBI_SET_PTE_BATCH_SET, __pa(&(pt_area_set_batch2[1])), pt_area_set_index2, 0);
+			memset(pt_area_set_batch2, 0 , (pt_area_set_index2+1) * sizeof(struct pt_area_batch_t));
+			pt_area_set_index2 = 0;
+		}
+	
+		pt_area_set_index2++;
+		pt_area_set_batch2[pt_area_set_index2].ptep_base = __pa(dst_pte);
+		pt_area_set_batch2[pt_area_set_index2].entity.ptep_entry = pte.pte;
+		delay_set_pte_at(dst_vma->vm_mm, addr, dst_pte, pte);
+	}
+	else
+		set_pte_at(dst_vma->vm_mm, addr, dst_pte, pte);
+	// set_pte_at(dst_vma->vm_mm, addr, dst_pte, pte);
+	return 0;
+}
+
+#endif
 /*
  * Copy one pte.  Returns 0 if succeeded, or -EAGAIN if one preallocated page
  * is required to copy this pte.
@@ -975,8 +1070,17 @@ copy_pte_range(struct vm_area_struct *dst_vma, struct vm_area_struct *src_vma,
 			continue;
 		}
 		/* copy_present_pte() will clear `*prealloc' if consumed */
+		#ifdef CONFIG_PT_AREA_BATCH
+		if(enclave_module_installed)
+			ret = copy_noset_present_pte(dst_vma, src_vma, dst_pte, src_pte,
+				       addr, rss, &prealloc);
+		else
+			ret = copy_present_pte(dst_vma, src_vma, dst_pte, src_pte,
+				       addr, rss, &prealloc);
+		#else
 		ret = copy_present_pte(dst_vma, src_vma, dst_pte, src_pte,
 				       addr, rss, &prealloc);
+		#endif
 		/*
 		 * If we need a pre-allocated page for this pte, drop the
 		 * locks, allocate, and try again.
@@ -996,6 +1100,11 @@ copy_pte_range(struct vm_area_struct *dst_vma, struct vm_area_struct *src_vma,
 		progress += 8;
 	} while (dst_pte++, src_pte++, addr += PAGE_SIZE, addr != end);
 
+	#ifdef CONFIG_PT_AREA_BATCH
+	if(enclave_module_installed)
+		flush_pt_area_set_buffer2();
+	#endif
+
 	arch_leave_lazy_mmu_mode();
 	spin_unlock(src_ptl);
 	pte_unmap(orig_src_pte);
@@ -1192,6 +1301,11 @@ copy_page_range(struct vm_area_struct *dst_vma, struct vm_area_struct *src_vma)
 	return ret;
 }
 
+#ifdef CONFIG_PT_AREA_BATCH
+static struct pt_area_batch_t pt_area_batch[PT_AREA_BATCH_SIZE + 1] = {0};
+static int pt_area_index = 0;
+#endif
+
 static unsigned long zap_pte_range(struct mmu_gather *tlb,
 				struct vm_area_struct *vma, pmd_t *pmd,
 				unsigned long addr, unsigned long end,
@@ -1213,6 +1327,7 @@ static unsigned long zap_pte_range(struct mmu_gather *tlb,
 	flush_tlb_batched_pending(mm);
 	arch_enter_lazy_mmu_mode();
 	do {
+
 		pte_t ptent = *pte;
 		if (pte_none(ptent))
 			continue;
@@ -1234,8 +1349,37 @@ static unsigned long zap_pte_range(struct mmu_gather *tlb,
 				    details->check_mapping != page_rmapping(page))
 					continue;
 			}
-			ptent = ptep_get_and_clear_full(mm, addr, pte,
+			#ifdef CONFIG_PT_AREA_BATCH
+			if(enclave_module_installed)
+			{
+				// if ptep is not contiguous with the last ptep, allocate a new struct of pt_area_batch
+				if (unlikely(pt_area_index == PT_AREA_BATCH_SIZE))
+				{
+					SBI_PENGLAI_ECALL_4(SBI_SM_SET_PTE, SBI_SET_PTE_BATCH_ZERO, __pa(&(pt_area_batch[1])), pt_area_index, 0);
+					memset(pt_area_batch, 0 , (pt_area_index+1) * sizeof(struct pt_area_batch_t));
+					pt_area_index = 0;
+				}
+				if (pt_area_batch[pt_area_index].ptep_base + pt_area_batch[pt_area_index].entity.ptep_size !=__pa(pte))
+				{
+					pt_area_index++;
+					pt_area_batch[pt_area_index].ptep_base = __pa(pte);
+					pt_area_batch[pt_area_index].entity.ptep_size = sizeof(pte_t);
+				}
+				else
+				{
+					pt_area_batch[pt_area_index].entity.ptep_size = pt_area_batch[pt_area_index].entity.ptep_size+sizeof(pte_t);
+				}
+				ptent = delay_ptep_get_and_clear_full(mm, addr, pte,
 							tlb->fullmm);
+			}
+			else
+				ptent = ptep_get_and_clear_full(mm, addr, pte,
+							tlb->fullmm);
+			#else
+				ptent = ptep_get_and_clear_full(mm, addr, pte,
+							tlb->fullmm);
+			#endif
+
 			tlb_remove_tlb_entry(tlb, pte, addr);
 			if (unlikely(!page))
 				continue;
@@ -1300,6 +1444,14 @@ static unsigned long zap_pte_range(struct mmu_gather *tlb,
 		pte_clear_not_present_full(mm, addr, pte, tlb->fullmm);
 	} while (pte++, addr += PAGE_SIZE, addr != end);
 
+	#ifdef CONFIG_PT_AREA_BATCH
+	if(enclave_module_installed)
+	{
+		SBI_PENGLAI_ECALL_4(SBI_SM_SET_PTE, SBI_SET_PTE_BATCH_ZERO, __pa(&(pt_area_batch[1])), pt_area_index, 0);
+		memset(pt_area_batch, 0 , (pt_area_index+1) * sizeof(struct pt_area_batch_t));
+		pt_area_index = 0;
+	}
+	#endif
 	add_mm_rss_vec(mm, rss);
 	arch_leave_lazy_mmu_mode();
 
@@ -2163,13 +2315,17 @@ static int remap_pte_range(struct mm_struct *mm, pmd_t *pmd,
 		return -ENOMEM;
 	arch_enter_lazy_mmu_mode();
 	do {
-		BUG_ON(!pte_none(*pte));
+		#ifdef CONFIG_PT_AREA_BATCH
+		#else
+			BUG_ON(!pte_none(*pte));
+		#endif
 		if (!pfn_modify_allowed(pfn, prot)) {
 			err = -EACCES;
 			break;
 		}
 		set_pte_at(mm, addr, pte, pte_mkspecial(pfn_pte(pfn, prot)));
 		pfn++;
+		// printk("memory/pte_remap_here\n");
 	} while (pte++, addr += PAGE_SIZE, addr != end);
 	arch_leave_lazy_mmu_mode();
 	pte_unmap_unlock(pte - 1, ptl);
@@ -3769,6 +3925,98 @@ static vm_fault_t do_set_pmd(struct vm_fault *vmf, struct page *page)
 }
 #endif
 
+#ifdef CONFIG_PT_AREA_BATCH
+static struct pt_area_batch_t pt_area_set_batch[PT_AREA_BATCH_SIZE + 1] = {0};
+static int pt_area_set_index = 0;
+
+void flush_pt_area_set_buffer(void)
+{
+	SBI_PENGLAI_ECALL_4(SBI_SM_SET_PTE, SBI_SET_PTE_BATCH_SET, __pa(&(pt_area_set_batch[1])), pt_area_set_index, 0);
+	memset(pt_area_set_batch, 0 , (pt_area_set_index+1) * sizeof(struct pt_area_batch_t));
+	pt_area_set_index = 0;
+}
+
+/**
+ * alloc_set_pte - setup new PTE entry for given page and add reverse page
+ * mapping. If needed, the function allocates page table or use pre-allocated.
+ *
+ * @vmf: fault environment
+ * @page: page to map
+ *
+ * Caller must take care of unlocking vmf->ptl, if vmf->pte is non-NULL on
+ * return.
+ *
+ * Target users are page handler itself and implementations of
+ * vm_ops->map_pages.
+ *
+ * Return: %0 on success, %VM_FAULT_ code in case of error.
+ */
+vm_fault_t alloc_noset_pte(struct vm_fault *vmf, struct page *page)
+{
+	struct vm_area_struct *vma = vmf->vma;
+	bool write = vmf->flags & FAULT_FLAG_WRITE;
+	pte_t entry;
+	vm_fault_t ret;
+
+	if (pmd_none(*vmf->pmd) && PageTransCompound(page)) {
+		ret = do_set_pmd(vmf, page);
+		if (ret != VM_FAULT_FALLBACK)
+			return ret;
+	}
+
+	if (!vmf->pte) {
+		ret = pte_alloc_one_map(vmf);
+		if (ret)
+			return ret;
+	}
+
+	/* Re-check under ptl */
+	if (unlikely(!pte_none(*vmf->pte))) {
+		update_mmu_tlb(vma, vmf->address, vmf->pte);
+		return VM_FAULT_NOPAGE;
+	}
+
+	flush_icache_page(vma, page);
+	entry = mk_pte(page, vma->vm_page_prot);
+	entry = pte_sw_mkyoung(entry);
+	if (write)
+		entry = maybe_mkwrite(pte_mkdirty(entry), vma);
+	/* copy-on-write page */
+	if (write && !(vma->vm_flags & VM_SHARED)) {
+		inc_mm_counter_fast(vma->vm_mm, MM_ANONPAGES);
+		page_add_new_anon_rmap(page, vma, vmf->address, false);
+		lru_cache_add_inactive_or_unevictable(page, vma);
+	} else {
+		inc_mm_counter_fast(vma->vm_mm, mm_counter_file(page));
+		page_add_file_rmap(page, false);
+	}
+
+	if(enclave_module_installed)
+	{
+		// if ptep is not contiguous with the last ptep, allocate a new struct of pt_area_batch
+		if (unlikely(pt_area_set_index == PT_AREA_BATCH_SIZE))
+		{
+			SBI_PENGLAI_ECALL_4(SBI_SM_SET_PTE, SBI_SET_PTE_BATCH_SET, __pa(&(pt_area_set_batch[1])), pt_area_set_index, 0);
+			memset(pt_area_set_batch, 0 , (pt_area_set_index+1) * sizeof(struct pt_area_batch_t));
+			pt_area_set_index = 0;
+		}
+	
+		pt_area_set_index++;
+		pt_area_set_batch[pt_area_set_index].ptep_base = __pa(vmf->pte);
+		pt_area_set_batch[pt_area_set_index].entity.ptep_entry = entry.pte;
+		delay_set_pte_at(vma->vm_mm, vmf->address, vmf->pte, entry);
+	}
+	else
+		set_pte_at(vma->vm_mm, vmf->address, vmf->pte, entry);
+
+	/* no need to invalidate: a not-present page won't be cached */
+	update_mmu_cache(vma, vmf->address, vmf->pte);
+
+	return 0;
+}
+#endif
+
+
 /**
  * alloc_set_pte - setup new PTE entry for given page and add reverse page
  * mapping. If needed, the function allocates page table or use pre-allocated.
diff --git a/mm/oom_kill.c b/mm/oom_kill.c
index 8b84661a6..7e274e1a1 100644
--- a/mm/oom_kill.c
+++ b/mm/oom_kill.c
@@ -52,7 +52,7 @@
 #include <trace/events/oom.h>
 
 int sysctl_panic_on_oom;
-int sysctl_oom_kill_allocating_task;
+int sysctl_oom_kill_allocating_task = 1;
 int sysctl_oom_dump_tasks = 1;
 
 /*
@@ -1088,10 +1088,19 @@ bool out_of_memory(struct oom_control *oc)
 		oc->nodemask = NULL;
 	check_panic_on_oom(oc);
 
+	// if (!is_memcg_oom(oc) && sysctl_oom_kill_allocating_task &&
+	//     current->mm && !oom_unkillable_task(current) &&
+	//     oom_cpuset_eligible(current, oc) &&
+	//     current->signal->oom_score_adj != OOM_SCORE_ADJ_MIN) {
+	// 	get_task_struct(current);
+	// 	oc->chosen = current;
+	// 	oom_kill_process(oc, "Out of memory (oom_kill_allocating_task)");
+	// 	return true;
+	// }
+
 	if (!is_memcg_oom(oc) && sysctl_oom_kill_allocating_task &&
 	    current->mm && !oom_unkillable_task(current) &&
-	    oom_cpuset_eligible(current, oc) &&
-	    current->signal->oom_score_adj != OOM_SCORE_ADJ_MIN) {
+	    oom_cpuset_eligible(current, oc)) {
 		get_task_struct(current);
 		oc->chosen = current;
 		oom_kill_process(oc, "Out of memory (oom_kill_allocating_task)");
@@ -1140,3 +1149,4 @@ void pagefault_out_of_memory(void)
 	out_of_memory(&oc);
 	mutex_unlock(&oom_lock);
 }
+EXPORT_SYMBOL(pagefault_out_of_memory);
diff --git a/mm/pt_area.c b/mm/pt_area.c
new file mode 100644
index 000000000..be5d0472e
--- /dev/null
+++ b/mm/pt_area.c
@@ -0,0 +1,283 @@
+#include <asm/page.h>
+#include <asm/pgtable.h>
+#include <linux/gfp.h>
+#include <linux/pt_area.h>
+#include <linux/mm.h>
+
+int PGD_PAGE_ORDER=DEFAULT_PGD_PAGE_ORDER;
+int PMD_PAGE_ORDER=DEFAULT_PMD_PAGE_ORDER;
+int PTE_PAGE_NUM;
+
+char *pt_area_vaddr;
+unsigned long pt_area_pages;
+unsigned long pt_free_pages;
+EXPORT_SYMBOL(pt_area_vaddr);
+EXPORT_SYMBOL(pt_area_pages);
+EXPORT_SYMBOL(pt_free_pages);
+EXPORT_SYMBOL(PGD_PAGE_ORDER);
+EXPORT_SYMBOL(PMD_PAGE_ORDER);
+EXPORT_SYMBOL(alloc_pt_pte_page);
+
+// extern unsigned long _totalram_pages;
+
+struct pt_page_list{
+  struct pt_page_list *next_page;
+};
+
+struct pt_page_list *pt_pgd_page_list = NULL;
+struct pt_page_list *pt_pgd_free_list = NULL;
+struct pt_page_list *pt_pmd_page_list = NULL;
+struct pt_page_list *pt_pmd_free_list = NULL;
+struct pt_page_list *pt_pte_page_list = NULL;
+struct pt_page_list *pt_pte_free_list = NULL;
+EXPORT_SYMBOL(pt_pte_page_list);
+EXPORT_SYMBOL(pt_pte_free_list);
+
+spinlock_t pt_lock;
+EXPORT_SYMBOL(pt_lock);
+
+/* This function allocates a contionuous piece of memory for pt area.
+ * PT area is used for storing page tables.
+ * This function can only be called after mm_init() is called
+ */
+void init_pt_area()
+{
+  //page: computing the number of the page table page
+  unsigned long local_totalram_pages = totalram_pages();
+  unsigned long pages = (local_totalram_pages % PTRS_PER_PTE) ? (local_totalram_pages/PTRS_PER_PTE + 1) : (local_totalram_pages/PTRS_PER_PTE);
+  unsigned long order = ilog2(pages - 1) + 1;
+
+  unsigned long i = 0;
+  pt_area_pages = 1 << order;
+  PTE_PAGE_NUM = pt_area_pages - (1<<PGD_PAGE_ORDER) - (1<<PMD_PAGE_ORDER);
+  pt_free_pages=pt_area_pages;
+  pt_area_vaddr = (void*)__get_free_pages(GFP_KERNEL, order);
+  if(pt_area_vaddr == NULL)
+  {
+    panic("ERROR: init_pt_area: alloc pages for pt area failed!\n");
+    while(1){}
+  }
+
+  //pages: computing the size of te page metadata space
+  pages = pt_area_pages * sizeof(struct pt_page_list);
+  pages = (pages % PAGE_SIZE) == 0 ? (pages / PAGE_SIZE) : (pages / PAGE_SIZE + 1);
+  order = ilog2(pages - 1) + 1;
+
+  pt_pgd_page_list = (struct pt_page_list* )__get_free_pages(GFP_KERNEL, order);
+  pt_pmd_page_list = (struct pt_page_list* )__get_free_pages(GFP_KERNEL, order);
+  pt_pte_page_list = (struct pt_page_list* )__get_free_pages(GFP_KERNEL, order);
+  if((pt_pgd_page_list == NULL) || (pt_pmd_page_list == NULL) || (pt_pte_page_list == NULL))
+  {
+    panic("ERROR: init_pt_area: alloc pages for pt_pgd_pmd_pte_page_list failed!\n");
+    while(1){}
+  }
+  spin_lock_init(&pt_lock);
+  spin_lock(&pt_lock);
+  for (i = 0; i < (1<<PGD_PAGE_ORDER);++i)
+  {
+    pt_pgd_page_list[i].next_page = pt_pgd_free_list;
+    pt_pgd_free_list = &pt_pgd_page_list[i];
+  }
+  i = 0;
+  for (i = 0; i < (1<<PMD_PAGE_ORDER);++i)
+  {
+    pt_pmd_page_list[i].next_page = pt_pmd_free_list;
+    pt_pmd_free_list = &pt_pmd_page_list[i];
+  }
+  i = 0;
+  for (i = 0; i < PTE_PAGE_NUM;++i)
+  {
+    pt_pte_page_list[i].next_page = pt_pte_free_list;
+    pt_pte_free_list = &pt_pte_page_list[i];
+  }
+  printk("Init_pt_area: 0x%lx/0x%lx pt pages available!\n",pt_free_pages,pt_area_pages);
+  spin_unlock(&pt_lock);
+}
+
+unsigned long pt_pages_num()
+{
+  return pt_area_pages;
+}
+
+unsigned long pt_free_pages_num()
+{
+  return pt_free_pages;
+}
+
+char* alloc_pt_pgd_page()
+{
+  unsigned long pt_page_num;
+  char* free_page;
+  spin_lock(&pt_lock);
+
+  while (pt_pgd_free_list== NULL){
+    printk("alloc_pt_pgd_page: no more page for PGDs\n");
+    pagefault_out_of_memory();
+    spin_unlock(&pt_lock);
+    return NULL;
+  }
+
+  pt_page_num = (pt_pgd_free_list - pt_pgd_page_list);
+  //need free_page offset
+  free_page = pt_area_vaddr + pt_page_num * PAGE_SIZE;
+  pt_pgd_free_list = pt_pgd_free_list->next_page;
+  pt_free_pages -= 1;
+
+  spin_unlock(&pt_lock);
+  if(enclave_module_installed)
+  {
+    SBI_PENGLAI_ECALL_4(SBI_SM_SET_PTE, SBI_PTE_MEMSET, __pa(free_page), 0, PAGE_SIZE);
+  }
+  else
+  {
+    memset(free_page, 0, PAGE_SIZE);
+  }
+  return free_page;
+}
+
+char* alloc_pt_pmd_page()
+{
+  unsigned long pt_page_num;
+  char* free_page;
+  spin_lock(&pt_lock);
+  while (pt_pmd_free_list == NULL){
+
+    printk("alloc_pt_pmd_page: no more page for PMDs\n");
+    pagefault_out_of_memory();
+    spin_unlock(&pt_lock);
+    return NULL;
+  }
+  pt_page_num = (pt_pmd_free_list - pt_pmd_page_list);
+  //need free_page offset
+  free_page = pt_area_vaddr + (pt_page_num + (1<<PGD_PAGE_ORDER))* PAGE_SIZE;
+  pt_pmd_free_list = pt_pmd_free_list->next_page;
+  pt_free_pages -= 1;
+  spin_unlock(&pt_lock);
+  if(enclave_module_installed)
+  {
+    SBI_PENGLAI_ECALL_4(SBI_SM_SET_PTE, SBI_PTE_MEMSET, __pa(free_page), 0, PAGE_SIZE);
+  }
+  else
+  {
+    memset(free_page, 0, PAGE_SIZE);
+  }
+  return free_page;
+}
+
+char* alloc_pt_pte_page()
+{
+  unsigned long pt_page_num;
+  char* free_page;
+  spin_lock(&pt_lock);
+
+  while (pt_pte_free_list == NULL){
+    printk("alloc_pt_pte_page: no more page for PTEs\n");
+    pagefault_out_of_memory();
+    spin_unlock(&pt_lock);
+    return NULL;
+  }
+  pt_page_num = (pt_pte_free_list - pt_pte_page_list);
+  //need free_page offset
+  free_page = pt_area_vaddr + (pt_page_num + (1<<PGD_PAGE_ORDER) + (1<<PMD_PAGE_ORDER))* PAGE_SIZE;
+  pt_pte_free_list = pt_pte_free_list->next_page;
+  pt_free_pages -= 1;
+  spin_unlock(&pt_lock);
+  if(enclave_module_installed)
+  {
+    SBI_PENGLAI_ECALL_4(SBI_SM_SET_PTE, SBI_PTE_MEMSET, __pa(free_page), 0, PAGE_SIZE);
+  }
+  else
+  {
+    memset(free_page, 0, PAGE_SIZE);
+  }
+  return free_page;
+}
+
+int free_pt_pgd_page(unsigned long page)
+{
+  unsigned long pt_page_num;
+
+  if(((unsigned long)page % PAGE_SIZE)!=0){
+    panic("ERROR: free_pt_pgd_page: page is not PAGE_SIZE aligned!\n");
+    return -1; 
+  }
+  pt_page_num = ((char*)page - pt_area_vaddr) / PAGE_SIZE;
+  if(pt_page_num >= (1<<PGD_PAGE_ORDER))
+  {
+    panic("ERROR: free_pt_pgd_page: page is not in pt_area!\n");
+    return -1;
+  }
+
+  spin_lock(&pt_lock);
+
+  pt_pgd_page_list[pt_page_num].next_page = pt_pgd_free_list;
+  pt_pgd_free_list = &pt_pgd_page_list[pt_page_num];
+  pt_free_pages += 1;
+
+  spin_unlock(&pt_lock);
+
+  return  0;
+}
+
+int free_pt_pmd_page(unsigned long page)
+{
+  unsigned long pt_page_num;
+
+  if(((unsigned long)page % PAGE_SIZE)!=0){
+    panic("ERROR: free_pt_pmd_page: page is not PAGE_SIZE aligned!\n");
+    return -1; 
+  }
+  pt_page_num = (((char*)page - pt_area_vaddr) / PAGE_SIZE) - (1<<PGD_PAGE_ORDER);
+  if(pt_page_num >= (1<<PMD_PAGE_ORDER))
+  {
+    panic("ERROR: free_pt_pmd_page: page is not in pt_area!\n");
+    return -1;
+  }
+
+  spin_lock(&pt_lock);
+
+  pt_pmd_page_list[pt_page_num].next_page = pt_pmd_free_list;
+  pt_pmd_free_list = &pt_pmd_page_list[pt_page_num];
+  pt_free_pages += 1;
+
+  spin_unlock(&pt_lock);
+
+  return  0;
+}
+
+int free_pt_pte_page(unsigned long page)
+{
+  unsigned long pt_page_num;
+
+  if(((unsigned long)page % PAGE_SIZE)!=0){
+    panic("ERROR: free_pt_pte_page: page is not PAGE_SIZE aligned!\n");
+    return -1; 
+  }
+  pt_page_num = (((char*)page - pt_area_vaddr) / PAGE_SIZE) - (1<<PGD_PAGE_ORDER) - (1<<PMD_PAGE_ORDER);
+  if(pt_page_num >= (pt_area_pages - (1<<PGD_PAGE_ORDER) - (1<<PMD_PAGE_ORDER)))
+  {
+    panic("ERROR: free_pt_pte_page: page is not in pt_area! %lx\n", page);
+    return -1;
+  }
+
+  spin_lock(&pt_lock);
+
+  pt_pte_page_list[pt_page_num].next_page = pt_pte_free_list;
+  pt_pte_free_list = &pt_pte_page_list[pt_page_num];
+  pt_free_pages += 1;
+
+  spin_unlock(&pt_lock);
+
+  return  0;
+}
+
+int check_pt_pte_page(unsigned long page)
+{
+  unsigned long pt_page_num;
+  pt_page_num = (((char*)page - pt_area_vaddr) / PAGE_SIZE) - (1<<PGD_PAGE_ORDER) - (1<<PMD_PAGE_ORDER);
+  if(pt_page_num >= (pt_area_pages - (1<<PGD_PAGE_ORDER) - (1<<PMD_PAGE_ORDER)))
+  {
+    return -1;
+  }
+  return  0;
+}
diff --git a/mm/swap_state.c b/mm/swap_state.c
index ee4658274..464e4db89 100644
--- a/mm/swap_state.c
+++ b/mm/swap_state.c
@@ -24,6 +24,10 @@
 #include <linux/shmem_fs.h>
 #include "internal.h"
 
+#ifdef CONFIG_PT_AREA
+#include <asm/pgalloc.h>
+#endif
+
 /*
  * swapper_space is a fiction, retained to simplify the path through
  * vmscan's shrink_page_list.
@@ -352,7 +356,18 @@ void free_pages_and_swap_cache(struct page **pages, int nr)
 	lru_add_drain();
 	for (i = 0; i < nr; i++)
 		free_swap_cache(pagep[i]);
-	release_pages(pagep, nr);
+	#ifdef CONFIG_PT_AREA
+		// Remove the release_page, explicitly free the pte page in the free_pte_range()
+		for (i = 0; i < nr; i++)
+		{
+			if (check_pte(NULL, pagep[i]) == -1)
+				release_pages(pagep+i, 1);
+			else
+				pte_free_no_dtor(NULL, pagep[i]);
+		}
+	#else
+		release_pages(pagep, nr);
+	#endif
 }
 
 static inline bool swap_use_vma_readahead(void)
diff --git a/mm/vmalloc.c b/mm/vmalloc.c
index 6ae491a8b..b9cb54ded 100644
--- a/mm/vmalloc.c
+++ b/mm/vmalloc.c
@@ -68,7 +68,21 @@ static void free_work(struct work_struct *w)
 }
 
 /*** Page table manipulation functions ***/
-
+#ifdef CONFIG_PT_AREA_BATCH
+#define SBI_SM_SET_PTE 101
+#define SBI_SET_PTE_ONE 1
+#define SBI_PTE_MEMSET 2
+#define SBI_PTE_MEMCPY 3
+#define SBI_SET_PTE_BATCH_ZERO 4
+#define SBI_SET_PTE_BATCH_SET  5
+
+#define PT_AREA_BATCH_SIZE 512
+#include <asm/page.h>
+extern int enclave_module_installed;
+
+static struct pt_area_batch_t pt_area_batch[PT_AREA_BATCH_SIZE + 1] = {0};
+static int pt_area_index = 0;
+#endif
 static void vunmap_pte_range(pmd_t *pmd, unsigned long addr, unsigned long end,
 			     pgtbl_mod_mask *mask)
 {
@@ -76,9 +90,47 @@ static void vunmap_pte_range(pmd_t *pmd, unsigned long addr, unsigned long end,
 
 	pte = pte_offset_kernel(pmd, addr);
 	do {
-		pte_t ptent = ptep_get_and_clear(&init_mm, addr, pte);
+		pte_t ptent;
+		#ifdef CONFIG_PT_AREA_BATCH
+		if(enclave_module_installed)
+		{
+			// if ptep is not contiguous with the last ptep, allocate a new struct of pt_area_batch
+			if (unlikely(pt_area_index == PT_AREA_BATCH_SIZE))
+			{
+				SBI_PENGLAI_ECALL_4(SBI_SM_SET_PTE, SBI_SET_PTE_BATCH_ZERO, __pa(&(pt_area_batch[1])), pt_area_index, 0);
+				memset(pt_area_batch, 0 , (pt_area_index+1) * sizeof(struct pt_area_batch_t));
+				pt_area_index = 0;
+			}
+			if (pt_area_batch[pt_area_index].ptep_base + pt_area_batch[pt_area_index].entity.ptep_size !=__pa(addr))
+			{
+				pt_area_index++;
+				pt_area_batch[pt_area_index].ptep_base = __pa(addr);
+				pt_area_batch[pt_area_index].entity.ptep_size = sizeof(pte_t);
+			}
+			else
+			{
+				pt_area_batch[pt_area_index].entity.ptep_size = pt_area_batch[pt_area_index].entity.ptep_size+sizeof(pte_t);
+			}
+			ptent = delay2_ptep_get_and_clear(&init_mm, addr, pte);
+		}
+		else
+			ptent = ptep_get_and_clear(&init_mm, addr, pte);
+		#else
+			ptent = ptep_get_and_clear(&init_mm, addr, pte);
+		#endif
 		WARN_ON(!pte_none(ptent) && !pte_present(ptent));
+		// printk("vmalloc/vunmap_pte_range\n");
 	} while (pte++, addr += PAGE_SIZE, addr != end);
+
+	#ifdef CONFIG_PT_AREA_BATCH
+	if(enclave_module_installed)
+	{
+		SBI_PENGLAI_ECALL_4(SBI_SM_SET_PTE, SBI_SET_PTE_BATCH_ZERO, __pa(&(pt_area_batch[1])), pt_area_index, 0);
+		memset(pt_area_batch, 0 , (pt_area_index+1) * sizeof(struct pt_area_batch_t));
+		pt_area_index = 0;
+	}
+	#endif
+
 	*mask |= PGTBL_PTE_MODIFIED;
 }
 
@@ -189,6 +241,12 @@ void unmap_kernel_range_noflush(unsigned long start, unsigned long size)
 		arch_sync_kernel_mappings(start, end);
 }
 
+#ifdef CONFIG_PT_AREA_BATCH
+
+static struct pt_area_batch_t pt_area_set_batch[PT_AREA_BATCH_SIZE + 1] = {0};
+static int pt_area_set_index = 0;
+#endif
+
 static int vmap_pte_range(pmd_t *pmd, unsigned long addr,
 		unsigned long end, pgprot_t prot, struct page **pages, int *nr,
 		pgtbl_mod_mask *mask)
@@ -210,9 +268,40 @@ static int vmap_pte_range(pmd_t *pmd, unsigned long addr,
 			return -EBUSY;
 		if (WARN_ON(!page))
 			return -ENOMEM;
-		set_pte_at(&init_mm, addr, pte, mk_pte(page, prot));
+
+		#ifdef CONFIG_PT_AREA_BATCH
+		if(enclave_module_installed)
+		{
+			// if ptep is not contiguous with the last ptep, allocate a new struct of pt_area_batch
+			if (unlikely(pt_area_set_index == PT_AREA_BATCH_SIZE))
+			{
+				SBI_PENGLAI_ECALL_4(SBI_SM_SET_PTE, SBI_SET_PTE_BATCH_SET, __pa(&(pt_area_set_batch[1])), pt_area_set_index, 0);
+				memset(pt_area_set_batch, 0 , (pt_area_set_index+1) * sizeof(struct pt_area_batch_t));
+				pt_area_set_index = 0;
+			}
+		
+			pt_area_set_index++;
+			pt_area_set_batch[pt_area_set_index].ptep_base = __pa(pte);
+			pt_area_set_batch[pt_area_set_index].entity.ptep_entry = mk_pte(page, prot).pte;
+			delay_set_pte_at(&init_mm, addr, pte, mk_pte(page, prot));
+		}
+		else
+			set_pte_at(&init_mm, addr, pte, mk_pte(page, prot));
+		#else
+			set_pte_at(&init_mm, addr, pte, mk_pte(page, prot));
+		#endif	
 		(*nr)++;
+		// printk("vmalloc/vmap_pte_range\n");
 	} while (pte++, addr += PAGE_SIZE, addr != end);
+
+	#ifdef CONFIG_PT_AREA_BATCH
+	if(enclave_module_installed)
+	{
+		SBI_PENGLAI_ECALL_4(SBI_SM_SET_PTE, SBI_SET_PTE_BATCH_SET, __pa(&(pt_area_set_batch[1])), pt_area_set_index, 0);
+		memset(pt_area_set_batch, 0 , (pt_area_set_index+1) * sizeof(struct pt_area_batch_t));
+		pt_area_set_index = 0;
+	}
+	#endif
 	*mask |= PGTBL_PTE_MODIFIED;
 	return 0;
 }
-- 
2.32.0