[NDTensors] [ITensors] Simplify block sparse multithreading implementation (#1049)

- Rewrite and reorganize sequential and multithreaded block sparse contraction, making use of `Folds.foreach` to clean up the code, make some parts more generic to sequential vs. threaded execution, as well as reorganizing to make it easier to try out different algorithms and backends.
- New benchmark for sequential and threaded block sparse contraction, based on the application of a 2-site effective Hamiltonian with tensor sparsity structure saved from a hybrid momentum space Hubbard model DMRG calculation with bond dimension 3000. The indices and fluxes of the tensors in the contraction are saved as HDF5 files so that random ITensors with the same block structure can be reproduced for benchmarking.
- Allow specifying a random number generator when constructing random ITensors and MPS by passing an RNG as the first input. Defaults to the default global RNG from Julia.
- Allow overloading `orthogonalize!` by dispatching on a projected MPO type in DMRG, which makes it easier to implement MPI parallelized sums of MPOs in ITensor/ITensorParallel.jl#15.

Author: mtfishman

.github/workflows/benchmark.yml

@@ -6,6 +6,8 @@ on:
 jobs:
   Benchmark:
     runs-on: ubuntu-latest
+    env:
+      JULIA_NUM_THREADS: 2
     steps:
       - uses: actions/checkout@v2
       - uses: julia-actions/setup-julia@latest

.github/workflows/format_check.yml

@@ -1,7 +1,7 @@
 name: Format check
 on:
   push:
-    branches: [master]
+    branches: [main]
     tags: [v*]
   pull_request:
 

ITensorGPU/README.md

@@ -1,6 +1,6 @@
 # ITensorGPU: Intelligent Tensors with GPU acceleration
 
-[![codecov](https://codecov.io/gh/ITensor/ITensorGPU.jl/branch/master/graph/badge.svg)](https://codecov.io/gh/ITensor/ITensorGPU.jl)
+[![codecov](https://codecov.io/gh/ITensor/ITensorGPU.jl/branch/main/graph/badge.svg)](https://codecov.io/gh/ITensor/ITensorGPU.jl)
 
 [![gitlab-ci](https://gitlab.com/JuliaGPU/ITensorGPU-jl/badges/master/pipeline.svg)](https://gitlab.com/JuliaGPU/ITensorGPU-jl/commits/master)
 

NDTensors/Project.toml

@@ -7,12 +7,15 @@ version = "0.1.46"
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
 Dictionaries = "85a47980-9c8c-11e8-2b9f-f7ca1fa99fb4"
+FLoops = "cc61a311-1640-44b5-9fba-1b764f453329"
+Folds = "41a02a25-b8f0-4f67-bc48-60067656b558"
 Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
 HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 SimpleTraits = "699a6c99-e7fa-54fc-8d76-47d257e15c1d"
+SplitApplyCombine = "03a91e81-4c3e-53e1-a0a4-9c0c8f19dd66"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Strided = "5e0ebb24-38b0-5f93-81fe-25c709ecae67"
 TimerOutputs = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f"
@@ -22,10 +25,13 @@ TupleTools = "9d95972d-f1c8-5527-a6e0-b4b365fa01f6"
 Adapt = "3.5"
 Compat = "2.1, 3, 4"
 Dictionaries = "0.3.5"
+FLoops = "0.2.1"
+Folds = "0.2.8"
 Functors = "0.2, 0.3, 0.4"
 HDF5 = "0.14, 0.15, 0.16"
 Requires = "1.1"
 SimpleTraits = "0.9.4"
+SplitApplyCombine = "1.2.2"
 StaticArrays = "0.12, 1.0"
 Strided = "0.3, 1"
 TimerOutputs = "0.5.5"

NDTensors/src/NDTensors.jl

@@ -4,13 +4,16 @@ using Adapt
 using Base.Threads
 using Compat
 using Dictionaries
+using FLoops
+using Folds
 using Random
 using LinearAlgebra
 using StaticArrays
 using Functors
 using HDF5
 using Requires
 using SimpleTraits
+using SplitApplyCombine
 using Strided
 using TimerOutputs
 using TupleTools
@@ -28,8 +31,9 @@ include("imports.jl")
 include("exports.jl")
 
 #####################################
-# DenseTensor and DiagTensor
+# General functionality
 #
+include("algorithm.jl")
 include("aliasstyle.jl")
 include("abstractarray/set_types.jl")
 include("abstractarray/to_shape.jl")
@@ -46,6 +50,10 @@ include("tensor/similar.jl")
 include("adapt.jl")
 include("generic_tensor_operations.jl")
 include("contraction_logic.jl")
+
+#####################################
+# DenseTensor and DiagTensor
+#
 include("dense/dense.jl")
 #include("dense/adapt.jl")
 include("fill.jl")
@@ -65,6 +73,12 @@ include("blocksparse/block.jl")
 include("blocksparse/blockoffsets.jl")
 include("blocksparse/blocksparse.jl")
 include("blocksparse/blocksparsetensor.jl")
+include("blocksparse/contract.jl")
+include("blocksparse/contract_utilities.jl")
+include("blocksparse/contract_generic.jl")
+include("blocksparse/contract_sequential.jl")
+include("blocksparse/contract_folds.jl")
+include("blocksparse/contract_threads.jl")
 include("blocksparse/diagblocksparse.jl")
 include("blocksparse/similar.jl")
 include("blocksparse/combiner.jl")
@@ -91,7 +105,7 @@ const timer = TimerOutput()
 # Optional block sparse multithreading
 #
 
-include("blas_get_num_threads.jl")
+blas_get_num_threads() = BLAS.get_num_threads()
 
 const _using_threaded_blocksparse = Ref(false)
 
@@ -134,9 +148,9 @@ function _enable_threaded_blocksparse()
         "WARNING: You are trying to enable block sparse multithreading, but you have started Julia with only a single thread. You can start Julia with `N` threads with `julia -t N`, and check the number of threads Julia can use with `Threads.nthreads()`. Your system has $(Sys.CPU_THREADS) threads available to use, which you can determine by running `Sys.CPU_THREADS`.\n",
       )
     end
-    if blas_get_num_threads() > 1 && Threads.nthreads() > 1
+    if BLAS.get_num_threads() > 1 && Threads.nthreads() > 1
       println(
-        "WARNING: You are enabling block sparse multithreading, but BLAS $(BLAS.vendor()) is currently set to use $(blas_get_num_threads()) threads. When using block sparse multithreading, we recommend setting BLAS to use only a single thread, otherwise you may see suboptimal performance. You can set it with `using LinearAlgebra; BLAS.set_num_threads(1)`.\n",
+        "WARNING: You are enabling block sparse multithreading, but your BLAS configuration $(BLAS.get_config()) is currently set to use $(BLAS.get_num_threads()) threads. When using block sparse multithreading, we recommend setting BLAS to use only a single thread, otherwise you may see suboptimal performance. You can set it with `using LinearAlgebra; BLAS.set_num_threads(1)`.\n",
       )
     end
     if Strided.get_num_threads() > 1

src/algorithm.jl NDTensors/src/algorithm.jlsrc/algorithm.jl renamed to NDTensors/src/algorithm.jl

@@ -12,9 +12,9 @@ use the same interface.
 struct Algorithm{Alg} end
 
 Algorithm(s) = Algorithm{Symbol(s)}()
-algorithm(::Algorithm{Alg}) where {Alg} = string(Alg)
+algorithm_string(::Algorithm{Alg}) where {Alg} = string(Alg)
 
-show(io::IO, alg::Algorithm) = print(io, "Algorithm type ", algorithm(alg))
+show(io::IO, alg::Algorithm) = print(io, "Algorithm type ", algorithm_string(alg))
 print(io::IO, ::Algorithm{Alg}) where {Alg} = print(io, Alg)
 
 """

NDTensors/src/blas_get_num_threads.jl

@@ -10,6 +10,12 @@ const BlockOffsets{N} = Dictionary{Block{N},Int}
 
 BlockOffset(block::Block{N}, offset::Int) where {N} = BlockOffset{N}(block, offset)
 
+ndims(::Blocks{N}) where {N} = N
+ndims(::BlockOffset{N}) where {N} = N
+ndims(::BlockOffsets{N}) where {N} = N
+
+blocktype(bofs::BlockOffsets) = keytype(bofs)
+
 nzblock(bof::BlockOffset) = first(bof)
 
 offset(bof::BlockOffset) = last(bof)

NDTensors/src/blocksparse/blockoffsets.jl

@@ -74,9 +74,15 @@ end
 #
 
 function randn(
-  ::Type{<:BlockSparse{ElT}}, blockoffsets::BlockOffsets, dim::Integer
+  StorageT::Type{<:BlockSparse{ElT}}, blockoffsets::BlockOffsets, dim::Integer
 ) where {ElT<:Number}
-  return BlockSparse(randn(ElT, dim), blockoffsets)
+  return randn(Random.default_rng(), StorageT, blockoffsets, dim)
+end
+
+function randn(
+  rng::AbstractRNG, ::Type{<:BlockSparse{ElT}}, blockoffsets::BlockOffsets, dim::Integer
+) where {ElT<:Number}
+  return BlockSparse(randn(rng, ElT, dim), blockoffsets)
 end
 
 #function BlockSparse{ElR}(data::VecT,offsets) where {ElR,VecT<:AbstractVector{ElT}} where {ElT}