EnzymeAD / Enzyme.jl

Julia bindings for the Enzyme automatic differentiator
https://enzyme.mit.edu
MIT License
428 stars 59 forks source link

Enzyme rules for `partialsort!` #1373

Closed jgreener64 closed 3 months ago

jgreener64 commented 3 months ago

In fixing https://github.com/EnzymeAD/Enzyme.jl/issues/880 I saw the need for an Enzyme rule for partialsort!, which fixes that issue and may be useful elsewhere since we have a rule for sort!.

codecov-commenter commented 3 months ago

Codecov Report

Attention: Patch coverage is 0% with 53 lines in your changes are missing coverage. Please review.

Project coverage is 74.73%. Comparing base (fe036b9) to head (82da102). Report is 3 commits behind head on main.

Files Patch % Lines
src/internal_rules.jl 0.00% 53 Missing :warning:

:exclamation: Your organization needs to install the Codecov GitHub app to enable full functionality.

Additional details and impacted files ```diff @@ Coverage Diff @@ ## main #1373 +/- ## ========================================== - Coverage 75.41% 74.73% -0.68% ========================================== Files 35 35 Lines 10664 10716 +52 ========================================== - Hits 8042 8009 -33 - Misses 2622 2707 +85 ```

:umbrella: View full report in Codecov by Sentry.
:loudspeaker: Have feedback on the report? Share it here.

jgreener64 commented 3 months ago

I can reproduce the Julia 1.7 CI failure locally, but I don't immediately see how to fix it.

With this PR and Julia 1.7.0 on Windows:

using Enzyme, Test

function f(x)
    a = [2.0, 2.5, x, x / 2]
    y = partialsort(a, 1:2)
    return sum(y)
end

@test autodiff(Forward, f, Duplicated(1.5, 1.0))[1] == 1.5 # Passes

@test autodiff(Forward, f, BatchDuplicated(1.5, (1.0, 2.0)))[1] == (var"1"=1.5, var"2"=3.0)
Assertion failed: Tracked.size() == BaseNumbers.size(), file /cygdrive/c/buildbot/worker/package_win64/build/src/llvm-late-gc-lowering.cpp, line 870

