Cython close to hand-crafted C code for generators
Mittwoch, April 6th, 2011I did a couple of experiments compiling itertools with the new generator support in Cython. In CPython, the itertools module is actually written in hand tuned C and does very little computation in its generators, so I knew it would be hard to reach with generated code. But Cython does a pretty good job.
Something as trivial as chain() is exactly as fast as in the C implementation, but compared to the more than 60 lines of C code, it is certainly a lot more readable in Cython:
def chain(*iterables):
"""Make an iterator that returns elements from the first iterable
until it is exhausted, then proceeds to the next iterable, until
all of the iterables are exhausted. Used for treating consecutive
sequences as a single sequence.
"""
for it in iterables:
for element in it:
yield element
Other functions, like islice(), are faster in C, partly because CPython actually takes a couple of shortcuts, e.g. by only looking up the iterator slot method once. You cannot do that in Python code, and I wanted to keep the implementation compatible with regular Python. Specifically, the C speed advantage for islice() is currently about 30-50% in general, although the Cython implementation can also be up to 10% faster for some cases, e.g. when extracting only a couple of items from the middle of a longer sequence. The C implementation is about 90 lines, here is the Cython implementation:
import sys
import cython
# Python 2/3 compatibility
_max_size = cython.declare(cython.Py_ssize_t,
getattr(sys, "maxsize", getattr(sys, "maxint", None)))
@cython.locals(i=cython.Py_ssize_t, nexti=cython.Py_ssize_t,
start=cython.Py_ssize_t, stop=cython.Py_ssize_t, step=cython.Py_ssize_t)
def islice(iterable, *args):
"""Make an iterator that returns selected elements from the
iterable. If start is non-zero, then elements from the iterable
are skipped until start is reached. Afterward, elements are
returned consecutively unless step is set higher than one which
results in items being skipped. If stop is None, then iteration
continues until the iterator is exhausted, if at all; otherwise,
it stops at the specified position. Unlike regular slicing,
islice() does not support negative values for start, stop, or
step. Can be used to extract related fields from data where the
internal structure has been flattened (for example, a multi-line
report may list a name field on every third line).
"""
s = slice(*args)
start = s.start or 0
stop = s.stop or _max_size
step = s.step or 1
if start < 0:
raise ValueError("...")
if step < 1:
raise ValueError("...")
if start >= stop:
return
nexti = start
for i, element in enumerate(iterable):
if i == nexti:
yield element
nexti += step
if nexti >= stop or nexti < 0:
return
Here is one that is conceptually quite simple: count(). I had to optimise it quite a bit, because the iteration code in the C code is extremely tight. Even the tuned version below runs about 10% slower than the hand tuned C version, which is about 230 lines long.
@cython.locals(i=cython.Py_ssize_t)
def count(n=0):
"""Make an iterator that returns consecutive integers starting
with n. If not specified n defaults to zero. Often used as an argument to imap()
to generate consecutive data points. Also, used with zip() to add
sequence numbers.
"""
try:
i = n
except OverflowError:
i = _max_size # skip i-loop
else:
n = _max_size # first value after i-loop
while i < _max_size:
yield i
i += 1
while True:
yield n
n += 1
Note that all of the above generators execute in the order of microseconds, so even a slow-down of 50% will likely not be measurable in real world code.
So far, I did not try any of the more fancy functions in itertools (those that actually do something). The Cython project has announced a Google Summer of Code project with exactly the intent to rewrite some of the C stdlib modules of CPython in pure Python code with Cython compiler hints. So I leave this exercise to interested readers for now.
