gh-129231: group executable JIT code in memory

Author: diegorusso

Python/jit.c

@@ -501,8 +501,8 @@ _PyJIT_Compile(_PyExecutorObject *executor, const _PyUOpInstruction trace[], siz
     // Round up to the nearest page:
     size_t page_size = get_page_size();
     assert((page_size & (page_size - 1)) == 0);
-    size_t padding = page_size - ((code_size + data_size + state.trampolines.size) & (page_size - 1));
-    size_t total_size = code_size + data_size + state.trampolines.size + padding;
+    size_t padding = page_size - ((code_size + state.trampolines.size + data_size) & (page_size - 1));
+    size_t total_size = code_size + state.trampolines.size + data_size  + padding;
     unsigned char *memory = jit_alloc(total_size);
     if (memory == NULL) {
         return -1;
@@ -516,8 +516,8 @@ _PyJIT_Compile(_PyExecutorObject *executor, const _PyUOpInstruction trace[], siz
     }
     // Loop again to emit the code:
     unsigned char *code = memory;
-    unsigned char *data = memory + code_size;
-    state.trampolines.mem = memory + code_size + data_size;
+    state.trampolines.mem = memory + code_size;
+    unsigned char *data = memory + code_size + state.trampolines.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).
@@ -539,7 +539,7 @@ _PyJIT_Compile(_PyExecutorObject *executor, const _PyUOpInstruction trace[], siz
     code += group->code_size;
     data += group->data_size;
     assert(code == memory + code_size);
-    assert(data == memory + code_size + data_size);
+    assert(data == memory + code_size + state.trampolines.size + data_size);
 #ifdef MAP_JIT
     pthread_jit_write_protect_np(1);
 #endif

Tools/jit/_stencils.py

@@ -364,7 +364,7 @@ def _get_trampoline_mask(self) -> str:
 
     def as_c(self, opname: str) -> str:
         """Dump this hole as a StencilGroup initializer."""
-        return f"{{emit_{opname}, {len(self.code.body)}, {len(self.data.body)}, {self._get_trampoline_mask()}}}"
+        return f"{{emit_{opname}, {len(self.code.body)}, {self._get_trampoline_mask()}, {len(self.data.body)}}}"
 
 
 def symbol_to_value(symbol: str) -> tuple[HoleValue, str | None]:

Tools/jit/_writer.py

@@ -18,8 +18,8 @@ def _dump_footer(
     yield "        unsigned char *code, unsigned char *data, _PyExecutorObject *executor,"
     yield "        const _PyUOpInstruction *instruction, jit_state *state);"
     yield "    size_t code_size;"
-    yield "    size_t data_size;"
     yield "    symbol_mask trampoline_mask;"
+    yield "    size_t data_size;"
     yield "} StencilGroup;"
     yield ""
     yield f"static const StencilGroup shim = {groups['shim'].as_c('shim')};"