Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Merged
merged 9 commits into from
Sep 24, 2024

Conversation

Pangoraw
Copy link
Contributor

@Pangoraw Pangoraw commented Dec 5, 2023

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.

Pangoraw and others added 2 commits December 5, 2023 22:27
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
Copy link
Member

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?

@Pangoraw Pangoraw marked this pull request as ready for review December 6, 2023 07:55
@Pangoraw Pangoraw marked this pull request as draft December 6, 2023 08:25
@oxinabox oxinabox self-requested a review December 7, 2023 03:16
docs/src/limitations.md Outdated Show resolved Hide resolved
Comment on lines -127 to -131
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
""")
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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?)

Comment on lines 325 to 334
# 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
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 Outdated Show resolved Hide resolved
Copy link
Member

@oxinabox oxinabox left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nice work.
See my comments

Pangoraw and others added 3 commits January 2, 2024 18:28
This requires updating the IRTools version to include the better ssa
conversion for try catch branches.
@Pangoraw

This comment was marked as resolved.

test/compiler.jl Outdated
Comment on lines 276 to 282
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
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this code seems unused and incomplete

@Pangoraw Pangoraw marked this pull request as ready for review January 29, 2024 22:54
@ToucheSir
Copy link
Member

@oxinabox is this still good to merge?

@ToucheSir ToucheSir mentioned this pull request Sep 22, 2024
2 tasks
@oxinabox
Copy link
Member

oxinabox commented Sep 24, 2024

I see no reason why it wouldn't be.
I thought it was merged back in January

@ToucheSir ToucheSir merged commit 406a6f5 into FluxML:master Sep 24, 2024
10 of 13 checks passed
@Pangoraw Pangoraw deleted the try_catch branch September 24, 2024 13:28
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants