diff --git a/include/trace/hooks/mm.h b/include/trace/hooks/mm.h
index 236aaeeff274..785ec3da0aa4 100644
--- a/include/trace/hooks/mm.h
+++ b/include/trace/hooks/mm.h
@@ -110,14 +110,14 @@ DECLARE_HOOK(android_vh_reclaim_pages_plug,
 	TP_PROTO(bool *do_plug),
 	TP_ARGS(do_plug));
 DECLARE_HOOK(android_vh_zap_pte_range_tlb_start,
-	TP_PROTO(void *unused),
-	TP_ARGS(unused));
+	TP_PROTO(void *ret),
+	TP_ARGS(ret));
 DECLARE_HOOK(android_vh_zap_pte_range_tlb_force_flush,
 	TP_PROTO(struct page *page, bool *flush),
 	TP_ARGS(page, flush));
 DECLARE_HOOK(android_vh_zap_pte_range_tlb_end,
-	TP_PROTO(void *unused),
-	TP_ARGS(unused));
+	TP_PROTO(void *ret),
+	TP_ARGS(ret));
 DECLARE_HOOK(android_vh_skip_lru_disable,
 	TP_PROTO(bool *skip),
 	TP_ARGS(skip));
diff --git a/mm/memory.c b/mm/memory.c
index f0fa06835537..83b715ed6577 100644
--- a/mm/memory.c
+++ b/mm/memory.c
@@ -1230,10 +1230,11 @@ static unsigned long zap_pte_range(struct mmu_gather *tlb,
 	pte_t *start_pte;
 	pte_t *pte;
 	swp_entry_t entry;
+	int v_ret = 0;
 
 	tlb_change_page_size(tlb, PAGE_SIZE);
 again:
-	trace_android_vh_zap_pte_range_tlb_start(NULL);
+	trace_android_vh_zap_pte_range_tlb_start(&v_ret);
 	init_rss_vec(rss);
 	start_pte = pte_offset_map_lock(mm, pmd, addr, &ptl);
 	pte = start_pte;
@@ -1349,7 +1350,7 @@ again:
 		tlb_flush_mmu(tlb);
 	}
 
-	trace_android_vh_zap_pte_range_tlb_end(NULL);
+	trace_android_vh_zap_pte_range_tlb_end(&v_ret);
 	if (addr != end) {
 		cond_resched();
 		goto again;