gh-131798: Split FOR_ITER_GEN into smaller uops

[Sacul0457]

I'm not sure how the optimizer knows the symbolic type on the first guard but it seems to work?

• Issue: Better uop coverage in the JIT optimizer #131798

[Fidget-Spinner]

I'm not sure how the optimizer knows the symbolic type on the first guard but it seems to work?

Uh oh. you just caught a bug in the JIT optimizer. Would you like to fix it?

The problem is that this line is wrong https://github.com/python/cpython/blob/main/Python/optimizer_symbols.c#L807 it should not be returning PyGen_Type, but rather NULL. Instead, _Py_uop_sym_get_probable_type should be the one returning PyGen_Type.

If you want to fix it, just move the lines around amend this small test to !, and copy the same line and change it to probable type https://github.com/python/cpython/blob/main/Python/optimizer_symbols.c#L2214

[Sacul0457]

Sure!
Should I fix it in this PR or create a new issue and open a separate PR?

[Fidget-Spinner]

Please open an issue and fix it in that PR. Thank you.

[Fidget-Spinner]

@Sacul0457 this might not be worth it, but splitting up _CHECK_AND_ALLOCATE_OBJECT into the CHECK and _ALLOCATE should be worth it (we can remove the redundant checks through same thing as _GUARD_TYPE_VERSION). I just merged a PR into main that should make this easier. Would you like to work on that instead?

[Sacul0457]

Sure :)

[Sacul0457]

Hi @Fidget-Spinner!
Does this look right?
For _CHECK_OBJECT it should look something like this

op(_CHECK_OBJECT, (type_version/2, callable, self_or_null, unused[oparg] -- callable, self_or_null, unused[oparg])) {
    PyObject *callable_o = PyStackRef_AsPyObjectBorrow(callable);
    EXIT_IF(!PyStackRef_IsNull(self_or_null));
    // ...
    // ...
    EXIT_IF(FT_ATOMIC_LOAD_UINT32_RELAXED(tp->tp_version_tag) != type_version);
}

Then _ALLOCATE_OBJECT would be everything else.

Then in optimizer_bytecodes.c:

op(_CHECK_OBJECT, (type_version/2, callable, self_or_null, unused[oparg] -- callable, self_or_null, unused[oparg])) {
    PyObject *probable_callable = sym_get_probable_value(callable);
    PyTypeObject *tp = (PyTypeObject *)probable_callable;
    if (tp->tp_version_tag == type_version && sym_is_not_null(self_or_null)) {
        // If the type version has not changed since we last saw it,
        // then we know this __init__ is definitely the same one as in the cache.
        ADD_OP(_NOP, 0, 0);
    }
    else {
        // do what _GUARD_TYPE_VERSION did in its else case
        sym_set_non_null(ctx, callable);
        sym_set_non_null(ctx, self_or_null);
    }
}

[Fidget-Spinner]

@Sacul0457 you need to do similar to what we're currently doing in _GUARD_TYPE_VERSION, which is to set the type watcher if it doesn't match, and leave the first type check around.

The refactored uop looks correct assuming you copied everything in correctly. Please open a PR.