Type inference with a functional programing style

General Usage Performance
1

I’m having trouble with type inference. I’d like to be able to write this sort of thing and have it run quickly:

function silly_last(vs::AbstractVector)
    x = first(vs)

    map(vs) do v
        x = v
    end

    x
end

Without having to resort to

function clever_last(vs::AbstractVector)
    x = first(vs)

    for v in vs
        x = v
    end

    x
end

Ultimately, I think my problem comes down to this issue:

function foo()
    x = 1

    bar(v::Int) = x = v

    x
end

@code_warntype foo()
Variables
  #self#::Core.Const(foo)
  bar::var"#bar#380"
  x@_3::Core.Box
  x@_4::Union{}

Body::Any
1 ─      (x@_3 = Core.Box())
│        Core.setfield!(x@_3, :contents, 1)
│        (bar = %new(Main.:(var"#bar#380"), x@_3))
│   %4 = Core.isdefined(x@_3, :contents)::Bool
└──      goto #3 if not %4
2 ─      goto #4
3 ─      Core.NewvarNode(:(x@_4))
└──      x@_4
4 ┄ %9 = Core.getfield(x@_3, :contents)::Any
└──      return %9

The type of x has to be a Core.Box to accommodate anything that bar puts in it. So I’m going to ask a common question: why can’t the compiler figure out the type of x?

2

I think the problem is more that the inner function is both capturing the value of x and modifying it, and modifying a captured variable (which, thus, is not an argument of the function), is like modifying a variable in global scope (I’m not saying that the compiler in these specific examples could not do better, but in any case seems to be a risky operation). You can solve the problem by simply changing the label used in the inner function, separating clearly the labels of the captured values from the variables of the function scope. For example (not sure if a good example):

Modifying the captured value:

julia> function foo(v)
           x = 1
           bar(v::Int) = x = v
           x
       end
foo (generic function with 2 methods)

julia> @code_warntype foo([1,2])
Variables
  #self#::Core.Const(foo)
  v::Vector{Int64}
  bar::var"#bar#14"
  x@_4::Core.Box
  x@_5::Union{}

Body::Any
1 ─      (x@_4 = Core.Box())
│        Core.setfield!(x@_4, :contents, 1)
│        (bar = %new(Main.:(var"#bar#14"), x@_4))
│   %4 = Core.isdefined(x@_4, :contents)::Bool
└──      goto #3 if not %4
2 ─      goto #4
3 ─      Core.NewvarNode(:(x@_5))
└──      x@_5
4 ┄ %9 = Core.getfield(x@_4, :contents)::Any
└──      return %9

Not modifying the captured value:

julia> function foo(v)
           x = 1
           bar(v::Int) = y = v
           x = bar(v)
           x
       end
foo (generic function with 2 methods)

julia> @code_warntype foo([1,2])
Variables
  #self#::Core.Const(foo)
  v::Vector{Int64}
  bar::var"#bar#15"
  x::Int64

Body::Union{}
1 ─     (x = 1)
│       (bar = %new(Main.:(var"#bar#15")))
│       (x = (bar)(v))
└──     Core.Const(:(return x))

julia>

I would say that having closures modifying captured values makes the code hard to follow, besides these problems.

3 Likes
3

I’m aware of that, but unfortunately, my silly_last relies on the side effects.

In those examples, yes they do. And yet they also have a place. In a map(x) do ... end block the closure is clearly only used in map and easy for me to reason about: the stuff in the closure happens once for every item in the collection and then the closure is gone forever. It would be nice if the compiler could reason about this the same way supported functional programming as well as the syntax does (see also tail-call optimization).

1 Like
4

I got accustomed to avoid that pattern, so that would be my advice. Maybe it is a good idea to show an example where the pattern really seems helpful (not an artificially simple case where obvious solutions, like the for loop, are good enough). With such an example you may get better advice from more experienced people.

1 Like
5

Thanks! I’ll wait till I stumble upon one in the wild and report back.

1 Like
6

