// Copyright © 2023-2025 Advanced Micro Devices, Inc.
// SPDX-License-Identifier: MIT

#include "../shim.bwd_kernel_dq.h"
#include <aotriton/_internal/triton_kernel.h>
#include <aotriton/_internal/kernel_cluster.h>
#include <aotriton/cpp_tune.h>
#include <string_view>
#ifndef NDEBUG
#include <iostream>
#endif

#define CURRENT_ENTRY_PUBLIC Autotune_bwd_kernel_dq__A9__F218

#define ARRAY_SIZE(array)  (sizeof(array) / sizeof(array[0]))

namespace { // Anonymous namespace

using namespace std::literals::string_view_literals;

#if AOTRITON_BUILD_FOR_TUNING
// PSels and Copts in JSON String
static const char* kernel_psels[] = {
  R"xyzw({"BLOCK_M": 16, "BLOCK_N": 16})xyzw",
  R"xyzw({"BLOCK_M": 16, "BLOCK_N": 16})xyzw",
  R"xyzw({"BLOCK_M": 16, "BLOCK_N": 16})xyzw",
  R"xyzw({"BLOCK_M": 32, "BLOCK_N": 16})xyzw",
  R"xyzw({"BLOCK_M": 32, "BLOCK_N": 16})xyzw",
  R"xyzw({"BLOCK_M": 32, "BLOCK_N": 16})xyzw",
  R"xyzw({"BLOCK_M": 32, "BLOCK_N": 16})xyzw",
  R"xyzw({"BLOCK_M": 32, "BLOCK_N": 16})xyzw",
  R"xyzw({"BLOCK_M": 32, "BLOCK_N": 16})xyzw",
  R"xyzw({"BLOCK_M": 32, "BLOCK_N": 16})xyzw",
  R"xyzw({"BLOCK_M": 32, "BLOCK_N": 32})xyzw",
  R"xyzw({"BLOCK_M": 32, "BLOCK_N": 32})xyzw",
  R"xyzw({"BLOCK_M": 32, "BLOCK_N": 32})xyzw",
  R"xyzw({"BLOCK_M": 32, "BLOCK_N": 32})xyzw",
  R"xyzw({"BLOCK_M": 64, "BLOCK_N": 16})xyzw",
  R"xyzw({"BLOCK_M": 64, "BLOCK_N": 16})xyzw",
  R"xyzw({"BLOCK_M": 64, "BLOCK_N": 16})xyzw",
  R"xyzw({"BLOCK_M": 64, "BLOCK_N": 16})xyzw",
  R"xyzw({"BLOCK_M": 64, "BLOCK_N": 32})xyzw",
  R"xyzw({"BLOCK_M": 64, "BLOCK_N": 32})xyzw",
  R"xyzw({"BLOCK_M": 64, "BLOCK_N": 32})xyzw",
  R"xyzw({"BLOCK_M": 64, "BLOCK_N": 32})xyzw",
  R"xyzw({"BLOCK_M": 64, "BLOCK_N": 64})xyzw",
  R"xyzw({"BLOCK_M": 64, "BLOCK_N": 64})xyzw",
  R"xyzw({"BLOCK_M": 64, "BLOCK_N": 64})xyzw",
  R"xyzw({"BLOCK_M": 64, "BLOCK_N": 64})xyzw"
};
static const char* kernel_copts[] = {
  R"xyzw({"waves_per_eu": 1, "num_warps": 2, "num_stages": 1})xyzw",
  R"xyzw({"waves_per_eu": 3, "num_warps": 2, "num_stages": 1})xyzw",
  R"xyzw({"waves_per_eu": 4, "num_warps": 2, "num_stages": 1})xyzw",
  R"xyzw({"waves_per_eu": 1, "num_warps": 2, "num_stages": 1})xyzw",
  R"xyzw({"waves_per_eu": 1, "num_warps": 4, "num_stages": 1})xyzw",
  R"xyzw({"waves_per_eu": 2, "num_warps": 2, "num_stages": 1})xyzw",
  R"xyzw({"waves_per_eu": 3, "num_warps": 2, "num_stages": 1})xyzw",
  R"xyzw({"waves_per_eu": 3, "num_warps": 4, "num_stages": 1})xyzw",
  R"xyzw({"waves_per_eu": 4, "num_warps": 2, "num_stages": 1})xyzw",
  R"xyzw({"waves_per_eu": 4, "num_warps": 4, "num_stages": 1})xyzw",
  R"xyzw({"waves_per_eu": 1, "num_warps": 4, "num_stages": 1})xyzw",
  R"xyzw({"waves_per_eu": 2, "num_warps": 4, "num_stages": 1})xyzw",
  R"xyzw({"waves_per_eu": 3, "num_warps": 4, "num_stages": 1})xyzw",
  R"xyzw({"waves_per_eu": 4, "num_warps": 4, "num_stages": 1})xyzw",
  R"xyzw({"waves_per_eu": 1, "num_warps": 4, "num_stages": 1})xyzw",
  R"xyzw({"waves_per_eu": 2, "num_warps": 4, "num_stages": 1})xyzw",
  R"xyzw({"waves_per_eu": 3, "num_warps": 4, "num_stages": 1})xyzw",
  R"xyzw({"waves_per_eu": 4, "num_warps": 4, "num_stages": 1})xyzw",
  R"xyzw({"waves_per_eu": 1, "num_warps": 4, "num_stages": 1})xyzw",
  R"xyzw({"waves_per_eu": 2, "num_warps": 4, "num_stages": 1})xyzw",
  R"xyzw({"waves_per_eu": 3, "num_warps": 4, "num_stages": 1})xyzw",
  R"xyzw({"waves_per_eu": 4, "num_warps": 4, "num_stages": 1})xyzw",
  R"xyzw({"waves_per_eu": 1, "num_warps": 4, "num_stages": 1})xyzw",
  R"xyzw({"waves_per_eu": 2, "num_warps": 4, "num_stages": 1})xyzw",
  R"xyzw({"waves_per_eu": 3, "num_warps": 4, "num_stages": 1})xyzw",
  R"xyzw({"waves_per_eu": 4, "num_warps": 4, "num_stages": 1})xyzw"
};
#endif

struct PerfFields {
  int16_t BLOCK_M;
    int16_t BLOCK_N;
};

static PerfFields image_perf_list [] = {
    { .BLOCK_M = 16, .BLOCK_N = 16 },
    { .BLOCK_M = 16, .BLOCK_N = 16 },
    { .BLOCK_M = 16, .BLOCK_N = 16 },
    { .BLOCK_M = 32, .BLOCK_N = 16 },
    { .BLOCK_M = 32, .BLOCK_N = 16 },
    { .BLOCK_M = 32, .BLOCK_N = 16 },
    { .BLOCK_M = 32, .BLOCK_N = 16 },
    { .BLOCK_M = 32, .BLOCK_N = 16 },
    { .BLOCK_M = 32, .BLOCK_N = 16 },
    { .BLOCK_M = 32, .BLOCK_N = 16 },
    { .BLOCK_M = 32, .BLOCK_N = 32 },
    { .BLOCK_M = 32, .BLOCK_N = 32 },
    { .BLOCK_M = 32, .BLOCK_N = 32 },
    { .BLOCK_M = 32, .BLOCK_N = 32 },
    { .BLOCK_M = 64, .BLOCK_N = 16 },
    { .BLOCK_M = 64, .BLOCK_N = 16 },
    { .BLOCK_M = 64, .BLOCK_N = 16 },
    { .BLOCK_M = 64, .BLOCK_N = 16 },
    { .BLOCK_M = 64, .BLOCK_N = 32 },
    { .BLOCK_M = 64, .BLOCK_N = 32 },
    { .BLOCK_M = 64, .BLOCK_N = 32 },
    { .BLOCK_M = 64, .BLOCK_N = 32 },
    { .BLOCK_M = 64, .BLOCK_N = 64 },
    { .BLOCK_M = 64, .BLOCK_N = 64 },
    { .BLOCK_M = 64, .BLOCK_N = 64 },
    { .BLOCK_M = 64, .BLOCK_N = 64 }
};

// u8R generates char8_t which is poorly supported almost everywhere.
constexpr pstring_view PACKAGE_PATH
#if defined(_WIN32)
{ LR"xyzw(amd-gfx120x/flash/bwd_kernel_dq/FONLY__＊bf16@16_32_3_F_T_0___gfx120x)xyzw" };
#else
{ R"xyzw(amd-gfx120x/flash/bwd_kernel_dq/FONLY__＊bf16@16_32_3_F_T_0___gfx120x)xyzw" };
#endif
constexpr std::string_view FUNC_NAME { R"xyzw(＊bf16@16_32_3_F_T_0)xyzw" };
constexpr std::string_view ARCH_NAME { R"xyzw(gfx1201)xyzw" };

// Checksum can be confirmed with `echo -n '<string>' | b2sum -l 64`
// For example:
//   $ echo -n 'amd-gfx110x/flash/attn_fwd/FONLY__^bf16@16,128,False,False,0,0___gfx1100__P__32_32_0_2_False__CO__wave3_warp2_stg1-Gpu-gfx1100' | b2sum -l 64
//   c4b51ee645d79580  -
static AOTRITON_NS::TritonKernelCompactMeta meta_list[] = {
    { 0xdd87b678u, 0x9f9026bfu, 0, 46 }, // dd87b6789f9026bf = b2sum -l 64 <<< amd-gfx120x/flash/bwd_kernel_dq/FONLY__＊bf16@16_32_3_F_T_0___gfx120x__P__16_16__CO__wave1_warp2_stg1--Arch_gfx1201
    { 0x2ce3c275u, 0x6a6ba6c3u, 0, 6 }, // 2ce3c2756a6ba6c3 = b2sum -l 64 <<< amd-gfx120x/flash/bwd_kernel_dq/FONLY__＊bf16@16_32_3_F_T_0___gfx120x__P__16_16__CO__wave3_warp2_stg1--Arch_gfx1201
    { 0x2e31cc44u, 0xbc304c53u, 0, 23 }, // 2e31cc44bc304c53 = b2sum -l 64 <<< amd-gfx120x/flash/bwd_kernel_dq/FONLY__＊bf16@16_32_3_F_T_0___gfx120x__P__16_16__CO__wave4_warp2_stg1--Arch_gfx1201
    { 0xa863397fu, 0x4e8a558fu, 40, 46 }, // a863397f4e8a558f = b2sum -l 64 <<< amd-gfx120x/flash/bwd_kernel_dq/FONLY__＊bf16@16_32_3_F_T_0___gfx120x__P__32_16__CO__wave1_warp2_stg1--Arch_gfx1201
    { 0x8245c6d2u, 0x244c928bu, 40, 138 }, // 8245c6d2244c928b = b2sum -l 64 <<< amd-gfx120x/flash/bwd_kernel_dq/FONLY__＊bf16@16_32_3_F_T_0___gfx120x__P__32_16__CO__wave1_warp4_stg1--Arch_gfx1201
    { 0x591d415bu, 0xd871c326u, 40, 63 }, // 591d415bd871c326 = b2sum -l 64 <<< amd-gfx120x/flash/bwd_kernel_dq/FONLY__＊bf16@16_32_3_F_T_0___gfx120x__P__32_16__CO__wave2_warp2_stg1--Arch_gfx1201
    { 0x5e9a4397u, 0x4c40be34u, 40, 6 }, // 5e9a43974c40be34 = b2sum -l 64 <<< amd-gfx120x/flash/bwd_kernel_dq/FONLY__＊bf16@16_32_3_F_T_0___gfx120x__P__32_16__CO__wave3_warp2_stg1--Arch_gfx1201
    { 0xc03e9f3fu, 0xf0ce7b8fu, 40, 155 }, // c03e9f3ff0ce7b8f = b2sum -l 64 <<< amd-gfx120x/flash/bwd_kernel_dq/FONLY__＊bf16@16_32_3_F_T_0___gfx120x__P__32_16__CO__wave3_warp4_stg1--Arch_gfx1201
    { 0x7bf9c3afu, 0x88f4a3dau, 40, 23 }, // 7bf9c3af88f4a3da = b2sum -l 64 <<< amd-gfx120x/flash/bwd_kernel_dq/FONLY__＊bf16@16_32_3_F_T_0___gfx120x__P__32_16__CO__wave4_warp2_stg1--Arch_gfx1201
    { 0x6fd2a9c0u, 0xbe4d1803u, 40, 103 }, // 6fd2a9c0be4d1803 = b2sum -l 64 <<< amd-gfx120x/flash/bwd_kernel_dq/FONLY__＊bf16@16_32_3_F_T_0___gfx120x__P__32_16__CO__wave4_warp4_stg1--Arch_gfx1201
    { 0xb0730719u, 0xb2776af6u, 97, 138 }, // b0730719b2776af6 = b2sum -l 64 <<< amd-gfx120x/flash/bwd_kernel_dq/FONLY__＊bf16@16_32_3_F_T_0___gfx120x__P__32_32__CO__wave1_warp4_stg1--Arch_gfx1201
    { 0xc454d30du, 0x577e785eu, 97, 80 }, // c454d30d577e785e = b2sum -l 64 <<< amd-gfx120x/flash/bwd_kernel_dq/FONLY__＊bf16@16_32_3_F_T_0___gfx120x__P__32_32__CO__wave2_warp4_stg1--Arch_gfx1201
    { 0x18b2e652u, 0x347a72f8u, 97, 155 }, // 18b2e652347a72f8 = b2sum -l 64 <<< amd-gfx120x/flash/bwd_kernel_dq/FONLY__＊bf16@16_32_3_F_T_0___gfx120x__P__32_32__CO__wave3_warp4_stg1--Arch_gfx1201
    { 0xe6018e25u, 0xe273ef4bu, 97, 103 }, // e6018e25e273ef4b = b2sum -l 64 <<< amd-gfx120x/flash/bwd_kernel_dq/FONLY__＊bf16@16_32_3_F_T_0___gfx120x__P__32_32__CO__wave4_warp4_stg1--Arch_gfx1201
    { 0x75b8c5dbu, 0x232ae15cu, 120, 138 }, // 75b8c5db232ae15c = b2sum -l 64 <<< amd-gfx120x/flash/bwd_kernel_dq/FONLY__＊bf16@16_32_3_F_T_0___gfx120x__P__64_16__CO__wave1_warp4_stg1--Arch_gfx1201
    { 0x11e28826u, 0x41a77778u, 120, 80 }, // 11e2882641a77778 = b2sum -l 64 <<< amd-gfx120x/flash/bwd_kernel_dq/FONLY__＊bf16@16_32_3_F_T_0___gfx120x__P__64_16__CO__wave2_warp4_stg1--Arch_gfx1201
    { 0x7acef473u, 0x795d5d92u, 120, 155 }, // 7acef473795d5d92 = b2sum -l 64 <<< amd-gfx120x/flash/bwd_kernel_dq/FONLY__＊bf16@16_32_3_F_T_0___gfx120x__P__64_16__CO__wave3_warp4_stg1--Arch_gfx1201
    { 0xaf5984f1u, 0xdcab6a5cu, 120, 103 }, // af5984f1dcab6a5c = b2sum -l 64 <<< amd-gfx120x/flash/bwd_kernel_dq/FONLY__＊bf16@16_32_3_F_T_0___gfx120x__P__64_16__CO__wave4_warp4_stg1--Arch_gfx1201
    { 0x8866b852u, 0xa4d1a126u, 126, 138 }, // 8866b852a4d1a126 = b2sum -l 64 <<< amd-gfx120x/flash/bwd_kernel_dq/FONLY__＊bf16@16_32_3_F_T_0___gfx120x__P__64_32__CO__wave1_warp4_stg1--Arch_gfx1201
    { 0x1636f778u, 0x6822217cu, 126, 80 }, // 1636f7786822217c = b2sum -l 64 <<< amd-gfx120x/flash/bwd_kernel_dq/FONLY__＊bf16@16_32_3_F_T_0___gfx120x__P__64_32__CO__wave2_warp4_stg1--Arch_gfx1201
    { 0xd9200e61u, 0x5d522c73u, 126, 155 }, // d9200e615d522c73 = b2sum -l 64 <<< amd-gfx120x/flash/bwd_kernel_dq/FONLY__＊bf16@16_32_3_F_T_0___gfx120x__P__64_32__CO__wave3_warp4_stg1--Arch_gfx1201
    { 0x08b2e9c2u, 0x91260d4au, 126, 103 }, // 08b2e9c291260d4a = b2sum -l 64 <<< amd-gfx120x/flash/bwd_kernel_dq/FONLY__＊bf16@16_32_3_F_T_0___gfx120x__P__64_32__CO__wave4_warp4_stg1--Arch_gfx1201
    { 0x4c251644u, 0xb33b5b3fu, 132, 138 }, // 4c251644b33b5b3f = b2sum -l 64 <<< amd-gfx120x/flash/bwd_kernel_dq/FONLY__＊bf16@16_32_3_F_T_0___gfx120x__P__64_64__CO__wave1_warp4_stg1--Arch_gfx1201
    { 0x4ad53588u, 0xa1b4df21u, 132, 80 }, // 4ad53588a1b4df21 = b2sum -l 64 <<< amd-gfx120x/flash/bwd_kernel_dq/FONLY__＊bf16@16_32_3_F_T_0___gfx120x__P__64_64__CO__wave2_warp4_stg1--Arch_gfx1201
    { 0x5e91c7ffu, 0x909838d5u, 132, 155 }, // 5e91c7ff909838d5 = b2sum -l 64 <<< amd-gfx120x/flash/bwd_kernel_dq/FONLY__＊bf16@16_32_3_F_T_0___gfx120x__P__64_64__CO__wave3_warp4_stg1--Arch_gfx1201
    { 0xa2f9d1d8u, 0xf8265cd1u, 132, 103 }, // a2f9d1d8f8265cd1 = b2sum -l 64 <<< amd-gfx120x/flash/bwd_kernel_dq/FONLY__＊bf16@16_32_3_F_T_0___gfx120x__P__64_64__CO__wave4_warp4_stg1--Arch_gfx1201
};

static constexpr int kTotalNumKernels = ARRAY_SIZE(meta_list);

using AOTRITON_NS::v3::flash::autotune::bwd_kernel_dq_packed_string;

static AOTRITON_NS::TritonKernelCluster<kTotalNumKernels>
kernel_cluster(meta_list,
               bwd_kernel_dq_packed_string);

static int8_t lut[1][10][10] =
{
// GPU gfx1201_mod0
{{ 0, 2, 1, 2, 4, 2, 9, 7, 7, 9},
 { 0,11,10,11,10,12,12,11,12,13},
 { 1,13,11,10,10,11,25,24,23,25},
 { 8, 6, 3,21,19,18,21,18,21,18},
 { 5,18,18,21, 6, 5,25,25,22,24},
 {15, 5, 8, 6,20,19,19,21,18,24},
 {17,16,16,16,19,17,17,21,23,24},
 { 6,14, 8,19,20,16,14,24,24,25},
 { 3, 5, 5,18,14,17,15, 8,25,25},
 { 3,20,15,15,14, 6,15,17,15,14}}
// End of GPU gfx1201_mod0
}
;

}; // End of anonymous namespace

