diff --git a/drivers/staging/media/atomisp/pci/atomisp2/hmm/hmm_bo.c b/drivers/staging/media/atomisp/pci/atomisp2/hmm/hmm_bo.c
index 40ac3582fb7aba8ee9b58c8b8805626af20a1f4c..11162f595fc745846a11dd56d275e6b631ab249a 100644
--- a/drivers/staging/media/atomisp/pci/atomisp2/hmm/hmm_bo.c
+++ b/drivers/staging/media/atomisp/pci/atomisp2/hmm/hmm_bo.c
@@ -36,12 +36,13 @@
 #include <linux/string.h>
 #include <linux/list.h>
 #include <linux/errno.h>
-#include <asm/cacheflush.h>
 #include <linux/io.h>
 #include <asm/current.h>
 #include <linux/sched/signal.h>
 #include <linux/file.h>
 
+#include <asm/set_memory.h>
+
 #include "atomisp_internal.h"
 #include "hmm/hmm_common.h"
 #include "hmm/hmm_pool.h"
diff --git a/drivers/staging/media/atomisp/pci/atomisp2/hmm/hmm_dynamic_pool.c b/drivers/staging/media/atomisp/pci/atomisp2/hmm/hmm_dynamic_pool.c
index 639b8cdf7a5e79f2bd7bb223d01d07193306cbab..19e0e9ee37de2bbde025f21a71baea1fe75a4b54 100644
--- a/drivers/staging/media/atomisp/pci/atomisp2/hmm/hmm_dynamic_pool.c
+++ b/drivers/staging/media/atomisp/pci/atomisp2/hmm/hmm_dynamic_pool.c
@@ -27,7 +27,7 @@
 #include <linux/types.h>
 #include <linux/mm.h>
 
-#include "asm/cacheflush.h"
+#include <asm/set_memory.h>
 
 #include "atomisp_internal.h"
 
diff --git a/drivers/staging/media/atomisp/pci/atomisp2/hmm/hmm_reserved_pool.c b/drivers/staging/media/atomisp/pci/atomisp2/hmm/hmm_reserved_pool.c
index 4000c05652e1f7d24e947bc6dee82aafe96d837f..bf6586805f7fa4f75f289384f2e85ed5eabdd4fc 100644
--- a/drivers/staging/media/atomisp/pci/atomisp2/hmm/hmm_reserved_pool.c
+++ b/drivers/staging/media/atomisp/pci/atomisp2/hmm/hmm_reserved_pool.c
@@ -27,7 +27,8 @@
 #include <linux/types.h>
 #include <linux/mm.h>
 
-#include "asm/cacheflush.h"
+#include <asm/set_memory.h>
+
 #include "atomisp_internal.h"
 #include "hmm/hmm_pool.h"
 
diff --git a/drivers/staging/media/atomisp/pci/atomisp2/mmu/isp_mmu.c b/drivers/staging/media/atomisp/pci/atomisp2/mmu/isp_mmu.c
index 2009e3a11b860728326087fda36d400ff005e8dd..706bd43e8b1b7b863f88a7f2850bc03bbc8bca21 100644
--- a/drivers/staging/media/atomisp/pci/atomisp2/mmu/isp_mmu.c
+++ b/drivers/staging/media/atomisp/pci/atomisp2/mmu/isp_mmu.c
@@ -30,13 +30,16 @@
 #include <linux/slab.h>		/* for kmalloc */
 #include <linux/list.h>
 #include <linux/io.h>
-#include <asm/cacheflush.h>
 #include <linux/module.h>
 #include <linux/moduleparam.h>
 #include <linux/string.h>
 #include <linux/errno.h>
 #include <linux/sizes.h>
 
+#ifdef CONFIG_X86
+#include <asm/set_memory.h>
+#endif
+
 #include "atomisp_internal.h"
 #include "mmu/isp_mmu.h"