slab: trace for each slab object [1/1]

PD#TV-8287

Problem:
Slab memleak is hard to track;

Solution:
Add slab trace to fix it

Verify:
P212

Change-Id: Ie1c44f28d5539c7bf71f9825c617c36af65d8058
Signed-off-by: Tao Zeng <tao.zeng@amlogic.com>

5 files changed, 799 insertions, 48 deletions

diff --git a/drivers/amlogic/memory_ext/Kconfig b/drivers/amlogic/memory_ext/Kconfig
index 0da0422..b37301d 100644
--- a/drivers/amlogic/memory_ext/Kconfig
+++ b/drivers/amlogic/memory_ext/Kconfig
@@ -20,6 +20,17 @@ config AMLOGIC_PAGE_TRACE
 		with allocate page count information of each caller functions from
 		/proc/pagetrace
 
+config AMLOGIC_SLAB_TRACE
+	bool "Amlogic trace for slab usage"
+	depends on SLUB
+	depends on AMLOGIC_PAGE_TRACE
+	default y
+	help
+		Amlogic slab trace will record function address of caller for slab
+		allocate/free(kmalloc-xxxx only). trace information is stored in
+		a rb tree. And can be shown with allocate size information of
+		each caller functions from /proc/slabtrace
+
 config AMLOGIC_RAMDUMP
 	bool "Amlogic RAM DUMP support"
 	depends on AMLOGIC_MEMORY_EXTEND
diff --git a/drivers/amlogic/memory_ext/page_trace.c b/drivers/amlogic/memory_ext/page_trace.c
index 4350dbc..08f8d2c 100644
--- a/drivers/amlogic/memory_ext/page_trace.c
+++ b/drivers/amlogic/memory_ext/page_trace.c
@@ -15,7 +15,6 @@
  *
  */
 
-#include <linux/amlogic/page_trace.h>
 #include <linux/gfp.h>
 #include <linux/proc_fs.h>
 #include <linux/kallsyms.h>
@@ -29,6 +28,8 @@
 #include <linux/list.h>
 #include <linux/uaccess.h>
 #include <linux/vmalloc.h>
+#include <linux/mm.h>
+#include <linux/amlogic/page_trace.h>
 #include <asm/stacktrace.h>
 #include <asm/sections.h>
 
@@ -53,7 +54,7 @@
 static bool merge_function = 1;
 static int page_trace_filter = 64; /* not print size < page_trace_filter */
 unsigned int cma_alloc_trace;
-static struct proc_dir_entry *dentry;
+static struct proc_dir_entry *d_pagetrace;
 #ifndef CONFIG_64BIT
 struct page_trace *trace_buffer;
 static unsigned long ptrace_size;
@@ -497,6 +498,41 @@ unsigned long find_back_trace(void)
 	return 0;
 }
 
