/*
 * Copyright (c) 2024, 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 <catch2/generators/catch_generators.hpp>

#include "../utils/helpers.h"

#include <cudnn_frontend.h>
#include <cudnn_frontend_utils.h>

namespace BlackwellNVFP4MXFP8BlockScaleMatmul {

namespace fe = cudnn_frontend;

struct TestParams {
    int64_t b = -1;
    int64_t m = -1;
    int64_t n = -1;
    int64_t k = -1;

    int32_t block_size = -1;

    cudnn_frontend::DataType_t datatype_a = cudnn_frontend::DataType_t::NOT_SET;
    cudnn_frontend::DataType_t datatype_b = cudnn_frontend::DataType_t::NOT_SET;

    cudnn_frontend::DataType_t datatype_block_scale = cudnn_frontend::DataType_t::NOT_SET;

    cudnn_frontend::DataType_t datatype_d = cudnn_frontend::DataType_t::NOT_SET;

    TestParams(int64_t b_,
               int64_t m_,
               int64_t n_,
               int64_t k_,
               int32_t block_size_,
               cudnn_frontend::DataType_t datatype_a_,
               cudnn_frontend::DataType_t datatype_b_,
               cudnn_frontend::DataType_t datatype_block_scale_,
               cudnn_frontend::DataType_t datatype_d_) {
        b = b_;
        m = m_;
        n = n_;
        k = k_;

        block_size = block_size_;

        datatype_a           = datatype_a_;
        datatype_b           = datatype_b_;
        datatype_block_scale = datatype_block_scale_;
        datatype_d           = datatype_d_;
    }
};

// Note: For nvfp4/mxfp8 block scale matmul, scale factors need to be organized in 128x4 layout for best performance.
//       Details of the layout can be found https://docs.nvidia.com/cuda/cublas/#d-block-scaling-factors-layout
TEST_CASE("Blackwell Block Scale Matmul", "[matmul][graph][FP4]") {
#if (CUDNN_VERSION < 90700)
    SKIP("Matmul with block scaling is not supported in cudnn versions prior to 9.7.0");
#endif
    if (check_device_arch_newer_than("blackwell") == false) {
        SKIP("Hardware accelerated NVFP4/MXFP8 block scale matmul requires Blackwell and up");
    }

    auto test_params = GENERATE(
        // TestParams(1,
        //            128,
        //            128,
        //            64,
        //            16,
        //            cudnn_frontend::DataType_t::FP4_E2M1,
        //            cudnn_frontend::DataType_t::FP4_E2M1,
        //            cudnn_frontend::DataType_t::FP8_E4M3,
        //            cudnn_frontend::DataType_t::FLOAT),
        TestParams(1,
                   128,
                   128,
                   64,
                   16,
                   cudnn_frontend::DataType_t::FP4_E2M1,
                   cudnn_frontend::DataType_t::FP4_E2M1,
                   cudnn_frontend::DataType_t::FP8_E4M3,
                   cudnn_frontend::DataType_t::HALF),
        TestParams(1,
                   128,
                   128,
                   64,
                   16,
                   cudnn_frontend::DataType_t::FP4_E2M1,
                   cudnn_frontend::DataType_t::FP4_E2M1,
                   cudnn_frontend::DataType_t::FP8_E4M3,
                   cudnn_frontend::DataType_t::BFLOAT16),
        TestParams(3,
                   137,
                   268,
                   160,
                   16,
                   cudnn_frontend::DataType_t::FP4_E2M1,
                   cudnn_frontend::DataType_t::FP4_E2M1,
                   cudnn_frontend::DataType_t::FP8_E4M3,
                   cudnn_frontend::DataType_t::FLOAT),
        TestParams(3,
                   137,
                   272,
                   160,
                   16,
                   cudnn_frontend::DataType_t::FP4_E2M1,
                   cudnn_frontend::DataType_t::FP4_E2M1,
                   cudnn_frontend::DataType_t::FP8_E4M3,
                   cudnn_frontend::DataType_t::HALF),
        TestParams(3,
                   137,
                   272,
                   160,
                   16,
                   cudnn_frontend::DataType_t::FP4_E2M1,
                   cudnn_frontend::DataType_t::FP4_E2M1,
                   cudnn_frontend::DataType_t::FP8_E4M3,
                   cudnn_frontend::DataType_t::BFLOAT16),
        TestParams(1,
                   128,
                   128,
                   128,
                   32,
                   cudnn_frontend::DataType_t::FP8_E4M3,
                   cudnn_frontend::DataType_t::FP8_E4M3,
                   cudnn_frontend::DataType_t::FP8_E8M0,
                   cudnn_frontend::DataType_t::FLOAT),
        TestParams(1,
                   128,
                   128,
                   128,
                   32,
                   cudnn_frontend::DataType_t::FP8_E4M3,
                   cudnn_frontend::DataType_t::FP8_E4M3,
                   cudnn_frontend::DataType_t::FP8_E8M0,
                   cudnn_frontend::DataType_t::HALF),
        TestParams(1,
                   128,
                   128,
                   128,
                   32,
                   cudnn_frontend::DataType_t::FP8_E4M3,
                   cudnn_frontend::DataType_t::FP8_E4M3,
                   cudnn_frontend::DataType_t::FP8_E8M0,
                   cudnn_frontend::DataType_t::BFLOAT16),
        TestParams(1,
                   128,
                   128,
                   128,
                   32,
                   cudnn_frontend::DataType_t::FP8_E4M3,
                   cudnn_frontend::DataType_t::FP8_E5M2,
                   cudnn_frontend::DataType_t::FP8_E8M0,
                   cudnn_frontend::DataType_t::FLOAT),
        TestParams(1,
                   128,
                   128,
                   128,
                   32,
                   cudnn_frontend::DataType_t::FP8_E4M3,
                   cudnn_frontend::DataType_t::FP8_E5M2,
                   cudnn_frontend::DataType_t::FP8_E8M0,
                   cudnn_frontend::DataType_t::HALF),
        TestParams(1,
                   128,
                   128,
                   128,
                   32,
                   cudnn_frontend::DataType_t::FP8_E4M3,
                   cudnn_frontend::DataType_t::FP8_E5M2,
                   cudnn_frontend::DataType_t::FP8_E8M0,
                   cudnn_frontend::DataType_t::BFLOAT16),
        TestParams(1,
                   128,
                   128,
                   128,
                   32,
                   cudnn_frontend::DataType_t::FP8_E5M2,
                   cudnn_frontend::DataType_t::FP8_E4M3,
                   cudnn_frontend::DataType_t::FP8_E8M0,
                   cudnn_frontend::DataType_t::FLOAT),
        TestParams(1,
                   128,
                   128,
                   128,
                   32,
                   cudnn_frontend::DataType_t::FP8_E5M2,
                   cudnn_frontend::DataType_t::FP8_E4M3,
                   cudnn_frontend::DataType_t::FP8_E8M0,
                   cudnn_frontend::DataType_t::HALF),
        TestParams(1,
                   128,
                   128,
                   128,
                   32,
                   cudnn_frontend::DataType_t::FP8_E5M2,
                   cudnn_frontend::DataType_t::FP8_E4M3,
                   cudnn_frontend::DataType_t::FP8_E8M0,
                   cudnn_frontend::DataType_t::BFLOAT16),
        TestParams(3,
                   137,
                   268,
                   160,
                   32,
                   cudnn_frontend::DataType_t::FP8_E4M3,
                   cudnn_frontend::DataType_t::FP8_E4M3,
                   cudnn_frontend::DataType_t::FP8_E8M0,
                   cudnn_frontend::DataType_t::FLOAT),
        TestParams(3,
                   137,
                   272,
                   160,
                   32,
                   cudnn_frontend::DataType_t::FP8_E4M3,
                   cudnn_frontend::DataType_t::FP8_E4M3,
                   cudnn_frontend::DataType_t::FP8_E8M0,
                   cudnn_frontend::DataType_t::HALF),
        TestParams(3,
                   137,
                   272,
                   160,
                   32,
                   cudnn_frontend::DataType_t::FP8_E4M3,
                   cudnn_frontend::DataType_t::FP8_E4M3,
                   cudnn_frontend::DataType_t::FP8_E8M0,
                   cudnn_frontend::DataType_t::BFLOAT16),
        TestParams(3,
                   137,
                   268,
                   160,
                   32,
                   cudnn_frontend::DataType_t::FP8_E4M3,
                   cudnn_frontend::DataType_t::FP8_E5M2,
                   cudnn_frontend::DataType_t::FP8_E8M0,
                   cudnn_frontend::DataType_t::FLOAT),
        TestParams(3,
                   137,
                   272,
                   160,
                   32,
                   cudnn_frontend::DataType_t::FP8_E4M3,
                   cudnn_frontend::DataType_t::FP8_E5M2,
                   cudnn_frontend::DataType_t::FP8_E8M0,
                   cudnn_frontend::DataType_t::HALF),
        TestParams(3,
                   137,
                   272,
                   160,
                   32,
                   cudnn_frontend::DataType_t::FP8_E4M3,
                   cudnn_frontend::DataType_t::FP8_E5M2,
                   cudnn_frontend::DataType_t::FP8_E8M0,
                   cudnn_frontend::DataType_t::BFLOAT16),
        TestParams(3,
                   137,
                   268,
                   160,
                   32,
                   cudnn_frontend::DataType_t::FP8_E5M2,
                   cudnn_frontend::DataType_t::FP8_E4M3,
                   cudnn_frontend::DataType_t::FP8_E8M0,
                   cudnn_frontend::DataType_t::FLOAT),
        TestParams(3,
                   137,
                   272,
                   160,
                   32,
                   cudnn_frontend::DataType_t::FP8_E5M2,
                   cudnn_frontend::DataType_t::FP8_E4M3,
                   cudnn_frontend::DataType_t::FP8_E8M0,
                   cudnn_frontend::DataType_t::HALF),
        TestParams(3,
                   137,
                   272,
                   160,
                   32,
                   cudnn_frontend::DataType_t::FP8_E5M2,
                   cudnn_frontend::DataType_t::FP8_E4M3,
                   cudnn_frontend::DataType_t::FP8_E8M0,
                   cudnn_frontend::DataType_t::BFLOAT16));

    int64_t const b = test_params.b;
    int64_t const m = test_params.m;
    int64_t const n = test_params.n;
    int64_t const k = test_params.k;

    auto datatype_a = test_params.datatype_a;
    auto datatype_b = test_params.datatype_b;

    auto datatype_block_scale = test_params.datatype_block_scale;

    auto block_size = test_params.block_size;

    auto datatype_d = test_params.datatype_d;

    Surface<int8_t> tensor_a_gpu(div_up(b * m * k * cudnn_frontend::detail::get_element_size_in_bits(datatype_a), 8));
    Surface<int8_t> tensor_b_gpu(div_up(b * k * n * cudnn_frontend::detail::get_element_size_in_bits(datatype_b), 8));

    static constexpr int indestructible_128x4_block_m_n = 128;
    static constexpr int indestructible_128x4_block_k   = 4;

    int64_t block_scale_dim_m = div_up(m, indestructible_128x4_block_m_n) * indestructible_128x4_block_m_n;
    int64_t block_scale_dim_n = div_up(n, indestructible_128x4_block_m_n) * indestructible_128x4_block_m_n;
    int64_t block_scale_dim_k =
        div_up(div_up(k, block_size), indestructible_128x4_block_k) * indestructible_128x4_block_k;

    Surface<int8_t> block_descale_a_gpu(
        div_up(b * block_scale_dim_m * block_scale_dim_k *
                   cudnn_frontend::detail::get_element_size_in_bits(datatype_block_scale),
               8));
    Surface<int8_t> block_descale_b_gpu(
        div_up(b * block_scale_dim_n * block_scale_dim_k *
                   cudnn_frontend::detail::get_element_size_in_bits(datatype_block_scale),
               8));

    Surface<int8_t> tensor_d_gpu(div_up(b * m * n * cudnn_frontend::detail::get_element_size_in_bits(datatype_d), 8));

    fe::graph::Graph graph{};

    graph.set_intermediate_data_type(fe::DataType_t::FLOAT);
    graph.set_compute_data_type(fe::DataType_t::FLOAT);

    auto tensor_a = graph.tensor(fe::graph::Tensor_attributes()
                                     .set_name("tensor_a")
                                     .set_data_type(datatype_a)
                                     .set_dim({b, m, k})
                                     .set_stride({m * k, k, 1}));

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

    auto block_descale_a = graph.tensor(fe::graph::Tensor_attributes()
                                            .set_name("block_descale_a")
                                            .set_data_type(datatype_block_scale)
                                            .set_dim({b, block_scale_dim_m, block_scale_dim_k})
                                            .set_stride({block_scale_dim_m * block_scale_dim_k, block_scale_dim_k, 1})
                                            .set_reordering_type(cudnn_frontend::TensorReordering_t::F8_128x4));

    auto block_descale_b = graph.tensor(fe::graph::Tensor_attributes()
                                            .set_name("block_descale_b")
                                            .set_data_type(datatype_block_scale)
                                            .set_dim({b, block_scale_dim_k, block_scale_dim_n})
                                            .set_stride({block_scale_dim_n * block_scale_dim_k, 1, block_scale_dim_k})
                                            .set_reordering_type(cudnn_frontend::TensorReordering_t::F8_128x4));

    auto dequantize_attr_a = fe::graph::Block_scale_dequantize_attributes().set_block_size({1, block_size});

    auto dequan_tensor_a = graph.block_scale_dequantize(tensor_a, block_descale_a, dequantize_attr_a);

    auto dequantize_attr_b = fe::graph::Block_scale_dequantize_attributes().set_block_size({block_size, 1});

    auto dequan_tensor_b = graph.block_scale_dequantize(tensor_b, block_descale_b, dequantize_attr_b);

    auto matmul_attributes =
        fe::graph::Matmul_attributes().set_name("GEMM").set_compute_data_type(fe::DataType_t::FLOAT);

    auto tensor_d = graph.matmul(dequan_tensor_a, dequan_tensor_b, matmul_attributes);

    tensor_d->set_data_type(datatype_d);
    tensor_d->set_is_virtual(false);

    std::cout << graph << std::endl;

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

    // Create a unique_ptr for the cuDNN handle
    auto handle_ptr = create_cudnn_handle();
    auto handle     = *handle_ptr;

    REQUIRE(graph.build_operation_graph(handle).is_good());
    REQUIRE(graph.create_execution_plans({fe::HeurMode_t::A}).is_good());

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

    REQUIRE(graph.build_plans(fe::BuildPlanPolicy_t::HEURISTICS_CHOICE).is_good());

    int64_t workspace_size = 0;
    REQUIRE(graph.get_workspace_size(workspace_size).is_good());
    Surface<int8_t> workspace(workspace_size);

    std::unordered_map<std::shared_ptr<fe::graph::Tensor_attributes>, void*> variant_pack = {
        {tensor_a, tensor_a_gpu.devPtr},
        {tensor_b, tensor_b_gpu.devPtr},
        {block_descale_a, block_descale_a_gpu.devPtr},
        {block_descale_b, block_descale_b_gpu.devPtr},
        {tensor_d, tensor_d_gpu.devPtr}};

    REQUIRE(graph.execute(handle, variant_pack, workspace.devPtr).is_good());
}

TEST_CASE("Blackwell Block Scale Matmul Quantize", "[matmul][graph][FP4]") {
#if (CUDNN_VERSION < 91400)
    SKIP("Matmul with block scaling quantize is not supported in cudnn versions prior to 9.14.0");
#endif
    if (get_compute_capability() < 100 || get_compute_capability() >= 110) {
        SKIP("Matmul with block scaling quantize is only supported on compute capability 100-109");
    }

    auto test_params = GENERATE(TestParams(1,
                                           137,
                                           272,
                                           256,
                                           16,
                                           cudnn_frontend::DataType_t::FP4_E2M1,
                                           cudnn_frontend::DataType_t::FP4_E2M1,
                                           cudnn_frontend::DataType_t::FP8_E4M3,
                                           cudnn_frontend::DataType_t::FP4_E2M1));

    int64_t const b = test_params.b;
    int64_t const m = test_params.m;
    int64_t const n = test_params.n;
    int64_t const k = test_params.k;

    auto datatype_a = test_params.datatype_a;
    auto datatype_b = test_params.datatype_b;

    auto datatype_block_scale = test_params.datatype_block_scale;

    auto block_size = test_params.block_size;

    auto datatype_d = test_params.datatype_d;

    Surface<int8_t> tensor_a_gpu(div_up(b * m * k * cudnn_frontend::detail::get_element_size_in_bits(datatype_a), 8));
    Surface<int8_t> tensor_b_gpu(div_up(b * k * n * cudnn_frontend::detail::get_element_size_in_bits(datatype_b), 8));

    static constexpr int indestructible_128x4_block_m_n = 128;
    static constexpr int indestructible_128x4_block_k   = 4;

    int64_t block_scale_dim_m = div_up(m, indestructible_128x4_block_m_n) * indestructible_128x4_block_m_n;
    int64_t block_scale_dim_n = div_up(n, indestructible_128x4_block_m_n) * indestructible_128x4_block_m_n;
    int64_t block_scale_dim_k =
        div_up(div_up(k, block_size), indestructible_128x4_block_k) * indestructible_128x4_block_k;
    int64_t block_scale_dim_out_m = block_scale_dim_m;
    int64_t block_scale_dim_out_n =
        div_up(div_up(n, block_size), indestructible_128x4_block_k) * indestructible_128x4_block_k;

    Surface<int8_t> block_descale_a_gpu(
        div_up(b * block_scale_dim_m * block_scale_dim_k *
                   cudnn_frontend::detail::get_element_size_in_bits(datatype_block_scale),
               8));
    Surface<int8_t> block_descale_b_gpu(
        div_up(b * block_scale_dim_n * block_scale_dim_k *
                   cudnn_frontend::detail::get_element_size_in_bits(datatype_block_scale),
               8));

    Surface<int8_t> tensor_d_gpu(div_up(b * m * n * cudnn_frontend::detail::get_element_size_in_bits(datatype_d), 8));

    Surface<int8_t> block_scale_gpu(div_up(b * block_scale_dim_out_m * block_scale_dim_out_n *
                                               cudnn_frontend::detail::get_element_size_in_bits(datatype_block_scale),
                                           8));

    fe::graph::Graph graph{};

    graph.set_intermediate_data_type(fe::DataType_t::FLOAT);
    graph.set_compute_data_type(fe::DataType_t::FLOAT);

    auto tensor_a = graph.tensor(fe::graph::Tensor_attributes()
                                     .set_name("tensor_a")
                                     .set_data_type(datatype_a)
                                     .set_dim({b, m, k})
                                     .set_stride({m * k, k, 1}));

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

    auto block_descale_a = graph.tensor(fe::graph::Tensor_attributes()
                                            .set_name("block_descale_a")
                                            .set_data_type(datatype_block_scale)
                                            .set_dim({b, block_scale_dim_m, block_scale_dim_k})
                                            .set_stride({block_scale_dim_m * block_scale_dim_k, block_scale_dim_k, 1})
                                            .set_reordering_type(cudnn_frontend::TensorReordering_t::F8_128x4));

    auto block_descale_b = graph.tensor(fe::graph::Tensor_attributes()
                                            .set_name("block_descale_b")
                                            .set_data_type(datatype_block_scale)
                                            .set_dim({b, block_scale_dim_k, block_scale_dim_n})
                                            .set_stride({block_scale_dim_n * block_scale_dim_k, 1, block_scale_dim_k})
                                            .set_reordering_type(cudnn_frontend::TensorReordering_t::F8_128x4));

    auto dequantize_attr_a = fe::graph::Block_scale_dequantize_attributes().set_block_size({1, block_size});

    auto dequan_tensor_a = graph.block_scale_dequantize(tensor_a, block_descale_a, dequantize_attr_a);

    auto dequantize_attr_b = fe::graph::Block_scale_dequantize_attributes().set_block_size({block_size, 1});

    auto dequan_tensor_b = graph.block_scale_dequantize(tensor_b, block_descale_b, dequantize_attr_b);

    auto matmul_attributes =
        fe::graph::Matmul_attributes().set_name("GEMM").set_compute_data_type(fe::DataType_t::FLOAT);

    auto tensor_c = graph.matmul(dequan_tensor_a, dequan_tensor_b, matmul_attributes);

    auto quantize_attr =
        fe::graph::Block_scale_quantize_attributes().set_block_size(block_size).set_axis(2).set_transpose(false);

    auto [tensor_d, block_scale] = graph.block_scale_quantize(tensor_c, quantize_attr);

    tensor_d->set_output(true).set_data_type(datatype_d);
    block_scale->set_output(true)
        .set_dim({b, block_scale_dim_out_m, block_scale_dim_out_n})
        .set_stride({block_scale_dim_out_m * block_scale_dim_out_n, block_scale_dim_out_n, 1})
        .set_data_type(datatype_block_scale)
        .set_reordering_type(cudnn_frontend::TensorReordering_t::F8_128x4);

    std::cout << graph << std::endl;

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

    // Create a unique_ptr for the cuDNN handle
    auto handle_ptr = create_cudnn_handle();
    auto handle     = *handle_ptr;

    REQUIRE(graph.build_operation_graph(handle).is_good());
    REQUIRE(graph.create_execution_plans({fe::HeurMode_t::A}).is_good());

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

    REQUIRE(graph.build_plans(fe::BuildPlanPolicy_t::HEURISTICS_CHOICE).is_good());

    int64_t workspace_size = 0;
    REQUIRE(graph.get_workspace_size(workspace_size).is_good());
    Surface<int8_t> workspace(workspace_size);

    std::unordered_map<std::shared_ptr<fe::graph::Tensor_attributes>, void*> variant_pack = {
        {tensor_a, tensor_a_gpu.devPtr},
        {tensor_b, tensor_b_gpu.devPtr},
        {block_descale_a, block_descale_a_gpu.devPtr},
        {block_descale_b, block_descale_b_gpu.devPtr},
        {tensor_d, tensor_d_gpu.devPtr},
        {block_scale, block_scale_gpu.devPtr}};

    REQUIRE(graph.execute(handle, variant_pack, workspace.devPtr).is_good());
}

TEST_CASE("Block Scale Matmul Swiglu", "[matmul][graph][FP4]") {
#if (CUDNN_VERSION < 91500)
    SKIP("Block Scale Matmul with Swiglu epilogue fusion is not supported in cudnn versions prior to 9.15.0");
#endif

    if (check_device_arch_newer_than("blackwell") == false) {
        SKIP("Hardware accelerated NVFP4/MXFP8 block scale matmul requires Blackwell and up");
    }

    auto test_params = GENERATE(TestParams(1,
                                           256,
                                           256,
                                           256,
                                           16,
                                           cudnn_frontend::DataType_t::FP4_E2M1,
                                           cudnn_frontend::DataType_t::FP4_E2M1,
                                           cudnn_frontend::DataType_t::FP8_E4M3,
                                           cudnn_frontend::DataType_t::FP4_E2M1));

    int64_t const b = test_params.b;
    int64_t const m = test_params.m;
    int64_t const n = test_params.n;
    int64_t const k = test_params.k;

    auto datatype_a = test_params.datatype_a;
    auto datatype_b = test_params.datatype_b;

    auto datatype_block_scale = test_params.datatype_block_scale;

    auto block_size = test_params.block_size;

    auto datatype_d = test_params.datatype_d;

    Surface<int8_t> tensor_a_gpu(div_up(b * m * k * cudnn_frontend::detail::get_element_size_in_bits(datatype_a), 8));
    Surface<int8_t> tensor_b0_gpu(div_up(b * k * n * cudnn_frontend::detail::get_element_size_in_bits(datatype_b), 8));
    Surface<int8_t> tensor_b1_gpu(div_up(b * k * n * cudnn_frontend::detail::get_element_size_in_bits(datatype_b), 8));

    static constexpr int indestructible_128x4_block_m_n = 128;
    static constexpr int indestructible_128x4_block_k   = 4;

    int64_t block_scale_dim_m = div_up(m, indestructible_128x4_block_m_n) * indestructible_128x4_block_m_n;
    int64_t block_scale_dim_n = div_up(n, indestructible_128x4_block_m_n) * indestructible_128x4_block_m_n;
    int64_t block_scale_dim_k =
        div_up(div_up(k, block_size), indestructible_128x4_block_k) * indestructible_128x4_block_k;

    Surface<int8_t> block_descale_a_gpu(
        div_up(b * block_scale_dim_m * block_scale_dim_k *
                   cudnn_frontend::detail::get_element_size_in_bits(datatype_block_scale),
               8));
    Surface<int8_t> block_descale_b0_gpu(
        div_up(b * block_scale_dim_n * block_scale_dim_k *
                   cudnn_frontend::detail::get_element_size_in_bits(datatype_block_scale),
               8));
    Surface<int8_t> block_descale_b1_gpu(
        div_up(b * block_scale_dim_n * block_scale_dim_k *
                   cudnn_frontend::detail::get_element_size_in_bits(datatype_block_scale),
               8));

    Surface<int8_t> tensor_d_gpu(div_up(b * m * n * cudnn_frontend::detail::get_element_size_in_bits(datatype_d), 8));

    fe::graph::Graph graph{};

    graph.set_intermediate_data_type(fe::DataType_t::FLOAT);
    graph.set_compute_data_type(fe::DataType_t::FLOAT);

    auto tensor_a = graph.tensor(fe::graph::Tensor_attributes()
                                     .set_name("tensor_a")
                                     .set_data_type(datatype_a)
                                     .set_dim({b, m, k})
                                     .set_stride({m * k, k, 1}));

    auto tensor_b0 = graph.tensor(fe::graph::Tensor_attributes()
                                      .set_name("tensor_b")
                                      .set_data_type(datatype_b)
                                      .set_dim({b, k, n})
                                      .set_stride({k * n, 1, k}));

    auto tensor_b1 = graph.tensor(fe::graph::Tensor_attributes()
                                      .set_name("tensor_b")
                                      .set_data_type(datatype_b)
                                      .set_dim({b, k, n})
                                      .set_stride({k * n, 1, k}));

    auto block_descale_a = graph.tensor(fe::graph::Tensor_attributes()
                                            .set_name("block_descale_a")
                                            .set_data_type(datatype_block_scale)
                                            .set_dim({b, block_scale_dim_m, block_scale_dim_k})
                                            .set_stride({block_scale_dim_m * block_scale_dim_k, block_scale_dim_k, 1})
                                            .set_reordering_type(cudnn_frontend::TensorReordering_t::F8_128x4));

    auto block_descale_b0 = graph.tensor(fe::graph::Tensor_attributes()
                                             .set_name("block_descale_b")
                                             .set_data_type(datatype_block_scale)
                                             .set_dim({b, block_scale_dim_k, block_scale_dim_n})
                                             .set_stride({block_scale_dim_n * block_scale_dim_k, 1, block_scale_dim_k})
                                             .set_reordering_type(cudnn_frontend::TensorReordering_t::F8_128x4));

    auto block_descale_b1 = graph.tensor(fe::graph::Tensor_attributes()
                                             .set_name("block_descale_b")
                                             .set_data_type(datatype_block_scale)
                                             .set_dim({b, block_scale_dim_k, block_scale_dim_n})
                                             .set_stride({block_scale_dim_n * block_scale_dim_k, 1, block_scale_dim_k})
                                             .set_reordering_type(cudnn_frontend::TensorReordering_t::F8_128x4));

    auto dequantize_attr_a = fe::graph::Block_scale_dequantize_attributes().set_block_size({1, block_size});

    auto dequan_tensor_a = graph.block_scale_dequantize(tensor_a, block_descale_a, dequantize_attr_a);

    auto dequantize_attr_b0 = fe::graph::Block_scale_dequantize_attributes().set_block_size({block_size, 1});

    auto dequan_tensor_b0 = graph.block_scale_dequantize(tensor_b0, block_descale_b0, dequantize_attr_b0);

    auto dequantize_attr_b1 = fe::graph::Block_scale_dequantize_attributes().set_block_size({block_size, 1});

    auto dequan_tensor_b1 = graph.block_scale_dequantize(tensor_b1, block_descale_b1, dequantize_attr_b1);

    auto matmul0_attributes =
        fe::graph::Matmul_attributes().set_name("GEMM").set_compute_data_type(fe::DataType_t::FLOAT);

    auto matmul1_attributes =
        fe::graph::Matmul_attributes().set_name("GEMM").set_compute_data_type(fe::DataType_t::FLOAT);

    auto tensor_c0 = graph.matmul(dequan_tensor_a, dequan_tensor_b0, matmul0_attributes);

    auto tensor_c1 = graph.matmul(dequan_tensor_a, dequan_tensor_b1, matmul1_attributes);

    auto swish_attributes = fe::graph::Pointwise_attributes()
                                .set_name("swish")
                                .set_mode(fe::PointwiseMode_t::SWISH_FWD)
                                .set_compute_data_type(fe::DataType_t::FLOAT);

    auto tensor_after_swish = graph.pointwise(tensor_c0, swish_attributes);

    auto mul_attributes = fe::graph::Pointwise_attributes()
                              .set_name("mul")
                              .set_mode(fe::PointwiseMode_t::MUL)
                              .set_compute_data_type(fe::DataType_t::FLOAT);

    auto tensor_d = graph.pointwise(tensor_after_swish, tensor_c1, mul_attributes);

    tensor_d->set_output(true).set_data_type(datatype_d);

    std::cout << graph << std::endl;

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

    // Create a unique_ptr for the cuDNN handle
    auto handle_ptr = create_cudnn_handle();
    auto handle     = *handle_ptr;

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

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

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

    REQUIRE(graph.build_plans(fe::BuildPlanPolicy_t::HEURISTICS_CHOICE).is_good());

    int64_t workspace_size = 0;
    REQUIRE(graph.get_workspace_size(workspace_size).is_good());
    Surface<int8_t> workspace(workspace_size);

    std::unordered_map<std::shared_ptr<fe::graph::Tensor_attributes>, void*> variant_pack = {
        {tensor_a, tensor_a_gpu.devPtr},
        {tensor_b0, tensor_b0_gpu.devPtr},
        {tensor_b1, tensor_b1_gpu.devPtr},
        {block_descale_a, block_descale_a_gpu.devPtr},
        {block_descale_b0, block_descale_b0_gpu.devPtr},
        {block_descale_b1, block_descale_b1_gpu.devPtr},
        {tensor_d, tensor_d_gpu.devPtr}};

    REQUIRE(graph.execute(handle, variant_pack, workspace.devPtr).is_good());
}

TEST_CASE("Blackwell Block Scale Matmul dynamic shape overrides", "[matmul][graph][dynamic_shape]") {
#if (CUDNN_VERSION < 92100)
    SKIP("Dynamic shape with overrides is not supported in cudnn versions prior to 9.21.0");
#endif

    if (get_compute_capability() < 100 || get_compute_capability() >= 110) {
        SKIP("Hardware accelerated NVFP4/MXFP8 block scale matmul requires Blackwell and up");
    }

    namespace fe = cudnn_frontend;

    constexpr int A_UID    = 1;
    constexpr int SF_A_UID = 2;
    constexpr int B_UID    = 3;
    constexpr int SF_B_UID = 4;
    constexpr int C_UID    = 5;

    static constexpr int indestructible_128x4_block_m_n = 128;
    static constexpr int indestructible_128x4_block_k   = 4;

    int block_size = 16;

    struct matmul_shapes {
        int64_t b, m, n, k;
    };

    matmul_shapes matmul_cache_shape     = {1, 1024, 1024, 1024};
    matmul_shapes matmul_dynamic_shape[] = {
        {2, 1024, 1024, 1024},
        {2, 2048, 2048, 2048},
    };

    constexpr int matmul_dynamic_shape_count = sizeof(matmul_dynamic_shape) / sizeof(matmul_cache_shape);

    auto handle_ptr = create_cudnn_handle();
    auto handle     = *handle_ptr;

    // build graph and execution plan with a fake shape
    auto graph = std::make_shared<fe::graph::Graph>();

    graph->set_intermediate_data_type(fe::DataType_t::FLOAT)
        .set_compute_data_type(fe::DataType_t::FLOAT)
        .set_override_shape_enabled(true);

    int64_t block_scale_dim_m_cache =
        div_up(matmul_cache_shape.m, indestructible_128x4_block_m_n) * indestructible_128x4_block_m_n;
    int64_t block_scale_dim_n_cache =
        div_up(matmul_cache_shape.n, indestructible_128x4_block_m_n) * indestructible_128x4_block_m_n;
    int64_t block_scale_dim_k_cache =
        div_up(div_up(matmul_cache_shape.k, block_size), indestructible_128x4_block_k) * indestructible_128x4_block_k;

    auto A = graph->tensor(fe::graph::Tensor_attributes()
                               .set_name("A")
                               .set_uid(A_UID)
                               .set_dim({matmul_cache_shape.b, matmul_cache_shape.m, matmul_cache_shape.k})
                               .set_stride({matmul_cache_shape.m * matmul_cache_shape.k, matmul_cache_shape.k, 1})
                               .set_data_type(fe::DataType_t::FP4_E2M1));

    auto SF_A =
        graph->tensor(fe::graph::Tensor_attributes()
                          .set_name("SF_A")
                          .set_uid(SF_A_UID)
                          .set_dim({matmul_cache_shape.b, block_scale_dim_m_cache, block_scale_dim_k_cache})
                          .set_stride({block_scale_dim_m_cache * block_scale_dim_k_cache, block_scale_dim_k_cache, 1})
                          .set_data_type(fe::DataType_t::FP8_E4M3)
                          .set_reordering_type(cudnn_frontend::TensorReordering_t::F8_128x4));

    auto dequantize_attr_a = fe::graph::Block_scale_dequantize_attributes().set_block_size({1, block_size});
    auto dequan_tensor_a   = graph->block_scale_dequantize(A, SF_A, dequantize_attr_a);

    auto B = graph->tensor(fe::graph::Tensor_attributes()
                               .set_name("B")
                               .set_uid(B_UID)
                               .set_dim({matmul_cache_shape.b, matmul_cache_shape.k, matmul_cache_shape.n})
                               .set_stride({matmul_cache_shape.n * matmul_cache_shape.k, 1, matmul_cache_shape.k})
                               .set_data_type(fe::DataType_t::FP4_E2M1));

    auto SF_B =
        graph->tensor(fe::graph::Tensor_attributes()
                          .set_name("SF_B")
                          .set_uid(SF_B_UID)
                          .set_dim({matmul_cache_shape.b, block_scale_dim_k_cache, block_scale_dim_n_cache})
                          .set_stride({block_scale_dim_n_cache * block_scale_dim_k_cache, 1, block_scale_dim_k_cache})
                          .set_data_type(fe::DataType_t::FP8_E4M3)
                          .set_reordering_type(cudnn_frontend::TensorReordering_t::F8_128x4));

    auto dequantize_attr_b = fe::graph::Block_scale_dequantize_attributes().set_block_size({block_size, 1});
    auto dequan_tensor_b   = graph->block_scale_dequantize(B, SF_B, dequantize_attr_b);

    auto C = graph->matmul(
        dequan_tensor_a, dequan_tensor_b, fe::graph::Matmul_attributes().set_compute_data_type(fe::DataType_t::FLOAT));
    C->set_uid(C_UID).set_output(true).set_data_type(fe::DataType_t::BFLOAT16);

    // For dynamic shape, recommend to query fallback plan to get a general good performance
    // Heuristics Mode A is recommended if the dynamic problem shapes are similar in size
    REQUIRE(graph->build(handle, {fe::HeurMode_t::A, fe::HeurMode_t::FALLBACK}).is_good());

    // run graph with dynamic shapes
    for (int idx_shape = 0; idx_shape < matmul_dynamic_shape_count; ++idx_shape) {
        int64_t block_scale_dim_m =
            div_up(matmul_dynamic_shape[idx_shape].m, indestructible_128x4_block_m_n) * indestructible_128x4_block_m_n;
        int64_t block_scale_dim_n =
            div_up(matmul_dynamic_shape[idx_shape].n, indestructible_128x4_block_m_n) * indestructible_128x4_block_m_n;
        int64_t block_scale_dim_k =
            div_up(div_up(matmul_dynamic_shape[idx_shape].k, block_size), indestructible_128x4_block_k) *
            indestructible_128x4_block_k;

        std::vector<int64_t> override_uids                = {A_UID, SF_A_UID, B_UID, SF_B_UID, C_UID};
        std::vector<std::vector<int64_t>> override_shapes = {
            {matmul_dynamic_shape[idx_shape].b, matmul_dynamic_shape[idx_shape].m, matmul_dynamic_shape[idx_shape].k},
            {matmul_dynamic_shape[idx_shape].b, block_scale_dim_m, block_scale_dim_k},
            {matmul_dynamic_shape[idx_shape].b, matmul_dynamic_shape[idx_shape].k, matmul_dynamic_shape[idx_shape].n},
            {matmul_dynamic_shape[idx_shape].b, block_scale_dim_k, block_scale_dim_n},
            {matmul_dynamic_shape[idx_shape].b, matmul_dynamic_shape[idx_shape].m, matmul_dynamic_shape[idx_shape].n}};
        std::vector<std::vector<int64_t>> override_strides = {
            {matmul_dynamic_shape[idx_shape].m * matmul_dynamic_shape[idx_shape].k,
             matmul_dynamic_shape[idx_shape].k,
             1},
            {block_scale_dim_m * block_scale_dim_k, block_scale_dim_k, 1},
            {matmul_dynamic_shape[idx_shape].n * matmul_dynamic_shape[idx_shape].k,
             1,
             matmul_dynamic_shape[idx_shape].k},
            {block_scale_dim_n * block_scale_dim_k, 1, block_scale_dim_k},
            {matmul_dynamic_shape[idx_shape].m * matmul_dynamic_shape[idx_shape].n,
             matmul_dynamic_shape[idx_shape].n,
             1}};

        Surface<int8_t> A_gpu(div_up(matmul_dynamic_shape[idx_shape].b * matmul_dynamic_shape[idx_shape].m *
                                         matmul_dynamic_shape[idx_shape].k *
                                         cudnn_frontend::detail::get_element_size_in_bits(fe::DataType_t::FP4_E2M1),
                                     8));
        Surface<int8_t> SF_A_gpu(div_up(matmul_dynamic_shape[idx_shape].b * block_scale_dim_m * block_scale_dim_k *
                                            cudnn_frontend::detail::get_element_size_in_bits(fe::DataType_t::FP8_E4M3),
                                        8));
        Surface<int8_t> B_gpu(div_up(matmul_dynamic_shape[idx_shape].b * matmul_dynamic_shape[idx_shape].k *
                                         matmul_dynamic_shape[idx_shape].n *
                                         cudnn_frontend::detail::get_element_size_in_bits(fe::DataType_t::FP4_E2M1),
                                     8));
        Surface<int8_t> SF_B_gpu(div_up(matmul_dynamic_shape[idx_shape].b * block_scale_dim_k * block_scale_dim_n *
                                            cudnn_frontend::detail::get_element_size_in_bits(fe::DataType_t::FP8_E4M3),
                                        8));
        Surface<int8_t> C_gpu(div_up(matmul_dynamic_shape[idx_shape].b * matmul_dynamic_shape[idx_shape].m *
                                         matmul_dynamic_shape[idx_shape].n *
                                         cudnn_frontend::detail::get_element_size_in_bits(fe::DataType_t::BFLOAT16),
                                     8));

        std::unordered_map<fe::graph::Tensor_attributes::uid_t, void*> variant_pack = {{A_UID, A_gpu.devPtr},
                                                                                       {SF_A_UID, SF_A_gpu.devPtr},
                                                                                       {B_UID, B_gpu.devPtr},
                                                                                       {SF_B_UID, SF_B_gpu.devPtr},
                                                                                       {C_UID, C_gpu.devPtr}};

        int64_t workspace_size = 0;
        if (cudnn_frontend::detail::get_backend_version() >= 92300) {
            REQUIRE(graph->get_workspace_size(handle, workspace_size, override_uids, override_shapes, override_strides)
                        .is_good());
        } else {
            REQUIRE(graph->get_workspace_size(workspace_size).is_good());
        }
        Surface<int8_t> workspace(workspace_size);

        REQUIRE(graph->execute(handle, variant_pack, workspace.devPtr, override_uids, override_shapes, override_strides)
                    .is_good());

        CUDA_CHECK(cudaDeviceSynchronize());
    }
}

}  // namespace BlackwellNVFP4MXFP8BlockScaleMatmul