Are "inline assignments" inside loop predicates bad for type inference?

Internals & Design
, , , ,
1

I was playing with JET and noticed that it does not like notation like while (i=f())!==nothing ... while it is perfectly happy with i=f(); while i!==nothing .... f is a function that returns Union{Int,Nothing}. It seems the “inline assignment” inside of the while predicate is the reason for this and I was hoping I can get more feedback concerning the root cause. Is this a bad way to write julia code? Here is a full example with two naive versions of Base.invperm:

using JET
using Random

function inline_assign_while_loop(permutation::Vector{Int})
    indices = Int[]
    v = 1
    while (i=findfirst(==(v),permutation)) !== nothing
        push!(indices, i)
        v += 1
    end
    indices
end

function normal_assign_while_loop(permutation::Vector{Int})
    indices = Int[]
    v = 1
    i = findfirst(==(v),permutation)
    while i !== nothing
        push!(indices, i)
        v += 1
        i = findfirst(==(v),permutation)
    end
    indices
end

test_perm = randperm(10)
@assert inline_assign_while_loop(test_perm) == normal_assign_while_loop(test_perm) == invperm(test_perm)

# A possible error reported by JET because i!==nothing is not inferred
@report_call inline_assign_while_loop(test_perm) 

# No possible error reported by JET
@report_call normal_assign_while_loop(test_perm)

@code_warntype is also happier with normal_assign_while_loop, so I assume the issue is not with JET, rather with the type inference.

My question is: Is this something that Julia’s type inference should be doing better or is there something that I did explicitly wrong (or both)?

I prefer using the problematic inline assignment inside the while predicate because that way I need to write it only once (instead of writing it once before the loop and once inside the loop body). I know this is an inefficient quadratic algorithm for invperm, I am using it just as an example.

2

A very small modification makes the error go away for me (on a somewhat recent master - which version did you test?):

function inline_assign_while_loop(permutation::Vector{Int})
    indices = Int[]
    v = 1
-    while (i=findfirst(==(v),permutation)) !== nothing
+    while ((i=findfirst(==(v),permutation)); i !== nothing)
        push!(indices, i)
        v += 1
    end
    indices
end

If we take a look at the lowered representation, the difference is extremely subtle:

julia> @code_lowered inline_assign_while_loop(test_perm)
CodeInfo(
1 ─      indices = Base.getindex(Main.Int)
└──      v = 1
2 ┄ %3 = (==)(v)
│   %4 = Main.findfirst(%3, permutation)
│        i = %4
│   %6 = %4 !== Main.nothing
└──      goto #4 if not %6
3 ─      Main.push!(indices, i)
│        v = v + 1
└──      goto #2
4 ─      return indices
)

julia> @code_lowered inline_assign_while_loop_extra(test_perm)
CodeInfo(
1 ─      indices = Base.getindex(Main.Int)
└──      v = 1
2 ┄ %3 = (==)(v)
│        i = Main.findfirst(%3, permutation)
│   %5 = i !== Main.nothing
└──      goto #4 if not %5
3 ─      Main.push!(indices, i)
│        v = v + 1
└──      goto #2
4 ─      return indices

The reason for the warning by JET seems to be that the SSA slot %4 is checked against nothing, not i. So the issue does not lie in inference per se, but in propagating the information we have about %4 to i as well. This seems very interesting, so I’ll open an issue.

EDIT: Issue opened here.

4 Likes
3

I was also using “fairly recent master” 1.9.0-DEV.487 (2022-05-08). Thank you, this was really clear and interesting explanation. I am surprised that i and slot %4 are two different entities… I thought the whole point of SSA was to have a canonical notation in which they are the same entity (but my understanding is superficial and vague).

4

That was a ridiculously quick turnaround! Thank you, it was very impressive! Fixed now by @aviatesk in inference: improve `ssa_def_slot` tracking by aviatesk · Pull Request #45524 · JuliaLang/julia · GitHub

1 Like
5

Thx for pointing that out, that kind of issue is very helpful to improve our compiler. Also happy to know you’re using JET🥰

2 Likes