feat(python): add get_engine_and_knobs_at_index for structured plan pinning

include/cudnn_frontend/graph_interface.h

@@ -1007,6 +1007,15 @@ class Graph : public ICudnn, public INode {
         return ret_val;
     }
 
+    // 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,
+                                  std::unordered_map<KnobType_t, int64_t> &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);

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;

include/cudnn_frontend/plans.h

@@ -2,8 +2,11 @@
 
 #include <optional>
 #include <string>
+#include <unordered_map>
+#include <utility>
 #include <vector>
 
+#include "knobs.h"
 #include "../cudnn_frontend_EngineConfig.h"
 #include "../cudnn_frontend_Logging.h"
 #include "graph_helpers.h"
@@ -477,6 +480,28 @@ class Execution_plan_list {
         return {error_code_t::OK, ""};
     }
 
+    // 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<KnobType_t, int64_t>& knobs) const {
+        RETURN_CUDNN_FRONTEND_ERROR_IF(index < 0 || index >= static_cast<int64_t>(engine_configs.size()),
+                                       error_code_t::GRAPH_EXECUTION_FAILED,
+                                       "Plan index " + std::to_string(index) + " is invalid.");
+        std::vector<std::pair<cudnnBackendKnobType_t, int64_t>> 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.

include/cudnn_frontend_utils.h

@@ -24,8 +24,10 @@
 #include <exception>
 #include <optional>
 #include <string>
+#include <algorithm>
 #include <variant>
 #include <vector>
+#include <utility>
 #include <iomanip>
 #include <sstream>
 #include <cmath>
@@ -2594,6 +2596,8 @@ std::string static get_engine_tag(ManagedOpaqueDescriptor const config) {
         return "INVALID_ENGINE_NAME_KNOB_COUNT";
     }
 
+    std::vector<std::pair<cudnnBackendKnobType_t, int64_t>> knob_choices;
+    knob_choices.reserve(static_cast<size_t>(numKnobs));
     for (size_t idx = 0; idx < static_cast<size_t>(numKnobs); ++idx) {
         const cudnnBackendDescriptor_t& knob = extractedKnobs_[idx];
         cudnnBackendKnobType_t type          = CUDNN_KNOB_TYPE_COUNTS;
@@ -2606,11 +2610,94 @@ 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();
 }
 
+// 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<std::pair<cudnnBackendKnobType_t, int64_t>>& 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<ManagedOpaqueDescriptor, CUDNN_KNOB_TYPE_COUNTS> 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<cudnnBackendDescriptor_t, CUDNN_KNOB_TYPE_COUNTS> 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<size_t>(numKnobs));
+    for (size_t idx = 0; idx < static_cast<size_t>(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

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_<cudnn_frontend::Knob, std::shared_ptr<cudnn_frontend::Knob>>(m, "knob")
         .def(py::init<cudnn_frontend::KnobType_t, int64_t, int64_t, int64_t>(),

python/pygraph/pygraph.cpp

@@ -807,6 +807,15 @@ PyGraph::get_plan_name_at_index(int64_t index) {
     return plan_name;
 }
 
+std::pair<int64_t, std::unordered_map<KnobType_t, int64_t>>
+PyGraph::get_engine_and_knobs_at_index(int64_t index) {
+    int64_t engine_id = -1;
+    std::unordered_map<KnobType_t, int64_t> 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<int64_t>
 default_vector(void) {
     return {};
@@ -1325,6 +1334,16 @@ 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 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.
+                )pbdoc")
         .def("_execute",
              &PyGraph::execute,
              py::arg("var_pack"),

python/pygraph/pygraph.h

@@ -798,6 +798,9 @@ class PyGraph {
     std::string
     get_plan_name_at_index(int64_t index);
 
+    std::pair<int64_t, std::unordered_map<KnobType_t, int64_t>>
+    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}

test/cpp/CMakeLists.txt

@@ -25,6 +25,7 @@ add_executable(
     validate.cpp
     version.cpp
     tensor.cpp
+    get_engine_and_knobs.cpp
 )
 
 if (MSVC)

test/cpp/get_engine_and_knobs.cpp

@@ -0,0 +1,110 @@
+/*
+ * 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 <catch2/catch_test_macros.hpp>
+
+#include <cudnn_frontend.h>
+
+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<fe::graph::Graph>
+make_matmul_graph(cudnnHandle_t handle) {
+    auto graph = std::make_shared<fe::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}));
+    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<fe::KnobType_t, int64_t> 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());
+
+    // 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);
+        REQUIRE(pinned->create_execution_plan(engine_id, knobs).is_good());
+
+        if (pinned->build_plans().is_good()) {
+            REQUIRE(pinned->get_execution_plan_count() == 1);
+
+            // Structured identity (engine + knob map) must match.
+            int64_t pinned_engine = -1;
+            std::unordered_map<fe::KnobType_t, int64_t> 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++;
+        }
+    }
+    REQUIRE(round_tripped > 0);
+
+    cudnnDestroy(handle);
+}