// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.

// This file is generated by compile_sparse_attention.py using triton AoT compiler

#pragma once
#include "contrib_ops/cuda/sparse/sparse_attention_v1/sparse_attention_common.h"

namespace onnxruntime {
namespace contrib {
namespace cuda {
namespace sparse_attention_v1 {

// launcher for: sparse_attention_bf16_sm80_64x0x64x0x64x2_warps4xstages2
Status sparse_attention_bf16_sm80_31acb592(SparseAttentionParams& params);

Status sparse_attention_bf16_sm80_64x0x64x0x64x2_warps4xstages2(SparseAttentionParams& params) {
  return sparse_attention_bf16_sm80_31acb592(params);
}

// load for: sparse_attention_bf16_sm80_64x0x64x0x64x2_warps4xstages2
void load_sparse_attention_bf16_sm80_31acb592();
void load_sparse_attention_bf16_sm80_64x0x64x0x64x2_warps4xstages2() {
  load_sparse_attention_bf16_sm80_31acb592();
}

// unload for: sparse_attention_bf16_sm80_64x0x64x0x64x2_warps4xstages2
void unload_sparse_attention_bf16_sm80_31acb592();
void unload_sparse_attention_bf16_sm80_64x0x64x0x64x2_warps4xstages2() {
  unload_sparse_attention_bf16_sm80_31acb592();
}

// launcher for: sparse_attention_bf16_sm80_64x1x64x1x64x2_warps4xstages2
Status sparse_attention_bf16_sm80_c777f3f5(SparseAttentionParams& params);

Status sparse_attention_bf16_sm80_64x1x64x1x64x2_warps4xstages2(SparseAttentionParams& params) {
  return sparse_attention_bf16_sm80_c777f3f5(params);
}

// load for: sparse_attention_bf16_sm80_64x1x64x1x64x2_warps4xstages2
void load_sparse_attention_bf16_sm80_c777f3f5();
void load_sparse_attention_bf16_sm80_64x1x64x1x64x2_warps4xstages2() {
  load_sparse_attention_bf16_sm80_c777f3f5();
}

// unload for: sparse_attention_bf16_sm80_64x1x64x1x64x2_warps4xstages2
void unload_sparse_attention_bf16_sm80_c777f3f5();
void unload_sparse_attention_bf16_sm80_64x1x64x1x64x2_warps4xstages2() {
  unload_sparse_attention_bf16_sm80_c777f3f5();
}

typedef Status (*kernel_func_t)(SparseAttentionParams& params);
kernel_func_t sparse_attention_bf16_sm80_kernels[] = {
    sparse_attention_bf16_sm80_64x0x64x0x64x2_warps4xstages2,
    sparse_attention_bf16_sm80_64x1x64x1x64x2_warps4xstages2,
};

int sparse_attention_bf16_sm80_get_num_algos(void) {
  return (int)sizeof(sparse_attention_bf16_sm80_kernels);
}

Status sparse_attention_bf16_sm80(SparseAttentionParams& params, int algo_id) {
  assert(algo_id < (int)sizeof(sparse_attention_bf16_sm80_kernels));
  return sparse_attention_bf16_sm80_kernels[algo_id](params);
}

void load_sparse_attention_bf16_sm80(void) {
  load_sparse_attention_bf16_sm80_64x0x64x0x64x2_warps4xstages2();
  load_sparse_attention_bf16_sm80_64x1x64x1x64x2_warps4xstages2();
}

void unload_sparse_attention_bf16_sm80(void) {
  unload_sparse_attention_bf16_sm80_64x0x64x0x64x2_warps4xstages2();
  unload_sparse_attention_bf16_sm80_64x1x64x1x64x2_warps4xstages2();
}

Status sparse_attention_bf16_sm80_default(SparseAttentionParams& params) {
  return sparse_attention_bf16_sm80(params, 0);
}

}  // namespace sparse_attention_v1
}  // namespace cuda
}  // namespace contrib
}  // namespace onnxruntime
