Bad performance for `sort!(CoSorter(os._data, os.scalars); alg=ThreadsX.QuickSort)`

[xiangjianqian]

I noticed that the function sort!(CoSorter(os._data, os.scalars); alg=ThreadsX.QuickSort) consumes a significant amount of runtime. I've found a potential way to improve this, but I'm unsure whether my approach is correct. Could someone help me verify it? Below is a short demonstration for the modification of the code (inside the MPOGraph function):

    use_perm = true
    t_sort = @elapsed begin
        resize!(os._data, length(os))
        resize!(os.scalars, length(os))
        n = length(os)
        if use_perm
            @inline function term_key(t::NTuple{N,ITensorMPOConstruction.OpID{Ti}}, maxn::Int, maxid::Int)::UInt128 where {N,Ti}
                stride = UInt128(maxid + 1)
                base = UInt128((maxn + 1) * (maxid + 1))
                s = UInt128(0)
                @inbounds for i in 1:N
                    ti = t[i]
                    digit = UInt128(Int(ti.n)) * stride + UInt128(Int(ti.id))
                    s = s * base + digit
                end
                return s
            end
            keys = Vector{UInt128}(undef, n)
            @inbounds for i in eachindex(os._data)
                keys[i] = term_key(os._data[i],n,length(os.op_cache_vec[1]))
            end
            perm = sortperm(keys)
        else
            perm = sortperm(os._data; alg=Base.Sort.QuickSort)
            # sort!(ITensorMPOConstruction.CoSorter(os._data, os.scalars); alg=Base.Sort.QuickSort)
        end
    end

    local nnz = 0
    t_combine = @elapsed begin
        # permute!(os._data, perm)
        # permute!(os.scalars, perm)
        # # # println("sorted terms: ", os._data)
        # nnz = 0
        # for i in eachindex(os)
        #     if i < length(os) && os._data[i] == os._data[i + 1]
        #         os.scalars[i + 1] += os.scalars[i]
        #         os.scalars[i] = 0
        #     elseif abs(os.scalars[i]) > os.abs_tol
        #         nnz += 1
        #         os.scalars[nnz] = os.scalars[i]
        #         os._data[nnz] = os._data[i]
        #     end
        # end
        # Alternative combine approach that avoids branching and may be more efficient for large numbers of terms:
        # Iterate in sorted-permutation order and write compacted result in place.
        x= Int[]
        for i in 1:length(perm)
            if i < length(os) && os._data[perm[i]] == os._data[perm[i+1]]
                os.scalars[perm[i+1]] += os.scalars[perm[i]]
                os.scalars[perm[i]] = 0
            elseif abs(os.scalars[perm[i]]) > os.abs_tol
                push!(x, perm[i])
                nnz += 1
            end
        end
        os.scalars = os.scalars[x]
        os._data = os._data[x]
    end

[corbett5]

Hi!

Do you have an example where this sorting takes up a large fraction of the run time? For the cases I've looked at (mostly momentum space Fermi-Hubbard), sorting takes up about 2% of the total run time.

And how many threads are you running with? If you're using just a single thread, likely using alg=Base.Sort.QuickSort is better.

I originally played around with using a permutation based algorithm, but found this to be faster and use less memory than my implementation, which was less sophisticated than yours.

[corbett5]

I am not able to get your code to pass the tests. It appears to produce correct MPOs, but with a much larger bond dimension. For the fermi-hubbard-ks.jl example, the bond dimension has increased by almost 10x. This is likely because the terms are not actually sorted. This appears to be an issue with using the term_key function. With use_perm = true the following assert fails when added between the two statements

permute!(os._data, perm)
@assert issorted(os._data)

With use_perm = false this does not fail.

Regardless, I tried out the performance on the 6x6 system from fermi-hubbard-tc.jl. Below are data for just the time taken to perform the sorting. It's not a direct comparison, since your "new sorting" code also performs the de-duplication, but this only takes about 150ms in the existing code, so the comparison is similar. For parity, I added a Threads.@threads to the loop generating the keys and changed sortperm to use ThreadsX.QuickSort. Even with these additions your change is only modestly faster than the single-threaded sorting but still much slower than the threaded version.

Section          ncalls     time    %tot     avg     alloc    %tot      avg
───────────────────────────────────────────────────────────────────────────
sort! 1 thread        1    4.88s   ????%   4.88s   1.17GiB   ????%  1.17GiB
sort! 8 threads       1    746ms   ????%   746ms   1.17GiB   ????%  1.17GiB
sort! Base.QuickSort  1    4.71s   ????%   4.71s     0.00B   ????%    0.00B
new sorting 8 threads 1    3.20s   ????%   3.20s    592MiB   ????%   592MiB
───────────────────────────────────────────────────────────────────────────

But again, this is a single example, and perhaps you have a different example where you algorithm is faster, though please ensure it is correct by running the tests. This can be done by entering the pkg mode and pkg> test ITensorMPOConstruction, or more directly, but less comprehensively with using Revise, ITensorMPOConstruction; include("test/test-MPOConstruction.jl")

