diff --git a/arch/x86/include/asm/sgx.h b/arch/x86/include/asm/sgx.h
index fe1a2ba..063ff8f 100644
--- a/arch/x86/include/asm/sgx.h
+++ b/arch/x86/include/asm/sgx.h
@@ -372,4 +372,16 @@ unsigned long sgx_epc_page_to_pfn(struct sgx_epc_page *entry);
 void *sgx_kmap_epc_page(struct sgx_epc_page *entry);
 void sgx_kunmap_epc_page(void *epc_page_vaddr);
 
+enum sgx_alloc_flags {
+	SGX_ALLOC_ATOMIC	= BIT(0),
+};
+
+enum sgx_free_flags {
+	SGX_FREE_RETURN_ERROR	= BIT(0),
+	SGX_FREE_NO_EREMOVE	= BIT(1),
+};
+
+struct sgx_epc_page *sgx_alloc_page(unsigned int flags);
+int sgx_free_page(struct sgx_epc_page *entry, unsigned int flags);
+
 #endif /* _ASM_X86_SGX_H */
diff --git a/drivers/platform/x86/intel_sgx.h b/drivers/platform/x86/intel_sgx.h
index b9b81d6..f9fc708 100644
--- a/drivers/platform/x86/intel_sgx.h
+++ b/drivers/platform/x86/intel_sgx.h
@@ -216,6 +216,8 @@ void sgx_unpin_mm(struct sgx_encl *encl);
 void sgx_invalidate(struct sgx_encl *encl);
 int sgx_find_encl(struct mm_struct *mm, unsigned long addr,
 		  struct vm_area_struct **vma);
