6 files changed, 37 insertions, 34 deletions

diff --git a/drivers/gpu/nvgpu/common/mm/bitmap_allocator.c b/drivers/gpu/nvgpu/common/mm/bitmap_allocator.c
index 6fc508d6..6e3bad6f 100644
--- a/drivers/gpu/nvgpu/common/mm/bitmap_allocator.c
+++ b/drivers/gpu/nvgpu/common/mm/bitmap_allocator.c
@@ -19,6 +19,7 @@
 #include <linux/bitops.h>
 
 #include <nvgpu/allocator.h>
+#include <nvgpu/kmem.h>
 
 #include "bitmap_allocator_priv.h"
 
@@ -248,12 +249,11 @@ static u64 nvgpu_bitmap_alloc(struct nvgpu_allocator *__a, u64 len)
 
         /*
          * Only do meta-data storage if we are allowed to allocate storage for
-         * that meta-data. The issue with using kmalloc() and friends is that
+         * that meta-data. The issue with using malloc and friends is that
          * in latency and success critical paths an alloc_page() call can either
-         * sleep for potentially a long time or, assuming GFP_ATOMIC, fail.
-         * Since we might not want either of these possibilities assume that the
-         * caller will keep what data it needs around to successfully free this
-         * allocation.
+         * sleep for potentially a long time or fail. Since we might not want
+         * either of these possibilities assume that the caller will keep what
+         * data it needs around to successfully free this allocation.
          */
         if (!(a->flags & GPU_ALLOC_NO_ALLOC_PAGE) &&
             __nvgpu_bitmap_store_alloc(a, addr, blks * a->blk_size))
@@ -332,8 +332,8 @@ static void nvgpu_bitmap_alloc_destroy(struct nvgpu_allocator *__a)
         }
 
         nvgpu_kmem_cache_destroy(a->meta_data_cache);
