Skip to content
Prev Previous commit
Next Next commit
Cleanup
  • Loading branch information
brandtbucher committed Jul 16, 2025
commit a322ad46cdf45e5a94ef92e476a2e47d64e26a1e
80 changes: 49 additions & 31 deletions Lib/test/test_jit_stencils.py
Original file line number Diff line number Diff line change
@@ -1,49 +1,67 @@

import pathlib
import shlex
import sys
import sysconfig
import tempfile
import test.support
import unittest

import test.support.script_helper

import unittest

_CPYTHON = pathlib.Path(test.support.REPO_ROOT).resolve()
Comment thread
tomasr8 marked this conversation as resolved.
_TOOLS_JIT = _CPYTHON / "Tools" / "jit"
_TOOLS_JIT_TEST = _TOOLS_JIT / "test"
_TOOLS_JIT_TEST_TEST_EXECUTOR_CASES_C_H = _TOOLS_JIT_TEST / "test_executor_cases.c.h"
_TOOLS_JIT_BUILD_PY = _TOOLS_JIT / "build.py"

@unittest.skipIf(test.support.Py_DEBUG, "XXX")
@unittest.skipUnless(sys._jit.is_available(), "XXX")
@unittest.skipIf(test.support.Py_GIL_DISABLED, "XXX")
@unittest.skipUnless(sysconfig.is_python_build(), "XXX")

@test.support.cpython_only
@unittest.skipIf(test.support.Py_DEBUG, "Debug stencils aren't tested.")
@unittest.skipIf(test.support.Py_GIL_DISABLED, "Free-threaded stencils aren't tested.")
@unittest.skipUnless(sysconfig.is_python_build(), "Requires a local Python build.")
class TestJITStencils(unittest.TestCase):

def _build_jit_stencils(self, target: str) -> str:
with tempfile.TemporaryDirectory() as work:

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think os_helper.temp_dir is a bit more common in tests but either is fine :)

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hm, I'll take a look!

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is it, like, better? Haha.

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Couldn't tell you 😅 but it seems quite common. Though both are used so it probably doesn't matter :)

jit_stencils_h = pathlib.Path(work, f"jit_stencils-{target}.h").resolve()
pyconfig_h = pathlib.Path(sysconfig.get_config_h_filename()).resolve()
result, args = test.support.script_helper.run_python_until_end(
_TOOLS_JIT_BUILD_PY,
"--input-file", _TOOLS_JIT_TEST_TEST_EXECUTOR_CASES_C_H,

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is just for my understanding. is this needed in case we need to regenerate this during development and we need to test it, right?

"--output-dir", jit_stencils_h.parent,
"--pyconfig-dir", pyconfig_h.parent,
target,
__isolated=False,
)
if result.rc:
self.skipTest(f"Build failed: {shlex.join(map(str, args))}")
body = jit_stencils_h.read_text()
# Strip out two lines of header comments:
_, _, body = body.split("\n", 2)
return body

def _check_jit_stencils(
self, expected: str, actual: str, test_jit_stencils_h: pathlib.Path
) -> None:
try:
self.assertEqual(expected.strip("\n"), actual.strip("\n"))
except AssertionError as e:
# Make it easy to re-validate the expected output:
relative = test_jit_stencils_h.relative_to(_CPYTHON)
message = f"If this is expected, replace {relative} with:"
banner = "=" * len(message)
e.add_note("\n".join([banner, message, banner]))
e.add_note(actual)
raise

