/*
 * Copyright (c) 2023, 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>

TEST_CASE("Tensor attributes", "[tensor][serialize]") {
    namespace fe = cudnn_frontend;

    auto tensor_attributes = fe::graph::Tensor_attributes()
                                 .set_name("image")
                                 .set_dim({4, 32, 16, 16})
                                 .set_stride({32 * 16 * 16, 1, 32 * 16, 32})
                                 .set_is_virtual(true)
                                 .set_is_pass_by_value(true)
                                 .set_uid(12312)
                                 .set_reordering_type(fe::TensorReordering_t::F16x16)
                                 .set_data_type(fe::DataType_t::HALF);

    json j                              = tensor_attributes;
    auto tensor_attributes_deserialized = j;

    REQUIRE(tensor_attributes_deserialized == tensor_attributes);
}

TEST_CASE("Context serialization", "[context][serialize]") {
    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)
        .set_sm_count(24);

    json j = graph;

    std::cout << j << std::endl;

    fe::graph::Graph graph_deserialized;

    REQUIRE(graph_deserialized.deserialize(j).is_good());

    json j2 = graph_deserialized;

    REQUIRE(j == j2);

    auto status = fe::graph::Graph().deserialize(j, true);
    REQUIRE(status.is_bad());
    REQUIRE(status.get_message().find("enforce_precompiled") != std::string::npos);

    REQUIRE(graph.validate().is_good());
}

TEST_CASE("Conv fprop attributes", "[conv_fprop][serialize]") {
    namespace fe = cudnn_frontend;

    auto x = std::make_shared<fe::graph::Tensor_attributes>();
    x->set_name("image")
        .set_dim({4, 32, 16, 16})
        .set_stride({32 * 16 * 16, 1, 32 * 16, 32})
        .set_is_virtual(true)
        .set_is_pass_by_value(true)
        .set_uid(12312)
        .set_reordering_type(fe::TensorReordering_t::F16x16)
        .set_data_type(fe::DataType_t::HALF);

    auto conv_fprop_attributes = fe::graph::Conv_fprop_attributes()
                                     .set_name("conv_fprop")
                                     .set_padding({1, 1})
                                     .set_stride({1, 1})
                                     .set_dilation({1, 1})
                                     .set_compute_data_type(fe::DataType_t::FLOAT);

    json j                                  = conv_fprop_attributes;
    auto conv_fprop_attributes_deserialized = j;

    REQUIRE(conv_fprop_attributes_deserialized == conv_fprop_attributes);
}

TEST_CASE("Graph key", "[graph][key]") {
    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 X = graph.tensor(
        fe::graph::Tensor_attributes().set_name("image").set_dim({4, 16, 64}).set_stride({16 * 64, 1, 16}));
    auto Y = graph.tensor(
        fe::graph::Tensor_attributes().set_name("filter").set_dim({4, 64, 32}).set_stride({32 * 64, 1, 64}));

    fe::graph::Matmul_attributes matmul;
    auto Z = graph.matmul(X, Y, matmul);

    auto scale_options = fe::graph::Pointwise_attributes().set_mode(fe::PointwiseMode_t::MUL);
    auto S             = graph.tensor(
        fe::graph::Tensor_attributes().set_name("scale").set_dim({4, 16, 32}).set_stride({16 * 32, 32, 1}));
    auto scale_output = graph.pointwise(Z, S, scale_options);

    auto bias_options = fe::graph::Pointwise_attributes().set_mode(fe::PointwiseMode_t::ADD);
    auto B =
        graph.tensor(fe::graph::Tensor_attributes().set_name("bias").set_dim({4, 16, 32}).set_stride({16 * 32, 32, 1}));
    auto bias_output = graph.pointwise(scale_output, B, bias_options);

    auto relu_options = fe::graph::Pointwise_attributes().set_mode(fe::PointwiseMode_t::RELU_FWD);
    auto O            = graph.pointwise(bias_output, relu_options);
    O->set_output(true);

    cudnnHandle_t handle;
    cudnnCreate(&handle);

    REQUIRE(graph.validate().is_good());

    REQUIRE(graph.build_operation_graph(handle).is_good());
    auto key = graph.key();

    REQUIRE(graph.create_execution_plans({fe::HeurMode_t::A}).is_good());
    REQUIRE(key == graph.key());

    REQUIRE(graph.check_support().is_good());
    REQUIRE(key == graph.key());

    REQUIRE(graph.build_plans().is_good());
    REQUIRE(key == graph.key());

    cudnnDestroy(handle);
}

TEST_CASE("Graph key dynamic shape", "[graph][key][dynamic_shape]") {
    namespace fe = cudnn_frontend;
    if (cudnnGetCudartVersion() < 12000) {
        SKIP("Test requires cuda toolkit 12.0 or above");
    }
    // clang-format off
    struct {
        int64_t b,    m,    n,    k;
    } shapes[] = {
        {       4,   16,   32,   64},
        {       8,   32,   64,  128},
    };
    // clang-format on

    constexpr int shapes_count = sizeof(shapes) / sizeof(shapes[0]);
    size_t key                 = 0;  // Save key between runs to verify that dim and stride information is deleted

    for (int idx_shape = 0; idx_shape < shapes_count; idx_shape++) {
        auto b = shapes[idx_shape].b;
        auto m = shapes[idx_shape].m;
        auto n = shapes[idx_shape].n;
        auto k = shapes[idx_shape].k;

        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)
            .set_dynamic_shape_enabled(true);

        auto X =
            graph.tensor(fe::graph::Tensor_attributes().set_name("image").set_dim({b, m, k}).set_stride({m * k, 1, m}));
        auto Y = graph.tensor(
            fe::graph::Tensor_attributes().set_name("filter").set_dim({b, k, n}).set_stride({n * k, 1, k}));

        fe::graph::Matmul_attributes matmul;
        auto Z = graph.matmul(X, Y, matmul);

        auto scale_options = fe::graph::Pointwise_attributes().set_mode(fe::PointwiseMode_t::MUL);
        auto S =
            graph.tensor(fe::graph::Tensor_attributes().set_name("scale").set_dim({b, m, n}).set_stride({m * n, n, 1}));
        auto scale_output = graph.pointwise(Z, S, scale_options);

        auto bias_options = fe::graph::Pointwise_attributes().set_mode(fe::PointwiseMode_t::ADD);
        auto B =
            graph.tensor(fe::graph::Tensor_attributes().set_name("bias").set_dim({b, m, n}).set_stride({m * n, n, 1}));
        auto bias_output = graph.pointwise(scale_output, B, bias_options);

        auto relu_options = fe::graph::Pointwise_attributes().set_mode(fe::PointwiseMode_t::RELU_FWD);
        auto O            = graph.pointwise(bias_output, relu_options);
        O->set_output(true);

        cudnnHandle_t handle;
        cudnnCreate(&handle);

        auto status = graph.validate();
        if (cudnnGetVersion() >= 90400) {
            REQUIRE(status.is_good());
        } else {
            REQUIRE(status.is_bad());
            SKIP("Dynamic shapes not supported pre 9.4");
        }

        REQUIRE(graph.build_operation_graph(handle).is_good());

        if (!key) {
            key = graph.key();
        }

        REQUIRE(graph.create_execution_plans({fe::HeurMode_t::A}).is_good());
        REQUIRE(key == graph.key());

        REQUIRE(graph.check_support().is_good());
        REQUIRE(key == graph.key());

        REQUIRE(graph.build_plans().is_good());
        REQUIRE(key == graph.key());

        cudnnDestroy(handle);
    }
}

TEST_CASE("Matmul fp8 fusion", "[graph][serialize]") {
    namespace fe = cudnn_frontend;
    // matmul problem size
    int64_t const b = 16;
    int64_t const m = 32;
    int64_t const n = 64;
    int64_t const k = 128;

    fe::graph::Graph graph{};

    // Create the two non-virtual input tensors A and B.
    // There are read from global memory.
    auto A_attributes = fe::graph::Tensor_attributes()
                            .set_name("A")
                            .set_dim({b, m, k})
                            .set_stride({m * k, k, 1})
                            .set_data_type(fe::DataType_t::FP8_E4M3);
    auto A = graph.tensor(A_attributes);

    auto B_attributes = fe::graph::Tensor_attributes()
                            .set_name("B")
                            .set_dim({b, k, n})
                            .set_stride({k * n, 1, k})
                            .set_data_type(fe::DataType_t::FP8_E4M3);
    auto B = graph.tensor(B_attributes);

    auto A_descale_attributes = fe::graph::Tensor_attributes()
                                    .set_name("descale0")
                                    .set_dim({1, 1, 1})
                                    .set_stride({1, 1, 1})
                                    .set_data_type(fe::DataType_t::FLOAT);
    auto B_descale_attributes = fe::graph::Tensor_attributes()
                                    .set_name("descale1")
                                    .set_dim({1, 1, 1})
                                    .set_stride({1, 1, 1})
                                    .set_data_type(fe::DataType_t::FLOAT);

    auto A_descale = graph.tensor(A_descale_attributes);
    auto B_descale = graph.tensor(B_descale_attributes);

    auto matmul_attributes =
        // fe::graph::Matmul_attributes().set_name("GEMM").set_compute_data_type(fe::DataType_t::FLOAT);
        fe::graph::Matmul_attributes().set_name("GEMM").set_compute_data_type(fe::DataType_t::FLOAT);
    auto C = graph.matmul(A, B, matmul_attributes);
    C->set_data_type(fe::DataType_t::FLOAT);

    // Add scale_A operation
    auto pw_0_attributes = fe::graph::Pointwise_attributes()
                               //    .set_name("pw0_Mul")
                               .set_mode(fe::PointwiseMode_t::MUL)
                               .set_compute_data_type(fe::DataType_t::FLOAT);
    auto C_after_pw_0 = graph.pointwise(C, A_descale, pw_0_attributes);
    C_after_pw_0->set_data_type(fe::DataType_t::FLOAT);

    // Add descale_B operation
    auto pw_1_attributes = fe::graph::Pointwise_attributes()
                               //    .set_name("pw1_Mul")
                               .set_mode(fe::PointwiseMode_t::MUL)
                               .set_compute_data_type(fe::DataType_t::FLOAT);
    auto C_after_pw_1 = graph.pointwise(C_after_pw_0, B_descale, pw_1_attributes);
    C_after_pw_1->set_output(true).set_data_type(fe::DataType_t::BFLOAT16);

    json j = graph;

    std::cout << j << std::endl;

    fe::graph::Graph graph_deserialized;

    REQUIRE(graph_deserialized.deserialize(j).is_good());

    json j2 = graph_deserialized;

    REQUIRE(j == j2);

    REQUIRE(graph.validate().is_good());

    std::cout << "Validating deserialized graph" << std::endl;

    cudnnHandle_t handle;  // Handle to use during deserialize and execute

    cudnnCreate(&handle);

    REQUIRE(graph_deserialized.validate().is_good());

    REQUIRE(graph_deserialized.build_operation_graph(handle).is_good());

    cudnnDestroy(handle);
}

TEST_CASE("conv graph serialization", "[graph][serialize]") {
    namespace fe = cudnn_frontend;

    fe::graph::Graph graph;

    auto x = graph.tensor(fe::graph::Tensor_attributes());
    x->set_name("image")
        .set_dim({4, 32, 16, 16})
        .set_stride({32 * 16 * 16, 1, 32 * 16, 32})
        .set_is_virtual(false)
        .set_is_pass_by_value(false)
        .set_reordering_type(fe::TensorReordering_t::NONE)
        .set_data_type(fe::DataType_t::HALF);

    auto w = graph.tensor(fe::graph::Tensor_attributes());
    w->set_name("weight")
        .set_dim({64, 32, 3, 3})
        .set_stride({32 * 3 * 3, 1, 32 * 3, 32})
        .set_is_virtual(false)
        .set_is_pass_by_value(false)
        .set_reordering_type(fe::TensorReordering_t::NONE)
        .set_data_type(fe::DataType_t::HALF);

    auto conv_fprop_attributes = fe::graph::Conv_fprop_attributes()
                                     .set_name("conv_fprop")
                                     .set_padding({1, 1})
                                     .set_stride({1, 1})
                                     .set_dilation({1, 1})
                                     .set_compute_data_type(fe::DataType_t::FLOAT);

    auto y = graph.conv_fprop(x, w, conv_fprop_attributes);

    auto b = graph.tensor(fe::graph::Tensor_attributes());
    b->set_name("bias")
        .set_dim({1, 64, 1, 1})
        .set_stride({64, 1, 64, 64})
        .set_is_virtual(false)
        .set_is_pass_by_value(false)
        .set_reordering_type(fe::TensorReordering_t::NONE)
        .set_data_type(fe::DataType_t::HALF);

    auto pointwise_attributes = fe::graph::Pointwise_attributes()
                                    .set_name("bias")
                                    .set_mode(fe::PointwiseMode_t::ADD)
                                    .set_compute_data_type(fe::DataType_t::FLOAT);

    auto o = graph.pointwise(y, b, pointwise_attributes);

    auto reduction_attributes = fe::graph::Reduction_attributes()
                                    .set_name("reduction")
                                    .set_mode(fe::ReductionMode_t::ADD)
                                    .set_compute_data_type(fe::DataType_t::FLOAT);
    auto r = graph.reduction(o, reduction_attributes);

    r->set_output(true).set_data_type(fe::DataType_t::HALF);

    json j = graph;

    fe::graph::Graph graph_deserialized;

    REQUIRE(graph_deserialized.deserialize(j).is_good());

    json j2 = graph_deserialized;

    REQUIRE(j == j2);

    REQUIRE(graph_deserialized.validate().is_good());
}

TEST_CASE("sdpa graph serialization", "[graph][serialize]") {
    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);

    int64_t b    = 3;     // batch size
    int64_t h    = 4;     // head dim
    int64_t s_q  = 1024;  // q tensor is padded to this seq length
    int64_t s_kv = 1024;  // k and v tensor is padded to this seq length
    int64_t d    = 128;   // hidden dim

    auto Q = graph.tensor(fe::graph::Tensor_attributes()
                              .set_name("Q")
                              .set_dim({b, h, s_q, d})
                              .set_stride({3 * h * d, 3 * d, 3 * b * h * d, 1}));
    auto K = graph.tensor(fe::graph::Tensor_attributes()
                              .set_name("K")
                              .set_dim({b, h, s_kv, d})
                              .set_stride({3 * h * d, 3 * d, 3 * b * h * d, 1}));
    auto V = graph.tensor(fe::graph::Tensor_attributes()
                              .set_name("V")
                              .set_dim({b, h, s_kv, d})
                              .set_stride({3 * h * d, 3 * d, 3 * b * h * d, 1}));

    auto attn_scale = graph.tensor(fe::graph::Tensor_attributes()
                                       .set_name("attn_scale")
                                       .set_dim({1, 1, 1, 1})
                                       .set_stride({1, 1, 1, 1})
                                       .set_is_pass_by_value(true)
                                       .set_data_type(fe::DataType_t::FLOAT));

    auto seed   = graph.tensor(fe::graph::Tensor_attributes()
                                 .set_name("Seed")
                                 .set_dim({1, 1, 1, 1})
                                 .set_stride({1, 1, 1, 1})
                                 .set_data_type(fe::DataType_t::INT32));
    auto offset = graph.tensor(fe::graph::Tensor_attributes()
                                   .set_name("Offset")
                                   .set_dim({1, 1, 1, 1})
                                   .set_stride({1, 1, 1, 1})
                                   .set_data_type(fe::DataType_t::INT32));

    auto bias = graph.tensor(fe::graph::Tensor_attributes()
                                 .set_name("bias")
                                 .set_dim({b, 1, s_q, s_kv})
                                 .set_stride({s_q * s_kv, s_q * s_kv, s_kv, 1}));

    auto seq_q  = graph.tensor(fe::graph::Tensor_attributes()
                                  .set_name("seq_q")
                                  .set_dim({b, 1, 1, 1})
                                  .set_stride({1, 1, 1, 1})
                                  .set_data_type(fe::DataType_t::INT32));
    auto seq_kv = graph.tensor(fe::graph::Tensor_attributes()
                                   .set_name("seq_kv")
                                   .set_dim({b, 1, 1, 1})
                                   .set_stride({1, 1, 1, 1})
                                   .set_data_type(fe::DataType_t::INT32));

    auto sdpa_options = fe::graph::SDPA_attributes()
                            .set_name("flash_attention")
                            .set_generate_stats(true)
                            .set_attn_scale(attn_scale)
                            .set_alibi_mask(true)
                            .set_diagonal_band_right_bound(0)
                            .set_dropout(0.1f, seed, offset)
                            .set_bias(bias)
                            .set_padding_mask(true)
                            .set_seq_len_q(seq_q)
                            .set_seq_len_kv(seq_kv);

    auto [O, stats] = graph.sdpa(Q, K, V, sdpa_options);

    O->set_output(true).set_dim({b, h, s_q, d}).set_stride({h * d, d, b * h * d, 1});
    stats->set_output(true).set_data_type(fe::DataType_t::FLOAT);

    json j = graph;

    fe::graph::Graph graph_deserialized;
    REQUIRE(graph_deserialized.deserialize(j).is_good());
    json j2 = graph_deserialized;

    REQUIRE(j == j2);
}

TEST_CASE("sdpa backward graph serialization", "[graph][serialize]") {
    namespace fe = cudnn_frontend;

    int64_t b    = 3;     // batch size
    int64_t h    = 4;     // head dim
    int64_t s_q  = 1024;  // q tensor is padded to this seq length
    int64_t s_kv = 1024;  // k and v tensor is padded to this seq length
    int64_t d    = 128;   // hidden dim

    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);

    std::shared_ptr<fe::graph::Tensor_attributes> bias, dropout_seed, dropout_offset;

    auto q = graph.tensor(
        fe::graph::Tensor_attributes().set_name("Q").set_dim({b, h, s_q, d}).set_stride({h * s_q * d, s_q * d, d, 1}));
    auto k = graph.tensor(fe::graph::Tensor_attributes()
                              .set_name("K")
                              .set_dim({b, h, s_kv, d})
                              .set_stride({h * s_kv * d, s_kv * d, d, 1}));
    auto v = graph.tensor(fe::graph::Tensor_attributes()
                              .set_name("V")
                              .set_dim({b, h, s_kv, d})
                              .set_stride({h * s_kv * d, s_kv * d, d, 1}));
    auto o = graph.tensor(
        fe::graph::Tensor_attributes().set_name("O").set_dim({b, h, s_q, d}).set_stride({h * s_q * d, s_q * d, d, 1}));
    auto dO = graph.tensor(
        fe::graph::Tensor_attributes().set_name("dO").set_dim({b, h, s_q, d}).set_stride({h * s_q * d, s_q * d, d, 1}));
    auto stats = graph.tensor(fe::graph::Tensor_attributes()
                                  .set_name("stats")
                                  .set_dim({b, h, s_q, 1})
                                  .set_stride({h * s_q, s_q, 1, 1})
                                  .set_data_type(fe::DataType_t::FLOAT));

    auto attn_scale = graph.tensor(fe::graph::Tensor_attributes()
                                       .set_name("attn_scale")
                                       .set_dim({1, 1, 1, 1})
                                       .set_stride({1, 1, 1, 1})
                                       .set_is_pass_by_value(true)
                                       .set_data_type(fe::DataType_t::FLOAT));

    bias = graph.tensor(fe::graph::Tensor_attributes()
                            .set_name("bias")
                            .set_dim({b, 1, s_q, s_kv})
                            .set_stride({s_q * s_kv, s_q * s_kv, s_kv, 1}));

    dropout_seed   = graph.tensor(fe::graph::Tensor_attributes()
                                    .set_name("Seed")
                                    .set_dim({1, 1, 1, 1})
                                    .set_stride({1, 1, 1, 1})
                                    .set_data_type(fe::DataType_t::INT32));
    dropout_offset = graph.tensor(fe::graph::Tensor_attributes()
                                      .set_name("Offset")
                                      .set_dim({1, 1, 1, 1})
                                      .set_stride({1, 1, 1, 1})
                                      .set_data_type(fe::DataType_t::INT32));

    auto sdpa_backward_options = fe::graph::SDPA_backward_attributes()
                                     .set_name("flash_attention_backward")
                                     .set_diagonal_alignment(cudnn_frontend::DiagonalAlignment_t::TOP_LEFT)
                                     .set_diagonal_band_right_bound(0)
                                     .set_attn_scale(attn_scale)
                                     .set_bias(bias)
                                     .set_dropout(0.1f, dropout_seed, dropout_offset);

    auto [dQ, dK, dV] = graph.sdpa_backward(q, k, v, o, dO, stats, sdpa_backward_options);

    dQ->set_output(true).set_dim({b, h, s_q, d}).set_stride({h * s_q * d, s_q * d, d, 1});
    dK->set_output(true).set_dim({b, h, s_kv, d}).set_stride({h * s_kv * d, s_kv * d, d, 1});
    dV->set_output(true).set_dim({b, h, s_kv, d}).set_stride({h * s_kv * d, s_kv * d, d, 1});

    json j = graph;
    fe::graph::Graph graph_deserialized;
    REQUIRE(graph_deserialized.deserialize(j).is_good());
    json j2 = graph_deserialized;

    REQUIRE(j == j2);

    REQUIRE(graph_deserialized.validate().is_good());
}

// Round-trips a plan through binary serialize/deserialize and confirms the variant
// pack template is populated on the deserialized graph. The check verifies that
// deserialize(handle, ...) eagerly prepares the template (instead of waiting for
// first execute).
TEST_CASE("Plan deserialize prepares variant pack template", "[graph][serialize][deserialize]") {
    namespace fe = cudnn_frontend;

    constexpr int64_t a_uid = 1, b_uid = 2, c_uid = 3;

    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(a_uid));
    auto B = graph.tensor(
        fe::graph::Tensor_attributes().set_name("B").set_dim({4, 64, 32}).set_stride({64 * 32, 32, 1}).set_uid(b_uid));

    auto C = graph.matmul(A, B, fe::graph::Matmul_attributes().set_name("matmul"));
    C->set_output(true).set_uid(c_uid);

    cudnnHandle_t handle;
    cudnnCreate(&handle);

    REQUIRE(graph.build(handle, {fe::HeurMode_t::A}).is_good());

    std::vector<uint8_t> serialized_data;
    REQUIRE(graph.serialize(serialized_data).is_good());

    fe::graph::Graph graph_deserialized;
    REQUIRE(graph_deserialized.deserialize(handle, serialized_data, true).is_good());

    // Variant pack template should already be populated; no execute needed.
    auto const user_uids = graph_deserialized.get_variant_pack_uids_sorted();
    REQUIRE(user_uids == std::vector<int64_t>{a_uid, b_uid, c_uid});

    cudnnDestroy(handle);
}

// Exercises the run_warmup=false fast path of deserialize(handle, ...). Skipping
// the throwaway warmup capture must not change the resulting plan: the variant
// pack template is built by prepare_variant_pack_template(), which is independent
// of warmup, so the deserialized graph is still fully usable.
//
// Note the additive argument order: deserialize(handle, data, enforce_precompiled, run_warmup).
// Both bools are spelled explicitly below so the positional meaning is unambiguous.
TEST_CASE("Plan deserialize with run_warmup=false still prepares template", "[graph][serialize][deserialize]") {
    namespace fe = cudnn_frontend;

    constexpr int64_t a_uid = 1, b_uid = 2, c_uid = 3;

    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(a_uid));
    auto B = graph.tensor(
        fe::graph::Tensor_attributes().set_name("B").set_dim({4, 64, 32}).set_stride({64 * 32, 32, 1}).set_uid(b_uid));

    auto C = graph.matmul(A, B, fe::graph::Matmul_attributes().set_name("matmul"));
    C->set_output(true).set_uid(c_uid);

    cudnnHandle_t handle;
    cudnnCreate(&handle);

    REQUIRE(graph.build(handle, {fe::HeurMode_t::A}).is_good());

    // serialize the graph
    std::vector<uint8_t> serialized_data;
    REQUIRE(graph.serialize(serialized_data).is_good());

    // expected uids
    std::vector<int64_t> const expected_uids{a_uid, b_uid, c_uid};

    // test the blob overload, run_warmup=false.
    // Check for success and variant pack uids are the same as the expected uids
    SECTION("blob overload, run_warmup=false") {
        fe::graph::Graph graph_deserialized;
        REQUIRE(
            graph_deserialized.deserialize(handle, serialized_data, /*enforce_precompiled=*/false, /*run_warmup=*/false)
                .is_good());
        // check the variant pack uids are the same as the expected uids
        REQUIRE(graph_deserialized.get_variant_pack_uids_sorted() == expected_uids);
        // tensor metadata must still resolve with warmup skipped (deserialized_tensor_properties)
        fe::graph::Tensor_attributes queried;
        REQUIRE(graph_deserialized.query_tensor_attributes_of_uid(a_uid, queried).is_good());
    }

    // test the json overload, run_warmup=false.
    // same assertion via the pre-parsed-json overload, covering the path that avoids a second from_ubjson.
    SECTION("json overload, run_warmup=false") {
        json const j = json::from_ubjson(serialized_data);
        fe::graph::Graph graph_deserialized;
        auto const status =
            graph_deserialized.deserialize(handle, j, /*enforce_precompiled=*/false, /*run_warmup=*/false);
        REQUIRE(status.is_good());
        REQUIRE(graph_deserialized.get_variant_pack_uids_sorted() == expected_uids);
    }

    // test the run_warmup=false matches the default warmup=true template
    // Confirms warmup is purely a priming step with no effect on the executable plan.
    SECTION("run_warmup=false matches the default warmup=true template") {
        fe::graph::Graph warmed;
        REQUIRE(
            warmed.deserialize(handle, serialized_data, /*enforce_precompiled=*/false, /*run_warmup=*/true).is_good());

        fe::graph::Graph skipped;
        auto const skipped_status =
            skipped.deserialize(handle, serialized_data, /*enforce_precompiled=*/false, /*run_warmup=*/false);
        REQUIRE(skipped_status.is_good());

        // The plan-side state the fast execute path relies on is identical either way.
        REQUIRE(warmed.get_variant_pack_uids_sorted() == skipped.get_variant_pack_uids_sorted());
        REQUIRE(skipped.get_variant_pack_uids_sorted() == expected_uids);
    }

    cudnnDestroy(handle);
}
