rm typed mode

Author: MikeInnes

.travis.yml

@@ -13,11 +13,6 @@ notifications:
 git:
   depth: 99999999
 
-env:
-  matrix:
-    - ZYGOTE_TYPED=true
-    - ZYGOTE_TYPED=false
-
 ## uncomment the following lines to allow failures on nightly julia
 ## (tests will run but not make your overall status red)
 matrix:

src/Zygote.jl

@@ -5,15 +5,6 @@ using LinearAlgebra: copytri!, AbstractTriangular
 
 import ZygoteRules: @adjoint, @adjoint!, AContext, adjoint, _pullback, pullback
 
-# This flag enables Zygote to grab extra type inference information during
-# compiles. When control flow is present, this can give gradient code a
-# performance boost.
-
-# HOWEVER, this is not Jameson-approved, nor well supported by the compiler, and
-# has several caveats. Recursion will cause inference to stack overflow.
-# Gradient redefinitions may result in ugly type errors. And Jameson *will* know.
-const usetyped = get(ENV, "ZYGOTE_TYPED", false) == "true"
-
 using IRTools
 using MacroTools, Requires
 using MacroTools: @forward
@@ -49,7 +40,7 @@ include("profiler/Profile.jl")
   include("flux.jl")
 end
 
-precompile() = usetyped || include(joinpath(@__DIR__, "precompile.jl"))
+precompile() = include(joinpath(@__DIR__, "precompile.jl"))
 
 # precompile()
 @init Requires.isprecompiling() || precompile()

src/compiler/emit.jl

@@ -22,7 +22,7 @@ end
 
 # Emit
 
-xstack(T) = stmt(Expr(:call, Vector{T}), type = Vector{T})
+xstack(T) = Expr(:call, Vector{T})
 
 function alphauses(b)
   us = Set{Alpha}()
@@ -45,24 +45,17 @@ function forward_stacks!(adj, F)
     if runonce(b)
       push!(recs, Variable(α))
     else
-      T = exprtype(pr, Variable(α))
-      stk = pushfirst!(pr, xstack(T))
+      stk = pushfirst!(pr, xstack(Any))
       push!(recs, stk)
       push!(b, xcall(Zygote, :_push!, stk, Variable(α)))
     end
     push!(stks, (b.id, alpha(α)))
   end
   args = arguments(pr)[3:end]
-  T = Tuple{concrete.(exprtype.((pr,), recs))...}
-  isconcretetype(T) || (T = Any)
   rec = push!(pr, xtuple(recs...))
-  if usetyped && length(pr.blocks) > 1
-    rec = push!(pr, Expr(:call, Pullback{F,T}, rec))
-  else
-    P = length(pr.blocks) == 1 ? Pullback{F} : Pullback{F,Any}
-    # P = Pullback{F,Any} # reduce specialisation
-    rec = push!(pr, Expr(:call, P, rec))
-  end
+  P = length(pr.blocks) == 1 ? Pullback{F} : Pullback{F,Any}
+  # P = Pullback{F,Any} # reduce specialisation
+  rec = push!(pr, Expr(:call, P, rec))
   ret = xtuple(pr.blocks[end].branches[end].args[1], rec)
   ret = push!(pr, ret)
   pr.blocks[end].branches[end].args[1] = ret
@@ -102,18 +95,8 @@ end
 
 varargs(m::Method, n) = m.isva ? n - m.nargs + 1 : nothing
 
-meta(T) = (usetyped ? IRTools.typed_meta : IRTools.meta)(T)
-
-function getmeta(T)
-  m = meta(T)
-  (usetyped && m != nothing) || return m
-  any(x -> isexpr(x, :goto, :gotoifnot), m.code.code) || return IRTools.meta(T)
-  return m
-end
-
 function _lookup_grad(T)
-  (m = getmeta(T)) == nothing && return
-  m isa IRTools.TypedMeta && m.ret == Union{} && return
+  (m = meta(T)) == nothing && return
   va = varargs(m.method, length(T.parameters))
   forw, back = stacks!(Adjoint(IR(m), varargs = va, normalise = false), T)
   m, forw, back

src/compiler/reverse.jl

@@ -1,7 +1,7 @@
 using IRTools: IR, Variable, Pipe, xcall, var, prewalk, postwalk,
   blocks, predecessors, successors, argument!, arguments, branches,
