Fix compile time bug (#2556)

Author: wsmoses

src/compiler.jl

@@ -482,8 +482,7 @@ function nested_codegen!(mode::API.CDerivativeMode, mod::LLVM.Module, @nospecial
     nested_codegen!(mode, mod, funcspec, world)
 end
 
-function prepare_llvm(mod::LLVM.Module, job, meta)
-    interp = GPUCompiler.get_interpreter(job)
+function prepare_llvm(interp, mod::LLVM.Module, job, meta)
     for f in functions(mod)
         attributes = function_attributes(f)
         push!(attributes, StringAttribute("enzymejl_world", string(job.world)))
@@ -1001,7 +1000,7 @@ end
     return
 end
 
-function set_module_types!(mod::LLVM.Module, primalf::Union{Nothing, LLVM.Function}, job, edges, run_enzyme, mode::API.CDerivativeMode)
+function set_module_types!(interp, mod::LLVM.Module, primalf::Union{Nothing, LLVM.Function}, job, edges, run_enzyme, mode::API.CDerivativeMode)
 
     for f in functions(mod)
         mi, RT = enzyme_custom_extract_mi(f, false)
@@ -1141,7 +1140,6 @@ function set_module_types!(mod::LLVM.Module, primalf::Union{Nothing, LLVM.Functi
     custom = Dict{String,LLVM.API.LLVMLinkage}()
 
     world = job.world
-    interp = GPUCompiler.get_interpreter(job)
     method_table = Core.Compiler.method_table(interp)
 
     state = HandlerState(
@@ -1203,7 +1201,8 @@ function nested_codegen!(
     GPUCompiler.prepare_job!(job)
     otherMod, meta = GPUCompiler.emit_llvm(job)
 
-    prepare_llvm(otherMod, job, meta)
+    interp = GPUCompiler.get_interpreter(job)
+    prepare_llvm(interp, otherMod, job, meta)
 
     entry = name(meta.entry)
 
@@ -1221,12 +1220,12 @@ function nested_codegen!(
         LLVM.run!(pm, otherMod)
     end
 
-    check_ir(job, otherMod)
+    check_ir(interp, job, otherMod)
 
     # Skipped inline of blas
 
     run_enzyme = false
-    set_module_types!(otherMod, nothing, job, edges, run_enzyme, mode)
+    set_module_types!(interp, otherMod, nothing, job, edges, run_enzyme, mode)
 
     # Apply first stage of optimization's so that this module is at the same stage as `mod`
     optimize!(otherMod, JIT.get_tm())
@@ -2233,6 +2232,7 @@ const DumpPostWrap = Ref(false)
 
 function enzyme!(
     job::CompilerJob,
+    interp,
     mod::LLVM.Module,
     primalf::LLVM.Function,
     @nospecialize(TT::Type),
@@ -2251,7 +2251,6 @@ function enzyme!(
         API.EnzymeDumpModuleRef(mod.ref)
     end
     world = job.world
-    interp = GPUCompiler.get_interpreter(job)
     rt = job.config.params.rt
     runtimeActivity = job.config.params.runtimeActivity
     strongZero = job.config.params.strongZero
@@ -4290,14 +4289,14 @@ function GPUCompiler.compile_unhooked(output::Symbol, job::CompilerJob{<:EnzymeT
         # ??? entry_abi
     )
     primal_job = CompilerJob(primal, primal_config, job.world)
-
     @safe_debug "Emit LLVM with" primal_job
     GPUCompiler.prepare_job!(primal_job)
     mod, meta = GPUCompiler.emit_llvm(primal_job)
     edges = Any[]
     mod_to_edges[mod] = edges
 
-    prepare_llvm(mod, primal_job, meta)
+    primal_interp = GPUCompiler.get_interpreter(primal_job)
+    prepare_llvm(primal_interp, mod, primal_job, meta)
     for f in functions(mod)
         permit_inlining!(f)
     end
@@ -4311,7 +4310,8 @@ function GPUCompiler.compile_unhooked(output::Symbol, job::CompilerJob{<:EnzymeT
     if DumpPreCheck[]
         API.EnzymeDumpModuleRef(mod.ref)
     end
-    check_ir(job, mod)
+    interp = GPUCompiler.get_interpreter(job)
+    check_ir(interp, job, mod)
 
     disableFallback = String[]
 
@@ -4420,7 +4420,7 @@ function GPUCompiler.compile_unhooked(output::Symbol, job::CompilerJob{<:EnzymeT
         end
     end
 
-    custom, state = set_module_types!(mod, primalf, job, edges, params.run_enzyme, mode)
+    custom, state = set_module_types!(interp, mod, primalf, job, edges, params.run_enzyme, mode)
 
     primalf = state.primalf
     must_wrap = state.must_wrap
@@ -4899,6 +4899,7 @@ end
 
         adjointf, augmented_primalf, TapeType = enzyme!(
             job,
+	    interp,
             mod,
             primalf,
             TT,

src/compiler/interpreter.jl

@@ -138,14 +138,26 @@ struct EnzymeInterpreter{T} <: AbstractInterpreter
     handler::T
 end
 
-
+const SigCache = Dict{Tuple, Dict{UInt, Base.IdSet{Type}}}()
 function get_rule_signatures(f, TT, world)
+    subdict = if haskey(SigCache, (f, TT))
+       SigCache[(f, TT)]
+    else
+       tmp = Dict{UInt, Base.IdSet{Type}}()
+       SigCache[(f, TT)] = tmp
+       tmp
+    end
+    if haskey(subdict, world)
+       return subdict[world]
+    end
     fwdrules_meths = Base._methods(f, TT, -1, world)::Vector
     sigs = Type[]
     for rule in fwdrules_meths
         push!(sigs, (rule::Core.MethodMatch).method.sig)
     end
-    return Base.IdSet{Type}(sigs)
+    result = Base.IdSet{Type}(sigs)
+    subdict[world] = result
+    return result
 end
 
 function rule_sigs_equal(a, b)

src/compiler/validation.jl

@@ -177,15 +177,15 @@ function restore_lookups(mod::LLVM.Module)::Nothing
     end
 end
 
-function check_ir(@nospecialize(job::CompilerJob), mod::LLVM.Module)
-    errors = check_ir!(job, IRError[], mod)
+function check_ir(interp, @nospecialize(job::CompilerJob), mod::LLVM.Module)
+    errors = check_ir!(interp, job, IRError[], mod)
     unique!(errors)
     if !isempty(errors)
         throw(InvalidIRError(job, errors))
     end
 end
 
-function check_ir!(@nospecialize(job::CompilerJob), errors::Vector{IRError}, mod::LLVM.Module)
+function check_ir!(interp, @nospecialize(job::CompilerJob), errors::Vector{IRError}, mod::LLVM.Module)
     imported = Set(String[])
     if haskey(functions(mod), "malloc")
         f = functions(mod)["malloc"]
@@ -209,7 +209,7 @@ function check_ir!(@nospecialize(job::CompilerJob), errors::Vector{IRError}, mod
         if in(f, del)
             continue
         end
-        check_ir!(job, errors, imported, f, del, mod)
+        check_ir!(interp, job, errors, imported, f, del, mod)
     end
     for d in del
         LLVM.API.LLVMDeleteFunction(d)
@@ -220,7 +220,7 @@ function check_ir!(@nospecialize(job::CompilerJob), errors::Vector{IRError}, mod
         if in(f, del)
             continue
         end
-        check_ir!(job, errors, imported, f, del, mod)
+        check_ir!(interp, job, errors, imported, f, del, mod)
     end
     for d in del
         LLVM.API.LLVMDeleteFunction(d)
@@ -229,7 +229,7 @@ function check_ir!(@nospecialize(job::CompilerJob), errors::Vector{IRError}, mod
     return errors
 end
 
-function check_ir!(@nospecialize(job::CompilerJob), errors::Vector{IRError}, imported::Set{String}, f::LLVM.Function, deletedfns::Vector{LLVM.Function}, mod::LLVM.Module)
+function check_ir!(interp, @nospecialize(job::CompilerJob), errors::Vector{IRError}, imported::Set{String}, f::LLVM.Function, deletedfns::Vector{LLVM.Function}, mod::LLVM.Module)
     calls = LLVM.CallInst[]
     isInline = API.EnzymeGetCLBool(cglobal((:EnzymeInline, API.libEnzyme))) != 0
     mod = LLVM.parent(f)
@@ -466,7 +466,7 @@ function check_ir!(@nospecialize(job::CompilerJob), errors::Vector{IRError}, imp
 
     while length(calls) > 0
         inst = pop!(calls)
-        check_ir!(job, errors, imported, inst, calls, mod)
+        check_ir!(interp, job, errors, imported, inst, calls, mod)
     end
     return errors
 end
@@ -610,9 +610,8 @@ end
 import GPUCompiler:
     DYNAMIC_CALL, DELAYED_BINDING, RUNTIME_FUNCTION, UNKNOWN_FUNCTION, POINTER_FUNCTION
 import GPUCompiler: backtrace, isintrinsic
-function check_ir!(@nospecialize(job::CompilerJob), errors::Vector{IRError}, imported::Set{String}, inst::LLVM.CallInst, calls::Vector{LLVM.CallInst}, mod::LLVM.Module)
+function check_ir!(interp, @nospecialize(job::CompilerJob), errors::Vector{IRError}, imported::Set{String}, inst::LLVM.CallInst, calls::Vector{LLVM.CallInst}, mod::LLVM.Module)
     world = job.world
-    interp = GPUCompiler.get_interpreter(job)
     method_table = Core.Compiler.method_table(interp)
     bt = backtrace(inst)
     dest = called_operand(inst)