Possible bug in type inference? (4 line MWE)

Hello!

I came across a weird type inference failure when using generators. Here is a MWE.

function example()
    col = 2
    ss = (col for _ in 1:5)
    col += 3
    return col, ss
end

function example2()
    col = 2
    ss = (col for _ in 1:5)
    return col, ss
end

@code_warntype example()
@code_warntype example2()

example has type instability:

Variables
  #self#::Core.Compiler.Const(example, false)
  #611::var"#611#612"
  ss::Base.Generator{UnitRange{Int64},var"#611#612"}
  col@_4::Core.Box
  col@_5::Union{}
  col@_6::Union{}

Body::Tuple{Any,Base.Generator{UnitRange{Int64},var"#611#612"}}
1 ─       (col@_4 = Core.Box())
β”‚         Core.setfield!(col@_4, :contents, 2)
β”‚         (#611 = %new(Main.:(var"#611#612"), col@_4))
β”‚   %4  = #611::var"#611#612"
β”‚   %5  = (1:5)::Core.Compiler.Const(1:5, false)
β”‚         (ss = Base.Generator(%4, %5))
β”‚   %7  = Core.isdefined(col@_4, :contents)::Bool
└──       goto #3 if not %7
2 ─       goto #4
3 ─       Core.NewvarNode(:(col@_5))
└──       col@_5
4 β”„ %12 = Core.getfield(col@_4, :contents)::Any
β”‚   %13 = (%12 + 3)::Any
β”‚         Core.setfield!(col@_4, :contents, %13)
β”‚   %15 = Core.isdefined(col@_4, :contents)::Bool
└──       goto #6 if not %15
5 ─       goto #7
6 ─       Core.NewvarNode(:(col@_6))
└──       col@_6
7 β”„ %20 = Core.getfield(col@_4, :contents)::Any
β”‚   %21 = Core.tuple(%20, ss::Core.Compiler.PartialStruct(Base.Generator{UnitRange{Int64},var"#611#612"}, Any[var"#611#612", Core.Compiler.Const(1:5, false)]))::Core.Compiler.PartialStruct(Tuple{Any,Base.Generator{UnitRange{Int64},var"#611#612"}}, Any[Any, Core.Compiler.PartialStruct(Base.Generator{UnitRange{Int64},var"#611#612"}, Any[var"#611#612", Core.Compiler.Const(1:5, false)])])
└──       return %21

However, example2 does not:

Variables
  #self#::Core.Compiler.Const(example2, false)
  #613::var"#613#614"{Int64}
  col::Int64
  ss::Base.Generator{UnitRange{Int64},var"#613#614"{Int64}}

Body::Tuple{Int64,Base.Generator{UnitRange{Int64},var"#613#614"{Int64}}}
1 ─      (col = 2)
β”‚   %2 = Main.:(var"#613#614")::Core.Compiler.Const(var"#613#614", false)
β”‚   %3 = Core.typeof(col::Core.Compiler.Const(2, false))::Core.Compiler.Const(Int64, false)
β”‚   %4 = Core.apply_type(%2, %3)::Core.Compiler.Const(var"#613#614"{Int64}, false)
β”‚        (#613 = %new(%4, col::Core.Compiler.Const(2, false)))
β”‚   %6 = #613::Core.Compiler.Const(var"#613#614"{Int64}(2), false)::Core.Compiler.Const(var"#613#614"{Int64}(2), false)
β”‚   %7 = (1:5)::Core.Compiler.Const(1:5, false)
β”‚        (ss = Base.Generator(%6, %7))
β”‚   %9 = Core.tuple(col::Core.Compiler.Const(2, false), ss::Core.Compiler.Const(Base.Generator{UnitRange{Int64},var"#613#614"{Int64}}(var"#613#614"{Int64}(2), 1:5), false))::Core.Compiler.Const((2, Base.Generator{UnitRange{Int64},var"#613#614"{Int64}}(var"#613#614"{Int64}(2), 1:5)), false)
└──      return %9

I cannot understand this in terms of usual type instability rules. Could this be a julia type inference bug?

This is not a bug. The problem is that col is not constant in example (but it is in example2), so it must be boxed because col must have the same value both in ss and in example (see the col@_4::Core.Box line in the output of @code_warntype example()).

You can solve this problem using a let block as described in the documentation:

function example()
    col = 2
    ss = let col = col
		(col for _ in 1:5)
	end
    col += 3
    return col, ss
end
3 Likes

Ah I see! I was misunderstanding what generators were supposed to. I thought

col = 2
a = (col for _ in 1:3)

will evaluate the value for col at the time of creation, such that it would be equivalent to writing a = (2 for _ in 1:3). But I guess I should really think of col as a captured variable just like in the case of local functions. Is that correct?

Exactly. Indeed, a generator implicitly creates an anonymous function to apply to each element of an iterable.

1 Like

Great, thanks!

Incidentally, is there any good reference for learning generators, closure, captured variables etc. in Julia? I come across these whenever I try to do something complicated, but I have not been able to find good documentation for their behavior, as they seem rather scattered in the documentation.

Many thanks!

When you create a generator, it basically calls Base.Generator:

julia> Meta.@lower (fun(x) for x in y)
:($(Expr(:thunk, CodeInfo(
    @ none within `top-level scope'
1 ─ %1 = Base.Generator(fun, y)
└──      return %1
))))

So you can call Base.Generator yourself:

help?> Base.Generator
  Generator(f, iter)

  Given a function f and an iterator iter, construct an iterator that yields the values of f applied to the elements of iter. The syntax for constructing
  an instance of this type is f(x) for x in iter [if cond(x)::Bool]. The [if cond(x)::Bool] expression is optional and acts as a "guard", effectively
  filtering out values where the condition is false.

  julia> g = (abs2(x) for x in 1:5 if x != 3);

  julia> for x in g
             println(x)
         end
  1
  4
  16
  25

  julia> collect(g)
  4-element Array{Int64,1}:
    1
    4
   16
   25

So, in your original code, you may write something like (note the let block):

function example()
	col = 2
	f = let col = col
		_ -> col
	end
	ss = Base.Generator(f, 1:5)
	col += 3
	return col, ss
end

Edit: A good discussion on closures and their performance is the documentation that I mentioned earlier.

3 Likes

Great, thank you!

Well, it’s kind of a bug. In theory the compiler could do escape analysis and figure out that the capture dies before the modification happens and thus doesn’t need to be captured in the first place. It’s just that the analysis is hard.

4 Likes

But in this case the function returns the generator. The captured variable changes between the creation of the generator and its use. Of course, a similar argument could be applied in this case: the generator does not use the original value of the captured variable, so the compiler could, in theory, figure out that it only needs the updated value of col.

1 Like

Yes, I misread the example. You are absolutely right that the analysis is still possible but harder this way.

1 Like