diff --git a/include/cudnn_frontend/graph_interface.h b/include/cudnn_frontend/graph_interface.h index 7cac72beb..65fae132a 100644 --- a/include/cudnn_frontend/graph_interface.h +++ b/include/cudnn_frontend/graph_interface.h @@ -1563,12 +1563,30 @@ class Graph : public ICudnn, public INode { return {error_code_t::OK, ""}; } + /** + * @brief Serialize graph execution-plan payload, optionally with graph structure. + * + * When @p serialize_structure is true, the payload includes structural graph metadata + * (context/nodes/tensors) in addition to plan fields. When false, the payload is + * plan-only and is suitable for handle-based plan reload paths that do not rebuild + * graph structure. Plan execution remains valid in plan-only mode. + * + * Recommendation: if the consumer does not call explicit structural deserialize + * (deserialize(const json&, ...)), prefer serialize_structure=false. + * + * @param data Output UBJSON payload. + * @param serialize_structure Include structural graph metadata when true. + * @return error_t OK on success, otherwise an error code describing the failure. + */ error_t - serialize(std::vector &data) const { + serialize(std::vector &data, bool serialize_structure = true) const { CUDNN_FE_LOG_BANNER(" SERIALIZE PLAN "); #ifndef CUDNN_FRONTEND_SKIP_JSON_LIB json j; - serialize(j); + // Optionally serialize the graph structure (nodes/tensors). + if (serialize_structure) { + serialize(j); + } auto const candidate = plans.candidate; auto execution_plan = plans.execution_plans[candidate]; @@ -1581,7 +1599,12 @@ class Graph : public ICudnn, public INode { j["behavior_notes"] = plans.behavior_notes; std::unordered_map tensor_to_pass_by_value; - CHECK_CUDNN_FRONTEND_ERROR(collect_pass_by_value_tensors_subtree(tensor_to_pass_by_value)); + // Pass-by-value data lives in the cached member restored on deserialize. + if (sub_nodes.empty()) { + tensor_to_pass_by_value = cached_pass_by_value; + } else { + CHECK_CUDNN_FRONTEND_ERROR(collect_pass_by_value_tensors_subtree(tensor_to_pass_by_value)); + } // Convert pass_by_values to JSON (unordered_map with numeric keys needs manual conversion) json pass_by_values_json = json::object(); @@ -1593,8 +1616,14 @@ class Graph : public ICudnn, public INode { j["pass_by_values"] = pass_by_values_json; std::unordered_map>> workspace_modifications; - int64_t workspace_offset = 0; - CHECK_CUDNN_FRONTEND_ERROR(collect_tensors_in_workspace_subtree(workspace_modifications, workspace_offset)); + + // Workspace modifications are cached when the graph is deserialized, walked otherwise. + if (sub_nodes.empty()) { + workspace_modifications = cached_workspace_modifications; + } else { + int64_t workspace_offset = 0; + CHECK_CUDNN_FRONTEND_ERROR(collect_tensors_in_workspace_subtree(workspace_modifications, workspace_offset)); + } // Convert workspace_modifications to JSON (nlohmann::json doesn't support std::tuple directly) json workspace_modifications_json = json::object(); @@ -1622,6 +1651,7 @@ class Graph : public ICudnn, public INode { return {error_code_t::OK, ""}; #else CUDNN_FRONTEND_UNUSED(data); + CUDNN_FRONTEND_UNUSED(serialize_structure); return {error_code_t::GRAPH_NOT_SUPPORTED, "unavailable when compiled with CUDNN_FRONTEND_SKIP_JSON_LIB"}; #endif } diff --git a/test/cpp/serialize.cpp b/test/cpp/serialize.cpp index 24b1e74b8..c7253de9c 100644 --- a/test/cpp/serialize.cpp +++ b/test/cpp/serialize.cpp @@ -678,3 +678,107 @@ TEST_CASE("Plan deserialize with run_warmup=false still prepares template", "[gr cudnnDestroy(handle); } + +// serialize(data, serialize_structure=false) omits the graph structure +// (nodes/tensors) while keeping the plan reloadable via deserialize(handle, ...). +// The default (true) still emits the structure, so existing callers are unaffected. +TEST_CASE("serialize_structure flag controls structural payload", "[graph][serialize][deserialize]") { + namespace fe = cudnn_frontend; + + fe::graph::Graph graph; + graph.set_io_data_type(fe::DataType_t::HALF) + .set_intermediate_data_type(fe::DataType_t::FLOAT) + .set_compute_data_type(fe::DataType_t::FLOAT); + + auto A = graph.tensor( + fe::graph::Tensor_attributes().set_name("A").set_dim({4, 16, 64}).set_stride({16 * 64, 64, 1}).set_uid(1)); + auto B = graph.tensor( + fe::graph::Tensor_attributes().set_name("B").set_dim({4, 64, 32}).set_stride({64 * 32, 32, 1}).set_uid(2)); + auto C = graph.matmul(A, B, fe::graph::Matmul_attributes().set_name("matmul")); + C->set_output(true).set_uid(3); + + cudnnHandle_t handle; + cudnnCreate(&handle); + + REQUIRE(graph.build(handle, {fe::HeurMode_t::A}).is_good()); + + std::vector with_structure, without_structure; + REQUIRE(graph.serialize(with_structure).is_good()); // default: serialize_structure=true + REQUIRE(graph.serialize(without_structure, /*serialize_structure=*/false).is_good()); + + json const j_with = json::from_ubjson(with_structure); + json const j_without = json::from_ubjson(without_structure); + + // Default emits the structure; opting out drops it and shrinks the blob. + REQUIRE(j_with.contains("nodes")); + REQUIRE_FALSE(j_without.contains("nodes")); + REQUIRE(without_structure.size() < with_structure.size()); + + // Both remain reloadable through the plan path with identical variant packs. + auto const expected_uids = graph.get_variant_pack_uids_sorted(); + for (auto const &blob : {with_structure, without_structure}) { + fe::graph::Graph reloaded; + REQUIRE(reloaded.deserialize(handle, blob, /*enforce_precompiled=*/false, /*run_warmup=*/false).is_good()); + REQUIRE(reloaded.get_variant_pack_uids_sorted() == expected_uids); + } + + cudnnDestroy(handle); +} + +// A graph loaded via deserialize(handle, ...) might have no node subtree, +// so serialize() must source its pass-by-value and workspace modifications from the +// cached members instead of walking the (now empty) subtree. +TEST_CASE("Plan re-serialize preserves pass-by-value and workspace modifications", "[graph][serialize][deserialize]") { + namespace fe = cudnn_frontend; + + fe::graph::Graph graph; + graph.set_io_data_type(fe::DataType_t::HALF) + .set_intermediate_data_type(fe::DataType_t::FLOAT) + .set_compute_data_type(fe::DataType_t::FLOAT); + + constexpr int64_t N = 4; + auto X = graph.tensor(fe::graph::Tensor_attributes() + .set_name("X") + .set_dim({N, N, N}) + .set_stride({N * N, N, 1}) + .set_data_type(fe::DataType_t::HALF) + .set_uid(1)); + auto scalar = graph.tensor(5.0f); + scalar->set_name("scalar").set_uid(2); + auto Y = graph.pointwise(X, + scalar, + fe::graph::Pointwise_attributes() + .set_name("add") + .set_mode(fe::PointwiseMode_t::ADD) + .set_compute_data_type(fe::DataType_t::FLOAT)); + Y->set_output(true).set_data_type(fe::DataType_t::HALF).set_uid(3); + + cudnnHandle_t handle; + cudnnCreate(&handle); + + REQUIRE(graph.build(handle, {fe::HeurMode_t::A}).is_good()); + + std::vector blob; + REQUIRE(graph.serialize(blob, /*serialize_structure=*/false).is_good()); + + // The fresh graph must actually carry pass-by-value data, else the test is moot. + json const j_fresh = json::from_ubjson(blob); + REQUIRE(j_fresh.contains("pass_by_values")); + REQUIRE_FALSE(j_fresh["pass_by_values"].empty()); + + auto const expected_uids = graph.get_variant_pack_uids_sorted(); + + fe::graph::Graph reloaded; + REQUIRE(reloaded.deserialize(handle, blob, /*enforce_precompiled=*/false, /*run_warmup=*/false).is_good()); + REQUIRE(reloaded.get_variant_pack_uids_sorted() == expected_uids); + + // Re-serialize the plan-only graph; the cached maps must survive intact. + std::vector next; + REQUIRE(reloaded.serialize(next, /*serialize_structure=*/false).is_good()); + + json const j_next = json::from_ubjson(next); + REQUIRE(j_next["pass_by_values"] == j_fresh["pass_by_values"]); + REQUIRE(j_next["workspace_modifications"] == j_fresh["workspace_modifications"]); + + cudnnDestroy(handle); +}