FluxML / Mjolnir.jl

A little less conversation, a little more abstraction
Other
87 stars 13 forks source link

@trace broken with some loops #27

Open zenna opened 3 years ago

zenna commented 3 years ago

I've found that when I trace a function where the number of iterations of a loop is passed in as an ::Int, things break. for instance;

julia> function m(n)
         for i = 1:n
           nothing
         end
       end
m (generic function with 1 method)

julia> @trace m(::Int)
ERROR: Tracing Error: Unrecognised expression $(QuoteNode(Int64))
in Tuple{typeof(m),Int64}
in Tuple{typeof(iterate),UnitRange{Int64}}
Stacktrace:
 [1] trace(::Any, ::Any, ::Vararg{Any,N} where N) at /Users/zenna/.julia/packages/Mjolnir/7FUVN/src/trace.jl:268
 [2] top-level scope at REPL[51]:1
caused by [exception 1]
Unrecognised expression $(QuoteNode(Int64))
Stacktrace:
 [1] error(::String) at ./error.jl:33
 [2] step!(::Mjolnir.Inference) at /Users/zenna/.julia/packages/Mjolnir/7FUVN/src/infer.jl:200
 [3] infer!(::Mjolnir.Inference) at /Users/zenna/.julia/packages/Mjolnir/7FUVN/src/infer.jl:225
 [4] abstract!(::Any, ::Any, ::Any) at /Users/zenna/.julia/packages/Mjolnir/7FUVN/src/trace.jl:136
 [5] trace!(::Mjolnir.Trace, ::Any, ::Any) at /Users/zenna/.julia/packages/Mjolnir/7FUVN/src/trace.jl:220
 [6] tracecall!(::Mjolnir.Trace, ::Any, ::Any, ::Vararg{Any,N} where N) at /Users/zenna/.julia/packages/Mjolnir/7FUVN/src/trace.jl:231
 [7] traceblock!(::Mjolnir.Trace, ::Any, ::Any) at /Users/zenna/.julia/packages/Mjolnir/7FUVN/src/trace.jl:189
 [8] trace!(::Mjolnir.Trace, ::Any, ::Any) at /Users/zenna/.julia/packages/Mjolnir/7FUVN/src/trace.jl:213
 [9] tracecall!(::Mjolnir.Trace, ::Any, ::Any, ::Vararg{Any,N} where N) at /Users/zenna/.julia/packages/Mjolnir/7FUVN/src/trace.jl:231
 [10] trace(::Any, ::Any, ::Vararg{Any,N} where N) at /Users/zenna/.julia/packages/Mjolnir/7FUVN/src/trace.jl:263
 [11] top-level scope at REPL[51]:1
femtomc commented 3 years ago

@zenna If you annotate the arg type in the function def, does the same issue occur? Would try myself but away from computer.

zenna commented 3 years ago

@femtomc Yes. Perhaps this is the same issue as #7 ? Rewriting it to use a while loop does not produce the same error.

femtomc commented 3 years ago

That's likely because the lowered IR is much simpler for the while version:

function m(n)
       i = 0
       while i <= n
       nothing
       i += 1
       end
 end
1: (%1, %2)
  br 2 (0)
2: (%3)
  %4 = %3 <= %2
  br 4 unless %4
  br 3
3:
  %5 = Main.nothing
  %6 = %3 + 1
  br 2 (%6)
4:
  return nothing

compared to

1: (%1, %2)
  %3 = 1:%2
  %4 = Base.iterate(%3)
  %5 = %4 === nothing
  %6 = Base.not_int(%5)
  br 3 unless %6
  br 2 (%4)
2: (%7)
  %8 = Core.getfield(%7, 1)
  %9 = Core.getfield(%7, 2)
  %10 = Main.nothing
  %11 = Base.iterate(%3, %9)
  %12 = %11 === nothing
  %13 = Base.not_int(%12)
  br 3 unless %13
  br 2 (%11)
3:
  return nothing
zenna commented 3 years ago

Right, and in particular the type of %4 will be a Union

femtomc commented 3 years ago

I'm trying to fill in the primitives in Basic right now to see if I can get it working.

zenna commented 3 years ago

Ok, I have not looked at the internals of Mjolnir but the older issues point to no support for Unions, so on the face of it it seems to be a bigger job that defining more primitives?

femtomc commented 3 years ago

Yep - think you're spot on.

femtomc commented 3 years ago

Although this particular issue seems to be one which is easily fixed without primitive definitions or handling unions.

The IR which causes an issue is:

1: (%1 :: const(first), %2 :: const(2:3))
  %3 = $(QuoteNode(Int64)) :: Union{}
  %4 = Base.getproperty(%2, :start) :: Union{}
  %5 = Base.convert(%3, %4) :: Union{}
  return %5

and it's just because Mjolnir can't handle the $ expression.

femtomc commented 3 years ago

@zenna Here ya go:

1: (%1 :: const(m), %2 :: Int64)
  %3 = (Colon())(1, %2) :: UnitRange{Int64}
  %4 = (getfield)(%3, :start) :: Int64
  %5 = (getfield)(%3, :stop) :: Int64
  %6 = (>)(%4, %5) :: Bool
  br 2 unless %6
  br 3
2:
  br 3
3:
  return nothing

I just added https://github.com/femtomc/Mjolnir.jl/blob/ad46e7cc08286153ee4b099a076ac492dd1eb45b/src/infer.jl#L199 to allow QuoteNode instances to fall through.

But don't ask me if I know if this is the right fix. At least it solves some issue.