GH-106360: Support very basic superblock introspection (#106422)

* Add len() and indexing support to uop superblocks.

Include/cpython/optimizer.h

@@ -38,6 +38,8 @@ PyAPI_FUNC(void) PyUnstable_SetOptimizer(_PyOptimizerObject* optimizer);
 
 PyAPI_FUNC(_PyOptimizerObject *) PyUnstable_GetOptimizer(void);
 
+PyAPI_FUNC(_PyExecutorObject *)PyUnstable_GetExecutor(PyCodeObject *code, int offset);
+
 struct _PyInterpreterFrame *
 _PyOptimizer_BackEdge(struct _PyInterpreterFrame *frame, _Py_CODEUNIT *src, _Py_CODEUNIT *dest, PyObject **stack_pointer);
 

Include/internal/pycore_opcode.h

@@ -242,8 +242,11 @@ const uint8_t _PyOpcode_Deopt[256] = {
 };
 #endif   // NEED_OPCODE_TABLES
 
-#ifdef Py_DEBUG
-static const char *const _PyOpcode_OpName[268] = {
+
+extern const char *const _PyOpcode_OpName[268];
+
+#ifdef NEED_OPCODE_TABLES
+const char *const _PyOpcode_OpName[268] = {
     [CACHE] = "CACHE",
     [POP_TOP] = "POP_TOP",
     [PUSH_NULL] = "PUSH_NULL",
@@ -513,7 +516,7 @@ static const char *const _PyOpcode_OpName[268] = {
     [STORE_FAST_MAYBE_NULL] = "STORE_FAST_MAYBE_NULL",
     [LOAD_CLOSURE] = "LOAD_CLOSURE",
 };
-#endif
+#endif   // NEED_OPCODE_TABLES
 
 #define EXTRA_CASES \
     case 184: \

Lib/test/test_capi/test_misc.py

@@ -2415,5 +2415,28 @@ class TestOptimizerAPI(unittest.TestCase):
             self.assertEqual(opt.get_count(), 10)
 
 
+class TestUops(unittest.TestCase):
+
+    def test_basic_loop(self):
+
+        def testfunc(x):
+            i = 0
+            while i < x:
+                i += 1
+
+        testfunc(1000)
+
+        ex = None
+        for offset in range(0, 100, 2):
+            try:
+                ex = _testinternalcapi.get_executor(testfunc.__code__, offset)
+                break
+            except ValueError:
+                pass
+        if ex is None:
+            return
+        self.assertIn("SAVE_IP", str(ex))
+
+
 if __name__ == "__main__":
     unittest.main()

Modules/_testinternalcapi.c

@@ -858,6 +858,26 @@ get_optimizer(PyObject *self, PyObject *Py_UNUSED(ignored))
     return opt;
 }
 
+static PyObject *
+get_executor(PyObject *self, PyObject *const *args, Py_ssize_t nargs)
+{
+
+    if (!_PyArg_CheckPositional("get_executor", nargs, 2, 2)) {
+        return NULL;
+    }
+    PyObject *code = args[0];
+    PyObject *offset = args[1];
+    long ioffset = PyLong_AsLong(offset);
+    if (ioffset == -1 && PyErr_Occurred()) {
+        return NULL;
+    }
+    if (!PyCode_Check(code)) {
+         PyErr_SetString(PyExc_TypeError, "first argument must be a code object");
+        return NULL;
+    }
+    return (PyObject *)PyUnstable_GetExecutor((PyCodeObject *)code, ioffset);
+}
+
 static int _pending_callback(void *arg)
 {
     /* we assume the argument is callable object to which we own a reference */
@@ -1326,6 +1346,7 @@ static PyMethodDef module_functions[] = {
     {"iframe_getlasti", iframe_getlasti, METH_O, NULL},
     {"get_optimizer", get_optimizer,  METH_NOARGS, NULL},
     {"set_optimizer", set_optimizer,  METH_O, NULL},
+    {"get_executor", _PyCFunction_CAST(get_executor),  METH_FASTCALL, NULL},
     {"get_counter_optimizer", get_counter_optimizer, METH_NOARGS, NULL},
     {"get_uop_optimizer", get_uop_optimizer, METH_NOARGS, NULL},
     {"pending_threadfunc", _PyCFunction_CAST(pending_threadfunc),

Python/opcode_metadata.h

@@ -942,9 +942,7 @@ struct opcode_macro_expansion {
 #ifndef NEED_OPCODE_METADATA
 extern const struct opcode_metadata _PyOpcode_opcode_metadata[512];
 extern const struct opcode_macro_expansion _PyOpcode_macro_expansion[256];
-#ifdef Py_DEBUG
 extern const char * const _PyOpcode_uop_name[512];
-#endif
 #else
 const struct opcode_metadata _PyOpcode_opcode_metadata[512] = {
     [NOP] = { true, INSTR_FMT_IX, 0 },
@@ -1265,7 +1263,7 @@ const struct opcode_macro_expansion _PyOpcode_macro_expansion[256] = {
     [COPY] = { .nuops = 1, .uops = { { COPY, 0, 0 } } },
     [SWAP] = { .nuops = 1, .uops = { { SWAP, 0, 0 } } },
 };
-#ifdef Py_DEBUG
+#ifdef NEED_OPCODE_METADATA
 const char * const _PyOpcode_uop_name[512] = {
     [300] = "EXIT_TRACE",
     [301] = "SAVE_IP",
@@ -1282,5 +1280,5 @@ const char * const _PyOpcode_uop_name[512] = {
     [312] = "_LOAD_LOCALS",
     [313] = "_LOAD_FROM_DICT_OR_GLOBALS",
 };
-#endif
+#endif // NEED_OPCODE_METADATA
 #endif

Python/optimizer.c

@@ -188,6 +188,23 @@ jump_to_destination:
     return frame;
 }
 
+_PyExecutorObject *
+PyUnstable_GetExecutor(PyCodeObject *code, int offset)
+{
+    int code_len = (int)Py_SIZE(code);
+    for (int i = 0 ; i < code_len;) {
+        if (_PyCode_CODE(code)[i].op.code == ENTER_EXECUTOR && i*2 == offset) {
+            int oparg = _PyCode_CODE(code)[i].op.arg;
+            _PyExecutorObject *res = code->co_executors->executors[oparg];
+            Py_INCREF(res);
+            return res;
+        }
+        i += _PyInstruction_GetLength(code, i);
+    }
+    PyErr_SetString(PyExc_ValueError, "no executor at given offset");
+    return NULL;
+}
+
 /** Test support **/
 
 
@@ -287,6 +304,58 @@ uop_dealloc(_PyUOpExecutorObject *self) {
     PyObject_Free(self);
 }
 
+static const char *
+uop_name(int index) {
+    if (index < EXIT_TRACE) {
+        return _PyOpcode_OpName[index];
+    }
+    return _PyOpcode_uop_name[index];
+}
+
+static Py_ssize_t
+uop_len(_PyUOpExecutorObject *self)
+{
+    int count = 1;
+    for (; count < _Py_UOP_MAX_TRACE_LENGTH; count++) {
+        if (self->trace[count-1].opcode == EXIT_TRACE) {
+            break;
+        }
+    }
+    return count;
+}
+
+static PyObject *
+uop_item(_PyUOpExecutorObject *self, Py_ssize_t index)
+{
+    for (int i = 0; i < _Py_UOP_MAX_TRACE_LENGTH; i++) {
+        if (self->trace[i].opcode == EXIT_TRACE) {
+            break;
+        }
+        if (i != index) {
+            continue;
+        }
+        const char *name = uop_name(self->trace[i].opcode);
+        PyObject *oname = _PyUnicode_FromASCII(name, strlen(name));
+        if (oname == NULL) {
+            return NULL;
+        }
+        PyObject *operand = PyLong_FromUnsignedLongLong(self->trace[i].operand);
+        if (operand == NULL) {
+            Py_DECREF(oname);
+            return NULL;
+        }
+        PyObject *args[2] = { oname, operand };
+        return _PyTuple_FromArraySteal(args, 2);
+    }
+    PyErr_SetNone(PyExc_IndexError);
+    return NULL;
+}
+
+PySequenceMethods uop_as_sequence = {
+    .sq_length = (lenfunc)uop_len,
+    .sq_item = (ssizeargfunc)uop_item,
+};
+
 static PyTypeObject UOpExecutor_Type = {
     PyVarObject_HEAD_INIT(&PyType_Type, 0)
     .tp_name = "uop_executor",
@@ -294,6 +363,7 @@ static PyTypeObject UOpExecutor_Type = {
     .tp_itemsize = 0,
     .tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_DISALLOW_INSTANTIATION,
     .tp_dealloc = (destructor)uop_dealloc,
+    .tp_as_sequence = &uop_as_sequence,
 };
 
 static int

Tools/build/generate_opcode_h.py

@@ -184,14 +184,15 @@ def main(opcode_py,
         fobj.write(f"#define ENABLE_SPECIALIZATION {int(ENABLE_SPECIALIZATION)}")
 
         iobj.write("\n")
-        iobj.write("#ifdef Py_DEBUG\n")
-        iobj.write(f"static const char *const _PyOpcode_OpName[{NUM_OPCODES}] = {{\n")
+        iobj.write(f"\nextern const char *const _PyOpcode_OpName[{NUM_OPCODES}];\n")
+        iobj.write("\n#ifdef NEED_OPCODE_TABLES\n")
+        iobj.write(f"const char *const _PyOpcode_OpName[{NUM_OPCODES}] = {{\n")
         for op, name in enumerate(opname_including_specialized):
             if name[0] != "<":
                 op = name
             iobj.write(f'''    [{op}] = "{name}",\n''')
         iobj.write("};\n")
-        iobj.write("#endif\n")
+        iobj.write("#endif   // NEED_OPCODE_TABLES\n")
 
         iobj.write("\n")
         iobj.write("#define EXTRA_CASES \\\n")

Tools/cases_generator/generate_cases.py

@@ -1224,9 +1224,7 @@ class Analyzer:
             self.out.emit("#ifndef NEED_OPCODE_METADATA")
             self.out.emit("extern const struct opcode_metadata _PyOpcode_opcode_metadata[512];")
             self.out.emit("extern const struct opcode_macro_expansion _PyOpcode_macro_expansion[256];")
-            self.out.emit("#ifdef Py_DEBUG")
             self.out.emit("extern const char * const _PyOpcode_uop_name[512];")
-            self.out.emit("#endif")
             self.out.emit("#else")
 
             self.out.emit("const struct opcode_metadata _PyOpcode_opcode_metadata[512] = {")
@@ -1273,10 +1271,10 @@ class Analyzer:
                         case _:
                             typing.assert_never(thing)
 
-            self.out.emit("#ifdef Py_DEBUG")
+            self.out.emit("#ifdef NEED_OPCODE_METADATA")
             with self.out.block("const char * const _PyOpcode_uop_name[512] =", ";"):
                 self.write_uop_items(lambda name, counter: f"[{counter}] = \"{name}\",")
-            self.out.emit("#endif")
+            self.out.emit("#endif // NEED_OPCODE_METADATA")
 
             self.out.emit("#endif")