namespace AOTRITON_NS::v3::flash::autotune {

// using AOTRITON_NS::v2::flash::BwdKernelDqContext;

void CURRENT_ENTRY_PUBLIC(BwdKernelDqContext& context, int mod_number) {
#if AOTRITON_BUILD_FOR_TUNING
    int preferred_index = context._has_preferred_kernel;
    context._total_number_of_kernels = kTotalNumKernels;
#ifndef NDEBUG
    std::cerr << "Autotune_bwd_kernel_dq__A9__F218 "
              << "kTotalNumKernels = " << kTotalNumKernels << " "
              << "_has_preferred_kernel = " << preferred_index << " "
              << std::endl;
#endif
    if (preferred_index != -1) {
        if (preferred_index >= kTotalNumKernels)
            return ;
        context.kernel_on_device = kernel_cluster.get(preferred_index);
        context.pp_args_index = 4;
        context.package_path = PACKAGE_PATH;
        context.func_name = FUNC_NAME;
        context.arch_name = ARCH_NAME;
        context._preferred_kernel_psels = kernel_psels[preferred_index];
        context._preferred_kernel_copts = kernel_copts[preferred_index];
        const auto& perf = image_perf_list[preferred_index];
        context.BLOCK_M = perf.BLOCK_M;
    context.BLOCK_N = perf.BLOCK_N;
        return ;
    }
#endif
    auto kernel_index = bwd_kernel_dq__lut_lambda__0(*context.params, mod_number, lut);
    if (kernel_index < 0) {
      return ;
    }
    context.kernel_on_device = kernel_cluster.get(kernel_index);
    context.pp_args_index = 4;
    context.package_path = PACKAGE_PATH;
    context.func_name = FUNC_NAME;
    context.arch_name = ARCH_NAME;
#ifndef NDEBUG
    std::cerr << __FILE__ << " kernel_index = " << int(kernel_index) << std::endl;
#endif
    const auto& perf = image_perf_list[kernel_index];
    context.BLOCK_M = perf.BLOCK_M;
    context.BLOCK_N = perf.BLOCK_N;
}

#undef CURRENT_ENTRY_PUBLIC
}

// Human-readable Signature 
// Q = "*bf16:16"
// sm_scale = "fp32"
// L = "*fp32:16"
// D = "*fp32:16"
// stride_bz = 0
// stride_dbz = 0
// num_head_q = "i32"
// cu_seqlens_q = "*i32:16"
// num_seqlens = "i32"
// dropout_p = 0
// philox_seed_ptr = 0
// philox_offset1 = 0
// philox_offset2 = 0
// Window_left = "i32"
// BLOCK_DMODEL = 32
// CAUSAL_TYPE = 3
// ENABLE_DROPOUT = False
// PADDED_HEAD = True
// BIAS_TYPE = 0

// SELECT * FROM FLASH$bwd_kernel_dq WHERE gpu IN ('gfx1201_mod0') AND inputs$Q_dtype = 'torch.bfloat16' AND inputs$BLOCK_DMODEL = 32 AND inputs$CAUSAL_TYPE = 3 AND inputs$ENABLE_DROPOUT = False AND inputs$PADDED_HEAD = False AND inputs$BIAS_TYPE = 0

// vim: set fileencoding=utf-8

