[PATCH 13/14] zram: add dictionary support to zstd backend

From: Sergey Senozhatsky
Date: Fri May 03 2024 - 05:22:17 EST


This adds support for pre-trained zstd dictionaries [1]
Dictionary is loaded once (per-config) and then loaded to Cctx
and Dctx by reference, so we don't allocate extra memory.

The patch is a little non-trivial, as it seems that noone
ever attempted to use dictionaries in the linux kernel
port of zstd.

It also uses GFP_KERNEL gfp in Cctx customAlloc(). We probably
would want to do something about it. Either make sure that we
always (somehow) fully setup all Cctx contexts from non-atomic
context before we attempt to use them, come up with some sort
of custom allocator or stop calling zcomp_compress() from atomic
context.

[1] https://github.com/facebook/zstd/blob/dev/programs/zstd.1.md#dictionary-builder

Signed-off-by: Sergey Senozhatsky <senozhatsky@xxxxxxxxxxxx>
---
drivers/block/zram/backend_zstd.c | 119 ++++++++++++++++++++++++------
1 file changed, 96 insertions(+), 23 deletions(-)

diff --git a/drivers/block/zram/backend_zstd.c b/drivers/block/zram/backend_zstd.c
index b2fb94902bef..6220c154e54e 100644
--- a/drivers/block/zram/backend_zstd.c
+++ b/drivers/block/zram/backend_zstd.c
@@ -12,23 +12,47 @@ struct zstd_ctx {
zstd_dctx *dctx;
void *cctx_mem;
void *dctx_mem;
+ ZSTD_customMem cctx_cmem;
+ ZSTD_customMem dctx_cmem;
+ ZSTD_CDict *cdict;
+ ZSTD_DDict *ddict;
s32 level;
};

+/*
+ * Cctx allocator.customAlloc() is called from zcom_compress(), which is
+ * called under local-lock (per-CPU compression stream), so we need to
+ * use GFP_ATOMIC here.
+ */
+static void *zstd_cctx_alloc(void *opaque, size_t size)
+{
+ return kvzalloc(size, GFP_ATOMIC);
+}
+
+static void *zstd_dctx_alloc(void *opaque, size_t size)
+{
+ return kvzalloc(size, GFP_KERNEL);
+}
+
+static void zstd_ctx_free(void *opaque, void *address)
+{
+ kvfree(address);
+}
+
static void zstd_destroy(void *ctx)
{
struct zstd_ctx *zctx = ctx;

vfree(zctx->cctx_mem);
vfree(zctx->dctx_mem);
+ ZSTD_freeCDict(zctx->cdict);
+ ZSTD_freeDDict(zctx->ddict);
kfree(zctx);
}

static void *zstd_create(struct zcomp_config *config)
{
- zstd_parameters params;
struct zstd_ctx *ctx;
- size_t sz;

ctx = kzalloc(sizeof(*ctx), GFP_KERNEL);
if (!ctx)
@@ -39,24 +63,64 @@ static void *zstd_create(struct zcomp_config *config)
else
ctx->level = ZSTD_defaultCLevel();

- params = zstd_get_params(ctx->level, PAGE_SIZE);
- sz = zstd_cctx_workspace_bound(&params.cParams);
- ctx->cctx_mem = vzalloc(sz);
- if (!ctx->cctx_mem)
- goto error;
-
- ctx->cctx = zstd_init_cctx(ctx->cctx_mem, sz);
- if (!ctx->cctx)
- goto error;
-
- sz = zstd_dctx_workspace_bound();
- ctx->dctx_mem = vzalloc(sz);
- if (!ctx->dctx_mem)
- goto error;
-
- ctx->dctx = zstd_init_dctx(ctx->dctx_mem, sz);
- if (!ctx->dctx)
- goto error;
+ ctx->cctx_cmem.customAlloc = zstd_cctx_alloc;
+ ctx->cctx_cmem.customFree = zstd_ctx_free;
+ ctx->dctx_cmem.customAlloc = zstd_dctx_alloc;
+ ctx->dctx_cmem.customFree = zstd_ctx_free;
+
+ if (config->dict_sz == 0) {
+ zstd_parameters params;
+ size_t sz;
+
+ params = zstd_get_params(ctx->level, PAGE_SIZE);
+ sz = zstd_cctx_workspace_bound(&params.cParams);
+ ctx->cctx_mem = vzalloc(sz);
+ if (!ctx->cctx_mem)
+ goto error;
+
+ ctx->cctx = zstd_init_cctx(ctx->cctx_mem, sz);
+ if (!ctx->cctx)
+ goto error;
+
+ sz = zstd_dctx_workspace_bound();
+ ctx->dctx_mem = vzalloc(sz);
+ if (!ctx->dctx_mem)
+ goto error;
+
+ ctx->dctx = zstd_init_dctx(ctx->dctx_mem, sz);
+ if (!ctx->dctx)
+ goto error;
+ } else {
+ ZSTD_compressionParameters params;
+
+ ctx->cctx = ZSTD_createCCtx_advanced(ctx->cctx_cmem);
+ if (!ctx->cctx)
+ goto error;
+
+ ctx->dctx = ZSTD_createDCtx_advanced(ctx->dctx_cmem);
+ if (!ctx->dctx)
+ goto error;
+
+ params = ZSTD_getCParams(ctx->level, PAGE_SIZE,
+ config->dict_sz);
+
+ ctx->cdict = ZSTD_createCDict_advanced(config->dict,
+ config->dict_sz,
+ ZSTD_dlm_byRef,
+ ZSTD_dct_auto,
+ params,
+ ctx->cctx_cmem);
+ if (!ctx->cdict)
+ goto error;
+
+ ctx->ddict = ZSTD_createDDict_advanced(config->dict,
+ config->dict_sz,
+ ZSTD_dlm_byRef,
+ ZSTD_dct_auto,
+ ctx->dctx_cmem);
+ if (!ctx->ddict)
+ goto error;
+ }

return ctx;

@@ -72,8 +136,12 @@ static int zstd_compress(void *ctx, const unsigned char *src,
const zstd_parameters params = zstd_get_params(zctx->level, PAGE_SIZE);
size_t ret;

- ret = zstd_compress_cctx(zctx->cctx, dst, *dst_len,
- src, PAGE_SIZE, &params);
+ if (!zctx->cdict)
+ ret = zstd_compress_cctx(zctx->cctx, dst, *dst_len,
+ src, PAGE_SIZE, &params);
+ else
+ ret = ZSTD_compress_usingCDict(zctx->cctx, dst, *dst_len,
+ src, PAGE_SIZE, zctx->cdict);
if (zstd_is_error(ret))
return -EINVAL;
*dst_len = ret;
@@ -86,7 +154,12 @@ static int zstd_decompress(void *ctx, const unsigned char *src, size_t src_len,
struct zstd_ctx *zctx = ctx;
size_t ret;

- ret = zstd_decompress_dctx(zctx->dctx, dst, PAGE_SIZE, src, src_len);
+ if (!zctx->ddict)
+ ret = zstd_decompress_dctx(zctx->dctx, dst, PAGE_SIZE,
+ src, src_len);
+ else
+ ret = ZSTD_decompress_usingDDict(zctx->dctx, dst, PAGE_SIZE,
+ src, src_len, zctx->ddict);
if (zstd_is_error(ret))
return -EINVAL;
return 0;
--
2.45.0.rc1.225.g2a3ae87e7f-goog