I would like to mutate a fixed size array inside a guvectorized method. Because closure variables are frozen I have no choice but to create it inside the method itself which kill performances. Can someone think of a workaround?
Hi @fdv1
Do you have a piece of code that demonstrates, or roughly outlines, what you want to do? That’d probably help with working out if there’s something that can be done to help. Thanks 
Sorry for not providing this in the first place, here it is.
n = 10
@guvectorize([(float64[:], float64)], '(r) ->()')
def test(state, result):
a = [0.] * n
a[0] = 1.
for i in range(1, n):
a[i] = a[i - 1] * np.exp(state[i])
# compute result using a ...
Hi,
I can’t really help with your closure problem, but one thing that I see here is that you are using a list for a.
I would think that you should already see a performance increase if you make a an actual array
n = 10
@guvectorize([(float64[:], float64)], '(r) ->()')
def test(state, result):
a = np.zeros(n)
a[0] = 1.
for i in range(1, n):
a[i] = a[i - 1] * np.exp(state[i])
# compute result using a ...
Beyond that, I am not sure that it would generally be advisable to have a shared memory location for a across function calls. That smells a bit like an accident waiting to happen :’)
If you really need to share a, maybe one could hack a little bit and pass a as an input to the function. That time you could allocate memory once and then pass it around.
However, this does not exactly seem to be encouraged: Creating NumPy universal functions — Numba 0.52.0.dev0+274.g626b40e-py3.7-linux-x86_64.egg documentation
The array a is simply an intermediate variable, I rewrote my code to make my intention clearer. The issue with this version is that it requires to compute all the states upfront (they won’t fit in memory).
# f is an arbitrary method which takes an array of double
@jit
def f(a):
return np.sum(a)
@jit
def test(states):
result = np.zeros(states.shape[0])
a = np.zeros(n)
a[0] = 1.
for j in range(states.shape[0]):
state = states[j]
for i in range(1, n):
a[i] = a[i - 1] * np.exp(state[i])
result[j] = f(a)
return result
Mhm, I am not sure I can follow. Are you saying that states would be too large for memory if you compute them in advance and then pass them to test? If yes, then I fail to see how vectorizing the function would help with that problem. (In that case could you compute the current state on demand inside the j loop instead of passing them to test ?)
states would be too large indeed. I’d rather not jit the outer loop as I am also doing some other things which don’t need to be jited. Here what my code looks like.
for i in range(nb_simulations):
# doing some complicated but not computationally intensive things here...
state = compute_state(state) # compute_state cannot be merged with test (already vectorized using numpy methods)
results = test(state)
I’m afraid that I don’t entirely understand what’s going on in the above. Where does the fixed sized array that needs mutating come into this? I’m not sure where the closure is? and why elect for guvectorize? does the dimensionality of state vary?
If the generation of all of the states is causing undue memory pressure, if the allocation for a is loop invariant then perhaps allocate it outside the loop in nb_simulations and pass it in to test each time at which point you could zero it out and reuse it? Something like this:
from numba import njit, guvectorize, float64, void
import numpy as np
n = 10
nb_simulations = 4
@njit
def f(a):
return np.sum(a)
@guvectorize([void(float64[:], float64[:], float64[:])], '(r),(DUMMY)->()')
def test(state, temp, result):
temp[:] = 0.
temp[0] = 1.
for i in range(1, n):
temp[i] = temp[i - 1] * np.exp(state[i])
result[()] = f(temp)
def compute_state(i):
return i * 1e-6 * np.arange(n * 4).reshape((4, n))
buf = np.empty((n,))
result = np.zeros((nb_simulations,))
for i in range(nb_simulations):
state = compute_state(i)
test(state, buf, result)
print(result)
there’s a lot of dimensions in the above that seem unusual, in part because I’m not sure there’s enough information in the previous examples to work it out all, but hopefully this’ll provide a starting point.
Thanks for bearing with me, I’m really embarassed of not making myself clearer. I just realized that many of the dimensions are not relevant to this problem.
Basically I need to be able to use a buffer inside a vectorized method without allocating it every time. One possibility would be to pass buffer as argument, but that’s not terribly elegant nor efficient.
from numba import vectorize, float64, njit
n = 10
@njit
def f(a):
return np.sum(a)
@vectorize([float64(float64, float64)])
def f(x, y):
buffer = np.zeros(n)
for i in range(1, n):
buffer[i] = max(buffer[i - 1] * x, y)
return f(buffer)
print(f(np.array((1.1, 2.3)), np.array((1.3, 1.2))))
Thanks, I see. You can do this, but it’s “dangerous” as it’s baking the runtime address of a global array allocation into the generated machine code, this would definitely prevent caching. The following is based on:
- python - How to pass array pointer to Numba function? - Stack Overflow where I answered a similar question previously for david-zwicker
- which came from a gitter chat here
- xref my original gist showing how to do it.
from numba import vectorize, float64, njit, carray, types
from numba.extending import intrinsic
import ctypes
import numpy as np
n = 10
global_buf = np.zeros(n)
addr = global_buf.ctypes.data
@intrinsic
def addr2ptr(typingctx, src):
sig = types.voidptr(src)
def codegen(cgctx, builder, sig, args):
return builder.inttoptr(args[0], cgctx.get_value_type(sig.return_type))
return sig, codegen
@njit
def f(a):
return np.sum(a)
@vectorize([float64(float64, float64)])
def vectf(x, y):
buf = carray(addr2ptr(addr), global_buf.shape, dtype=global_buf.dtype)
buf[:] = 0.
for i in range(1, n):
buf[i] = max(buf[i - 1] * x, y)
return f(buf)
print(vectf(np.array((1.1, 2.3)), np.array((1.3, 1.2))))
# Run the original that allocates a buffer to check
@vectorize([float64(float64, float64)])
def original_f(x, y):
buffer = np.zeros(n)
for i in range(1, n):
buffer[i] = max(buffer[i - 1] * x, y)
return f(buffer)
print(original_f(np.array((1.1, 2.3)), np.array((1.3, 1.2))))
I’d really consider the merits of doing this vs. the, maybe inelegant, but much safer and clearer, passing of the scratch space as an argument.
Addendum… should also note that something needs to keep the buffer alive throughout the duration of the JIT function calls. The compiled function is relying on that address being valid.
Thank you very much for your answer! I agree with you this seems a borderline dangerous. It would be much nicer if we could mutate a closure variable instead (hence my initial question ;-).
It would indeed be nice to support globals as non-constant but unfortunately that’s an incredibly tricky problem to solve, and as such is a limitation of Numba at present. Passing in a scratch space perhaps isn’t so bad given the performance gain it gets you. Also, if n is always 10, you could just unroll the loop, all the temporaries would be on the stack then?
You can also use this trick to generate a pseudo-dynamic unroll, so long as the range is resolvable as compile time constant:
from numba import njit, types, literal_unroll, literally, vectorize, float64
from numba.extending import intrinsic, overload
import ctypes
import numpy as np
n = 10
def fixed_range(start, stop):
pass
@overload(fixed_range, prefer_literal=True)
def ol_fixed_range(start, stop):
fixed = tuple(range(start.literal_value, stop.literal_value))
return lambda start, stop: fixed
@njit
def f(a):
return np.sum(a)
@njit
def wrap(x, y):
buffer = np.zeros(n)
for i in literal_unroll(fixed_range(1, n)):
buffer[i] = max(buffer[i - 1] * x, y)
return f(buffer)
@vectorize([float64(float64, float64)])
def original_f(x, y):
return wrap(x, y)
print(original_f(np.array((1.1, 2.3)), np.array((1.3, 1.2))))
print(wrap.inspect_asm(wrap.signatures[0]))
Note: need to figure out how to get rid of the buffer still
Thanks for your suggestion! It will take me a bit of time to digest it 
.
I wouldn’t have suspected it would be so tricky to make globals mutable indeed (even copies?)