+void sgx_get_tgid_ctx(struct sgx_tgid_ctx *ctx);
+void sgx_put_tgid_ctx(struct sgx_tgid_ctx *ctx);
 
 enum sgx_fault_flags {
 	SGX_FAULT_RESERVE	= BIT(0),
@@ -237,20 +239,8 @@ extern struct mutex sgx_tgid_ctx_mutex;
 extern struct list_head sgx_tgid_ctx_list;
 extern struct task_struct *ksgxswapd_tsk;
 
-enum sgx_alloc_flags {
-	SGX_ALLOC_ATOMIC	= BIT(0),
-};
-
-enum sgx_free_flags {
-	SGX_FREE_RETURN_ERROR	= BIT(0),
-};
-
 int ksgxswapd(void *p);
 int sgx_page_cache_init(struct sgx_epc_bank *banks, int nr_banks);
 void sgx_page_cache_teardown(struct sgx_epc_bank *banks, int nr_banks);
-struct sgx_epc_page *sgx_alloc_page(struct sgx_tgid_ctx *tgid_epc_cnt,
-				    unsigned int flags);
-int sgx_free_page(struct sgx_epc_page *entry, struct sgx_encl *encl,
-		  unsigned int flags);
 
 #endif /* __ARCH_X86_INTEL_SGX_H__ */
diff --git a/drivers/platform/x86/intel_sgx_ioctl.c b/drivers/platform/x86/intel_sgx_ioctl.c
index a8729c3..5ff647f 100644
--- a/drivers/platform/x86/intel_sgx_ioctl.c
+++ b/drivers/platform/x86/intel_sgx_ioctl.c
@@ -216,12 +216,15 @@ static bool sgx_process_add_page_req(struct sgx_add_page_req *req)
 	struct vm_area_struct *vma;
 	int ret;
 
-	epc_page = sgx_alloc_page(encl->tgid_ctx, 0);
+	epc_page = sgx_alloc_page(0);
 	if (IS_ERR(epc_page))
 		return false;
 
+	sgx_get_tgid_ctx(encl->tgid_ctx);
+
 	if (!sgx_pin_mm(encl)) {
-		sgx_free_page(epc_page, encl, 0);
+		sgx_free_page(epc_page, 0);
+		sgx_put_tgid_ctx(encl->tgid_ctx);
 		return false;
 	}
 
@@ -275,7 +278,8 @@ static bool sgx_process_add_page_req(struct sgx_add_page_req *req)
 	sgx_unpin_mm(encl);
 	return true;
 out:
-	sgx_free_page(epc_page, encl, 0);
+	sgx_free_page(epc_page, 0);
+	sgx_put_tgid_ctx(encl->tgid_ctx);
 	mutex_unlock(&encl->lock);
 	sgx_unpin_mm(encl);
 	return false;
@@ -405,17 +409,19 @@ static int sgx_init_page(struct sgx_encl *encl,
 		if (!va_page)
 			return -ENOMEM;
 
-		epc_page = sgx_alloc_page(encl->tgid_ctx, 0);
+		epc_page = sgx_alloc_page(0);
 		if (IS_ERR(epc_page)) {
 			kfree(va_page);
 			return PTR_ERR(epc_page);
 		}
+		sgx_get_tgid_ctx(encl->tgid_ctx);
 
 		vaddr = sgx_kmap_epc_page(epc_page);
 		if (!vaddr) {
 			sgx_warn(encl, "kmap of a new VA page failed %d\n",
 				 ret);
-			sgx_free_page(epc_page, encl, 0);
+			sgx_free_page(epc_page, 0);
+			sgx_put_tgid_ctx(encl->tgid_ctx);
 			kfree(va_page);
 			return -EFAULT;
 		}
@@ -425,7 +431,8 @@ static int sgx_init_page(struct sgx_encl *encl,
 
 		if (ret) {
 			sgx_warn(encl, "EPA returned %d\n", ret);
-			sgx_free_page(epc_page, encl, 0);
+			sgx_free_page(epc_page, 0);
+			sgx_put_tgid_ctx(encl->tgid_ctx);
 			kfree(va_page);
 			return -EFAULT;
 		}
@@ -539,12 +546,14 @@ static long sgx_ioc_enclave_create(struct file *filep, unsigned int cmd,
 	encl->backing = backing;
 	encl->pcmd = pcmd;
 
-	secs_epc = sgx_alloc_page(encl->tgid_ctx, 0);
+	secs_epc = sgx_alloc_page(0);
 	if (IS_ERR(secs_epc)) {
 		ret = PTR_ERR(secs_epc);
 		secs_epc = NULL;
 		goto out;
 	}
+	/* don't need to sgx_get_tgid_ctx(encl->tgid_ctx) as encl->tgid_ctx is
+	 * NULL now */
 
 	ret = sgx_add_to_tgid_ctx(encl);
 	if (ret)
diff --git a/drivers/platform/x86/intel_sgx_page_cache.c b/drivers/platform/x86/intel_sgx_page_cache.c
index 3f9ba9c..71569e7 100644
--- a/drivers/platform/x86/intel_sgx_page_cache.c
+++ b/drivers/platform/x86/intel_sgx_page_cache.c
@@ -330,7 +330,8 @@ static void sgx_evict_page(struct sgx_encl_page *entry,
 			   struct sgx_encl *encl)
 {
 	sgx_ewb(encl, entry);
-	sgx_free_page(entry->epc_page, encl, 0);
+	sgx_free_page(entry->epc_page, 0);
+	sgx_put_tgid_ctx(encl->tgid_ctx);
 	entry->epc_page = NULL;
 	entry->flags &= ~SGX_ENCL_PAGE_RESERVED;
 }
@@ -554,18 +555,15 @@ static struct sgx_epc_page *sgx_alloc_page_fast(void)
  *
  * Return: an EPC page or an error code
  */
-struct sgx_epc_page *sgx_alloc_page(struct sgx_tgid_ctx *ctx,
-				    unsigned int flags)
+struct sgx_epc_page *sgx_alloc_page(unsigned int flags)
 {
 	struct sgx_epc_page *entry;
 
 	for ( ; ; ) {
 		entry = sgx_alloc_page_fast();
-		if (entry) {
-			if (ctx)
-				atomic_inc(&ctx->epc_cnt);
+		if (entry)
 			break;
-		} else if (flags & SGX_ALLOC_ATOMIC) {
+		else if (flags & SGX_ALLOC_ATOMIC) {
 			entry = ERR_PTR(-EBUSY);
 			break;
 		}
@@ -584,6 +582,7 @@ struct sgx_epc_page *sgx_alloc_page(struct sgx_tgid_ctx *ctx,
 
 	return entry;
 }
+EXPORT_SYMBOL_GPL(sgx_alloc_page);
 
 /**
  * sgx_free_page - free an EPC page
@@ -591,25 +590,29 @@ struct sgx_epc_page *sgx_alloc_page(struct sgx_tgid_ctx *ctx,
  * @encl:	the enclave who owns the EPC page
  * @flags:	free flags
  */
-int sgx_free_page(struct sgx_epc_page *entry, struct sgx_encl *encl,
-		  unsigned int flags)
+int sgx_free_page(struct sgx_epc_page *entry, unsigned int flags)
 {
 	void *epc;
-	int ret;
+	int ret = 0;
 
-	epc = sgx_kmap_epc_page(entry);
-	ret = __eremove(epc);
-	sgx_kunmap_epc_page(epc);
+	if (!(flags & SGX_FREE_NO_EREMOVE)) {
+		epc = sgx_kmap_epc_page(entry);
+		ret = __eremove(epc);
+		sgx_kunmap_epc_page(epc);
+	}
 
 	if (ret) {
 		if (flags & SGX_FREE_RETURN_ERROR)
 			return ret;
 
-		sgx_crit(encl, "EREMOVE returned %d\n", ret);
+		/*
+		 * sgx_crit is removed because it takes encl as parameter,
+		 * which sgx_free_page doesn't take anymore.
+		 * TODO: remove encl from sgx_{print} functions and call
+		 * sgx_crit again.
+		 */
 	}
 
-	atomic_dec(&encl->tgid_ctx->epc_cnt);
-
 	spin_lock(&sgx_free_list_lock);
 	list_add(&entry->free_list, &sgx_free_list);
 	sgx_nr_free_pages++;
@@ -617,3 +620,4 @@ int sgx_free_page(struct sgx_epc_page *entry, struct sgx_encl *encl,
 
 	return 0;
 }
+EXPORT_SYMBOL_GPL(sgx_free_page);
diff --git a/drivers/platform/x86/intel_sgx_util.c b/drivers/platform/x86/intel_sgx_util.c
index c63193d..96b3770 100644
--- a/drivers/platform/x86/intel_sgx_util.c
+++ b/drivers/platform/x86/intel_sgx_util.c
@@ -372,12 +372,13 @@ static struct sgx_encl_page *sgx_do_fault(struct vm_area_struct *vma,
 		goto out;
 	}
 
-	epc_page = sgx_alloc_page(encl->tgid_ctx, SGX_ALLOC_ATOMIC);
+	epc_page = sgx_alloc_page(SGX_ALLOC_ATOMIC);
 	if (IS_ERR(epc_page)) {
 		rc = PTR_ERR(epc_page);
 		epc_page = NULL;
 		goto out;
 	}
+	sgx_get_tgid_ctx(encl->tgid_ctx);
 
 	if (encl->flags & SGX_ENCL_DEAD) {
 		rc = -EFAULT;
@@ -406,12 +407,13 @@ static struct sgx_encl_page *sgx_do_fault(struct vm_area_struct *vma,
 
 	/* If SECS is evicted then reload it first */
 	if (encl->flags & SGX_ENCL_SECS_EVICTED) {
-		secs_epc_page = sgx_alloc_page(encl->tgid_ctx, SGX_ALLOC_ATOMIC);
+		secs_epc_page = sgx_alloc_page(SGX_ALLOC_ATOMIC);
 		if (IS_ERR(secs_epc_page)) {
 			rc = PTR_ERR(secs_epc_page);
 			secs_epc_page = NULL;
 			goto out;
 		}
+		sgx_get_tgid_ctx(encl->tgid_ctx);
 
 		rc = sgx_eldu(encl, &encl->secs_page, secs_epc_page, true);
 		if (rc)
@@ -446,10 +448,14 @@ static struct sgx_encl_page *sgx_do_fault(struct vm_area_struct *vma,
 	list_add_tail(&entry->load_list, &encl->load_list);
 out:
 	mutex_unlock(&encl->lock);
-	if (epc_page)
-		sgx_free_page(epc_page, encl, 0);
-	if (secs_epc_page)
-		sgx_free_page(secs_epc_page, encl, 0);
+	if (epc_page) {
+		sgx_free_page(epc_page, 0);
+		sgx_put_tgid_ctx(encl->tgid_ctx);
+	}
+	if (secs_epc_page) {
+		sgx_free_page(secs_epc_page, 0);
+		sgx_put_tgid_ctx(encl->tgid_ctx);
+	}
 	return rc ? ERR_PTR(rc) : entry;
 }
 
@@ -489,7 +495,8 @@ void sgx_encl_release(struct kref *ref)
 		entry = *slot;
 		if (entry->epc_page) {
 			list_del(&entry->load_list);
-			sgx_free_page(entry->epc_page, encl, 0);
+			sgx_free_page(entry->epc_page, 0);
+			sgx_put_tgid_ctx(encl->tgid_ctx);
 		}
 		radix_tree_delete(&encl->page_tree, entry->addr >> PAGE_SHIFT);
 		kfree(entry);
@@ -499,12 +506,15 @@ void sgx_encl_release(struct kref *ref)
 		va_page = list_first_entry(&encl->va_pages,
 					   struct sgx_va_page, list);
 		list_del(&va_page->list);
-		sgx_free_page(va_page->epc_page, encl, 0);
+		sgx_free_page(va_page->epc_page, 0);
+		sgx_put_tgid_ctx(encl->tgid_ctx);
 		kfree(va_page);
 	}
 
-	if (encl->secs_page.epc_page)
-		sgx_free_page(encl->secs_page.epc_page, encl, 0);
+	if (encl->secs_page.epc_page) {
+		sgx_free_page(encl->secs_page.epc_page, 0);
+		sgx_put_tgid_ctx(encl->tgid_ctx);
+	}
 
 	encl->secs_page.epc_page = NULL;
 
@@ -519,3 +529,13 @@ void sgx_encl_release(struct kref *ref)
 
 	kfree(encl);
 }
+
+void sgx_get_tgid_ctx(struct sgx_tgid_ctx *ctx)
+{
+	atomic_inc(&ctx->epc_cnt);
+}
+
+void sgx_put_tgid_ctx(struct sgx_tgid_ctx *ctx)
+{
+	atomic_dec(&ctx->epc_cnt);
+}
