gpu: nvgpu: Move kmem_caches to allocator

Instead of using a single static kmem_cache for each type of
data structure the allocators may want to allocate each
allocator now has its own instance of the kmem_cache. This is
done so that each GPU driver instance can accurately track how
much memory it is using.

In order to support this on older kernels a new NVGPU API has
been made,

  nvgpu_kmem_cache_create(struct gk20a *g, size_t size)

To handle the possibility that caches cannot be created with
the same name.

This patch also fixes numerous places where kfree() was wrongly
used to free kmem_cache allocs.

Bug 1799159
Bug 1823380

Change-Id: Id674f9a5445fde3f95db65ad6bf3ea990444603d
Signed-off-by: Alex Waterman <alexw@nvidia.com>
Reviewed-on: http://git-master/r/1283826
Reviewed-by: svccoveritychecker <svccoveritychecker@nvidia.com>
Reviewed-by: Terje Bergstrom <tbergstrom@nvidia.com>
GVS: Gerrit_Virtual_Submit

9 files changed, 211 insertions, 68 deletions

diff --git a/drivers/gpu/nvgpu/Makefile.nvgpu b/drivers/gpu/nvgpu/Makefile.nvgpu
index c3c5fa8c..9a061493 100644
--- a/drivers/gpu/nvgpu/Makefile.nvgpu
+++ b/drivers/gpu/nvgpu/Makefile.nvgpu
@@ -24,6 +24,7 @@ endif
 obj-$(CONFIG_GK20A) := nvgpu.o
 
 nvgpu-y := \
+        common/linux/kmem.o \
         common/linux/timers.o \
         common/mm/nvgpu_allocator.o \
         common/mm/bitmap_allocator.o \
diff --git a/drivers/gpu/nvgpu/common/linux/kmem.c b/drivers/gpu/nvgpu/common/linux/kmem.c
new file mode 100644
index 00000000..24e0ca5d
--- /dev/null
+++ b/drivers/gpu/nvgpu/common/linux/kmem.c
@@ -0,0 +1,83 @@
+/*
+ * Copyright (c) 2017, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * This program is free software; you can redistribute it and/or modify it
+ * under the terms and conditions of the GNU General Public License,
+ * version 2, as published by the Free Software Foundation.
+ *
+ * This program is distributed in the hope it will be useful, but WITHOUT
+ * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
+ * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License for
+ * more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+
+#include <linux/kernel.h>
+#include <linux/slab.h>
+#include <linux/atomic.h>
+
+#include <nvgpu/kmem.h>
+
+/*
+ * Statically declared because this needs to be shared across all nvgpu driver
+ * instances. This makes sure that all kmem caches are _definitely_ uniquely
+ * named.
+ */
+static atomic_t kmem_cache_id;
+
+/*
+ * Linux specific version of the nvgpu_kmem_cache struct. This type is
+ * completely opaque to the rest of the driver.
+ */
+struct nvgpu_kmem_cache {
+        struct gk20a *g;
+        struct kmem_cache *cache;
+
+        /*
+         * Memory to hold the kmem_cache unique name. Only necessary on our
+         * k3.10 kernel when not using the SLUB allocator but it's easier to
+         * just carry this on to newer kernels.
+         */
+        char name[128];
+};
+
+struct nvgpu_kmem_cache *nvgpu_kmem_cache_create(struct gk20a *g, size_t size)
+{
+        struct nvgpu_kmem_cache *cache =
+                kzalloc(sizeof(struct nvgpu_kmem_cache), GFP_KERNEL);
+
+        if (!cache)
+                return NULL;
+
+        cache->g = g;
+
+        snprintf(cache->name, sizeof(cache->name),
+                 "nvgpu-cache-0x%p-%d-%d", g, (int)size,
+                 atomic_inc_return(&kmem_cache_id));
+        cache->cache = kmem_cache_create(cache->name,
+                                         size, size, 0, NULL);
+        if (!cache->cache) {
+                kfree(cache);
+                return NULL;
+        }
+
+        return cache;
+}
+
+void nvgpu_kmem_cache_destroy(struct nvgpu_kmem_cache *cache)
+{
+        kmem_cache_destroy(cache->cache);
+        kfree(cache);
+}
+
+void *nvgpu_kmem_cache_alloc(struct nvgpu_kmem_cache *cache)
+{
+        return kmem_cache_alloc(cache->cache, GFP_KERNEL);
+}
+
+void nvgpu_kmem_cache_free(struct nvgpu_kmem_cache *cache, void *ptr)
+{
+        kmem_cache_free(cache->cache, ptr);
+}
diff --git a/drivers/gpu/nvgpu/common/mm/bitmap_allocator.c b/drivers/gpu/nvgpu/common/mm/bitmap_allocator.c
index 5042980f..6fc508d6 100644
--- a/drivers/gpu/nvgpu/common/mm/bitmap_allocator.c
+++ b/drivers/gpu/nvgpu/common/mm/bitmap_allocator.c
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2016, NVIDIA CORPORATION.  All rights reserved.
+ * Copyright (c) 2016-2017, NVIDIA CORPORATION.  All rights reserved.
  *
  * This program is free software; you can redistribute it and/or modify it
  * under the terms and conditions of the GNU General Public License,
