5 files changed, 51 insertions, 45 deletions

diff --git a/include/linux/idr.h b/include/linux/idr.h
index f58c0a3addc3..2027c7aba50d 100644
--- a/include/linux/idr.h
+++ b/include/linux/idr.h
@@ -14,6 +14,7 @@
 
 #include <linux/radix-tree.h>
 #include <linux/gfp.h>
+#include <linux/percpu.h>
 
 struct idr {
         struct radix_tree_root  idr_rt;
@@ -171,9 +172,10 @@ struct ida_bitmap {
         unsigned long           bitmap[IDA_BITMAP_LONGS];
 };
 
+DECLARE_PER_CPU(struct ida_bitmap *, ida_bitmap);
+
 struct ida {
         struct radix_tree_root  ida_rt;
-        struct ida_bitmap       *free_bitmap;
 };
 
 #define IDA_INIT        {                                               \
@@ -193,7 +195,6 @@ void ida_simple_remove(struct ida *ida, unsigned int id);
 static inline void ida_init(struct ida *ida)
 {
         INIT_RADIX_TREE(&ida->ida_rt, IDR_RT_MARKER | GFP_NOWAIT);
-        ida->free_bitmap = NULL;
 }
 
 /**
diff --git a/lib/idr.c b/lib/idr.c
index b87056e2cc4c..2abd7769c430 100644
--- a/lib/idr.c
+++ b/lib/idr.c
@@ -4,6 +4,7 @@
 #include <linux/slab.h>
 #include <linux/spinlock.h>
 
+DEFINE_PER_CPU(struct ida_bitmap *, ida_bitmap);
 static DEFINE_SPINLOCK(simple_ida_lock);
 
 /**
@@ -193,38 +194,6 @@ EXPORT_SYMBOL(idr_replace);
  * limitation, it should be quite straightforward to raise the maximum.
  */
 
-/**
- * ida_pre_get - reserve resources for ida allocation
- * @ida: ida handle
- * @gfp: memory allocation flags
- *
- * This function should be called before calling ida_get_new_above().  If it
- * is unable to allocate memory, it will return %0.  On success, it returns %1.
- */
-int ida_pre_get(struct ida *ida, gfp_t gfp)
-{
-        struct ida_bitmap *bitmap;
-
-        /*
-         * This looks weird, but the IDA API has no preload_end() equivalent.
-         * Instead, ida_get_new() can return -EAGAIN, prompting the caller
-         * to return to the ida_pre_get() step.
-         */
-        idr_preload(gfp);
-        idr_preload_end();
-
-        if (!ida->free_bitmap) {
-                bitmap = kmalloc(sizeof(struct ida_bitmap), gfp);
-                if (!bitmap)
-                        return 0;
-                bitmap = xchg(&ida->free_bitmap, bitmap);
-                kfree(bitmap);
-        }
-
-        return 1;
-}
-EXPORT_SYMBOL(ida_pre_get);
-
 #define IDA_MAX (0x80000000U / IDA_BITMAP_BITS)
 
 /**
@@ -292,10 +261,9 @@ int ida_get_new_above(struct ida *ida, int start, int *id)
                         new += bit;
                         if (new < 0)
                                 return -ENOSPC;
-                        bitmap = ida->free_bitmap;
+                        bitmap = this_cpu_xchg(ida_bitmap, NULL);
                         if (!bitmap)
                                 return -EAGAIN;
-                        ida->free_bitmap = NULL;
                         memset(bitmap, 0, sizeof(*bitmap));
                         __set_bit(bit, bitmap->bitmap);
                         radix_tree_iter_replace(root, &iter, slot, bitmap);
@@ -361,9 +329,6 @@ void ida_destroy(struct ida *ida)
                 kfree(bitmap);
                 radix_tree_iter_delete(&ida->ida_rt, &iter, slot);
         }
-
-        kfree(ida->free_bitmap);
-        ida->free_bitmap = NULL;
 }
 EXPORT_SYMBOL(ida_destroy);
 
diff --git a/lib/radix-tree.c b/lib/radix-tree.c
index eaea14b8f2ca..7b9f8515033e 100644
--- a/lib/radix-tree.c
+++ b/lib/radix-tree.c
@@ -70,6 +70,14 @@ static struct kmem_cache *radix_tree_node_cachep;
 #define IDR_PRELOAD_SIZE        (IDR_MAX_PATH * 2 - 1)
 
 /*
+ * The IDA is even shorter since it uses a bitmap at the last level.
+ */
+#define IDA_INDEX_BITS          (8 * sizeof(int) - 1 - ilog2(IDA_BITMAP_BITS))
+#define IDA_MAX_PATH            (DIV_ROUND_UP(IDA_INDEX_BITS, \
+                                                RADIX_TREE_MAP_SHIFT))
+#define IDA_PRELOAD_SIZE        (IDA_MAX_PATH * 2 - 1)
+
+/*
  * Per-cpu pool of preloaded nodes
  */
 struct radix_tree_preload {
@@ -346,9 +354,8 @@ static void dump_ida_node(void *entry, unsigned long index)
 static void ida_dump(struct ida *ida)
 {
         struct radix_tree_root *root = &ida->ida_rt;
-        pr_debug("ida: %p %p free %d bitmap %p\n", ida, root->rnode,
-                                root->gfp_mask >> ROOT_TAG_SHIFT,
-                                ida->free_bitmap);
+        pr_debug("ida: %p node %p free %d\n", ida, root->rnode,
+                                root->gfp_mask >> ROOT_TAG_SHIFT);
         dump_ida_node(root->rnode, 0);
 }
 #endif
@@ -2080,6 +2087,36 @@ void idr_preload(gfp_t gfp_mask)
 }
 EXPORT_SYMBOL(idr_preload);
 
