From c2528f467133ac54c086cad31a4ec67689f374f0 Mon Sep 17 00:00:00 2001 From: Yang Xu Date: Thu, 28 May 2026 19:46:59 -0700 Subject: [PATCH 1/4] feat(python): add get_engine_and_knobs_at_index for structured plan pinning get_plan_name_at_index returns a formatted "engN_kT=V" tag built from the engine global index and knob choices. Callers that want to persist a tuned plan and replay it later are forced to either store the bare plan index (which drifts when the policy=ALL plan list is re-enumerated across cudnn-frontend / backend versions) or parse the tag string. Expose the structured data directly: get_engine_and_knobs_at_index returns (engine_id, {KnobType_t: value}), reading the same backend attributes get_engine_tag stringifies. The result feeds straight into create_execution_plan(engine_id, knobs) to rebuild the exact same kernel on a fresh graph without a heuristics query. - detail::get_engine_id_and_knobs (cudnn_frontend_utils.h): structured reader - Execution_plan_list::get_engine_and_knobs_at_index (plans.h) - Graph::get_engine_and_knobs_at_index (graph_interface.h) - PyGraph binding (pygraph.h/.cpp) Co-Authored-By: Claude Opus 4.8 (1M context) --- include/cudnn_frontend/graph_interface.h | 11 ++++ include/cudnn_frontend/plans.h | 26 ++++++++ include/cudnn_frontend_utils.h | 81 ++++++++++++++++++++++++ python/pygraph/pygraph.cpp | 24 +++++++ python/pygraph/pygraph.h | 3 + 5 files changed, 145 insertions(+) diff --git a/include/cudnn_frontend/graph_interface.h b/include/cudnn_frontend/graph_interface.h index 32efdf6d8..44d6a43fc 100644 --- a/include/cudnn_frontend/graph_interface.h +++ b/include/cudnn_frontend/graph_interface.h @@ -1007,6 +1007,17 @@ class Graph : public ICudnn, public INode { return ret_val; } + // Structured counterpart of get_plan_name_at_index(): returns the engine + // global index and (KnobType_t, value) choices for the plan, suitable for + // pinning the same kernel later via create_execution_plan() without a + // heuristics query. + error_t + get_engine_and_knobs_at_index(int64_t plan_index, + int64_t &engine_id, + std::unordered_map &knobs) const { + return plans.get_engine_and_knobs_at_index(plan_index, engine_id, knobs); + } + error_t get_workspace_size(int64_t &cudnn_workspace_size) const { return get_workspace_size_plan_at_index(plans.candidate, cudnn_workspace_size); diff --git a/include/cudnn_frontend/plans.h b/include/cudnn_frontend/plans.h index 64f3f05f1..d48ba6e43 100644 --- a/include/cudnn_frontend/plans.h +++ b/include/cudnn_frontend/plans.h @@ -2,8 +2,11 @@ #include #include +#include +#include #include +#include "knobs.h" #include "../cudnn_frontend_EngineConfig.h" #include "../cudnn_frontend_Logging.h" #include "graph_helpers.h" @@ -477,6 +480,29 @@ class Execution_plan_list { return {error_code_t::OK, ""}; } + // Structured counterpart of get_name_at_index(): returns the engine global + // index and the (KnobType_t, value) choices for the plan at `index`. + // The returned pair can be fed straight back into + // Graph::create_execution_plan(engine_id, knobs) to rebuild the exact same + // kernel without a heuristics query -- and, unlike a positional plan index, + // it survives plan-list re-enumeration across cudnn-frontend versions. + error_t + get_engine_and_knobs_at_index(int64_t index, + int64_t& engine_id, + std::unordered_map& knobs) const { + std::vector> backend_knobs; + auto status = detail::get_engine_id_and_knobs(engine_configs[index], engine_id, backend_knobs); + RETURN_CUDNN_FRONTEND_ERROR_IF( + status != CUDNN_STATUS_SUCCESS, + error_code_t::CUDNN_BACKEND_API_FAILED, + "Failed to query engine id / knob choices for plan at index " + std::to_string(index)); + knobs.clear(); + for (auto const& backend_knob : backend_knobs) { + knobs[convert_from_backend_knob_type(backend_knob.first)] = backend_knob.second; + } + return {error_code_t::OK, ""}; + } + error_t check_support_at_index(int64_t index) { // Ignore if the engine config was deselected. diff --git a/include/cudnn_frontend_utils.h b/include/cudnn_frontend_utils.h index 33e7c3a5e..d1245c65a 100644 --- a/include/cudnn_frontend_utils.h +++ b/include/cudnn_frontend_utils.h @@ -26,6 +26,7 @@ #include #include #include +#include #include #include #include @@ -2611,6 +2612,86 @@ std::string static get_engine_tag(ManagedOpaqueDescriptor const config) { return tag.str(); } +// Structured counterpart of get_engine_tag(): returns the engine's global +// index and its (knob type, value) choices instead of a formatted string. +// Reads the exact same backend attributes get_engine_tag() stringifies, so a +// caller can pin a specific plan via Graph::create_execution_plan() (which +// takes an engine id + knob choices) on a freshly built graph -- skipping the +// heuristics query entirely -- instead of parsing the tag string or relying +// on a positional plan index that can drift across library versions. +cudnnStatus_t static get_engine_id_and_knobs(ManagedOpaqueDescriptor const config, + int64_t& engineId, + std::vector>& knobs) { + engineId = -1; + knobs.clear(); + + ManagedOpaqueDescriptor extractedEngine = make_shared_backend_pointer(CUDNN_BACKEND_ENGINE_DESCRIPTOR); + if (extractedEngine->get_status() != CUDNN_STATUS_SUCCESS) { + return extractedEngine->get_status(); + } + cudnnBackendDescriptor_t extractedEngine_ = extractedEngine->get_backend_descriptor(); + + int64_t elemCount = 0; + auto status = detail::get_attribute(config->get_backend_descriptor(), + CUDNN_ATTR_ENGINECFG_ENGINE, + CUDNN_TYPE_BACKEND_DESCRIPTOR, + 1, + &elemCount, + &extractedEngine_); + if (status != CUDNN_STATUS_SUCCESS) { + return status; + } + + std::array extractedKnobs{{nullptr}}; + for (auto& knob : extractedKnobs) { + knob = make_shared_backend_pointer(CUDNN_BACKEND_KNOB_CHOICE_DESCRIPTOR); + if (knob->get_status() != CUDNN_STATUS_SUCCESS) { + return knob->get_status(); + } + } + std::array extractedKnobs_{{nullptr}}; + for (std::uint32_t i = 0; i < extractedKnobs.size(); i++) { + extractedKnobs_[i] = extractedKnobs[i]->get_backend_descriptor(); + } + + status = detail::get_attribute( + extractedEngine_, CUDNN_ATTR_ENGINE_GLOBAL_INDEX, CUDNN_TYPE_INT64, 1, &elemCount, &engineId); + if (status != CUDNN_STATUS_SUCCESS) { + return status; + } + + int64_t numKnobs = 0; + status = detail::get_attribute(config->get_backend_descriptor(), + CUDNN_ATTR_ENGINECFG_KNOB_CHOICES, + CUDNN_TYPE_BACKEND_DESCRIPTOR, + CUDNN_KNOB_TYPE_COUNTS, + &numKnobs, + &(extractedKnobs_[0])); + if (status != CUDNN_STATUS_SUCCESS) { + return status; + } + if (numKnobs > CUDNN_KNOB_TYPE_COUNTS) { + return CUDNN_STATUS_NOT_SUPPORTED; + } + + knobs.reserve(static_cast(numKnobs)); + for (size_t idx = 0; idx < static_cast(numKnobs); ++idx) { + const cudnnBackendDescriptor_t& knob = extractedKnobs_[idx]; + cudnnBackendKnobType_t type = CUDNN_KNOB_TYPE_COUNTS; + int64_t choice = -2; + status = detail::get_attribute(knob, CUDNN_ATTR_KNOB_CHOICE_KNOB_TYPE, CUDNN_TYPE_KNOB_TYPE, 1, nullptr, &type); + if (status != CUDNN_STATUS_SUCCESS) { + return status; + } + status = detail::get_attribute(knob, CUDNN_ATTR_KNOB_CHOICE_KNOB_VALUE, CUDNN_TYPE_INT64, 1, nullptr, &choice); + if (status != CUDNN_STATUS_SUCCESS) { + return status; + } + knobs.emplace_back(type, choice); + } + return CUDNN_STATUS_SUCCESS; +} + } // namespace detail } // namespace cudnn_frontend diff --git a/python/pygraph/pygraph.cpp b/python/pygraph/pygraph.cpp index eb09ec91a..c5843a59d 100644 --- a/python/pygraph/pygraph.cpp +++ b/python/pygraph/pygraph.cpp @@ -807,6 +807,15 @@ PyGraph::get_plan_name_at_index(int64_t index) { return plan_name; } +std::pair> +PyGraph::get_engine_and_knobs_at_index(int64_t index) { + int64_t engine_id = -1; + std::unordered_map knobs; + auto status = graph->get_engine_and_knobs_at_index(index, engine_id, knobs); + throw_if(status.is_bad(), status.get_code(), status.get_message()); + return {engine_id, knobs}; +} + std::vector default_vector(void) { return {}; @@ -1325,6 +1334,21 @@ init_pygraph_submodule(py::module_& m) { Args: index (int): The index of the plan to get workspace from. )pbdoc") + .def("get_engine_and_knobs_at_index", + &PyGraph::get_engine_and_knobs_at_index, + py::arg("index"), + R"pbdoc( + Get the engine id and knob choices for the plan at the given index. + + This is the structured counterpart of get_plan_name_at_index: the + returned (engine_id, {knob_type: value}) can be passed back to + create_execution_plan() to rebuild the exact same kernel without a + heuristics query, and survives plan re-enumeration across versions. + Args: + index (int): The index of the plan to query. + Returns: + tuple[int, dict[knob_type, int]]: engine global index and knob choices. + )pbdoc") .def("_execute", &PyGraph::execute, py::arg("var_pack"), diff --git a/python/pygraph/pygraph.h b/python/pygraph/pygraph.h index 83974aee6..1899f7ace 100644 --- a/python/pygraph/pygraph.h +++ b/python/pygraph/pygraph.h @@ -798,6 +798,9 @@ class PyGraph { std::string get_plan_name_at_index(int64_t index); + std::pair> + get_engine_and_knobs_at_index(int64_t index); + private: // Internal SDPA implementation - delegates to sdpa() or sdpa_fp8() based on mma_core_mode // return SDPA_outputs struct: {O, Stats, RNG_DUMP, Amax_S, Amax_O} From e86eb985e24a81e38ca7007140fc9553d7abc03f Mon Sep 17 00:00:00 2001 From: Yang Xu Date: Thu, 28 May 2026 22:14:34 -0700 Subject: [PATCH 2/4] address review: bounds-check index, add cpp unit test, trim comments - get_engine_and_knobs_at_index: reject out-of-range index (mirrors check_support_at_index) instead of indexing engine_configs OOB. - add test/cpp/get_engine_and_knobs.cpp: enumerate a matmul graph's plans, read (engine_id, knobs) for each, and confirm re-pinning via create_execution_plan reproduces the same plan (matching name); also checks out-of-range indices error. - trim the new doc comments to match neighboring style. Co-Authored-By: Claude Opus 4.8 (1M context) --- include/cudnn_frontend/graph_interface.h | 6 +- include/cudnn_frontend/plans.h | 11 ++- include/cudnn_frontend_utils.h | 9 +-- python/pygraph/pygraph.cpp | 11 +-- test/cpp/CMakeLists.txt | 1 + test/cpp/get_engine_and_knobs.cpp | 95 ++++++++++++++++++++++++ 6 files changed, 108 insertions(+), 25 deletions(-) create mode 100644 test/cpp/get_engine_and_knobs.cpp diff --git a/include/cudnn_frontend/graph_interface.h b/include/cudnn_frontend/graph_interface.h index 44d6a43fc..33dae23f2 100644 --- a/include/cudnn_frontend/graph_interface.h +++ b/include/cudnn_frontend/graph_interface.h @@ -1007,10 +1007,8 @@ class Graph : public ICudnn, public INode { return ret_val; } - // Structured counterpart of get_plan_name_at_index(): returns the engine - // global index and (KnobType_t, value) choices for the plan, suitable for - // pinning the same kernel later via create_execution_plan() without a - // heuristics query. + // Structured counterpart of get_plan_name_at_index(): engine index + knob + // choices, for replay via create_execution_plan(). error_t get_engine_and_knobs_at_index(int64_t plan_index, int64_t &engine_id, diff --git a/include/cudnn_frontend/plans.h b/include/cudnn_frontend/plans.h index d48ba6e43..de5192902 100644 --- a/include/cudnn_frontend/plans.h +++ b/include/cudnn_frontend/plans.h @@ -480,16 +480,15 @@ class Execution_plan_list { return {error_code_t::OK, ""}; } - // Structured counterpart of get_name_at_index(): returns the engine global - // index and the (KnobType_t, value) choices for the plan at `index`. - // The returned pair can be fed straight back into - // Graph::create_execution_plan(engine_id, knobs) to rebuild the exact same - // kernel without a heuristics query -- and, unlike a positional plan index, - // it survives plan-list re-enumeration across cudnn-frontend versions. + // Structured counterpart of get_name_at_index(): engine index + knob choices, + // suitable for replay via Graph::create_execution_plan(engine_id, knobs). error_t get_engine_and_knobs_at_index(int64_t index, int64_t& engine_id, std::unordered_map& knobs) const { + RETURN_CUDNN_FRONTEND_ERROR_IF(index < 0 || index >= static_cast(engine_configs.size()), + error_code_t::GRAPH_EXECUTION_FAILED, + "Plan index " + std::to_string(index) + " is invalid."); std::vector> backend_knobs; auto status = detail::get_engine_id_and_knobs(engine_configs[index], engine_id, backend_knobs); RETURN_CUDNN_FRONTEND_ERROR_IF( diff --git a/include/cudnn_frontend_utils.h b/include/cudnn_frontend_utils.h index d1245c65a..0e7f9a873 100644 --- a/include/cudnn_frontend_utils.h +++ b/include/cudnn_frontend_utils.h @@ -2612,13 +2612,8 @@ std::string static get_engine_tag(ManagedOpaqueDescriptor const config) { return tag.str(); } -// Structured counterpart of get_engine_tag(): returns the engine's global -// index and its (knob type, value) choices instead of a formatted string. -// Reads the exact same backend attributes get_engine_tag() stringifies, so a -// caller can pin a specific plan via Graph::create_execution_plan() (which -// takes an engine id + knob choices) on a freshly built graph -- skipping the -// heuristics query entirely -- instead of parsing the tag string or relying -// on a positional plan index that can drift across library versions. +// Structured counterpart of get_engine_tag(): engine global index + (knob type, +// value) choices instead of a formatted string. Reads the same backend attributes. cudnnStatus_t static get_engine_id_and_knobs(ManagedOpaqueDescriptor const config, int64_t& engineId, std::vector>& knobs) { diff --git a/python/pygraph/pygraph.cpp b/python/pygraph/pygraph.cpp index c5843a59d..c66a8c3dc 100644 --- a/python/pygraph/pygraph.cpp +++ b/python/pygraph/pygraph.cpp @@ -1338,16 +1338,11 @@ init_pygraph_submodule(py::module_& m) { &PyGraph::get_engine_and_knobs_at_index, py::arg("index"), R"pbdoc( - Get the engine id and knob choices for the plan at the given index. - - This is the structured counterpart of get_plan_name_at_index: the - returned (engine_id, {knob_type: value}) can be passed back to - create_execution_plan() to rebuild the exact same kernel without a - heuristics query, and survives plan re-enumeration across versions. + Get the engine id and knob choices for a plan at the given index. + Structured counterpart of get_plan_name_at_index; the result can be + passed to create_execution_plan() to rebuild the same kernel. Args: index (int): The index of the plan to query. - Returns: - tuple[int, dict[knob_type, int]]: engine global index and knob choices. )pbdoc") .def("_execute", &PyGraph::execute, diff --git a/test/cpp/CMakeLists.txt b/test/cpp/CMakeLists.txt index 27502948e..fc25ff4b4 100644 --- a/test/cpp/CMakeLists.txt +++ b/test/cpp/CMakeLists.txt @@ -25,6 +25,7 @@ add_executable( validate.cpp version.cpp tensor.cpp + get_engine_and_knobs.cpp ) if (MSVC) diff --git a/test/cpp/get_engine_and_knobs.cpp b/test/cpp/get_engine_and_knobs.cpp new file mode 100644 index 000000000..b5e577b1f --- /dev/null +++ b/test/cpp/get_engine_and_knobs.cpp @@ -0,0 +1,95 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a + * copy of this software and associated documentation files (the "Software"), + * to deal in the Software without restriction, including without limitation + * the rights to use, copy, modify, merge, publish, distribute, sublicense, + * and/or sell copies of the Software, and to permit persons to whom the + * Software is furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL + * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING + * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER + * DEALINGS IN THE SOFTWARE. + */ + +#include + +#include + +namespace { +namespace fe = cudnn_frontend; + +// Build a simple b x m x k @ b x k x n matmul graph (shared handle). +std::shared_ptr +make_matmul_graph(cudnnHandle_t handle) { + auto graph = std::make_shared(); + 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})); + auto B = + graph->tensor(fe::graph::Tensor_attributes().set_name("B").set_dim({4, 64, 32}).set_stride({64 * 32, 32, 1})); + auto C = graph->matmul(A, B, fe::graph::Matmul_attributes().set_name("matmul")); + C->set_output(true); + + REQUIRE(graph->validate().is_good()); + REQUIRE(graph->build_operation_graph(handle).is_good()); + return graph; +} +} // namespace + +TEST_CASE("get_engine_and_knobs_at_index round-trips via create_execution_plan", "[graph][knobs]") { + cudnnHandle_t handle; + cudnnCreate(&handle); + + auto graph = make_matmul_graph(handle); + REQUIRE(graph->create_execution_plans({fe::HeurMode_t::A}).is_good()); + REQUIRE(graph->check_support().is_good()); + REQUIRE(graph->build_plans(fe::BuildPlanPolicy_t::ALL).is_good()); + + auto const count = graph->get_execution_plan_count(); + REQUIRE(count > 0); + + // Out-of-range indices must error rather than read OOB. + int64_t engine_id = -1; + std::unordered_map knobs; + REQUIRE(graph->get_engine_and_knobs_at_index(-1, engine_id, knobs).is_bad()); + REQUIRE(graph->get_engine_and_knobs_at_index(count, engine_id, knobs).is_bad()); + + // The getter must succeed for every plan; pinning it back via + // create_execution_plan must reproduce the *same* kernel (matching name). + // Not every heuristic-enumerated engine is standalone-constructable, so the + // re-pin is best-effort -- but at least one must round-trip. + int64_t round_tripped = 0; + for (int64_t i = 0; i < count; i++) { + std::string name; + REQUIRE(graph->get_plan_name_at_index(i, name).is_good()); + + engine_id = -1; + knobs.clear(); + REQUIRE(graph->get_engine_and_knobs_at_index(i, engine_id, knobs).is_good()); + REQUIRE(engine_id >= 0); + + auto pinned = make_matmul_graph(handle); + if (pinned->create_execution_plan(engine_id, knobs).is_good() && pinned->build_plans().is_good()) { + REQUIRE(pinned->get_execution_plan_count() == 1); + std::string pinned_name; + REQUIRE(pinned->get_plan_name_at_index(0, pinned_name).is_good()); + REQUIRE(pinned_name == name); + round_tripped++; + } + } + REQUIRE(round_tripped > 0); + + cudnnDestroy(handle); +} From c5d04e646156f4a4738e8aae497fff5b70627de2 Mon Sep 17 00:00:00 2001 From: Yang Xu Date: Thu, 28 May 2026 22:36:21 -0700 Subject: [PATCH 3/4] knobs: add SWAP_AB / INPUT_TMA_ENABLE / OUTPUT_TMA_ENABLE to KnobType_t KnobType_t (and the to/from backend converters) stopped at WARP_SPEC_CFG (42), so engines using SWAP_AB (43, cuDNN 9.18), INPUT_TMA_ENABLE (44) or OUTPUT_TMA_ENABLE (45, cuDNN 9.22) had those knobs mapped to NOT_SET by convert_from_backend_knob_type. Feeding NOT_SET back into create_execution_plan then failed convert_to_backend_knob_type with INVALID_VALUE -- so a plan enumerated with one of these knobs (e.g. via get_engine_and_knobs_at_index) could not be pinned. Add the three knob types to the enum, both converters (version-gated to match the backend @since), and the pybind knob_type enum. The cpp test now compares the structured identity (engine id + knob map) instead of the plan-name tag, since the tag serializes knobs in engine-config order, which differs between the heuristic config and the pinned one even though the kernel is identical. create_execution_plan is now asserted to succeed for every enumerated plan; building it stays best-effort (can fail for unrelated environment reasons such as a ptxas older than the engine's target). Co-Authored-By: Claude Opus 4.8 (1M context) --- include/cudnn_frontend/knobs.h | 26 ++++++++++++++++++++++++++ python/properties.cpp | 5 ++++- test/cpp/get_engine_and_knobs.cpp | 28 +++++++++++++++++----------- 3 files changed, 47 insertions(+), 12 deletions(-) diff --git a/include/cudnn_frontend/knobs.h b/include/cudnn_frontend/knobs.h index ee4ee4c43..559643ac4 100644 --- a/include/cudnn_frontend/knobs.h +++ b/include/cudnn_frontend/knobs.h @@ -33,6 +33,9 @@ enum class KnobType_t { TILE_M, TILE_N, WARP_SPEC_CFG, + SWAP_AB, + INPUT_TMA_ENABLE, + OUTPUT_TMA_ENABLE, }; class Knob { @@ -143,6 +146,19 @@ convert_to_backend_knob_type(KnobType_t const knob_type, cudnnBackendKnobType_t& cudnn_knob_type = CUDNN_KNOB_TYPE_WARP_SPEC_CFG; return cudnnStatus_t::CUDNN_STATUS_SUCCESS; #endif +#if (CUDNN_VERSION >= 91800) + case KnobType_t::SWAP_AB: + cudnn_knob_type = CUDNN_KNOB_TYPE_SWAP_AB; + return cudnnStatus_t::CUDNN_STATUS_SUCCESS; +#endif +#if (CUDNN_VERSION >= 92200) + case KnobType_t::INPUT_TMA_ENABLE: + cudnn_knob_type = CUDNN_KNOB_TYPE_INPUT_TMA_ENABLE; + return cudnnStatus_t::CUDNN_STATUS_SUCCESS; + case KnobType_t::OUTPUT_TMA_ENABLE: + cudnn_knob_type = CUDNN_KNOB_TYPE_OUTPUT_TMA_ENABLE; + return cudnnStatus_t::CUDNN_STATUS_SUCCESS; +#endif #ifndef NO_DEFAULT_IN_SWITCH default: return cudnnStatus_t::CUDNN_STATUS_INVALID_VALUE; @@ -219,6 +235,16 @@ convert_from_backend_knob_type(cudnnBackendKnobType_t cudnn_knob_type) { return KnobType_t::TILE_N; case CUDNN_KNOB_TYPE_WARP_SPEC_CFG: return KnobType_t::WARP_SPEC_CFG; +#endif +#if (CUDNN_VERSION >= 91800) + case CUDNN_KNOB_TYPE_SWAP_AB: + return KnobType_t::SWAP_AB; +#endif +#if (CUDNN_VERSION >= 92200) + case CUDNN_KNOB_TYPE_INPUT_TMA_ENABLE: + return KnobType_t::INPUT_TMA_ENABLE; + case CUDNN_KNOB_TYPE_OUTPUT_TMA_ENABLE: + return KnobType_t::OUTPUT_TMA_ENABLE; #endif default: return KnobType_t::NOT_SET; diff --git a/python/properties.cpp b/python/properties.cpp index c3f6c84f2..0aad3b73a 100644 --- a/python/properties.cpp +++ b/python/properties.cpp @@ -231,7 +231,10 @@ init_properties(py::module_& m) { .value("SPLIT_P_SLC", cudnn_frontend::KnobType_t::SPLIT_P_SLC) .value("TILE_M", cudnn_frontend::KnobType_t::TILE_M) .value("TILE_N", cudnn_frontend::KnobType_t::TILE_N) - .value("WARP_SPEC_CFG", cudnn_frontend::KnobType_t::WARP_SPEC_CFG); + .value("WARP_SPEC_CFG", cudnn_frontend::KnobType_t::WARP_SPEC_CFG) + .value("SWAP_AB", cudnn_frontend::KnobType_t::SWAP_AB) + .value("INPUT_TMA_ENABLE", cudnn_frontend::KnobType_t::INPUT_TMA_ENABLE) + .value("OUTPUT_TMA_ENABLE", cudnn_frontend::KnobType_t::OUTPUT_TMA_ENABLE); py::class_>(m, "knob") .def(py::init(), diff --git a/test/cpp/get_engine_and_knobs.cpp b/test/cpp/get_engine_and_knobs.cpp index b5e577b1f..258159b83 100644 --- a/test/cpp/get_engine_and_knobs.cpp +++ b/test/cpp/get_engine_and_knobs.cpp @@ -66,26 +66,32 @@ TEST_CASE("get_engine_and_knobs_at_index round-trips via create_execution_plan", REQUIRE(graph->get_engine_and_knobs_at_index(-1, engine_id, knobs).is_bad()); REQUIRE(graph->get_engine_and_knobs_at_index(count, engine_id, knobs).is_bad()); - // The getter must succeed for every plan; pinning it back via - // create_execution_plan must reproduce the *same* kernel (matching name). - // Not every heuristic-enumerated engine is standalone-constructable, so the - // re-pin is best-effort -- but at least one must round-trip. + // For every plan: the reported (engine, knobs) must feed straight back into + // create_execution_plan -- i.e. every backend knob type the engine uses is + // representable as a KnobType_t. Building the pinned plan is best-effort + // (it can fail for environment reasons, e.g. a ptxas older than the + // engine's target); when it succeeds it must reproduce the same plan. int64_t round_tripped = 0; for (int64_t i = 0; i < count; i++) { - std::string name; - REQUIRE(graph->get_plan_name_at_index(i, name).is_good()); - engine_id = -1; knobs.clear(); REQUIRE(graph->get_engine_and_knobs_at_index(i, engine_id, knobs).is_good()); REQUIRE(engine_id >= 0); auto pinned = make_matmul_graph(handle); - if (pinned->create_execution_plan(engine_id, knobs).is_good() && pinned->build_plans().is_good()) { + REQUIRE(pinned->create_execution_plan(engine_id, knobs).is_good()); + + if (pinned->build_plans().is_good()) { REQUIRE(pinned->get_execution_plan_count() == 1); - std::string pinned_name; - REQUIRE(pinned->get_plan_name_at_index(0, pinned_name).is_good()); - REQUIRE(pinned_name == name); + // Compare the structured identity (engine + knob map), not the tag + // string: the tag serializes knobs in engine-config order, which + // differs between the heuristic and pinned configs, but the engine + // id + knob values are what define the kernel. + int64_t pinned_engine = -1; + std::unordered_map pinned_knobs; + REQUIRE(pinned->get_engine_and_knobs_at_index(0, pinned_engine, pinned_knobs).is_good()); + REQUIRE(pinned_engine == engine_id); + REQUIRE(pinned_knobs == knobs); round_tripped++; } } From bda90e5a9e83e67c3031c5b5220d1932a219dc2a Mon Sep 17 00:00:00 2001 From: Yang Xu Date: Fri, 29 May 2026 00:32:05 -0700 Subject: [PATCH 4/4] make get_engine_tag deterministic: sort knob choices by type The plan-name tag was built by iterating CUDNN_ATTR_ENGINECFG_KNOB_CHOICES in stored order, which differs between the heuristics path and create_execution_plan (set_knob_choices iterates a std::unordered_map). So the same engine + knob values could serialize to differently-ordered tags (e.g. eng11_k2=29_k27=0...k43=0 vs eng11_k43=0_k38=0...k2=29) -- the kernel is identical but the string isn't a stable id. Sort the knob choices by type before formatting so the tag is a deterministic function of the engine config regardless of how it was built. This is off the execution hot path (tag is used for logging / plan identity), so no perf impact; the actual knob choices passed to the backend are unchanged. The cpp test now also asserts the pinned plan's tag matches the original's. Co-Authored-By: Claude Opus 4.8 (1M context) --- include/cudnn_frontend_utils.h | 11 +++++++++++ test/cpp/get_engine_and_knobs.cpp | 17 +++++++++++++---- 2 files changed, 24 insertions(+), 4 deletions(-) diff --git a/include/cudnn_frontend_utils.h b/include/cudnn_frontend_utils.h index 0e7f9a873..9cc8aa0d2 100644 --- a/include/cudnn_frontend_utils.h +++ b/include/cudnn_frontend_utils.h @@ -24,6 +24,7 @@ #include #include #include +#include #include #include #include @@ -2595,6 +2596,8 @@ std::string static get_engine_tag(ManagedOpaqueDescriptor const config) { return "INVALID_ENGINE_NAME_KNOB_COUNT"; } + std::vector> knob_choices; + knob_choices.reserve(static_cast(numKnobs)); for (size_t idx = 0; idx < static_cast(numKnobs); ++idx) { const cudnnBackendDescriptor_t& knob = extractedKnobs_[idx]; cudnnBackendKnobType_t type = CUDNN_KNOB_TYPE_COUNTS; @@ -2607,6 +2610,14 @@ std::string static get_engine_tag(ManagedOpaqueDescriptor const config) { if (status != CUDNN_STATUS_SUCCESS) { return "INVALID_ENGINE_NAME_KNOB_CHOICE_KNOB_VALUE"; } + knob_choices.emplace_back(type, choice); + } + // Sort by knob type so the tag is a deterministic function of the engine + // config -- the knob-choice array order differs between the heuristics path + // and create_execution_plan (which iterates an unordered_map), but the + // engine + knob values are identical. + std::sort(knob_choices.begin(), knob_choices.end()); + for (auto const& [type, choice] : knob_choices) { tag << "_k" << type << "=" << choice; } return tag.str(); diff --git a/test/cpp/get_engine_and_knobs.cpp b/test/cpp/get_engine_and_knobs.cpp index 258159b83..c54e9f4c7 100644 --- a/test/cpp/get_engine_and_knobs.cpp +++ b/test/cpp/get_engine_and_knobs.cpp @@ -73,6 +73,9 @@ TEST_CASE("get_engine_and_knobs_at_index round-trips via create_execution_plan", // engine's target); when it succeeds it must reproduce the same plan. int64_t round_tripped = 0; for (int64_t i = 0; i < count; i++) { + std::string name; + REQUIRE(graph->get_plan_name_at_index(i, name).is_good()); + engine_id = -1; knobs.clear(); REQUIRE(graph->get_engine_and_knobs_at_index(i, engine_id, knobs).is_good()); @@ -83,15 +86,21 @@ TEST_CASE("get_engine_and_knobs_at_index round-trips via create_execution_plan", if (pinned->build_plans().is_good()) { REQUIRE(pinned->get_execution_plan_count() == 1); - // Compare the structured identity (engine + knob map), not the tag - // string: the tag serializes knobs in engine-config order, which - // differs between the heuristic and pinned configs, but the engine - // id + knob values are what define the kernel. + + // Structured identity (engine + knob map) must match. int64_t pinned_engine = -1; std::unordered_map pinned_knobs; REQUIRE(pinned->get_engine_and_knobs_at_index(0, pinned_engine, pinned_knobs).is_good()); REQUIRE(pinned_engine == engine_id); REQUIRE(pinned_knobs == knobs); + + // The plan-name tag is canonical (get_engine_tag sorts knobs by + // type), so it must match too despite the two configs storing + // knob choices in different orders. + std::string pinned_name; + REQUIRE(pinned->get_plan_name_at_index(0, pinned_name).is_good()); + REQUIRE(pinned_name == name); + round_tripped++; } }