-  exprtype, insertafter!, finish, expand!, prune!, substitute!, substitute,
-  block, block!, branch!, return!, stmt
+  insertafter!, finish, expand!, prune!, substitute!, substitute,
+  block, block!, branch!, return!, stmt, meta
 using Base: @get!
 
 @inline tuple_va(N, xs) = xs
@@ -75,7 +75,7 @@ function instrument_global!(ir, v, ex)
   else
     ir[v] = prewalk(ex) do x
       istrackable(x) || return x
-      insert!(ir, v, stmt(xcall(Zygote, :unwrap, QuoteNode(x), x), type = exprtype(x)))
+      insert!(ir, v, xcall(Zygote, :unwrap, QuoteNode(x), x))
     end
   end
 end
@@ -125,16 +125,6 @@ ignored_f(ir, f::Variable) = ignored_f(get(ir, f, nothing))
 ignored(ir, ex) = isexpr(ex, :call) && ignored_f(ir, ex.args[1])
 ignored(ir, ex::Variable) = ignored(ir, ir[ex])
 
-# TODO: remove this once we don't mess with type inference
-function _forward_type(Ts)
-  usetyped || return Any
-  all(T -> isconcretetype(T) || T <: DataType, Ts) || return Any
-  T = Core.Compiler.return_type(_pullback, Tuple{Context,Ts...})
-  return T == Union{} ? Any : T
-end
-
-isvalidtype(jT, yT) = jT <: Tuple && length(jT.parameters) == 2 && jT.parameters[1] <: yT
-
 function primal(ir::IR)
   pr = Pipe(ir)
   pbs = Dict{Variable,Variable}()
@@ -143,23 +133,12 @@ function primal(ir::IR)
   for (v, st) in pr
     ex = st.expr
     if isexpr(ex, :call) && !ignored(ir, ex)
-      yT = exprtype(ir, v)
-      T = _forward_type(exprtype.((ir,), ex.args))
-      if yT == Any || isvalidtype(T, yT)
-        yJ = insert!(pr, v, stmt(xcall(Zygote, :_pullback, cx, ex.args...),
-                                 line = ir[v].line))
-        pr[v] = xgetindex(yJ, 1)
-        J = insertafter!(pr, v, stmt(xgetindex(yJ, 2),
-                                     type = T == Any ? Any : T.parameters[2],
-                                     line = ir[v].line))
-        pbs[v] = substitute(pr, J)
-      else
-        yJ = insert!(pr, v, xcall(Zygote, :_pullback, cx, ex.args...))
-        y =  insert!(pr, v, xgetindex(yJ, 1))
-        J =  insert!(pr, v, stmt(xgetindex(yJ, 2), line = ir[v].line))
-        pr[v] = xcall(Zygote, :typeassert, y, yT)
-        pbs[v] = substitute(pr, J)
-      end
+      yJ = insert!(pr, v, stmt(xcall(Zygote, :_pullback, cx, ex.args...),
+                               line = ir[v].line))
+      pr[v] = xgetindex(yJ, 1)
+      J = insertafter!(pr, v, stmt(xgetindex(yJ, 2),
+                                   line = ir[v].line))
+      pbs[v] = substitute(pr, J)
     end
   end
   pr = finish(pr)

src/flux.jl

@@ -1,7 +1,5 @@
 using .Tracker: TrackedArray, TrackedReal
 
-if !usetyped
-  unwrap(x::Union{TrackedArray,TrackedReal}) = Tracker.data(x)
-end
+unwrap(x::Union{TrackedArray,TrackedReal}) = Tracker.data(x)
 
 pullback(f, ps::Tracker.Params) = pullback(f, Params(ps))

test/compiler.jl

@@ -19,8 +19,6 @@ end
 bad(x) = x
 @adjoint bad(x) = x, Δ -> error("bad")
 
-Zygote.usetyped && Zygote.refresh()
-
 function badly(x)
   x = x + 1
   x = bad(x)
@@ -33,7 +31,7 @@ y, back = pullback(badly, 2)
 bt = try back(1) catch e stacktrace(catch_backtrace()) end
 
 @test trace_contains(bt, nothing, "compiler.jl", 20)
-@test trace_contains(bt, :badly, "compiler.jl", 26)
+@test trace_contains(bt, :badly, "compiler.jl", 24)
 
 # Type inference checks
 
