17 lines
965 B
C++
17 lines
965 B
C++
#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");
|
|
}
|