GH-113464: A copy-and-patch JIT compiler

.github/workflows/jit.yml

@@ -0,0 +1,111 @@
+name: JIT
+on:
+  - pull_request
+  - push
+jobs:
+  jit:
+    # Comment this line to run JIT CI:
+    # if: false  # XXX: Uncomment before merging.
+    name: ${{ matrix.target }} (${{ matrix.debug && 'Debug' || 'Release' }})
+    runs-on: ${{ matrix.runner }}
+    strategy:
+      fail-fast: false
+      matrix:
+        target:
+          - i686-pc-windows-msvc/msvc
+          - x86_64-pc-windows-msvc/msvc
+          - x86_64-apple-darwin/clang
+          - x86_64-unknown-linux-gnu/gcc
+          - x86_64-unknown-linux-gnu/clang
+          - aarch64-unknown-linux-gnu/gcc
+          - aarch64-unknown-linux-gnu/clang
+        debug:
+          - true
+          - false
+        llvm:
+          - 16
+        include:
+          - target: i686-pc-windows-msvc/msvc
+            architecture: Win32
+            runner: windows-latest
+            compiler: msvc
+          - target: x86_64-pc-windows-msvc/msvc
+            architecture: x64
+            runner: windows-latest
+            compiler: msvc
+          - target: x86_64-apple-darwin/clang
+            architecture: x86_64
+            runner: macos-latest
+            compiler: clang
+            exclude: test_embed
+          - target: x86_64-unknown-linux-gnu/gcc
+            architecture: x86_64
+            runner: ubuntu-latest
+            compiler: gcc
+          - target: x86_64-unknown-linux-gnu/clang
+            architecture: x86_64
+            runner: ubuntu-latest
+            compiler: clang
+          - target: aarch64-unknown-linux-gnu/gcc
+            architecture: aarch64
+            runner: ubuntu-latest
+            compiler: gcc
+            # These fail because of emulation, not because of the JIT:
+            exclude: test_unix_events test_init test_process_pool test_shutdown test_multiprocessing_fork test_cmd_line test_faulthandler test_os test_perf_profiler test_posix test_signal test_socket test_subprocess test_threading test_venv
+          - target: aarch64-unknown-linux-gnu/clang
+            architecture: aarch64
+            runner: ubuntu-latest
+            compiler: clang
+            # These fail because of emulation, not because of the JIT:
+            exclude: test_unix_events test_init test_process_pool test_shutdown test_multiprocessing_fork test_cmd_line test_faulthandler test_os test_perf_profiler test_posix test_signal test_socket test_subprocess test_threading test_venv
+    env:
+      CC: ${{ matrix.compiler }}
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: '3.11'
+
+      - name: Windows
+        if: runner.os == 'Windows'
+        run: |
+          choco install llvm --allow-downgrade --version ${{ matrix.llvm }}
+          ./PCbuild/build.bat --experimental-jit ${{ matrix.debug && '-d' || '--pgo' }} -p ${{ matrix.architecture }}
+          ./PCbuild/rt.bat ${{ matrix.debug && '-d' }} -p ${{ matrix.architecture }} -q --exclude ${{ matrix.exclude }} --multiprocess 0 --timeout 3600 --verbose2 --verbose3
+
+      - name: macOS
+        if: runner.os == 'macOS'
+        run: |
+          brew install llvm@${{ matrix.llvm }}
+          export SDKROOT="$(xcrun --show-sdk-path)"
+          ./configure --enable-experimental-jit ${{ matrix.debug && '--with-pydebug' || '--enable-optimizations --with-lto' }}
+          make all --jobs 3
+          ./python.exe -m test --exclude ${{ matrix.exclude }} --multiprocess 0 --timeout 3600 --verbose2 --verbose3
+
+      - name: Native Linux
+        if: runner.os == 'Linux' && matrix.architecture == 'x86_64'
+        run: |
+          sudo bash -c "$(wget -O - https://apt.llvm.org/llvm.sh)" ./llvm.sh ${{ matrix.llvm }}
+          export PATH="$(llvm-config-${{ matrix.llvm }} --bindir):$PATH"
+          ./configure --enable-experimental-jit ${{ matrix.debug && '--with-pydebug' || '--enable-optimizations --with-lto' }}
+          make all --jobs 2
+          ./python -m test --exclude ${{ matrix.exclude }} --multiprocess 0 --timeout 3600 --verbose2 --verbose3
+      - name: Emulated Linux
+        if: runner.os == 'Linux' && matrix.architecture != 'x86_64'
+        run: |
+          sudo bash -c "$(wget -O - https://apt.llvm.org/llvm.sh)" ./llvm.sh ${{ matrix.llvm }}
+          export PATH="$(llvm-config-${{ matrix.llvm }} --bindir):$PATH"
+          ./configure --prefix="$(pwd)/../build"
+          make install --jobs 2
+          make clean --jobs 2
+          export HOST=${{ matrix.architecture }}-linux-gnu
+          sudo apt install --yes "gcc-$HOST" qemu-user
+          ${{ !matrix.debug && matrix.compiler == 'clang' && './configure --enable-optimizations' || '' }}
+          ${{ !matrix.debug && matrix.compiler == 'clang' && 'make profile-run-stamp --jobs 2' || '' }}
+          export CC="${{ matrix.compiler == 'clang' && 'clang --target=$HOST' || '$HOST-gcc' }}"
+          export CPP="$CC --preprocess"
+          export HOSTRUNNER=qemu-${{ matrix.architecture }}
+          export QEMU_LD_PREFIX="/usr/$HOST"
+          ./configure --enable-experimental-jit ${{ matrix.debug && '--with-pydebug' || '--enable-optimizations --with-lto' }} --build=x86_64-linux-gnu --host="$HOST" --with-build-python=../build/bin/python3 --with-pkg-config=no ac_cv_buggy_getaddrinfo=no ac_cv_file__dev_ptc=no ac_cv_file__dev_ptmx=yes
+          make all --jobs 2
+          ./python -m test --exclude ${{ matrix.exclude }} --multiprocess 0 --timeout 3600 --verbose2 --verbose3

