Skip to content

Commit

Permalink
Handle Expr(:boundscheck)
Browse files Browse the repository at this point in the history
  • Loading branch information
ToucheSir committed Oct 11, 2023
1 parent cf7f7d0 commit efcc64b
Show file tree
Hide file tree
Showing 2 changed files with 13 additions and 0 deletions.
4 changes: 4 additions & 0 deletions src/compiler/reverse.jl
Original file line number Diff line number Diff line change
Expand Up @@ -132,6 +132,10 @@ function instrument(ir::IR)
elseif isexpr(ex, :(=))
@assert ex.args[1] isa GlobalRef
pr[v] = xcall(Zygote, :global_set, QuoteNode(ex.args[1]), ex.args[2])
elseif isexpr(ex, :boundscheck)
# Expr(:boundscheck) now appears in common Julia code paths, so we need to handle it.
# For correctness sake, fix to true like https://github.com/dfdx/Umlaut.jl/issues/34.
pr[v] = true
else
ex = instrument_new!(pr, v, ex)
ex = instrument_literals!(pr, v, ex)
Expand Down
9 changes: 9 additions & 0 deletions test/compiler.jl
Original file line number Diff line number Diff line change
Expand Up @@ -225,3 +225,12 @@ end

# issue 897
@test gradient(x -> sum(norm, collect(eachcol(x))), ones(3, 400))[1] fill(0.5773502691896258, 3, 400)

# Tests adapted from https://github.com/dfdx/Umlaut.jl/pull/35
@eval _boundscheck_foo(x) = ifelse($(Expr(:boundscheck)), 2x, x)

@testset "Meta Expr handling" begin
y, (dx,) = withgradient(_boundscheck_foo, 1)
@test y == 2
@test dx == 2
end

0 comments on commit efcc64b

Please sign in to comment.