#pragma once

#include <torch/torch.h>

void mak_run(at::Tensor prog, at::Tensor barrier, int64_t pos,
             int64_t step_slot, int64_t max_k);

void mak_run_steps(at::Tensor prog, at::Tensor barrier, int64_t pos0,
                   int64_t steps, int64_t slot0, int64_t max_k);

void mak_run_seq(at::Tensor prog, at::Tensor barrier, int64_t pos0,
                 int64_t steps, int64_t slot0, int64_t prompt_len,
                 int64_t max_k, int64_t chunk_m);

at::Tensor mak_run_phased(at::Tensor prog, int64_t pos, int64_t step_slot,
                          bool timed, int64_t max_k);

void mak_run_batch(at::Tensor prog, at::Tensor barrier, at::Tensor pos_b,
                   int64_t steps, int64_t slot0, int64_t max_k,
                   int64_t kv_bstride, int64_t stage_elems);

int64_t mak_num_blocks(at::Tensor prog, int64_t max_k);

int64_t mak_batch_maxb();

double mak_bw_per_sm(at::Tensor probe);