drm/ttm: add set_pages_wb for handling page order more than zero

index a8b2bfa..cdbb731 100644
@@ -223,6 +223,17 @@ static struct kobj_type ttm_pool_kobj_type = {
 static struct ttm_pool_manager *_manager;
 
 #ifndef CONFIG_X86
+static int set_pages_wb(struct page *page, int numpages)
+{
+#if IS_ENABLED(CONFIG_AGP)
+	int i;
+
+	for (i = 0; i < numpages; i++)
+		unmap_page_from_agp(page++);
+#endif
+	return 0;
+}
+
 static int set_pages_array_wb(struct page **pages, int addrinarray)
 {
 #if IS_ENABLED(CONFIG_AGP)

index cdbb731..25b0fa5 100644
@@ -296,13 +296,23 @@ static struct ttm_page_pool *ttm_get_pool(int flags, bool huge,
 }
 
 /* set memory back to wb and free the pages. */
-static void ttm_pages_put(struct page *pages[], unsigned npages)
+static void ttm_pages_put(struct page *pages[], unsigned npages,
+		unsigned int order)
 {
-	unsigned i;
-	if (set_pages_array_wb(pages, npages))
-		pr_err("Failed to set %d pages to wb!\n", npages);
-	for (i = 0; i < npages; ++i)
-		__free_page(pages[i]);
+	unsigned int i, pages_nr = (1 << order);
+
+	if (order == 0) {
+		if (set_pages_array_wb(pages, npages))
+			pr_err("Failed to set %d pages to wb!\n", npages);
+	}
+
+	for (i = 0; i < npages; ++i) {
+		if (order > 0) {
+			if (set_pages_wb(pages[i], pages_nr))
+				pr_err("Failed to set %d pages to wb!\n", pages_nr);
+		}
+		__free_pages(pages[i], order);
+	}
 }
 
 static void ttm_pool_update_free_locked(struct ttm_page_pool *pool,

@@ -365,7 +375,7 @@ static int ttm_page_pool_free(struct ttm_page_pool *pool, unsigned nr_free,
 			 */
 			spin_unlock_irqrestore(&pool->lock, irq_flags);
 
-			ttm_pages_put(pages_to_free, freed_pages);
+			ttm_pages_put(pages_to_free, freed_pages, pool->order);
 			if (likely(nr_free != FREE_ALL_PAGES))
 				nr_free -= freed_pages;
 

@@ -400,7 +410,7 @@ static int ttm_page_pool_free(struct ttm_page_pool *pool, unsigned nr_free,
 	spin_unlock_irqrestore(&pool->lock, irq_flags);
 
 	if (freed_pages)
-		ttm_pages_put(pages_to_free, freed_pages);
+		ttm_pages_put(pages_to_free, freed_pages, pool->order);
 out:
 	if (pages_to_free != static_buf)
 		kfree(pages_to_free);

index 25b0fa5..1543532 100644
@@ -442,17 +442,19 @@ ttm_pool_shrink_scan(struct shrinker *shrink, struct shrink_control *sc)
 	/* select start pool in round robin fashion */
 	for (i = 0; i < NUM_POOLS; ++i) {
 		unsigned nr_free = shrink_pages;
+		unsigned page_nr;
+
 		if (shrink_pages == 0)
 			break;
 
 		pool = &_manager->pools[(i + pool_offset)%NUM_POOLS];
+		page_nr = (1 << pool->order);
 		/* OK to use static buffer since global mutex is held. */
-		nr_free_pool = (nr_free >> pool->order);
-		if (nr_free_pool == 0)
-			continue;
-
+		nr_free_pool = roundup(nr_free, page_nr) >> pool->order;
 		shrink_pages = ttm_page_pool_free(pool, nr_free_pool, true);
-		freed += ((nr_free_pool - shrink_pages) << pool->order);
+		freed += (nr_free_pool - shrink_pages) << pool->order;
+		if (freed >= sc->nr_to_scan)
+			break;
 	}
 	mutex_unlock(&lock);
 	return freed;