@@ -22,9 +22,6 @@
 
 #include "bitmap_allocator_priv.h"
 
-static struct kmem_cache *meta_data_cache;      /* slab cache for meta data. */
-static DEFINE_MUTEX(meta_data_cache_lock);
-
 static u64 nvgpu_bitmap_alloc_length(struct nvgpu_allocator *a)
 {
         struct nvgpu_bitmap_allocator *ba = a->priv;
@@ -195,7 +192,7 @@ static int __nvgpu_bitmap_store_alloc(struct nvgpu_bitmap_allocator *a,
                                       u64 addr, u64 len)
 {
         struct nvgpu_bitmap_alloc *alloc =
-                kmem_cache_alloc(meta_data_cache, GFP_KERNEL);
+                nvgpu_kmem_cache_alloc(a->meta_data_cache);
 
         if (!alloc)
                 return -ENOMEM;
@@ -312,7 +309,8 @@ static void nvgpu_bitmap_free(struct nvgpu_allocator *__a, u64 addr)
         a->bytes_freed += alloc->length;
 
 done:
-        kfree(alloc);
+        if (a->meta_data_cache && alloc)
+                nvgpu_kmem_cache_free(a->meta_data_cache, alloc);
         alloc_unlock(__a);
 }
 
@@ -330,9 +328,10 @@ static void nvgpu_bitmap_alloc_destroy(struct nvgpu_allocator *__a)
                                      alloc_entry);
 
                 rb_erase(node, &a->allocs);
-                kfree(alloc);
+                nvgpu_kmem_cache_free(a->meta_data_cache, alloc);
         }
 
+        nvgpu_kmem_cache_destroy(a->meta_data_cache);
         kfree(a->bitmap);
         kfree(a);
 }
