Implement multi-threading using OhMyThreads and make it differentiable (

#70)

Implementation of parallelization over unit cells, both in forwards and in backwards passes, using [OhMyThreads.jl](https://github.com/JuliaFolds2/OhMyThreads.jl).
Additional changes to CI to fix various things.

---
Co-authored-by: Lukas Devos <[email protected]>

.github/workflows/CI.yml

@@ -6,6 +6,8 @@ on:
       - 'main'
       - 'release-'
     tags: '*'
+    paths-ignore:
+      - 'docs/**'
   pull_request:
   workflow_dispatch:
 
@@ -15,47 +17,26 @@ concurrency:
   cancel-in-progress: ${{ startsWith(github.ref, 'refs/pull/') }}
 
 jobs:
-  test:
-    name: Julia ${{ matrix.version }} - ${{ matrix.os }} - ${{ matrix.group }}
-    runs-on: ${{ matrix.os }}
+  tests:
+    name: "Tests"
     strategy:
       fail-fast: false
       matrix:
         version:
           - 'lts'
-          - '1' # automatically expands to the latest stable 1.x release of Julia
-        os:
-          - ubuntu-latest
-          - macOS-latest
-          - windows-latest
+          - '1'
         group:
           - ctmrg
           - boundarymps
           - examples
-    steps:
-      - uses: actions/checkout@v4
-      - uses: julia-actions/setup-julia@v2
-        with:
-          version: ${{ matrix.version }}
-      - uses: actions/cache@v4
-        env:
-          cache-name: cache-artifacts
-        with:
-          path: ~/.julia/artifacts
-          key: ${{ runner.os }}-test-${{ env.cache-name }}-${{ hashFiles('**/Project.toml') }}
-          restore-keys: |
-            ${{ runner.os }}-test-${{ env.cache-name }}-
-            ${{ runner.os }}-test-
-            ${{ runner.os }}-
-      - uses: julia-actions/julia-buildpkg@latest
-      - uses: julia-actions/julia-runtest@latest
-        env:
-          JULIA_NUM_THREADS: 4
-          GROUP: ${{ matrix.group }}
-      - uses: julia-actions/julia-processcoverage@v1
-      - uses: codecov/codecov-action@v4
-        if: always()
-        env:
-          CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }}
-        with:
-          file: lcov.info
+          - utility
+        os:
+          - ubuntu-latest
+          - macOS-latest
+          - windows-latest
+    uses: "QuantumKitHub/.github/.github/workflows/tests.yml@main"
+    with:
+      group: "${{ matrix.group }}"
+      julia-version: "${{ matrix.version }}"
+      os: "${{ matrix.os }}"
+    secrets: inherit

.github/workflows/FormatCheck.yml

@@ -1,47 +1,14 @@
-name: FormatCheck
+name: "Format Check"
 
 on:
   push:
     branches:
       - 'main'
       - 'master'
-      - 'release-'
     tags: '*'
   pull_request:
 
 jobs:
-  build:
-    runs-on: ${{ matrix.os }}
-    strategy:
-      matrix:
-        version:
-          - '1' # automatically expands to the latest stable 1.x release of Julia
-        os:
-          - ubuntu-latest
-        arch:
-          - x64
-    steps:
-      - uses: julia-actions/setup-julia@latest
-        with:
-          version: ${{ matrix.version }}
-          arch: ${{ matrix.arch }}
-
-      - uses: actions/checkout@v4
-      - name: Install JuliaFormatter and format
-        # This will use the latest version by default but you can set the version like so:
-        #
-        # julia  -e 'using Pkg; Pkg.add(PackageSpec(name="JuliaFormatter", version="0.13.0"))'
-        run: |
-          julia  -e 'using Pkg; Pkg.add(PackageSpec(name="JuliaFormatter"))'
-          julia  -e 'using JuliaFormatter; format(".", verbose=true)'
-      - name: Format check
-        run: |
-          julia -e '
-          out = Cmd(`git diff --name-only`) |> read |> String
-          if out == ""
-              exit(0)
-          else
-              @error "Some files have not been formatted !!!"
-              write(stdout, out)
-              exit(1)
-          end'
+  format-check:
+    name: "Format Check"
+    uses: "QuantumKitHub/.github/.github/workflows/formatcheck.yml@main"

