Allow user to specify memory limit for dictionary training

facebook · Dec 14, 2021 · 5cd8292 · 5cd8292
1 parent 57383d2
commit 5cd8292
Show file tree

Hide file tree

Showing 6 changed files with 233 additions and 54 deletions.
diff --git a/programs/dibio.c b/programs/dibio.c
@@ -309,7 +309,7 @@ static fileStats DiB_fileStats(const char** fileNamesTable, int nbFiles, size_t
 int DiB_trainFromFiles(const char* dictFileName, size_t maxDictSize,
                        const char** fileNamesTable, int nbFiles, size_t chunkSize,
                        ZDICT_legacy_params_t* params, ZDICT_cover_params_t* coverParams,
-                       ZDICT_fastCover_params_t* fastCoverParams, int optimize)
+                       ZDICT_fastCover_params_t* fastCoverParams, int optimize, unsigned memLimit)
 {
     fileStats fs;
     size_t* sampleSizes; /* vector of sample sizes. Each sample can be up to SAMPLESIZE_MAX */
@@ -341,6 +341,11 @@ int DiB_trainFromFiles(const char* dictFileName, size_t maxDictSize,
         /* Limit the size of the training data to 2GB */
         /* TODO: there is opportunity to stop DiB_fileStats() early when the data limit is reached */
         loadedSize = (size_t)MIN( MIN((S64)maxMem, fs.totalSizeToLoad), MAX_SAMPLES_SIZE );
+        if (memLimit != 0) {
+            DISPLAYLEVEL(2, "!  Warning : setting manual memory limit for dictionary training data at %u MB \n",
+                (unsigned)(memLimit / (1 MB)));
+            loadedSize = (size_t)MIN(loadedSize, memLimit);
+        }
         srcBuffer = malloc(loadedSize+NOISELENGTH);
         sampleSizes = (size_t*)malloc(fs.nbSamples * sizeof(size_t));
     }

diff --git a/programs/dibio.h b/programs/dibio.h
@@ -34,6 +34,6 @@
 int DiB_trainFromFiles(const char* dictFileName, size_t maxDictSize,
                        const char** fileNamesTable, int nbFiles, size_t chunkSize,
                        ZDICT_legacy_params_t* params, ZDICT_cover_params_t* coverParams,
-                       ZDICT_fastCover_params_t* fastCoverParams, int optimize);
+                       ZDICT_fastCover_params_t* fastCoverParams, int optimize, unsigned memLimit);
 
 #endif