gh-120953: Add UOp Paircounts to JIT when Built with PyStats

Include/cpython/optimizer.h

@@ -72,6 +72,7 @@ typedef struct _PyExecutorObject {
     uint32_t exit_count;
     uint32_t code_size;
     size_t jit_size;
+    uint16_t last_uop;
     void *jit_code;
     void *jit_side_entry;
     _PyExitData exits[1];

Python/optimizer.c

@@ -1215,6 +1215,7 @@ make_executor_from_uops(_PyUOpInstruction *buffer, int length, const _PyBloomFil
     executor->jit_code = NULL;
     executor->jit_side_entry = NULL;
     executor->jit_size = 0;
+    executor->last_uop = 0;
     if (_PyJIT_Compile(executor, executor->trace, length)) {
         Py_DECREF(executor);
         return NULL;

Tools/jit/_targets.py

@@ -44,11 +44,13 @@ class _Target(typing.Generic[_S, _R]):
     stable: bool = False
     debug: bool = False
     verbose: bool = False
+    pystats: bool = False
 
     def _compute_digest(self, out: pathlib.Path) -> str:
         hasher = hashlib.sha256()
         hasher.update(self.triple.encode())
         hasher.update(self.debug.to_bytes())
+        hasher.update(self.pystats.to_bytes())
         # These dependencies are also reflected in _JITSources in regen.targets:
         hasher.update(PYTHON_EXECUTOR_CASES_C_H.read_bytes())
         hasher.update((out / "pyconfig.h").read_bytes())
@@ -142,6 +144,8 @@ async def _compile(
             "-std=c11",
             *self.args,
         ]
+        if self.pystats:
+            args.append("-DPy_STATS=1")
         if self.ghccc:
             # This is a bit of an ugly workaround, but it makes the code much
             # smaller and faster, so it's worth it. We want to use the GHC

Tools/jit/build.py

@@ -13,6 +13,9 @@
     parser.add_argument(
         "target", type=_targets.get_target, help="a PEP 11 target triple to compile for"
     )
+    parser.add_argument(
+        "--pystats", action="store_true", help="compile for a pystats build of Python"
+    )
     parser.add_argument(
         "-d", "--debug", action="store_true", help="compile for a debug build of Python"
     )
@@ -25,4 +28,5 @@
     args = parser.parse_args()
     args.target.debug = args.debug
     args.target.verbose = args.verbose
+    args.target.pystats = args.pystats
     args.target.build(pathlib.Path.cwd(), comment=comment, force=args.force)

Tools/jit/template.c

@@ -4,6 +4,7 @@
 #include "pycore_call.h"
 #include "pycore_ceval.h"
 #include "pycore_cell.h"
+#include "pycore_code.h"
 #include "pycore_dict.h"
 #include "pycore_emscripten_signal.h"
 #include "pycore_intrinsics.h"
@@ -104,6 +105,11 @@ _JIT_ENTRY(_PyInterpreterFrame *frame, PyObject **stack_pointer, PyThreadState *
     OPT_STAT_INC(uops_executed);
     UOP_STAT_INC(uopcode, execution_count);
 
+#ifdef Py_STATS
+    UOP_PAIR_INC(uopcode, current_executor->last_uop);
+    current_executor->last_uop = uopcode;
+#endif
+
     // The actual instruction definitions (only one will be used):
     if (uopcode == _JUMP_TO_TOP) {
         PATCH_JUMP(_JIT_TOP);

configure.ac

@@ -1798,7 +1798,12 @@ AS_VAR_IF([jit_flags],
            AS_VAR_IF([Py_DEBUG],
                      [true],
                      [AS_VAR_APPEND([REGEN_JIT_COMMAND], [" --debug"])],
-                     [])])
+                     [])
+           AS_VAR_IF([enable_pystats],
+                     [yes],
+                     [AS_VAR_APPEND([REGEN_JIT_COMMAND], [" --pystats"])],
+                     [])
+                    ])
 AC_SUBST([REGEN_JIT_COMMAND])
 AC_SUBST([JIT_STENCILS_H])
 AC_MSG_RESULT([$tier2_flags $jit_flags])