Merge pull request #469 from maleadt/tb/tsmod_ret

Support returning values from thread-safe module callbacks.

Author: maleadt

.github/workflows/enzyme.yml

@@ -53,9 +53,11 @@ jobs:
           using Pkg
           try
             # force it to use this PR's version of the package
-            Pkg.develop(PackageSpec(path="."))  # resolver may fail with main deps
+            ## resolver may fail with main deps
+            Pkg.develop(PackageSpec(path="."))
             Pkg.update()
-            Pkg.test(coverage=true)  # resolver may fail with test time deps
+            ## # resolver may fail with test time deps
+            Pkg.test(; coverage=true, julia_args=`--depwarn=no`)
           catch err
             err isa Pkg.Resolve.ResolverError || rethrow()
             # If we can't resolve that means this is incompatible by SemVer and this is fine

src/executionengine/ts_module.jl

@@ -58,7 +58,10 @@ function dispose(mod::ThreadSafeModule)
 end
 
 mutable struct ThreadSafeModuleCallback
+    ret::Ref{Any}
     callback
+
+    ThreadSafeModuleCallback(callback) = new(Ref{Any}(), callback)
 end
 
 function tsm_callback(data::Ptr{Cvoid}, ref::API.LLVMModuleRef)
@@ -67,7 +70,7 @@ function tsm_callback(data::Ptr{Cvoid}, ref::API.LLVMModuleRef)
     ctx = context(mod)
     activate(ctx)
     try
-        cb.callback(Module(ref))
+        cb.ret = cb.callback(Module(ref))
     catch err
         msg = sprint(Base.display_error, err, Base.catch_backtrace())
         return API.LLVMCreateStringError(msg)
@@ -91,5 +94,5 @@ function (mod::ThreadSafeModule)(f)
             @cfunction(tsm_callback, API.LLVMErrorRef, (Ptr{Cvoid}, API.LLVMModuleRef)),
             Base.pointer_from_objref(cb))
     end
-    return mod
+    cb.ret[]
 end

test/orc_tests.jl

@@ -20,11 +20,9 @@ end
         @test_throws LLVMException ts_mod() do mod
             error("Error")
         end
-        run = Ref{Bool}(false)
-        ts_mod() do mod
-            run[] = true
+        @test ts_mod() do mod
+            true
         end
-        @test run[]
     end
 end