GH-126910: Build/link the JIT shim in the Python interpreter

Include/internal/pycore_ceval.h

@@ -121,18 +121,11 @@ _PyEval_EvalFrame(PyThreadState *tstate, _PyInterpreterFrame *frame, int throwfl
 }
 
 #ifdef _Py_TIER2
-#ifdef _Py_JIT
-_Py_CODEUNIT *_Py_LazyJitShim(
-    struct _PyExecutorObject *current_executor, _PyInterpreterFrame *frame,
-    _PyStackRef *stack_pointer, PyThreadState *tstate
-);
-#else
 _Py_CODEUNIT *_PyTier2Interpreter(
     struct _PyExecutorObject *current_executor, _PyInterpreterFrame *frame,
     _PyStackRef *stack_pointer, PyThreadState *tstate
 );
 #endif
-#endif
 
 extern _PyJitEntryFuncPtr _Py_jit_entry;
 

Include/internal/pycore_jit.h

@@ -23,9 +23,13 @@ typedef _Py_CODEUNIT *(*jit_func)(
     _PyStackRef _tos_cache0, _PyStackRef _tos_cache1, _PyStackRef _tos_cache2
 );
 
+_Py_CODEUNIT *_PyJIT(
+    _PyExecutorObject *executor, _PyInterpreterFrame *frame,
+    _PyStackRef *stack_pointer, PyThreadState *tstate
+);
+
 int _PyJIT_Compile(_PyExecutorObject *executor, const _PyUOpInstruction *trace, size_t length);
 void _PyJIT_Free(_PyExecutorObject *executor);
-void _PyJIT_Fini(void);
 PyAPI_FUNC(int) _PyJIT_AddressInJitCode(PyInterpreterState *interp, uintptr_t addr);
 
 #endif  // _Py_JIT

Makefile.pre.in

@@ -290,6 +290,7 @@ LDLIBRARYDIR=   @LDLIBRARYDIR@
 INSTSONAME=	@INSTSONAME@
 LIBRARY_DEPS=	@LIBRARY_DEPS@
 LINK_PYTHON_DEPS=@LINK_PYTHON_DEPS@
+JIT_OBJS=	@JIT_SHIM_O@
 PY_ENABLE_SHARED=	@PY_ENABLE_SHARED@
 STATIC_LIBPYTHON=	@STATIC_LIBPYTHON@
 
@@ -469,6 +470,7 @@ PYTHON_OBJS=	\
 		Python/instruction_sequence.o \
 		Python/intrinsics.o \
 		Python/jit.o \
+		$(JIT_OBJS) \
 		Python/legacy_tracing.o \
 		Python/lock.o \
 		Python/marshal.o \
@@ -3203,21 +3205,37 @@ Python/emscripten_trampoline_inner.wasm: $(srcdir)/Python/emscripten_trampoline_
 Python/emscripten_trampoline_wasm.c: Python/emscripten_trampoline_inner.wasm
 	$(PYTHON_FOR_REGEN) $(srcdir)/Platforms/emscripten/prepare_external_wasm.py $< $@ getWasmTrampolineModule
 
