Outlining of expression bodies into own function

[KristofferC]

A common "trick" to deal with heterogeneous data that is looped over is to introduce a function barrier that acts as an "inference point" so that dynamic dispatch is only done once per element. When working on the element, the type is known and if the work is significant enough, the cost of dynamic dispatch is negligible. While this usually works well, it is pretty annoying to have to split out the new function, pass all the arguments to the new function etc. Therefore, I wonder if it is possible to introduce a syntax / feature whereby you could achieve this "outlining" without having to manually do it. The result should be the equivalent of moving the body of, for example a for loop, into it's own function.

The code below shows three different implementation of summing over elements of an array that contain 3 vectors with different element types.

function f1(vs)
    z = 0.0
    for v in vs
        for x in v
            z += x
        end
    end
    return z
end

function f2(vs)
    z = 0.0
    for v in vs
        (v -> begin
            for x in v
                z += x
            end
        end)(v)
    end
    return z
end

function f3(vs)
    z = 0.0
    for v in vs
        z += _f3(v)
    end
    return z
end

function _f3(v)
    _z = 0.0
    for x in v
        _z += x
    end
    return _z
end


using BenchmarkTools
vs = Any[rand(10^5), rand(Int, 10^5), rand(Float32, 10^5)];

@btime f1(vs);
@btime f2(vs);
@btime f3(vs);

with the result:

julia> @btime f1(vs);
  53.785 ms (1198470 allocations: 22.86 MiB)

julia> @btime f2(vs);
  10.569 ms (600004 allocations: 9.16 MiB)

julia> @btime f3(vs);
  254.428 μs (6 allocations: 96 bytes)

Using a closure seems faster than the normal loop but is clearly much slower than explicitly outlined version.

I realize this possible ties closely into closure lowering so if this is considered a dup of #15276, please close.

[yuyichao]

I don't think this is a dup but this can possibly use closures if #15276 is fixed. It can also be a typeinf pass though doing that in compiler is known to be much harder than inlining....

[TsurHerman]

slight modification to f2 brings it to the same level of f3

function f2(vs)
    z = Ref(0.0)
    for v in vs
        (v -> begin
            @inbounds for x in v
                z.x += x
            end
        end)(v)
    end
    return z.x
end

taking the anonymous function outside the loop , makes things slightly better (less allocations)

function f2(vs)
    z = Ref(0.0)
    lambda = v -> begin
        @inbounds for x in v
            z.x += x
        end
    end
    for v in vs
        lambda(v)
    end
    return z.x
end

[Orbots]

Would be great the closure worked best ( and automatically ). I tried the FastClosures.jl macro, which wraps the closures in a let block. It's on par with f3.

using FastClosures

function f4(vs)
    z = 0.0
    for v in vs
        z = (@closure ((v,z) -> begin
             for x in v
                 z += x
             end
             return z
        end))(v,z)
    end
    return z
end

results:

julia> @btime f1(vs);
  38.970 ms (1198470 allocations: 22.86 MiB)

julia> @btime f2(vs);
  15.061 ms (600004 allocations: 9.16 MiB)

julia> @btime f3(vs);
  244.302 μs (6 allocations: 96 bytes)

julia> @btime f4(vs);
  244.102 μs (3 allocations: 48 bytes)

[StefanKarpinski]

Yes, we'll definitely get there, but this kind of pure optimization work is largely on hold while we get all the breaking stuff finished and wrapped up for 1.0.

[KristofferC]

Could probably be done in a package.