+void blk_set_preempt_only(struct request_queue *q, bool preempt_only)
+{
+	blk_mq_freeze_queue(q);
+	if (preempt_only)
+		queue_flag_set_unlocked(QUEUE_FLAG_PREEMPT_ONLY, q);
+	else
+		queue_flag_clear_unlocked(QUEUE_FLAG_PREEMPT_ONLY, q);
+	blk_mq_unfreeze_queue(q);
+}
+EXPORT_SYMBOL(blk_set_preempt_only);

What do you need the queue freeze for?

+		/*
+		 * preempt_only flag has to be set after queue is frozen,
+		 * so it can be checked here lockless and safely
+		 */
+		if (blk_queue_preempt_only(q)) {

We can always check a single bit flag safely, so I really don't
understand that comment.

+			if (!(flags & BLK_REQ_PREEMPT))
+				goto slow_path;
+		}
+
 		if (percpu_ref_tryget_live(&q->q_usage_counter))
 			return 0;
-
+ slow_path:

Also this looks a very spaghetti, why not:


	if (!blk_queue_preempt_only(q) || (flags & BLK_MQ_REQ_PREEMPT)) {
		if (percpu_ref_tryget_live(&q->q_usage_counter))
			return 0;
	}