Type inference in closures

General Usage
1

Is there any way to make Julia properly infer types in this code:

function f()
    x = 0
    l = Threads.SpinLock()

    Threads.@threads for i in 1:Threads.nthreads()
        Threads.lock(l)
        x = x + i
        Threads.unlock(l)
    end
    x
end

Currently @code_warntype f() indicates that x is boxed because @threads creates a closure. I tried let trick, but it does not work in this case.

Does anyone have some advice?

2

Using a reference to x allocates (on creating the Ref), but it comes up clean in @code_warntype:

f() = f!(Ref(0))
function f!(x::Ref{Int})
    x[] = 0
    l = Threads.SpinLock()
    Threads.@threads for i in 1:Threads.nthreads()
        Threads.lock(l)
        x[] = x[] + i
        Threads.unlock(l)
    end
    x[]
end

Seems a little faster.

Maybe there’s a better solution.

Also, the let trick doesn’t really work here, because then the x in the closure isn’t updating the x outside of it, and I just get 0.

2 Likes
3

Good solution if we know that we are working with Int. Thank you.
However, actually this was a simplified example from a package where the type can be arbitrary.

Then I avoid Ref on Julia 0.7 as Ref(a) where a is an array does not work correctly. Maybe using Base.RefValue although it is not exported could be the general approach?

4

For more complicated examples, I’d just replace the Ref with a mutable struct that holds everything.
If something can be entirely arbitrary, I guess make that field parametric?

mutable struct DefinitelyNotARef{T}
    ref::T
end
2 Likes
5

If the let trick is used you put x as the last expression in let block so it gets returned (but it does not solve the problem here though).

1 Like
6

Thanks - this is what I ended up doing in the end :).

1 Like
8

Ah, of course! Should’ve of thought of that.