[ITensors] Use task_local_storage for Index ID RNG (#1646)

Author: mtfishman

Project.toml

@@ -1,7 +1,7 @@
 name = "ITensors"
 uuid = "9136182c-28ba-11e9-034c-db9fb085ebd5"
 authors = ["Matthew Fishman <[email protected]>", "Miles Stoudenmire <[email protected]>"]
-version = "0.9.3"
+version = "0.9.4"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

src/ITensors.jl

@@ -155,8 +155,4 @@ include("deprecated.jl")
 include("argsdict/argsdict.jl")
 include("packagecompile/compile.jl")
 include("developer_tools.jl")
-
-function __init__()
-  return resize!(empty!(INDEX_ID_RNGs), Threads.nthreads()) # ensures that we didn't save a bad object
-end
 end

src/index.jl

@@ -1,26 +1,14 @@
 using NDTensors: NDTensors, sim
 using .QuantumNumbers: QuantumNumbers, Arrow, In, Neither, Out
+using Random: Xoshiro
 using .TagSets:
   TagSets, TagSet, @ts_str, addtags, commontags, hastags, removetags, replacetags
 
 #const IDType = UInt128
 const IDType = UInt64
 
-# Custom RNG for Index id
-# Vector of RNGs, one for each thread
-const INDEX_ID_RNGs = MersenneTwister[]
-@inline index_id_rng() = index_id_rng(Threads.threadid())
-@noinline function index_id_rng(tid::Int)
-  0 < tid <= length(INDEX_ID_RNGs) || _index_id_rng_length_assert()
-  if @inbounds isassigned(INDEX_ID_RNGs, tid)
-    @inbounds MT = INDEX_ID_RNGs[tid]
-  else
-    MT = MersenneTwister()
-    @inbounds INDEX_ID_RNGs[tid] = MT
-  end
-  return MT
-end
-@noinline _index_id_rng_length_assert() = @assert false "0 < tid <= length(INDEX_ID_RNGs)"
+const _INDEX_ID_RNG_KEY = :ITensors_index_id_rng_bLeTZeEsme4bG3vD
+index_id_rng() = get!(task_local_storage(), _INDEX_ID_RNG_KEY, Xoshiro())::Xoshiro
 
 """
 An `Index` represents a single tensor index with fixed dimension `dim`. Copies of an Index compare equal unless their

test/ext/ITensorsTensorOperationsExt/runtests.jl

@@ -150,10 +150,10 @@ end
     tmp * As[n]
     allocations_right_associative_pairwise += @allocated tmp = tmp * As[n]
   end
-  @test allocations_right_associative_pairwise ≈ allocations_right_associative_1 rtol = 0.1
-  @test allocations_right_associative_pairwise ≈ allocations_right_associative_2 rtol = 0.1
+  @test allocations_right_associative_pairwise ≈ allocations_right_associative_1 rtol = 0.2
+  @test allocations_right_associative_pairwise ≈ allocations_right_associative_2 rtol = 0.2
   @test allocations_right_associative_pairwise ≈ allocations_right_associative_3 rtol = 0.2
-  @test allocations_right_associative_pairwise ≈ allocations_right_associative_4 rtol = 0.1
+  @test allocations_right_associative_pairwise ≈ allocations_right_associative_4 rtol = 0.2
 
   @test allocations_right_associative_1 < allocations_left_associative
   @test allocations_right_associative_2 < allocations_left_associative