diff --git a/Documentation/admin-guide/cgroup-v2.rst b/Documentation/admin-guide/cgroup-v2.rst
index aed195a71cbf1..d883f4edb1bd8 100644
--- a/Documentation/admin-guide/cgroup-v2.rst
+++ b/Documentation/admin-guide/cgroup-v2.rst
@@ -3257,7 +3257,7 @@ the configuration, the bio may be executed at a lower priority and if
 the writeback session is holding shared resources, e.g. a journal
 entry, may lead to priority inversion.  There is no one easy solution
 for the problem.  Filesystems can try to work around specific problem
-cases by skipping wbc_init_bio() and using bio_associate_blkg()
+cases by skipping wbc_init_bio() and using bio_associate_blkcg()
 directly.
 
 
diff --git a/block/bfq-cgroup.c b/block/bfq-cgroup.c
index e82ff03bda02e..720f11b389acf 100644
--- a/block/bfq-cgroup.c
+++ b/block/bfq-cgroup.c
@@ -363,11 +363,13 @@ void bfqg_and_blkg_put(struct bfq_group *bfqg)
 
 void bfqg_stats_update_legacy_io(struct request_queue *q, struct request *rq)
 {
-	struct bfq_group *bfqg = blkg_to_bfqg(rq->bio->bi_blkg);
+	struct blkcg_gq *blkg = bio_blkg_lookup(rq->bio);
+	struct bfq_group *bfqg;
 
-	if (!bfqg)
+	if (!blkg)
 		return;
 
+	bfqg = blkg_to_bfqg(blkg);
 	blkg_rwstat_add(&bfqg->stats.bytes, rq->cmd_flags, blk_rq_bytes(rq));
 	blkg_rwstat_add(&bfqg->stats.ios, rq->cmd_flags, 1);
 }
@@ -606,7 +608,7 @@ static void bfq_link_bfqg(struct bfq_data *bfqd, struct bfq_group *bfqg)
 
 struct bfq_group *bfq_bio_bfqg(struct bfq_data *bfqd, struct bio *bio)
 {
-	struct blkcg_gq *blkg = bio->bi_blkg;
+	struct blkcg_gq *blkg = bio_blkg_lookup(bio);
 	struct bfq_group *bfqg;
 
 	while (blkg) {
@@ -614,15 +616,17 @@ struct bfq_group *bfq_bio_bfqg(struct bfq_data *bfqd, struct bio *bio)
 			blkg = blkg->parent;
 			continue;
 		}
+
 		bfqg = blkg_to_bfqg(blkg);
 		if (bfqg->pd.online) {
-			bio_associate_blkg_from_css(bio, &blkg->blkcg->css);
+			bio_associate_blkcg_from_css(bio, &blkg->blkcg->css);
 			return bfqg;
 		}
 		blkg = blkg->parent;
 	}
-	bio_associate_blkg_from_css(bio,
-				&bfqg_to_blkg(bfqd->root_group)->blkcg->css);
+
+	blkg = bfqg_to_blkg(bfqd->root_group);
+	bio_associate_blkcg_from_css(bio, &blkg->blkcg->css);
 	return bfqd->root_group;
 }
 
diff --git a/block/bfq-iosched.c b/block/bfq-iosched.c
index 0f75301b3115f..d4f657bc7a4e5 100644
--- a/block/bfq-iosched.c
+++ b/block/bfq-iosched.c
@@ -128,6 +128,7 @@
 
 #include "elevator.h"
 #include "blk.h"
+#include "blk-cgroup.h"
 #include "blk-mq.h"
 #include "blk-mq-sched.h"
 #include "bfq-iosched.h"
