util: add util_idalloc_sparse, solving the excessive virtual memory usage

The code comment in the header file describes how it works.

Reviewed-by: Jesse Natalie <jenatali@microsoft.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/30106>

src/util/u_idalloc.c

@@ -181,6 +181,10 @@ util_idalloc_reserve(struct util_idalloc *buf, unsigned id)
    buf->num_set_elements = MAX2(buf->num_set_elements, idx + 1);
 }
 
+/*********************************************
+ * util_idalloc_mt
+ *********************************************/
+
 void
 util_idalloc_mt_init(struct util_idalloc_mt *buf,
                      unsigned initial_num_ids, bool skip_zero)
@@ -228,3 +232,81 @@ util_idalloc_mt_free(struct util_idalloc_mt *buf, unsigned id)
    util_idalloc_free(&buf->buf, id);
    simple_mtx_unlock(&buf->mutex);
 }
+
+/*********************************************
+ * util_idalloc_sparse
+ *********************************************/
+
+void
+util_idalloc_sparse_init(struct util_idalloc_sparse *buf)
+{
+   static_assert(IS_POT_NONZERO(ARRAY_SIZE(buf->segment)),
+         "buf->segment[] must have a power of two number of elements");
+
+   for (unsigned i = 0; i < ARRAY_SIZE(buf->segment); i++)
+      util_idalloc_init(&buf->segment[i], 1);
+}
+
+void
+util_idalloc_sparse_fini(struct util_idalloc_sparse *buf)
+{
+   for (unsigned i = 0; i < ARRAY_SIZE(buf->segment); i++)
+      util_idalloc_fini(&buf->segment[i]);
+}
+
+unsigned
+util_idalloc_sparse_alloc(struct util_idalloc_sparse *buf)
+{
+   unsigned max_ids = UTIL_IDALLOC_MAX_IDS_PER_SEGMENT(buf);
+
+   for (unsigned i = 0; i < ARRAY_SIZE(buf->segment); i++) {
+      if (buf->segment[i].lowest_free_idx <
+          UTIL_IDALLOC_MAX_ELEMS_PER_SEGMENT(buf))
+         return max_ids * i + util_idalloc_alloc(&buf->segment[i]);
+   }
+
+   fprintf(stderr, "mesa: util_idalloc_sparse_alloc: "
+                   "all 2^32 IDs are used, this shouldn't happen\n");
+   assert(0);
+   return 0;
+}
+
+unsigned
+util_idalloc_sparse_alloc_range(struct util_idalloc_sparse *buf, unsigned num)
+{
+   unsigned max_ids = UTIL_IDALLOC_MAX_IDS_PER_SEGMENT(buf);
+   unsigned num_elems = DIV_ROUND_UP(num, 32);
+
+   /* TODO: This doesn't try to find a range that spans 2 different segments */
+   for (unsigned i = 0; i < ARRAY_SIZE(buf->segment); i++) {
+      if (buf->segment[i].lowest_free_idx + num_elems <=
+          UTIL_IDALLOC_MAX_ELEMS_PER_SEGMENT(buf)) {
+         unsigned base = util_idalloc_alloc_range(&buf->segment[i], num);
+
+         if (base + num <= max_ids)
+            return max_ids * i + base;
+
+         /* Back off the allocation and try again with the next segment. */
+         for (unsigned i = 0; i < num; i++)
+            util_idalloc_free(&buf->segment[i], base + i);
+      }
+   }
+
+   fprintf(stderr, "mesa: util_idalloc_sparse_alloc_range: can't find a free consecutive range of IDs\n");
+   assert(0);
+   return 0;
+}
+
+void
+util_idalloc_sparse_free(struct util_idalloc_sparse *buf, unsigned id)
+{
+   unsigned max_ids = UTIL_IDALLOC_MAX_IDS_PER_SEGMENT(buf);
+   util_idalloc_free(&buf->segment[id / max_ids], id % max_ids);
+}
+
+void
+util_idalloc_sparse_reserve(struct util_idalloc_sparse *buf, unsigned id)
+{
+   unsigned max_ids = UTIL_IDALLOC_MAX_IDS_PER_SEGMENT(buf);
+   util_idalloc_reserve(&buf->segment[id / max_ids], id % max_ids);
+}

src/util/u_idalloc.h

@@ -108,6 +108,51 @@ util_idalloc_mt_alloc(struct util_idalloc_mt *buf);
 void
 util_idalloc_mt_free(struct util_idalloc_mt *buf, unsigned id);
 
+/* util_idalloc_sparse: The 32-bit ID range is divided into separately managed
+ * segments. This reduces virtual memory usage when IDs are sparse.
+ * It's done by layering util_idalloc_sparse on top of util_idalloc.
+ *
+ * If the last ID is allocated:
+ * - "util_idalloc" occupies 512 MB of virtual memory
+ * - "util_idalloc_sparse" occupies only 512 KB of virtual memory
+ */
+struct util_idalloc_sparse {
+   struct util_idalloc segment[1024];
+};
+
+#define UTIL_IDALLOC_MAX_IDS_PER_SEGMENT(buf) \
+   ((uint32_t)(BITFIELD64_BIT(32) / ARRAY_SIZE((buf)->segment)))
+
+#define UTIL_IDALLOC_MAX_ELEMS_PER_SEGMENT(buf) \
+   (UTIL_IDALLOC_MAX_IDS_PER_SEGMENT(buf) / 32)
+
+void
+util_idalloc_sparse_init(struct util_idalloc_sparse *buf);
+
+void
+util_idalloc_sparse_fini(struct util_idalloc_sparse *buf);
+
+unsigned
+util_idalloc_sparse_alloc(struct util_idalloc_sparse *buf);
+
+unsigned
+util_idalloc_sparse_alloc_range(struct util_idalloc_sparse *buf, unsigned num);
+
+void
+util_idalloc_sparse_free(struct util_idalloc_sparse *buf, unsigned id);
+
+void
+util_idalloc_sparse_reserve(struct util_idalloc_sparse *buf, unsigned id);
+
+/* This allows freeing IDs while iterating, excluding ID=0. */
+#define util_idalloc_sparse_foreach_no_zero_safe(buf, id) \
+   for (uint32_t _s = 0; _s < ARRAY_SIZE((buf)->segment); _s++) \
+      for (uint32_t _i = 0, _bit, id, _count = (buf)->segment[_s].num_set_elements, \
+            _mask = _count ? (buf)->segment[_s].data[0] & ~0x1 : 0; \
+           _i < _count; _mask = ++_i < _count ? (buf)->segment[_s].data[_i] : 0) \
+         while (_mask) \
+            if ((_bit = u_bit_scan(&_mask), id = _s * UTIL_IDALLOC_MAX_IDS_PER_SEGMENT(buf) + _i * 32 + _bit), \
+                (buf)->segment[_s].data[_i] & BITFIELD_BIT(_bit))
 
 #ifdef __cplusplus
 }