Here is a possible solution. It appears to be type-stable but has an unnecessary allocation. This solution assumes that all the entries of the vector have the same type as the first entry.

 function last1(vs::AbstractVector)
       last1_(vs,first(vs))
 end

function last1_(vs::AbstractVector{T}, vsfirst::T) where T
       x = Array{T,0}(undef)
       x[] = vsfirst
       map(vs) do v
           let x = x
               x[] = v
           end
       end
       x[]
    end
7

The Right Way to write a generic loop in a functional style is foldl not map.

julia> right(_, x) = x;

julia> foldl(right, 1:3)
3

It’s free from the inference issue.

julia> using Test

julia> @inferred foldl(right, 1:3)
3
3 Likes
8

Interesting, although isn’t foldl here more akin to reduce?

1 Like
9

The short answer is “because you’re writing to a location in memory outside of your current scope”.

I don’t know the long answer, but it probably has to do with escape analysis :slight_smile: Another guess would be that the compiler can’t properly see “into” the inner function from the outside and see that it will always assign a Int to the location in the outer scope. That’s just a guess though.

In any case, you can keep your pattern and type stability:

function better_last(vs::AbstractVector)
    x = Ref(first(vs))

    map(vs) do v
        x[] = v
    end

    x[]
end

And now it’s type stable:

julia> @code_warntype better_last([1,2,4,4])
MethodInstance for better_last(::Vector{Int64})
  from better_last(vs::AbstractVector) in Main at REPL[3]:1
Arguments
  #self#::Core.Const(better_last)
  vs::Vector{Int64}
Locals
  #3::var"#3#4"{Base.RefValue{Int64}}
  x::Base.RefValue{Int64}
