Segmentation fault with generator function

[krono-i2]

I'm trying to jit-ing a function in nopython environment.
Find below the function definition.

@jit('int64[:](int64, int64)', nopython=True)
def numba_set_partitions_helper(n, k):
    if k == 1:
        yield np.zeros(n, dtype = np.int_)
    elif n == k:
        yield np.arange(n, dtype = np.int_)
    else:
        for p in numba_set_partitions_helper(n - 1, k - 1):
            yield np.append(0, p + 1)
        for p in numba_set_partitions_helper(n - 1, k):
            for i in range(k):
                yield np.append(i, p)
    return np.array([0], dtype=np.int_)

The output is.

Fatal Python error: Segmentation fault

[apmasell]

Confirmed. Complete reproducer:

import numba as nb
import numpy as np

@nb.jit('int64[:](int64, int64)', nopython=True)
def numba_set_partitions_helper(n, k):
    if k == 1:
        yield np.zeros(n, dtype = np.int_)
    elif n == k:
        yield np.arange(n, dtype = np.int_)
    else:
        for p in numba_set_partitions_helper(n - 1, k - 1):
            yield np.append(0, p + 1)
        for p in numba_set_partitions_helper(n - 1, k):
            for i in range(k):
                yield np.append(i, p)
    return np.array([0], dtype=np.int_)

for x in numba_set_partitions_helper(3, 2):
   print(x)

[apmasell]

Further debugging shows that there's some kind of stack corruption going on here. Any recursive generator seems to trigger the problem.

Simplifying to this:

import numba as nb
import numpy as np

@nb.jit('int64(int64)', nopython=True)
def countamajig(n):
    if n < 1:
        yield 0
    else:
        print( countamajig(n - 1))
        for p in range(0, n):
            yield p
    return 0

for x in countamajig(3):
  print(x)

triggers a fault:

==16823== Process terminating with default action of signal 11 (SIGSEGV)
==16823==  Access not within mapped region at address 0x0
==16823==    at 0x1B05B333: generator_iternext (in /home/andre/miniconda3/envs/numba/lib/python3.10/site-packages/numba/_dynfunc.cpython-310-x86_64-linux-gnu.so)
...

So, as a stop gap, if you can write a non-recursive implementation, it should avoid the problem.

[apmasell]

Alternative workaround is to use forceobj=True in the decorator as this will trigger creating the generator as a Python object and it will be heap allocated.

[apmasell]

After discussion, generators are very fragile and we think the best course of action is to disallow nested generators (including recursive ones).