def test_jit_stencils(self):
self.maxDiff = None
found = False
pyconfig_dir = pathlib.Path(sysconfig.get_config_h_filename()).parent
with tempfile.TemporaryDirectory() as work:
output_dir = pathlib.Path(work).resolve()
for test_jit_stencils_h in sorted(_TOOLS_JIT_TEST.glob("test_jit_stencils-*.h")):
>
jit_stencils_h = output_dir / f"jit_stencils-{target}.h"
with self.subTest(target):
# relative = jit_stencils_h.relative_to(_CPYTHON)
result, args = test.support.script_helper.run_python_until_end(
_TOOLS_JIT_BUILD_PY,
"--input-file", _TOOLS_JIT_TEST / "test_executor_cases.c.h",
"--output-dir", output_dir,
"--pyconfig-dir", pyconfig_dir,
target,
__isolated=False
)
if result.rc:
self.skipTest(shlex.join(map(str, args)))
found = True
expected = test_jit_stencils_h.read_text()
actual = "".join(jit_stencils_h.read_text().splitlines(True)[3:])
self.assertEqual(expected, actual)
self.assertTrue(found, "No JIT stencil tests run!")
for test_jit_stencils_h in _TOOLS_JIT_TEST.glob("test_jit_stencils-*.h"):
>
with self.subTest(target):
expected = test_jit_stencils_h.read_text()
actual = self._build_jit_stencils(target)
found = True
self._check_jit_stencils(expected, actual, test_jit_stencils_h)
# This is a local build. If the JIT is available, at least one test should run:
assert found or not sys._jit.is_available(), "No JIT stencils built!"
7 changes: 7 additions & 0 deletions Tools/jit/_targets.py
Original file line number Diff line number Diff line change
Expand Up @@ -566,36 +566,43 @@ def get_target(host: str) -> _COFF32 | _COFF64 | _ELF | _MachO:
optimizer: type[_optimizers.Optimizer]
target: _COFF32 | _COFF64 | _ELF | _MachO
if re.fullmatch(r"aarch64-apple-darwin.*", host):
host = "aarch64-apple-darwin"
condition = "defined(__aarch64__) && defined(__APPLE__)"
optimizer = _optimizers.OptimizerAArch64
condition, optimizer=optimizer)
elif re.fullmatch(r"aarch64-pc-windows-msvc", host):
host = "aarch64-pc-windows-msvc"
args = ["-fms-runtime-lib=dll", "-fplt"]
condition = "defined(_M_ARM64)"
optimizer = _optimizers.OptimizerAArch64
condition, args=args, optimizer=optimizer)
elif re.fullmatch(r"aarch64-.*-linux-gnu", host):
host = "aarch64-unknown-linux-gnu"
# -mno-outline-atomics: Keep intrinsics from being emitted.
args = ["-fpic", "-mno-outline-atomics"]
condition = "defined(__aarch64__) && defined(__linux__)"
optimizer = _optimizers.OptimizerAArch64
condition, args=args, optimizer=optimizer)
elif re.fullmatch(r"i686-pc-windows-msvc", host):
host = "i686-pc-windows-msvc"
# -Wno-ignored-attributes: __attribute__((preserve_none)) is not supported here.
args = ["-DPy_NO_ENABLE_SHARED", "-Wno-ignored-attributes"]
optimizer = _optimizers.OptimizerX86
condition = "defined(_M_IX86)"
condition, args=args, optimizer=optimizer)
elif re.fullmatch(r"x86_64-apple-darwin.*", host):
host = "x86_64-apple-darwin"
condition = "defined(__x86_64__) && defined(__APPLE__)"
optimizer = _optimizers.OptimizerX86
condition, optimizer=optimizer)
elif re.fullmatch(r"x86_64-pc-windows-msvc", host):
host = "x86_64-pc-windows-msvc"
args = ["-fms-runtime-lib=dll"]
condition = "defined(_M_X64)"
optimizer = _optimizers.OptimizerX86
condition, args=args, optimizer=optimizer)
elif re.fullmatch(r"x86_64-.*-linux-gnu", host):
host = "x86_64-unknown-linux-gnu"
args = ["-fno-pic", "-mcmodel=medium", "-mlarge-data-threshold=0"]
condition = "defined(__x86_64__) && defined(__linux__)"
optimizer = _optimizers.OptimizerX86
Expand Down
28 changes: 15 additions & 13 deletions Tools/jit/test/test_executor_cases.c.h
Original file line number Diff line number Diff line change
@@ -1,27 +1,29 @@
case 0: {
// Zero-length jumps should be removed:
break;
}

