#include <torch/extension.h>

#include <cstdint>

torch::Tensor launch_handoff_probe(int64_t blocks, int64_t epochs, bool payload);
torch::Tensor launch_mbarrier_handoff_probe(int64_t blocks, int64_t epochs, bool payload);
std::vector<torch::Tensor> launch_tensor_issue_probe(int64_t blocks, int64_t iterations, int64_t mode);
torch::Tensor launch_register_probe(int64_t blocks, int64_t iterations, int64_t role);
torch::Tensor probe_kernel_attributes(int64_t dynamic_shared_bytes);

PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) {
  module.def("handoff", &launch_handoff_probe, "VEA-B bounded handoff probe");
  module.def("handoff_mbarrier", &launch_mbarrier_handoff_probe, "VEA-B mbarrier handoff probe");
  module.def("tensor_issue", &launch_tensor_issue_probe, "VEA-B INT8/FP8 issue probe");
  module.def("register_probe", &launch_register_probe, "VEA-B role register probe");
  module.def("attributes", &probe_kernel_attributes, "VEA-B kernel attributes");
}