+JIT_SHIM_BUILD_OBJS=	@JIT_SHIM_BUILD_O@
+JIT_BUILD_TARGETS=	jit_stencils.h @JIT_STENCILS_H@ $(JIT_SHIM_BUILD_OBJS)
+JIT_TARGETS=	$(JIT_BUILD_TARGETS) $(filter-out $(JIT_SHIM_BUILD_OBJS),$(JIT_OBJS))
+JIT_GENERATED_STAMP=	.jit-stamp
+
 JIT_DEPS = \
 		$(srcdir)/Tools/jit/*.c \
+		$(srcdir)/Tools/jit/*.h \
 		$(srcdir)/Tools/jit/*.py \
 		$(srcdir)/Python/executor_cases.c.h \
 		pyconfig.h
 
-jit_stencils.h @JIT_STENCILS_H@: $(JIT_DEPS)
+$(JIT_GENERATED_STAMP): $(JIT_DEPS)
 	@REGEN_JIT_COMMAND@
+	@touch $@
+
+$(JIT_BUILD_TARGETS): $(JIT_GENERATED_STAMP)
+	@if test ! -f "$@"; then \
+		rm -f $(JIT_GENERATED_STAMP); \
+		$(MAKE) $(JIT_GENERATED_STAMP); \
+		test -f "$@"; \
+	fi
+
+jit_shim-universal2-apple-darwin.o: jit_shim-aarch64-apple-darwin.o jit_shim-x86_64-apple-darwin.o
+	lipo -create -output $@ jit_shim-aarch64-apple-darwin.o jit_shim-x86_64-apple-darwin.o
 
 Python/jit.o: $(srcdir)/Python/jit.c @JIT_STENCILS_H@
 	$(CC) -c $(PY_CORE_CFLAGS) -o $@ $<
 
 .PHONY: regen-jit
-regen-jit:
-	@REGEN_JIT_COMMAND@
+regen-jit: $(JIT_TARGETS)
 
 # Some make's put the object file in the current directory
 .c.o:
@@ -3341,7 +3359,7 @@ clean-profile: clean-retain-profile clean-bolt
 # gh-141808: The JIT stencils are deliberately kept in clean-profile
 .PHONY: clean-jit-stencils
 clean-jit-stencils:
-	-rm -f jit_stencils*.h
+	-rm -f $(JIT_TARGETS) $(JIT_GENERATED_STAMP) jit_stencils*.h jit_shim*.o
 
 .PHONY: clean
 clean: clean-profile clean-jit-stencils

PCbuild/pyproject.props

@@ -12,8 +12,9 @@
     <IntDir>$(IntDir.Replace(`\\`, `\`))</IntDir>
     <GeneratedFrozenModulesDir>$(Py_IntDir)\$(MajorVersionNumber)$(MinorVersionNumber)_frozen\</GeneratedFrozenModulesDir>
     <GeneratedZlibNgDir>$(Py_IntDir)\$(MajorVersionNumber)$(MinorVersionNumber)$(ArchName)_$(Configuration)\zlib-ng\</GeneratedZlibNgDir>
-    <GeneratedJitStencilsDir>$(Py_IntDir)\$(MajorVersionNumber)$(MinorVersionNumber)_$(Configuration)</GeneratedJitStencilsDir>
-    <GeneratedJitStencilsDir Condition="$(Configuration) == 'PGUpdate'">$(Py_IntDir)\$(MajorVersionNumber)$(MinorVersionNumber)_PGInstrument</GeneratedJitStencilsDir>
+    <GeneratedJitStencilsDir>$(Py_IntDir)\$(MajorVersionNumber)$(MinorVersionNumber)_$(Configuration)\</GeneratedJitStencilsDir>
+    <GeneratedJitStencilsDir Condition="$(Configuration) == 'PGUpdate'">$(Py_IntDir)\$(MajorVersionNumber)$(MinorVersionNumber)_PGInstrument\</GeneratedJitStencilsDir>
+    <GeneratedJitStencilsDir>$(GeneratedJitStencilsDir.Replace(`\\`, `\`))</GeneratedJitStencilsDir>
     <TargetName Condition="'$(TargetName)' == ''">$(ProjectName)</TargetName>
     <TargetName>$(TargetName)$(PyDebugExt)</TargetName>
     <GenerateManifest>false</GenerateManifest>

PCbuild/pythoncore.vcxproj

@@ -115,6 +115,9 @@
     <Link>
       <AdditionalDependencies>version.lib;ws2_32.lib;pathcch.lib;bcrypt.lib;%(AdditionalDependencies)</AdditionalDependencies>
       <AdditionalDependencies Condition="$(IncludeExternals)">zlib-ng$(PyDebugExt).lib;%(AdditionalDependencies)</AdditionalDependencies>
+      <AdditionalDependencies Condition="'$(UseJIT)' == 'true' and $(Platform) == 'ARM64'">$(GeneratedJitStencilsDir)jit_shim-aarch64-pc-windows-msvc.o;%(AdditionalDependencies)</AdditionalDependencies>
+      <AdditionalDependencies Condition="'$(UseJIT)' == 'true' and $(Platform) == 'Win32'">$(GeneratedJitStencilsDir)jit_shim-i686-pc-windows-msvc.o;%(AdditionalDependencies)</AdditionalDependencies>
+      <AdditionalDependencies Condition="'$(UseJIT)' == 'true' and $(Platform) == 'x64'">$(GeneratedJitStencilsDir)jit_shim-x86_64-pc-windows-msvc.o;%(AdditionalDependencies)</AdditionalDependencies>
     </Link>
   </ItemDefinitionGroup>
   <ItemGroup>

PCbuild/regen.targets

@@ -35,6 +35,9 @@
     <_JITOutputs Include="$(GeneratedJitStencilsDir)jit_stencils-aarch64-pc-windows-msvc.h" Condition="$(Platform) == 'ARM64'"/>
     <_JITOutputs Include="$(GeneratedJitStencilsDir)jit_stencils-i686-pc-windows-msvc.h" Condition="$(Platform) == 'Win32'"/>
     <_JITOutputs Include="$(GeneratedJitStencilsDir)jit_stencils-x86_64-pc-windows-msvc.h" Condition="$(Platform) == 'x64'"/>
+    <_JITOutputs Include="$(GeneratedJitStencilsDir)jit_shim-aarch64-pc-windows-msvc.o" Condition="$(Platform) == 'ARM64'"/>
+    <_JITOutputs Include="$(GeneratedJitStencilsDir)jit_shim-i686-pc-windows-msvc.o" Condition="$(Platform) == 'Win32'"/>
+    <_JITOutputs Include="$(GeneratedJitStencilsDir)jit_shim-x86_64-pc-windows-msvc.o" Condition="$(Platform) == 'x64'"/>
     <_CasesSources Include="$(PySourcePath)Python\bytecodes.c;$(PySourcePath)Python\optimizer_bytecodes.c;"/>
     <_CasesOutputs Include="$(PySourcePath)Python\generated_cases.c.h;$(PySourcePath)Include\opcode_ids.h;$(PySourcePath)Include\internal\pycore_uop_ids.h;$(PySourcePath)Python\opcode_targets.h;$(PySourcePath)Include\internal\pycore_opcode_metadata.h;$(PySourcePath)Include\internal\pycore_uop_metadata.h;$(PySourcePath)Python\optimizer_cases.c.h;$(PySourcePath)Lib\_opcode_metadata.py"/>
     <_SbomSources Include="$(PySourcePath)PCbuild\get_externals.bat" />
@@ -129,7 +132,7 @@
       <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) --output-dir "$(GeneratedJitStencilsDir)" --pyconfig-dir "$(PySourcePath)PC" --llvm-version="$(LLVM_VERSION)" --llvm-tools-install-dir="$(LLVM_TOOLS_INSTALL_DIR)"'/>
+    <Exec Command='$(PythonForBuild) "$(PySourcePath)Tools\jit\build.py" $(JITArgs) --output-dir "$(GeneratedJitStencilsDir.TrimEnd(`\`))" --pyconfig-dir "$(PySourcePath)PC" --llvm-version="$(LLVM_VERSION)" --llvm-tools-install-dir="$(LLVM_TOOLS_INSTALL_DIR)"'/>
   </Target>
   <Target Name="_CleanJIT" AfterTargets="Clean">
     <Delete Files="@(_JITOutputs)"/>

Python/ceval.c

@@ -1305,7 +1305,7 @@ _PyEval_EvalFrameDefault(PyThreadState *tstate, _PyInterpreterFrame *frame, int
 }
 #ifdef _Py_TIER2
 #ifdef _Py_JIT
-_PyJitEntryFuncPtr _Py_jit_entry = _Py_LazyJitShim;
+_PyJitEntryFuncPtr _Py_jit_entry = _PyJIT;
 #else
 _PyJitEntryFuncPtr _Py_jit_entry = _PyTier2Interpreter;
 #endif

Python/jit.c

@@ -60,8 +60,6 @@ jit_error(const char *message)
     PyErr_Format(PyExc_RuntimeWarning, "JIT %s (%d)", message, hint);
 }
 
-static size_t _Py_jit_shim_size = 0;
-
 static int
 address_in_executor_array(_PyExecutorObject **ptrs, size_t count, uintptr_t addr)
 {
@@ -104,13 +102,6 @@ _PyJIT_AddressInJitCode(PyInterpreterState *interp, uintptr_t addr)
     if (interp == NULL) {
         return 0;
     }
-    if (_Py_jit_entry != _Py_LazyJitShim && _Py_jit_shim_size != 0) {
-        uintptr_t start = (uintptr_t)_Py_jit_entry;
-        uintptr_t end = start + _Py_jit_shim_size;
-        if (addr >= start && addr < end) {
-            return 1;
-        }
-    }
     if (address_in_executor_array(interp->executor_ptrs, interp->executor_count, addr)) {
         return 1;
     }
@@ -727,75 +718,6 @@ _PyJIT_Compile(_PyExecutorObject *executor, const _PyUOpInstruction trace[], siz
     return 0;
 }
 
-/* One-off compilation of the jit entry shim
- * We compile this once only as it effectively a normal
- * function, but we need to use the JIT because it needs
- * to understand the jit-specific calling convention.
- * Don't forget to call _PyJIT_Fini later!
- */
-static _PyJitEntryFuncPtr
-compile_shim(void)
-{
-    _PyExecutorObject dummy;
-    const StencilGroup *group;
-    size_t code_size = 0;
-    size_t data_size = 0;
-    jit_state state = {0};
-    group = &shim;
-    code_size += group->code_size;
-    data_size += group->data_size;
-    combine_symbol_mask(group->trampoline_mask, state.trampolines.mask);
-    combine_symbol_mask(group->got_mask, state.got_symbols.mask);
-    // Round up to the nearest page:
-    size_t page_size = get_page_size();
-    assert((page_size & (page_size - 1)) == 0);
-    size_t code_padding = DATA_ALIGN - ((code_size + state.trampolines.size) & (DATA_ALIGN - 1));
-    size_t padding = page_size - ((code_size + state.trampolines.size + code_padding + data_size + state.got_symbols.size) & (page_size - 1));
-    size_t total_size = code_size + state.trampolines.size + code_padding + data_size + state.got_symbols.size + padding;
-    unsigned char *memory = jit_alloc(total_size);
-    if (memory == NULL) {
-        return NULL;
-    }
-    unsigned char *code = memory;
-    state.trampolines.mem = memory + code_size;
-    unsigned char *data = memory + code_size + state.trampolines.size + code_padding;
-    state.got_symbols.mem = data + data_size;
-    // Compile the shim, which handles converting between the native
-    // calling convention and the calling convention used by jitted code
-    // (which may be different for efficiency reasons).
-    group = &shim;
-    group->emit(code, data, &dummy, NULL, &state);
-    code += group->code_size;
-    data += group->data_size;
-    assert(code == memory + code_size);
-    assert(data == memory + code_size + state.trampolines.size + code_padding + data_size);
-    if (mark_executable(memory, total_size)) {
-        jit_free(memory, total_size);
-        return NULL;
-    }
-    _Py_jit_shim_size = total_size;
-    return (_PyJitEntryFuncPtr)memory;
-}
-
-static PyMutex lazy_jit_mutex = { 0 };
-
-_Py_CODEUNIT *
-_Py_LazyJitShim(
-    _PyExecutorObject *executor, _PyInterpreterFrame *frame, _PyStackRef *stack_pointer, PyThreadState *tstate
-) {
-    PyMutex_Lock(&lazy_jit_mutex);
-    if (_Py_jit_entry == _Py_LazyJitShim) {
-        _PyJitEntryFuncPtr shim = compile_shim();
-        if (shim == NULL) {
-            PyMutex_Unlock(&lazy_jit_mutex);
-            Py_FatalError("Cannot allocate core JIT code");
-        }
-        _Py_jit_entry = shim;
-    }
-    PyMutex_Unlock(&lazy_jit_mutex);
-    return _Py_jit_entry(executor, frame, stack_pointer, tstate);
-}
-
 // Free executor's memory allocated with _PyJIT_Compile
 void
 _PyJIT_Free(_PyExecutorObject *executor)
