gpu: nvgpu: use nvgpu rbtree for bitmap allocator

Use nvgpu rbtree instead of linux rbtree for bitmap allocator
Move to use nvgpu_rbtree_node structure and
nvgpu_rbtree_* APIs

Jira NVGPU-13

Change-Id: I25543cd737ed16728a67b42f70c4d233b8cc2a63
Signed-off-by: Deepak Nibade <dnibade@nvidia.com>
Reviewed-on: http://git-master/r/1457859
Reviewed-by: svccoveritychecker <svccoveritychecker@nvidia.com>
GVS: Gerrit_Virtual_Submit
Reviewed-by: Bharat Nihalani <bnihalani@nvidia.com>

2 files changed, 26 insertions, 42 deletions

diff --git a/drivers/gpu/nvgpu/common/mm/bitmap_allocator.c b/drivers/gpu/nvgpu/common/mm/bitmap_allocator.c
index c65f33c3..4b44bd7e 100644
--- a/drivers/gpu/nvgpu/common/mm/bitmap_allocator.c
+++ b/drivers/gpu/nvgpu/common/mm/bitmap_allocator.c
@@ -132,27 +132,10 @@ static void nvgpu_bitmap_free_fixed(struct nvgpu_allocator *__a,
 static void insert_alloc_metadata(struct nvgpu_bitmap_allocator *a,
                                   struct nvgpu_bitmap_alloc *alloc)
 {
-        struct rb_node **new = &a->allocs.rb_node;
-        struct rb_node *parent = NULL;
-        struct nvgpu_bitmap_alloc *tmp;
-
-        while (*new) {
-                tmp = container_of(*new, struct nvgpu_bitmap_alloc,
-                                   alloc_entry);
-
-                parent = *new;
-                if (alloc->base < tmp->base)
-                        new = &((*new)->rb_left);
-                else if (alloc->base > tmp->base)
-                        new = &((*new)->rb_right);
-                else {
-                        WARN_ON("Duplicate entries in RB alloc tree!\n");
-                        return;
-                }
-        }
+        alloc->alloc_entry.key_start = alloc->base;
+        alloc->alloc_entry.key_end = alloc->base + alloc->length;
 
-        rb_link_node(&alloc->alloc_entry, parent, new);
-        rb_insert_color(&alloc->alloc_entry, &a->allocs);
+        nvgpu_rbtree_insert(&alloc->alloc_entry, &a->allocs);
 }
 
 /*
@@ -161,25 +144,16 @@ static void insert_alloc_metadata(struct nvgpu_bitmap_allocator *a,
 static struct nvgpu_bitmap_alloc *find_alloc_metadata(
         struct nvgpu_bitmap_allocator *a, u64 addr)
 {
-        struct rb_node *node = a->allocs.rb_node;
         struct nvgpu_bitmap_alloc *alloc;
+        struct nvgpu_rbtree_node *node = NULL;
 
-        while (node) {
-                alloc = container_of(node, struct nvgpu_bitmap_alloc,
-                                     alloc_entry);
-
-                if (addr < alloc->base)
-                        node = node->rb_left;
-                else if (addr > alloc->base)
-                        node = node->rb_right;
-                else
-                        break;
-        }
-
+        nvgpu_rbtree_search(addr, &node, a->allocs);
         if (!node)
                 return NULL;
 
-        rb_erase(node, &a->allocs);
+        alloc = nvgpu_bitmap_alloc_from_rbtree_node(node);
+
+        nvgpu_rbtree_unlink(node, &a->allocs);
 
         return alloc;
 }
@@ -316,17 +290,19 @@ static void nvgpu_bitmap_alloc_destroy(struct nvgpu_allocator *__a)
 {
         struct nvgpu_bitmap_allocator *a = bitmap_allocator(__a);
         struct nvgpu_bitmap_alloc *alloc;
-        struct rb_node *node;
+        struct nvgpu_rbtree_node *node;
 
         /*
          * Kill any outstanding allocations.
          */
-        while ((node = rb_first(&a->allocs)) != NULL) {
-                alloc = container_of(node, struct nvgpu_bitmap_alloc,
-                                     alloc_entry);
+        nvgpu_rbtree_enum_start(0, &node, a->allocs);
+        while (node) {
+                alloc = nvgpu_bitmap_alloc_from_rbtree_node(node);
 
-                rb_erase(node, &a->allocs);
+                nvgpu_rbtree_unlink(node, &a->allocs);
                 nvgpu_kmem_cache_free(a->meta_data_cache, alloc);
+
+                nvgpu_rbtree_enum_start(0, &node, a->allocs);
         }
 
         nvgpu_kmem_cache_destroy(a->meta_data_cache);
@@ -419,6 +395,7 @@ int nvgpu_bitmap_allocator_init(struct gk20a *g, struct nvgpu_allocator *__a,
         a->num_bits = length >> a->blk_shift;
         a->bit_offs = a->base >> a->blk_shift;
         a->flags = flags;
+        a->allocs = NULL;
 
         a->bitmap = nvgpu_kcalloc(g, BITS_TO_LONGS(a->num_bits),
                                   sizeof(*a->bitmap));
diff --git a/drivers/gpu/nvgpu/common/mm/bitmap_allocator_priv.h b/drivers/gpu/nvgpu/common/mm/bitmap_allocator_priv.h
index 95780202..4f0b3012 100644
--- a/drivers/gpu/nvgpu/common/mm/bitmap_allocator_priv.h
+++ b/drivers/gpu/nvgpu/common/mm/bitmap_allocator_priv.h
@@ -17,8 +17,8 @@
 #ifndef BITMAP_ALLOCATOR_PRIV_H
 #define BITMAP_ALLOCATOR_PRIV_H
 
-#include <linux/rbtree.h>
 
+#include <nvgpu/rbtree.h>
 #include <nvgpu/kmem.h>
 
 struct nvgpu_allocator;
@@ -43,7 +43,7 @@ struct nvgpu_bitmap_allocator {
         u64 next_blk;
 
         unsigned long *bitmap;          /* The actual bitmap! */
-        struct rb_root allocs;          /* Tree of outstanding allocations. */
+        struct nvgpu_rbtree_node *allocs;  /* Tree of outstanding allocations */
 
         struct nvgpu_kmem_cache *meta_data_cache;
 
@@ -61,7 +61,14 @@ struct nvgpu_bitmap_allocator {
 struct nvgpu_bitmap_alloc {
         u64 base;
         u64 length;
-        struct rb_node alloc_entry;     /* RB tree of allocations. */
+        struct nvgpu_rbtree_node alloc_entry;   /* RB tree of allocations. */
+};
+
+static inline struct nvgpu_bitmap_alloc *
+nvgpu_bitmap_alloc_from_rbtree_node(struct nvgpu_rbtree_node *node)
+{
+        return (struct nvgpu_bitmap_alloc *)
+        ((uintptr_t)node - offsetof(struct nvgpu_bitmap_alloc, alloc_entry));
 };
 
 static inline struct nvgpu_bitmap_allocator *bitmap_allocator(