@@ -76,8 +74,3 @@ y, back = @test_inferred pullback(x->x[1], (5,:a))
 
 y, back = @test_inferred pullback(((a,b),) -> a, (5, 10))
 @test_inferred back(1)
-
-# Checks that use control flow
-if Zygote.usetyped
-  include("typed.jl")
-end

test/complex.jl

@@ -7,10 +7,8 @@ using Zygote, Test
 
 @test gradient(a -> real((a*conj(a))), 0.3im)[1] == 0.6im
 @test gradient(a -> real((a.*conj(a))), 0.3im)[1] == 0.6im
-if !Zygote.usetyped
-  @test gradient(a -> real(([a].*conj([a])))[], 0.3im)[1] == 0.6im #TODO
-  @test gradient(a -> real(([a].*conj.([a])))[], 0.3im)[1] == 0.6im #TODO
-end
+@test gradient(a -> real(([a].*conj([a])))[], 0.3im)[1] == 0.6im
+@test gradient(a -> real(([a].*conj.([a])))[], 0.3im)[1] == 0.6im
 @test gradient(a -> real.(([a].*conj.([a])))[], 0.3im)[1] == 0.6im
 
 fs_C_to_R = (real,

test/features.jl

@@ -137,9 +137,7 @@ end == (4,)
 
 pow_rec(x, n) = n == 0 ? 1 : x*pow_rec(x, n-1)
 
-if !Zygote.usetyped
-  @test gradient(pow_rec, 2, 3) == (12, nothing)
-end
+@test gradient(pow_rec, 2, 3) == (12, nothing)
 
 # For nested AD, until we support errors
 function grad(f, args...)
@@ -187,17 +185,14 @@ y, back = pullback(() -> layer(x), Params([W]))
   sum(H)
 end[1] == 1
 
-# FIXME
-if !Zygote.usetyped
-  @test gradient(2) do x
-    if x < 0
-      throw("foo")
-    end
-    return x*5
-  end[1] == 5
+@test gradient(2) do x
+  if x < 0
+    throw("foo")
+  end
+  return x*5
+end[1] == 5
 
-  @test gradient(x -> one(eltype(x)), rand(10))[1] == nothing
-end
+@test gradient(x -> one(eltype(x)), rand(10))[1] == nothing
 
 # Thre-way control flow merge
 @test gradient(1) do x
@@ -214,15 +209,11 @@ grad_closure(x) = 2x
 
 Zygote.@adjoint (f::typeof(grad_closure))(x) = f(x), Δ -> (1, 2)
 
-Zygote.usetyped && Zygote.refresh()
-
 @test gradient((f, x) -> f(x), grad_closure, 5) == (1, 2)
 
-if !Zygote.usetyped
-  invokable(x) = 2x
-  invokable(x::Integer) = 3x
-  @test gradient(x -> invoke(invokable, Tuple{Any}, x), 5) == (2,)
-end
+invokable(x) = 2x
+invokable(x::Integer) = 3x
+@test gradient(x -> invoke(invokable, Tuple{Any}, x), 5) == (2,)
 
 y, back = Zygote.pullback(x->tuple(x...), [1, 2, 3])
 @test back((1, 1, 1)) == ((1,1,1),)

test/gradcheck.jl

@@ -121,7 +121,7 @@ end
     bar = (x) -> x*y
     sum(map(bar, 1:5))
   end
-  Zygote.usetyped || @test gradtest(foo, 3) #TODO
+  @test gradtest(foo, 3)
   @test gradient(v -> sum([x for x in v]), [1.1,2.2,3.3]) == ([1, 1, 1],)
 end
 
@@ -711,9 +711,7 @@ end
 end
 
 @testset "broadcast" begin
-  if !Zygote.usetyped
-    @test gradient(x -> sum(sin.(x)), Diagonal(randn(3)))[1][2] == 1
-  end
+  @test gradient(x -> sum(sin.(x)), Diagonal(randn(3)))[1][2] == 1
 end
 
 using Zygote: Buffer

test/runtests.jl

@@ -1,12 +1,6 @@
 using Zygote, Test
 using Zygote: gradient
 
-if Zygote.usetyped
-  @info "Testing Zygote in type-hacks mode."
-else
-  @info "Testing Zygote in normal mode."
-end
-
 @testset "Zygote" begin
 
 @testset "Features" begin