Make sure that we pick up new definition of typetree

Author: vchuravy

src/compiler.jl

@@ -1078,7 +1078,7 @@ function set_module_types!(mod::LLVM.Module, primalf::Union{Nothing, LLVM.Functi
 
                 byref = arg.cc
 
-                rest = copy(typetree(arg.typ, ctx, dl))
+                rest = copy(typetree_total(job, typetree, arg.typ, ctx, dl))
 
                 if byref == GPUCompiler.BITS_REF || byref == GPUCompiler.MUT_REF
                     # adjust first path to size of type since if arg.typ is {[-1]:Int}, that doesn't mean the broader
@@ -1103,7 +1103,7 @@ function set_module_types!(mod::LLVM.Module, primalf::Union{Nothing, LLVM.Functi
             if sret !== nothing
                 idx = 0
                 if !in(0, parmsRemoved)
-                    rest = typetree(sret, ctx, dl)
+                    rest = typetree_total(job, sret, ctx, dl)
                     push!(
                         parameter_attributes(f, idx + 1),
                         StringAttribute("enzyme_type", string(rest)),
@@ -1125,12 +1125,12 @@ function set_module_types!(mod::LLVM.Module, primalf::Union{Nothing, LLVM.Functi
                LLVM.return_type(LLVM.function_type(f)) != LLVM.VoidType()
                 @assert !retRemoved
                 rest = if llRT == Ptr{RT}
-                    typeTree = copy(typetree(RT, ctx, dl))
+                    typeTree = copy(typetree_total(job, RT, ctx, dl))
                     merge!(typeTree, TypeTree(API.DT_Pointer, ctx))
                     only!(typeTree, -1)
                     typeTree
                 else
-                    typetree(RT, ctx, dl)
+                    typetree_total(job, RT, ctx, dl)
                 end
                 push!(return_attributes(f), StringAttribute("enzyme_type", string(rest)))
             end

src/typetree.jl

@@ -195,6 +195,16 @@ end
 
 const TypeTreeTable = IdDict{Any,Union{Nothing,TypeTree}}
 
+"""
+    typetree_total(job, T, ctx, dl, seen=TypeTreeTable())
+
+A wrapper around `typetree` that ensures the call happens in the correct world for GPUCompiler.
+Useful when using typetree from a generated function since typetree is user-extendable.
+"""
+function typetree_total(@nospecialize(job::GPUCompiler.CompilerJob), @nospecialize(T), ctx, dl, seen=TypeTreeTable())
+    return Core._call_in_world_total(job.world, typetree, T, ctx, dl)
+end
+
 """
     function typetree(T, ctx, dl, seen=TypeTreeTable())