case 1: {
if (CURRENT_OPARG()) {
JUMP_TO_JUMP_TARGET();
// -Os duplicates less code than -O3:
PyAPI_DATA(bool) sausage;
PyAPI_DATA(bool) spammed;
PyAPI_FUNC(void) order_eggs_and_bacon(void);
PyAPI_FUNC(void) order_eggs_sausage_and_bacon(void);
if (!sausage) {
order_eggs_and_bacon();
}
else {
order_eggs_sausage_and_bacon();
}
spammed = false;
break;
}

case 2: {
if (CURRENT_OPARG()) {
// The assembly optimizer inverts hot branches:
PyAPI_DATA(bool) spam;
if (spam) {
JUMP_TO_ERROR();
}
break;
}

case 3: {
GOTO_TIER_ONE((void *)CURRENT_OPERAND0() + CURRENT_TARGET());
break;
}

case 4: {
GOTO_TIER_TWO((void *)CURRENT_OPERAND1());
break;
}
150 changes: 60 additions & 90 deletions Tools/jit/test/test_jit_stencils-aarch64-unknown-linux-gnu.h
Original file line number Diff line number Diff line change
Expand Up @@ -61,109 +61,81 @@ emit_1(
const _PyUOpInstruction *instruction, jit_state *state)
{
// 0000000000000000 <_JIT_ENTRY>:
// 0: 90000008 adrp x8, 0x0 <_JIT_ENTRY>
// 0000000000000000: R_AARCH64_ADR_GOT_PAGE _JIT_OPARG
// 4: f9400108 ldr x8, [x8]
// 0000000000000004: R_AARCH64_LD64_GOT_LO12_NC _JIT_OPARG
// 8: 72003d1f tst w8, #0xffff
// c: 54000040 b.eq 0x14 <_JIT_ENTRY+0x14>
// 10: 14000000 b 0x10 <_JIT_ENTRY+0x10>
// 0000000000000010: R_AARCH64_JUMP26 _JIT_JUMP_TARGET
const unsigned char code_body[20] = {
// 0: a9bf7bfd stp x29, x30, [sp, #-0x10]!
// 4: 90000008 adrp x8, 0x0 <_JIT_ENTRY>
// 0000000000000004: R_AARCH64_ADR_GOT_PAGE sausage
// 8: 910003fd mov x29, sp
// c: f9400108 ldr x8, [x8]
// 000000000000000c: R_AARCH64_LD64_GOT_LO12_NC sausage
// 10: 39400108 ldrb w8, [x8]
// 14: 36000088 tbz w8, #0x0, 0x24 <_JIT_ENTRY+0x24>
// 18: 90000008 adrp x8, 0x0 <_JIT_ENTRY>
// 0000000000000018: R_AARCH64_ADR_GOT_PAGE order_eggs_sausage_and_bacon
// 1c: f9400108 ldr x8, [x8]
// 000000000000001c: R_AARCH64_LD64_GOT_LO12_NC order_eggs_sausage_and_bacon
// 20: 14000003 b 0x2c <_JIT_ENTRY+0x2c>
// 24: 90000008 adrp x8, 0x0 <_JIT_ENTRY>
// 0000000000000024: R_AARCH64_ADR_GOT_PAGE order_eggs_and_bacon
// 28: f9400108 ldr x8, [x8]
// 0000000000000028: R_AARCH64_LD64_GOT_LO12_NC order_eggs_and_bacon
// 2c: d63f0100 blr x8
// 30: 90000008 adrp x8, 0x0 <_JIT_ENTRY>
// 0000000000000030: R_AARCH64_ADR_GOT_PAGE spammed
// 34: f9400108 ldr x8, [x8]
// 0000000000000034: R_AARCH64_LD64_GOT_LO12_NC spammed
// 38: 3900011f strb wzr, [x8]
// 3c: a8c17bfd ldp x29, x30, [sp], #0x10
const unsigned char code_body[64] = {
0xfd, 0x7b, 0xbf, 0xa9, 0x08, 0x00, 0x00, 0x90,
0xfd, 0x03, 0x00, 0x91, 0x08, 0x01, 0x40, 0xf9,
0x08, 0x01, 0x40, 0x39, 0x88, 0x00, 0x00, 0x36,
0x08, 0x00, 0x00, 0x90, 0x08, 0x01, 0x40, 0xf9,
0x1f, 0x3d, 0x00, 0x72, 0x40, 0x00, 0x00, 0x54,
0x00, 0x00, 0x00, 0x14,
};
// 0: OPARG
patch_64(data + 0x0, instruction->oparg);
memcpy(code, code_body, sizeof(code_body));
patch_aarch64_33rx(code + 0x0, (uintptr_t)data);
patch_aarch64_26r(code + 0x10, state->instruction_starts[instruction->jump_target]);
}

void
emit_2(
unsigned char *code, unsigned char *data, _PyExecutorObject *executor,
const _PyUOpInstruction *instruction, jit_state *state)
{
// 0000000000000000 <_JIT_ENTRY>:
// 0: 90000008 adrp x8, 0x0 <_JIT_ENTRY>
// 0000000000000000: R_AARCH64_ADR_GOT_PAGE _JIT_OPARG
// 4: f9400108 ldr x8, [x8]
// 0000000000000004: R_AARCH64_LD64_GOT_LO12_NC _JIT_OPARG
// 8: 72003d1f tst w8, #0xffff
// c: 54000040 b.eq 0x14 <_JIT_ENTRY+0x14>
// 10: 14000000 b 0x10 <_JIT_ENTRY+0x10>
// 0000000000000010: R_AARCH64_JUMP26 _JIT_ERROR_TARGET
const unsigned char code_body[20] = {
0x03, 0x00, 0x00, 0x14, 0x08, 0x00, 0x00, 0x90,
0x08, 0x01, 0x40, 0xf9, 0x00, 0x01, 0x3f, 0xd6,
0x08, 0x00, 0x00, 0x90, 0x08, 0x01, 0x40, 0xf9,
0x1f, 0x3d, 0x00, 0x72, 0x40, 0x00, 0x00, 0x54,
0x00, 0x00, 0x00, 0x14,
0x1f, 0x01, 0x00, 0x39, 0xfd, 0x7b, 0xc1, 0xa8,
};
// 0: OPARG
patch_64(data + 0x0, instruction->oparg);
// 0: &sausage+0x0
// 8: &order_eggs_sausage_and_bacon+0x0
// 10: &order_eggs_and_bacon+0x0
// 18: &spammed+0x0
patch_64(data + 0x0, (uintptr_t)&sausage);
patch_64(data + 0x8, (uintptr_t)&order_eggs_sausage_and_bacon);
patch_64(data + 0x10, (uintptr_t)&order_eggs_and_bacon);
patch_64(data + 0x18, (uintptr_t)&spammed);
memcpy(code, code_body, sizeof(code_body));
patch_aarch64_33rx(code + 0x0, (uintptr_t)data);
patch_aarch64_26r(code + 0x10, state->instruction_starts[instruction->error_target]);
patch_aarch64_21rx(code + 0x4, (uintptr_t)data);
patch_aarch64_12x(code + 0xc, (uintptr_t)data);
patch_aarch64_33rx(code + 0x18, (uintptr_t)data + 0x8);
patch_aarch64_33rx(code + 0x24, (uintptr_t)data + 0x10);
patch_aarch64_33rx(code + 0x30, (uintptr_t)data + 0x18);
}

void
emit_3(
unsigned char *code, unsigned char *data, _PyExecutorObject *executor,
const _PyUOpInstruction *instruction, jit_state *state)
{
// 0000000000000000 <_JIT_ENTRY>:
// 0: 90000008 adrp x8, 0x0 <_JIT_ENTRY>
// 0000000000000000: R_AARCH64_ADR_GOT_PAGE _JIT_TARGET
// 4: 90000009 adrp x9, 0x0 <_JIT_ENTRY>
// 0000000000000004: R_AARCH64_ADR_GOT_PAGE _JIT_OPERAND0
// 8: f9400108 ldr x8, [x8]
// 0000000000000008: R_AARCH64_LD64_GOT_LO12_NC _JIT_TARGET
// c: f9400129 ldr x9, [x9]
// 000000000000000c: R_AARCH64_LD64_GOT_LO12_NC _JIT_OPERAND0
// 10: f9008adf str xzr, [x22, #0x110]
// 14: f9002295 str x21, [x20, #0x40]
// 18: 8b284120 add x0, x9, w8, uxtw
// 1c: d65f03c0 ret
const unsigned char code_body[32] = {
0x08, 0x00, 0x00, 0x90, 0x09, 0x00, 0x00, 0x90,
0x08, 0x01, 0x40, 0xf9, 0x29, 0x01, 0x40, 0xf9,
0xdf, 0x8a, 0x00, 0xf9, 0x95, 0x22, 0x00, 0xf9,
0x20, 0x41, 0x28, 0x8b, 0xc0, 0x03, 0x5f, 0xd6,
};
// 0: TARGET
// 8: OPERAND0
patch_64(data + 0x0, instruction->target);
patch_64(data + 0x8, instruction->operand0);
memcpy(code, code_body, sizeof(code_body));
patch_aarch64_21rx(code + 0x0, (uintptr_t)data);
patch_aarch64_21rx(code + 0x4, (uintptr_t)data + 0x8);
patch_aarch64_12x(code + 0x8, (uintptr_t)data);
patch_aarch64_12x(code + 0xc, (uintptr_t)data + 0x8);
}

void
emit_4(
emit_2(
unsigned char *code, unsigned char *data, _PyExecutorObject *executor,
const _PyUOpInstruction *instruction, jit_state *state)
{
// 0000000000000000 <_JIT_ENTRY>:
// 0: 90000008 adrp x8, 0x0 <_JIT_ENTRY>
// 0000000000000000: R_AARCH64_ADR_GOT_PAGE _JIT_OPERAND1
// 0000000000000000: R_AARCH64_ADR_GOT_PAGE spam
// 4: f9400108 ldr x8, [x8]
// 0000000000000004: R_AARCH64_LD64_GOT_LO12_NC _JIT_OPERAND1
// 8: f9403d00 ldr x0, [x8, #0x78]
// c: f9008ac8 str x8, [x22, #0x110]
// 10: d61f0000 br x0
const unsigned char code_body[20] = {
// 0000000000000004: R_AARCH64_LD64_GOT_LO12_NC spam
// 8: 39400108 ldrb w8, [x8]
// c: 7100051f cmp w8, #0x1
// 10: 54000041 b.ne 0x18 <_JIT_ENTRY+0x18>
// 14: 14000000 b 0x14 <_JIT_ENTRY+0x14>
// 0000000000000014: R_AARCH64_JUMP26 _JIT_ERROR_TARGET
const unsigned char code_body[24] = {
0x08, 0x00, 0x00, 0x90, 0x08, 0x01, 0x40, 0xf9,
0x00, 0x3d, 0x40, 0xf9, 0xc8, 0x8a, 0x00, 0xf9,
0x00, 0x00, 0x1f, 0xd6,
0x08, 0x01, 0x40, 0x39, 0x1f, 0x05, 0x00, 0x71,
0x41, 0x00, 0x00, 0x54, 0x00, 0x00, 0x00, 0x14,
};
// 0: OPERAND1
patch_64(data + 0x0, instruction->operand1);
// 0: &spam+0x0
patch_64(data + 0x0, (uintptr_t)&spam);
memcpy(code, code_body, sizeof(code_body));
patch_aarch64_33rx(code + 0x0, (uintptr_t)data);
patch_aarch64_26r(code + 0x14, state->instruction_starts[instruction->error_target]);
}

static_assert(SYMBOL_MASK_WORDS >= 1, "SYMBOL_MASK_WORDS too small");
Expand All @@ -181,10 +153,8 @@ static const StencilGroup shim = {emit_shim, 104, 0, {0}};

static const StencilGroup stencil_groups[MAX_UOP_ID + 1] = {
[0] = {emit_0, 0, 0, {0}},
[1] = {emit_1, 20, 8, {0}},
[2] = {emit_2, 20, 8, {0}},
[3] = {emit_3, 32, 16, {0}},
[4] = {emit_4, 20, 8, {0}},
[1] = {emit_1, 64, 32, {0}},
[2] = {emit_2, 24, 8, {0}},
};

static const void * const symbols_map[1] = {
Expand Down
Loading