ZS_FULL
 };
 
+/*
+ * number of size_classes
+ */
+static int zs_size_classes;
+
 /*
  * We assign a page to ZS_ALMOST_EMPTY fullness group when:
  *     n <= N / f, where
 };
 
 struct zs_pool {
-       struct size_class *size_class[ZS_SIZE_CLASSES];
+       struct size_class **size_class;
 
        gfp_t flags;    /* allocation flags used when growing pool */
        atomic_long_t pages_allocated;
         */
        if (area->vm_buf)
                return 0;
-       area->vm_buf = (char *)__get_free_page(GFP_KERNEL);
+       area->vm_buf = kmalloc(ZS_MAX_ALLOC_SIZE, GFP_KERNEL);
        if (!area->vm_buf)
                return -ENOMEM;
        return 0;
 
 static inline void __zs_cpu_down(struct mapping_area *area)
 {
-       if (area->vm_buf)
-               free_page((unsigned long)area->vm_buf);
+       kfree(area->vm_buf);
        area->vm_buf = NULL;
 }
 
        return notifier_to_errno(ret);
 }
 
+static void init_zs_size_classes(void)
+{
+       int nr;
+
+       nr = (ZS_MAX_ALLOC_SIZE - ZS_MIN_ALLOC_SIZE) / ZS_SIZE_CLASS_DELTA + 1;
+       if ((ZS_MAX_ALLOC_SIZE - ZS_MIN_ALLOC_SIZE) % ZS_SIZE_CLASS_DELTA)
+               nr += 1;
+
+       zs_size_classes = nr;
+}
+
 static void __exit zs_exit(void)
 {
 #ifdef CONFIG_ZPOOL
                return ret;
        }
 
+       init_zs_size_classes();
+
 #ifdef CONFIG_ZPOOL
        zpool_register_driver(&zs_zpool_driver);
 #endif
        if (!pool)
                return NULL;
 
+       pool->size_class = kcalloc(zs_size_classes, sizeof(struct size_class *),
+                       GFP_KERNEL);
+       if (!pool->size_class) {
+               kfree(pool);
+               return NULL;
+       }
+
        /*
         * Iterate reversly, because, size of size_class that we want to use
         * for merging should be larger or equal to current size.
         */
-       for (i = ZS_SIZE_CLASSES - 1; i >= 0; i--) {
+       for (i = zs_size_classes - 1; i >= 0; i--) {
                int size;
                int pages_per_zspage;
                struct size_class *class;
 {
        int i;
 
-       for (i = 0; i < ZS_SIZE_CLASSES; i++) {
+       for (i = 0; i < zs_size_classes; i++) {
                int fg;
                struct size_class *class = pool->size_class[i];
 
                }
                kfree(class);
        }
+
+       kfree(pool->size_class);
        kfree(pool);
 }
 EXPORT_SYMBOL_GPL(zs_destroy_pool);