@@ -2452,15 +2453,16 @@ static bool bfq_bio_merge(struct request_queue *q, struct bio *bio,
 	struct request *free = NULL;
 	bool ret;
 
+#ifdef CONFIG_BFQ_GROUP_IOSCHED
+	/* blkg creation takes q->queue_lock, so do it before bfqd->lock. */
+	if (bic)
+		bio_blkg(bio);
+#endif
+
 	spin_lock_irq(&bfqd->lock);
 
 	if (bic) {
-		/*
-		 * Make sure cgroup info is uptodate for current process before
-		 * considering the merge.
-		 */
 		bfq_bic_update_cgroup(bic, bio);
-
 		bfqd->bio_bfqq = bic_to_bfqq(bic, op_is_sync(bio->bi_opf),
 					     bfq_actuator_index(bfqd, bio));
 	} else {
@@ -6245,6 +6247,13 @@ static void bfq_insert_request(struct blk_mq_hw_ctx *hctx, struct request *rq,
 	LIST_HEAD(free);
 
 #ifdef CONFIG_BFQ_GROUP_IOSCHED
+	/*
+	 * Pin the blkg used to look up bfqg.  If this is the first IO for
+	 * the blkcg on this queue, create the bfqg before holding bfqd->lock.
+	 */
+	if (rq->bio && !bio_flagged(rq->bio, BIO_BLKG_REF))
+		bio_blkg(rq->bio);
+
 	if (!cgroup_subsys_on_dfl(io_cgrp_subsys) && rq->bio)
 		bfqg_stats_update_legacy_io(q, rq);
 #endif
diff --git a/block/bio.c b/block/bio.c
index 6a2f6fc3413e8..db33c993c2964 100644
--- a/block/bio.c
+++ b/block/bio.c
@@ -179,12 +179,7 @@ static inline gfp_t try_alloc_gfp(gfp_t gfp)
 
 void bio_uninit(struct bio *bio)
 {
-#ifdef CONFIG_BLK_CGROUP
-	if (bio->bi_blkg) {
-		blkg_put(bio->bi_blkg);
-		bio->bi_blkg = NULL;
-	}
-#endif
+	bio_clear_blkcg(bio);
 	if (bio_integrity(bio))
 		bio_integrity_free(bio);
 
@@ -233,10 +228,10 @@ void bio_init(struct bio *bio, struct block_device *bdev, struct bio_vec *table,
 	bio->bi_end_io = NULL;
 	bio->bi_private = NULL;
 #ifdef CONFIG_BLK_CGROUP
-	bio->bi_blkg = NULL;
+	bio->bi_blkcg = NULL;
 	bio->issue_time_ns = 0;
 	if (bdev)
-		bio_associate_blkg(bio);
+		bio_associate_blkcg(bio);
 #ifdef CONFIG_BLK_CGROUP_IOCOST
 	bio->bi_iocost_cost = 0;
 #endif
@@ -281,7 +276,7 @@ void bio_reset(struct bio *bio, struct block_device *bdev, blk_opf_t opf)
 	bio->bi_io_vec = bv;
 	bio->bi_bdev = bdev;
 	if (bio->bi_bdev)
-		bio_associate_blkg(bio);
+		bio_associate_blkcg(bio);
 	bio->bi_opf = opf;
 }
 EXPORT_SYMBOL(bio_reset);
@@ -865,7 +860,7 @@ static int __bio_clone(struct bio *bio, struct bio *bio_src, gfp_t gfp)
 		if (bio->bi_bdev == bio_src->bi_bdev &&
 		    bio_flagged(bio_src, BIO_REMAPPED))
 			bio_set_flag(bio, BIO_REMAPPED);
-		bio_clone_blkg_association(bio, bio_src);
+		bio_clone_blkcg_association(bio, bio_src);
 	}
 
 	if (bio_crypt_clone(bio, bio_src, gfp) < 0)
@@ -1803,17 +1798,12 @@ void bio_endio(struct bio *bio)
 		goto again;
 	}
 
-#ifdef CONFIG_BLK_CGROUP
 	/*
 	 * Release cgroup info.  We shouldn't have to do this here, but quite
 	 * a few callers of bio_init fail to call bio_uninit, so we cover up
 	 * for that here at least for now.
 	 */
-	if (bio->bi_blkg) {
-		blkg_put(bio->bi_blkg);
-		bio->bi_blkg = NULL;
-	}
-#endif
+	bio_clear_blkcg(bio);
 
 	if (bio->bi_end_io)
 		bio->bi_end_io(bio);
diff --git a/block/blk-cgroup-fc-appid.c b/block/blk-cgroup-fc-appid.c
index 3ec21333f393b..7589c62099898 100644
--- a/block/blk-cgroup-fc-appid.c
+++ b/block/blk-cgroup-fc-appid.c
@@ -50,8 +50,14 @@ EXPORT_SYMBOL_GPL(blkcg_set_fc_appid);
  */
 char *blkcg_get_fc_appid(struct bio *bio)
 {
-	if (!bio->bi_blkg || bio->bi_blkg->blkcg->fc_app_id[0] == '\0')
+	struct blkcg *blkcg = bio_blkcg(bio);
+
+	if (!blkcg)
+		return NULL;
+
+	if (blkcg->fc_app_id[0] == '\0')
 		return NULL;
-	return bio->bi_blkg->blkcg->fc_app_id;
+
+	return blkcg->fc_app_id;
 }
 EXPORT_SYMBOL_GPL(blkcg_get_fc_appid);
diff --git a/block/blk-cgroup.c b/block/blk-cgroup.c
index d9676126c5b5d..e7405fe5a9e25 100644
--- a/block/blk-cgroup.c
+++ b/block/blk-cgroup.c
@@ -63,6 +63,13 @@ static DEFINE_RAW_SPINLOCK(blkg_stat_lock);
 
 #define BLKG_DESTROY_BATCH_SIZE  64
 
+const struct rhashtable_params blkg_hash_params = {
+	.key_len		= sizeof_field(struct blkcg_gq, blkcg_id),
+	.key_offset		= offsetof(struct blkcg_gq, blkcg_id),
+	.head_offset		= offsetof(struct blkcg_gq, q_hash_node),
+	.automatic_shrinking	= true,
+};
+
 /*
  * Lockless lists for tracking IO stats update
  *
@@ -170,10 +177,6 @@ static void __blkg_release(struct rcu_head *rcu)
 {
 	struct blkcg_gq *blkg = container_of(rcu, struct blkcg_gq, rcu_head);
 
-#ifdef CONFIG_BLK_CGROUP_PUNT_BIO
-	WARN_ON(!bio_list_empty(&blkg->async_bios));
-#endif
-
 	blkg_free(blkg);
 }
 
@@ -191,6 +194,16 @@ static void blkg_release(struct percpu_ref *ref)
 	struct blkcg *blkcg = blkg->blkcg;
 	int cpu;
 
+	/*
+	 * A blkg that was never inserted into q->blkg_list has no hash
+	 * entry.  This happens when allocation or creation fails before
+	 * rhashtable_insert_fast() succeeds.
+	 */
+	if (!list_empty(&blkg->q_node))
+		WARN_ON_ONCE(rhashtable_remove_fast(&blkg->q->blkg_hash,
+						    &blkg->q_hash_node,
+						    blkg_hash_params));
+
 	/*
 	 * Flush all the non-empty percpu lockless lists before releasing
 	 * us, given these stat belongs to us.
@@ -206,19 +219,18 @@ static void blkg_release(struct percpu_ref *ref)
 #ifdef CONFIG_BLK_CGROUP_PUNT_BIO
 static struct workqueue_struct *blkcg_punt_bio_wq;
 
-static void blkg_async_bio_workfn(struct work_struct *work)
+static void blkcg_async_bio_workfn(struct work_struct *work)
 {
-	struct blkcg_gq *blkg = container_of(work, struct blkcg_gq,
-					     async_bio_work);
+	struct blkcg *blkcg = container_of(work, struct blkcg, async_bio_work);
 	struct bio_list bios = BIO_EMPTY_LIST;
 	struct bio *bio;
 	struct blk_plug plug;
 	bool need_plug = false;
 
-	/* as long as there are pending bios, @blkg can't go away */
-	spin_lock(&blkg->async_bio_lock);
-	bio_list_merge_init(&bios, &blkg->async_bios);
-	spin_unlock(&blkg->async_bio_lock);
+	/* as long as there are pending bios, @blkcg can't go away */
+	spin_lock(&blkcg->async_bio_lock);
+	bio_list_merge_init(&bios, &blkcg->async_bios);
+	spin_unlock(&blkcg->async_bio_lock);
 
 	/* start plug only when bio_list contains at least 2 bios */
 	if (bios.head && bios.head->bi_next) {
@@ -239,15 +251,15 @@ static void blkg_async_bio_workfn(struct work_struct *work)
  */
 void blkcg_punt_bio_submit(struct bio *bio)
 {
-	struct blkcg_gq *blkg = bio->bi_blkg;
+	struct blkcg *blkcg = bio_blkcg(bio);
 
-	if (blkg->parent) {
-		spin_lock(&blkg->async_bio_lock);
-		bio_list_add(&blkg->async_bios, bio);
-		spin_unlock(&blkg->async_bio_lock);
-		queue_work(blkcg_punt_bio_wq, &blkg->async_bio_work);
+	if (blkcg && cgroup_parent(blkcg->css.cgroup)) {
+		spin_lock(&blkcg->async_bio_lock);
+		bio_list_add(&blkcg->async_bios, bio);
+		spin_unlock(&blkcg->async_bio_lock);
+		queue_work(blkcg_punt_bio_wq, &blkcg->async_bio_work);
 	} else {
-		/* never bounce for the root cgroup */
+		/* Never bounce if there is no non-root blkcg to queue on. */
 		submit_bio(bio);
 	}
 }
@@ -275,9 +287,13 @@ subsys_initcall(blkcg_punt_bio_init);
  */
 struct cgroup_subsys_state *bio_blkcg_css(struct bio *bio)
 {
-	if (!bio || !bio->bi_blkg)
+	struct blkcg *blkcg;
+
+	if (!bio)
 		return NULL;
-	return &bio->bi_blkg->blkcg->css;
+
+	blkcg = bio_blkcg(bio);
+	return blkcg ? &blkcg->css : NULL;
 }
 EXPORT_SYMBOL_GPL(bio_blkcg_css);
 
@@ -324,12 +340,8 @@ static struct blkcg_gq *blkg_alloc(struct blkcg *blkcg, struct gendisk *disk,
 	blkg->q = disk->queue;
 	INIT_LIST_HEAD(&blkg->q_node);
 	blkg->blkcg = blkcg;
+	blkg->blkcg_id = blkcg->css.id;
 	blkg->iostat.blkg = blkg;
-#ifdef CONFIG_BLK_CGROUP_PUNT_BIO
-	spin_lock_init(&blkg->async_bio_lock);
-	bio_list_init(&blkg->async_bios);
-	INIT_WORK(&blkg->async_bio_work, blkg_async_bio_workfn);
-#endif
 
 	u64_stats_init(&blkg->iostat.sync);
 	for_each_possible_cpu(cpu) {
@@ -420,7 +432,8 @@ static struct blkcg_gq *blkg_create(struct blkcg *blkcg, struct gendisk *disk,
 
 	/* insert */
 	spin_lock(&blkcg->lock);
-	ret = radix_tree_insert(&blkcg->blkg_tree, disk->queue->id, blkg);
+	ret = rhashtable_insert_fast(&disk->queue->blkg_hash,
+				     &blkg->q_hash_node, blkg_hash_params);
 	if (likely(!ret)) {
 		hlist_add_head_rcu(&blkg->blkcg_node, &blkcg->blkg_list);
 		list_add(&blkg->q_node, &disk->queue->blkg_list);
@@ -451,6 +464,17 @@ static struct blkcg_gq *blkg_create(struct blkcg *blkcg, struct gendisk *disk,
 	return ERR_PTR(ret);
 }
 
+/*
+ * The root blkg holds a live reference while the disk is active, so walking
+ * the parent chain always finds a blkg which can be pinned.
+ */
+static struct blkcg_gq *blkg_lookup_tryget(struct blkcg_gq *blkg)
+{
+	while (!blkg_tryget(blkg))
+		blkg = blkg->parent;
+	return blkg;
+}
+
 /**
  * blkg_lookup_create - lookup blkg, try to create one if not there
  * @blkcg: blkcg of interest
@@ -458,11 +482,13 @@ static struct blkcg_gq *blkg_create(struct blkcg *blkcg, struct gendisk *disk,
  *
  * Lookup blkg for the @blkcg - @disk pair.  If it doesn't exist, try to
  * create one.  blkg creation is performed recursively from blkcg_root such
- * that all non-root blkg's have access to the parent blkg.  This function
- * should be called under RCU read lock and takes @disk->queue->queue_lock.
+ * that all non-root blkg's have access to the parent blkg.
  *
- * Returns the blkg or the closest blkg if blkg_create() fails as it walks
- * down from root.
+ * Must be called with @disk->queue->queue_lock held.
+ *
+ * Returns the closest blkg with an extra reference acquired.  If
+ * blkg_create() fails while walking down from root, the returned blkg may
+ * belong to an ancestor of @blkcg.  This function never returns %NULL.
  */
 static struct blkcg_gq *blkg_lookup_create(struct blkcg *blkcg,
 		struct gendisk *disk)
@@ -470,12 +496,12 @@ static struct blkcg_gq *blkg_lookup_create(struct blkcg *blkcg,
 	struct request_queue *q = disk->queue;
 	struct blkcg_gq *blkg;
 
+	lockdep_assert_held(&q->queue_lock);
+
 	rcu_read_lock();
 	blkg = blkg_lookup(blkcg, q);
 	if (blkg) {
-		if (blkcg != &blkcg_root &&
-		    blkg != rcu_dereference(blkcg->blkg_hint))
-			rcu_assign_pointer(blkcg->blkg_hint, blkg);
+		blkg = blkg_lookup_tryget(blkg);
 		rcu_read_unlock();
 		return blkg;
 	}
@@ -511,7 +537,7 @@ static struct blkcg_gq *blkg_lookup_create(struct blkcg *blkcg,
 			break;
 	}
 
-	return blkg;
+	return blkg_lookup_tryget(blkg);
 }
 
 static void blkg_destroy(struct blkcg_gq *blkg)
@@ -543,17 +569,8 @@ static void blkg_destroy(struct blkcg_gq *blkg)
 
 	blkg->online = false;
 
-	radix_tree_delete(&blkcg->blkg_tree, blkg->q->id);
 	hlist_del_init_rcu(&blkg->blkcg_node);
 
-	/*
-	 * Both setting lookup hint to and clearing it from @blkg are done
-	 * under queue_lock.  If it's not pointing to @blkg now, it never
-	 * will.  Hint assignment itself can race safely.
-	 */
-	if (rcu_access_pointer(blkcg->blkg_hint) == blkg)
-		rcu_assign_pointer(blkcg->blkg_hint, NULL);
-
 	/*
 	 * Put the reference taken at the time of creation so that when all
 	 * queues are gone, group can be destroyed.
@@ -876,18 +893,12 @@ int blkg_conf_prep(struct blkcg *blkcg, const struct blkcg_policy *pol,
 			goto fail_exit;
 		}
 
-		if (radix_tree_preload(GFP_KERNEL)) {
-			blkg_free(new_blkg);
-			ret = -ENOMEM;
-			goto fail_exit;
-		}
-
 		spin_lock_irq(&q->queue_lock);
 
 		if (!blkcg_policy_enabled(q, pol)) {
 			blkg_free(new_blkg);
 			ret = -EOPNOTSUPP;
-			goto fail_preloaded;
+			goto fail_unlock;
 		}
 
 		blkg = blkg_lookup(pos, q);
@@ -897,12 +908,10 @@ int blkg_conf_prep(struct blkcg *blkcg, const struct blkcg_policy *pol,
 			blkg = blkg_create(pos, disk, new_blkg);
 			if (IS_ERR(blkg)) {
 				ret = PTR_ERR(blkg);
-				goto fail_preloaded;
+				goto fail_unlock;
 			}
 		}
 
-		radix_tree_preload_end();
-
 		if (pos == blkcg)
 			goto success;
 	}
@@ -911,8 +920,6 @@ int blkg_conf_prep(struct blkcg *blkcg, const struct blkcg_policy *pol,
 	ctx->blkg = blkg;
 	return 0;
 
-fail_preloaded:
-	radix_tree_preload_end();
 fail_unlock:
 	spin_unlock_irq(&q->queue_lock);
 fail_exit:
@@ -1357,6 +1364,9 @@ static void blkcg_css_free(struct cgroup_subsys_state *css)
 
 	mutex_unlock(&blkcg_pol_mutex);
 
+#ifdef CONFIG_BLK_CGROUP_PUNT_BIO
+	WARN_ON(!bio_list_empty(&blkcg->async_bios));
+#endif
 	free_percpu(blkcg->lhead);
 	kfree(blkcg);
 }
@@ -1404,8 +1414,12 @@ blkcg_css_alloc(struct cgroup_subsys_state *parent_css)
 
 	spin_lock_init(&blkcg->lock);
 	refcount_set(&blkcg->online_pin, 1);
-	INIT_RADIX_TREE(&blkcg->blkg_tree, GFP_NOWAIT);
 	INIT_HLIST_HEAD(&blkcg->blkg_list);
+#ifdef CONFIG_BLK_CGROUP_PUNT_BIO
+	spin_lock_init(&blkcg->async_bio_lock);
+	bio_list_init(&blkcg->async_bios);
+	INIT_WORK(&blkcg->async_bio_work, blkcg_async_bio_workfn);
+#endif
 #ifdef CONFIG_CGROUP_WRITEBACK
 	INIT_LIST_HEAD(&blkcg->cgwb_list);
 #endif
@@ -1441,17 +1455,22 @@ static int blkcg_css_online(struct cgroup_subsys_state *css)
 	return 0;
 }
 
-void blkg_init_queue(struct request_queue *q)
+int blkg_init_queue(struct request_queue *q)
 {
 	INIT_LIST_HEAD(&q->blkg_list);
 	mutex_init(&q->blkcg_mutex);
+	return rhashtable_init(&q->blkg_hash, &blkg_hash_params);
+}
+
+void blkg_exit_queue(struct request_queue *q)
+{
+	rhashtable_destroy(&q->blkg_hash);
 }
 
 int blkcg_init_disk(struct gendisk *disk)
 {
 	struct request_queue *q = disk->queue;
 	struct blkcg_gq *new_blkg, *blkg;
-	bool preloaded;
 
 	/*
 	 * If the queue is shared across disk rebind (e.g., SCSI), the
@@ -1469,8 +1488,6 @@ int blkcg_init_disk(struct gendisk *disk)
 	if (!new_blkg)
 		return -ENOMEM;
 
-	preloaded = !radix_tree_preload(GFP_KERNEL);
-
 	/* Make sure the root blkg exists. */
 	/* spin_lock_irq can serve as RCU read-side critical section. */
 	spin_lock_irq(&q->queue_lock);
@@ -1480,15 +1497,10 @@ int blkcg_init_disk(struct gendisk *disk)
 	q->root_blkg = blkg;
 	spin_unlock_irq(&q->queue_lock);
 
-	if (preloaded)
-		radix_tree_preload_end();
-
 	return 0;
 
 err_unlock:
 	spin_unlock_irq(&q->queue_lock);
-	if (preloaded)
-		radix_tree_preload_end();
 	return PTR_ERR(blkg);
 }
 
@@ -2046,129 +2058,165 @@ void blkcg_add_delay(struct blkcg_gq *blkg, u64 now, u64 delta)
 	atomic64_add(delta, &blkg->delay_nsec);
 }
 
-static inline struct blkcg_gq *blkg_lookup_tryget(struct blkcg_gq *blkg)
+/*
+ * Return the blkg pinned by @bio through BIO_BLKG_REF.  The returned blkg is
+ * already owned by @bio and no extra reference is acquired.
+ */
+static struct blkcg_gq *bio_pinned_blkg(struct bio *bio)
 {
-retry:
-	if (blkg_tryget(blkg))
-		return blkg;
+	struct request_queue *q = bdev_get_queue(bio->bi_bdev);
+	struct blkcg_gq *blkg;
 
-	blkg = blkg->parent;
-	if (blkg)
-		goto retry;
+	rcu_read_lock();
+	blkg = blkg_lookup_any(bio_blkcg(bio), q);
+	rcu_read_unlock();
 
-	return NULL;
+	WARN_ON_ONCE(!blkg);
+	return blkg;
+}
+
+static void bio_set_blkg_ref(struct bio *bio, struct blkcg_gq *blkg)
+{
+	if (bio_blkcg(bio) != blkg->blkcg) {
+		css_get(&blkg->blkcg->css);
+		bio_clear_blkcg(bio);
+		bio->bi_blkcg = blkg->blkcg;
+	}
+	bio_set_flag(bio, BIO_BLKG_REF);
 }
+
 /**
- * blkg_tryget_closest - try and get a blkg ref on the closet blkg
+ * bio_blkg_lookup - look up a blkg associated with a bio
  * @bio: target bio
- * @css: target css
  *
- * As the failure mode here is to walk up the blkg tree, this ensure that the
- * blkg->parent pointers are always valid.  This returns the blkg that it ended
- * up taking a reference on or %NULL if no reference was taken.
+ * Return the blkg already pinned by @bio, or %NULL if @bio doesn't own a blkg
+ * reference.  Call bio_blkg() instead when a missing blkg should be created.
  */
-static inline struct blkcg_gq *blkg_tryget_closest(struct bio *bio,
-		struct cgroup_subsys_state *css)
+struct blkcg_gq *bio_blkg_lookup(struct bio *bio)
 {
-	struct request_queue *q = bio->bi_bdev->bd_queue;
-	struct blkcg *blkcg = css_to_blkcg(css);
+	if (!bio_flagged(bio, BIO_BLKG_REF))
+		return NULL;
+	return bio_pinned_blkg(bio);
+}
+EXPORT_SYMBOL_GPL(bio_blkg_lookup);
+
+/**
+ * bio_put_blkg_ref - drop the blkg reference pinned by a bio
+ * @bio: target bio
+ *
+ * Drop the bio-owned blkg reference acquired by bio_blkg(), if any.
+ */
+void bio_put_blkg_ref(struct bio *bio)
+{
+	if (bio_flagged(bio, BIO_BLKG_REF)) {
+		struct blkcg_gq *blkg = bio_pinned_blkg(bio);
+
+		if (blkg)
+			blkg_put(blkg);
+		bio_clear_flag(bio, BIO_BLKG_REF);
+	}
+}
+EXPORT_SYMBOL_GPL(bio_put_blkg_ref);
+
+/**
+ * bio_blkg - look up the blkg associated with a bio
+ * @bio: target bio
+ *
+ * Look up the queue-local blkg for @bio's current device and blkcg.  If this
+ * is the first policy use of @bio, create the missing blkg hierarchy if
+ * necessary, pin the closest available blkg, and update @bio's blkcg
+ * association if allocation falls back to an ancestor.
+ */
+struct blkcg_gq *bio_blkg(struct bio *bio)
+{
+	struct blkcg *blkcg = bio_blkcg(bio);
+	struct gendisk *disk;
+	struct request_queue *q;
 	struct blkcg_gq *blkg;
 
-	rcu_read_lock();
-	blkg = blkg_lookup(blkcg, q);
-	if (likely(blkg))
-		blkg = blkg_lookup_tryget(blkg);
-	rcu_read_unlock();
+	if (!blkcg || !bio->bi_bdev)
+		return NULL;
 
-	if (blkg)
-		return blkg;
+	if (bio_flagged(bio, BIO_BLKG_REF))
+		return bio_pinned_blkg(bio);
+
+	disk = bio->bi_bdev->bd_disk;
+	q = disk->queue;
 
-	/*
-	 * Fast path failed, we're probably issuing IO in this cgroup the first
-	 * time, hold lock to create new blkg.
-	 */
 	spin_lock_irq(&q->queue_lock);
-	blkg = blkg_lookup_create(blkcg, bio->bi_bdev->bd_disk);
-	if (blkg)
-		blkg = blkg_lookup_tryget(blkg);
+	blkg = blkg_lookup_create(blkcg, disk);
 	spin_unlock_irq(&q->queue_lock);
 
+	bio_set_blkg_ref(bio, blkg);
 	return blkg;
 }
+EXPORT_SYMBOL_GPL(bio_blkg);
 
 /**
- * bio_associate_blkg_from_css - associate a bio with a specified css
+ * bio_associate_blkcg_from_css - associate a bio with a specified css
  * @bio: target bio
  * @css: target css
  *
- * Associate @bio with the blkg found by combining the css's blkg and the
- * request_queue of the @bio.  An association failure is handled by walking up
- * the blkg tree.  Therefore, the blkg associated can be anything between @blkg
- * and q->root_blkg.  This situation only happens when a cgroup is dying and
- * then the remaining bios will spill to the closest alive blkg.
+ * Associate @bio with the blkcg found from @css.  The queue-local blkg is
+ * created and pinned by bio_blkg() when blkcg policies need it.  If @css
+ * can't be referenced online, associate @bio with the root blkcg instead.
  *
- * A reference will be taken on the blkg and will be released when @bio is
+ * A reference will be taken on the blkcg and will be released when @bio is
  * freed.
  */
-void bio_associate_blkg_from_css(struct bio *bio,
-				 struct cgroup_subsys_state *css)
+void bio_associate_blkcg_from_css(struct bio *bio,
+				  struct cgroup_subsys_state *css)
 {
-	if (bio->bi_blkg)
-		blkg_put(bio->bi_blkg);
+	struct blkcg *blkcg;
 
-	if (css && css->parent) {
-		bio->bi_blkg = blkg_tryget_closest(bio, css);
-	} else {
-		blkg_get(bdev_get_queue(bio->bi_bdev)->root_blkg);
-		bio->bi_blkg = bdev_get_queue(bio->bi_bdev)->root_blkg;
+	if (!css || !css_tryget_online(css)) {
+		css = &blkcg_root.css;
+		css_get(css);
 	}
+
+	blkcg = css_to_blkcg(css);
+	if (bio_blkcg(bio) == blkcg) {
+		css_put(css);
+		return;
+	}
+
+	bio_clear_blkcg(bio);
+	bio->bi_blkcg = blkcg;
 }
-EXPORT_SYMBOL_GPL(bio_associate_blkg_from_css);
+EXPORT_SYMBOL_GPL(bio_associate_blkcg_from_css);
 
 /**
- * bio_associate_blkg - associate a bio with a blkg
+ * bio_associate_blkcg - associate a bio with a blkcg
  * @bio: target bio
  *
- * Associate @bio with the blkg found from the bio's css and request_queue.
- * If one is not found, bio_lookup_blkg() creates the blkg.  If a blkg is
- * already associated, the css is reused and association redone as the
- * request_queue may have changed.
+ * Associate @bio with the blkcg found from the bio's css.  If a blkcg is
+ * already associated, keep it as blkcg association is not queue-local.
  */
-void bio_associate_blkg(struct bio *bio)
+void bio_associate_blkcg(struct bio *bio)
 {
-	struct cgroup_subsys_state *css;
-
 	if (blk_op_is_passthrough(bio->bi_opf))
 		return;
 
-	if (bio->bi_blkg) {
-		css = bio_blkcg_css(bio);
-		bio_associate_blkg_from_css(bio, css);
-	} else {
-		rcu_read_lock();
-		css = blkcg_css();
-		if (!css_tryget_online(css))
-			css = NULL;
-		rcu_read_unlock();
+	if (bio_blkcg(bio))
+		return;
 
-		bio_associate_blkg_from_css(bio, css);
-		if (css)
-			css_put(css);
-	}
+	rcu_read_lock();
+	bio_associate_blkcg_from_css(bio, blkcg_css());
+	rcu_read_unlock();
 }
-EXPORT_SYMBOL_GPL(bio_associate_blkg);
+EXPORT_SYMBOL_GPL(bio_associate_blkcg);
 
 /**
- * bio_clone_blkg_association - clone blkg association from src to dst bio
+ * bio_clone_blkcg_association - clone blkcg association from src to dst bio
  * @dst: destination bio
  * @src: source bio
  */
-void bio_clone_blkg_association(struct bio *dst, struct bio *src)
+void bio_clone_blkcg_association(struct bio *dst, struct bio *src)
 {
-	if (src->bi_blkg)
-		bio_associate_blkg_from_css(dst, bio_blkcg_css(src));
+	if (bio_blkcg(src))
+		bio_associate_blkcg_from_css(dst, bio_blkcg_css(src));
 }
-EXPORT_SYMBOL_GPL(bio_clone_blkg_association);
+EXPORT_SYMBOL_GPL(bio_clone_blkcg_association);
 
 static int blk_cgroup_io_type(struct bio *bio)
 {
@@ -2181,20 +2229,27 @@ static int blk_cgroup_io_type(struct bio *bio)
 
 void blk_cgroup_bio_start(struct bio *bio)
 {
-	struct blkcg *blkcg = bio->bi_blkg->blkcg;
+	struct blkcg *blkcg = bio_blkcg(bio);
+	struct blkcg_gq *blkg;
 	int rwd = blk_cgroup_io_type(bio), cpu;
 	struct blkg_iostat_set *bis;
 	unsigned long flags;
 
 	if (!cgroup_subsys_on_dfl(io_cgrp_subsys))
 		return;
+	if (!blkcg)
+		return;
 
 	/* Root-level stats are sourced from system-wide IO stats */
 	if (!cgroup_parent(blkcg->css.cgroup))
 		return;
 
+	blkg = bio_blkg_lookup(bio);
+	if (!blkg)
+		return;
+
 	cpu = get_cpu();
-	bis = per_cpu_ptr(bio->bi_blkg->iostat_cpu, cpu);
+	bis = per_cpu_ptr(blkg->iostat_cpu, cpu);
 	flags = u64_stats_update_begin_irqsave(&bis->sync);
 
 	/*
diff --git a/block/blk-cgroup.h b/block/blk-cgroup.h
index 615390f751aa2..08e54d9430edf 100644
--- a/block/blk-cgroup.h
+++ b/block/blk-cgroup.h
@@ -19,6 +19,8 @@
 #include <linux/kthread.h>
 #include <linux/blk-mq.h>
 #include <linux/llist.h>
+#include <linux/rhashtable.h>
+#include <linux/rcupdate.h>
 #include "blk.h"
 
 struct blkcg_gq;
@@ -56,9 +58,11 @@ struct blkg_iostat_set {
 struct blkcg_gq {
 	/* Pointer to the associated request_queue */
 	struct request_queue		*q;
+	struct rhash_head		q_hash_node;
 	struct list_head		q_node;
 	struct hlist_node		blkcg_node;
 	struct blkcg			*blkcg;
+	int				blkcg_id;
 
 	/* all non-root blkcg_gq's are guaranteed to have access to parent */
 	struct blkcg_gq			*parent;
@@ -73,14 +77,7 @@ struct blkcg_gq {
 	struct blkg_iostat_set		iostat;
 
 	struct blkg_policy_data		*pd[BLKCG_MAX_POLS];
-#ifdef CONFIG_BLK_CGROUP_PUNT_BIO
-	spinlock_t			async_bio_lock;
-	struct bio_list			async_bios;
-#endif
-	union {
-		struct work_struct	async_bio_work;
-		struct work_struct	free_work;
-	};
+	struct work_struct		free_work;
 
 	atomic_t			use_delay;
 	atomic64_t			delay_nsec;
@@ -98,8 +95,6 @@ struct blkcg {
 	/* If there is block congestion on this cgroup. */
 	atomic_t			congestion_count;
 
-	struct radix_tree_root		blkg_tree;
-	struct blkcg_gq	__rcu		*blkg_hint;
 	struct hlist_head		blkg_list;
 
 	struct blkcg_policy_data	*cpd[BLKCG_MAX_POLS];
@@ -111,6 +106,11 @@ struct blkcg {
 	 */
 	struct llist_head __percpu	*lhead;
 
+#ifdef CONFIG_BLK_CGROUP_PUNT_BIO
+	spinlock_t			async_bio_lock; /* protects async_bios */
+	struct bio_list			async_bios;
+	struct work_struct		async_bio_work;
+#endif
 #ifdef CONFIG_BLK_CGROUP_FC_APPID
 	char                            fc_app_id[FC_APPID_LEN];
 #endif
@@ -124,6 +124,11 @@ static inline struct blkcg *css_to_blkcg(struct cgroup_subsys_state *css)
 	return css ? container_of(css, struct blkcg, css) : NULL;
 }
 
+static inline struct blkcg *bio_blkcg(struct bio *bio)
+{
+	return bio->bi_blkcg;
+}
+
 /*
  * A blkcg_gq (blkg) is association between a block cgroup (blkcg) and a
  * request_queue (q).  This is used by blkcg policies which need to track
@@ -192,8 +197,10 @@ struct blkcg_policy {
 
 extern struct blkcg blkcg_root;
 extern bool blkcg_debug_stats;
+extern const struct rhashtable_params blkg_hash_params;
 
-void blkg_init_queue(struct request_queue *q);
+int blkg_init_queue(struct request_queue *q);
+void blkg_exit_queue(struct request_queue *q);
 int blkcg_init_disk(struct gendisk *disk);
 void blkcg_exit_disk(struct gendisk *disk);
 
@@ -249,13 +256,36 @@ static inline bool bio_issue_as_root_blkg(struct bio *bio)
 }
 
 /**
- * blkg_lookup - lookup blkg for the specified blkcg - q pair
+ * blkg_lookup_any - lookup any blkg for the specified blkcg - q pair
  * @blkcg: blkcg of interest
  * @q: request_queue of interest
  *
- * Lookup blkg for the @blkcg - @q pair.
+ * Lookup a blkg for the @blkcg - @q pair, whether it is online or dying.
  *
- * Must be called in a RCU critical section.
+ * Must be called with either RCU read lock or queue_lock held.
+ *
+ * This does not acquire a reference.  The caller must already hold one, or
+ * have the blkg pinned by I/O.
+ */
+static inline struct blkcg_gq *blkg_lookup_any(struct blkcg *blkcg,
+					       struct request_queue *q)
+{
+	RCU_LOCKDEP_WARN(!rcu_read_lock_held() &&
+			 !lockdep_is_held(&q->queue_lock),
+			 "blkg_lookup_any() requires an RCU read lock or queue_lock");
+
+	return rhashtable_lookup(&q->blkg_hash, &blkcg->css.id,
+				 blkg_hash_params);
+}
+
+/**
+ * blkg_lookup - lookup an online blkg for the specified blkcg - q pair
+ * @blkcg: blkcg of interest
+ * @q: request_queue of interest
+ *
+ * Lookup an online blkg for the @blkcg - @q pair.
+ *
+ * Must be called with either RCU read lock or queue_lock held.
  */
 static inline struct blkcg_gq *blkg_lookup(struct blkcg *blkcg,
 					   struct request_queue *q)
@@ -265,17 +295,15 @@ static inline struct blkcg_gq *blkg_lookup(struct blkcg *blkcg,
 	if (blkcg == &blkcg_root)
 		return q->root_blkg;
 
-	blkg = rcu_dereference_check(blkcg->blkg_hint,
-			lockdep_is_held(&q->queue_lock));
-	if (blkg && blkg->q == q)
-		return blkg;
-
-	blkg = radix_tree_lookup(&blkcg->blkg_tree, q->id);
-	if (blkg && blkg->q != q)
+	blkg = blkg_lookup_any(blkcg, q);
+	if (blkg && !READ_ONCE(blkg->online))
 		blkg = NULL;
 	return blkg;
 }
 
+struct blkcg_gq *bio_blkg_lookup(struct bio *bio);
+struct blkcg_gq *bio_blkg(struct bio *bio);
+
 /**
  * blkg_to_pd - get policy private data
  * @blkg: blkg of interest
@@ -343,6 +371,18 @@ static inline void blkg_put(struct blkcg_gq *blkg)
 	percpu_ref_put(&blkg->refcnt);
 }
 
+static inline void bio_clear_blkcg(struct bio *bio)
+{
+	struct blkcg *blkcg = bio_blkcg(bio);
+
+	bio_put_blkg_ref(bio);
+
+	if (blkcg) {
+		css_put(&blkcg->css);
+		bio->bi_blkcg = NULL;
+	}
+}
+
 /**
  * blkg_for_each_descendant_pre - pre-order walk of a blkg's descendants
  * @d_blkg: loop cursor pointing to the current descendant
@@ -455,7 +495,7 @@ static inline void blkcg_clear_delay(struct blkcg_gq *blkg)
  */
 static inline bool blk_cgroup_mergeable(struct request *rq, struct bio *bio)
 {
-	return rq->bio->bi_blkg == bio->bi_blkg &&
+	return bio_blkcg(rq->bio) == bio_blkcg(bio) &&
 		bio_issue_as_root_blkg(rq->bio) == bio_issue_as_root_blkg(bio);
 }
 
@@ -481,8 +521,13 @@ struct blkcg_policy {
 struct blkcg {
 };
 
+static inline struct blkcg *bio_blkcg(struct bio *bio) { return NULL; }
+static inline struct blkcg_gq *bio_blkg_lookup(struct bio *bio) { return NULL; }
+static inline struct blkcg_gq *bio_blkg(struct bio *bio) { return NULL; }
+static inline struct blkcg_gq *blkg_lookup_any(struct blkcg *blkcg, void *key) { return NULL; }
 static inline struct blkcg_gq *blkg_lookup(struct blkcg *blkcg, void *key) { return NULL; }
-static inline void blkg_init_queue(struct request_queue *q) { }
+static inline int blkg_init_queue(struct request_queue *q) { return 0; }
+static inline void blkg_exit_queue(struct request_queue *q) { }
 static inline int blkcg_init_disk(struct gendisk *disk) { return 0; }
 static inline void blkcg_exit_disk(struct gendisk *disk) { }
 static inline int blkcg_policy_register(struct blkcg_policy *pol) { return 0; }
@@ -497,6 +542,7 @@ static inline struct blkg_policy_data *blkg_to_pd(struct blkcg_gq *blkg,
 static inline struct blkcg_gq *pd_to_blkg(struct blkg_policy_data *pd) { return NULL; }
 static inline void blkg_get(struct blkcg_gq *blkg) { }
 static inline void blkg_put(struct blkcg_gq *blkg) { }
+static inline void bio_clear_blkcg(struct bio *bio) { }
 static inline void blk_cgroup_bio_start(struct bio *bio) { }
 static inline bool blk_cgroup_mergeable(struct request *rq, struct bio *bio) { return true; }
 
diff --git a/block/blk-core.c b/block/blk-core.c
index 365641266c9e8..7063e72465400 100644
--- a/block/blk-core.c
+++ b/block/blk-core.c
@@ -303,6 +303,7 @@ static void blk_free_queue(struct request_queue *q)
 	if (queue_is_mq(q))
 		blk_mq_release(q);
 
+	blkg_exit_queue(q);
 	ida_free(&blk_queue_ida, q->id);
 	lockdep_unregister_key(&q->io_lock_cls_key);
 	lockdep_unregister_key(&q->q_lock_cls_key);
@@ -481,7 +482,9 @@ struct request_queue *blk_alloc_queue(struct queue_limits *lim, int node_id)
 	init_waitqueue_head(&q->mq_freeze_wq);
 	mutex_init(&q->mq_freeze_lock);
 
-	blkg_init_queue(q);
+	error = blkg_init_queue(q);
+	if (error)
+		goto fail_stats;
 
 	/*
 	 * Init percpu_ref in atomic mode so that it's faster to shutdown.
@@ -491,7 +494,7 @@ struct request_queue *blk_alloc_queue(struct queue_limits *lim, int node_id)
 				blk_queue_usage_counter_release,
 				PERCPU_REF_INIT_ATOMIC, GFP_KERNEL);
 	if (error)
-		goto fail_stats;
+		goto fail_blkg;
 	lockdep_register_key(&q->io_lock_cls_key);
 	lockdep_register_key(&q->q_lock_cls_key);
 	lockdep_init_map(&q->io_lockdep_map, "&q->q_usage_counter(io)",
@@ -510,6 +513,8 @@ struct request_queue *blk_alloc_queue(struct queue_limits *lim, int node_id)
 
 	return q;
 
+fail_blkg:
+	blkg_exit_queue(q);
 fail_stats:
 	blk_free_queue_stats(q->stats);
 fail_id:
diff --git a/block/blk-crypto-fallback.c b/block/blk-crypto-fallback.c
index 2a8f40a65158b..a348c74d53297 100644
--- a/block/blk-crypto-fallback.c
+++ b/block/blk-crypto-fallback.c
@@ -187,7 +187,7 @@ static struct bio *blk_crypto_alloc_enc_bio(struct bio *bio_src,
 	bio->bi_write_hint	= bio_src->bi_write_hint;
 	bio->bi_write_stream	= bio_src->bi_write_stream;
 	bio->bi_iter.bi_sector	= bio_src->bi_iter.bi_sector;
-	bio_clone_blkg_association(bio, bio_src);
+	bio_clone_blkcg_association(bio, bio_src);
 
 	/*
 	 * Move page array up in the allocated memory for the bio vecs as far as
diff --git a/block/blk-iocost.c b/block/blk-iocost.c
index 8b2aeba2e1e39..3d1679c1101d1 100644
--- a/block/blk-iocost.c
+++ b/block/blk-iocost.c
@@ -2686,7 +2686,7 @@ iocg_handle_over_budget(struct rq_qos *rqos, struct ioc_gq *iocg,
 
 static void ioc_rqos_throttle(struct rq_qos *rqos, struct bio *bio)
 {
-	struct blkcg_gq *blkg = bio->bi_blkg;
+	struct blkcg_gq *blkg = bio_blkg(bio);
 	struct ioc *ioc = rqos_to_ioc(rqos);
 	struct ioc_gq *iocg = blkg_to_iocg(blkg);
 	struct ioc_now now;
@@ -2775,7 +2775,7 @@ static void ioc_rqos_throttle(struct rq_qos *rqos, struct bio *bio)
 static void ioc_rqos_merge(struct rq_qos *rqos, struct request *rq,
 			   struct bio *bio)
 {
-	struct ioc_gq *iocg = blkg_to_iocg(bio->bi_blkg);
+	struct ioc_gq *iocg = blkg_to_iocg(bio_blkg(bio));
 	struct ioc *ioc = rqos_to_ioc(rqos);
 	sector_t bio_end = bio_end_sector(bio);
 	struct ioc_now now;
@@ -2833,9 +2833,13 @@ static void ioc_rqos_merge(struct rq_qos *rqos, struct request *rq,
 
 static void ioc_rqos_done_bio(struct rq_qos *rqos, struct bio *bio)
 {
-	struct ioc_gq *iocg = blkg_to_iocg(bio->bi_blkg);
+	struct ioc_gq *iocg;
+
+	if (!bio->bi_iocost_cost)
+		return;
 
-	if (iocg && bio->bi_iocost_cost)
+	iocg = blkg_to_iocg(bio_blkg_lookup(bio));
+	if (iocg)
 		atomic64_add(bio->bi_iocost_cost, &iocg->done_vtime);
 }
 
diff --git a/block/blk-iolatency.c b/block/blk-iolatency.c
index cef02b6c5fa91..7ad18a538d7e2 100644
--- a/block/blk-iolatency.c
+++ b/block/blk-iolatency.c
@@ -463,7 +463,7 @@ static void check_scale_change(struct iolatency_grp *iolat)
 static void blkcg_iolatency_throttle(struct rq_qos *rqos, struct bio *bio)
 {
 	struct blk_iolatency *blkiolat = BLKIOLATENCY(rqos);
-	struct blkcg_gq *blkg = bio->bi_blkg;
+	struct blkcg_gq *blkg = bio_blkg(bio);
 	bool issue_as_root = bio_issue_as_root_blkg(bio);
 
 	if (!blkiolat->enabled)
@@ -590,11 +590,14 @@ static void blkcg_iolatency_done_bio(struct rq_qos *rqos, struct bio *bio)
 	bool issue_as_root = bio_issue_as_root_blkg(bio);
 	int inflight = 0;
 
-	blkg = bio->bi_blkg;
-	if (!blkg || !bio_flagged(bio, BIO_QOS_THROTTLED))
+	if (!bio_flagged(bio, BIO_QOS_THROTTLED))
 		return;
 
-	iolat = blkg_to_lat(bio->bi_blkg);
+	blkg = bio_blkg_lookup(bio);
+	if (!blkg)
+		return;
+
+	iolat = blkg_to_lat(blkg);
 	if (!iolat)
 		return;
 
diff --git a/block/blk-ioprio.c b/block/blk-ioprio.c
index 8fa8bca350629..5e3f8e49550da 100644
--- a/block/blk-ioprio.c
+++ b/block/blk-ioprio.c
@@ -132,7 +132,7 @@ static struct blkcg_policy ioprio_policy = {
 
 void blkcg_set_ioprio(struct bio *bio)
 {
-	struct ioprio_blkcg *blkcg = blkcg_to_ioprio_blkcg(bio->bi_blkg->blkcg);
+	struct ioprio_blkcg *blkcg = blkcg_to_ioprio_blkcg(bio_blkcg(bio));
 	u16 prio;
 
 	if (!blkcg || blkcg->prio_policy == POLICY_NO_CHANGE)
diff --git a/block/blk-throttle.c b/block/blk-throttle.c
index ffc3b70065d4b..3828c38579002 100644
--- a/block/blk-throttle.c
+++ b/block/blk-throttle.c
@@ -1792,7 +1792,7 @@ static bool tg_within_limit(struct throtl_grp *tg, struct bio *bio, bool rw)
 bool __blk_throtl_bio(struct bio *bio)
 {
 	struct request_queue *q = bdev_get_queue(bio->bi_bdev);
-	struct blkcg_gq *blkg = bio->bi_blkg;
+	struct blkcg_gq *blkg = bio_blkg(bio);
 	struct throtl_qnode *qn = NULL;
 	struct throtl_grp *tg = blkg_to_tg(blkg);
 	struct throtl_service_queue *sq;
diff --git a/block/blk-throttle.h b/block/blk-throttle.h
index 9d7a42c039a15..609a126c7ccb2 100644
--- a/block/blk-throttle.h
+++ b/block/blk-throttle.h
@@ -173,7 +173,7 @@ static inline bool blk_should_throtl(struct bio *bio)
 	if (!blk_throtl_activated(bio->bi_bdev->bd_queue))
 		return false;
 
-	tg = blkg_to_tg(bio->bi_blkg);
+	tg = blkg_to_tg(bio_blkg(bio));
 	if (!cgroup_subsys_on_dfl(io_cgrp_subsys)) {
 		if (!bio_flagged(bio, BIO_CGROUP_ACCT)) {
 			bio_set_flag(bio, BIO_CGROUP_ACCT);
diff --git a/drivers/md/bcache/request.c b/drivers/md/bcache/request.c
index 3fa3b13a410f4..c0f945b8d9413 100644
--- a/drivers/md/bcache/request.c
+++ b/drivers/md/bcache/request.c
@@ -848,7 +848,7 @@ static CLOSURE_CALLBACK(cached_dev_read_done)
 		s->iop.bio->bi_iter.bi_sector =
 			s->cache_miss->bi_iter.bi_sector;
 		s->iop.bio->bi_iter.bi_size = s->insert_bio_sectors << 9;
-		bio_clone_blkg_association(s->iop.bio, s->cache_miss);
+		bio_clone_blkcg_association(s->iop.bio, s->cache_miss);
 		bch_bio_map(s->iop.bio, NULL);
 
 		bio_copy_data(s->cache_miss, s->iop.bio);
diff --git a/drivers/md/dm.c b/drivers/md/dm.c
index d413bfaf35278..cd68eec77f5a1 100644
--- a/drivers/md/dm.c
+++ b/drivers/md/dm.c
@@ -1373,7 +1373,7 @@ void dm_submit_bio_remap(struct bio *clone, struct bio *tgt_clone)
 	if (!tgt_clone)
 		tgt_clone = clone;
 
-	bio_clone_blkg_association(tgt_clone, io->orig_bio);
+	bio_clone_blkcg_association(tgt_clone, io->orig_bio);
 
 	/*
 	 * Account io->origin_bio to DM dev on behalf of target
diff --git a/drivers/md/md.c b/drivers/md/md.c
index d1465bcd86c81..af55f8efa46bc 100644
--- a/drivers/md/md.c
+++ b/drivers/md/md.c
@@ -9355,7 +9355,7 @@ void md_submit_discard_bio(struct mddev *mddev, struct md_rdev *rdev,
 		return;
 
 	bio_chain(discard_bio, bio);
-	bio_clone_blkg_association(discard_bio, bio);
+	bio_clone_blkcg_association(discard_bio, bio);
 	mddev_trace_remap(mddev, discard_bio, bio->bi_iter.bi_sector);
 	submit_bio_noacct(discard_bio);
 }
diff --git a/drivers/nvdimm/nd_virtio.c b/drivers/nvdimm/nd_virtio.c
index 4176046627beb..54e4adb0ccb7d 100644
--- a/drivers/nvdimm/nd_virtio.c
+++ b/drivers/nvdimm/nd_virtio.c
@@ -121,7 +121,7 @@ int async_pmem_flush(struct nd_region *nd_region, struct bio *bio)
 
 		if (!child)
 			return -ENOMEM;
-		bio_clone_blkg_association(child, bio);
+		bio_clone_blkcg_association(child, bio);
 		child->bi_iter.bi_sector = -1;
 		bio_chain(child, bio);
 		submit_bio(child);
diff --git a/fs/gfs2/lops.c b/fs/gfs2/lops.c
index 6dabe73ad790d..6512dbd9516fc 100644
--- a/fs/gfs2/lops.c
+++ b/fs/gfs2/lops.c
@@ -484,7 +484,7 @@ static struct bio *gfs2_chain_bio(struct bio *prev, unsigned int nr_iovecs,
 	struct bio *new;
 
 	new = bio_alloc(prev->bi_bdev, nr_iovecs, opf, GFP_NOIO);
-	bio_clone_blkg_association(new, prev);
+	bio_clone_blkcg_association(new, prev);
 	new->bi_iter.bi_sector = sector;
 	bio_chain(new, prev);
 	submit_bio(prev);
@@ -1114,4 +1114,3 @@ const struct gfs2_log_operations *gfs2_log_ops[] = {
 	&gfs2_revoke_lops,
 	NULL,
 };
-
diff --git a/include/linux/bio.h b/include/linux/bio.h
index 8f33f717b14f5..fee0310457677 100644
--- a/include/linux/bio.h
+++ b/include/linux/bio.h
@@ -503,18 +503,20 @@ static inline void bio_release_pages(struct bio *bio, bool mark_dirty)
 	disk_devt((bio)->bi_bdev->bd_disk)
 
 #ifdef CONFIG_BLK_CGROUP
-void bio_associate_blkg(struct bio *bio);
-void bio_associate_blkg_from_css(struct bio *bio,
-				 struct cgroup_subsys_state *css);
-void bio_clone_blkg_association(struct bio *dst, struct bio *src);
+void bio_associate_blkcg(struct bio *bio);
+void bio_associate_blkcg_from_css(struct bio *bio,
+				  struct cgroup_subsys_state *css);
+void bio_clone_blkcg_association(struct bio *dst, struct bio *src);
+void bio_put_blkg_ref(struct bio *bio);
 void blkcg_punt_bio_submit(struct bio *bio);
 #else	/* CONFIG_BLK_CGROUP */
-static inline void bio_associate_blkg(struct bio *bio) { }
-static inline void bio_associate_blkg_from_css(struct bio *bio,
-					       struct cgroup_subsys_state *css)
+static inline void bio_associate_blkcg(struct bio *bio) { }
+static inline void bio_associate_blkcg_from_css(struct bio *bio,
+						struct cgroup_subsys_state *css)
 { }
-static inline void bio_clone_blkg_association(struct bio *dst,
-					      struct bio *src) { }
+static inline void bio_clone_blkcg_association(struct bio *dst,
+					       struct bio *src) { }
+static inline void bio_put_blkg_ref(struct bio *bio) { }
 static inline void blkcg_punt_bio_submit(struct bio *bio)
 {
 	submit_bio(bio);
@@ -524,10 +526,12 @@ static inline void blkcg_punt_bio_submit(struct bio *bio)
 static inline void bio_set_dev(struct bio *bio, struct block_device *bdev)
 {
 	bio_clear_flag(bio, BIO_REMAPPED);
-	if (bio->bi_bdev != bdev)
+	if (bio->bi_bdev != bdev) {
+		bio_put_blkg_ref(bio);
 		bio_clear_flag(bio, BIO_BPS_THROTTLED);
+	}
 	bio->bi_bdev = bdev;
-	bio_associate_blkg(bio);
+	bio_associate_blkcg(bio);
 }
 
 /*
diff --git a/include/linux/blk_types.h b/include/linux/blk_types.h
index 5a725a0cd35f4..bc2da8c6fc38f 100644
--- a/include/linux/blk_types.h
+++ b/include/linux/blk_types.h
@@ -246,12 +246,10 @@ struct bio {
 	void			*bi_private;
 #ifdef CONFIG_BLK_CGROUP
 	/*
-	 * Represents the association of the css and request_queue for the bio.
-	 * If a bio goes direct to device, it will not have a blkg as it will
-	 * not have a request_queue associated with it.  The reference is put
-	 * on release of the bio.
+	 * Represents the blkcg css association for the bio.  The reference is
+	 * put on release of the bio.
 	 */
-	struct blkcg_gq		*bi_blkg;
+	struct blkcg		*bi_blkcg;
 	/* Time that this bio was issued. */
 	u64			issue_time_ns;
 #ifdef CONFIG_BLK_CGROUP_IOCOST
@@ -309,6 +307,7 @@ enum {
 	BIO_TRACE_COMPLETION,	/* bio_endio() should trace the final completion
 				 * of this bio. */
 	BIO_CGROUP_ACCT,	/* has been accounted to a cgroup */
+	BIO_BLKG_REF,		/* bio pins the associated blkg */
 	BIO_QOS_THROTTLED,	/* bio went through rq_qos throttle path */
 	/*
 	 * This bio has completed bps throttling at the single tg granularity,
diff --git a/include/linux/blkdev.h b/include/linux/blkdev.h
index dbb549cdfb778..928cd69b9a356 100644
--- a/include/linux/blkdev.h
+++ b/include/linux/blkdev.h
@@ -27,6 +27,7 @@
 #include <linux/xarray.h>
 #include <linux/file.h>
 #include <linux/lockdep.h>
+#include <linux/rhashtable-types.h>
 
 struct module;
 struct request_queue;
@@ -579,6 +580,7 @@ struct request_queue {
 #ifdef CONFIG_BLK_CGROUP
 	DECLARE_BITMAP		(blkcg_pols, BLKCG_MAX_POLS);
 	struct blkcg_gq		*root_blkg;
+	struct rhashtable	blkg_hash;
 	struct list_head	blkg_list;
 	struct mutex		blkcg_mutex;
 #endif
diff --git a/include/linux/writeback.h b/include/linux/writeback.h
index 62552a2ce5b9e..4f869fe9cc90f 100644
--- a/include/linux/writeback.h
+++ b/include/linux/writeback.h
@@ -262,7 +262,7 @@ static inline void wbc_init_bio(struct writeback_control *wbc, struct bio *bio)
 	 * regular writeback instead of writing things out itself.
 	 */
 	if (wbc->wb)
-		bio_associate_blkg_from_css(bio, wbc->wb->blkcg_css);
+		bio_associate_blkcg_from_css(bio, wbc->wb->blkcg_css);
 }
 
 void inode_switch_wbs_work_fn(struct work_struct *work);
diff --git a/mm/page_io.c b/mm/page_io.c
index b23f494fcc83d..a4a1b8a8a0e38 100644
--- a/mm/page_io.c
+++ b/mm/page_io.c
@@ -316,7 +316,7 @@ static inline void count_swpout_vm_event(struct folio *folio)
 }
 
 #if defined(CONFIG_MEMCG) && defined(CONFIG_BLK_CGROUP)
-static void bio_associate_blkg_from_page(struct bio *bio, struct folio *folio)
+static void bio_associate_blkcg_from_page(struct bio *bio, struct folio *folio)
 {
 	struct cgroup_subsys_state *css;
 	struct mem_cgroup *memcg;
@@ -331,12 +331,15 @@ static void bio_associate_blkg_from_page(struct bio *bio, struct folio *folio)
 		css = NULL;
 	rcu_read_unlock();
 
-	bio_associate_blkg_from_css(bio, css);
+	bio_associate_blkcg_from_css(bio, css);
 	if (css)
 		css_put(css);
 }
 #else
-#define bio_associate_blkg_from_page(bio, folio)		do { } while (0)
+static inline void bio_associate_blkcg_from_page(struct bio *bio,
+						 struct folio *folio)
+{
+}
 #endif /* CONFIG_MEMCG && CONFIG_BLK_CGROUP */
 
 struct swap_iocb {
@@ -436,7 +439,7 @@ static void swap_writepage_bdev_sync(struct folio *folio,
 	bio.bi_iter.bi_sector = swap_folio_sector(folio);
 	bio_add_folio_nofail(&bio, folio, folio_size(folio), 0);
 
-	bio_associate_blkg_from_page(&bio, folio);
+	bio_associate_blkcg_from_page(&bio, folio);
 	count_swpout_vm_event(folio);
 
 	folio_start_writeback(folio);
@@ -456,7 +459,7 @@ static void swap_writepage_bdev_async(struct folio *folio,
 	bio->bi_end_io = end_swap_bio_write;
 	bio_add_folio_nofail(bio, folio, folio_size(folio), 0);
 
-	bio_associate_blkg_from_page(bio, folio);
+	bio_associate_blkcg_from_page(bio, folio);
 	count_swpout_vm_event(folio);
 	folio_start_writeback(folio);
 	folio_unlock(folio);