sbitmap.c 7.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348
  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. #include <linux/random.h>
  18. #include <linux/sbitmap.h>
  19. int sbitmap_init_node(struct sbitmap *sb, unsigned int depth, int shift,
  20. gfp_t flags, int node)
  21. {
  22. unsigned int bits_per_word;
  23. unsigned int i;
  24. if (shift < 0) {
  25. shift = ilog2(BITS_PER_LONG);
  26. /*
  27. * If the bitmap is small, shrink the number of bits per word so
  28. * we spread over a few cachelines, at least. If less than 4
  29. * bits, just forget about it, it's not going to work optimally
  30. * anyway.
  31. */
  32. if (depth >= 4) {
  33. while ((4U << shift) > depth)
  34. shift--;
  35. }
  36. }
  37. bits_per_word = 1U << shift;
  38. if (bits_per_word > BITS_PER_LONG)
  39. return -EINVAL;
  40. sb->shift = shift;
  41. sb->depth = depth;
  42. sb->map_nr = DIV_ROUND_UP(sb->depth, bits_per_word);
  43. if (depth == 0) {
  44. sb->map = NULL;
  45. return 0;
  46. }
  47. sb->map = kzalloc_node(sb->map_nr * sizeof(*sb->map), flags, node);
  48. if (!sb->map)
  49. return -ENOMEM;
  50. for (i = 0; i < sb->map_nr; i++) {
  51. sb->map[i].depth = min(depth, bits_per_word);
  52. depth -= sb->map[i].depth;
  53. }
  54. return 0;
  55. }
  56. EXPORT_SYMBOL_GPL(sbitmap_init_node);
  57. void sbitmap_resize(struct sbitmap *sb, unsigned int depth)
  58. {
  59. unsigned int bits_per_word = 1U << sb->shift;
  60. unsigned int i;
  61. sb->depth = depth;
  62. sb->map_nr = DIV_ROUND_UP(sb->depth, bits_per_word);
  63. for (i = 0; i < sb->map_nr; i++) {
  64. sb->map[i].depth = min(depth, bits_per_word);
  65. depth -= sb->map[i].depth;
  66. }
  67. }
  68. EXPORT_SYMBOL_GPL(sbitmap_resize);
  69. static int __sbitmap_get_word(struct sbitmap_word *word, unsigned int hint,
  70. bool wrap)
  71. {
  72. unsigned int orig_hint = hint;
  73. int nr;
  74. while (1) {
  75. nr = find_next_zero_bit(&word->word, word->depth, hint);
  76. if (unlikely(nr >= word->depth)) {
  77. /*
  78. * We started with an offset, and we didn't reset the
  79. * offset to 0 in a failure case, so start from 0 to
  80. * exhaust the map.
  81. */
  82. if (orig_hint && hint && wrap) {
  83. hint = orig_hint = 0;
  84. continue;
  85. }
  86. return -1;
  87. }
  88. if (!test_and_set_bit(nr, &word->word))
  89. break;
  90. hint = nr + 1;
  91. if (hint >= word->depth - 1)
  92. hint = 0;
  93. }
  94. return nr;
  95. }
  96. int sbitmap_get(struct sbitmap *sb, unsigned int alloc_hint, bool round_robin)
  97. {
  98. unsigned int i, index;
  99. int nr = -1;
  100. index = SB_NR_TO_INDEX(sb, alloc_hint);
  101. for (i = 0; i < sb->map_nr; i++) {
  102. nr = __sbitmap_get_word(&sb->map[index],
  103. SB_NR_TO_BIT(sb, alloc_hint),
  104. !round_robin);
  105. if (nr != -1) {
  106. nr += index << sb->shift;
  107. break;
  108. }
  109. /* Jump to next index. */
  110. index++;
  111. alloc_hint = index << sb->shift;
  112. if (index >= sb->map_nr) {
  113. index = 0;
  114. alloc_hint = 0;
  115. }
  116. }
  117. return nr;
  118. }
  119. EXPORT_SYMBOL_GPL(sbitmap_get);
  120. bool sbitmap_any_bit_set(const struct sbitmap *sb)
  121. {
  122. unsigned int i;
  123. for (i = 0; i < sb->map_nr; i++) {
  124. if (sb->map[i].word)
  125. return true;
  126. }
  127. return false;
  128. }
  129. EXPORT_SYMBOL_GPL(sbitmap_any_bit_set);
  130. bool sbitmap_any_bit_clear(const struct sbitmap *sb)
  131. {
  132. unsigned int i;
  133. for (i = 0; i < sb->map_nr; i++) {
  134. const struct sbitmap_word *word = &sb->map[i];
  135. unsigned long ret;
  136. ret = find_first_zero_bit(&word->word, word->depth);
  137. if (ret < word->depth)
  138. return true;
  139. }
  140. return false;
  141. }
  142. EXPORT_SYMBOL_GPL(sbitmap_any_bit_clear);
  143. unsigned int sbitmap_weight(const struct sbitmap *sb)
  144. {
  145. unsigned int i, weight = 0;
  146. for (i = 0; i < sb->map_nr; i++) {
  147. const struct sbitmap_word *word = &sb->map[i];
  148. weight += bitmap_weight(&word->word, word->depth);
  149. }
  150. return weight;
  151. }
  152. EXPORT_SYMBOL_GPL(sbitmap_weight);
  153. static unsigned int sbq_calc_wake_batch(unsigned int depth)
  154. {
  155. unsigned int wake_batch;
  156. /*
  157. * For each batch, we wake up one queue. We need to make sure that our
  158. * batch size is small enough that the full depth of the bitmap is
  159. * enough to wake up all of the queues.
  160. */
  161. wake_batch = SBQ_WAKE_BATCH;
  162. if (wake_batch > depth / SBQ_WAIT_QUEUES)
  163. wake_batch = max(1U, depth / SBQ_WAIT_QUEUES);
  164. return wake_batch;
  165. }
  166. int sbitmap_queue_init_node(struct sbitmap_queue *sbq, unsigned int depth,
  167. int shift, bool round_robin, gfp_t flags, int node)
  168. {
  169. int ret;
  170. int i;
  171. ret = sbitmap_init_node(&sbq->sb, depth, shift, flags, node);
  172. if (ret)
  173. return ret;
  174. sbq->alloc_hint = alloc_percpu_gfp(unsigned int, flags);
  175. if (!sbq->alloc_hint) {
  176. sbitmap_free(&sbq->sb);
  177. return -ENOMEM;
  178. }
  179. if (depth && !round_robin) {
  180. for_each_possible_cpu(i)
  181. *per_cpu_ptr(sbq->alloc_hint, i) = prandom_u32() % depth;
  182. }
  183. sbq->wake_batch = sbq_calc_wake_batch(depth);
  184. atomic_set(&sbq->wake_index, 0);
  185. sbq->ws = kzalloc_node(SBQ_WAIT_QUEUES * sizeof(*sbq->ws), flags, node);
  186. if (!sbq->ws) {
  187. free_percpu(sbq->alloc_hint);
  188. sbitmap_free(&sbq->sb);
  189. return -ENOMEM;
  190. }
  191. for (i = 0; i < SBQ_WAIT_QUEUES; i++) {
  192. init_waitqueue_head(&sbq->ws[i].wait);
  193. atomic_set(&sbq->ws[i].wait_cnt, sbq->wake_batch);
  194. }
  195. sbq->round_robin = round_robin;
  196. return 0;
  197. }
  198. EXPORT_SYMBOL_GPL(sbitmap_queue_init_node);
  199. void sbitmap_queue_resize(struct sbitmap_queue *sbq, unsigned int depth)
  200. {
  201. sbq->wake_batch = sbq_calc_wake_batch(depth);
  202. sbitmap_resize(&sbq->sb, depth);
  203. }
  204. EXPORT_SYMBOL_GPL(sbitmap_queue_resize);
  205. int __sbitmap_queue_get(struct sbitmap_queue *sbq)
  206. {
  207. unsigned int hint, depth;
  208. int nr;
  209. hint = this_cpu_read(*sbq->alloc_hint);
  210. depth = READ_ONCE(sbq->sb.depth);
  211. if (unlikely(hint >= depth)) {
  212. hint = depth ? prandom_u32() % depth : 0;
  213. this_cpu_write(*sbq->alloc_hint, hint);
  214. }
  215. nr = sbitmap_get(&sbq->sb, hint, sbq->round_robin);
  216. if (nr == -1) {
  217. /* If the map is full, a hint won't do us much good. */
  218. this_cpu_write(*sbq->alloc_hint, 0);
  219. } else if (nr == hint || unlikely(sbq->round_robin)) {
  220. /* Only update the hint if we used it. */
  221. hint = nr + 1;
  222. if (hint >= depth - 1)
  223. hint = 0;
  224. this_cpu_write(*sbq->alloc_hint, hint);
  225. }
  226. return nr;
  227. }
  228. EXPORT_SYMBOL_GPL(__sbitmap_queue_get);
  229. static struct sbq_wait_state *sbq_wake_ptr(struct sbitmap_queue *sbq)
  230. {
  231. int i, wake_index;
  232. wake_index = atomic_read(&sbq->wake_index);
  233. for (i = 0; i < SBQ_WAIT_QUEUES; i++) {
  234. struct sbq_wait_state *ws = &sbq->ws[wake_index];
  235. if (waitqueue_active(&ws->wait)) {
  236. int o = atomic_read(&sbq->wake_index);
  237. if (wake_index != o)
  238. atomic_cmpxchg(&sbq->wake_index, o, wake_index);
  239. return ws;
  240. }
  241. wake_index = sbq_index_inc(wake_index);
  242. }
  243. return NULL;
  244. }
  245. static void sbq_wake_up(struct sbitmap_queue *sbq)
  246. {
  247. struct sbq_wait_state *ws;
  248. int wait_cnt;
  249. /* Ensure that the wait list checks occur after clear_bit(). */
  250. smp_mb();
  251. ws = sbq_wake_ptr(sbq);
  252. if (!ws)
  253. return;
  254. wait_cnt = atomic_dec_return(&ws->wait_cnt);
  255. if (unlikely(wait_cnt < 0))
  256. wait_cnt = atomic_inc_return(&ws->wait_cnt);
  257. if (wait_cnt == 0) {
  258. atomic_add(sbq->wake_batch, &ws->wait_cnt);
  259. sbq_index_atomic_inc(&sbq->wake_index);
  260. wake_up(&ws->wait);
  261. }
  262. }
  263. void sbitmap_queue_clear(struct sbitmap_queue *sbq, unsigned int nr,
  264. unsigned int cpu)
  265. {
  266. sbitmap_clear_bit(&sbq->sb, nr);
  267. sbq_wake_up(sbq);
  268. if (likely(!sbq->round_robin && nr < sbq->sb.depth))
  269. *per_cpu_ptr(sbq->alloc_hint, cpu) = nr;
  270. }
  271. EXPORT_SYMBOL_GPL(sbitmap_queue_clear);
  272. void sbitmap_queue_wake_all(struct sbitmap_queue *sbq)
  273. {
  274. int i, wake_index;
  275. /*
  276. * Make sure all changes prior to this are visible from other CPUs.
  277. */
  278. smp_mb();
  279. wake_index = atomic_read(&sbq->wake_index);
  280. for (i = 0; i < SBQ_WAIT_QUEUES; i++) {
  281. struct sbq_wait_state *ws = &sbq->ws[wake_index];
  282. if (waitqueue_active(&ws->wait))
  283. wake_up(&ws->wait);
  284. wake_index = sbq_index_inc(wake_index);
  285. }
  286. }
  287. EXPORT_SYMBOL_GPL(sbitmap_queue_wake_all);