Issue with terminators in normalise!

[gdalle]

@vchuravy spotted an issue with Mooncake's custom normalisation. Here's an MWE:

const CC = Core.Compiler
using Mooncake

function func(x)
    y = sin(x)
    if x[1] > 0
        z = cos(y)
    else
        z = sin(y)
    end
    return z
end

interp = Mooncake.get_interpreter()
ir = Base.code_ircode(func, (Float64,))[1][1]
sig = Tuple{typeof(func),Float64}
primal_ir, _ = Mooncake.lookup_ir(interp, sig)
_, spnames = Mooncake.is_vararg_and_sparam_names(sig)
primal_ir = Mooncake.normalise!(primal_ir, spnames)
CC.verify_ir(primal_ir)

Here's the error

julia> CC.verify_ir(primal_ir)
Block 1 successors (Array{Int64, 1}(dims=(2,), mem=Memory{Int64}(2, 0x739f1089a780)[3, 2])), does not match fall-through terminator %1 (Expr(:invoke, sin(Float64) from sin(T) where {T<:Union{Float32, Float64}}, Main.sin, Core.Argument(n=2)))::Float64
5  1 ─ %1  = invoke Main.sin(_2::Float64)::Float64
   2 ─       nothing::Nothing
   3 ┄       goto #4
6  4 ─ %4  = (Mooncake.IntrinsicsWrappers.lt_float)(0.0, _2)::Bool
   │   %5  = (Mooncake.IntrinsicsWrappers.or_int)(%4, false)::Bool
   └──       goto #6 if not %5
7  5 ─ %7  = invoke Main.cos(%1::Float64)::Float64
   └──       goto #7
9  6 ─ %9  = invoke Main.sin(%1::Float64)::Float64
11 7 ┄ %10 = φ (#5 => %7, #6 => %9)::Float64
   └──       return %10
   ERROR: 
Stacktrace:
 [1] error(s::String)
   @ Core.Compiler ./error.jl:35
 [2] verify_ir(ir::Core.Compiler.IRCode, print::Bool, allow_frontend_forms::Bool, 𝕃ₒ::Core.Compiler.PartialsLattice{…})
   @ Core.Compiler ./compiler/ssair/verify.jl:233
 [3] verify_ir(ir::Core.Compiler.IRCode)
   @ Core.Compiler ./compiler/ssair/verify.jl:98
 [4] top-level scope
   @ ~/Work/GitHub/Julia/Mooncake.jl/test/playground.jl:20
Some type information was truncated. Use `show(err)` to see complete types.

And here's the proposed fix: replace

Mooncake.jl/src/interpreter/ir_utils.jl

Line 335 in 0f37c07

return val === true ? nothing : GotoNode(stmt.dest)

with

            return val === true ? GotoIfNot(val, stmt.dest) : GotoNode(stmt.dest)

I don't understand anything but I dutifully convey the message.

[willtebbutt]

Ooo this is fun. Thanks @vchuravy for spotting this, and @gdalle for opening.

I'm kind of surprised that this hasn't caused us problems already, as it looks like something that really ought to. I'm not running CC.verify_ir after the normalisation pass, but I'll try adding that in and see if it complains for some of the existing test cases.

[vchuravy]

Other thing to note that I found

Mooncake.jl/src/interpreter/ir_normalisation.jl

Line 60 in 5736d3a

statements[m] = replace_uses_with!(stmt, def, val)

odd.

In Julia SSA IR form replace_all_uses_with is just statements[n] = val. And if you are walking the IR like that, there are better ways. You can use userefs to walk any stmt and visit + replace all arguments.