[PATCH 4/4] slub: apply new queue_percpu_work_on() interface

[Marcelo Tosatti <mtosatti@redhat.com>]

Make use of the new qpw_{un,}lock*() and queue_percpu_work_on()
interface to improve performance & latency on PREEMPT_RT kernels.

For functions that may be scheduled in a different cpu, replace
local_{un,}lock*() by qpw_{un,}lock*(), and replace schedule_work_on() by
queue_percpu_work_on(). The same happens for flush_work() and
flush_percpu_work().

This change requires allocation of qpw_structs instead of a work_structs,
and changing parameters of a few functions to include the cpu parameter.

This should bring no relevant performance impact on non-RT kernels:
For functions that may be scheduled in a different cpu, the local_*lock's
this_cpu_ptr() becomes a per_cpu_ptr(smp_processor_id()).

Signed-off-by: Leonardo Bras <leobras@redhat.com>
Signed-off-by: Marcelo Tosatti <mtosatti@redhat.com>

---
 mm/slub.c |  218 ++++++++++++++++++++++++++++++++++++++++----------------------
 1 file changed, 142 insertions(+), 76 deletions(-)

Index: slab/mm/slub.c
===================================================================
--- slab.orig/mm/slub.c
+++ slab/mm/slub.c
@@ -49,6 +49,7 @@
 #include <linux/irq_work.h>
 #include <linux/kprobes.h>
 #include <linux/debugfs.h>
+#include <linux/qpw.h>
 #include <trace/events/kmem.h>
 
 #include "internal.h"
@@ -128,7 +129,7 @@
  *   For debug caches, all allocations are forced to go through a list_lock
  *   protected region to serialize against concurrent validation.
  *
- *   cpu_sheaves->lock (local_trylock)
+ *   cpu_sheaves->lock (qpw_trylock)
  *
  *   This lock protects fastpath operations on the percpu sheaves. On !RT it
  *   only disables preemption and does no atomic operations. As long as the main
@@ -156,7 +157,7 @@
  *   Interrupts are disabled as part of list_lock or barn lock operations, or
  *   around the slab_lock operation, in order to make the slab allocator safe
  *   to use in the context of an irq.
- *   Preemption is disabled as part of local_trylock operations.
+ *   Preemption is disabled as part of qpw_trylock operations.
  *   kmalloc_nolock() and kfree_nolock() are safe in NMI context but see
  *   their limitations.
  *
