Constant propagation and inference of constructed types

Context: I had a parameterized struct type with convenience constructors that input some of the parameters as arguments. Based on previous experience, I expected that if such constructors are provided with constant values for those arguments, the constants will be propagated and the type of the constructed object will be inferred. To my surprise it was not, even though that approach had generally worked for me previously.

Here’s an illustration of how constant propagation sometimes works and sometimes doesn’t, even for similar-looking code. First, a couple of cases for which constant propagation does work:

struct F1{T} end
f1(T) = F1{T}()
g1_float() = f1(Float64)
g1_tuple() = f1((1,2,3))

julia> @code_warntype g1_float()
Variables
  #self#::Core.Compiler.Const(g1_float, false)
Body::F1{Float64}
1 ─ %1 = Main.f1(Main.Float64)::Core.Compiler.Const(F1{Float64}(), false)
└──      return %1

julia> @code_warntype g1_tuple()
Variables
  #self#::Core.Compiler.Const(g1_tuple, false)
Body::F1{(1, 2, 3)}
1 ─ %1 = Core.tuple(1, 2, 3)::Core.Compiler.Const((1, 2, 3), false)
β”‚   %2 = Main.f1(%1)::Core.Compiler.Const(F1{(1, 2, 3)}(), false)
└──      return %2

So far so good – constant floats and tuples can be propagated to types. Now, we add a second type parameter:

struct F2{T,X} end
f2(T, x) = F2{T,typeof(x)}()
g2_float(x) = f2(Float64, x)
g2_tuple(x) = f2((1,2,3), x)

julia> @code_warntype g2_float('a')
Variables
  #self#::Core.Compiler.Const(g2_float, false)
  x::Char
Body::F2{Float64,Char}
1 ─ %1 = Main.f2(Main.Float64, x)::Core.Compiler.Const(F2{Float64,Char}(), false)
└──      return %1

julia> @code_warntype g2_tuple('a')
Variables
  #self#::Core.Compiler.Const(g2_tuple, false)
  x::Char
Body::F2{_A,Char} where _A
1 ─ %1 = Core.tuple(1, 2, 3)::Core.Compiler.Const((1, 2, 3), false)
β”‚   %2 = Main.f2(%1, x)::F2{_A,Char} where _A
└──      return %2

For some reason, adding a second argument (or type parameter) made the compiler forget how to propagate the tuple to the output type, but not the float.

It gets stranger. Suppose we don’t ask for an instance, just for the type of object that would’ve been created:

f2_type(T, x) = F2{T,typeof(x)}
g2_type_tuple(x) = f3_type((1,2,3), x)

julia> @code_warntype g2_type_tuple('a')
Variables
  #self#::Core.Compiler.Const(g2_type_tuple, false)
  x::Char
Body::Type{F2{(1, 2, 3),Char}}
1 ─ %1 = Core.tuple(1, 2, 3)::Core.Compiler.Const((1, 2, 3), false)
β”‚   %2 = Main.f2_type(%1, x)::Core.Compiler.Const(F2{(1, 2, 3),Char}, false)
└──      return %2

The compiler remembered how to propagate the tuple again! This shows that the compiler actually has the ability to infer the output type of f2, but is just unwilling to.

Finally, we can obtain the desired behavior by wrapping the tuple in a Val:

f2_val(::Val{T}, x) where {T} = F2{T, typeof(x)}()
g2_val_tuple(x) = f2_type((1,2,3), x)


julia> @code_warntype g2_val_tuple('a')
Variables
  #self#::Core.Compiler.Const(g2_val_tuple, false)
  x::Char
Body::Type{F2{(1, 2, 3),Char}}
1 ─ %1 = Core.tuple(1, 2, 3)::Core.Compiler.Const((1, 2, 3), false)
β”‚   %2 = Main.f2_type(%1, x)::Core.Compiler.Const(F2{(1, 2, 3),Char}, false)
└──      return %2

The heuristics regarding what will and what will not be inferred by the compiler seem to me quite complex, almost haphazard. I realize that the compiler has the prerogative to infer as much or as little as it wants, but it is helpful when writing code to have some intuition about what is likely to be performant and what is not. With behavior like this, I seem to spend a significant amount of time rewriting code trying different approaches until I stumble upon the β€œmagic” way of writing things so that the compiler will propagate constants and infer types.

Is there a rational explanation for the behavior above? Could compiler type inference be made to work more consistently/predictably? (BTW, this is not meant to be a criticism of the developers, since Julia is overall quite impressive and I know a lot of effort has gone into the compiler.)

1 Like