GH-113464: Add a JIT backend for tier 2 (GH-113465)

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). See Tools/jit/README.md for more information on how to install the required build-time tooling.
python · Jan 29, 2024 · f6d9e59 · f6d9e59
1 parent f7c05d7
commit f6d9e59
Show file tree

Hide file tree

Showing 29 changed files with 1,738 additions and 5 deletions.
diff --git a/.github/workflows/jit.yml b/.github/workflows/jit.yml
@@ -0,0 +1,112 @@
+name: JIT
+on:
+  pull_request:
+    paths: '**jit**'
+  push:
+    paths: '**jit**'
+  workflow_dispatch:
+jobs:
+  jit:
+    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 --no-progress --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 4
+          ./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 4
+          make clean --jobs 4
+          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 4' || '' }}
+          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 4
+          ./python -m test --exclude ${{ matrix.exclude }} --multiprocess 0 --timeout 3600 --verbose2 --verbose3
diff --git a/.github/workflows/mypy.yml b/.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",
         ]

diff --git a/.gitignore b/.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

diff --git a/Include/cpython/optimizer.h b/Include/cpython/optimizer.h
@@ -39,6 +39,8 @@ typedef struct {
 typedef struct _PyExecutorObject {
     PyObject_VAR_HEAD
     _PyVMData vm_data; /* Used by the VM, but opaque to the optimizer */
+    void *jit_code;
+    size_t jit_size;
     _PyUOpInstruction trace[1];
 } _PyExecutorObject;
 

diff --git a/Include/internal/pycore_jit.h b/Include/internal/pycore_jit.h
@@ -0,0 +1,25 @@
+#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
+
+typedef _Py_CODEUNIT *(*jit_func)(_PyInterpreterFrame *frame, PyObject **stack_pointer, PyThreadState *tstate);
+
+int _PyJIT_Compile(_PyExecutorObject *executor, _PyUOpInstruction *trace, size_t length);
+void _PyJIT_Free(_PyExecutorObject *executor);
+
+#endif  // _Py_JIT
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif // !Py_INTERNAL_JIT_H
diff --git a/Include/internal/pycore_object.h b/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

diff --git a/Makefile.pre.in b/Makefile.pre.in
@@ -433,6 +433,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 \
@@ -1365,7 +1366,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"
@@ -1846,6 +1847,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 \
@@ -2641,6 +2643,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 $@ $<
@@ -2733,6 +2741,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

diff --git a/Misc/NEWS.d/next/Core and Builtins/2023-12-24-03-25-28.gh-issue-113464.dvjQmA.rst b/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>`_
diff --git a/PCbuild/_freeze_module.vcxproj b/PCbuild/_freeze_module.vcxproj
@@ -224,6 +224,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" />

diff --git a/PCbuild/_freeze_module.vcxproj.filters b/PCbuild/_freeze_module.vcxproj.filters
@@ -250,6 +250,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>

diff --git a/PCbuild/build.bat b/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

diff --git a/PCbuild/pythoncore.vcxproj b/PCbuild/pythoncore.vcxproj
@@ -104,6 +104,7 @@
       <AdditionalIncludeDirectories Condition="$(IncludeExternals)">$(zlibDir);%(AdditionalIncludeDirectories)</AdditionalIncludeDirectories>
       <PreprocessorDefinitions>_USRDLL;Py_BUILD_CORE;Py_BUILD_CORE_BUILTIN;Py_ENABLE_SHARED;MS_DLL_ID="$(SysWinVer)";%(PreprocessorDefinitions)</PreprocessorDefinitions>
       <PreprocessorDefinitions Condition="$(IncludeExternals)">_Py_HAVE_ZLIB;%(PreprocessorDefinitions)</PreprocessorDefinitions>
+      <PreprocessorDefinitions Condition="'$(UseJIT)' == 'true'">_Py_JIT;%(PreprocessorDefinitions)</PreprocessorDefinitions>
     </ClCompile>
     <Link>
       <AdditionalDependencies>version.lib;ws2_32.lib;pathcch.lib;bcrypt.lib;%(AdditionalDependencies)</AdditionalDependencies>
@@ -247,6 +248,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" />
@@ -585,6 +587,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" />

diff --git a/PCbuild/pythoncore.vcxproj.filters b/PCbuild/pythoncore.vcxproj.filters
@@ -669,6 +669,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>
@@ -1337,6 +1340,9 @@
     <ClCompile Include="..\Python\instrumentation.c">
       <Filter>Source Files</Filter>
     </ClCompile>
+    <ClCompile Include="..\Python\jit.c">
+      <Filter>Python</Filter>
+    </ClCompile>
     <ClCompile Include="..\Python\legacy_tracing.c">
       <Filter>Source Files</Filter>
     </ClCompile>

diff --git a/PCbuild/regen.targets b/PCbuild/regen.targets
@@ -28,6 +28,9 @@
     </_TokenOutputs>
     <_KeywordSources Include="$(PySourcePath)Grammar\python.gram;$(PySourcePath)Grammar\Tokens" />
     <_KeywordOutputs Include="$(PySourcePath)Lib\keyword.py" />
+    <!-- Taken from _Target._compute_digest in Tools\jit\_targets.py: -->
+    <_JITSources Include="$(PySourcePath)Python\executor_cases.c.h;$(GeneratedPyConfigDir)pyconfig.h;$(PySourcePath)Tools\jit\**"/>
+    <_JITOutputs Include="$(GeneratedPyConfigDir)jit_stencils.h"/>
   </ItemGroup>
 
   <Target Name="_TouchRegenSources" Condition="$(ForceRegen) == 'true'">
@@ -76,10 +79,28 @@
     <Exec Command="$(PythonForBuild) Tools\build\generate_global_objects.py"
           WorkingDirectory="$(PySourcePath)" />
   </Target>
+
+  <Target Name="_RegenJIT"
+          Condition="'$(UseJIT)' == 'true'"
+          DependsOnTargets="_UpdatePyconfig;FindPythonForBuild"
+          Inputs="@(_JITSources)"
+          Outputs="@(_JITOutputs)">
+    <PropertyGroup>
+      <JITArgs Condition="$(Platform) == 'ARM64'">aarch64-pc-windows-msvc</JITArgs>
+      <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="$(GeneratedPyConfigDir)"/>
+  </Target>
 
-  <Target Name="Regen"
+  <Target Name="_RegenNoPGUpdate"
           Condition="$(Configuration) != 'PGUpdate'"
           DependsOnTargets="_TouchRegenSources;_RegenPegen;_RegenAST_H;_RegenTokens;_RegenKeywords;_RegenGlobalObjects">
+  </Target>
+
+  <Target Name="Regen" DependsOnTargets="_RegenNoPGUpdate;_RegenJIT">
     <Message Text="Generated sources are up to date" Importance="high" />
   </Target>