+/**
+ * ida_pre_get - reserve resources for ida allocation
+ * @ida: ida handle
+ * @gfp: memory allocation flags
+ *
+ * This function should be called before calling ida_get_new_above().  If it
+ * is unable to allocate memory, it will return %0.  On success, it returns %1.
+ */
+int ida_pre_get(struct ida *ida, gfp_t gfp)
+{
+        __radix_tree_preload(gfp, IDA_PRELOAD_SIZE);
+        /*
+         * The IDA API has no preload_end() equivalent.  Instead,
+         * ida_get_new() can return -EAGAIN, prompting the caller
+         * to return to the ida_pre_get() step.
+         */
+        preempt_enable();
+
+        if (!this_cpu_read(ida_bitmap)) {
+                struct ida_bitmap *bitmap = kmalloc(sizeof(*bitmap), gfp);
+                if (!bitmap)
+                        return 0;
+                bitmap = this_cpu_cmpxchg(ida_bitmap, NULL, bitmap);
+                kfree(bitmap);
+        }
+
+        return 1;
+}
+EXPORT_SYMBOL(ida_pre_get);
+
 void **idr_get_free(struct radix_tree_root *root,
                         struct radix_tree_iter *iter, gfp_t gfp, int end)
 {
@@ -2219,6 +2256,8 @@ static int radix_tree_cpu_dead(unsigned int cpu)
                 kmem_cache_free(radix_tree_node_cachep, node);
                 rtp->nr--;
         }
+        kfree(per_cpu(ida_bitmap, cpu));
+        per_cpu(ida_bitmap, cpu) = NULL;
         return 0;
 }
 
diff --git a/tools/testing/radix-tree/linux/kernel.h b/tools/testing/radix-tree/linux/kernel.h
index 63fce553781a..677b8c0f60f9 100644
--- a/tools/testing/radix-tree/linux/kernel.h
+++ b/tools/testing/radix-tree/linux/kernel.h
@@ -24,6 +24,4 @@
 
 #define ARRAY_SIZE(arr) (sizeof(arr) / sizeof((arr)[0]))
 
-#define xchg(ptr, x)    uatomic_xchg(ptr, x)
-
 #endif /* _KERNEL_H */
diff --git a/tools/testing/radix-tree/linux/percpu.h b/tools/testing/radix-tree/linux/percpu.h
index 5837f1d56f17..3ea01a1a88c2 100644
--- a/tools/testing/radix-tree/linux/percpu.h
+++ b/tools/testing/radix-tree/linux/percpu.h
@@ -1,7 +1,10 @@
-
+#define DECLARE_PER_CPU(type, val) extern type val
 #define DEFINE_PER_CPU(type, val) type val
 
 #define __get_cpu_var(var)      var
 #define this_cpu_ptr(var)       var
+#define this_cpu_read(var)      var
+#define this_cpu_xchg(var, val)         uatomic_xchg(&var, val)
+#define this_cpu_cmpxchg(var, old, new) uatomic_cmpxchg(&var, old, new)
 #define per_cpu_ptr(ptr, cpu)   ({ (void)(cpu); (ptr); })
 #define per_cpu(var, cpu)       (*per_cpu_ptr(&(var), cpu))