Project.toml

@@ -13,6 +13,7 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LoggingExtras = "e6f89c97-d47a-5376-807f-9c37f3926c36"
 MPSKit = "bb1c41ca-d63c-52ed-829e-0820dda26502"
 MPSKitModels = "ca635005-6f8c-4cd1-b51d-8491250ef2ab"
+OhMyThreads = "67456a42-1dca-4109-a031-0a68de7e3ad5"
 OptimKit = "77e91f04-9b3b-57a6-a776-40b61faaebe0"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"

src/PEPSKit.jl

@@ -11,8 +11,10 @@ using LoggingExtras
 using MPSKit: loginit!, logiter!, logfinish!, logcancel!
 using MPSKitModels
 using FiniteDifferences
+using OhMyThreads
 
 include("utility/util.jl")
+include("utility/diffable_threads.jl")
 include("utility/svd.jl")
 include("utility/rotations.jl")
 include("utility/diffset.jl")
@@ -51,22 +53,45 @@ include("utility/symmetrization.jl")
     module Defaults
         const ctmrg_maxiter = 100
         const ctmrg_miniter = 4
-        const ctmrg_tol = 1e-12
-        const fpgrad_maxiter = 100
+        const ctmrg_tol = 1e-8
+        const fpgrad_maxiter = 30
         const fpgrad_tol = 1e-6
+        const ctmrgscheme = :simultaneous
+        const reuse_env = true
+        const trscheme = FixedSpaceTruncation()
+        const fwd_alg = TensorKit.SVD()
+        const rrule_alg = Arnoldi(; tol=1e-2fpgrad_tol, krylovdim=48, verbosity=-1)
+        const svd_alg = SVDAdjoint(; fwd_alg, rrule_alg)
+        const optimizer = LBFGS(32; maxiter=100, gradtol=1e-4, verbosity=2)
+        const gradient_linsolver = KrylovKit.BiCGStab(;
+            maxiter=Defaults.fpgrad_maxiter, tol=Defaults.fpgrad_tol
+        )
+        const iterscheme = :fixed
+        const gradient_alg = LinSolver(; solver=gradient_linsolver, iterscheme)
+        const scheduler = Ref{Scheduler}(Threads.nthreads() == 1 ? SerialScheduler() : DynamicScheduler())
     end
 
 Module containing default values that represent typical algorithm parameters.
 
