diff --git a/cputlb.c b/cputlb.c
index dd149a28bc..bf1d50adde 100644
--- a/cputlb.c
+++ b/cputlb.c
@@ -293,8 +293,9 @@ static inline void tlb_set_dirty1(CPUTLBEntry *tlb_entry, target_ulong vaddr)
 
 /* update the TLB corresponding to virtual page vaddr
    so that it is no longer dirty */
-void tlb_set_dirty(CPUArchState *env, target_ulong vaddr)
+void tlb_set_dirty(CPUState *cpu, target_ulong vaddr)
 {
+    CPUArchState *env = cpu->env_ptr;
     int i;
     int mmu_idx;
 
diff --git a/exec.c b/exec.c
index 1fa27f5c08..f760f44911 100644
--- a/exec.c
+++ b/exec.c
@@ -1922,8 +1922,7 @@ static void notdirty_mem_write(void *opaque, hwaddr ram_addr,
     /* we remove the notdirty callback only if the code has been
        flushed */
     if (!cpu_physical_memory_is_clean(ram_addr)) {
-        CPUArchState *env = current_cpu->env_ptr;
-        tlb_set_dirty(env, current_cpu->mem_io_vaddr);
+        tlb_set_dirty(current_cpu, current_cpu->mem_io_vaddr);
     }
 }
 
diff --git a/include/exec/cputlb.h b/include/exec/cputlb.h
index c3aaa30efa..7ad5c9a467 100644
--- a/include/exec/cputlb.h
+++ b/include/exec/cputlb.h
@@ -26,7 +26,7 @@ void tlb_unprotect_code(ram_addr_t ram_addr);
 void tlb_reset_dirty_range(CPUTLBEntry *tlb_entry, uintptr_t start,
                            uintptr_t length);
 void tlb_reset_dirty(CPUState *cpu, ram_addr_t start1, ram_addr_t length);
-void tlb_set_dirty(CPUArchState *env, target_ulong vaddr);
+void tlb_set_dirty(CPUState *cpu, target_ulong vaddr);
 extern int tlb_flush_count;
 
 /* exec.c */