.github/workflows/mypy.yml

@@ -12,6 +12,7 @@ on:
       - "Tools/build/generate_sbom.py"
       - "Tools/cases_generator/**"
       - "Tools/clinic/**"
+      - "Tools/jit/**"
       - "Tools/peg_generator/**"
       - "Tools/requirements-dev.txt"
       - "Tools/wasm/**"
@@ -38,6 +39,7 @@ jobs:
           "Tools/build/",
           "Tools/cases_generator",
           "Tools/clinic",
+          "Tools/jit",
           "Tools/peg_generator",
           "Tools/wasm",
         ]

.gitignore

@@ -126,6 +126,7 @@ Tools/unicode/data/
 # hendrikmuhs/ccache-action@v1
 /.ccache
 /cross-build/
+/jit_stencils.h
 /platform
 /profile-clean-stamp
 /profile-run-stamp

Include/cpython/optimizer.h

@@ -75,7 +75,7 @@ extern void _Py_Executors_InvalidateAll(PyInterpreterState *interp);
 
 /* For testing */
 PyAPI_FUNC(PyObject *)PyUnstable_Optimizer_NewCounter(void);
-PyAPI_FUNC(PyObject *)PyUnstable_Optimizer_NewUOpOptimizer(void);
+PyAPI_FUNC(PyObject *)PyUnstable_Optimizer_NewUOpOptimizer(int jit);
 
 #define OPTIMIZER_BITS_IN_COUNTER 4
 /* Minimum of 16 additional executions before retry */

Include/internal/pycore_jit.h

@@ -0,0 +1,27 @@
+#ifndef Py_INTERNAL_JIT_H
+#define Py_INTERNAL_JIT_H
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+#ifndef Py_BUILD_CORE
+#  error "this header requires Py_BUILD_CORE define"
+#endif
+
+#ifdef _Py_JIT
+
+#include "pycore_uops.h"
+
+typedef _Py_CODEUNIT *(*jit_func)(_PyInterpreterFrame *frame, PyObject **stack_pointer, PyThreadState *tstate);
+
+int _PyJIT_Compile(_PyUOpExecutorObject *executor);
+void _PyJIT_Free(_PyUOpExecutorObject *executor);
+
+#endif  // _Py_JIT
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif // !Py_INTERNAL_JIT_H

Include/internal/pycore_object.h