@@ -812,22 +734,4 @@ _PyJIT_Free(_PyExecutorObject *executor)
     }
 }
 
-// Free shim memory allocated with compile_shim
-void
-_PyJIT_Fini(void)
-{
-    PyMutex_Lock(&lazy_jit_mutex);
-    unsigned char *memory = (unsigned char *)_Py_jit_entry;
-    size_t size = _Py_jit_shim_size;
-    if (size) {
-        _Py_jit_entry = _Py_LazyJitShim;
-        _Py_jit_shim_size = 0;
-        if (jit_free(memory, size)) {
-            PyErr_FormatUnraisable("Exception ignored while "
-                                   "freeing JIT entry code");
-        }
-    }
-    PyMutex_Unlock(&lazy_jit_mutex);
-}
-
 #endif  // _Py_JIT

Python/pylifecycle.c

@@ -37,9 +37,6 @@
 #include "pycore_uniqueid.h"      // _PyObject_FinalizeUniqueIdPool()
 #include "pycore_warnings.h"      // _PyWarnings_InitState()
 #include "pycore_weakref.h"       // _PyWeakref_GET_REF()
-#ifdef _Py_JIT
-#include "pycore_jit.h"           // _PyJIT_Fini()
-#endif
 
 #if defined(PYMALLOC_USE_HUGEPAGES) && defined(MS_WINDOWS)
 #include <Windows.h>
@@ -2531,11 +2528,6 @@ _Py_Finalize(_PyRuntimeState *runtime)
 
     finalize_interp_clear(tstate);
 
-#ifdef _Py_JIT
-    /* Free JIT shim memory */
-    _PyJIT_Fini();
-#endif
-
 #ifdef Py_TRACE_REFS
     /* Display addresses (& refcnts) of all objects still alive.
      * An address can be used to find the repr of the object, printed

Python/pystate.c

@@ -489,11 +489,6 @@ free_interpreter(PyInterpreterState *interp)
 static inline int check_interpreter_whence(long);
 #endif
 
-extern _Py_CODEUNIT *
-_Py_LazyJitShim(
-    struct _PyExecutorObject *exec, _PyInterpreterFrame *frame, _PyStackRef *stack_pointer, PyThreadState *tstate
-);
-
 /* Get the interpreter state to a minimal consistent state.
    Further init happens in pylifecycle.c before it can be used.
    All fields not initialized here are expected to be zeroed out,