Skip to content

Fix conditional wrappers for vararg runtime calls #3290

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.

Sign up for GitHub

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

Draft
ChrisRackauckas-Claude wants to merge 1 commit into
base: main
Choose a base branch
from
+40 −5
Draft

Fix conditional wrappers for vararg runtime calls #3290

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 14 additions & 4 deletions src/gradientutils.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -106,9 +106,13 @@ function set_reverse_block!(gutils::GradientUtils, block::LLVM.BasicBlock)
return API.EnzymeGradientUtilsSetReverseBlock(gutils, block)
end

function get_or_insert_conditional_execute!(fn::LLVM.Function; force_run=false, need_result=true, preprocess=nothing, postprocess=nothing, postprocess_const=nothing, cmpidx::Int = 1)
function get_or_insert_conditional_execute!(fn::LLVM.Function; force_run=false, need_result=true, preprocess=nothing, postprocess=nothing, postprocess_const=nothing, cmpidx::Int = 1, arg_tys=nothing)
FT0 = LLVM.function_type(fn)
ptys = LLVM.parameters(FT0)
if length(ptys) < cmpidx
@assert arg_tys !== nothing
ptys = collect(LLVM.LLVMType, arg_tys)
end
insert!(ptys, 1, ptys[cmpidx])

void_rt = LLVM.return_type(FT0) == LLVM.VoidType() || !need_result
Expand All @@ -134,6 +138,9 @@ function get_or_insert_conditional_execute!(fn::LLVM.Function; force_run=false,
if postprocess_const !== nothing
newname = newname * ".poc_$(postprocess_const)"
end
if arg_tys !== nothing
newname = newname * ".argtys$(hash(string.(arg_tys)))"
end
newname = newname * LLVM.name(fn)
cfn, _ = get_function!(mod, newname, FT)
if isempty(blocks(cfn))
Expand All @@ -151,7 +158,7 @@ function get_or_insert_conditional_execute!(fn::LLVM.Function; force_run=false,

if force_run
if preprocess !== nothing
ppr = preprocess(builder, args)
ppr = preprocess(builder, rparms)
end
res = call!(builder, FT0, fn, rparms)
callconv!(res, callconv(fn))
Expand Down Expand Up @@ -275,14 +282,17 @@ function call_same_with_inverted_arg_if_active!(
valTys[cmpidx] = API.VT_Both
end
args = collect(LLVM.Value, args)
prefn = LLVM.called_operand(orig)::LLVM.Function
arg_tys = if length(LLVM.parameters(LLVM.function_type(prefn))) < cmpidx
LLVM.LLVMType[value_type(arg) for arg in args]
end
insert!(args, 1, new_from_original(gutils, origops[cmpidx]))
newval = nothing
if value_type(orig) != LLVM.VoidType() && postprocess_const === nothing && need_result
newval = new_from_original(gutils, orig)
insert!(args, 1, newval)
end
prefn = LLVM.called_operand(orig)::LLVM.Function
condfn = get_or_insert_conditional_execute!(prefn; force_run, preprocess, postprocess, postprocess_const, need_result, cmpidx)
condfn = get_or_insert_conditional_execute!(prefn; force_run, preprocess, postprocess, postprocess_const, need_result, cmpidx, arg_tys)

res = LLVM.Value(
API.EnzymeGradientUtilsCallWithInvertedBundles(
Expand Down
27 changes: 26 additions & 1 deletion test/typeunstable.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -285,4 +285,29 @@ end
Enzyme.autodiff(set_runtime_activity(Reverse), Enzyme.Const(loss_3280), Enzyme.Active,
Enzyme.Duplicated(p, dp))
@test dp ≈ [1.0]
end
end

mutable struct TypeUnstableSetpropertyBox
x::Any
end

@noinline function setproperty_unstable!(box, name, x)::Any
setproperty!(box, name, x)
return box
end

function setproperty_unstable_fn(x)
box = TypeUnstableSetpropertyBox(nothing)
name = Base.inferencebarrier(:x)
setproperty_unstable!(box, name, x[1])
return box.x
end

@testset "Forward type-unstable setproperty!" begin
res = Enzyme.autodiff(
set_runtime_activity(Forward),
setproperty_unstable_fn,
Duplicated([2.0], [3.0]),
)
@test res[1] ≈ 3.0
end