signal (22): SIGABRT
in expression starting at REPL[4]:1
crt_sig_handler at /cygdrive/c/buildbot/worker/package_win64/build/src\signals-win.c:92
raise at C:\WINDOWS\System32\msvcrt.dll (unknown line)
abort at C:\WINDOWS\System32\msvcrt.dll (unknown line)
assert at C:\WINDOWS\System32\msvcrt.dll (unknown line)
NumberAllBase at /cygdrive/c/buildbot/worker/package_win64/build/src\llvm-late-gc-lowering.cpp:870
NumberBase at /cygdrive/c/buildbot/worker/package_win64/build/src\llvm-late-gc-lowering.cpp:787
Number at /cygdrive/c/buildbot/worker/package_win64/build/src\llvm-late-gc-lowering.cpp:804
LocalScan at /cygdrive/c/buildbot/worker/package_win64/build/src\llvm-late-gc-lowering.cpp:1581
runOnFunction at /cygdrive/c/buildbot/worker/package_win64/build/src\llvm-late-gc-lowering.cpp:2654
_ZN4llvm13FPPassManager13runOnFunctionERNS_8FunctionE at C:\Users\Joe\Documents\soft\Julia\Julia-1.7.0\bin\libLLVM.dll (unknown line)
_ZN4llvm13FPPassManager11runOnModuleERNS_6ModuleE at C:\Users\Joe\Documents\soft\Julia\Julia-1.7.0\bin\libLLVM.dll (unknown line)
_ZN4llvm6legacy15PassManagerImpl3runERNS_6ModuleE at C:\Users\Joe\Documents\soft\Julia\Julia-1.7.0\bin\libLLVM.dll (unknown line)
LLVMRunPassManager at C:\Users\Joe\Documents\soft\Julia\Julia-1.7.0\bin\libLLVM.dll (unknown line)
LLVMRunPassManager at C:\Users\Joe\.julia\packages\LLVM\Od0DH\lib\12\libLLVM_h.jl:4741 [inlined]
run! at C:\Users\Joe\.julia\packages\LLVM\Od0DH\src\passmanager.jl:39 [inlined]
#433 at C:\Users\Joe\.julia\dev\Enzyme\src\compiler\optimize.jl:1709
#ModulePassManager#47 at C:\Users\Joe\.julia\packages\LLVM\Od0DH\src\passmanager.jl:33
ModulePassManager at C:\Users\Joe\.julia\packages\LLVM\Od0DH\src\passmanager.jl:31 [inlined]
post_optimze! at C:\Users\Joe\.julia\dev\Enzyme\src\compiler\optimize.jl:1706
post_optimze! at C:\Users\Joe\.julia\dev\Enzyme\src\compiler\optimize.jl:1685
unknown function (ip: 000000005bdbbfa8)
_thunk at C:\Users\Joe\.julia\dev\Enzyme\src\compiler.jl:5597
_thunk at C:\Users\Joe\.julia\dev\Enzyme\src\compiler.jl:5579 [inlined]
cached_compilation at C:\Users\Joe\.julia\dev\Enzyme\src\compiler.jl:5613 [inlined]
#525 at C:\Users\Joe\.julia\dev\Enzyme\src\compiler.jl:5679
JuliaContext at C:\Users\Joe\.julia\packages\GPUCompiler\YO8Uj\src\driver.jl:58
unknown function (ip: 0000000052d43700)
#s1763#524 at C:\Users\Joe\.julia\dev\Enzyme\src\compiler.jl:5631 [inlined]
#s1763#524 at .\none:0
GeneratedFunctionStub at .\boot.jl:580
jl_apply at /cygdrive/c/buildbot/worker/package_win64/build/src\julia.h:1788 [inlined]
jl_call_staged at /cygdrive/c/buildbot/worker/package_win64/build/src\method.c:431
jl_code_for_staged at /cygdrive/c/buildbot/worker/package_win64/build/src\method.c:482
get_staged at .\compiler\utilities.jl:111
retrieve_code_info at .\compiler\utilities.jl:123 [inlined]
InferenceState at .\compiler\inferencestate.jl:234
typeinf_edge at .\compiler\typeinfer.jl:814 [inlined]
abstract_call_method at .\compiler\abstractinterpretation.jl:504
abstract_call_gf_by_type at .\compiler\abstractinterpretation.jl:105
abstract_call_known at .\compiler\abstractinterpretation.jl:1342
abstract_call at .\compiler\abstractinterpretation.jl:1397
abstract_call at .\compiler\abstractinterpretation.jl:1382
abstract_eval_statement at .\compiler\abstractinterpretation.jl:1534
typeinf_local at .\compiler\abstractinterpretation.jl:1900
typeinf_nocycle at .\compiler\abstractinterpretation.jl:2014
_typeinf at .\compiler\typeinfer.jl:226
typeinf at .\compiler\typeinfer.jl:209
typeinf_edge at .\compiler\typeinfer.jl:823 [inlined]
abstract_call_method at .\compiler\abstractinterpretation.jl:504
abstract_call_gf_by_type at .\compiler\abstractinterpretation.jl:105
abstract_call_known at .\compiler\abstractinterpretation.jl:1342
abstract_call at .\compiler\abstractinterpretation.jl:1397
abstract_apply at .\compiler\abstractinterpretation.jl:987
abstract_call_known at .\compiler\abstractinterpretation.jl:1249
abstract_call at .\compiler\abstractinterpretation.jl:1397
abstract_call at .\compiler\abstractinterpretation.jl:1382
abstract_eval_statement at .\compiler\abstractinterpretation.jl:1534
typeinf_local at .\compiler\abstractinterpretation.jl:1918
typeinf_nocycle at .\compiler\abstractinterpretation.jl:2014
_typeinf at .\compiler\typeinfer.jl:226
typeinf at .\compiler\typeinfer.jl:209
typeinf_edge at .\compiler\typeinfer.jl:823 [inlined]
abstract_call_method at .\compiler\abstractinterpretation.jl:504
abstract_call_gf_by_type at .\compiler\abstractinterpretation.jl:105
abstract_call_known at .\compiler\abstractinterpretation.jl:1342
abstract_call at .\compiler\abstractinterpretation.jl:1397
abstract_apply at .\compiler\abstractinterpretation.jl:987
abstract_call_known at .\compiler\abstractinterpretation.jl:1249
abstract_call at .\compiler\abstractinterpretation.jl:1397
abstract_call at .\compiler\abstractinterpretation.jl:1382
abstract_eval_statement at .\compiler\abstractinterpretation.jl:1534
typeinf_local at .\compiler\abstractinterpretation.jl:1918
typeinf_nocycle at .\compiler\abstractinterpretation.jl:2014
_typeinf at .\compiler\typeinfer.jl:226
typeinf at .\compiler\typeinfer.jl:209
typeinf_ext at .\compiler\typeinfer.jl:909
typeinf_ext_toplevel at .\compiler\typeinfer.jl:942
typeinf_ext_toplevel at .\compiler\typeinfer.jl:938
jfptr_typeinf_ext_toplevel_14217.clone_1 at C:\Users\Joe\Documents\soft\Julia\Julia-1.7.0\lib\julia\sys.dll (unknown line)
_jl_invoke at /cygdrive/c/buildbot/worker/package_win64/build/src\gf.c:2247 [inlined]
jl_apply_generic at /cygdrive/c/buildbot/worker/package_win64/build/src\gf.c:2429 [inlined]
jl_apply at /cygdrive/c/buildbot/worker/package_win64/build/src\julia.h:1788 [inlined]
jl_type_infer at /cygdrive/c/buildbot/worker/package_win64/build/src\gf.c:295
jl_generate_fptr at /cygdrive/c/buildbot/worker/package_win64/build/src\jitlayers.cpp:338
jl_compile_method_internal at /cygdrive/c/buildbot/worker/package_win64/build/src\gf.c:1980
jl_compile_method_internal at /cygdrive/c/buildbot/worker/package_win64/build/src\gf.c:1934 [inlined]
_jl_invoke at /cygdrive/c/buildbot/worker/package_win64/build/src\gf.c:2239 [inlined]
jl_apply_generic at /cygdrive/c/buildbot/worker/package_win64/build/src\gf.c:2429
jl_apply at /cygdrive/c/buildbot/worker/package_win64/build/src\julia.h:1788 [inlined]
do_call at /cygdrive/c/buildbot/worker/package_win64/build/src\interpreter.c:126
eval_value at /cygdrive/c/buildbot/worker/package_win64/build/src\interpreter.c:215
eval_stmt_value at /cygdrive/c/buildbot/worker/package_win64/build/src\interpreter.c:166 [inlined]
eval_body at /cygdrive/c/buildbot/worker/package_win64/build/src\interpreter.c:583
eval_body at /cygdrive/c/buildbot/worker/package_win64/build/src\interpreter.c:516
jl_interpret_toplevel_thunk at /cygdrive/c/buildbot/worker/package_win64/build/src\interpreter.c:731
jl_toplevel_eval_flex at /cygdrive/c/buildbot/worker/package_win64/build/src\toplevel.c:885
jl_toplevel_eval_flex at /cygdrive/c/buildbot/worker/package_win64/build/src\toplevel.c:830
jl_toplevel_eval at /cygdrive/c/buildbot/worker/package_win64/build/src\toplevel.c:894 [inlined]
jl_toplevel_eval_in at /cygdrive/c/buildbot/worker/package_win64/build/src\toplevel.c:944
eval at .\boot.jl:373 [inlined]
eval_user_input at C:\buildbot\worker\package_win64\build\usr\share\julia\stdlib\v1.7\REPL\src\REPL.jl:150
repl_backend_loop at C:\buildbot\worker\package_win64\build\usr\share\julia\stdlib\v1.7\REPL\src\REPL.jl:244
start_repl_backend at C:\buildbot\worker\package_win64\build\usr\share\julia\stdlib\v1.7\REPL\src\REPL.jl:229
#run_repl#47 at C:\buildbot\worker\package_win64\build\usr\share\julia\stdlib\v1.7\REPL\src\REPL.jl:362
run_repl at C:\buildbot\worker\package_win64\build\usr\share\julia\stdlib\v1.7\REPL\src\REPL.jl:349
#930 at .\client.jl:394
jfptr_YY.930_37269.clone_1 at C:\Users\Joe\Documents\soft\Julia\Julia-1.7.0\lib\julia\sys.dll (unknown line)
jl_apply at /cygdrive/c/buildbot/worker/package_win64/build/src\julia.h:1788 [inlined]
jl_f__call_latest at /cygdrive/c/buildbot/worker/package_win64/build/src\builtins.c:757
#invokelatest#2 at .\essentials.jl:716 [inlined]
invokelatest at .\essentials.jl:714 [inlined]
run_main_repl at .\client.jl:379
exec_options at .\client.jl:309
_start at .\client.jl:495
jfptr__start_41894.clone_1 at C:\Users\Joe\Documents\soft\Julia\Julia-1.7.0\lib\julia\sys.dll (unknown line)
jl_apply at /cygdrive/c/buildbot/worker/package_win64/build/src\julia.h:1788 [inlined]
true_main at /cygdrive/c/buildbot/worker/package_win64/build/src\jlapi.c:559
jl_repl_entrypoint at /cygdrive/c/buildbot/worker/package_win64/build/src\jlapi.c:701
mainCRTStartup at /cygdrive/c/buildbot/worker/package_win64/build/cli\loader_exe.c:42
BaseThreadInitThunk at C:\WINDOWS\System32\KERNEL32.DLL (unknown line)
RtlUserThreadStart at C:\WINDOWS\SYSTEM32\ntdll.dll (unknown line)
Allocations: 76365419 (Pool: 76315715; Big: 49704); GC: 78
wsmoses commented 3 months ago

@jgreener64 I think this is a bug in an older julia.

If you just version gate the test, I think it's fine

jgreener64 commented 3 months ago

I think the remaining test failures are unrelated.