From 5fec824e2f8578519b54633effc07acdaa00d72e Mon Sep 17 00:00:00 2001 From: Gabriel Baraldi Date: Mon, 2 Oct 2023 09:01:32 -0300 Subject: [PATCH 01/10] Change allocation function logic to be negative. --- src/allocfunc.jl | 54 ++++++++++++------------------------------------ 1 file changed, 13 insertions(+), 41 deletions(-) diff --git a/src/allocfunc.jl b/src/allocfunc.jl index c4a1c6e..4ce0704 100644 --- a/src/allocfunc.jl +++ b/src/allocfunc.jl @@ -1,49 +1,21 @@ # List of methods to location of arg which is the mi/function, then start of args const generic_method_offsets = Dict{String, Tuple{Int,Int}}(("jl_f__apply_latest" => (2,3), "ijl_f__apply_latest" => (2,3), "jl_f__call_latest" => (2,3), "ijl_f__call_latest" => (2,3), "jl_f_invoke" => (2,3), "jl_invoke" => (1,3), "jl_apply_generic" => (1,2), "ijl_f_invoke" => (2,3), "ijl_invoke" => (1,3), "ijl_apply_generic" => (1,2))) -const alloc_funcs = ( - "ijl_f__apply_latest", "ijl_f__call_latest", "ijl_f_invoke", "ijl_invoke", "ijl_apply_generic", - "ijl_alloc_array_1d", "ijl_alloc_array_2d", "ijl_alloc_array_3d", - "ijl_new_array", - "ijl_array_copy", - "ijl_alloc_string", - "ijl_in_threaded_region", "ijl_enter_threaded_region", "ijl_exit_threaded_region", "ijl_set_task_tid", "ijl_new_task", - "ijl_array_grow_beg", - "ijl_array_grow_end", - "ijl_array_grow_at", - "ijl_array_del_beg", - "ijl_array_del_end", - "ijl_array_del_at", - "ijl_gc_add_finalizer_th", - "ijl_symbol_n", "ijl_", - "ijl_reshape_array", "ijl_reshape_array", - "ijl_matching_methods", "ijl_matching_methods", - "ijl_array_sizehint", "ijl_array_sizehint", - "ijl_get_keyword_sorter", "ijl_get_keyword_sorter", - "ijl_ptr_to_array", - "ijl_box_float32", - "ijl_box_float64", - "ijl_box_int16", - "ijl_box_int32", - "ijl_box_int64", - "ijl_box_int8", - "ijl_box_slotnumber", - "ijl_box_ssavalue", - "ijl_box_uint16", - "ijl_box_uint32", - "ijl_box_uint64", - "ijl_box_uint8", - "ijl_box_uint8pointer", - "ijl_box_voidpointer", - "ijl_ptr_to_array_1d", - "ijl_eqtable_get", "ijl_eqtable_get", - "ijl_get_nth_field_checked", - "ijl_gc_alloc_typed", "ijl_gc_pool_alloc", "ijl_gc_big_alloc", - "ijl_gc_pool_alloc_instrumented", "ijl_gc_big_alloc_instrumented" - ) + +const known_nonalloc_funcs = [ + "jl_egal__unboxed", "ijl_egal__unboxed", + "jl_lock_value", "ijl_lock_value", + "jl_unlock_value", "ijl_unlock_value", + "jl_get_nth_field_noalloc", "ijl_get_nth_field_noalloc", +] function is_alloc_function(name) - name in alloc_funcs + maybe_alloc = occursin(r"(ijl_|jl_).*", name) + if maybe_alloc + name in known_nonalloc_funcs && return false + return true + end + return false end function guess_julia_type(val::LLVM.Value, typeof=true) From dd2c5d19f92cbad0293b7288c7347a9d7571e7e1 Mon Sep 17 00:00:00 2001 From: Gabriel Baraldi Date: Mon, 2 Oct 2023 11:11:42 -0300 Subject: [PATCH 02/10] Cleanup ccall code --- src/allocfunc.jl | 32 +++++++++++++++++++++++++++++--- 1 file changed, 29 insertions(+), 3 deletions(-) diff --git a/src/allocfunc.jl b/src/allocfunc.jl index 4ce0704..1467905 100644 --- a/src/allocfunc.jl +++ b/src/allocfunc.jl @@ -2,17 +2,23 @@ const generic_method_offsets = Dict{String, Tuple{Int,Int}}(("jl_f__apply_latest" => (2,3), "ijl_f__apply_latest" => (2,3), "jl_f__call_latest" => (2,3), "ijl_f__call_latest" => (2,3), "jl_f_invoke" => (2,3), "jl_invoke" => (1,3), "jl_apply_generic" => (1,2), "ijl_f_invoke" => (2,3), "ijl_invoke" => (1,3), "ijl_apply_generic" => (1,2))) -const known_nonalloc_funcs = [ +const known_nonalloc_funcs = ( "jl_egal__unboxed", "ijl_egal__unboxed", "jl_lock_value", "ijl_lock_value", "jl_unlock_value", "ijl_unlock_value", "jl_get_nth_field_noalloc", "ijl_get_nth_field_noalloc", -] + "jl_load_and_lookup","ijl_load_and_lookup", + "jl_lazy_load_and_lookup","ijl_lazy_load_and_lookup", + "jl_box_bool","ijl_box_bool", + "jl_box_int8","ijl_box_int8", + "jl_box_uint8","ijl_box_uint8", + r"(ijl|jl)_unbox.*" +) function is_alloc_function(name) maybe_alloc = occursin(r"(ijl_|jl_).*", name) if maybe_alloc - name in known_nonalloc_funcs && return false + any(x->contains(str, x), known_nonalloc_funcs) && return false return true end return false @@ -82,6 +88,26 @@ function rename_ir!(job, inst::LLVM.CallInst) method_table = Core.Compiler.method_table(interp) dest = called_operand(inst) + if isa(dest, LLVM.LoadInst) + fptr = LLVM.Value(LLVM.LLVM.API.LLVMGetOperand(dest, 0)) + if occursin("bitcast", string(dest)) + fn_got = LLVM.Value(LLVM.LLVM.API.LLVMGetOperand(fptr, 0)) + fname = name(fn_got) + if startswith(fname, "jlplt_") + fname = fname[7:end] + fname = replace(fname, r"_\d+_got$" => "") + end + mod = LLVM.parent(LLVM.parent(LLVM.parent(inst))) + lfn = LLVM.API.LLVMGetNamedFunction(mod, fname) + if lfn == C_NULL + lfn = LLVM.API.LLVMAddFunction(mod, Symbol(fname), LLVM.API.LLVMGetCalledFunctionType(inst)) + else + lfn = LLVM.API.LLVMConstBitCast(lfn, LLVM.PointerType(LLVM.FunctionType(LLVM.API.LLVMGetCalledFunctionType(inst)))) + end + LLVM.API.LLVMSetOperand(inst, LLVM.API.LLVMGetNumOperands(inst)-1, lfn) + end + end + if isa(dest, ConstantExpr) # Enzyme should be able to handle these # detect calls to literal pointers and replace with function name, if possible From 06eb22835b6bf5622a618998047c384c3482ff27 Mon Sep 17 00:00:00 2001 From: Gabriel Baraldi Date: Wed, 4 Oct 2023 09:25:31 -0300 Subject: [PATCH 03/10] Fix typo --- src/allocfunc.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/allocfunc.jl b/src/allocfunc.jl index 1467905..344b0d5 100644 --- a/src/allocfunc.jl +++ b/src/allocfunc.jl @@ -18,7 +18,7 @@ const known_nonalloc_funcs = ( function is_alloc_function(name) maybe_alloc = occursin(r"(ijl_|jl_).*", name) if maybe_alloc - any(x->contains(str, x), known_nonalloc_funcs) && return false + any(x->contains(name, x), known_nonalloc_funcs) && return false return true end return false @@ -107,7 +107,7 @@ function rename_ir!(job, inst::LLVM.CallInst) LLVM.API.LLVMSetOperand(inst, LLVM.API.LLVMGetNumOperands(inst)-1, lfn) end end - + if isa(dest, ConstantExpr) # Enzyme should be able to handle these # detect calls to literal pointers and replace with function name, if possible From 78f50e3fda1210fa7400715a013450e17506e86d Mon Sep 17 00:00:00 2001 From: Gabriel Baraldi Date: Wed, 4 Oct 2023 16:05:20 -0300 Subject: [PATCH 04/10] Apply sugestions from review --- src/allocfunc.jl | 34 +++++++++++++++------------------- 1 file changed, 15 insertions(+), 19 deletions(-) diff --git a/src/allocfunc.jl b/src/allocfunc.jl index 344b0d5..c4d80fe 100644 --- a/src/allocfunc.jl +++ b/src/allocfunc.jl @@ -1,5 +1,5 @@ # List of methods to location of arg which is the mi/function, then start of args -const generic_method_offsets = Dict{String, Tuple{Int,Int}}(("jl_f__apply_latest" => (2,3), "ijl_f__apply_latest" => (2,3), "jl_f__call_latest" => (2,3), "ijl_f__call_latest" => (2,3), "jl_f_invoke" => (2,3), "jl_invoke" => (1,3), "jl_apply_generic" => (1,2), "ijl_f_invoke" => (2,3), "ijl_invoke" => (1,3), "ijl_apply_generic" => (1,2))) +const generic_method_offsets = Dict{String,Tuple{Int,Int}}(("jl_f__apply_latest" => (2, 3), "ijl_f__apply_latest" => (2, 3), "jl_f__call_latest" => (2, 3), "ijl_f__call_latest" => (2, 3), "jl_f_invoke" => (2, 3), "jl_invoke" => (1, 3), "jl_apply_generic" => (1, 2), "ijl_f_invoke" => (2, 3), "ijl_invoke" => (1, 3), "ijl_apply_generic" => (1, 2))) const known_nonalloc_funcs = ( @@ -7,18 +7,18 @@ const known_nonalloc_funcs = ( "jl_lock_value", "ijl_lock_value", "jl_unlock_value", "ijl_unlock_value", "jl_get_nth_field_noalloc", "ijl_get_nth_field_noalloc", - "jl_load_and_lookup","ijl_load_and_lookup", - "jl_lazy_load_and_lookup","ijl_lazy_load_and_lookup", - "jl_box_bool","ijl_box_bool", - "jl_box_int8","ijl_box_int8", - "jl_box_uint8","ijl_box_uint8", + "jl_load_and_lookup", "ijl_load_and_lookup", + "jl_lazy_load_and_lookup", "ijl_lazy_load_and_lookup", + "jl_box_bool", "ijl_box_bool", + "jl_box_int8", "ijl_box_int8", + "jl_box_uint8", "ijl_box_uint8", r"(ijl|jl)_unbox.*" ) function is_alloc_function(name) maybe_alloc = occursin(r"(ijl_|jl_).*", name) if maybe_alloc - any(x->contains(name, x), known_nonalloc_funcs) && return false + any(x -> contains(name, x), known_nonalloc_funcs) && return false return true end return false @@ -93,18 +93,16 @@ function rename_ir!(job, inst::LLVM.CallInst) if occursin("bitcast", string(dest)) fn_got = LLVM.Value(LLVM.LLVM.API.LLVMGetOperand(fptr, 0)) fname = name(fn_got) - if startswith(fname, "jlplt_") + if startswith(fname, "jlplt_") && endswith(fname, "_got") fname = fname[7:end] fname = replace(fname, r"_\d+_got$" => "") + mod = LLVM.parent(LLVM.parent(LLVM.parent(inst))) + lfn = LLVM.API.LLVMGetNamedFunction(mod, fname) + if lfn == C_NULL + lfn = LLVM.API.LLVMAddFunction(mod, Symbol(fname), LLVM.API.LLVMGetCalledFunctionType(inst)) + end + LLVM.API.LLVMSetOperand(inst, LLVM.API.LLVMGetNumOperands(inst) - 1, lfn) end - mod = LLVM.parent(LLVM.parent(LLVM.parent(inst))) - lfn = LLVM.API.LLVMGetNamedFunction(mod, fname) - if lfn == C_NULL - lfn = LLVM.API.LLVMAddFunction(mod, Symbol(fname), LLVM.API.LLVMGetCalledFunctionType(inst)) - else - lfn = LLVM.API.LLVMConstBitCast(lfn, LLVM.PointerType(LLVM.FunctionType(LLVM.API.LLVMGetCalledFunctionType(inst)))) - end - LLVM.API.LLVMSetOperand(inst, LLVM.API.LLVMGetNumOperands(inst)-1, lfn) end end @@ -130,10 +128,8 @@ function rename_ir!(job, inst::LLVM.CallInst) lfn = LLVM.API.LLVMGetNamedFunction(mod, fn_str) if lfn == C_NULL lfn = LLVM.API.LLVMAddFunction(mod, fn, LLVM.API.LLVMGetCalledFunctionType(inst)) - else - lfn = LLVM.API.LLVMConstBitCast(lfn, LLVM.PointerType(LLVM.FunctionType(LLVM.API.LLVMGetCalledFunctionType(inst)))) end - LLVM.API.LLVMSetOperand(inst, LLVM.API.LLVMGetNumOperands(inst)-1, lfn) + LLVM.API.LLVMSetOperand(inst, LLVM.API.LLVMGetNumOperands(inst) - 1, lfn) end end end From 27dd0793fd531246e8ee7f09438ee8be016ca7dc Mon Sep 17 00:00:00 2001 From: Gabriel Baraldi Date: Wed, 4 Oct 2023 16:14:04 -0300 Subject: [PATCH 05/10] Add try catch funcs to list --- src/allocfunc.jl | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/allocfunc.jl b/src/allocfunc.jl index c4d80fe..5a3f22e 100644 --- a/src/allocfunc.jl +++ b/src/allocfunc.jl @@ -12,7 +12,11 @@ const known_nonalloc_funcs = ( "jl_box_bool", "ijl_box_bool", "jl_box_int8", "ijl_box_int8", "jl_box_uint8", "ijl_box_uint8", - r"(ijl|jl)_unbox.*" + r"(ijl|jl)_unbox.*", + "jl_excstack_state", "ijl_excstack_state", + "jl_restore_excstack", "ijl_restore_excstack", + "jl_enter_handler", "ijl_enter_handler", + "jl_pop_handler", "ijl_pop_handler", ) function is_alloc_function(name) From 923e04cd610013dd54c922043668e4b6e81f9741 Mon Sep 17 00:00:00 2001 From: Gabriel Baraldi Date: Wed, 4 Oct 2023 16:21:59 -0300 Subject: [PATCH 06/10] Add some more tests --- test/runtests.jl | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/test/runtests.jl b/test/runtests.jl index 921b579..97329bb 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -9,6 +9,12 @@ function alloc_in_catch() return Int64[] end +function same_ccall() + a = Array{Int}(undef,5,5) + b = Array{Int}(undef,5,5) + a,b +end + @testset "AllocCheck.jl" begin @test length(check_allocs(mod, (Float64,Float64))) == 0 @test length(check_allocs(sin, (Float64,); ignore_throw=false)) > 0 @@ -17,4 +23,7 @@ end @test length(check_allocs(alloc_in_catch, (); ignore_throw=false)) == 2 @test length(check_allocs(alloc_in_catch, (); ignore_throw=true)) == 1 + + @test length(check_allocs(same_ccall, (), ignore_throw=false)) == 2 + @test length(check_allocs(same_ccall, (), ignore_throw=true)) == 2 end From c666843d49a247fd0cde1022c81c7e41c7f877d4 Mon Sep 17 00:00:00 2001 From: Gabriel Baraldi Date: Wed, 11 Oct 2023 13:24:02 -0300 Subject: [PATCH 07/10] Also add some functions that may throw --- src/AllocCheck.jl | 4 ++-- src/allocfunc.jl | 21 ++++++++++++++++++--- 2 files changed, 20 insertions(+), 5 deletions(-) diff --git a/src/AllocCheck.jl b/src/AllocCheck.jl index 98b15c0..f3cdd44 100644 --- a/src/AllocCheck.jl +++ b/src/AllocCheck.jl @@ -187,7 +187,7 @@ function rename_calls_and_throws!(f::LLVM.Function, job) end # `catch`: Add pseudo-edge from any_catch - if name(decl) == "__sigsetjmp" + if name(decl) == "__sigsetjmp" || name(decl) == "sigsetjmp" icmp_ = user(only(uses(inst))) # Asserts one usage @assert icmp_ isa LLVM.ICmpInst br_ = user(only(uses(icmp_))) # Asserts one usage @@ -269,7 +269,7 @@ function check_allocs(@nospecialize(func), @nospecialize(types); entry_abi=:spec catch_only = dominates(domtree, catch_, inst) ignore_throw && catch_only && continue - if is_alloc_function(name(decl)) + if is_alloc_function(name(decl), ignore_throw) bt = backtrace_(inst; compiled) alloc = AllocInstance(inst, bt) push!(allocs, alloc) diff --git a/src/allocfunc.jl b/src/allocfunc.jl index 5a3f22e..523ac23 100644 --- a/src/allocfunc.jl +++ b/src/allocfunc.jl @@ -19,11 +19,26 @@ const known_nonalloc_funcs = ( "jl_pop_handler", "ijl_pop_handler", ) -function is_alloc_function(name) +const known_alloc_with_throw_funcs = ( + "jl_f_ifelse", "ijl_f_ifelse", + "jl_f_typeassert", "ijl_f_typeassert", + "jl_f_isa", "ijl_f_isa", + "jl_f_issubtype", "ijl_f_issubtype", + "jl_f_is", "ijl_f_is", + "jl_f_typeof", "ijl_f_typeof", + "jl_f_sizeof", "ijl_f_sizeof", + "jl_f_throw", "ijl_f_throw", +) + +function is_alloc_function(name, ignore_throw) maybe_alloc = occursin(r"(ijl_|jl_).*", name) if maybe_alloc - any(x -> contains(name, x), known_nonalloc_funcs) && return false - return true + has_alloc = false + if ignore_throw + has_alloc = any(x -> contains(name, x), known_alloc_with_throw_funcs) + end + has_alloc |= !any(x -> contains(name, x), known_nonalloc_funcs) + return has_alloc end return false end From 77d22d283cbd2445581cd1e2952b3ceb6a0fd69d Mon Sep 17 00:00:00 2001 From: Gabriel Baraldi Date: Wed, 11 Oct 2023 14:25:21 -0300 Subject: [PATCH 08/10] Change logic to make it clearer + test --- src/allocfunc.jl | 13 ++++++++----- test/runtests.jl | 3 +++ 2 files changed, 11 insertions(+), 5 deletions(-) diff --git a/src/allocfunc.jl b/src/allocfunc.jl index 523ac23..250ea76 100644 --- a/src/allocfunc.jl +++ b/src/allocfunc.jl @@ -28,17 +28,20 @@ const known_alloc_with_throw_funcs = ( "jl_f_typeof", "ijl_f_typeof", "jl_f_sizeof", "ijl_f_sizeof", "jl_f_throw", "ijl_f_throw", + "jl_f__svec_ref", "ijl_f__svec_ref", ) function is_alloc_function(name, ignore_throw) maybe_alloc = occursin(r"(ijl_|jl_).*", name) if maybe_alloc - has_alloc = false - if ignore_throw - has_alloc = any(x -> contains(name, x), known_alloc_with_throw_funcs) + is_throw_func = any(x -> contains(name, x), known_alloc_with_throw_funcs) + if is_throw_func && ignore_throw + return false + else + return true end - has_alloc |= !any(x -> contains(name, x), known_nonalloc_funcs) - return has_alloc + any(x -> contains(name, x), known_nonalloc_funcs) || return true + return false end return false end diff --git a/test/runtests.jl b/test/runtests.jl index 97329bb..9262568 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -26,4 +26,7 @@ end @test length(check_allocs(same_ccall, (), ignore_throw=false)) == 2 @test length(check_allocs(same_ccall, (), ignore_throw=true)) == 2 + + @test length(check_allocs(first, (Core.SimpleVector,); ignore_throw = false)) == 3 + @test length(check_allocs(first, (Core.SimpleVector,); ignore_throw = true)) == 0 end From b3dd48d3bb564d35e2e9a4bb68ff516d0ebe5f84 Mon Sep 17 00:00:00 2001 From: Gabriel Baraldi Date: Wed, 11 Oct 2023 14:32:17 -0300 Subject: [PATCH 09/10] Actually fix the logic --- src/allocfunc.jl | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/src/allocfunc.jl b/src/allocfunc.jl index 250ea76..1d93384 100644 --- a/src/allocfunc.jl +++ b/src/allocfunc.jl @@ -35,10 +35,12 @@ function is_alloc_function(name, ignore_throw) maybe_alloc = occursin(r"(ijl_|jl_).*", name) if maybe_alloc is_throw_func = any(x -> contains(name, x), known_alloc_with_throw_funcs) - if is_throw_func && ignore_throw - return false - else - return true + if is_throw_func + if ignore_throw + return false + else + return true + end end any(x -> contains(name, x), known_nonalloc_funcs) || return true return false From d65bc608c5cd287ef938381a9bcc81eab73ed6ba Mon Sep 17 00:00:00 2001 From: Gabriel Baraldi Date: Wed, 11 Oct 2023 16:28:57 -0300 Subject: [PATCH 10/10] Apply suggestions from code review --- src/allocfunc.jl | 13 +++++-------- 1 file changed, 5 insertions(+), 8 deletions(-) diff --git a/src/allocfunc.jl b/src/allocfunc.jl index 1d93384..828e99c 100644 --- a/src/allocfunc.jl +++ b/src/allocfunc.jl @@ -17,16 +17,13 @@ const known_nonalloc_funcs = ( "jl_restore_excstack", "ijl_restore_excstack", "jl_enter_handler", "ijl_enter_handler", "jl_pop_handler", "ijl_pop_handler", + "jl_f_typeof", "ijl_f_typeof", ) const known_alloc_with_throw_funcs = ( "jl_f_ifelse", "ijl_f_ifelse", "jl_f_typeassert", "ijl_f_typeassert", - "jl_f_isa", "ijl_f_isa", - "jl_f_issubtype", "ijl_f_issubtype", "jl_f_is", "ijl_f_is", - "jl_f_typeof", "ijl_f_typeof", - "jl_f_sizeof", "ijl_f_sizeof", "jl_f_throw", "ijl_f_throw", "jl_f__svec_ref", "ijl_f__svec_ref", ) @@ -114,12 +111,12 @@ function rename_ir!(job, inst::LLVM.CallInst) if isa(dest, LLVM.LoadInst) fptr = LLVM.Value(LLVM.LLVM.API.LLVMGetOperand(dest, 0)) - if occursin("bitcast", string(dest)) + if fptr isa LLVM.ConstantExpr && opcode(fptr) == LLVM.API.LLVMBitCast fn_got = LLVM.Value(LLVM.LLVM.API.LLVMGetOperand(fptr, 0)) fname = name(fn_got) - if startswith(fname, "jlplt_") && endswith(fname, "_got") - fname = fname[7:end] - fname = replace(fname, r"_\d+_got$" => "") + match_ = match(r"^jlplt_(.*)_\d+_got$", fname) + if match_ !== nothing + fname = match_[1] mod = LLVM.parent(LLVM.parent(LLVM.parent(inst))) lfn = LLVM.API.LLVMGetNamedFunction(mod, fname) if lfn == C_NULL