[xiangjianqian]

Sorry, I made some mistakes when constructing os, which contains a lot of duplicate terms. I've revised the code, and found my new_sorting is only slightly faster than simply using sort!(CoSorter(os._data, os.scalars)).

The term_key function was also incorrect. I've corrected it as well, and it now yields the right permutation and the correct MPO.

Thank you very much for your help.

---

Besides, I'm wondering if there is any way to cache certain calculations to speed up the code. I currently need to construct the MPO multiple times for the Hamiltonian:

$$ H = \sum_{ij\alpha\beta} T_{ij}^{\alpha\beta} c_{i\alpha}^\dagger c_{j\beta} + \sum_{ijkl\alpha\beta\gamma\delta} V_{ijkl}^{\alpha\beta\gamma\delta} c_{i\alpha}^\dagger c_{j\beta}^\dagger c_{k\gamma} c_{l\delta} $$

Each run only changes the entries of the T and V matrices, while the structure of the terms remains the same. Do you have any suggestions for caching intermediate results to avoid recomputing from scratch every time? Any advice would be greatly appreciated.

[corbett5]

Is there any structure to the Hamiltonian, or does the bond dimension of the MPO scale as $O(N^2)$ as it would for random coefficients? What is the system size? How many different MPOs do you want to construct?

With some work, assuming that the sparsity structure of T and V remain fixed, you could essentially write your own implementation of MPO_new that would just create the initial graph MPOGraph(os). Then it should be possible to copy and then modify the graph when you want to change the coefficients, and pass the modified graph to resume_MPO_construction. However, I don't think there is a good way to cache any part of the decomposition algorithm.

In principle, it would not be an insurmountable amount of work to add support for the symbolic bipartite graph algorithm from https://aip.scitation.org/doi/10.1063/5.0018149, in which case you could construct a symbolic MPO once, and then substitute in the coefficients as many times as you want. If you are interested in this I would check out Renormalizer (https://github.com/shuaigroup/Renormalizer) or Block2 (https://github.com/block-hczhai/block2-preview), both of which have a bipartite graph construction algorithm, but I am not sure if their implementations support symbolic construction or not.

[xiangjianqian]

I am considering an ab initio electronic Hamiltonian whose MPO bond dimension scales as $O(N^2)$ for around 100~200 orbitals. I need to construct the MPO about 10 times for different T and V. Perhaps I can manually track all the operators during the DMRG sweep process, rather than constructing MPO.

[corbett5]

I anticipate tracking the individual operators during DMRG, as in using dmrg(Hs::Vector{MPO}, psi0::MPS; kwargs...), will be intractable.

For these types of electronic structure Hamiltonians, the bipartite graph algorithm is the best choice anyways, since the MPOs cannot be compressed and the bipartite algorithm produces MPOs of greater sparsity (see https://itensor.github.io/ITensorMPOConstruction.jl/dev/on-sparsity/)

[corbett5]

Well, in terms of speed of constructing a single MPO I expect this library to be the best. I tried constructing a random electronic structure MPO (from https://github.com/ITensor/ITensorMPOConstruction.jl/blob/main/examples/electronic-structure.jl, and https://github.com/ITensor/ITensorMPOConstruction.jl/blob/main/docs/plot-generators/block2-plots.py) and for 40 electron sites (i.e. 80 spin-orbitals) ITensorMPOConstruction takes 33s while Block2 takes 120s, not counting the time to describe the Hamiltonian (i.e. OpIDSum).

[xiangjianqian]

Thank you very much for your help again.

[corbett5]

@xiangjianqian I'm working on adding a bipartite graph approach, it looks to be much faster for the electronic structure Hamiltonian. I will put up a PR within the next few days and you can try it out. I don't anticipate having symbolic construction working, nor do I know if it will let you scale up to 200 electron sites (is this what you meant by 200 orbitals?) but it's a step in the right direction.

[xiangjianqian]

Thanks again for your help! I actually need to go up to a 10×10 system, which corresponds to 100 electron sites (or 200 orbitals when spin degree of freedom is included). I'm really excited to see how much time can be saved under the new approach!

[corbett5]

Even without the symbolic construction I think that will be achievable.

[corbett5]

I've got the bipartite construction working. Constructing the MPO for 100 electron sites takes about 10 minutes on my laptop. It might be another week until I get the documentation in a stage to merge it, since I made some major changes. But if you want to try it out now you can check out this commit from the corbett/bipartite branch : cbd7262

Specifically you'll want to pass the following arguments (see examples/electronic-structure.jl):

MPO_new(os, sites; alg="VC", basis_op_cache_vec=os.op_cache_vec, splitblocks=true, check_for_errors=false, checkflux=false)

[corbett5]

Ok, it's been merged. I plan to add symbolic construction in another week or so, which should speed things even more, but hopefully it's already tractable.