diff --git a/arch/riscv/include/asm/set_memory.h b/arch/riscv/include/asm/set_memory.h
index ec11001c3fe042..2b9f1b17022e58 100644
--- a/arch/riscv/include/asm/set_memory.h
+++ b/arch/riscv/include/asm/set_memory.h
@@ -16,6 +16,10 @@ int set_memory_rw(unsigned long addr, int numpages);
 int set_memory_x(unsigned long addr, int numpages);
 int set_memory_nx(unsigned long addr, int numpages);
 int set_memory_rw_nx(unsigned long addr, int numpages);
+
+#if defined(CONFIG_32BIT) && defined(CONFIG_STRICT_KERNEL_RWX)
+int set_memory_rw_nx_by_mm(unsigned long addr, int numpages, struct mm_struct *mm);
+#endif
 static __always_inline int set_kernel_memory(char *startp, char *endp,
 					     int (*set_memory)(unsigned long start,
 							       int num_pages))
@@ -32,6 +36,14 @@ static inline int set_memory_rw(unsigned long addr, int numpages) { return 0; }
 static inline int set_memory_x(unsigned long addr, int numpages) { return 0; }
 static inline int set_memory_nx(unsigned long addr, int numpages) { return 0; }
 static inline int set_memory_rw_nx(unsigned long addr, int numpages) { return 0; }
+
+#if defined(CONFIG_32BIT) && defined(CONFIG_STRICT_KERNEL_RWX)
+static inline int set_memory_rw_nx_by_mm(unsigned long addr,
+				int numpages, struct mm_struct *mm)
+{
+	return 0;
+}
+#endif
 static inline int set_kernel_memory(char *startp, char *endp,
 				    int (*set_memory)(unsigned long start,
 						      int num_pages))
diff --git a/arch/riscv/kernel/setup.c b/arch/riscv/kernel/setup.c
index 8fd6c02353d434..1750f789207568 100644
--- a/arch/riscv/kernel/setup.c
+++ b/arch/riscv/kernel/setup.c
@@ -342,13 +342,43 @@ static int __init topology_init(void)
 }
 subsys_initcall(topology_init);
 