@@ -417,7 +418,7 @@ struct slab_sheaf {
 };
 
 struct slub_percpu_sheaves {
-	local_trylock_t lock;
+	qpw_trylock_t lock;
 	struct slab_sheaf *main; /* never NULL when unlocked */
 	struct slab_sheaf *spare; /* empty or full, may be NULL */
 	struct slab_sheaf *rcu_free; /* for batching kfree_rcu() */
@@ -479,7 +480,7 @@ static nodemask_t slab_nodes;
 static struct workqueue_struct *flushwq;
 
 struct slub_flush_work {
-	struct work_struct work;
+	struct qpw_struct qpw;
 	struct kmem_cache *s;
 	bool skip;
 };
@@ -2826,7 +2827,7 @@ static void __kmem_cache_free_bulk(struc
  *
  * returns true if at least partially flushed
  */
-static bool sheaf_flush_main(struct kmem_cache *s)
+static bool sheaf_flush_main(struct kmem_cache *s, int cpu)
 {
 	struct slub_percpu_sheaves *pcs;
 	unsigned int batch, remaining;
@@ -2835,10 +2836,10 @@ static bool sheaf_flush_main(struct kmem
 	bool ret = false;
 
 next_batch:
-	if (!local_trylock(&s->cpu_sheaves->lock))
+	if (!qpw_trylock(&s->cpu_sheaves->lock, cpu))
 		return ret;
 
-	pcs = this_cpu_ptr(s->cpu_sheaves);
+	pcs = per_cpu_ptr(s->cpu_sheaves, cpu);
 	sheaf = pcs->main;
 
 	batch = min(PCS_BATCH_MAX, sheaf->size);
@@ -2848,7 +2849,7 @@ next_batch:
 
 	remaining = sheaf->size;
 
-	local_unlock(&s->cpu_sheaves->lock);
+	qpw_unlock(&s->cpu_sheaves->lock, cpu);
 
 	__kmem_cache_free_bulk(s, batch, &objects[0]);
 
@@ -2932,13 +2933,13 @@ static void rcu_free_sheaf_nobarn(struct
  * flushing operations are rare so let's keep it simple and flush to slabs
  * directly, skipping the barn
  */
-static void pcs_flush_all(struct kmem_cache *s)
+static void pcs_flush_all(struct kmem_cache *s, int cpu)
 {
 	struct slub_percpu_sheaves *pcs;
 	struct slab_sheaf *spare, *rcu_free;
 
-	local_lock(&s->cpu_sheaves->lock);
-	pcs = this_cpu_ptr(s->cpu_sheaves);
+	qpw_lock(&s->cpu_sheaves->lock, cpu);
+	pcs = per_cpu_ptr(s->cpu_sheaves, cpu);
 
 	spare = pcs->spare;
 	pcs->spare = NULL;
@@ -2946,7 +2947,7 @@ static void pcs_flush_all(struct kmem_ca
 	rcu_free = pcs->rcu_free;
 	pcs->rcu_free = NULL;
 
-	local_unlock(&s->cpu_sheaves->lock);
+	qpw_unlock(&s->cpu_sheaves->lock, cpu);
 
 	if (spare) {
 		sheaf_flush_unused(s, spare);
@@ -2956,7 +2957,7 @@ static void pcs_flush_all(struct kmem_ca
 	if (rcu_free)
 		call_rcu(&rcu_free->rcu_head, rcu_free_sheaf_nobarn);
 
-	sheaf_flush_main(s);
+	sheaf_flush_main(s, cpu);
 }
 
 static void __pcs_flush_all_cpu(struct kmem_cache *s, unsigned int cpu)
@@ -3881,13 +3882,13 @@ static void flush_cpu_sheaves(struct wor
 {
 	struct kmem_cache *s;
 	struct slub_flush_work *sfw;
+	int cpu = qpw_get_cpu(w);
 
-	sfw = container_of(w, struct slub_flush_work, work);
-
+	sfw = &per_cpu(slub_flush, cpu);
 	s = sfw->s;
 
 	if (cache_has_sheaves(s))
-		pcs_flush_all(s);
+		pcs_flush_all(s, cpu);
 }
 
 static void flush_all_cpus_locked(struct kmem_cache *s)
@@ -3904,17 +3905,17 @@ static void flush_all_cpus_locked(struct
 			sfw->skip = true;
 			continue;
 		}
-		INIT_WORK(&sfw->work, flush_cpu_sheaves);
+		INIT_QPW(&sfw->qpw, flush_cpu_sheaves, cpu);
 		sfw->skip = false;
 		sfw->s = s;
-		queue_work_on(cpu, flushwq, &sfw->work);
+		queue_percpu_work_on(cpu, flushwq, &sfw->qpw);
 	}
 
 	for_each_online_cpu(cpu) {
 		sfw = &per_cpu(slub_flush, cpu);
 		if (sfw->skip)
 			continue;
-		flush_work(&sfw->work);
+		flush_percpu_work(&sfw->qpw);
 	}
 
 	mutex_unlock(&flush_lock);
@@ -3933,17 +3934,18 @@ static void flush_rcu_sheaf(struct work_
 	struct slab_sheaf *rcu_free;
 	struct slub_flush_work *sfw;
 	struct kmem_cache *s;
+	int cpu = qpw_get_cpu(w);
 
-	sfw = container_of(w, struct slub_flush_work, work);
+	sfw = &per_cpu(slub_flush, cpu);
 	s = sfw->s;
 
-	local_lock(&s->cpu_sheaves->lock);
-	pcs = this_cpu_ptr(s->cpu_sheaves);
+	qpw_lock(&s->cpu_sheaves->lock, cpu);
+	pcs = per_cpu_ptr(s->cpu_sheaves, cpu);
 
 	rcu_free = pcs->rcu_free;
 	pcs->rcu_free = NULL;
 
-	local_unlock(&s->cpu_sheaves->lock);
+	qpw_unlock(&s->cpu_sheaves->lock, cpu);
 
 	if (rcu_free)
 		call_rcu(&rcu_free->rcu_head, rcu_free_sheaf_nobarn);
@@ -3968,14 +3970,14 @@ void flush_rcu_sheaves_on_cache(struct k
 		 * sure the __kfree_rcu_sheaf() finished its call_rcu()
 		 */
 
-		INIT_WORK(&sfw->work, flush_rcu_sheaf);
+		INIT_QPW(&sfw->qpw, flush_rcu_sheaf, cpu);
 		sfw->s = s;
-		queue_work_on(cpu, flushwq, &sfw->work);
+		queue_percpu_work_on(cpu, flushwq, &sfw->qpw);
 	}
 
 	for_each_online_cpu(cpu) {
 		sfw = &per_cpu(slub_flush, cpu);
-		flush_work(&sfw->work);
+		flush_percpu_work(&sfw->qpw);
 	}
 
 	mutex_unlock(&flush_lock);
@@ -4472,22 +4474,24 @@ bool slab_post_alloc_hook(struct kmem_ca
  *
  * Must be called with the cpu_sheaves local lock locked. If successful, returns
  * the pcs pointer and the local lock locked (possibly on a different cpu than
- * initially called). If not successful, returns NULL and the local lock
- * unlocked.
+ * initially called), and migration disabled. If not successful, returns NULL
+ * and the local lock unlocked, with migration enabled.
  */
 static struct slub_percpu_sheaves *
-__pcs_replace_empty_main(struct kmem_cache *s, struct slub_percpu_sheaves *pcs, gfp_t gfp)
+__pcs_replace_empty_main(struct kmem_cache *s, struct slub_percpu_sheaves *pcs, gfp_t gfp,
+			 int *cpu)
 {
 	struct slab_sheaf *empty = NULL;
 	struct slab_sheaf *full;
 	struct node_barn *barn;
 	bool can_alloc;
 
-	lockdep_assert_held(this_cpu_ptr(&s->cpu_sheaves->lock));
+	qpw_lockdep_assert_held(&s->cpu_sheaves->lock);
 
 	/* Bootstrap or debug cache, back off */
 	if (unlikely(!cache_has_sheaves(s))) {
-		local_unlock(&s->cpu_sheaves->lock);
+		qpw_unlock(&s->cpu_sheaves->lock, *cpu);
+		migrate_enable();
 		return NULL;
 	}
 
@@ -4498,7 +4502,8 @@ __pcs_replace_empty_main(struct kmem_cac
 
 	barn = get_barn(s);
 	if (!barn) {
-		local_unlock(&s->cpu_sheaves->lock);
+		qpw_unlock(&s->cpu_sheaves->lock, *cpu);
+		migrate_enable();
 		return NULL;
 	}
 
@@ -4524,7 +4529,8 @@ __pcs_replace_empty_main(struct kmem_cac
 		}
 	}
 
-	local_unlock(&s->cpu_sheaves->lock);
+	qpw_unlock(&s->cpu_sheaves->lock, *cpu);
+	migrate_enable();
 
 	if (!can_alloc)
 		return NULL;
@@ -4550,7 +4556,9 @@ __pcs_replace_empty_main(struct kmem_cac
 	 * we can reach here only when gfpflags_allow_blocking
 	 * so this must not be an irq
 	 */
-	local_lock(&s->cpu_sheaves->lock);
+	migrate_disable();
+	*cpu = smp_processor_id();
+	qpw_lock(&s->cpu_sheaves->lock, *cpu);
 	pcs = this_cpu_ptr(s->cpu_sheaves);
 
 	/*
@@ -4593,6 +4601,7 @@ void *alloc_from_pcs(struct kmem_cache *
 	struct slub_percpu_sheaves *pcs;
 	bool node_requested;
 	void *object;
+	int cpu;
 
 #ifdef CONFIG_NUMA
 	if (static_branch_unlikely(&strict_numa) &&
@@ -4627,13 +4636,17 @@ void *alloc_from_pcs(struct kmem_cache *
 		return NULL;
 	}
 
-	if (!local_trylock(&s->cpu_sheaves->lock))
+	migrate_disable();
+	cpu = smp_processor_id();
+	if (!qpw_trylock(&s->cpu_sheaves->lock, cpu)) {
+		migrate_enable();
 		return NULL;
+	}
 
 	pcs = this_cpu_ptr(s->cpu_sheaves);
 
 	if (unlikely(pcs->main->size == 0)) {
-		pcs = __pcs_replace_empty_main(s, pcs, gfp);
+		pcs = __pcs_replace_empty_main(s, pcs, gfp, &cpu);
 		if (unlikely(!pcs))
 			return NULL;
 	}
@@ -4647,7 +4660,8 @@ void *alloc_from_pcs(struct kmem_cache *
 		 * the current allocation or previous freeing process.
 		 */
 		if (page_to_nid(virt_to_page(object)) != node) {
-			local_unlock(&s->cpu_sheaves->lock);
+			qpw_unlock(&s->cpu_sheaves->lock, cpu);
+			migrate_enable();
 			stat(s, ALLOC_NODE_MISMATCH);
 			return NULL;
 		}
@@ -4655,7 +4669,8 @@ void *alloc_from_pcs(struct kmem_cache *
 
 	pcs->main->size--;
 
-	local_unlock(&s->cpu_sheaves->lock);
+	qpw_unlock(&s->cpu_sheaves->lock, cpu);
+	migrate_enable();
 
 	stat(s, ALLOC_FASTPATH);
 
@@ -4670,10 +4685,15 @@ unsigned int alloc_from_pcs_bulk(struct
 	struct slab_sheaf *main;
 	unsigned int allocated = 0;
 	unsigned int batch;
+	int cpu;
 
 next_batch:
-	if (!local_trylock(&s->cpu_sheaves->lock))
+	migrate_disable();
+	cpu = smp_processor_id();
+	if (!qpw_trylock(&s->cpu_sheaves->lock, cpu)) {
+		migrate_enable();
 		return allocated;
+	}
 
 	pcs = this_cpu_ptr(s->cpu_sheaves);
 
@@ -4683,7 +4703,8 @@ next_batch:
 		struct node_barn *barn;
 
 		if (unlikely(!cache_has_sheaves(s))) {
-			local_unlock(&s->cpu_sheaves->lock);
+			qpw_unlock(&s->cpu_sheaves->lock, cpu);
+			migrate_enable();
 			return allocated;
 		}
 
@@ -4694,7 +4715,8 @@ next_batch:
 
 		barn = get_barn(s);
 		if (!barn) {
-			local_unlock(&s->cpu_sheaves->lock);
+			qpw_unlock(&s->cpu_sheaves->lock, cpu);
+			migrate_enable();
 			return allocated;
 		}
 
@@ -4709,7 +4731,8 @@ next_batch:
 
 		stat(s, BARN_GET_FAIL);
 
-		local_unlock(&s->cpu_sheaves->lock);
+		qpw_unlock(&s->cpu_sheaves->lock, cpu);
+		migrate_enable();
 
 		/*
 		 * Once full sheaves in barn are depleted, let the bulk
@@ -4727,7 +4750,8 @@ do_alloc:
 	main->size -= batch;
 	memcpy(p, main->objects + main->size, batch * sizeof(void *));
 
-	local_unlock(&s->cpu_sheaves->lock);
+	qpw_unlock(&s->cpu_sheaves->lock, cpu);
+	migrate_enable();
 
 	stat_add(s, ALLOC_FASTPATH, batch);
 
@@ -4877,6 +4901,7 @@ kmem_cache_prefill_sheaf(struct kmem_cac
 	struct slub_percpu_sheaves *pcs;
 	struct slab_sheaf *sheaf = NULL;
 	struct node_barn *barn;
+	int cpu;
 
 	if (unlikely(!size))
 		return NULL;
@@ -4906,7 +4931,9 @@ kmem_cache_prefill_sheaf(struct kmem_cac
 		return sheaf;
 	}
 
-	local_lock(&s->cpu_sheaves->lock);
+	migrate_disable();
+	cpu = smp_processor_id();
+	qpw_lock(&s->cpu_sheaves->lock, cpu);
 	pcs = this_cpu_ptr(s->cpu_sheaves);
 
 	if (pcs->spare) {
@@ -4925,7 +4952,8 @@ kmem_cache_prefill_sheaf(struct kmem_cac
 			stat(s, BARN_GET_FAIL);
 	}
 
-	local_unlock(&s->cpu_sheaves->lock);
+	qpw_unlock(&s->cpu_sheaves->lock, cpu);
+	migrate_enable();
 
 
 	if (!sheaf)
@@ -4961,6 +4989,7 @@ void kmem_cache_return_sheaf(struct kmem
 {
 	struct slub_percpu_sheaves *pcs;
 	struct node_barn *barn;
+	int cpu;
 
 	if (unlikely((sheaf->capacity != s->sheaf_capacity)
 		     || sheaf->pfmemalloc)) {
@@ -4969,7 +4998,9 @@ void kmem_cache_return_sheaf(struct kmem
 		return;
 	}
 
-	local_lock(&s->cpu_sheaves->lock);
+	migrate_disable();
+	cpu = smp_processor_id();
+	qpw_lock(&s->cpu_sheaves->lock, cpu);
 	pcs = this_cpu_ptr(s->cpu_sheaves);
 	barn = get_barn(s);
 
@@ -4979,7 +5010,8 @@ void kmem_cache_return_sheaf(struct kmem
 		stat(s, SHEAF_RETURN_FAST);
 	}
 
-	local_unlock(&s->cpu_sheaves->lock);
+	qpw_unlock(&s->cpu_sheaves->lock, cpu);
+	migrate_enable();
 
 	if (!sheaf)
 		return;
@@ -5507,9 +5539,9 @@ slab_empty:
  */
 static void __pcs_install_empty_sheaf(struct kmem_cache *s,
 		struct slub_percpu_sheaves *pcs, struct slab_sheaf *empty,
-		struct node_barn *barn)
+		struct node_barn *barn, int cpu)
 {
-	lockdep_assert_held(this_cpu_ptr(&s->cpu_sheaves->lock));
+	qpw_lockdep_assert_held(&s->cpu_sheaves->lock);
 
 	/* This is what we expect to find if nobody interrupted us. */
 	if (likely(!pcs->spare)) {
@@ -5546,31 +5578,34 @@ static void __pcs_install_empty_sheaf(st
 /*
  * Replace the full main sheaf with a (at least partially) empty sheaf.
  *
- * Must be called with the cpu_sheaves local lock locked. If successful, returns
- * the pcs pointer and the local lock locked (possibly on a different cpu than
- * initially called). If not successful, returns NULL and the local lock
- * unlocked.
+ * Must be called with the cpu_sheaves local lock locked, and migration counter
+ * increased. If successful, returns the pcs pointer and the local lock locked
+ * (possibly on a different cpu than initially called), with migration counter
+ * increased. If not successful, returns NULL and the local lock unlocked,
+ * and migration counter decreased.
  */
 static struct slub_percpu_sheaves *
 __pcs_replace_full_main(struct kmem_cache *s, struct slub_percpu_sheaves *pcs,
-			bool allow_spin)
+			bool allow_spin, int *cpu)
 {
 	struct slab_sheaf *empty;
 	struct node_barn *barn;
 	bool put_fail;
 
 restart:
-	lockdep_assert_held(this_cpu_ptr(&s->cpu_sheaves->lock));
+	qpw_lockdep_assert_held(&s->cpu_sheaves->lock);
 
 	/* Bootstrap or debug cache, back off */
 	if (unlikely(!cache_has_sheaves(s))) {
-		local_unlock(&s->cpu_sheaves->lock);
+		qpw_unlock(&s->cpu_sheaves->lock, *cpu);
+		migrate_enable();
 		return NULL;
 	}
 
 	barn = get_barn(s);
 	if (!barn) {
-		local_unlock(&s->cpu_sheaves->lock);
+		qpw_unlock(&s->cpu_sheaves->lock, *cpu);
+		migrate_enable();
 		return NULL;
 	}
 
@@ -5607,7 +5642,8 @@ restart:
 		stat(s, BARN_PUT_FAIL);
 
 		pcs->spare = NULL;
-		local_unlock(&s->cpu_sheaves->lock);
+		qpw_unlock(&s->cpu_sheaves->lock, *cpu);
+		migrate_enable();
 
 		sheaf_flush_unused(s, to_flush);
 		empty = to_flush;
@@ -5623,7 +5659,8 @@ restart:
 	put_fail = true;
 
 alloc_empty:
-	local_unlock(&s->cpu_sheaves->lock);
+	qpw_unlock(&s->cpu_sheaves->lock, *cpu);
+	migrate_enable();
 
 	/*
 	 * alloc_empty_sheaf() doesn't support !allow_spin and it's
@@ -5640,11 +5677,17 @@ alloc_empty:
 	if (put_fail)
 		 stat(s, BARN_PUT_FAIL);
 
-	if (!sheaf_flush_main(s))
+	migrate_disable();
+	*cpu = smp_processor_id();
+	if (!sheaf_flush_main(s, *cpu)) {
+		migrate_enable();
 		return NULL;
+	}
 
-	if (!local_trylock(&s->cpu_sheaves->lock))
+	if (!qpw_trylock(&s->cpu_sheaves->lock, *cpu)) {
+		migrate_enable();
 		return NULL;
+	}
 
 	pcs = this_cpu_ptr(s->cpu_sheaves);
 
@@ -5659,13 +5702,14 @@ alloc_empty:
 	return pcs;
 
 got_empty:
-	if (!local_trylock(&s->cpu_sheaves->lock)) {
+	if (!qpw_trylock(&s->cpu_sheaves->lock, *cpu)) {
+		migrate_enable();
 		barn_put_empty_sheaf(barn, empty);
 		return NULL;
 	}
 
 	pcs = this_cpu_ptr(s->cpu_sheaves);
-	__pcs_install_empty_sheaf(s, pcs, empty, barn);
+	__pcs_install_empty_sheaf(s, pcs, empty, barn, *cpu);
 
 	return pcs;
 }
@@ -5678,22 +5722,28 @@ static __fastpath_inline
 bool free_to_pcs(struct kmem_cache *s, void *object, bool allow_spin)
 {
 	struct slub_percpu_sheaves *pcs;
+	int cpu;
 
-	if (!local_trylock(&s->cpu_sheaves->lock))
+	migrate_disable();
+	cpu = smp_processor_id();
+	if (!qpw_trylock(&s->cpu_sheaves->lock, cpu)) {
+		migrate_enable();
 		return false;
+	}
 
 	pcs = this_cpu_ptr(s->cpu_sheaves);
 
 	if (unlikely(pcs->main->size == s->sheaf_capacity)) {
 
-		pcs = __pcs_replace_full_main(s, pcs, allow_spin);
+		pcs = __pcs_replace_full_main(s, pcs, allow_spin, &cpu);
 		if (unlikely(!pcs))
 			return false;
 	}
 
 	pcs->main->objects[pcs->main->size++] = object;
 
-	local_unlock(&s->cpu_sheaves->lock);
+	qpw_unlock(&s->cpu_sheaves->lock, cpu);
+	migrate_enable();
 
 	stat(s, FREE_FASTPATH);
 
@@ -5777,14 +5827,19 @@ bool __kfree_rcu_sheaf(struct kmem_cache
 {
 	struct slub_percpu_sheaves *pcs;
 	struct slab_sheaf *rcu_sheaf;
+	int cpu;
 
 	if (WARN_ON_ONCE(IS_ENABLED(CONFIG_PREEMPT_RT)))
 		return false;
 
 	lock_map_acquire_try(&kfree_rcu_sheaf_map);
 
-	if (!local_trylock(&s->cpu_sheaves->lock))
+	migrate_disable();
+	cpu = smp_processor_id();
+	if (!qpw_trylock(&s->cpu_sheaves->lock, cpu)) {
+		migrate_enable();
 		goto fail;
+	}
 
 	pcs = this_cpu_ptr(s->cpu_sheaves);
 
@@ -5795,7 +5850,8 @@ bool __kfree_rcu_sheaf(struct kmem_cache
 
 		/* Bootstrap or debug cache, fall back */
 		if (unlikely(!cache_has_sheaves(s))) {
-			local_unlock(&s->cpu_sheaves->lock);
+			qpw_unlock(&s->cpu_sheaves->lock, cpu);
+			migrate_enable();
 			goto fail;
 		}
 
@@ -5807,7 +5863,8 @@ bool __kfree_rcu_sheaf(struct kmem_cache
 
 		barn = get_barn(s);
 		if (!barn) {
-			local_unlock(&s->cpu_sheaves->lock);
+			qpw_unlock(&s->cpu_sheaves->lock, cpu);
+			migrate_enable();
 			goto fail;
 		}
 
@@ -5818,15 +5875,18 @@ bool __kfree_rcu_sheaf(struct kmem_cache
 			goto do_free;
 		}
 
-		local_unlock(&s->cpu_sheaves->lock);
+		qpw_unlock(&s->cpu_sheaves->lock, cpu);
+		migrate_enable();
 
 		empty = alloc_empty_sheaf(s, GFP_NOWAIT);
 
 		if (!empty)
 			goto fail;
 
-		if (!local_trylock(&s->cpu_sheaves->lock)) {
+		migrate_disable();
+		if (!qpw_trylock(&s->cpu_sheaves->lock, cpu)) {
 			barn_put_empty_sheaf(barn, empty);
+			migrate_enable();
 			goto fail;
 		}
 
@@ -5862,7 +5922,8 @@ do_free:
 	if (rcu_sheaf)
 		call_rcu(&rcu_sheaf->rcu_head, rcu_free_sheaf);
 
-	local_unlock(&s->cpu_sheaves->lock);
+	qpw_unlock(&s->cpu_sheaves->lock, cpu);
+	migrate_enable();
 
 	stat(s, FREE_RCU_SHEAF);
 	lock_map_release(&kfree_rcu_sheaf_map);
@@ -5889,6 +5950,7 @@ static void free_to_pcs_bulk(struct kmem
 	void *remote_objects[PCS_BATCH_MAX];
 	unsigned int remote_nr = 0;
 	int node = numa_mem_id();
+	int cpu;
 
 next_remote_batch:
 	while (i < size) {
@@ -5918,7 +5980,9 @@ next_remote_batch:
 		goto flush_remote;
 
 next_batch:
-	if (!local_trylock(&s->cpu_sheaves->lock))
+	migrate_disable();
+	cpu = smp_processor_id();
+	if (!qpw_trylock(&s->cpu_sheaves->lock, cpu))
 		goto fallback;
 
 	pcs = this_cpu_ptr(s->cpu_sheaves);
@@ -5961,7 +6025,8 @@ do_free:
 	memcpy(main->objects + main->size, p, batch * sizeof(void *));
 	main->size += batch;
 
-	local_unlock(&s->cpu_sheaves->lock);
+	qpw_unlock(&s->cpu_sheaves->lock, cpu);
+	migrate_enable();
 
 	stat_add(s, FREE_FASTPATH, batch);
 
@@ -5977,7 +6042,8 @@ do_free:
 	return;
 
 no_empty:
-	local_unlock(&s->cpu_sheaves->lock);
+	qpw_unlock(&s->cpu_sheaves->lock, cpu);
+	migrate_enable();
 
 	/*
 	 * if we depleted all empty sheaves in the barn or there are too
@@ -7377,7 +7443,7 @@ static int init_percpu_sheaves(struct km
 
 		pcs = per_cpu_ptr(s->cpu_sheaves, cpu);
 
-		local_trylock_init(&pcs->lock);
+		qpw_trylock_init(&pcs->lock);
 
 		/*
 		 * Bootstrap sheaf has zero size so fast-path allocation fails.