-        kfree(a->bitmap);
-        kfree(a);
+        nvgpu_kfree(nvgpu_alloc_to_gpu(__a), a->bitmap);
+        nvgpu_kfree(nvgpu_alloc_to_gpu(__a), a);
 }
 
 static void nvgpu_bitmap_print_stats(struct nvgpu_allocator *__a,
@@ -397,11 +397,11 @@ int nvgpu_bitmap_allocator_init(struct gk20a *g, struct nvgpu_allocator *__a,
                 length -= blk_size;
         }
 
-        a = kzalloc(sizeof(struct nvgpu_bitmap_allocator), GFP_KERNEL);
+        a = nvgpu_kzalloc(g, sizeof(struct nvgpu_bitmap_allocator));
         if (!a)
                 return -ENOMEM;
 
-        err = __nvgpu_alloc_common_init(__a, name, a, false, &bitmap_ops);
+        err = __nvgpu_alloc_common_init(__a, g, name, a, false, &bitmap_ops);
         if (err)
                 goto fail;
 
@@ -422,8 +422,8 @@ int nvgpu_bitmap_allocator_init(struct gk20a *g, struct nvgpu_allocator *__a,
         a->bit_offs = a->base >> a->blk_shift;
         a->flags = flags;
 
-        a->bitmap = kcalloc(BITS_TO_LONGS(a->num_bits), sizeof(*a->bitmap),
-                            GFP_KERNEL);
+        a->bitmap = nvgpu_kcalloc(g, BITS_TO_LONGS(a->num_bits),
+                                  sizeof(*a->bitmap));
         if (!a->bitmap) {
                 err = -ENOMEM;
                 goto fail;
@@ -445,6 +445,6 @@ int nvgpu_bitmap_allocator_init(struct gk20a *g, struct nvgpu_allocator *__a,
 fail:
         if (a->meta_data_cache)
                 nvgpu_kmem_cache_destroy(a->meta_data_cache);
-        kfree(a);
+        nvgpu_kfree(g, a);
         return err;
 }
diff --git a/drivers/gpu/nvgpu/common/mm/buddy_allocator.c b/drivers/gpu/nvgpu/common/mm/buddy_allocator.c
index 6f4c670a..246be974 100644
--- a/drivers/gpu/nvgpu/common/mm/buddy_allocator.c
+++ b/drivers/gpu/nvgpu/common/mm/buddy_allocator.c
@@ -18,6 +18,7 @@
 #include <linux/slab.h>
 
 #include <nvgpu/allocator.h>
+#include <nvgpu/kmem.h>
 
 #include "gk20a/mm_gk20a.h"
 #include "gk20a/platform_gk20a.h"
@@ -304,7 +305,7 @@ static void nvgpu_buddy_allocator_destroy(struct nvgpu_allocator *__a)
         }
 
         nvgpu_kmem_cache_destroy(a->buddy_cache);
-        kfree(a);
+        nvgpu_kfree(nvgpu_alloc_to_gpu(__a), a);
 
         alloc_unlock(__a);
 }
@@ -809,7 +810,7 @@ static void __balloc_do_free_fixed(struct nvgpu_buddy_allocator *a,
                 balloc_coalesce(a, bud);
         }
 
-        kfree(falloc);
+        nvgpu_kfree(nvgpu_alloc_to_gpu(a->owner), falloc);
 }
 
 /*
@@ -893,7 +894,7 @@ static u64 __nvgpu_balloc_fixed_buddy(struct nvgpu_allocator *__a,
                         goto fail;
         }
 
-        falloc = kmalloc(sizeof(*falloc), GFP_KERNEL);
+        falloc = nvgpu_kmalloc(nvgpu_alloc_to_gpu(__a), sizeof(*falloc));
         if (!falloc)
                 goto fail;
 
@@ -932,7 +933,7 @@ static u64 __nvgpu_balloc_fixed_buddy(struct nvgpu_allocator *__a,
 fail_unlock:
         alloc_unlock(__a);
 fail:
-        kfree(falloc);
+        nvgpu_kfree(nvgpu_alloc_to_gpu(__a), falloc);
         nvgpu_alloc_trace_func_done();
         return 0;
 }
@@ -1261,11 +1262,11 @@ int __nvgpu_buddy_allocator_init(struct gk20a *g, struct nvgpu_allocator *__a,
         if (flags & GPU_ALLOC_GVA_SPACE && !vm)
                 return -EINVAL;
 
-        a = kzalloc(sizeof(struct nvgpu_buddy_allocator), GFP_KERNEL);
+        a = nvgpu_kzalloc(g, sizeof(struct nvgpu_buddy_allocator));
         if (!a)
                 return -ENOMEM;
 
-        err = __nvgpu_alloc_common_init(__a, name, a, false, &buddy_ops);
+        err = __nvgpu_alloc_common_init(__a, g, name, a, false, &buddy_ops);
         if (err)
                 goto fail;
 
@@ -1339,7 +1340,7 @@ int __nvgpu_buddy_allocator_init(struct gk20a *g, struct nvgpu_allocator *__a,
 fail:
         if (a->buddy_cache)
                 nvgpu_kmem_cache_destroy(a->buddy_cache);
-        kfree(a);
+        nvgpu_kfree(g, a);
         return err;
 }
 
diff --git a/drivers/gpu/nvgpu/common/mm/buddy_allocator_priv.h b/drivers/gpu/nvgpu/common/mm/buddy_allocator_priv.h
index 56aaea62..935b3f1c 100644
--- a/drivers/gpu/nvgpu/common/mm/buddy_allocator_priv.h
+++ b/drivers/gpu/nvgpu/common/mm/buddy_allocator_priv.h
@@ -20,8 +20,7 @@
 #include <linux/list.h>
 #include <linux/rbtree.h>
 
-#include <nvgpu/kmem.h>
-
+struct nvgpu_kmem_cache;
 struct nvgpu_allocator;
 struct vm_gk20a;
 
diff --git a/drivers/gpu/nvgpu/common/mm/lockless_allocator.c b/drivers/gpu/nvgpu/common/mm/lockless_allocator.c
index e3063a42..6fd9bc48 100644
--- a/drivers/gpu/nvgpu/common/mm/lockless_allocator.c
+++ b/drivers/gpu/nvgpu/common/mm/lockless_allocator.c
@@ -20,6 +20,7 @@
 #include <linux/atomic.h>
 
 #include <nvgpu/allocator.h>
+#include <nvgpu/kmem.h>
 
 #include "lockless_allocator_priv.h"
 
@@ -106,7 +107,7 @@ static void nvgpu_lockless_alloc_destroy(struct nvgpu_allocator *a)
         nvgpu_fini_alloc_debug(a);
 
         vfree(pa->next);
-        kfree(pa);
+        nvgpu_kfree(nvgpu_alloc_to_gpu(a), pa);
 }
 
 static void nvgpu_lockless_print_stats(struct nvgpu_allocator *a,
@@ -154,7 +155,7 @@ int nvgpu_lockless_allocator_init(struct gk20a *g, struct nvgpu_allocator *__a,
                 return -EINVAL;
 
         /*
-         * Ensure we have space for atleast one node & there's no overflow.
+         * Ensure we have space for at least one node & there's no overflow.
          * In order to control memory footprint, we require count < INT_MAX
          */
         count = length;
