1 #include <torch/csrc/instruction_counter/Module.h>
2 #include <torch/csrc/utils/pybind.h>
3 #include <cerrno>
4 #include <cstdint>
5 #include <cstdio>
6 #include <cstdlib>
7 #include <cstring>
8 #include <stdexcept>
9
10 #if defined(__linux__)
11 #include <linux/perf_event.h>
12 #include <sys/ioctl.h>
13 #include <sys/syscall.h>
14 #include <unistd.h>
15 #endif
16
17 namespace torch::instruction_counter {
18
start()19 long start() {
20 #if !defined(__linux__)
21 throw std::runtime_error("This systems seems not to be Linux");
22 #else
23
24 // Construct base perf_event_attr struct
25 perf_event_attr attr{};
26 memset(&attr, 0, sizeof(attr));
27 attr.size = sizeof(attr);
28 attr.exclude_kernel = 1;
29 attr.disabled = 1;
30 attr.exclude_hv = 1;
31 attr.sample_period = 0;
32 // Enable hardware counting
33 attr.type = PERF_TYPE_HARDWARE;
34 attr.config = PERF_COUNT_HW_INSTRUCTIONS;
35
36 long fd = syscall(SYS_perf_event_open, &attr, 0, -1, -1, 0);
37 if (fd == -1) {
38 fprintf(
39 stderr,
40 "Failed to open instruction count event: %s.\n",
41 strerror(errno));
42 return -1;
43 }
44 ioctl((int)fd, PERF_EVENT_IOC_RESET, 0); // Reset the counter
45 ioctl((int)fd, PERF_EVENT_IOC_ENABLE, 0); // Enable the counter
46 return fd;
47 #endif
48 }
49
end(int fd)50 uint64_t end(int fd) {
51 #if !defined(__linux__)
52 throw std::runtime_error("This systems seems not to be Linux");
53 #else
54 // Disable the event group
55 if (ioctl(fd, PERF_EVENT_IOC_DISABLE, PERF_IOC_FLAG_GROUP) == -1) {
56 fprintf(
57 stderr,
58 "Error disabling perf event (fd: %d): %s\n",
59 fd,
60 strerror(errno));
61 return -1;
62 }
63
64 uint64_t total_instructions = 0;
65
66 // Read results
67 long ret_val = read(fd, &total_instructions, sizeof(total_instructions));
68 if (ret_val == -1) {
69 fprintf(stderr, "Error reading perf event results: %s\n", strerror(errno));
70 return -1;
71 }
72
73 close(fd);
74 return total_instructions;
75 #endif
76 }
77
initModule(PyObject * module)78 void initModule(PyObject* module) {
79 auto m = py::handle(module).cast<py::module>();
80 auto instruction_counter = m.def_submodule(
81 "_instruction_counter", "instruction_counter related pybind.");
82 instruction_counter.def("start", start);
83 instruction_counter.def("end", end);
84 }
85
86 } // namespace torch::instruction_counter
87