Specialize FOR_ITER during GET_ITER

[sweeneyde]

FOR_ITER already assumes that the iterator on the stack is safe for (Py_TYPE(it)->tp_iternext)(it), but perhaps a specialization of FOR_ITER could assume the specific type of the iterator. This could be verified by the GET_ITER call and then the specialized FOR_ITER that follows it wouldn't have to do any checks.

Since multiple frames can be running the same code at overlapping times, the only way I can think of to make sure types don't conflict is to only adapt once, and never adapt again afterward. Something like this:

TARGET(GET_ITER_ADAPT) {
    _Py_CODEUNIT *this_instr = next_instr - 1;
    // only execute this instruction once.
    _Py_SET_OPCODE(*this_instr, GET_ITER);

    PyObject *iterable = TOP();
    PyObject *iter = PyObject_GetIter(iterable);
    Py_DECREF(iterable);
    SET_TOP(iter);
    if (iter == NULL) {
        goto error;
    }
    if (_Py_OPCODE(*next_instr) == FOR_ITER) {
        PyTypeObject *tp = Py_TYPE(iter);
        if (tp == &PyRangeIter_Type) {
            _Py_SET_OPCODE(*this_instr, GET_ITER_MAINTAIN_SPECIALIZED);
            _Py_SET_OPCODE(*next_instr, FOR_ITER_RANGE);
        }
        else if (tp == &PyListIter_Type) {
            _Py_SET_OPCODE(*this_instr, GET_ITER_MAINTAIN_SPECIALIZED);
            _Py_SET_OPCODE(*next_instr, FOR_ITER_LIST);
        }
    }
    DISPATCH();
}

static const PyTypeObject *const for_iter_types[256] = {
    [FOR_ITER_LIST] = &PyListIter_Type,
    [FOR_ITER_RANGE] = &PyRangeIter_Type,
};

TARGET(GET_ITER_MAINTAIN_SPECIALIZED) {
    PyObject *iterable = TOP();
    PyObject *iter = PyObject_GetIter(iterable);
    Py_DECREF(iterable);
    SET_TOP(iter);
    if (iter == NULL) {
        goto error;
    }
    _Py_CODEUNIT next_op = _Py_OPCODE(*next_instr);
    if (Py_TYPE(iter) == for_iter_types[next_op]) {
        NOTRACE_DISPATCH();
    }
    else {
        // deopt.
        _Py_CODEUNIT *this_instr = next_instr - 1;
        _Py_SET_OPCODE(*this_instr, GET_ITER);
        _Py_SET_OPCODE(*next_instr, FOR_ITER);
        NOTRACE_DISPATCH();
    }
}

TARGET(FOR_ITER_RANGE) {
    assert(Py_TYPE(TOP()) == &PyRangeIter_Type); // assert, not deopt_if
#if 0
    // XXX implement this.
#else
    JUMP_TO_INSTRUCTION(FOR_ITER);
#endif
}

TARGET(FOR_ITER_LIST) {
    assert(Py_TYPE(TOP()) == &PyListIter_Type); // assert, not deopt_if
#if 0
    // XXX implement this.
#else
    JUMP_TO_INSTRUCTION(FOR_ITER);
#endif
}

The downside: it ceases to be adaptive in some situations.
The upside: very little overhead (no caches, no counters, no DEOPT_IF checks). Not much more overhead than today in almost any situation.

An aside: is there a reason why comprehensions do comprehension_code(iter(x)) rather than doing comprehension_code(x) and having the GET_ITER in the comprehension's bytecode? Because if we implemented this idea, comprehensions would have to do their own GET_ITER in order to benefit from specialization.

Thoughts?

[sweeneyde]

I have a draft: python/cpython#32340

[sweeneyde]

If the loss of adaptiveness is a huge deal, maybe there's some sort of reference counting strategy that would work, where the instruction counts how many frames are currently executing the for loop, and so it can specialize if that gets down to zero.

[markshannon]

This seems complex compared to specializing FOR_ITER, and adds another special case.

Is DEOPT_IF(Py_TYPE(it) != _PyListIter_Type) that expensive?