Relooping issue with IR generated by Mjolnir

[tshort]

I was playing around a bit using Mjolnir to construct WebAssembly from Julia code. The conversion from IR to WebAssembly failed during relooping. Here's the code. It's based on similar code in XLA.jl/src/compile/convert.jl.

using Mjolnir
using Mjolnir: AType, Const, @abstract
using WebAssembly
using WebAssembly: i32, i64, f32, f64
using IRTools.All

struct Operations end

Defaults() = Mjolnir.Multi(Operations(), Mjolnir.Basic())

exprtype(ir, x) = IRTools.exprtype(ir, x, typeof = Const)

CoreTypes = Union{Int32, Int64, Float32, Float64}

for op in (+, *, -, /, ^, >, <, >=, <=, max)
  @eval @abstract Operations $op(a::Const{T}, b::Const{T}) where T<:CoreTypes =
    Const($op(a.value, b.value))
  @eval @abstract Operations $op(a::AType{T}, b::AType{T}) where T<:CoreTypes =
    Core.Compiler.return_type($op, Tuple{T,T})
end

wasmop(args, ::AType{typeof(-)}, a::AType{Int64}, b::AType{Int64}) = stmt(xcall(i64.sub, args[2:end]...), type = Int64)
wasmop(args, ::AType{typeof(>)}, a::AType{Int64}, b::AType{Int64}) = stmt(xcall(i64.gt_s, args[2:end]...), type = Int64)
wasmop(args, ::AType{typeof(*)}, a::AType{Float64}, b::AType{Float64}) = stmt(xcall(f64.mul, args[2:end]...), type = Float64)

function wasmop!(ir, v, args...; kw...)
  ir[v] = wasmop(ir[v].expr.args, args...; kw...)
end

function wasmops!(ir)
  for (v, st) in ir
    if st.expr isa Expr
      wasmop!(ir, v, exprtype.((ir,), st.expr.args)...)
    elseif st.expr isa IR
      ir[v] = wasmops!(st.expr)
    end
  end
  return ir
end

function wasmtypes!(ir)
  deletearg!(ir, 1)
  for bl in blocks(ir)
    bb = BasicBlock(bl)
    for (idx, T) in enumerate(bb.argtypes)
      bb.argtypes[idx] = WebAssembly.WType(T)
    end
    for (idx, st) in enumerate(bb.stmts)
      bb.stmts[idx] = Statement(st.expr, WebAssembly.WType(st.type), st.line)
    end
  end
end

function convert_wasm!(ir)
  wasmops!(ir)
  wasmtypes!(ir)
  return ir
end

function pow(x, n)
    r = one(x)
    while n > 0
        n -= 1
        r *= x
    end
    return r
end

ir = @trace Defaults() pow(Float64, Int)
convert_wasm!(ir)
@show ir

func = WebAssembly.irfunc(:pow, copy(ir))

The resulting IR looks reasonable for WebAssembly IR.

ir = 1: (%2 :: f64, %3 :: i64)
  %4 = (i64.gt_s)(%3, 0) :: i64
  br 3 (1.0) unless %4
  br 2 (%3, 1.0)
2: (%5 :: i64, %6 :: f64)
  %7 = (i64.sub)(%5, 1) :: i64
  %8 = (f64.mul)(%6, %2) :: f64
  %9 = (i64.gt_s)(%7, 0) :: i64
  br 3 (%8) unless %9
  br 2 (%7, %8)
3: (%10 :: f64)
  return %10

Here's the error message when trying to use irfunc:

ERROR: LoadError: MethodError: no method matching -(::Nothing, ::Int64)
Closest candidates are:
  -(::BigFloat, ::Union{Int16, Int32, Int64, Int8}) at mpfr.jl:425
  -(::BigInt, ::Union{Int16, Int32, Int64, Int8}) at gmp.jl:532
  -(::Missing, ::Number) at missing.jl:115
  ...
Stacktrace:
 [1] reloop(::IRTools.Inner.IR, ::IRTools.Inner.CFG) at /home/tshort/.julia/dev/WebAssembly/src/ir.jl:108
 [2] irfunc(::Symbol, ::IRTools.Inner.IR) at /home/tshort/.julia/dev/WebAssembly/src/ir.jl:129
 [3] top-level scope at /home/tshort/mj/pow.jl:93
 [4] include(::String) at ./client.jl:457
 [5] top-level scope at REPL[6]:1
in expression starting at /home/tshort/mj/pow.jl:93

I'm not sure the IR is correct or if I'm even on the right track here.

[terasakisatoshi]

I could reproduce same error by constructing IR manually like @trace pow(Float64, Int32) without using @tshort 's converter above.

https://gist.github.com/terasakisatoshi/81a97ae06323793512dd9be1a3b0eba7