@@ -162,11 +163,11 @@ int nvgpu_lockless_allocator_init(struct gk20a *g, struct nvgpu_allocator *__a,
         if (!base || !count || count > INT_MAX)
                 return -EINVAL;
 
-        a = kzalloc(sizeof(struct nvgpu_lockless_allocator), GFP_KERNEL);
+        a = nvgpu_kzalloc(g, sizeof(struct nvgpu_lockless_allocator));
         if (!a)
                 return -ENOMEM;
 
-        err = __nvgpu_alloc_common_init(__a, name, a, false, &pool_ops);
+        err = __nvgpu_alloc_common_init(__a, g, name, a, false, &pool_ops);
         if (err)
                 goto fail;
 
@@ -202,6 +203,6 @@ int nvgpu_lockless_allocator_init(struct gk20a *g, struct nvgpu_allocator *__a,
         return 0;
 
 fail:
-        kfree(a);
+        nvgpu_kfree(g, a);
         return err;
 }
diff --git a/drivers/gpu/nvgpu/common/mm/nvgpu_allocator.c b/drivers/gpu/nvgpu/common/mm/nvgpu_allocator.c
index 24a76a0b..02b7b48d 100644
--- a/drivers/gpu/nvgpu/common/mm/nvgpu_allocator.c
+++ b/drivers/gpu/nvgpu/common/mm/nvgpu_allocator.c
@@ -123,7 +123,7 @@ void nvgpu_alloc_destroy(struct nvgpu_allocator *a)
 /*
  * Handle the common init stuff for a nvgpu_allocator.
  */
-int __nvgpu_alloc_common_init(struct nvgpu_allocator *a,
+int __nvgpu_alloc_common_init(struct nvgpu_allocator *a, struct gk20a *g,
                               const char *name, void *priv, bool dbg,
                               const struct nvgpu_allocator_ops *ops)
 {
@@ -143,6 +143,7 @@ int __nvgpu_alloc_common_init(struct nvgpu_allocator *a,
         if (err)
                 return err;
 
+        a->g = g;
         a->ops = ops;
         a->priv = priv;
         a->debug = dbg;
diff --git a/drivers/gpu/nvgpu/common/mm/page_allocator.c b/drivers/gpu/nvgpu/common/mm/page_allocator.c
index 193decc9..7d2cedc9 100644
--- a/drivers/gpu/nvgpu/common/mm/page_allocator.c
+++ b/drivers/gpu/nvgpu/common/mm/page_allocator.c
@@ -21,6 +21,7 @@
 
 #include <nvgpu/allocator.h>
 #include <nvgpu/page_allocator.h>
+#include <nvgpu/kmem.h>
 
 #include "buddy_allocator_priv.h"
 
@@ -760,7 +761,7 @@ static void nvgpu_page_allocator_destroy(struct nvgpu_allocator *__a)
         struct nvgpu_page_allocator *a = page_allocator(__a);
 
         alloc_lock(__a);
-        kfree(a);
+        nvgpu_kfree(nvgpu_alloc_to_gpu(__a), a);
         __a->priv = NULL;
         alloc_unlock(__a);
 }
@@ -848,9 +849,9 @@ static int nvgpu_page_alloc_init_slabs(struct nvgpu_page_allocator *a)
         size_t nr_slabs = ilog2(a->page_size >> 12);
         unsigned int i;
 
-        a->slabs = kcalloc(nr_slabs,
-                           sizeof(struct page_alloc_slab),
-                           GFP_KERNEL);
+        a->slabs = nvgpu_kcalloc(nvgpu_alloc_to_gpu(a->owner),
+                                 nr_slabs,
+                                 sizeof(struct page_alloc_slab));
         if (!a->slabs)
                 return -ENOMEM;
         a->nr_slabs = nr_slabs;
@@ -881,11 +882,11 @@ int nvgpu_page_allocator_init(struct gk20a *g, struct nvgpu_allocator *__a,
         if (blk_size < SZ_4K)
                 return -EINVAL;
 
-        a = kzalloc(sizeof(struct nvgpu_page_allocator), GFP_KERNEL);
+        a = nvgpu_kzalloc(g, sizeof(struct nvgpu_page_allocator));
         if (!a)
                 return -ENOMEM;
 
-        err = __nvgpu_alloc_common_init(__a, name, a, false, &page_ops);
+        err = __nvgpu_alloc_common_init(__a, g, name, a, false, &page_ops);
         if (err)
                 goto fail;
 
@@ -938,6 +939,6 @@ fail:
                 nvgpu_kmem_cache_destroy(a->chunk_cache);
         if (a->slab_page_cache)
                 nvgpu_kmem_cache_destroy(a->slab_page_cache);
-        kfree(a);
+        nvgpu_kfree(g, a);
         return err;
 }