@@ -382,14 +381,6 @@ int nvgpu_bitmap_allocator_init(struct gk20a *g, struct nvgpu_allocator *__a,
         int err;
         struct nvgpu_bitmap_allocator *a;
 
-        mutex_lock(&meta_data_cache_lock);
-        if (!meta_data_cache)
-                meta_data_cache = KMEM_CACHE(nvgpu_bitmap_alloc, 0);
-        mutex_unlock(&meta_data_cache_lock);
-
-        if (!meta_data_cache)
-                return -ENOMEM;
-
         if (WARN_ON(blk_size & (blk_size - 1)))
                 return -EINVAL;
 
@@ -414,6 +405,15 @@ int nvgpu_bitmap_allocator_init(struct gk20a *g, struct nvgpu_allocator *__a,
         if (err)
                 goto fail;
 
+        if (!(flags & GPU_ALLOC_NO_ALLOC_PAGE)) {
+                a->meta_data_cache = nvgpu_kmem_cache_create(g,
+                                        sizeof(struct nvgpu_bitmap_alloc));
+                if (!a->meta_data_cache) {
+                        err = -ENOMEM;
+                        goto fail;
+                }
+        }
+
         a->base = base;
         a->length = length;
         a->blk_size = blk_size;
@@ -424,8 +424,10 @@ int nvgpu_bitmap_allocator_init(struct gk20a *g, struct nvgpu_allocator *__a,
 
         a->bitmap = kcalloc(BITS_TO_LONGS(a->num_bits), sizeof(*a->bitmap),
                             GFP_KERNEL);
-        if (!a->bitmap)
+        if (!a->bitmap) {
+                err = -ENOMEM;
                 goto fail;
+        }
 
         wmb();
         a->inited = true;
@@ -441,6 +443,8 @@ int nvgpu_bitmap_allocator_init(struct gk20a *g, struct nvgpu_allocator *__a,
         return 0;
 
 fail:
+        if (a->meta_data_cache)
+                nvgpu_kmem_cache_destroy(a->meta_data_cache);
         kfree(a);
         return err;
 }
diff --git a/drivers/gpu/nvgpu/common/mm/bitmap_allocator_priv.h b/drivers/gpu/nvgpu/common/mm/bitmap_allocator_priv.h
index 9802b9db..95780202 100644
--- a/drivers/gpu/nvgpu/common/mm/bitmap_allocator_priv.h
+++ b/drivers/gpu/nvgpu/common/mm/bitmap_allocator_priv.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2016, NVIDIA CORPORATION.  All rights reserved.
+ * Copyright (c) 2016-2017, NVIDIA CORPORATION.  All rights reserved.
  *
  * This program is free software; you can redistribute it and/or modify it
  * under the terms and conditions of the GNU General Public License,
@@ -19,6 +19,8 @@
 
 #include <linux/rbtree.h>
 
+#include <nvgpu/kmem.h>
+
 struct nvgpu_allocator;
 
 struct nvgpu_bitmap_allocator {
@@ -43,6 +45,8 @@ struct nvgpu_bitmap_allocator {
         unsigned long *bitmap;          /* The actual bitmap! */
         struct rb_root allocs;          /* Tree of outstanding allocations. */
 
+        struct nvgpu_kmem_cache *meta_data_cache;
+
         u64 flags;
 
         bool inited;
diff --git a/drivers/gpu/nvgpu/common/mm/buddy_allocator.c b/drivers/gpu/nvgpu/common/mm/buddy_allocator.c
index eee0b634..6f4c670a 100644
--- a/drivers/gpu/nvgpu/common/mm/buddy_allocator.c
+++ b/drivers/gpu/nvgpu/common/mm/buddy_allocator.c
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2016, NVIDIA CORPORATION.  All rights reserved.
+ * Copyright (c) 2016-2017, NVIDIA CORPORATION.  All rights reserved.
  *
  * This program is free software; you can redistribute it and/or modify it
  * under the terms and conditions of the GNU General Public License,
@@ -24,8 +24,6 @@
 
 #include "buddy_allocator_priv.h"
 
-static struct kmem_cache *buddy_cache;  /* slab cache for meta data. */
-
 /* Some other buddy allocator functions. */
 static struct nvgpu_buddy *balloc_free_buddy(struct nvgpu_buddy_allocator *a,
                                              u64 addr);
@@ -103,7 +101,7 @@ static struct nvgpu_buddy *balloc_new_buddy(struct nvgpu_buddy_allocator *a,
 {
         struct nvgpu_buddy *new_buddy;
 
-        new_buddy = kmem_cache_alloc(buddy_cache, GFP_KERNEL);
+        new_buddy = nvgpu_kmem_cache_alloc(a->buddy_cache);
         if (!new_buddy)
                 return NULL;
 
@@ -232,7 +230,7 @@ cleanup:
                         buddy = list_first_entry(balloc_get_order_list(a, i),
                                         struct nvgpu_buddy, buddy_entry);
                         balloc_blist_rem(a, buddy);
-                        kmem_cache_free(buddy_cache, buddy);
+                        nvgpu_kmem_cache_free(a->buddy_cache, buddy);
                 }
         }
 
@@ -285,7 +283,7 @@ static void nvgpu_buddy_allocator_destroy(struct nvgpu_allocator *__a)
                         bud = list_first_entry(balloc_get_order_list(a, i),
                                                struct nvgpu_buddy, buddy_entry);
                         balloc_blist_rem(a, bud);
-                        kmem_cache_free(buddy_cache, bud);
+                        nvgpu_kmem_cache_free(a->buddy_cache, bud);
                 }
 
                 if (a->buddy_list_len[i] != 0) {
@@ -305,6 +303,7 @@ static void nvgpu_buddy_allocator_destroy(struct nvgpu_allocator *__a)
                 }
         }
 
+        nvgpu_kmem_cache_destroy(a->buddy_cache);
         kfree(a);
 
         alloc_unlock(__a);
@@ -348,8 +347,8 @@ static void balloc_coalesce(struct nvgpu_buddy_allocator *a,
         balloc_coalesce(a, parent);
 
         /* Clean up the remains. */
-        kmem_cache_free(buddy_cache, b->buddy);
-        kmem_cache_free(buddy_cache, b);
+        nvgpu_kmem_cache_free(a->buddy_cache, b->buddy);
+        nvgpu_kmem_cache_free(a->buddy_cache, b);
 }
 
 /*
@@ -371,7 +370,7 @@ static int balloc_split_buddy(struct nvgpu_buddy_allocator *a,
 
         right = balloc_new_buddy(a, b, b->start + half, b->order - 1);
         if (!right) {
-                kmem_cache_free(buddy_cache, left);
+                nvgpu_kmem_cache_free(a->buddy_cache, left);
                 return -ENOMEM;
         }
 
@@ -783,7 +782,7 @@ err_and_cleanup:
 
                 __balloc_buddy_list_rem(a, bud);
                 balloc_free_buddy(a, bud->start);
-                kmem_cache_free(buddy_cache, bud);
+                nvgpu_kmem_cache_free(a->buddy_cache, bud);
         }
 
         return 0;
@@ -1307,10 +1306,8 @@ int __nvgpu_buddy_allocator_init(struct gk20a *g, struct nvgpu_allocator *__a,
         balloc_allocator_align(a);
         balloc_compute_max_order(a);
 
-        /* Shared buddy kmem_cache for all allocators. */
-        if (!buddy_cache)
-                buddy_cache = KMEM_CACHE(nvgpu_buddy, 0);
-        if (!buddy_cache) {
+        a->buddy_cache = nvgpu_kmem_cache_create(g, sizeof(struct nvgpu_buddy));
+        if (!a->buddy_cache) {
                 err = -ENOMEM;
                 goto fail;
         }
@@ -1340,6 +1337,8 @@ int __nvgpu_buddy_allocator_init(struct gk20a *g, struct nvgpu_allocator *__a,
         return 0;
 
 fail:
+        if (a->buddy_cache)
+                nvgpu_kmem_cache_destroy(a->buddy_cache);
         kfree(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 50a11f14..56aaea62 100644
--- a/drivers/gpu/nvgpu/common/mm/buddy_allocator_priv.h
+++ b/drivers/gpu/nvgpu/common/mm/buddy_allocator_priv.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2016, NVIDIA CORPORATION.  All rights reserved.
+ * Copyright (c) 2016-2017, NVIDIA CORPORATION.  All rights reserved.
  *
  * This program is free software; you can redistribute it and/or modify it
  * under the terms and conditions of the GNU General Public License,
@@ -20,6 +20,8 @@
 #include <linux/list.h>
 #include <linux/rbtree.h>
 
+#include <nvgpu/kmem.h>
+
 struct nvgpu_allocator;
 struct vm_gk20a;
 
@@ -126,6 +128,8 @@ struct nvgpu_buddy_allocator {
 
         struct list_head co_list;
 
+        struct nvgpu_kmem_cache *buddy_cache;
+
         /*
          * Impose an upper bound on the maximum order.
          */
diff --git a/drivers/gpu/nvgpu/common/mm/page_allocator.c b/drivers/gpu/nvgpu/common/mm/page_allocator.c
index 056c0c8f..193decc9 100644
--- a/drivers/gpu/nvgpu/common/mm/page_allocator.c
+++ b/drivers/gpu/nvgpu/common/mm/page_allocator.c
@@ -27,11 +27,6 @@
 #define palloc_dbg(a, fmt, arg...)                      \
         alloc_dbg(palloc_owner(a), fmt, ##arg)
 
-static struct kmem_cache *page_alloc_cache;
-static struct kmem_cache *page_alloc_chunk_cache;
-static struct kmem_cache *page_alloc_slab_page_cache;
-static DEFINE_MUTEX(meta_data_cache_lock);
-
 /*
  * Handle the book-keeping for these operations.
  */
@@ -147,10 +142,10 @@ static void __nvgpu_free_pages(struct nvgpu_page_allocator *a,
 
                 if (free_buddy_alloc)
                         nvgpu_free(&a->source_allocator, chunk->base);
-                kfree(chunk);
+                nvgpu_kmem_cache_free(a->chunk_cache, chunk);
         }
 
-        kmem_cache_free(page_alloc_cache, alloc);
+        nvgpu_kmem_cache_free(a->alloc_cache, alloc);
 }
 
 static int __insert_page_alloc(struct nvgpu_page_allocator *a,
@@ -213,7 +208,7 @@ static struct page_alloc_slab_page *alloc_slab_page(
 {
         struct page_alloc_slab_page *slab_page;
 
-        slab_page = kmem_cache_alloc(page_alloc_slab_page_cache, GFP_KERNEL);
+        slab_page = nvgpu_kmem_cache_alloc(a->slab_page_cache);
         if (!slab_page) {
                 palloc_dbg(a, "OOM: unable to alloc slab_page struct!\n");
                 return ERR_PTR(-ENOMEM);
@@ -223,7 +218,7 @@ static struct page_alloc_slab_page *alloc_slab_page(
 
         slab_page->page_addr = nvgpu_alloc(&a->source_allocator, a->page_size);
         if (!slab_page->page_addr) {
-                kfree(slab_page);
+                nvgpu_kmem_cache_free(a->slab_page_cache, slab_page);
                 palloc_dbg(a, "OOM: vidmem is full!\n");
                 return ERR_PTR(-ENOMEM);
         }
@@ -255,7 +250,7 @@ static void free_slab_page(struct nvgpu_page_allocator *a,
         nvgpu_free(&a->source_allocator, slab_page->page_addr);
         a->pages_freed++;
 
-        kmem_cache_free(page_alloc_slab_page_cache, slab_page);
+        nvgpu_kmem_cache_free(a->slab_page_cache, slab_page);
 }
 
 /*
@@ -352,12 +347,12 @@ static struct nvgpu_page_alloc *__nvgpu_alloc_slab(
         slab_nr = (int)ilog2(PAGE_ALIGN(len) >> 12);
         slab = &a->slabs[slab_nr];
 
-        alloc = kmem_cache_alloc(page_alloc_cache, GFP_KERNEL);
+        alloc = nvgpu_kmem_cache_alloc(a->alloc_cache);
         if (!alloc) {
                 palloc_dbg(a, "OOM: could not alloc page_alloc struct!\n");
                 goto fail;
         }
-        chunk = kmem_cache_alloc(page_alloc_chunk_cache, GFP_KERNEL);
+        chunk = nvgpu_kmem_cache_alloc(a->chunk_cache);
         if (!chunk) {
                 palloc_dbg(a, "OOM: could not alloc alloc_chunk struct!\n");
                 goto fail;
@@ -377,8 +372,10 @@ static struct nvgpu_page_alloc *__nvgpu_alloc_slab(
         return alloc;
 
 fail:
-        kfree(alloc);
-        kfree(chunk);
+        if (alloc)
+                nvgpu_kmem_cache_free(a->alloc_cache, alloc);
+        if (chunk)
+                nvgpu_kmem_cache_free(a->chunk_cache, chunk);
         return NULL;
 }
 
@@ -444,7 +441,7 @@ static struct nvgpu_page_alloc *__do_nvgpu_alloc_pages(
         u64 max_chunk_len = pages << a->page_shift;
         int i = 0;
 
-        alloc = kmem_cache_alloc(page_alloc_cache, GFP_KERNEL);
+        alloc = nvgpu_kmem_cache_alloc(a->alloc_cache);
         if (!alloc)
                 goto fail;
 
@@ -496,7 +493,7 @@ static struct nvgpu_page_alloc *__do_nvgpu_alloc_pages(
                         goto fail_cleanup;
                 }
 
-                c = kmem_cache_alloc(page_alloc_chunk_cache, GFP_KERNEL);
+                c = nvgpu_kmem_cache_alloc(a->chunk_cache);
                 if (!c) {
                         nvgpu_free(&a->source_allocator, chunk_addr);
                         goto fail_cleanup;
@@ -524,9 +521,9 @@ fail_cleanup:
                                      struct page_alloc_chunk, list_entry);
                 list_del(&c->list_entry);
                 nvgpu_free(&a->source_allocator, c->base);
-                kfree(c);
+                nvgpu_kmem_cache_free(a->chunk_cache, c);
         }
-        kfree(alloc);
+        nvgpu_kmem_cache_free(a->alloc_cache, alloc);
 fail:
         return ERR_PTR(-ENOMEM);
 }
@@ -653,8 +650,8 @@ static struct nvgpu_page_alloc *__nvgpu_alloc_pages_fixed(
         struct nvgpu_page_alloc *alloc;
         struct page_alloc_chunk *c;
 
-        alloc = kmem_cache_alloc(page_alloc_cache, GFP_KERNEL);
-        c = kmem_cache_alloc(page_alloc_chunk_cache, GFP_KERNEL);
+        alloc = nvgpu_kmem_cache_alloc(a->alloc_cache);
+        c = nvgpu_kmem_cache_alloc(a->chunk_cache);
         if (!alloc || !c)
                 goto fail;
 
@@ -675,8 +672,10 @@ static struct nvgpu_page_alloc *__nvgpu_alloc_pages_fixed(
         return alloc;
 
 fail:
-        kfree(c);
-        kfree(alloc);
+        if (c)
+                nvgpu_kmem_cache_free(a->chunk_cache, c);
+        if (alloc)
+                nvgpu_kmem_cache_free(a->alloc_cache, alloc);
         return ERR_PTR(-ENOMEM);
 }
 
@@ -879,19 +878,6 @@ int nvgpu_page_allocator_init(struct gk20a *g, struct nvgpu_allocator *__a,
         char buddy_name[sizeof(__a->name)];
         int err;
 
-        mutex_lock(&meta_data_cache_lock);
-        if (!page_alloc_cache)
-                page_alloc_cache = KMEM_CACHE(nvgpu_page_alloc, 0);
-        if (!page_alloc_chunk_cache)
-                page_alloc_chunk_cache = KMEM_CACHE(page_alloc_chunk, 0);
-        if (!page_alloc_slab_page_cache)
-                page_alloc_slab_page_cache =
-                        KMEM_CACHE(page_alloc_slab_page, 0);
-        mutex_unlock(&meta_data_cache_lock);
-
-        if (!page_alloc_cache || !page_alloc_chunk_cache)
-                return -ENOMEM;
-
         if (blk_size < SZ_4K)
                 return -EINVAL;
 
@@ -903,6 +889,17 @@ int nvgpu_page_allocator_init(struct gk20a *g, struct nvgpu_allocator *__a,
         if (err)
                 goto fail;
 
+        a->alloc_cache = nvgpu_kmem_cache_create(g,
+                                        sizeof(struct nvgpu_page_alloc));
+        a->chunk_cache = nvgpu_kmem_cache_create(g,
+                                        sizeof(struct page_alloc_chunk));
+        a->slab_page_cache = nvgpu_kmem_cache_create(g,
+                                        sizeof(struct page_alloc_slab_page));
+        if (!a->alloc_cache || !a->chunk_cache || !a->slab_page_cache) {
+                err = -ENOMEM;
+                goto fail;
+        }
+
         a->base = base;
         a->length = length;
         a->page_size = blk_size;
@@ -935,6 +932,12 @@ int nvgpu_page_allocator_init(struct gk20a *g, struct nvgpu_allocator *__a,
         return 0;
 
 fail:
+        if (a->alloc_cache)
+                nvgpu_kmem_cache_destroy(a->alloc_cache);
+        if (a->chunk_cache)
+                nvgpu_kmem_cache_destroy(a->chunk_cache);
+        if (a->slab_page_cache)
+                nvgpu_kmem_cache_destroy(a->slab_page_cache);
         kfree(a);
         return err;
 }
diff --git a/drivers/gpu/nvgpu/include/nvgpu/kmem.h b/drivers/gpu/nvgpu/include/nvgpu/kmem.h
new file mode 100644
index 00000000..3d983e77
--- /dev/null
+++ b/drivers/gpu/nvgpu/include/nvgpu/kmem.h
@@ -0,0 +1,40 @@
+/*
+ * Copyright (c) 2017, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * This program is free software; you can redistribute it and/or modify it
+ * under the terms and conditions of the GNU General Public License,
+ * version 2, as published by the Free Software Foundation.
+ *
+ * This program is distributed in the hope it will be useful, but WITHOUT
+ * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
+ * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License for
+ * more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+
+#ifndef NVGPU_KMEM_H
+#define NVGPU_KMEM_H
+
+struct gk20a;
+
+/*
+ * In Linux there is support for the notion of a kmem_cache. It gives better
+ * memory usage characteristics for lots of allocations of the same size. Think
+ * structs that get allocated over and over. Normal kmalloc() type routines
+ * typically round to the next power-of-2 since that's easy.
+ *
+ * But if we know the size ahead of time the packing for the allocations can be
+ * much better. This is the benefit of a slab allocator. This type hides the
+ * underlying kmem_cache (or absense thereof).
+ */
+struct nvgpu_kmem_cache;
+
+struct nvgpu_kmem_cache *nvgpu_kmem_cache_create(struct gk20a *g, size_t size);
+void nvgpu_kmem_cache_destroy(struct nvgpu_kmem_cache *cache);
+
+void *nvgpu_kmem_cache_alloc(struct nvgpu_kmem_cache *cache);
+void nvgpu_kmem_cache_free(struct nvgpu_kmem_cache *cache, void *ptr);
+
+#endif
diff --git a/drivers/gpu/nvgpu/include/nvgpu/page_allocator.h b/drivers/gpu/nvgpu/include/nvgpu/page_allocator.h
index 7c21c117..fa586dba 100644
--- a/drivers/gpu/nvgpu/include/nvgpu/page_allocator.h
+++ b/drivers/gpu/nvgpu/include/nvgpu/page_allocator.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2016, NVIDIA CORPORATION.  All rights reserved.
+ * Copyright (c) 2016-2017, NVIDIA CORPORATION.  All rights reserved.
  *
  * This program is free software; you can redistribute it and/or modify it
  * under the terms and conditions of the GNU General Public License,
@@ -21,6 +21,7 @@
 #include <linux/rbtree.h>
 
 #include <nvgpu/allocator.h>
+#include <nvgpu/kmem.h>
 
 struct nvgpu_allocator;
 
@@ -134,6 +135,10 @@ struct nvgpu_page_allocator {
         struct page_alloc_slab *slabs;
         int nr_slabs;
 
+        struct nvgpu_kmem_cache *alloc_cache;
+        struct nvgpu_kmem_cache *chunk_cache;
+        struct nvgpu_kmem_cache *slab_page_cache;
+
         u64 flags;
 
         /*