Handle unreachable blocks in the adjoint CFG

[Pangoraw]

The error described in #1118 and #1380, stopped happening in Julia 1.10-beta (maybe because of JuliaLang/julia#50943) but the gradient is now wrong.

julia> f(x) = @inbounds return x
f (generic function with 1 method)

julia> Zygote.gradient(f,1.)
(nothing,)

This PR adds unreachable branches at the end of blocks in the adjoint when those blocks are unreachable in the primal which fixes the issue in both 1.9 and 1.10 because it avoids implicit branches.

I also though of removing unreachable blocks altogether since some operations are invalid for these blocks in IRTools (see dominators for example) but it seemed to do a lot more work than this fix.

Fixes #1380
Fixes #1118

Note

It depends on FluxML/IRTools.jl#115.

[oxinabox]

this seems good to me, based on the new test. But I am not that familar with this bit of the code

[Pangoraw]

I had to add the Core.throw call because having an unreachable in a reachable path throws a compiler error. And it is possible to manually construct a pullback which reaches these adjoint blocks whose primal blocks are unreachable. This is specifically for all versions concerned by FluxML/IRTools.jl#115 (>= 1.10) since unreachable is emitted as a Core.throw call otherwise.

julia> using Zygote

julia> function f(x, cond)
           if cond
               return x
           else
               return 2x
           end
           return 3x
       end
f (generic function with 1 method)

julia> back = Zygote.Pullback{Tuple{typeof(f), Float64, Bool}, Any}(([Returns([1,2,3])], [Returns([1,2,3])], 0x03))
∂(f)

julia> back(1.)
ERROR: "unreachable"
Stacktrace:
 [1] f
   @ ./REPL[2]:7 [inlined]
 [2] (::Zygote.Pullback{Tuple{typeof(f), Float64, Bool}, Any})(Δ::Float64)
   @ Zygote ~/Projects/Zygote.jl/src/compiler/interface2.jl:0
 [3] top-level scope
   @ REPL[5]:1

[ToucheSir]

My understanding is that this may incur an allocation for the string "unreachable" in some methods. Is this reliably optimized out in practice? If not, it might make sense to define a custom immutable error type. That type could also track more information such as file and line number to help with troubleshooting.

[Pangoraw]

The string "unreachable" string literal is actually allocated once by the parser and permanently rooted.

julia> using Zygote

julia> f(x) = @inbounds return x;

julia> invalid_pull = Zygote.Pullback{Tuple{typeof(f), Float64}, Any}(0x02,);

julia> m = only(methods(Zygote.adjointcfg));
       root = m.roots[findfirst(==("unreachable"), m.roots)];

julia> unreachable_string = try invalid_pull(1.); catch exc; exc end
"unreachable"

julia> unreachable_string === root
true

One problem is that it may taint the effects of the pullback compared to the case without unreachable blocks (which we could have if we removed unreachable blocks from the primal):

julia> Core.Compiler.infer_effects(invalid_pull, (Float64,))
(!c,!e,!n,!t,!s,!m,+i)′

julia> g(x) = return x;

julia> _, pull = Zygote.pullback(g, 1.);
       Core.Compiler.infer_effects(pull.back, (Float64,))
(+c,+e,+n,+t,+s,+m,+i)

But any sufficiently complex pullbacks (involving Zygote.Stack) can throw too so it may not be a problem after all.

[ToucheSir]

A quick sanity check before I merge: does this pass tests locally on nightly for you?

[Pangoraw]

The added testset does, but the entire test suite fails with CUDA and a missing make_seed attribute just like in CI.

[ToucheSir]

Could you try temporarily commenting out the CUDA test imports and the test block in

Zygote.jl/test/runtests.jl

Lines 8 to 15 in 49a1184

	if has_cuda()
	@testset "CUDA tests" begin
	include("cuda.jl")
	end
	@info "CUDA tests have run"
	else
	@warn "CUDA not found - Skipping CUDA Tests"
	end

? Either locally or on this branch to run CI would be fine.

[Pangoraw]

I had to cherry-pick #1462 but there is one @optout and two grad checks tests failing. I don't know if those are related to this pr.