@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
cpp 204 lines 5.29 KB
Raw
1 #include <Python.h>
2
3 extern "C" {
4
5 #include <stdbool.h>
6 #include <time.h>
7
8 #if defined(__i386__) || defined(__x86_64__)
9 #include <cpuid.h>
10 #include <x86intrin.h>
11 #endif
12
13 #if defined(CLOCK_MONOTONIC_RAW)
14 #define TIMEOUT_CLOCK CLOCK_MONOTONIC_RAW
15 #else
16 #define TIMEOUT_CLOCK CLOCK_MONOTONIC
17 #endif
18
19 #define CPU_SUPPORT_NONE 0
20 #define CPU_SUPPORT_MONITORX 1
21
22 #define MWAITX_DEFAULT_TIMEOUT_CYCLES 1000000
23
24 typedef struct {
25 unsigned int cpu_support;
26 unsigned int max_monitor_line_size;
27 } spinloop_state_t;
28
29 static void determine_cpu_support(spinloop_state_t* state) {
30 state->cpu_support = CPU_SUPPORT_NONE;
31 state->max_monitor_line_size = 0;
32
33 #if defined(__i386__) || defined(__x86_64__)
34 unsigned int eax, ebx, ecx, edx;
35 if (__get_cpuid(0, &eax, &ebx, &ecx, &edx) == 1) {
36 // AMD CPU (possible monitorx/mwaitx support)
37 if (ebx == 0x68747541 && edx == 0x69746e65 && ecx == 0x444d4163) {
38 if (__get_cpuid(0x80000000, &eax, &ebx, &ecx, &edx) == 1 &&
39 eax >= 0x80000001 &&
40 __get_cpuid(0x80000001, &eax, &ebx, &ecx, &edx) == 1) {
41 if ((ecx & (1 << 29)) != 0) {
42 state->cpu_support = CPU_SUPPORT_MONITORX;
43 }
44 }
45 }
46 }
47
48 if (state->cpu_support == CPU_SUPPORT_MONITORX) {
49 if (__get_cpuid(5, &eax, &ebx, &ecx, &edx) == 1) {
50 state->max_monitor_line_size = ebx & 0xff;
51 }
52 }
53 #endif
54 }
55
56 static PyObject* method_spinloop(PyObject* self, PyObject* args,
57 PyObject* kwargs) {
58 Py_buffer buffer;
59 PyObject* callback;
60 double timeout = 0.;
61
62 spinloop_state_t* state = (spinloop_state_t*)PyModule_GetState(self);
63 if (state == NULL) {
64 PyErr_SetString(PyExc_TypeError, "Failed to retrieve module state!");
65 return NULL;
66 }
67
68 static const char* keywords[] = {"buffer", "callback", "timeout", NULL};
69 if (!PyArg_ParseTupleAndKeywords(args, kwargs, "y*O|d", (char**)keywords,
70 &buffer, &callback, &timeout)) {
71 return NULL;
72 }
73
74 if (!PyCallable_Check(callback)) {
75 PyErr_SetString(PyExc_TypeError, "callback parameter must be callable!");
76 PyBuffer_Release(&buffer);
77 return NULL;
78 }
79
80 struct timespec t_start;
81 if (clock_gettime(TIMEOUT_CLOCK, &t_start) != 0) {
82 PyErr_SetString(PyExc_RuntimeError, "clock_gettime() failed!");
83 PyBuffer_Release(&buffer);
84 return NULL;
85 }
86
87 bool result = false;
88 bool error = false;
89 bool have_timeout = (timeout > 1e-9);
90 unsigned int iteration = 0;
91 const bool buffer_qualifies = (buffer.len <= state->max_monitor_line_size);
92
93 while (true) {
94 PyObject* res = PyObject_CallNoArgs(callback);
95 if (res == NULL) {
96 error = true;
97 break;
98 }
99 int ok = (res == Py_True);
100 Py_DECREF(res);
101
102 if (ok) {
103 result = true;
104 break;
105 }
106
107 // Check timeout at most every 16 iterations to avoid clock_gettime and
108 // comparison cost
109 if (have_timeout && (iteration & 15u) == 0) {
110 struct timespec t_now;
111 if (clock_gettime(TIMEOUT_CLOCK, &t_now) != 0) {
112 PyErr_SetString(PyExc_RuntimeError, "clock_gettime() failed!");
113 error = true;
114 break;
115 }
116
117 const double elapsed = (double)(t_now.tv_sec - t_start.tv_sec) +
118 (t_now.tv_nsec - t_start.tv_nsec) * 1e-9;
119 if (elapsed >= timeout) {
120 result = false;
121 break;
122 }
123 }
124 ++iteration;
125
126 #if defined(__i386__) || defined(__x86_64__)
127 // monitorx + mwaitx with qualified buffer
128 if (buffer_qualifies && state->cpu_support == CPU_SUPPORT_MONITORX) {
129 _mm_monitorx(buffer.buf, 0, 0);
130
131 // Check once more in case the buffer has been modified while we were
132 // arming the monitor hardware
133 res = PyObject_CallNoArgs(callback);
134 if (res == NULL) {
135 error = true;
136 break;
137 }
138 ok = (res == Py_True);
139 Py_DECREF(res);
140
141 if (ok) {
142 result = true;
143 break;
144 }
145
146 // Run mwaitx with enabled timeout (bit 1). The actual timeout value
147 // is not very important, we just want to ensure we don't lock up
148 // here for too long.
149 Py_BEGIN_ALLOW_THREADS _mm_mwaitx((1 << 1), 0,
150 MWAITX_DEFAULT_TIMEOUT_CYCLES);
151 Py_END_ALLOW_THREADS
152 }
153
154 // Fallback: Busy poll
155 else {
156 #endif
157 // Give other threads a chance to be scheduled
158 Py_BEGIN_ALLOW_THREADS
159 #if defined(__i386__) || defined(__x86_64__)
160 __builtin_ia32_pause();
161 #elif defined(__aarch64__)
162 __asm__ volatile("yield" :: : "memory");
163 #endif
164 Py_END_ALLOW_THREADS
165 #if defined(__i386__) || defined(__x86_64__)
166 }
167 #endif
168 }
169
170 PyBuffer_Release(&buffer);
171
172 if (error) {
173 return NULL;
174 }
175
176 if (result) {
177 Py_RETURN_TRUE;
178 }
179
180 Py_RETURN_FALSE;
181 }
182
183 static PyMethodDef spinloop_methods[] = {
184 {"spinloop", (PyCFunction)method_spinloop, METH_VARARGS | METH_KEYWORDS,
185 "Wait for store with callback"},
186 {NULL, NULL, 0, NULL}};
187
188 static struct PyModuleDef spinloop_module = {
189 PyModuleDef_HEAD_INIT, "spinloop",
190 "Hardware-optimized spinloops for Python", sizeof(spinloop_state_t),
191 spinloop_methods};
192
193 PyMODINIT_FUNC PyInit_spinloop(void) {
194 PyObject* m = PyModule_Create(&spinloop_module);
195 if (m != NULL) {
196 spinloop_state_t* state = (spinloop_state_t*)PyModule_GetState(m);
197 if (state != NULL) {
198 determine_cpu_support(state);
199 }
200 }
201 return m;
202 }
203
204 } // extern "C"