review 2

EnzymeAD · Dec 14, 2024 · 3671ddb · 3671ddb
1 parent 4d27e39
commit 3671ddb
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 0 deletions.
diff --git a/src/Precompile.jl b/src/Precompile.jl
@@ -12,4 +12,5 @@ using PrecompileTools: @setup_workload, @compile_workload
         XLA.free_client(cpu)
         deinitialize_dialect()
     end
+    XLA.cpuclientcount[] = 0
 end
diff --git a/src/XLA.jl b/src/XLA.jl
@@ -41,8 +41,11 @@ end
 
 SetLogLevel(x) = @ccall MLIR.API.mlir_c.SetLogLevel(x::Cint)::Cvoid
 
+global cpuclientcount = Ref(0)
 # TODO synchronization when async is not working because `future` in `ConcreteRArray` is always `nothing`
 function CPUClient(asynchronous=false, node_id=0, num_nodes=1)
+    @assert cpuclientcount[] == 0
+    cpuclientcount[] += 1
     f = Libdl.dlsym(Reactant_jll.libReactantExtra_handle, "MakeCPUClient")
     client = ccall(f, Ptr{Cvoid}, (UInt, Cint, Cint), asynchronous, node_id, num_nodes)
     #client = @ccall MLIR.API.mlir_c.MakeCPUClient(asynchronous::UInt8, node_id::Cint, num_nodes::Cint)::Ptr{Cvoid}