Make sure that we pick up new definition of typetree

Author: vchuravy

src/compiler.jl

@@ -14,7 +14,7 @@ import Enzyme:
     eltype,
     API,
     TypeTree,
-    typetree,
+    typetree, typetree_total,
     TypeTreeTable,
     only!,
     shift!,
@@ -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, 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
@@ -2313,7 +2313,7 @@ function enzyme!(
         else
             error("illegal annotation type $T")
         end
-        typeTree = typetree(source_typ, ctx, dl, seen)
+        typeTree = typetree_total(job, source_typ, ctx, dl, seen)
         if isboxed
             typeTree = copy(typeTree)
             merge!(typeTree, TypeTree(API.DT_Pointer, ctx))
@@ -2355,7 +2355,7 @@ function enzyme!(
             in(Any, actualRetType.parameters)
         TypeTree()
     else
-        typeTree = typetree(actualRetType, ctx, dl, seen)
+        typeTree = typetree_total(job, actualRetType, ctx, dl, seen)
         if !isa(actualRetType, Union) && GPUCompiler.deserves_retbox(actualRetType)
             typeTree = copy(typeTree)
             merge!(typeTree, TypeTree(API.DT_Pointer, ctx))
@@ -3703,7 +3703,7 @@ function lower_convention(
                     metadata(sretPtr)["enzyme_inactive"] = MDNode(LLVM.Metadata[])
                 end
 
-                typeTree = copy(typetree(actualRetType, ctx, dl, seen))
+                typeTree = copy(typetree(job, actualRetType, ctx, dl, seen))
                 merge!(typeTree, TypeTree(API.DT_Pointer, ctx))
                 only!(typeTree, -1)
                 metadata(sretPtr)["enzyme_type"] = to_md(typeTree, ctx)

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())