gpu: nvgpu: use nvgpu rbtree for kmem allocator

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

Jira NVGPU-13

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

2 files changed, 48 insertions, 70 deletions

diff --git a/drivers/gpu/nvgpu/common/linux/kmem.c b/drivers/gpu/nvgpu/common/linux/kmem.c
index 819b3a11..4fe68830 100644
--- a/drivers/gpu/nvgpu/common/linux/kmem.c
+++ b/drivers/gpu/nvgpu/common/linux/kmem.c
@@ -16,7 +16,6 @@
 
 #include <linux/mutex.h>
 #include <linux/slab.h>
-#include <linux/rbtree.h>
 #include <linux/debugfs.h>
 #include <linux/spinlock.h>
 #include <linux/seq_file.h>
@@ -73,53 +72,26 @@ static void kmem_print_mem_alloc(struct gk20a *g,
 static int nvgpu_add_alloc(struct nvgpu_mem_alloc_tracker *tracker,
                            struct nvgpu_mem_alloc *alloc)
 {
-        struct rb_node **new = &tracker->allocs.rb_node;
-        struct rb_node *parent = NULL;
-
-        while (*new) {
-                struct nvgpu_mem_alloc *tmp = rb_entry(*new,
-                                                       struct nvgpu_mem_alloc,
-                                                       allocs_entry);
-
-                parent = *new;
-
-                if (alloc->addr < tmp->addr)
-                        new = &(*new)->rb_left;
-                else if (alloc->addr > tmp->addr)
-                        new = &(*new)->rb_right;
-                else
-                        return -EINVAL;
-        }
-
-        /* Put the new node there */
-        rb_link_node(&alloc->allocs_entry, parent, new);
-        rb_insert_color(&alloc->allocs_entry, &tracker->allocs);
+        alloc->allocs_entry.key_start = alloc->addr;
+        alloc->allocs_entry.key_end = alloc->addr + alloc->size;
 
+        nvgpu_rbtree_insert(&alloc->allocs_entry, &tracker->allocs);
         return 0;
 }
 
 static struct nvgpu_mem_alloc *nvgpu_rem_alloc(
         struct nvgpu_mem_alloc_tracker *tracker, u64 alloc_addr)
 {
-        struct rb_node *node = tracker->allocs.rb_node;
         struct nvgpu_mem_alloc *alloc;
+        struct nvgpu_rbtree_node *node = NULL;
 
-        while (node) {
-                alloc = container_of(node,
-                                     struct nvgpu_mem_alloc, allocs_entry);
-
-                if (alloc_addr < alloc->addr)
-                        node = node->rb_left;
-                else if (alloc_addr > alloc->addr)
-                        node = node->rb_right;
-                else
-                        break;
-        }
-
+        nvgpu_rbtree_search(alloc_addr, &node, tracker->allocs);
         if (!node)
                 return NULL;
 
-        rb_erase(node, &tracker->allocs);
+        alloc = nvgpu_mem_alloc_from_rbtree_node(node);
+
+        nvgpu_rbtree_unlink(node, &tracker->allocs);
 
         return alloc;
 }
@@ -417,7 +389,7 @@ static void print_histogram(struct nvgpu_mem_alloc_tracker *tracker,
         u64 nr_buckets;
         unsigned int *buckets;
         unsigned int total_allocs;
-        struct rb_node *node;
+        struct nvgpu_rbtree_node *node;
         static const char histogram_line[] =
                 "++++++++++++++++++++++++++++++++++++++++";
 
@@ -443,15 +415,13 @@ static void print_histogram(struct nvgpu_mem_alloc_tracker *tracker,
          * should go in. Round the size down to the nearest power of two to
          * find the right bucket.
          */
-        for (node = rb_first(&tracker->allocs);
-             node != NULL;
-             node = rb_next(node)) {
+        nvgpu_rbtree_enum_start(0, &node, tracker->allocs);
+        while (node) {
                 int b;
                 u64 bucket_min;
-                struct nvgpu_mem_alloc *alloc;
+                struct nvgpu_mem_alloc *alloc =
+                        nvgpu_mem_alloc_from_rbtree_node(node);
 
-                alloc = container_of(node, struct nvgpu_mem_alloc,
-                                     allocs_entry);
                 bucket_min = (u64)rounddown_pow_of_two(alloc->size);
                 if (bucket_min < tracker->min_alloc)
                         bucket_min = tracker->min_alloc;
@@ -469,6 +439,8 @@ static void print_histogram(struct nvgpu_mem_alloc_tracker *tracker,
                         b--;
 
                 buckets[b]++;
+
+                nvgpu_rbtree_enum_next(&node, node);
         }
 
         total_allocs = 0;
@@ -569,17 +541,16 @@ static int __kmem_traces_dump_tracker(struct gk20a *g,
                                       struct nvgpu_mem_alloc_tracker *tracker,
                                       struct seq_file *s)
 {
-        struct rb_node *node;
-
-        for (node = rb_first(&tracker->allocs);
-             node != NULL;
-             node = rb_next(node)) {
-                struct nvgpu_mem_alloc *alloc;
+        struct nvgpu_rbtree_node *node;
 
-                alloc = container_of(node, struct nvgpu_mem_alloc,
-                                     allocs_entry);
+        nvgpu_rbtree_enum_start(0, &node, tracker->allocs);
+        while (node) {
+                struct nvgpu_mem_alloc *alloc =
+                        nvgpu_mem_alloc_from_rbtree_node(node);
 
                 kmem_print_mem_alloc(g, alloc, s);
+
+                nvgpu_rbtree_enum_next(&node, node);
         }
 
         return 0;
@@ -647,21 +618,19 @@ static int __do_check_for_outstanding_allocs(
         struct nvgpu_mem_alloc_tracker *tracker,
         const char *type, bool silent)
 {
-        struct rb_node *node;
+        struct nvgpu_rbtree_node *node;
         int count = 0;
 
-        for (node = rb_first(&tracker->allocs);
-             node != NULL;
-             node = rb_next(node)) {
-                struct nvgpu_mem_alloc *alloc;
-
-                alloc = container_of(node, struct nvgpu_mem_alloc,
-                                     allocs_entry);
+        nvgpu_rbtree_enum_start(0, &node, tracker->allocs);
+        while (node) {
+                struct nvgpu_mem_alloc *alloc =
+                        nvgpu_mem_alloc_from_rbtree_node(node);
 
                 if (!silent)
                         kmem_print_mem_alloc(g, alloc, NULL);
 
                 count++;
+                nvgpu_rbtree_enum_next(&node, node);
         }
 
         return count;
@@ -690,17 +659,20 @@ static int check_for_outstanding_allocs(struct gk20a *g, bool silent)
 static void do_nvgpu_kmem_cleanup(struct nvgpu_mem_alloc_tracker *tracker,
                                   void (*force_free_func)(const void *))
 {
-        struct rb_node *node;
+        struct nvgpu_rbtree_node *node;
 
-        while ((node = rb_first(&tracker->allocs)) != NULL) {
-                struct nvgpu_mem_alloc *alloc;
+        nvgpu_rbtree_enum_start(0, &node, tracker->allocs);
+        while (node) {
+                struct nvgpu_mem_alloc *alloc =
+                        nvgpu_mem_alloc_from_rbtree_node(node);
 
-                alloc = container_of(node, struct nvgpu_mem_alloc,
-                                     allocs_entry);
                 if (force_free_func)
                         force_free_func((void *)alloc->addr);
 
+                nvgpu_rbtree_unlink(node, &tracker->allocs);
                 kfree(alloc);
+
+                nvgpu_rbtree_enum_start(0, &node, tracker->allocs);
         }
 }
 
@@ -772,8 +744,8 @@ int nvgpu_kmem_init(struct gk20a *g)
         g->vmallocs->name = "vmalloc";
         g->kmallocs->name = "kmalloc";
 
-        g->vmallocs->allocs = RB_ROOT;
-        g->kmallocs->allocs = RB_ROOT;
+        g->vmallocs->allocs = NULL;
+        g->kmallocs->allocs = NULL;
 
         mutex_init(&g->vmallocs->lock);
         mutex_init(&g->kmallocs->lock);
diff --git a/drivers/gpu/nvgpu/common/linux/kmem_priv.h b/drivers/gpu/nvgpu/common/linux/kmem_priv.h
index 5e38ad5d..cd58ac28 100644
--- a/drivers/gpu/nvgpu/common/linux/kmem_priv.h
+++ b/drivers/gpu/nvgpu/common/linux/kmem_priv.h
@@ -17,7 +17,7 @@
 #ifndef __KMEM_PRIV_H__
 #define __KMEM_PRIV_H__
 
-#include <linux/rbtree.h>
+#include <nvgpu/rbtree.h>
 
 #define __pstat(s, fmt, msg...)                         \
         do {                                            \
@@ -61,8 +61,14 @@ struct nvgpu_mem_alloc {
         unsigned long size;
         unsigned long real_size;
 
-        /* Ugh - linux specific. Will need to be abstracted. */
-        struct rb_node allocs_entry;
+        struct nvgpu_rbtree_node allocs_entry;
+};
+
+static inline struct nvgpu_mem_alloc *
+nvgpu_mem_alloc_from_rbtree_node(struct nvgpu_rbtree_node *node)
+{
+        return (struct nvgpu_mem_alloc *)
+        ((uintptr_t)node - offsetof(struct nvgpu_mem_alloc, allocs_entry));
 };
 
 /*
@@ -71,7 +77,7 @@ struct nvgpu_mem_alloc {
 struct nvgpu_mem_alloc_tracker {
         const char *name;
         struct nvgpu_kmem_cache *allocs_cache;
-        struct rb_root allocs;
+        struct nvgpu_rbtree_node *allocs;
         struct mutex lock;
 
         u64 bytes_alloced;