diff --git a/src/Precompile.jl b/src/Precompile.jl index acd8d4a8e..c3c86f55d 100644 --- 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 index 4000a17c1..18c5760b8 100644 --- 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}