Add zstd dict trainer with ZDICT_optimizeTrainFromBuffer_cover (#806) (#806)

src/openzl/compress/cgraph.c

@@ -859,6 +859,21 @@ ZL_Report ZL_Compressor_overrideGraphParams(
     return ZL_returnSuccess();
 }
 
+ZL_Report ZL_Compressor_overrideNodeParams(
+        ZL_Compressor* compressor,
+        ZL_NodeID node,
+        const ZL_NodeParameters* np)
+{
+    ZL_RESULT_DECLARE_SCOPE(size_t, compressor);
+    ZL_ERR_IF_NULL(
+            NM_getCNode(&compressor->nmgr, node),
+            node_invalid,
+            "Node must be registered in compressor");
+
+    ZL_ERR_IF_ERR(NM_overrideNodeParams(&compressor->nmgr, node, np));
+    return ZL_returnSuccess();
+}
+
 ZL_Report ZL_Compressor_overrideBaseGraph(
         ZL_Compressor* compressor,
         ZL_GraphID graph,

src/openzl/compress/cgraph.h

@@ -132,6 +132,22 @@ ZL_Report ZL_Compressor_overrideBaseGraph(
         ZL_GraphID graph,
         ZL_GraphID newBaseGraph);
 
+/**
+ * Warning: This is part of experimental API for compressor mutation.
+ *
+ * Requires that:
+ * @p node is a parameterized node registered in @p compressor
+ *
+ * Replaces the parameters of @p node with @p np.
+ * @note: This function does not do any additional work typically associated
+ * with ZL_Compressor_registerParameterizedNode(), including dictionary
+ * unpacking.
+ */
+ZL_Report ZL_Compressor_overrideNodeParams(
+        ZL_Compressor* compressor,
+        ZL_NodeID node,
+        const ZL_NodeParameters* np);
+
 /**
  * Look up a previously loaded dict by its ZL_DictID.
  * @param matDesc must match the materializer used when the dict was loaded.

src/openzl/compress/cnodes.c

@@ -439,3 +439,29 @@ void CTM_setDictIndex(CNodes_manager* ctm, CNodeID id, uint32_t index)
     ZL_ASSERT_LT(id.cnid, VECTOR_SIZE(ctm->cnodes));
     VECTOR_AT(ctm->cnodes, id.cnid).maybeDictIndex = index;
 }
+
+ZL_Report CTM_overrideNodeParams(
+        CNodes_manager* ctm,
+        CNodeID id,
+        const ZL_NodeParameters* np)
+{
+    ZL_RESULT_DECLARE_SCOPE_REPORT(ctm->opCtx);
+    ZL_ASSERT_NN(ctm);
+    ZL_ASSERT_NN(np);
+    ZL_ASSERT_LT(id.cnid, VECTOR_SIZE(ctm->cnodes));
+
+    if (ZL_UniqueID_isValid(&np->mparam.mparamID.id)) {
+        ZL_ERR(GENERIC, "MParam override not supported");
+    }
+
+    CNode* const cnode             = &VECTOR_AT(ctm->cnodes, id.cnid);
+    ZL_MIEncoderDesc* const trDesc = &cnode->transformDesc.publicDesc;
+    if (np->localParams) {
+        trDesc->localParams = *np->localParams;
+        ZL_ERR_IF_ERR(CTM_transferLocalParams(ctm, &trDesc->localParams));
+    }
+    if (ZL_UniqueID_isValid(&np->dictID.id)) {
+        trDesc->dictID = np->dictID;
+    }
+    return ZL_returnSuccess();
+}

src/openzl/compress/cnodes.h

@@ -81,6 +81,11 @@ CTM_registerStandardTransform(
 /// the dict's position within the compressor's bundle.
 void CTM_setDictIndex(CNodes_manager* ctm, CNodeID id, uint32_t index);
 
+ZL_Report CTM_overrideNodeParams(
+        CNodes_manager* ctm,
+        CNodeID id,
+        const ZL_NodeParameters* np);
+
 /**
  * Rolls back the registration of @p id
  * @warning This only works when @p id was the last node registered. If local

src/openzl/compress/nodemgr.c

@@ -153,6 +153,42 @@ NM_parameterizeNode(Nodes_manager* nmgr, const ZL_ParameterizedNodeDesc* desc)
             ZL_NodeID, NM_NodeID_fromCNodeID(ZL_RES_value(cnodeidResult)));
 }
 
+ZL_Report NM_overrideNodeParams(
+        Nodes_manager* nmgr,
+        ZL_NodeID node,
+        const ZL_NodeParameters* np)
+{
+    ZL_RESULT_DECLARE_SCOPE_REPORT(nmgr->opCtx);
+    ZL_ASSERT_NN(nmgr);
+    ZL_ASSERT_NN(np);
+
+    ZL_ERR_IF(
+            NM_isStandardNode(node),
+            node_invalid,
+            "Cannot replace standard node");
+
+    const CNodeID cnodeID = NM_CNodeID_fromNodeID(node);
+    ZL_ERR_IF_GE(
+            cnodeID.cnid,
+            CTM_nbCNodes(&nmgr->ctm),
+            node_invalid,
+            "Node must be registered");
+    const CNode* const cnode = CTM_getCNode(&nmgr->ctm, cnodeID);
+    ZL_ASSERT_NN(cnode);
+    ZL_ERR_IF_EQ(
+            cnode->baseNodeID.nid,
+            ZL_NODE_ILLEGAL.nid,
+            node_invalid,
+            "Node must be parameterized");
+    ZL_ERR_IF_NE(cnode->nodetype, node_internalTransform, node_invalid);
+
+    if (np->name) {
+        ZL_ERR(parameter_invalid, "Cannot replace the name of a node");
+    }
+    ZL_ERR_IF_ERR(CTM_overrideNodeParams(&nmgr->ctm, cnodeID, np));
+    return ZL_returnSuccess();
+}
+
 const CNode* NM_getCNode(const Nodes_manager* nmgr, ZL_NodeID nodeid)
 {
     if (NM_isStandardNode(nodeid)) {

src/openzl/compress/nodemgr.h

@@ -46,6 +46,11 @@ NM_registerStandardTransform(
 ZL_RESULT_OF(ZL_NodeID)
 NM_parameterizeNode(Nodes_manager* nmgr, const ZL_ParameterizedNodeDesc* desc);
 
+ZL_Report NM_overrideNodeParams(
+        Nodes_manager* nmgr,
+        ZL_NodeID node,
+        const ZL_NodeParameters* np);
+
 // Read Accessors
 
 /*

tools/training/BUCK

@@ -17,6 +17,7 @@ zs_cxxlibrary(
         "../ml_selector:ml_selector_trainer",
         "ace:automated_compressor_explorer",
         "clustering:clustering_graph_trainer",
+        "dict:base_dict_trainer",
         "graph_mutation:graph_mutation",
     ],
     exported_deps = [

tools/training/CMakeLists.txt

@@ -22,6 +22,7 @@ if (OPENZL_BUILD_TRAINING_TOOLS)
   set_property(TARGET tools_training PROPERTY POSITION_INDEPENDENT_CODE ON)
 
   target_include_directories(tools_training PUBLIC ${PROJECT_SOURCE_DIR})
+  target_compile_definitions(tools_training PUBLIC ZDICT_STATIC_LINKING_ONLY)
   apply_openzl_compile_options_to_target(tools_training)
     target_link_libraries(
         tools_training

tools/training/dict/BUCK

@@ -0,0 +1,29 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+
+load("../../../defs.bzl", "zs_cxxlibrary")
+
+oncall("data_compression")
+
+zs_cxxlibrary(
+    name = "base_dict_trainer",
+    srcs = [
+        "base_dict_trainer.cpp",
+        "zstd_dict_trainer.cpp",
+    ],
+    headers = [
+        "base_dict_trainer.h",
+        "zstd_dict_trainer.h",
+    ],
+    deps = [
+        "../..:logger",
+        "../../../cpp:openzl_cpp",
+        "../sample_collection:sample_collection",
+    ],
+    exported_deps = [
+        "..:train_common",
+        "../utils:training_utils",
+    ],
+    external_deps = [
+        "zstd",
+    ],
+)

tools/training/dict/base_dict_trainer.cpp

@@ -0,0 +1,174 @@
+// Copyright (c) Meta Platforms, Inc. and affiliates.
+
+#include "tools/training/dict/base_dict_trainer.h"
+
+#include <vector>
+
+#include "openzl/compress/cgraph.h"
+#include "openzl/dict/dict.h"
+#include "openzl/dict/dict_constants.h"
+#include "openzl/zl_compressor.h"
+#include "openzl/zl_reflection.h"
+
+#include "tools/logger/Logger.h"
+#include "tools/training/dict/zstd_dict_trainer.h"
+#include "tools/training/sample_collection/training_sample_collector.h"
+
+using namespace openzl::tools::logger;
+
+namespace openzl::training {
+
+namespace {
+
+/// Build the list of all registered DictTrainer instances.
+std::vector<std::unique_ptr<DictTrainer>> createAllTrainers()
+{
+    std::vector<std::unique_ptr<DictTrainer>> trainers;
+    trainers.push_back(std::make_unique<ZstdDictTrainer>());
+    return trainers;
+}
+
+} // anonymous namespace
+
+TrainedCandidate trainDictsForCandidate(
+        const std::vector<MultiInput>& inputs,
+        Compressor& compressor,
+        const TrainParams& trainParams)
+{
+    (void)trainParams;
+
+    TrainedCandidate candidate;
+    candidate.serializedCompressor = compressor.serialize();
+
+    // Ask each registered trainer to find its dict-requiring nodes.
+    auto trainers = createAllTrainers();
+
+    struct TrainerWork {
+        DictTrainer* trainer;
+        DictNodeInfo node;
+    };
+    std::vector<TrainerWork> work;
+    std::vector<std::string> nodeNames;
+
+    for (auto& trainer : trainers) {
+        auto nodes = trainer->findDictNodes(compressor);
+        for (auto& node : nodes) {
+            nodeNames.push_back(node.nodeName);
+            work.push_back(
+                    TrainerWork{
+                            .trainer = trainer.get(),
+                            .node    = std::move(node),
+                    });
+        }
+    }
+
+    if (work.empty()) {
+        return candidate;
+    }
+
+    Logger::log_c(
+            VERBOSE1,
+            "Dict training: found %zu nodes requiring dictionaries",
+            work.size());
+
+    // Use introspection hooks to collect the actual data flowing into
+    // each dict-requiring codec node.
+    auto cctx           = refCCtxForTraining(compressor);
+    auto samplesPerNode = collectInputStreams(inputs, {}, nodeNames, cctx);
+
+    for (auto& [trainer, node] : work) {
+        auto it = samplesPerNode.find(node.nodeName);
+        if (it == samplesPerNode.end() || it->second.empty()) {
+            Logger::log_c(
+                    VERBOSE1,
+                    "No samples collected for node %s, skipping",
+                    node.nodeName.c_str());
+            continue;
+        }
+
+        const auto& nodeSamples = it->second;
+        Logger::log_c(
+                VERBOSE1,
+                "Training dict for node %u (codec %u, name %s) "
+                "with %zu samples",
+                node.nodeID.nid,
+                node.codecID,
+                node.nodeName.c_str(),
+                nodeSamples.size());
+
+        // trainDict returns packed dict content ready for Dict_pack.
+        ZL_LocalParams localParams = ZL_Compressor_Node_getLocalParams(
+                compressor.get(), node.nodeID);
+        auto dictContent =
+                trainer->trainDict(nodeSamples, compressor, localParams);
+        if (!dictContent.has_value()) {
+            Logger::log_c(
+                    VERBOSE1,
+                    "Trainer declined to train dict for node %s, skipping",
+                    node.nodeName.c_str());
+            continue;
+        }
+
+        // Pack into the generic ZL_Dict wire format.
+        std::string packedDict(ZL_DICT_HEADER_SIZE + dictContent->size(), '\0');
+        ZL_Report report = Dict_pack(
+                packedDict.data(),
+                packedDict.size(),
+                ZL_DICT_ID_NULL,
+                node.codecID,
+                false,
+                dictContent->data(),
+                dictContent->size());
+        if (ZL_isError(report)) {
+            throw Exception("Dict_pack failed");
+        }
+        packedDict.resize(ZL_validResult(report));
+
+        ZL_DictID generatedDictID =
+                Dict_extractID(packedDict.data(), packedDict.size());
+
+        ZL_NodeParameters nodeParams = {
+            .dictID = generatedDictID,
+        };
+        compressor.unwrap(ZL_Compressor_overrideNodeParams(
+                compressor.get(), node.nodeID, &nodeParams));
+
+        candidate.dicts.push_back(
+                TrainedCandidate::DictEntry{
+                        .dictID     = generatedDictID,
+                        .packedDict = std::move(packedDict),
+                });
+
+        Logger::log_c(
+                VERBOSE1,
+                "Trained dict: %zu bytes content",
+                dictContent->size());
+    }
+
+    if (candidate.dicts.empty()) {
+        return candidate;
+    }
+
+    std::vector<ZL_DictID> allDictIDs;
+    allDictIDs.reserve(candidate.dicts.size());
+    for (auto& dict : candidate.dicts) {
+        allDictIDs.push_back(dict.dictID);
+    }
+    // Compute bundleID
+    auto bundleID =
+            ZL_DictBundle_genBundleID(allDictIDs.data(), allDictIDs.size());
+
+    candidate.serializedCompressor = compressor.serialize();
+
+    // Set bundle ID on the candidate and patch dict_bundle_id in the CBOR.
+    candidate.replaceBundleID(bundleID);
+
+    Logger::log_c(
+            VERBOSE1,
+            "Dict training complete: %zu dicts, CBOR patched with bundleID",
+            candidate.dicts.size());
+
+    return candidate;
+}
+
+} // namespace openzl::training