lib/sbitmap.c

   1 /*
   2  * Copyright (C) 2016 Facebook
   3  * Copyright (C) 2013-2014 Jens Axboe
   4  *
   5  * This program is free software; you can redistribute it and/or
   6  * modify it under the terms of the GNU General Public
   7  * License v2 as published by the Free Software Foundation.
   8  *
   9  * This program is distributed in the hope that it will be useful,
  10  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  11  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  12  * General Public License for more details.
  13  *
  14  * You should have received a copy of the GNU General Public License
  15  * along with this program.  If not, see <https://www.gnu.org/licenses/>.
  16  */
  17
  18 #include <linux/sbitmap.h>
  19
  20 int sbitmap_init_node(struct sbitmap *sb, unsigned int depth, int shift,
  21                       gfp_t flags, int node)
  22 {
  23         unsigned int bits_per_word;
  24         unsigned int i;
  25
  26         if (shift < 0) {
  27                 shift = ilog2(BITS_PER_LONG);
  28                 /*
  29                  * If the bitmap is small, shrink the number of bits per word so
  30                  * we spread over a few cachelines, at least. If less than 4
  31                  * bits, just forget about it, it's not going to work optimally
  32                  * anyway.
  33                  */
  34                 if (depth >= 4) {
  35                         while ((4U << shift) > depth)
  36                                 shift--;
  37                 }
  38         }
  39         bits_per_word = 1U << shift;
  40         if (bits_per_word > BITS_PER_LONG)
  41                 return -EINVAL;
  42
  43         sb->shift = shift;
  44         sb->depth = depth;
  45         sb->map_nr = DIV_ROUND_UP(sb->depth, bits_per_word);
  46
  47         if (depth == 0) {
  48                 sb->map = NULL;
  49                 return 0;
  50         }
  51
  52         sb->map = kzalloc_node(sb->map_nr * sizeof(*sb->map), flags, node);
  53         if (!sb->map)
  54                 return -ENOMEM;
  55
  56         for (i = 0; i < sb->map_nr; i++) {
  57                 sb->map[i].depth = min(depth, bits_per_word);
  58                 depth -= sb->map[i].depth;
  59         }
  60         return 0;
  61 }
  62 EXPORT_SYMBOL_GPL(sbitmap_init_node);
  63
  64 void sbitmap_resize(struct sbitmap *sb, unsigned int depth)
  65 {
  66         unsigned int bits_per_word = 1U << sb->shift;
  67         unsigned int i;
  68
  69         sb->depth = depth;
  70         sb->map_nr = DIV_ROUND_UP(sb->depth, bits_per_word);
  71
  72         for (i = 0; i < sb->map_nr; i++) {
  73                 sb->map[i].depth = min(depth, bits_per_word);
  74                 depth -= sb->map[i].depth;
  75         }
  76 }
  77 EXPORT_SYMBOL_GPL(sbitmap_resize);
  78
  79 static int __sbitmap_get_word(struct sbitmap_word *word, unsigned int hint,
  80                               bool wrap)
  81 {
  82         unsigned int orig_hint = hint;
  83         int nr;
  84
  85         while (1) {
  86                 nr = find_next_zero_bit(&word->word, word->depth, hint);
  87                 if (unlikely(nr >= word->depth)) {
  88                         /*
  89                          * We started with an offset, and we didn't reset the
  90                          * offset to 0 in a failure case, so start from 0 to
  91                          * exhaust the map.
  92                          */
  93                         if (orig_hint && hint && wrap) {
  94                                 hint = orig_hint = 0;
  95                                 continue;
  96                         }
  97                         return -1;
  98                 }
  99
 100                 if (!test_and_set_bit(nr, &word->word))
 101                         break;
 102
 103                 hint = nr + 1;
 104                 if (hint >= word->depth - 1)
 105                         hint = 0;
 106         }
 107
 108         return nr;
 109 }
 110
 111 int sbitmap_get(struct sbitmap *sb, unsigned int alloc_hint, bool round_robin)
 112 {
 113         unsigned int i, index;
 114         int nr = -1;
 115
 116         index = SB_NR_TO_INDEX(sb, alloc_hint);
 117
 118         for (i = 0; i < sb->map_nr; i++) {
 119                 nr = __sbitmap_get_word(&sb->map[index],
 120                                         SB_NR_TO_BIT(sb, alloc_hint),
 121                                         !round_robin);
 122                 if (nr != -1) {
 123                         nr += index << sb->shift;
 124                         break;
 125                 }
 126
 127                 /* Jump to next index. */
 128                 index++;
 129                 alloc_hint = index << sb->shift;
 130
 131                 if (index >= sb->map_nr) {
 132                         index = 0;
 133                         alloc_hint = 0;
 134                 }
 135         }
 136
 137         return nr;
 138 }
 139 EXPORT_SYMBOL_GPL(sbitmap_get);
 140
 141 bool sbitmap_any_bit_set(const struct sbitmap *sb)
 142 {
 143         unsigned int i;
 144
 145         for (i = 0; i < sb->map_nr; i++) {
 146                 if (sb->map[i].word)
 147                         return true;
 148         }
 149         return false;
 150 }
 151 EXPORT_SYMBOL_GPL(sbitmap_any_bit_set);
 152
 153 bool sbitmap_any_bit_clear(const struct sbitmap *sb)
 154 {
 155         unsigned int i;
 156
 157         for (i = 0; i < sb->map_nr; i++) {
 158                 const struct sbitmap_word *word = &sb->map[i];
 159                 unsigned long ret;
 160
 161                 ret = find_first_zero_bit(&word->word, word->depth);
 162                 if (ret < word->depth)
 163                         return true;
 164         }
 165         return false;
 166 }
 167 EXPORT_SYMBOL_GPL(sbitmap_any_bit_clear);
 168
 169 unsigned int sbitmap_weight(const struct sbitmap *sb)
 170 {
 171         unsigned int i, weight;
 172
 173         for (i = 0; i < sb->map_nr; i++) {
 174                 const struct sbitmap_word *word = &sb->map[i];
 175
 176                 weight += bitmap_weight(&word->word, word->depth);
 177         }
 178         return weight;
 179 }
 180 EXPORT_SYMBOL_GPL(sbitmap_weight);
 181
 182 static unsigned int sbq_calc_wake_batch(unsigned int depth)
 183 {
 184         unsigned int wake_batch;
 185
 186         /*
 187          * For each batch, we wake up one queue. We need to make sure that our
 188          * batch size is small enough that the full depth of the bitmap is
 189          * enough to wake up all of the queues.
 190          */
 191         wake_batch = SBQ_WAKE_BATCH;
 192         if (wake_batch > depth / SBQ_WAIT_QUEUES)
 193                 wake_batch = max(1U, depth / SBQ_WAIT_QUEUES);
 194
 195         return wake_batch;
 196 }
 197
 198 int sbitmap_queue_init_node(struct sbitmap_queue *sbq, unsigned int depth,
 199                             int shift, bool round_robin, gfp_t flags, int node)
 200 {
 201         int ret;
 202         int i;
 203
 204         ret = sbitmap_init_node(&sbq->sb, depth, shift, flags, node);
 205         if (ret)
 206                 return ret;
 207
 208         sbq->alloc_hint = alloc_percpu_gfp(unsigned int, flags);
 209         if (!sbq->alloc_hint) {
 210                 sbitmap_free(&sbq->sb);
 211                 return -ENOMEM;
 212         }
 213
 214         sbq->wake_batch = sbq_calc_wake_batch(depth);
 215         atomic_set(&sbq->wake_index, 0);
 216
 217         sbq->ws = kzalloc_node(SBQ_WAIT_QUEUES * sizeof(*sbq->ws), flags, node);
 218         if (!sbq->ws) {
 219                 free_percpu(sbq->alloc_hint);
 220                 sbitmap_free(&sbq->sb);
 221                 return -ENOMEM;
 222         }
 223
 224         for (i = 0; i < SBQ_WAIT_QUEUES; i++) {
 225                 init_waitqueue_head(&sbq->ws[i].wait);
 226                 atomic_set(&sbq->ws[i].wait_cnt, sbq->wake_batch);
 227         }
 228
 229         sbq->round_robin = round_robin;
 230         return 0;
 231 }
 232 EXPORT_SYMBOL_GPL(sbitmap_queue_init_node);
 233
 234 void sbitmap_queue_resize(struct sbitmap_queue *sbq, unsigned int depth)
 235 {
 236         sbq->wake_batch = sbq_calc_wake_batch(depth);
 237         sbitmap_resize(&sbq->sb, depth);
 238 }
 239 EXPORT_SYMBOL_GPL(sbitmap_queue_resize);
 240
 241 int __sbitmap_queue_get(struct sbitmap_queue *sbq)
 242 {
 243         unsigned int hint;
 244         int nr;
 245
 246         hint = this_cpu_read(*sbq->alloc_hint);
 247         nr = sbitmap_get(&sbq->sb, hint, sbq->round_robin);
 248
 249         if (nr == -1) {
 250                 /* If the map is full, a hint won't do us much good. */
 251                 this_cpu_write(*sbq->alloc_hint, 0);
 252         } else if (nr == hint || unlikely(sbq->round_robin)) {
 253                 /* Only update the hint if we used it. */
 254                 hint = nr + 1;
 255                 if (hint >= sbq->sb.depth - 1)
 256                         hint = 0;
 257                 this_cpu_write(*sbq->alloc_hint, hint);
 258         }
 259
 260         return nr;
 261 }
 262 EXPORT_SYMBOL_GPL(__sbitmap_queue_get);
 263
 264 static struct sbq_wait_state *sbq_wake_ptr(struct sbitmap_queue *sbq)
 265 {
 266         int i, wake_index;
 267
 268         wake_index = atomic_read(&sbq->wake_index);
 269         for (i = 0; i < SBQ_WAIT_QUEUES; i++) {
 270                 struct sbq_wait_state *ws = &sbq->ws[wake_index];
 271
 272                 if (waitqueue_active(&ws->wait)) {
 273                         int o = atomic_read(&sbq->wake_index);
 274
 275                         if (wake_index != o)
 276                                 atomic_cmpxchg(&sbq->wake_index, o, wake_index);
 277                         return ws;
 278                 }
 279
 280                 wake_index = sbq_index_inc(wake_index);
 281         }
 282
 283         return NULL;
 284 }
 285
 286 static void sbq_wake_up(struct sbitmap_queue *sbq)
 287 {
 288         struct sbq_wait_state *ws;
 289         int wait_cnt;
 290
 291         /* Ensure that the wait list checks occur after clear_bit(). */
 292         smp_mb();
 293
 294         ws = sbq_wake_ptr(sbq);
 295         if (!ws)
 296                 return;
 297
 298         wait_cnt = atomic_dec_return(&ws->wait_cnt);
 299         if (unlikely(wait_cnt < 0))
 300                 wait_cnt = atomic_inc_return(&ws->wait_cnt);
 301         if (wait_cnt == 0) {
 302                 atomic_add(sbq->wake_batch, &ws->wait_cnt);
 303                 sbq_index_atomic_inc(&sbq->wake_index);
 304                 wake_up(&ws->wait);
 305         }
 306 }
 307
 308 void sbitmap_queue_clear(struct sbitmap_queue *sbq, unsigned int nr,
 309                          unsigned int cpu)
 310 {
 311         sbitmap_clear_bit(&sbq->sb, nr);
 312         sbq_wake_up(sbq);
 313         if (likely(!sbq->round_robin))
 314                 *per_cpu_ptr(sbq->alloc_hint, cpu) = nr;
 315 }
 316 EXPORT_SYMBOL_GPL(sbitmap_queue_clear);
 317
 318 void sbitmap_queue_wake_all(struct sbitmap_queue *sbq)
 319 {
 320         int i, wake_index;
 321
 322         /*
 323          * Make sure all changes prior to this are visible from other CPUs.
 324          */
 325         smp_mb();
 326         wake_index = atomic_read(&sbq->wake_index);
 327         for (i = 0; i < SBQ_WAIT_QUEUES; i++) {
 328                 struct sbq_wait_state *ws = &sbq->ws[wake_index];
 329
 330                 if (waitqueue_active(&ws->wait))
 331                         wake_up(&ws->wait);
 332
 333                 wake_index = sbq_index_inc(wake_index);
 334         }
 335 }
 336 EXPORT_SYMBOL_GPL(sbitmap_queue_wake_all);