@@ -178,7 +178,7 @@ _Py_DECREF_SPECIALIZED(PyObject *op, const destructor destruct)
     }
     _Py_DECREF_STAT_INC();
 #ifdef Py_REF_DEBUG
-    _Py_DEC_REFTOTAL(_PyInterpreterState_GET());
+    _Py_DEC_REFTOTAL(PyInterpreterState_Get());
 #endif
     if (--op->ob_refcnt != 0) {
         assert(op->ob_refcnt > 0);
@@ -199,7 +199,7 @@ _Py_DECREF_NO_DEALLOC(PyObject *op)
     }
     _Py_DECREF_STAT_INC();
 #ifdef Py_REF_DEBUG
-    _Py_DEC_REFTOTAL(_PyInterpreterState_GET());
+    _Py_DEC_REFTOTAL(PyInterpreterState_Get());
 #endif
     op->ob_refcnt--;
 #ifdef Py_DEBUG

Include/internal/pycore_uops.h

@@ -21,9 +21,16 @@ typedef struct {
 
 typedef struct {
     _PyExecutorObject base;
+    void *jit_code;
+    size_t jit_size;
     _PyUOpInstruction trace[1];
 } _PyUOpExecutorObject;
 
+typedef struct {
+    _PyOptimizerObject base;
+    bool jit;
+} _PyUOpOptimizerObject;
+
 _Py_CODEUNIT *_PyUOpExecute(
     _PyExecutorObject *executor,
     _PyInterpreterFrame *frame,

Makefile.pre.in

@@ -430,6 +430,7 @@ PYTHON_OBJS=	\
 		Python/initconfig.o \
 		Python/instrumentation.o \
 		Python/intrinsics.o \
+		Python/jit.o \
 		Python/legacy_tracing.o \
 		Python/lock.o \
 		Python/marshal.o \
@@ -1361,7 +1362,7 @@ regen-unicodedata:
 regen-all: regen-cases regen-typeslots \
 	regen-token regen-ast regen-keyword regen-sre regen-frozen \
 	regen-pegen-metaparser regen-pegen regen-test-frozenmain \
-	regen-test-levenshtein regen-global-objects regen-sbom
+	regen-test-levenshtein regen-global-objects regen-sbom regen-jit
 	@echo
 	@echo "Note: make regen-stdlib-module-names, make regen-limited-abi, "
 	@echo "make regen-configure and make regen-unicodedata should be run manually"
@@ -1840,6 +1841,7 @@ PYTHON_HEADERS= \
 		$(srcdir)/Include/internal/pycore_initconfig.h \
 		$(srcdir)/Include/internal/pycore_interp.h \
 		$(srcdir)/Include/internal/pycore_intrinsics.h \
+		$(srcdir)/Include/internal/pycore_jit.h \
 		$(srcdir)/Include/internal/pycore_list.h \
 		$(srcdir)/Include/internal/pycore_llist.h \
 		$(srcdir)/Include/internal/pycore_lock.h \
@@ -2632,6 +2634,12 @@ config.status:	$(srcdir)/configure
 Python/asm_trampoline.o: $(srcdir)/Python/asm_trampoline.S
 	$(CC) -c $(PY_CORE_CFLAGS) -o $@ $<
 
+Python/jit.o: regen-jit
+
+.PHONY: regen-jit
+regen-jit:
+	@REGEN_JIT_COMMAND@
+
 # Some make's put the object file in the current directory
 .c.o:
 	$(CC) -c $(PY_CORE_CFLAGS) -o $@ $<
@@ -2724,6 +2732,7 @@ clean-retain-profile: pycremoval
 	-rm -f Python/deepfreeze/*.[co]
 	-rm -f Python/frozen_modules/*.h
 	-rm -f Python/frozen_modules/MANIFEST
+	-rm -f jit_stencils.h
 	-find build -type f -a ! -name '*.gc??' -exec rm -f {} ';'
 	-rm -f Include/pydtrace_probes.h
 	-rm -f profile-gen-stamp

Misc/NEWS.d/next/Core and Builtins/2023-12-24-03-25-28.gh-issue-113464.dvjQmA.rst

@@ -0,0 +1,4 @@
+Add an option (``--enable-experimental-jit`` for ``configure``-based builds
+or ``--experimental-jit`` for ``PCbuild``-based ones) to build an
+*experimental* just-in-time compiler, based on `copy-and-patch
+<https://fredrikbk.com/publications/copy-and-patch.pdf>`_

Modules/_testinternalcapi.c

@@ -968,7 +968,7 @@ get_counter_optimizer(PyObject *self, PyObject *arg)
 static PyObject *
 get_uop_optimizer(PyObject *self, PyObject *arg)
 {
-    return PyUnstable_Optimizer_NewUOpOptimizer();
+    return PyUnstable_Optimizer_NewUOpOptimizer(false);
 }
 
 static PyObject *

PCbuild/_freeze_module.vcxproj

@@ -221,6 +221,7 @@
     <ClCompile Include="..\Python\initconfig.c" />
     <ClCompile Include="..\Python\intrinsics.c" />
     <ClCompile Include="..\Python\instrumentation.c" />
+    <ClCompile Include="..\Python\jit.c" />
     <ClCompile Include="..\Python\legacy_tracing.c" />
     <ClCompile Include="..\Python\lock.c" />
     <ClCompile Include="..\Python\marshal.c" />
@@ -443,6 +444,14 @@
                       Overwrite="true"
                       Condition="$(PyConfigHText.TrimEnd()) != $(OldPyConfigH.TrimEnd())" />
   </Target>
+  <Target Name="_BuildJIT" BeforeTargets="PrepareForBuild" Condition="'$(UseJIT)' == 'true'" DependsOnTargets="_UpdatePyconfig;FindPythonForBuild">
+    <PropertyGroup>
+      <JITArgs Condition="$(Platform) == 'Win32'">i686-pc-windows-msvc</JITArgs>
+      <JITArgs Condition="$(Platform) == 'x64'">x86_64-pc-windows-msvc</JITArgs>
+      <JITArgs Condition="$(Configuration) == 'Debug'">$(JITArgs) --debug</JITArgs>
+    </PropertyGroup>
+    <Exec Command='$(PythonForBuild) "$(PySourcePath)Tools\jit\build.py" $(JITArgs)' WorkingDirectory="$(IntDir)"/>
+  </Target>
 
   <Target Name="_RebuildGetPath" AfterTargets="_RebuildFrozen" Condition="$(Configuration) != 'PGUpdate'">
     <Exec Command='"$(TargetPath)" "%(GetPath.ModName)" "%(GetPath.FullPath)" "%(GetPath.IntFile)"' />

PCbuild/_freeze_module.vcxproj.filters

@@ -241,6 +241,9 @@
     <ClCompile Include="..\Objects\iterobject.c">
       <Filter>Source Files</Filter>
     </ClCompile>
+    <ClCompile Include="..\Python\jit.c">
+      <Filter>Source Files</Filter>
+    </ClCompile>
     <ClCompile Include="..\Objects\listobject.c">
       <Filter>Source Files</Filter>
     </ClCompile>

PCbuild/build.bat

@@ -36,6 +36,7 @@ echo.                 overrides -c and -d
 echo.  --disable-gil  Enable experimental support for running without the GIL.
 echo.  --test-marker  Enable the test marker within the build.
 echo.  --regen        Regenerate all opcodes, grammar and tokens.
+echo.  --experimental-jit  Enable the experimental just-in-time compiler.
 echo.
 echo.Available flags to avoid building certain modules.
 echo.These flags have no effect if '-e' is not given:
@@ -85,6 +86,7 @@ if "%~1"=="--disable-gil" (set UseDisableGil=true) & shift & goto CheckOpts
 if "%~1"=="--test-marker" (set UseTestMarker=true) & shift & goto CheckOpts
 if "%~1"=="-V" shift & goto Version
 if "%~1"=="--regen" (set Regen=true) & shift & goto CheckOpts
+if "%~1"=="--experimental-jit" (set UseJIT=true) & shift & goto CheckOpts
 rem These use the actual property names used by MSBuild.  We could just let
 rem them in through the environment, but we specify them on the command line
 rem anyway for visibility so set defaults after this
@@ -176,6 +178,7 @@ echo on
  /p:IncludeSSL=%IncludeSSL% /p:IncludeTkinter=%IncludeTkinter%^
  /p:DisableGil=%UseDisableGil%^
  /p:UseTestMarker=%UseTestMarker% %GITProperty%^
+ /p:UseJIT=%UseJIT%^
  %1 %2 %3 %4 %5 %6 %7 %8 %9
 
 @echo off

PCbuild/pyproject.props

@@ -43,6 +43,7 @@
       <AdditionalIncludeDirectories>$(PySourcePath)Include;$(PySourcePath)Include\internal;$(PySourcePath)Include\internal\mimalloc;$(GeneratedPyConfigDir);$(PySourcePath)PC;%(AdditionalIncludeDirectories)</AdditionalIncludeDirectories>
       <PreprocessorDefinitions>WIN32;$(_Py3NamePreprocessorDefinition);$(_PlatformPreprocessorDefinition)$(_DebugPreprocessorDefinition)$(_PydPreprocessorDefinition)%(PreprocessorDefinitions)</PreprocessorDefinitions>
       <PreprocessorDefinitions Condition="'$(SupportPGO)' and ($(Configuration) == 'PGInstrument' or $(Configuration) == 'PGUpdate')">_Py_USING_PGO=1;%(PreprocessorDefinitions)</PreprocessorDefinitions>
+      <PreprocessorDefinitions Condition="'$(UseJIT)' == 'true'">_Py_JIT;%(PreprocessorDefinitions)</PreprocessorDefinitions>
 
       <Optimization>MaxSpeed</Optimization>
       <IntrinsicFunctions>true</IntrinsicFunctions>

PCbuild/pythoncore.vcxproj

@@ -246,6 +246,7 @@
     <ClInclude Include="..\Include\internal\pycore_initconfig.h" />
     <ClInclude Include="..\Include\internal\pycore_interp.h" />
     <ClInclude Include="..\Include\internal\pycore_intrinsics.h" />
+    <ClInclude Include="..\Include\internal\pycore_jit.h" />
     <ClInclude Include="..\Include\internal\pycore_list.h" />
     <ClInclude Include="..\Include\internal\pycore_llist.h" />
     <ClInclude Include="..\Include\internal\pycore_lock.h" />
@@ -578,6 +579,7 @@
     <ClCompile Include="..\Python\initconfig.c" />
     <ClCompile Include="..\Python\intrinsics.c" />
     <ClCompile Include="..\Python\instrumentation.c" />
+    <ClCompile Include="..\Python\jit.c" />
     <ClCompile Include="..\Python\legacy_tracing.c" />
     <ClCompile Include="..\Python\lock.c" />
     <ClCompile Include="..\Python\marshal.c" />
@@ -669,6 +671,14 @@
                       Overwrite="true"
                       Condition="$(PyConfigHText.TrimEnd()) != $(OldPyConfigH.TrimEnd())" />
   </Target>
+  <Target Name="_BuildJIT" BeforeTargets="PrepareForBuild" Condition="'$(UseJIT)' == 'true'" DependsOnTargets="_UpdatePyconfig;FindPythonForBuild">
+    <PropertyGroup>
+      <JITArgs Condition="$(Platform) == 'Win32'">i686-pc-windows-msvc</JITArgs>
+      <JITArgs Condition="$(Platform) == 'x64'">x86_64-pc-windows-msvc</JITArgs>
+      <JITArgs Condition="$(Configuration) == 'Debug'">$(JITArgs) --debug</JITArgs>
+    </PropertyGroup>
+    <Exec Command='$(PythonForBuild) "$(PySourcePath)Tools\jit\build.py" $(JITArgs)' WorkingDirectory="$(IntDir)"/>
+  </Target>
   <Target Name="_CopyPyconfig" Inputs="$(IntDir)pyconfig.h" Outputs="$(OutDir)pyconfig.h" AfterTargets="Build" DependsOnTargets="_UpdatePyconfig">
     <Copy SourceFiles="$(IntDir)pyconfig.h" DestinationFolder="$(OutDir)" />
   </Target>

PCbuild/pythoncore.vcxproj.filters

@@ -663,6 +663,9 @@
     <ClInclude Include="..\Include\internal\pycore_intrinsics.h">
       <Filter>Include\cpython</Filter>
     </ClInclude>
+    <ClInclude Include="..\Include\internal\pycore_jit.h">
+      <Filter>Include\internal</Filter>
+    </ClInclude>
     <ClInclude Include="..\Include\internal\pycore_list.h">
       <Filter>Include\internal</Filter>
     </ClInclude>