mypyc / mypyc

Compile type annotated Python to fast C extensions
1.74k stars 46 forks source link

Faster argument type-checking #1012

Open nickdrozd opened 1 year ago

nickdrozd commented 1 year ago

Feature

Faster compiled code and also less compiled code, especially for union types.

Pitch

Mypyc executes runtime type-checks on function arguments. Here is an example of the C code that implements a check:

    PyObject *arg_some_arg;
    if (PyLong_Check(obj_some_arg))
        arg_some_arg = obj_some_arg;
    else {
        arg_some_arg = NULL;
    }
    if (arg_some_arg != NULL) goto __LL1;

The jump at the end is the "good label", meaning that the check succeeded and the rest of the code can continue.

But this code is somewhat inefficient. If type-check function PyLong_Check succeeds, arg_some_arg gets the value obj_some_arg, and otherwise it gets NULL. Then arg_some_arg is unconditionally checked for NULL. But in the case that PyLong_Check has failed it will have just been explicitly set to NULL, so checking for NULL is extra work for nothing.

If I'm reading it right, this could be rewritten so that NULL is only checked when the type-check function succeeds:

    PyObject *arg_some_arg = NULL;
    if (PyLong_Check(obj_some_arg)) {
        arg_some_arg = obj_some_arg;
        if (arg_some_arg != NULL) goto __LL1;
    }

The more types an argument can be, the more unnecessary NULL checks there will be. Consider an argument like some_arg: int | str | float | list | dict, for which Mypyc generates this type-checking code:

    PyObject *arg_some_arg;
    if (PyLong_Check(obj_some_arg))
        arg_some_arg = obj_some_arg;
    else {
        arg_some_arg = NULL;
    }
    if (arg_some_arg != NULL) goto __LL1;
    if (PyUnicode_Check(obj_some_arg))
        arg_some_arg = obj_some_arg;
    else {
        arg_some_arg = NULL;
    }
    if (arg_some_arg != NULL) goto __LL1;
    if (CPyFloat_Check(obj_some_arg))
        arg_some_arg = obj_some_arg;
    else {
        arg_some_arg = NULL;
    }
    if (arg_some_arg != NULL) goto __LL1;
    if (PyList_Check(obj_some_arg))
        arg_some_arg = obj_some_arg;
    else {
        arg_some_arg = NULL;
    }
    if (arg_some_arg != NULL) goto __LL1;
    if (PyDict_Check(obj_some_arg))
        arg_some_arg = obj_some_arg;
    else {
        arg_some_arg = NULL;
    }
    if (arg_some_arg != NULL) goto __LL1;

The change I'm proposing here would result in faster compiled code and also less compiled code, and these benefits would more pronounced for union types.

(This is a small-scale optimization. I'm assuming that the type-checking code is mostly correct and necessary.)

nickdrozd commented 1 year ago

(Copied from https://github.com/python/mypy/issues/16040)

JukkaL commented 1 year ago

Yeah, we could do better here. C compilers are pretty good at optimizing away redundant checks, though, but at least simplifying the code could speed up compilation.

nickdrozd commented 1 year ago

@JukkaL Are you open to a PR to implement this change?

nickdrozd commented 1 year ago

The logic is handled by emit_cast, which is complicated and a little inefficient itself.

JukkaL commented 1 year ago

I'm open to a PR if it doesn't make the code significantly harder to read or maintain, and isn't too big. (This is something that I'd definitely like to improve eventually, but it's not a priority for me yet so I don't wish to dedicate much time at the moment on it.)