diff --git a/arch/s390/mm/page-states.c b/arch/s390/mm/page-states.c
index f92ec20..098923a 100644
--- a/arch/s390/mm/page-states.c
+++ b/arch/s390/mm/page-states.c
@@ -50,28 +50,64 @@
 		cmma_flag = 0;
 }
 
-void arch_free_page(struct page *page, int order)
+static inline void set_page_unstable(struct page *page, int order)
 {
 	int i, rc;
 
-	if (!cmma_flag)
-		return;
 	for (i = 0; i < (1 << order); i++)
 		asm volatile(".insn rrf,0xb9ab0000,%0,%1,%2,0"
 			     : "=&d" (rc)
-			     : "a" ((page_to_pfn(page) + i) << PAGE_SHIFT),
+			     : "a" (page_to_phys(page + i)),
 			       "i" (ESSA_SET_UNUSED));
 }
 
+void arch_free_page(struct page *page, int order)
+{
+	if (!cmma_flag)
+		return;
+	set_page_unstable(page, order);
+}
+
+static inline void set_page_stable(struct page *page, int order)
+{
+	int i, rc;
+
+	for (i = 0; i < (1 << order); i++)
+		asm volatile(".insn rrf,0xb9ab0000,%0,%1,%2,0"
+			     : "=&d" (rc)
+			     : "a" (page_to_phys(page + i)),
+			       "i" (ESSA_SET_STABLE));
+}
+
 void arch_alloc_page(struct page *page, int order)
 {
-	int i, rc;
+	if (!cmma_flag)
+		return;
+	set_page_stable(page, order);
+}
+
+void arch_set_page_states(int make_stable)
+{
+	unsigned long flags, order, t;
+	struct list_head *l;
+	struct page *page;
+	struct zone *zone;
 
 	if (!cmma_flag)
 		return;
-	for (i = 0; i < (1 << order); i++)
-		asm volatile(".insn rrf,0xb9ab0000,%0,%1,%2,0"
-			     : "=&d" (rc)
-			     : "a" ((page_to_pfn(page) + i) << PAGE_SHIFT),
-			       "i" (ESSA_SET_STABLE));
+	if (make_stable)
+		drain_local_pages(NULL);
+	for_each_populated_zone(zone) {
+		spin_lock_irqsave(&zone->lock, flags);
+		for_each_migratetype_order(order, t) {
+			list_for_each(l, &zone->free_area[order].free_list[t]) {
+				page = list_entry(l, struct page, lru);
+				if (make_stable)
+					set_page_stable(page, order);
+				else
+					set_page_unstable(page, order);
+			}
+		}
+		spin_unlock_irqrestore(&zone->lock, flags);
+	}
 }