Body::Int64
1 ─ %1 = Main.first(vs)::Int64
│        (x = Main.Ref(%1))
│   %3 = Main.:(var"#3#4")::Core.Const(var"#3#4")
│   %4 = Core.typeof(x)::Core.Const(Base.RefValue{Int64})
│   %5 = Core.apply_type(%3, %4)::Core.Const(var"#3#4"{Base.RefValue{Int64}})
│        (#3 = %new(%5, x))
│   %7 = #3::var"#3#4"{Base.RefValue{Int64}}
│        Main.map(%7, vs)
│   %9 = Base.getindex(x)::Int64
└──      return %9

However, performance can’t measure up (down?) to last, which just computes the last index and directly indexes into the array:

julia> @benchmark better_last(arr) setup=(arr=rand(UInt, 8192)) evals=1
BenchmarkTools.Trial: 10000 samples with 1 evaluation.
 Range (min … max):   3.357 μs … 937.875 μs  ┊ GC (min … max): 0.00% … 98.24%
 Time  (median):      7.968 μs               ┊ GC (median):    0.00%
 Time  (mean ± σ):   11.191 μs ±  26.418 μs  ┊ GC (mean ± σ):  7.68% ±  3.50%

     ██▄▅▅▅▃▁
  ▁▁▂████████▇▆▄▃▃▂▂▂▂▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▂▂▃▂▂▂▂▂▁▁▁▁▁▁▁▁▁▁▁ ▂
  3.36 μs         Histogram: frequency by time         37.4 μs <

 Memory estimate: 64.06 KiB, allocs estimate: 3.

julia> @benchmark last(arr) setup=(arr=rand(UInt, 8192)) evals=1
BenchmarkTools.Trial: 10000 samples with 1 evaluation.
 Range (min … max):  24.000 ns …  1.973 μs  ┊ GC (min … max): 0.00% … 0.00%
 Time  (median):     29.000 ns              ┊ GC (median):    0.00%
 Time  (mean ± σ):   33.343 ns ± 33.082 ns  ┊ GC (mean ± σ):  0.00% ± 0.00%

  ▂▆█▆▅▅▃▂▄▂▃▁▁▂▁▁▂▁▂  ▁                                      ▂
  ███████████████████████▇▇▇▇▆▅▆▅▆▅▅▆▁▄▄▄▁▅▃▅▆▄▄▃▅▅▅▆▅▆▅▅▆▅▅▆ █
  24 ns        Histogram: log(frequency) by time       105 ns <

 Memory estimate: 0 bytes, allocs estimate: 0.

The intuitive reasoning for this working is that the compiler now has a deterministic place to put data and using Ref makes the location where overwrites occur explicit (with x[]). Further, since that’s another call to getindex/setindex, it can do type inference to check whether types will match, since Ref has a type parameter to indicate what it holds while Box intentionally does not, as far as I know. I’m not 100% certain on this, but I’m fairly certain that using Ref when modifying captured variables should work every time.

Said differently, type inference sees that the variable gets captured by the inner function, so it has to Box it. This however loses type information when reading from it, which Ref does not suffer from. This leads to the interesting situation that a slightly modified version appears type stable:

julia> function unstable_last(vs::AbstractVector)
           x = Ref(first(vs))

           map(vs) do v
               x[] = string(v)
           end

           x[]
       end
unstable_last (generic function with 1 method)

julia> @code_warntype unstable_last([1,2,4,4])
MethodInstance for unstable_last(::Vector{Int64})
  from unstable_last(vs::AbstractVector) in Main at REPL[33]:1
Arguments
  #self#::Core.Const(unstable_last)
  vs::Vector{Int64}
Locals
  #9::var"#9#10"{Base.RefValue{Int64}}
  x::Base.RefValue{Int64}
Body::Int64
1 ─ %1 = Main.first(vs)::Int64
│        (x = Main.Ref(%1))
│   %3 = Main.:(var"#9#10")::Core.Const(var"#9#10")
│   %4 = Core.typeof(x)::Core.Const(Base.RefValue{Int64})
│   %5 = Core.apply_type(%3, %4)::Core.Const(var"#9#10"{Base.RefValue{Int64}})
│        (#9 = %new(%5, x))
│   %7 = #9::var"#9#10"{Base.RefValue{Int64}}
│        Main.map(%7, vs)
│   %9 = Base.getindex(x)::Int64
└──      return %9

But in fact crashes horribly when trying to run it, due to a type mismatch:

julia> unstable_last([1,2,4,4])
ERROR: MethodError: Cannot `convert` an object of type String to an object of type Int64

Luckily, JET.jl can still find it!

julia> using JET

julia> @report_call unstable_last([1,2,4,4])
═════ 1 possible error found ═════
┌ @ REPL[33]:4 Main.map(#9, vs)
│┌ @ abstractarray.jl:2853 Base.collect_similar(A, Base.Generator(f, A))
││┌ @ array.jl:713 Base._collect(cont, itr, Base.IteratorEltype(itr), Base.IteratorSize(itr))
│││┌ @ array.jl:804 y = Base.iterate(itr)
││││┌ @ generator.jl:47 Base.getproperty(g, :f)(Base.getindex(y, 1))
│││││┌ @ REPL[33]:5 Base.setindex!(Core.getfield(#self#, :x), Main.string(v))
││││││┌ @ refvalue.jl:57 Base.setproperty!(b, :x, x)
│││││││┌ @ Base.jl:39 Base.convert(Base.fieldtype(Base.typeof(x), f), v)
││││││││ no matching method found for call signature (Tuple{typeof(convert), Type{Int64}, String}): Base.convert(Base.fieldtype(Base.typeof(x::Base.RefValue{Int64})::Type{Base.RefValue{Int64}}, f::Symbol)::Type{Int64}, v::String)

You should also know that julia does not do/guarantee tail call optimizations - the reasons for which have to do with semantics.

1 Like
10

In Julia, reduce requires an associative operator. So, foldl is the functional representation of sequential loops, and reduce is the functional representation of parallel loops.

3 Likes
11

This is one of my least-favorite limitations of Julia (the issue number, 15276, is burned into my brain), but fixing it is very nontrivial. IIUC, Julia’s code-lowering step is going to have to gain access to some information about type-inference. Currently there’s a bright dividing line between these two and and it’s a one-way trip: whatever decisions lowering makes, you’re stuck with them. This issue is the result. My understanding is that it will basically take an entire rewrite of lowering to fix this.

In case you haven’t seen it: Performance Tips · The Julia Language.

5 Likes