Skip to content

Commit

Permalink
fix
Browse files Browse the repository at this point in the history
  • Loading branch information
wsmoses committed Dec 3, 2024
1 parent 2956b4d commit f2dec4c
Showing 1 changed file with 4 additions and 3 deletions.
7 changes: 4 additions & 3 deletions src/compiler.jl
Original file line number Diff line number Diff line change
Expand Up @@ -5261,7 +5261,7 @@ end
const DumpPostOpt = Ref(false)

# actual compilation
function _thunk(job, postopt::Bool = true)
function _thunk(job, postopt::Bool = true)::Tuple{LLVM.Module, String, Union{String, Nothing}, Type, String}
mod, meta = codegen(:llvm, job; optimize = false)
adjointf, augmented_primalf = meta.adjointf, meta.augmented_primalf

Expand Down Expand Up @@ -5315,15 +5315,16 @@ const cache_lock = ReentrantLock()
if obj === nothing
asm = _thunk(job)
obj = _link(job, asm...)
@show obj
if obj.adjoint isa Ptr{Nothing}
autodiff_cache[obj.adjoint] = (asm[2], asm[5])
end
if obj.primal isa Ptr{Nothing}
if obj.primal isa Ptr{Nothing} && asm[3] isa String
autodiff_cache[obj.primal] = (asm[3], asm[5])
end
cache[key] = obj
end
nothing
obj
finally
unlock(cache_lock)
end
Expand Down

0 comments on commit f2dec4c

Please sign in to comment.