GH-135379: Support limited scalar replacement for replicated uops in the JIT code generator. (GH-135563)

* Use it to support efficient specializations of COPY and SWAP in the JIT.
This commit is contained in:
Mark Shannon 2025-06-17 13:43:09 +01:00 committed by GitHub
parent a9e66a7c50
commit 8dd8b5c2f0
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
9 changed files with 313 additions and 198 deletions

View file

@ -4946,8 +4946,7 @@ dummy_func(
res = PyStackRef_FromPyObjectSteal(res_o);
}
pure inst(COPY, (bottom, unused[oparg-1] -- bottom, unused[oparg-1], top)) {
assert(oparg > 0);
pure replicate(1:4) inst(COPY, (bottom, unused[oparg-1] -- bottom, unused[oparg-1], top)) {
top = PyStackRef_DUP(bottom);
}
@ -4980,12 +4979,11 @@ dummy_func(
macro(BINARY_OP) = _SPECIALIZE_BINARY_OP + unused/4 + _BINARY_OP;
pure inst(SWAP, (bottom, unused[oparg-2], top --
pure replicate(2:4) inst(SWAP, (bottom, unused[oparg-2], top --
bottom, unused[oparg-2], top)) {
_PyStackRef temp = bottom;
bottom = top;
top = temp;
assert(oparg >= 2);
}
inst(INSTRUMENTED_LINE, ( -- )) {

View file

@ -6763,12 +6763,44 @@
break;
}
case _COPY_1: {
_PyStackRef bottom;
_PyStackRef top;
bottom = stack_pointer[-1];
top = PyStackRef_DUP(bottom);
stack_pointer[0] = top;
stack_pointer += 1;
assert(WITHIN_STACK_BOUNDS());
break;
}
case _COPY_2: {
_PyStackRef bottom;
_PyStackRef top;
bottom = stack_pointer[-2];
top = PyStackRef_DUP(bottom);
stack_pointer[0] = top;
stack_pointer += 1;
assert(WITHIN_STACK_BOUNDS());
break;
}
case _COPY_3: {
_PyStackRef bottom;
_PyStackRef top;
bottom = stack_pointer[-3];
top = PyStackRef_DUP(bottom);
stack_pointer[0] = top;
stack_pointer += 1;
assert(WITHIN_STACK_BOUNDS());
break;
}
case _COPY: {
_PyStackRef bottom;
_PyStackRef top;
oparg = CURRENT_OPARG();
bottom = stack_pointer[-1 - (oparg-1)];
assert(oparg > 0);
top = PyStackRef_DUP(bottom);
stack_pointer[0] = top;
stack_pointer += 1;
@ -6808,6 +6840,32 @@
break;
}
case _SWAP_2: {
_PyStackRef top;
_PyStackRef bottom;
top = stack_pointer[-1];
bottom = stack_pointer[-2];
_PyStackRef temp = bottom;
bottom = top;
top = temp;
stack_pointer[-2] = bottom;
stack_pointer[-1] = top;
break;
}
case _SWAP_3: {
_PyStackRef top;
_PyStackRef bottom;
top = stack_pointer[-1];
bottom = stack_pointer[-3];
_PyStackRef temp = bottom;
bottom = top;
top = temp;
stack_pointer[-3] = bottom;
stack_pointer[-1] = top;
break;
}
case _SWAP: {
_PyStackRef top;
_PyStackRef bottom;
@ -6817,7 +6875,6 @@
_PyStackRef temp = bottom;
bottom = top;
top = temp;
assert(oparg >= 2);
stack_pointer[-2 - (oparg-2)] = bottom;
stack_pointer[-1] = top;
break;

View file

@ -5228,7 +5228,6 @@
_PyStackRef bottom;
_PyStackRef top;
bottom = stack_pointer[-1 - (oparg-1)];
assert(oparg > 0);
top = PyStackRef_DUP(bottom);
stack_pointer[0] = top;
stack_pointer += 1;
@ -11568,7 +11567,6 @@
_PyStackRef temp = bottom;
bottom = top;
top = temp;
assert(oparg >= 2);
stack_pointer[-2 - (oparg-2)] = bottom;
stack_pointer[-1] = top;
DISPATCH();

View file

@ -1292,8 +1292,8 @@ uop_optimize(
for (int pc = 0; pc < length; pc++) {
int opcode = buffer[pc].opcode;
int oparg = buffer[pc].oparg;
if (oparg < _PyUop_Replication[opcode]) {
buffer[pc].opcode = opcode + oparg + 1;
if (oparg < _PyUop_Replication[opcode].stop && oparg >= _PyUop_Replication[opcode].start) {
buffer[pc].opcode = opcode + oparg + 1 - _PyUop_Replication[opcode].start;
assert(strncmp(_PyOpcode_uop_name[buffer[pc].opcode], _PyOpcode_uop_name[opcode], strlen(_PyOpcode_uop_name[opcode])) == 0);
}
else if (is_terminator(&buffer[pc])) {