Support try/catch on the happy (nothrow) path #1474

Pangoraw · 2023-12-05T21:44:54Z

This enables differentation support to try/catch when no exception is thrown on the try block.

The changes are twofolds:

Stop throwing when encountering :enter and :leave expressions, add an dedicated error message in the adjoint of the catch block instead.
Make sure that alphas in the first block are stored as stacks if it is a try block since they are not garanteed to run (an exception can be thrown before).

Here is an example of the behavior:

julia> using Zygote

julia> function foo(f)
         y = 1
         try
           y = f()
         catch
           y
         end
         y
       end
foo (generic function with 1 method)

julia> gradient(x -> foo(() -> x), 1) # 1
(1.0,)

julia> _, pull = pullback(foo, () -> 0//0)
(1, Zygote.var"#76#77"{Zygote.Pullback{Tuple{typeof(foo), var"#9#10"}, Any}}(∂(foo)))

julia> pull(1.)
ERROR: Can't differentiate function execution in catch block at #= REPL[4]:4 =#.
Stacktrace:
...

FluxML/IRTools.jl#117 is a companion change that fixes some of the error that can currently happen with try/catch blocks in IRTools.jl.

Supersedes #466.

another possibility would be to thread the pullback values as block parameters through the control flow with arguments coming from the catch block being `nothing`.

ToucheSir · 2023-12-06T01:40:06Z

I'm not really sure how to evaluate any risks from this change other than reading through the old PR. @oxinabox as the reviewer on that, would you be able to take a look here?

docs/src/limitations.md

src/compiler/reverse.jl

oxinabox · 2023-12-07T04:01:19Z

src/compiler/reverse.jl

-    elseif isexpr(ex, :enter, :leave)
-      error("""try/catch is not supported.
-            Refer to the Zygote documentation for fixes.
-            https://fluxml.ai/Zygote.jl/latest/limitations
-            """)


I believer we still need this on versions of julia that do not have enter and leave?
so maybe we keep this and add a && VERSION<v"1.10" (is that the right bounds?)

oxinabox · 2023-12-07T04:02:55Z

src/compiler/reverse.jl

+    # This is corresponds to a catch blocks which technically
+    # has predecessors but they are not modelled in the IRTools CFG.
+    # We put an error message at the beginning of said block.
+    if has_leave && isempty(predecessors(b)) && b.id != 1
+        _, f_stmt = first(b)
+        li = pr.ir.lines[f_stmt.line]
+        li = LineNumberNode(Int(li.line), li.file)
+        pushfirst!(rb, stmt(xcall(Base, :error,
+                                  "Can't differentiate function execution in catch block at $(li).")))
+    end


I don't think throwing the error in the catch block is right.
I think we need to put it at the throw location.
Because I assume we also do not support finally blocks?

Try finally is implemented in term of enter and leave which can roughly be mapped to the following pseudo code:

try foo() origin = 1 catch origin = 2 end if origin == 2 rethrow() end

I will improve the error message to include try/finally as another construct which is not supported when an error is thrown in the block.

test/compiler.jl

oxinabox

Nice work.
See my comments

Co-authored-by: Frames White <oxinabox@ucc.asn.au>

This requires updating the IRTools version to include the better ssa conversion for try catch branches.

oxinabox · 2024-01-22T02:29:22Z

test/compiler.jl

+if VERSION >= v"1.8"
+    # try/catch/else is invalid syntax prior to v1.8
+    eval(Meta.parse("""
+        function try_catch_else(cond, x)
+        end
+    """))
+end


this code seems unused and incomplete

Pangoraw and others added 2 commits December 5, 2023 22:27

enable try/catch support on the happy (nothrow) path

3ddf945

use stacks for possibly undefined pullbacks

8f416aa

another possibility would be to thread the pullback values as block parameters through the control flow with arguments coming from the catch block being `nothing`.

Pangoraw marked this pull request as ready for review December 6, 2023 07:55

Pangoraw marked this pull request as draft December 6, 2023 08:25

Pangoraw added 2 commits December 6, 2023 09:39

Update features.jl

0836d94

Update emit.jl

07ec290

oxinabox self-requested a review December 7, 2023 03:16