*    get different tags.
  */
 static u8 assign_tag(struct kmem_cache *cache, const void *object,
-                       bool init, bool krealloc)
+                       bool init, bool keep_tag)
 {
-       /* Reuse the same tag for krealloc'ed objects. */
-       if (krealloc)
+       /*
+        * 1. When an object is kmalloc()'ed, two hooks are called:
+        *    kasan_slab_alloc() and kasan_kmalloc(). We assign the
+        *    tag only in the first one.
+        * 2. We reuse the same tag for krealloc'ed objects.
+        */
+       if (keep_tag)
                return get_tag(object);
 
        /*
        return (void *)object;
 }
 
-void * __must_check kasan_slab_alloc(struct kmem_cache *cache, void *object,
-                                       gfp_t flags)
-{
-       return kasan_kmalloc(cache, object, cache->object_size, flags);
-}
-
 static inline bool shadow_invalid(u8 tag, s8 shadow_byte)
 {
        if (IS_ENABLED(CONFIG_KASAN_GENERIC))
 }
 
 static void *__kasan_kmalloc(struct kmem_cache *cache, const void *object,
-                               size_t size, gfp_t flags, bool krealloc)
+                               size_t size, gfp_t flags, bool keep_tag)
 {
        unsigned long redzone_start;
        unsigned long redzone_end;
                                KASAN_SHADOW_SCALE_SIZE);
 
        if (IS_ENABLED(CONFIG_KASAN_SW_TAGS))
-               tag = assign_tag(cache, object, false, krealloc);
+               tag = assign_tag(cache, object, false, keep_tag);
 
        /* Tag is ignored in set_tag without CONFIG_KASAN_SW_TAGS */
        kasan_unpoison_shadow(set_tag(object, tag), size);
        return set_tag(object, tag);
 }
 
+void * __must_check kasan_slab_alloc(struct kmem_cache *cache, void *object,
+                                       gfp_t flags)
+{
+       return __kasan_kmalloc(cache, object, cache->object_size, flags, false);
+}
+
 void * __must_check kasan_kmalloc(struct kmem_cache *cache, const void *object,
                                size_t size, gfp_t flags)
 {
-       return __kasan_kmalloc(cache, object, size, flags, false);
+       return __kasan_kmalloc(cache, object, size, flags, true);
 }
 EXPORT_SYMBOL(kasan_kmalloc);