+int save_obj_stack(unsigned long *stack, int depth)
+{
+	struct stackframe frame;
+	int ret, step = 0;
+
+#ifdef CONFIG_ARM64
+	frame.fp = (unsigned long)__builtin_frame_address(0);
+	frame.sp = current_stack_pointer;
+	frame.pc = _RET_IP_;
+#ifdef CONFIG_FUNCTION_GRAPH_TRACER
+	frame.graph = current->curr_ret_stack;
+#endif
+#else
+	frame.fp = (unsigned long)__builtin_frame_address(0);
+	frame.sp = current_stack_pointer;
+	frame.lr = (unsigned long)__builtin_return_address(0);
+	frame.pc = (unsigned long)save_obj_stack;
+#endif
+	while (step < depth) {
+	#ifdef CONFIG_ARM64
+		ret = unwind_frame(current, &frame);
+	#elif defined(CONFIG_ARM)
+		ret = unwind_frame(&frame);
+	#else	/* not supported */
+		ret = -1;
+	#endif
+		if (ret < 0)
+			return 0;
+		if (step >= 1)	/* ignore first function */
+			stack[step - 1] = frame.pc;
+		step++;
+	}
+	return 0;
+}
+
 #ifdef CONFIG_64BIT
 struct page_trace *find_page_base(struct page *page)
 {
@@ -600,7 +636,7 @@ void set_page_trace(struct page *page, int order, gfp_t flag, void *func)
 {
 	unsigned long ip;
 	struct page_trace *base;
-	unsigned int val;
+	unsigned int val, i;
 
 #ifdef CONFIG_64BIT
 	if (page) {
@@ -619,6 +655,18 @@ void set_page_trace(struct page *page, int order, gfp_t flag, void *func)
 		val = pack_ip(ip, order, flag);
 		base = find_page_base(page);
 		push_ip(base, (struct page_trace *)&val);
+		if (order) {
+			/* in order to easy get trace for high order alloc */
+			val = pack_ip(ip, 0, flag);
+			for (i = 1; i < (1 << order); i++) {
+			#ifdef CONFIG_64BIT
+				base = find_page_base(++page);
+			#else
+				base += (trace_step);
+			#endif
+				push_ip(base, (struct page_trace *)&val);
+			}
+		}
 	}
 }
 EXPORT_SYMBOL(set_page_trace);
@@ -742,8 +790,9 @@ static int mt_offset[] = {
 };
 
 struct page_summary {
+	struct rb_node entry;
 	unsigned long ip;
-	unsigned int cnt;
+	unsigned long cnt;
 };
 
 struct pagetrace_summary {
@@ -767,31 +816,44 @@ static unsigned long find_ip_base(unsigned long ip)
 
 static int find_page_ip(struct page_trace *trace,
 			struct page_summary *sum, int *o,
-			int range, int *mt_cnt)
+			int range, int *mt_cnt, int size,
+			struct rb_root *root)
 {
-	int i = 0;
 	int order;
 	unsigned long ip;
+	struct rb_node **link, *parent = NULL;
+	struct page_summary *tmp;
 
 	*o = 0;
 	ip = unpack_ip(trace);
+	if (!ip || !trace->ip_data)	/* invalid ip */
+		return 0;
+
 	order = trace->order;
-	for (i = 0; i < range; i++) {
-		if (sum[i].ip == ip) {
-			/* find */
-			sum[i].cnt += (1 << order);
+	link  = &root->rb_node;
+	while (*link) {
+		parent = *link;
+		tmp = rb_entry(parent, struct page_summary, entry);
+		if (ip == tmp->ip) { /* match */
+			tmp->cnt += ((1 << order) * size);
 			*o = order;
 			return 0;
-		}
-		if (sum[i].ip == 0) {			/* empty */
-			sum[i].cnt += (1 << order);
-			sum[i].ip = ip;
-			*o = order;
-			mt_cnt[trace->migrate_type]++;
-			return 0;
-		}
+		} else if (ip < tmp->ip)
+			link = &tmp->entry.rb_left;
+		else
+			link = &tmp->entry.rb_right;
 	}
-	return -ERANGE;
+	/* not found, get a new page summary */
+	if (mt_cnt[trace->migrate_type] >= range)
+		return -ERANGE;
+	tmp       = &sum[mt_cnt[trace->migrate_type]];
+	tmp->ip   = ip;
+	tmp->cnt += ((1 << order) * size);
+	*o        = order;
+	mt_cnt[trace->migrate_type]++;
+	rb_link_node(&tmp->entry, parent, link);
+	rb_insert_color(&tmp->entry, root);
+	return 0;
 }
 
 #define K(x)		((x) << (PAGE_SHIFT - 10))
@@ -811,7 +873,7 @@ static void show_page_trace(struct seq_file *m, struct zone *zone,
 	struct page_summary *p;
 	unsigned long total_mt, total_used = 0;
 
-	seq_printf(m, "%s             %s %s\n",
+	seq_printf(m, "%s            %s, %s\n",
 		  "count(KB)", "kaddr", "function");
 	seq_puts(m, "------------------------------\n");
 	for (j = 0; j < MIGRATE_TYPES; j++) {
@@ -828,7 +890,7 @@ static void show_page_trace(struct seq_file *m, struct zone *zone,
 				continue;
 
 			if (K(p[i].cnt) >= page_trace_filter) {
-				seq_printf(m, "%8d, %16lx, %pf\n",
+				seq_printf(m, "%8ld, %16lx, %pf\n",
 					   K(p[i].cnt), p[i].ip,
 					   (void *)p[i].ip);
 			}
@@ -848,30 +910,41 @@ static void show_page_trace(struct seq_file *m, struct zone *zone,
 	seq_puts(m, "------------------------------\n");
 }
 
+static int _merge_same_function(struct page_summary *p, int range)
+{
+	int j, k, real_used;
+
+	/* first, replace all ip to entry of each function */
+	for (j = 0; j < range; j++) {
+		if (!p[j].ip)	/* Not used */
+			break;
+		p[j].ip = find_ip_base(p[j].ip);
+	}
+
+	real_used = j;
+	/* second, loop and merge same ip */
+	for (j = 0; j < real_used; j++) {
+		for (k = j + 1; k < real_used; k++) {
+			if (p[k].ip != (-1ul) &&
+			    p[k].ip == p[j].ip) {
+				p[j].cnt += p[k].cnt;
+				p[k].ip  = (-1ul);
+				p[k].cnt = 0;
+			}
+		}
+	}
+	return real_used;
+}
+
 static void merge_same_function(struct page_summary *sum, int *mt_cnt)
 {
-	int i, j, k, range;
+	int i, range;
 	struct page_summary *p;
 
 	for (i = 0; i < MIGRATE_TYPES; i++) {
 		range = mt_cnt[i];
 		p = sum + mt_offset[i];
-
-		/* first, replace all ip to entry of each function */
-		for (j = 0; j < range; j++)
-			p[j].ip = find_ip_base(p[j].ip);
-
-		/* second, loop and merge same ip */
-		for (j = 0; j < range; j++) {
-			for (k = j + 1; k < range; k++) {
-				if (p[k].ip != (-1ul) &&
-				    p[k].ip == p[j].ip) {
-					p[j].cnt += p[k].cnt;
-					p[k].ip  = (-1ul);
-					p[k].cnt = 0;
-				}
-			}
-		}
+		_merge_same_function(p, range);
 	}
 }
 
@@ -885,7 +958,9 @@ static int update_page_trace(struct seq_file *m, struct zone *zone,
 	int    order;
 	struct page_trace *trace;
 	struct page_summary *p;
+	struct rb_root root[MIGRATE_TYPES];
 
+	memset(root, 0, sizeof(root));
 	/* loop whole zone */
 	for (pfn = start_pfn; pfn < end_pfn; pfn++) {
 		struct page *page;
@@ -913,7 +988,8 @@ static int update_page_trace(struct seq_file *m, struct zone *zone,
 		mt = trace->migrate_type;
 		p  = sum + mt_offset[mt];
 		ret = find_page_ip(trace, p, &order,
-				   mt_offset[mt + 1] - mt_offset[mt], mt_cnt);
+				   mt_offset[mt + 1] - mt_offset[mt],
+				   mt_cnt, 1, &root[mt]);
 		if (ret) {
 			pr_err("mt type:%d, out of range:%d\n",
 			       mt, mt_offset[mt + 1] - mt_offset[mt]);
@@ -1035,15 +1111,535 @@ static const struct file_operations pagetrace_file_ops = {
 	.release	= single_release,
 };
 
+#ifdef CONFIG_AMLOGIC_SLAB_TRACE
+/*---------------- part 2 for slab trace ---------------------*/
+#include <linux/jhash.h>
+
+#define SLAB_TRACE_FLAG		(__GFP_ZERO | __GFP_REPEAT | GFP_ATOMIC)
+
+static LIST_HEAD(st_root);
+static int slab_trace_en __read_mostly;
+static struct kmem_cache *slab_trace_cache;
+static struct slab_stack_master *stm;
+static struct proc_dir_entry *d_slabtrace;
+
+static int __init early_slab_trace_param(char *buf)
+{
+	if (!buf)
+		return -EINVAL;
+
+	if (strcmp(buf, "off") == 0)
+		slab_trace_en = false;
+	else if (strcmp(buf, "on") == 0)
+		slab_trace_en = true;
+
+	pr_info("slab_trace %sabled\n", slab_trace_en ? "dis" : "en");
+
+	return 0;
+}
+early_param("slab_trace", early_slab_trace_param);
+
+int __init slab_trace_init(void)
+{
+	struct slab_trace_group *group = NULL;
+	struct kmem_cache *cache;
+	int cache_size;
+	char buf[64] = {0};
+	int i;
+
+	if (!slab_trace_en)
+		return -EINVAL;
+
+	for (i = 0; i <= KMALLOC_SHIFT_HIGH; i++) {
+		cache = kmalloc_caches[i];
+		if (!cache || cache->size >= PAGE_SIZE)
+			continue;
+
+		sprintf(buf, "trace_%s", cache->name);
+		group = kzalloc(sizeof(*group), GFP_KERNEL);
+		if (!group)
+			goto nomem;
+
+		cache_size = PAGE_SIZE * (1 << get_cache_max_order(cache));
+		cache_size = (cache_size / cache->size) * sizeof(int);
+		group->ip_cache = kmem_cache_create(buf, cache_size, cache_size,
+						    SLAB_NOLEAKTRACE, NULL);
+		if (!group->ip_cache)
+			goto nomem;
+
+		spin_lock_init(&group->lock);
+		list_add(&group->list, &st_root);
+		group->object_size = cache->size;
+		cache->trace_group = group;
+		pr_debug("%s, trace group %p for %s, %d:%d, cache_size:%d:%d\n",
+			 __func__, group, cache->name,
+			 cache->size, cache->object_size,
+			 cache_size, get_cache_max_order(cache));
+	}
+	stm = kzalloc(sizeof(*stm), GFP_KERNEL);
+	stm->slab_stack_cache = KMEM_CACHE(slab_stack, SLAB_NOLEAKTRACE);
+	spin_lock_init(&stm->stack_lock);
+
+	slab_trace_cache = KMEM_CACHE(slab_trace, SLAB_NOLEAKTRACE);
+	WARN_ON(!slab_trace_cache);
+	pr_info("%s, create slab trace cache:%p\n",
+		__func__, slab_trace_cache);
+
+	return 0;
+nomem:
+	pr_err("%s, failed to create trace group for %s\n",
+		__func__, buf);
+	kfree(group);
+	return -ENOMEM;
+}
+
+/*
+ * This function must under protect of lock
+ */
+static struct slab_trace *find_slab_trace(struct slab_trace_group *group,
+					  unsigned long addr)
+{
+	struct rb_node *rb;
+	struct slab_trace *trace;
+
+	rb = group->root.rb_node;
+	while (rb) {
+		trace = rb_entry(rb, struct slab_trace, entry);
+		if (addr >= trace->s_addr && addr < trace->e_addr)
+			return trace;
+		if (addr < trace->s_addr)
+			rb = trace->entry.rb_left;
+		if (addr >= trace->e_addr)
+			rb = trace->entry.rb_right;
+	}
+	return NULL;
+}
+
+int slab_trace_add_page(struct page *page, int order,
+			struct kmem_cache *s, gfp_t flag)
+{
+	struct rb_node **link, *parent = NULL;
+	struct slab_trace *trace = NULL, *tmp;
+	struct slab_trace_group *group;
+	void *buf = NULL;
+	unsigned long addr, flags;
+	int obj_cnt;
+
+	if (!slab_trace_en || !page || !s || !s->trace_group)
+		return -EINVAL;
+
+	trace = kmem_cache_alloc(slab_trace_cache, SLAB_TRACE_FLAG);
+	if (!trace)
+		goto nomem;
+
+	obj_cnt = PAGE_SIZE * (1 << order) / s->size;
+	group = s->trace_group;
+	buf = kmem_cache_alloc(group->ip_cache, SLAB_TRACE_FLAG);
+	if (!buf)
+		goto nomem;
+
+	addr = (unsigned long)page_address(page);
+	trace->s_addr = addr;
+	trace->e_addr = addr + PAGE_SIZE * (1 << order);
+	trace->object_count = obj_cnt;
+	trace->object_ip = buf;
+	/*
+	 * insert it to rb_tree;
+	 */
+	spin_lock_irqsave(&group->lock, flags);
+	link = &group->root.rb_node;
+	while (*link) {
+		parent = *link;
+		tmp = rb_entry(parent, struct slab_trace, entry);
+		if (addr < tmp->s_addr)
+			link = &tmp->entry.rb_left;
+		else if (addr >= tmp->e_addr)
+			link = &tmp->entry.rb_right;
+	}
+	rb_link_node(&trace->entry, parent, link);
+	rb_insert_color(&trace->entry, &group->root);
+	group->trace_count++;
+	spin_unlock_irqrestore(&group->lock, flags);
+	pr_debug("%s, add:%lx-%lx, buf:%p, trace:%p to group:%p, %ld, obj:%d\n",
+		s->name, addr, trace->e_addr,
+		buf, trace, group, group->trace_count, obj_cnt);
+	return 0;
+
+nomem:
+	kfree(trace);
+	pr_err("%s, failed to trace obj %p for %s, trace:%p\n", __func__,
+		page_address(page), s->name, trace);
+	return -ENOMEM;
+}
+
+int slab_trace_remove_page(struct page *page, int order, struct kmem_cache *s)
+{
+	struct slab_trace *trace = NULL;
+	struct slab_trace_group *group;
+	unsigned long addr, flags;
+
+	if (!slab_trace_en || !page || !s || !s->trace_group)
+		return -EINVAL;
+
+	addr  = (unsigned long)page_address(page);
+	group = s->trace_group;
+	spin_lock_irqsave(&group->lock, flags);
+	trace = find_slab_trace(group, addr);
+	if (!trace) {
+		spin_unlock_irqrestore(&group->lock, flags);
+		return 0;
+	}
+	rb_erase(&trace->entry, &group->root);
+	group->trace_count--;
+	spin_unlock_irqrestore(&group->lock, flags);
+	WARN_ON((addr + PAGE_SIZE * (1 << order)) != trace->e_addr);
+	pr_debug("%s, rm: %lx-%lx, buf:%p, trace:%p to group:%p, %ld\n",
+		s->name, addr, trace->e_addr,
+		trace->object_ip, trace, group, group->trace_count);
+	kfree(trace->object_ip);
+	kfree(trace);
+	return 0;
+}
+
+static int record_stack(unsigned int hash, unsigned long *stack)
+{
+	struct rb_node **link, *parent = NULL;
+	struct slab_stack *tmp, *new;
+	unsigned long flags;
+
+	/* No matched hash, we need create another one */
+	new = kmem_cache_alloc(stm->slab_stack_cache, SLAB_TRACE_FLAG);
+	if (!new)
+		return -ENOMEM;
+
+	spin_lock_irqsave(&stm->stack_lock, flags);
+	link = &stm->stack_root.rb_node;
+	while (*link) {
+		parent = *link;
+		tmp = rb_entry(parent, struct slab_stack, entry);
+		if (hash == tmp->hash) {
+			tmp->use_cnt++;
+			/* hash match, but we need check stack same? */
+			if (memcmp(stack, tmp->stack, sizeof(tmp->stack))) {
+				int i;
+
+				pr_err("%s stack hash confilct:%x\n",
+				       __func__, hash);
+				for (i = 0; i < SLAB_STACK_DEP; i++) {
+					pr_err("%16lx %16lx, %pf, %pf\n",
+					       tmp->stack[i], stack[i],
+					       (void *)tmp->stack[i],
+					       (void *)stack[i]);
+				}
+			}
+			spin_unlock_irqrestore(&stm->stack_lock, flags);
+			kfree(new);
+			return 0;
+		} else if (hash < tmp->hash)
+			link = &tmp->entry.rb_left;
+		else
+			link = &tmp->entry.rb_right;
+	}
+	/* add to stack tree */
+	new->hash    = hash;
+	new->use_cnt = 1;
+	memcpy(new->stack, stack, sizeof(new->stack));
+	rb_link_node(&new->entry, parent, link);
+	rb_insert_color(&new->entry, &stm->stack_root);
+	stm->stack_cnt++;
+	spin_unlock_irqrestore(&stm->stack_lock, flags);
+
+	return 0;
+}
+
+static struct slab_stack *get_hash_stack(unsigned int hash)
+{
+	struct rb_node *rb;
+	struct slab_stack *stack;
+
+	rb = stm->stack_root.rb_node;
+	while (rb) {
+		stack = rb_entry(rb, struct slab_stack, entry);
+		if (hash == stack->hash)
+			return stack;
+
+		if (hash < stack->hash)
+			rb = stack->entry.rb_left;
+		if (hash > stack->hash)
+			rb = stack->entry.rb_right;
+	}
+	return NULL;
+}
+
+int slab_trace_mark_object(void *object, unsigned long ip,
+			   struct kmem_cache *s)
+{
+	struct slab_trace *trace = NULL;
+	struct slab_trace_group *group;
+	unsigned long addr, flags, index;
+	unsigned long stack[SLAB_STACK_DEP] = {0};
+	unsigned int hash;
+
+	if (!slab_trace_en || !object || !s || !s->trace_group)
+		return -EINVAL;
+
+	addr  = (unsigned long)object;
+	group = s->trace_group;
+	spin_lock_irqsave(&group->lock, flags);
+	trace = find_slab_trace(group, addr);
+	spin_unlock_irqrestore(&group->lock, flags);
+	if (!trace)
+		return -ENODEV;
+
+	group->total_obj_size += s->size;
+	index = (addr - trace->s_addr) / s->size;
+	WARN_ON(index >= trace->object_count);
+	if (save_obj_stack(stack, SLAB_STACK_DEP))
+		return -EINVAL;
+	hash = jhash2((unsigned int *)stack,
+		      sizeof(stack) / sizeof(int), 0x9747b28c);
+	record_stack(hash, stack);
+	trace->object_ip[index] = hash;
+	pr_debug("%s, mk object:%p,%lx, idx:%ld, trace:%p, group:%p,%ld, %pf\n",
+		s->name, object, trace->s_addr, index,
+		trace, group, group->total_obj_size, (void *)ip);
+	return 0;
+}
+
+int slab_trace_remove_object(void *object, struct kmem_cache *s)
+{
+	struct slab_trace *trace = NULL;
+	struct slab_trace_group *group;
+	unsigned long addr, flags, index;
+	unsigned int hash, need_free = 0;
+	struct slab_stack *ss;
+
+	if (!slab_trace_en || !object || !s || !s->trace_group)
+		return -EINVAL;
+
+	addr  = (unsigned long)object;
+	group = s->trace_group;
+	spin_lock_irqsave(&group->lock, flags);
+	trace = find_slab_trace(group, addr);
+	spin_unlock_irqrestore(&group->lock, flags);
+	if (!trace)
+		return -EINVAL;
+
+	group->total_obj_size -= s->size;
+	index = (addr - trace->s_addr) / s->size;
+	WARN_ON(index >= trace->object_count);
+
+	/* remove hashed stack */
+	hash = trace->object_ip[index];
+	spin_lock_irqsave(&stm->stack_lock, flags);
+	ss = get_hash_stack(hash);
+	if (ss) {
+		ss->use_cnt--;
+		if (!ss->use_cnt) {
+			rb_erase(&ss->entry, &stm->stack_root);
+			stm->stack_cnt--;
+			need_free = 1;
+		}
+	}
+	spin_unlock_irqrestore(&stm->stack_lock, flags);
+	trace->object_ip[index] = 0;
+	if (need_free)
+		kfree(ss);
+	pr_debug("%s, rm object: %p, %lx, idx:%ld, trace:%p, group:%p, %ld\n",
+		s->name, object, trace->s_addr, index,
+		trace, group, group->total_obj_size);
+	return 0;
+}
+
+/*
+ * functions to summary slab trace
+ */
+#define SLAB_TRACE_SHOW_CNT		1024
+
+static int find_slab_hash(unsigned int hash, struct page_summary *sum,
+			  int range, int *funcs, int size, struct rb_root *root)
+{
+	struct rb_node **link, *parent = NULL;
+	struct page_summary *tmp;
+
+	link  = &root->rb_node;
+	while (*link) {
+		parent = *link;
+		tmp = rb_entry(parent, struct page_summary, entry);
+		if (hash == tmp->ip) { /* match */
+			tmp->cnt += size;
+			return 0;
+		} else if (hash < tmp->ip)
+			link = &tmp->entry.rb_left;
+		else
+			link = &tmp->entry.rb_right;
+	}
+	/* not found, get a new page summary */
+	if (*funcs >= range)
+		return -ERANGE;
+	tmp       = &sum[*funcs];
+	tmp->ip   = hash;
+	tmp->cnt += size;
+	*funcs    = *funcs + 1;
+	rb_link_node(&tmp->entry, parent, link);
+	rb_insert_color(&tmp->entry, root);
+	return 0;
+}
+
+static int update_slab_trace(struct seq_file *m, struct slab_trace_group *group,
+			     struct page_summary *sum, unsigned long *tick,
+			     int remain)
+{
+	struct rb_node *rb;
+	struct slab_trace *trace;
+	unsigned long flags, time;
+	int i, r, funcs = 0;
+	struct rb_root root = RB_ROOT;
+
+	/* This may lock long time */
+	time = sched_clock();
+	spin_lock_irqsave(&group->lock, flags);
+	for (rb = rb_first(&group->root); rb; rb = rb_next(rb)) {
+		trace = rb_entry(rb, struct slab_trace, entry);
+		for (i = 0; i < trace->object_count; i++) {
+			if (!trace->object_ip[i])
+				continue;
+
+			r = find_slab_hash(trace->object_ip[i], sum, remain,
+					   &funcs, group->object_size, &root);
+			if (r) {
+				pr_err("slab trace cout is not enough\n");
+				spin_unlock_irqrestore(&group->lock, flags);
+				return -ERANGE;
+			}
+		}
+	}
+	spin_unlock_irqrestore(&group->lock, flags);
+	*tick = sched_clock() - time;
+	return funcs;
+}
+
+static void show_slab_trace(struct seq_file *m, struct page_summary *p,
+			    int count)
+{
+	int i, j;
+	unsigned long total = 0, flags;
+	struct slab_stack *stack;
+
+	seq_printf(m, "%s          %s, %s\n",
+		  "size(bytes)", "kaddr", "function");
+	seq_puts(m, "------------------------------\n");
+
+	sort(p, count, sizeof(*p), trace_cmp, NULL);
+
+	for (i = 0; i < count; i++) {
+		if (!p[i].cnt)	/* may be empty after merge */
+			continue;
+
+		total += p[i].cnt;
+		if (p[i].cnt >= page_trace_filter) {
+			spin_lock_irqsave(&stm->stack_lock, flags);
+			stack = get_hash_stack(p[i].ip);
+			spin_unlock_irqrestore(&stm->stack_lock, flags);
+			if (!stack)
+				continue;
+
+			seq_printf(m, "%8ld, %16lx, %pf\n",
+				   p[i].cnt, stack->stack[0],
+				   (void *)stack->stack[0]);
+			for (j = 1; j < SLAB_STACK_DEP; j++) {
+				seq_printf(m, "%8s  %16lx, %pf\n",
+					   " ", stack->stack[j],
+					   (void *)stack->stack[j]);
+			}
+		}
+	}
+	seq_printf(m, "total kmalloc slabs:%6ld, %9ld kB\n",
+		   total, total >> 10);
+	seq_puts(m, "------------------------------\n");
+}
+
+static int slabtrace_show(struct seq_file *m, void *arg)
+{
+	struct page_summary *sum, *p;
+	int ret = 0, remain, alloc;
+	struct slab_trace_group *group;
+	unsigned long ticks, group_time = 0, funcs = 0;
+
+	alloc = stm->stack_cnt + 200;
+	sum   = vzalloc(sizeof(struct page_summary) * alloc);
+	if (!sum)
+		return -ENOMEM;
+	m->private = sum;
+
+	/* update only once */
+	seq_puts(m, "==============================\n");
+	p      = sum;
+	remain = alloc;
+	ticks  = sched_clock();
+	list_for_each_entry(group, &st_root, list) {
+		ret = update_slab_trace(m, group, p, &group_time, remain);
+		seq_printf(m, "%s-%4d, trace:%8ld, total:%10ld, time:%12ld, f:%d\n",
+			   "slab kmalloc", group->object_size,
+			   group->trace_count, group->total_obj_size,
+			   group_time, ret);
+		if (ret < 0) {
+			seq_printf(m, "Error %d in slab %d\n",
+				   ret, group->object_size);
+			return -ERANGE;
+		}
+		funcs  += ret;
+		p      += ret;
+		remain -= ret;
+	}
+	seq_puts(m, "------------------------------\n");
+	show_slab_trace(m, sum, funcs);
+	ticks = sched_clock() - ticks;
+
+	seq_printf(m, "SHOW_CNT:%d, tick:%ld ns, funs:%ld\n",
+		   stm->stack_cnt, ticks, funcs);
+	seq_puts(m, "==============================\n");
+
+	vfree(sum);
+
+	return 0;
+}
+
+static int slabtrace_open(struct inode *inode, struct file *file)
+{
+	return single_open(file, slabtrace_show, NULL);
+}
+static const struct file_operations slabtrace_file_ops = {
+	.open		= slabtrace_open,
+	.read		= seq_read,
+	.llseek		= seq_lseek,
+	.release	= single_release,
+};
+#endif
+
+/*---------------- end for slab trace -----------------*/
+
 static int __init page_trace_module_init(void)
 {
 
-	dentry = proc_create("pagetrace", 0444, NULL, &pagetrace_file_ops);
-	if (IS_ERR_OR_NULL(dentry)) {
+	if (!page_trace_disable)
+		d_pagetrace = proc_create("pagetrace", 0444,
+					  NULL, &pagetrace_file_ops);
+	if (IS_ERR_OR_NULL(d_pagetrace)) {
 		pr_err("%s, create sysfs failed\n", __func__);
 		return -1;
 	}
 
+#ifdef CONFIG_AMLOGIC_SLAB_TRACE
+	if (slab_trace_en)
+		d_slabtrace = proc_create("slabtrace", 0444,
+					  NULL, &slabtrace_file_ops);
+	if (IS_ERR_OR_NULL(d_slabtrace)) {
+		pr_err("%s, create sysfs failed\n", __func__);
+		return -1;
+	}
+#endif
+
 #ifndef CONFIG_64BIT
 	if (!trace_buffer)
 		return -ENOMEM;
@@ -1054,8 +1650,12 @@ static int __init page_trace_module_init(void)
 
 static void __exit page_trace_module_exit(void)
 {
-	if (dentry)
-		proc_remove(dentry);
+	if (d_pagetrace)
+		proc_remove(d_pagetrace);
+#ifdef CONFIG_AMLOGIC_SLAB_TRACE
+	if (d_slabtrace)
+		proc_remove(d_slabtrace);
+#endif
 }
 module_init(page_trace_module_init);
 module_exit(page_trace_module_exit);
diff --git a/include/linux/amlogic/page_trace.h b/include/linux/amlogic/page_trace.h
index 300e7ef..f1c7a24 100644
--- a/include/linux/amlogic/page_trace.h
+++ b/include/linux/amlogic/page_trace.h
@@ -22,6 +22,7 @@
 #include <asm/stacktrace.h>
 #include <asm/sections.h>
 #include <linux/page-flags.h>
+#include <linux/slub_def.h>
 
 /*
  * bit map lay out for _ret_ip table
@@ -56,11 +57,70 @@ struct page;
 
 /* this struct should not larger than 32 bit */
 struct page_trace {
-	unsigned int ret_ip       :24;
-	unsigned int migrate_type : 3;
-	unsigned int module_flag  : 1;
-	unsigned int order        : 4;
+	union {
+		struct {
+			unsigned int ret_ip       :24;
+			unsigned int migrate_type : 3;
+			unsigned int module_flag  : 1;
+			unsigned int order        : 4;
+		};
+		unsigned int ip_data;
+	};
+};
+
+#ifdef CONFIG_AMLOGIC_SLAB_TRACE
+/*
+ * @entry:         rb tree for quick search/insert/delete
+ * @s_addr:        start address for this slab object
+ * @e_addr:        end address for this slab object
+ * @object_count:  how many objects in this slab obj
+ * @object_ip: a   array stores ip for each slab object
+ */
+struct slab_trace {
+	struct rb_node entry;
+	unsigned long s_addr;
+	unsigned long e_addr;
+	unsigned int object_count;
+	unsigned int *object_ip;
+};
+
+/*
+ * @trace_count:    how many slab_trace object we have used
+ * @total_obj_size: total object size according obj size
+ * @lock:           protection for rb tree update
+ * @list:           link to root list
+ * @root:           root for rb tree
+ */
+struct slab_trace_group {
+	unsigned long trace_count;
+	unsigned long total_obj_size;
+	unsigned int  object_size;
+	spinlock_t   lock;
+	struct list_head list;
+	struct kmem_cache *ip_cache;
+	struct rb_root root;
+};
+
+#define SLAB_STACK_DEP		7
+/*
+ * @hash: hash value for stack
+ * @entry: rb tree for quick search
+ * @stack: stack for object
+ */
+struct slab_stack {
+	unsigned int hash;
+	unsigned int use_cnt;
+	struct rb_node entry;
+	unsigned long stack[SLAB_STACK_DEP];
+};
+
+struct slab_stack_master {
+	int stack_cnt;
+	spinlock_t stack_lock;
+	struct kmem_cache *slab_stack_cache;
+	struct rb_root stack_root;
 };
+#endif
 
 #ifdef CONFIG_AMLOGIC_PAGE_TRACE
 extern unsigned int cma_alloc_trace;
@@ -74,6 +134,18 @@ extern struct page_trace *find_page_base(struct page *page);
 extern unsigned long find_back_trace(void);
 extern unsigned long get_page_trace(struct page *page);
 extern void show_data(unsigned long addr, int nbytes, const char *name);
+extern int save_obj_stack(unsigned long *stack, int depth);
+#ifdef CONFIG_AMLOGIC_SLAB_TRACE
+extern int slab_trace_init(void);
+extern int slab_trace_add_page(struct page *page, int order,
+			       struct kmem_cache *s, gfp_t flags);
+extern int slab_trace_remove_page(struct page *page, int order,
+				  struct kmem_cache *s);
+extern int slab_trace_mark_object(void *object, unsigned long ip,
+				  struct kmem_cache *s);
+extern int slab_trace_remove_object(void *object, struct kmem_cache *s);
+extern int get_cache_max_order(struct kmem_cache *s);
+#endif
 #else
 static inline unsigned long unpack_ip(struct page_trace *trace)
 {
@@ -100,10 +172,40 @@ static inline unsigned long get_page_trace(struct page *page)
 {
 	return 0;
 }
+static inline int slab_trace_init(void)
+{
+	return 0;
+}
+static inline int slab_trace_add_page(struct page *page, int order,
+				      struct kmem_cache *s, gfp_t flags)
+{
+	return 0;
+}
+static inline int slab_trace_remove_page(struct page *page, int order,
+					 struct kmem_cache *s)
+{
+	return 0;
+}
+static inline int slab_trace_mark_object(void *object, unsigned long ip,
+					 struct kmem_cache *s)
+{
+	return 0;
+}
+static inline int slab_trace_remove_object(void *object, struct kmem_cache *s)
+{
+	return 0;
+}
+static inline int get_cache_max_order(struct kmem_cache *s)
+{
+	return 0;
+}
+static inline int save_obj_stack(unsigned long *stack, int depth)
+{
+	return 0;
+}
 #endif
 
 #ifdef CONFIG_AMLOGIC_SLUB_DEBUG
-#include <linux/slub_def.h>
 extern int aml_slub_check_object(struct kmem_cache *s, void *p, void *q);
 extern void aml_get_slub_trace(struct kmem_cache *s, struct page *page,
 				gfp_t flags, int order);
diff --git a/include/linux/slub_def.h b/include/linux/slub_def.h
index b6a59e8..bd2a206 100644
--- a/include/linux/slub_def.h
+++ b/include/linux/slub_def.h
@@ -7,6 +7,10 @@
  * (C) 2007 SGI, Christoph Lameter
  */
 #include <linux/kobject.h>
+#ifdef CONFIG_AMLOGIC_PAGE_TRACE
+#include <linux/amlogic/page_trace.h>
+#endif
+
 
 enum stat_item {
 	ALLOC_FASTPATH,		/* Allocation from cpu slab */
@@ -108,6 +112,9 @@ struct kmem_cache {
 #ifdef CONFIG_KASAN
 	struct kasan_cache kasan_info;
 #endif
+#ifdef CONFIG_AMLOGIC_SLAB_TRACE
+	struct slab_trace_group *trace_group;
+#endif
 
 	struct kmem_cache_node *node[MAX_NUMNODES];
 };
diff --git a/mm/slub.c b/mm/slub.c
index ec0b8d8..f182706 100644
--- a/mm/slub.c
+++ b/mm/slub.c
@@ -1412,6 +1412,9 @@ static inline struct page *alloc_slab_page(struct kmem_cache *s,
 		page = NULL;
 	}
 
+#ifdef CONFIG_AMLOGIC_SLAB_TRACE
+	slab_trace_add_page(page, order, s, flags);
+#endif
 	return page;
 }
 
@@ -1664,6 +1667,9 @@ static void __free_slab(struct kmem_cache *s, struct page *page)
 	if (current->reclaim_state)
 		current->reclaim_state->reclaimed_slab += pages;
 	memcg_uncharge_slab(page, order, s);
+#ifdef CONFIG_AMLOGIC_SLAB_TRACE
+	slab_trace_remove_page(page, order, s);
+#endif
 	__free_pages(page, order);
 }
 
@@ -2561,6 +2567,9 @@ load_freelist:
 	VM_BUG_ON(!c->page->frozen);
 	c->freelist = get_freepointer(s, freelist);
 	c->tid = next_tid(c->tid);
+#ifdef CONFIG_AMLOGIC_SLAB_TRACE
+	slab_trace_mark_object(freelist, addr, s);
+#endif
 	return freelist;
 
 new_slab:
@@ -2592,6 +2601,9 @@ new_slab:
 	deactivate_slab(s, page, get_freepointer(s, freelist));
 	c->page = NULL;
 	c->freelist = NULL;
+#ifdef CONFIG_AMLOGIC_SLAB_TRACE
+	slab_trace_mark_object(freelist, addr, s);
+#endif
 	return freelist;
 }
 
@@ -2707,6 +2719,9 @@ redo:
 		}
 		prefetch_freepointer(s, next_object);
 		stat(s, ALLOC_FASTPATH);
+	#ifdef CONFIG_AMLOGIC_SLAB_TRACE
+		slab_trace_mark_object(object, addr, s);
+	#endif
 	}
 
 	if (unlikely(gfpflags & __GFP_ZERO) && object)
@@ -2934,6 +2949,9 @@ redo:
 	/* Same with comment on barrier() in slab_alloc_node() */
 	barrier();
 
+#ifdef CONFIG_AMLOGIC_SLAB_TRACE
+	slab_trace_remove_object(head, s);
+#endif
 	if (likely(page == c->page)) {
 		set_freepointer(s, tail_obj, c->freelist);
 
@@ -3028,6 +3046,9 @@ int build_detached_freelist(struct kmem_cache *s, size_t size,
 		if (unlikely(!PageSlab(page))) {
 			BUG_ON(!PageCompound(page));
 			kfree_hook(object);
+		#ifdef CONFIG_AMLOGIC_SLAB_TRACE
+			slab_trace_remove_page(page, compound_order(page), s);
+		#endif
 			__free_pages(page, compound_order(page));
 			p[size] = NULL; /* mark object processed */
 			return size;
@@ -3529,6 +3550,13 @@ static int calculate_sizes(struct kmem_cache *s, int forced_order)
 	return !!oo_objects(s->oo);
 }
 
+#ifdef CONFIG_AMLOGIC_SLAB_TRACE
+int get_cache_max_order(struct kmem_cache *s)
+{
+	return oo_order(s->oo);
+}
+#endif
+
 static int kmem_cache_open(struct kmem_cache *s, unsigned long flags)
 {
 	s->flags = kmem_cache_flags(s->size, flags, s->name, s->ctor);
@@ -4187,6 +4215,9 @@ void __init kmem_cache_init(void)
 	/* Now we can use the kmem_cache to allocate kmalloc slabs */
 	setup_kmalloc_cache_index_table();
 	create_kmalloc_caches(0);
+#ifdef CONFIG_AMLOGIC_SLAB_TRACE
+	slab_trace_init();
+#endif
 
 	/* Setup random freelists for each cache */
 	init_freelist_randomization();