Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
27 changes: 14 additions & 13 deletions src/compiler.jl
Original file line number Diff line number Diff line change
Expand Up @@ -1997,6 +1997,17 @@ end
include("rules/allocrules.jl")
include("rules/llvmrules.jl")

function add_one_in_place(x)
if x isa Base.RefValue
x[] = recursive_add(x[], default_adjoint(eltype(Core.Typeof(x))))
elseif x isa (Array{T,0} where T)
x[] = recursive_add(x[], default_adjoint(eltype(Core.Typeof(x))))
else
throw(EnzymeNonScalarReturnException(x, ""))
end
return nothing
end

for (k, v) in (
("enz_runtime_newtask_fwd", Enzyme.Compiler.runtime_newtask_fwd),
("enz_runtime_newtask_augfwd", Enzyme.Compiler.runtime_newtask_augfwd),
Expand All @@ -2018,6 +2029,7 @@ for (k, v) in (
("enz_runtime_jl_setfield_rev", Enzyme.Compiler.rt_jl_setfield_rev),
("enz_runtime_error_if_differentiable", Enzyme.Compiler.error_if_differentiable),
("enz_runtime_error_if_active", Enzyme.Compiler.error_if_active),
("enz_add_one_in_place", Enzyme.Compiler.add_one_in_place),
)
JuliaEnzymeNameMap[k] = v
end
Expand Down Expand Up @@ -5072,7 +5084,7 @@ end
if !(primal_target isa GPUCompiler.NativeCompilerTarget)
reinsert_gcmarker!(adjointf)
augmented_primalf !== nothing && reinsert_gcmarker!(augmented_primalf)
post_optimze!(mod, target_machine, false) #=machine=#
post_optimize!(mod, target_machine, false) #=machine=#
end

adjointf = functions(mod)[adjointf_name]
Expand Down Expand Up @@ -5236,17 +5248,6 @@ include("typeutils/recursive_add.jl")
end
end

function add_one_in_place(x)
if x isa Base.RefValue
x[] = recursive_add(x[], default_adjoint(eltype(Core.Typeof(x))))
elseif x isa (Array{T,0} where T)
x[] = recursive_add(x[], default_adjoint(eltype(Core.Typeof(x))))
else
throw(EnzymeNonScalarReturnException(x, ""))
end
return nothing
end

@generated function enzyme_call(
::Val{RawCall},
fptr::PT,
Expand Down Expand Up @@ -5814,7 +5815,7 @@ function _thunk(job, postopt::Bool = true)::Tuple{LLVM.Module, Vector{Any}, Stri
if DumpPrePostOpt[]
API.EnzymeDumpModuleRef(mod.ref)
end
post_optimze!(mod, JIT.get_tm())
post_optimize!(mod, JIT.get_tm())
if DumpPostOpt[]
API.EnzymeDumpModuleRef(mod.ref)
end
Expand Down
10 changes: 9 additions & 1 deletion src/compiler/optimize.jl
Original file line number Diff line number Diff line change
Expand Up @@ -718,7 +718,7 @@ function addJuliaLegalizationPasses!(pm::LLVM.ModulePassManager, tm::LLVM.Target
end
end

function post_optimze!(mod::LLVM.Module, tm::LLVM.TargetMachine, machine::Bool = true)
function post_optimize!(mod::LLVM.Module, tm::LLVM.TargetMachine, machine::Bool = true)
addr13NoAlias(mod)
removeDeadArgs!(mod, tm)
for f in collect(functions(mod))
Expand Down Expand Up @@ -764,6 +764,14 @@ function post_optimze!(mod::LLVM.Module, tm::LLVM.TargetMachine, machine::Bool =
LLVM.run!(pm, mod)
end
end
for f in functions(mod)
if isempty(blocks(f))
continue
end
if !has_fn_attr(f, StringAttribute("frame-pointer"))
push!(function_attributes(f), StringAttribute("frame-pointer", "all"))
end
end
# @safe_show "post_mod", mod
# flush(stdout)
# flush(stderr)
Expand Down
2 changes: 1 addition & 1 deletion src/compiler/reflection.jl
Original file line number Diff line number Diff line change
Expand Up @@ -74,7 +74,7 @@ function reflect(
mod, meta = GPUCompiler.codegen(:llvm, job) #= validate=false =#

if second_stage
post_optimze!(mod, JIT.get_tm())
post_optimize!(mod, JIT.get_tm())
end

llvmf = meta.adjointf
Expand Down
Loading