-void free_initmem(void)
+#if defined(CONFIG_32BIT) && defined(CONFIG_STRICT_KERNEL_RWX)
+static void set_kernel_mm_early(char *startp, char *endp,
+				int (*set_memory)(unsigned long start,
+				int num_pages, struct mm_struct *mm))
 {
-	if (IS_ENABLED(CONFIG_STRICT_KERNEL_RWX)) {
-		set_kernel_memory(lm_alias(__init_begin), lm_alias(__init_end), set_memory_rw_nx);
-		if (IS_ENABLED(CONFIG_64BIT))
-			set_kernel_memory(__init_begin, __init_end, set_memory_nx);
+	struct task_struct *t, *s;
+	unsigned long start = (unsigned long)startp;
+	unsigned long end = (unsigned long)endp;
+	int num_pages = PAGE_ALIGN(end - start) >> PAGE_SHIFT;
+
+	set_memory(start, num_pages, current->active_mm);
+	if (current->active_mm != &init_mm)
+		set_memory(start, num_pages, &init_mm);
+
+	for_each_process(t) {
+		if (t->flags & PF_KTHREAD)
+			continue;
+		for_each_thread(t, s) {
+			if (s->mm)
+				set_memory(start, num_pages, s->mm);
+		}
 	}
+}
+#endif
+
+void free_initmem(void)
+{
+#ifdef CONFIG_STRICT_KERNEL_RWX
+#ifdef CONFIG_32BIT
+	set_kernel_mm_early(lm_alias(__init_begin), lm_alias(__init_end),
+			    set_memory_rw_nx_by_mm);
+#else
+	set_kernel_memory(lm_alias(__init_begin), lm_alias(__init_end), set_memory_rw_nx);
+#endif
+	if (IS_ENABLED(CONFIG_64BIT))
+		set_kernel_memory(__init_begin, __init_end, set_memory_nx);
+#endif
 
 	free_initmem_default(POISON_FREE_INITMEM);
 }
diff --git a/arch/riscv/mm/pageattr.c b/arch/riscv/mm/pageattr.c
index 161d0b34c2cb28..c712b22c303867 100644
--- a/arch/riscv/mm/pageattr.c
+++ b/arch/riscv/mm/pageattr.c
@@ -106,7 +106,7 @@ static const struct mm_walk_ops pageattr_ops = {
 };
 
 static int __set_memory(unsigned long addr, int numpages, pgprot_t set_mask,
-			pgprot_t clear_mask)
+			pgprot_t clear_mask, struct mm_struct *mm)
 {
 	int ret;
 	unsigned long start = addr;
@@ -119,42 +119,50 @@ static int __set_memory(unsigned long addr, int numpages, pgprot_t set_mask,
 	if (!numpages)
 		return 0;
 
-	mmap_write_lock(&init_mm);
-	ret =  walk_page_range_novma(&init_mm, start, end, &pageattr_ops, NULL,
+	mmap_write_lock(mm);
+	ret =  walk_page_range_novma(mm, start, end, &pageattr_ops, NULL,
 				     &masks);
-	mmap_write_unlock(&init_mm);
+	mmap_write_unlock(mm);
 
 	flush_tlb_kernel_range(start, end);
 
 	return ret;
 }
 
+#if defined(CONFIG_32BIT) && defined(CONFIG_STRICT_KERNEL_RWX)
+int set_memory_rw_nx_by_mm(unsigned long addr, int numpages, struct mm_struct *mm)
+{
+	return __set_memory(addr, numpages, __pgprot(_PAGE_READ | _PAGE_WRITE),
+			    __pgprot(_PAGE_EXEC), mm);
+}
+#endif
+
 int set_memory_rw_nx(unsigned long addr, int numpages)
 {
 	return __set_memory(addr, numpages, __pgprot(_PAGE_READ | _PAGE_WRITE),
-			    __pgprot(_PAGE_EXEC));
+			    __pgprot(_PAGE_EXEC), &init_mm);
 }
 
 int set_memory_ro(unsigned long addr, int numpages)
 {
 	return __set_memory(addr, numpages, __pgprot(_PAGE_READ),
-			    __pgprot(_PAGE_WRITE));
+			    __pgprot(_PAGE_WRITE), &init_mm);
 }
 
 int set_memory_rw(unsigned long addr, int numpages)
 {
 	return __set_memory(addr, numpages, __pgprot(_PAGE_READ | _PAGE_WRITE),
-			    __pgprot(0));
+			    __pgprot(0), &init_mm);
 }
 
 int set_memory_x(unsigned long addr, int numpages)
 {
-	return __set_memory(addr, numpages, __pgprot(_PAGE_EXEC), __pgprot(0));
+	return __set_memory(addr, numpages, __pgprot(_PAGE_EXEC), __pgprot(0), &init_mm);
 }
 
 int set_memory_nx(unsigned long addr, int numpages)
 {
-	return __set_memory(addr, numpages, __pgprot(0), __pgprot(_PAGE_EXEC));
+	return __set_memory(addr, numpages, __pgprot(0), __pgprot(_PAGE_EXEC), &init_mm);
 }
 
 int set_direct_map_invalid_noflush(struct page *page)
@@ -199,10 +207,10 @@ void __kernel_map_pages(struct page *page, int numpages, int enable)
 
 	if (enable)
 		__set_memory((unsigned long)page_address(page), numpages,
-			     __pgprot(_PAGE_PRESENT), __pgprot(0));
+			     __pgprot(_PAGE_PRESENT), __pgprot(0), &init_mm);
 	else
 		__set_memory((unsigned long)page_address(page), numpages,
-			     __pgprot(0), __pgprot(_PAGE_PRESENT));
+			     __pgprot(0), __pgprot(_PAGE_PRESENT), &init_mm);
 }
 #endif