-- `ctmrg_maxiter = 100`: Maximal number of CTMRG iterations per run
-- `ctmrg_miniter = 4`: Minimal number of CTMRG carried out
-- `ctmrg_tol = 1e-12`: Tolerance checking singular value and norm convergence
-- `fpgrad_maxiter = 100`: Maximal number of iterations for computing the CTMRG fixed-point gradient
-- `fpgrad_tol = 1e-6`: Convergence tolerance for the fixed-point gradient iteration
-- `rrule_alg = Arnoldi(; tol=1e1ctmrg_tol, krylovdim=48, verbosity=-1)`: Default cotangent linear problem algorithm
+- `ctmrg_maxiter`: Maximal number of CTMRG iterations per run
+- `ctmrg_miniter`: Minimal number of CTMRG carried out
+- `ctmrg_tol`: Tolerance checking singular value and norm convergence
+- `fpgrad_maxiter`: Maximal number of iterations for computing the CTMRG fixed-point gradient
+- `fpgrad_tol`: Convergence tolerance for the fixed-point gradient iteration
+- `ctmrgscheme`: Scheme for growing, projecting and renormalizing CTMRG environments
+- `reuse_env`: If `true`, the current optimization step is initialized on the previous environment
+- `trscheme`: Truncation scheme for SVDs and other decompositions
+- `fwd_alg`: SVD algorithm that is used in the forward pass
+- `rrule_alg`: Reverse-rule for differentiating that SVD
+- `svd_alg`: Combination of `fwd_alg` and `rrule_alg`
+- `optimizer`: Optimization algorithm for PEPS ground-state optimization
+- `gradient_linsolver`: Default linear solver for the `LinSolver` gradient algorithm
+- `iterscheme`: Scheme for differentiating one CTMRG iteration
+- `gradient_alg`: Algorithm to compute the gradient fixed-point
+- `scheduler`: Multi-threading scheduler which can be accessed via `set_scheduler!`
 """
 module Defaults
-    using TensorKit, KrylovKit, OptimKit
+    using TensorKit, KrylovKit, OptimKit, OhMyThreads
     using PEPSKit: LinSolver, FixedSpaceTruncation, SVDAdjoint
     const ctmrg_maxiter = 100
     const ctmrg_miniter = 4
@@ -77,17 +102,63 @@ module Defaults
     const sparse = false
     const reuse_env = true
     const trscheme = FixedSpaceTruncation()
-    const iterscheme = :fixed
     const fwd_alg = TensorKit.SDD()
     const rrule_alg = Arnoldi(; tol=1e-2fpgrad_tol, krylovdim=48, verbosity=-1)
     const svd_alg = SVDAdjoint(; fwd_alg, rrule_alg)
     const optimizer = LBFGS(32; maxiter=100, gradtol=1e-4, verbosity=2)
     const gradient_linsolver = KrylovKit.BiCGStab(;
         maxiter=Defaults.fpgrad_maxiter, tol=Defaults.fpgrad_tol
     )
+    const iterscheme = :fixed
     const gradient_alg = LinSolver(; solver=gradient_linsolver, iterscheme)
+
+    # OhMyThreads scheduler defaults
+    const scheduler = Ref{Scheduler}()
+    """
+        set_scheduler!([scheduler]; kwargs...)
+
+    Set `OhMyThreads` multi-threading scheduler parameters.
+
+    The function either accepts a `scheduler` as an `OhMyThreads.Scheduler` or
+    as a symbol where the corresponding parameters are specificed as keyword arguments.
+    For instance, a static scheduler that uses four tasks with chunking enabled
+    can be set via
+    ```
+    set_scheduler!(StaticScheduler(; ntasks=4, chunking=true))
+    ```
+    or equivalently with 
+    ```
+    set_scheduler!(:static; ntasks=4, chunking=true)
+    ```
+    For a detailed description of all schedulers and their keyword arguments consult the
+    [`OhMyThreads` documentation](https://juliafolds2.github.io/OhMyThreads.jl/stable/refs/api/#Schedulers).
+
+    If no `scheduler` is passed and only kwargs are provided, the `DynamicScheduler`
+    constructor is used with the provided kwargs.
+
+    To reset the scheduler to its default value, one calls `set_scheduler!` without passing
+    arguments which then uses the default `DynamicScheduler()`. If the number of used threads is
+    just one it falls back to `SerialScheduler()`.
+    """
+    function set_scheduler!(sc=OhMyThreads.Implementation.NotGiven(); kwargs...)
+        if isempty(kwargs) && sc isa OhMyThreads.Implementation.NotGiven
+            scheduler[] = Threads.nthreads() == 1 ? SerialScheduler() : DynamicScheduler()
+        else
+            scheduler[] = OhMyThreads.Implementation._scheduler_from_userinput(
+                sc; kwargs...
+            )
+        end
+        return nothing
+    end
+    export set_scheduler!
+
+    function __init__()
+        return set_scheduler!()
+    end
 end
 
+using .Defaults: set_scheduler!
+export set_scheduler!
 export SVDAdjoint, IterSVD, NonTruncSVDAdjoint
 export FixedSpaceTruncation, ProjectorAlg, CTMRG, CTMRGEnv, correlation_length
 export LocalOperator