[ITensors][NDTensors] An attempt to address contract compile time

[kmp5VT]

Description

The lowest contract function in NDTensors is explicitly dependent on the number of tensor indices for tensors A, B, and C, for a generic tensor contraction C = A * B. Therefore when we change the number of modes for A, B or C, we need to recompile the full stack of contract functions.

Fixes #(issue)

Because most of the contractions fall into a matrix contraction (gemm), we could refactor the original ITensors into ITensors of order-2 and then follow the standard contract approach. This would mean that we only have to recompile the top-most function which converts order-N tensors into order-2 tensors.

Minimal demonstration of previous behavior

i = Index(5)
j = Index(10)
l = Index(3)
k = Index(20)
m = Index(30)

A = randomITensor(i,j)
B = randomITensor(j,k,l)

@time A * B  # Need to compile all of the contract functions
  3.910665 seconds (19.21 M allocations: 1.084 GiB, 8.31% gc time, 99.93% compilation time)

@time B * A  # This would also need to recompile all of the contract functions 
  0.866481 seconds (4.16 M allocations: 212.052 MiB, 3.91% gc time, 99.74% compilation time)

A = randomITensor(i,j,m)
@time A * B
1.517049 seconds (6.71 M allocations: 360.763 MiB, 4.15% gc time, 99.81% compilation time)

Minimal demonstration of new behavior

i = Index(5)
j = Index(10)
l = Index(3)
k = Index(20)
m = Index(30)

A = randomITensor(i,j)
B = randomITensor(j,k,l)

@time contract(A, B; matricize=true)  # Need to compile all of the contract functions
3.697179 seconds (17.44 M allocations: 1010.773 MiB, 7.74% gc time, 99.88% compilation time)
@time contract(B, A; matricize=true) # This only need to recompile the top most function
  0.238060 seconds (590.71 k allocations: 35.695 MiB, 2.87% gc time, 99.16% compilation time)

How Has This Been Tested?

Tests will be added in the contract section that use the matricize portion of contract and give the exact same result as the non-matricized portion

Checklist:

• My code follows the style guidelines of this project. Please run using JuliaFormatter; format(".") in the base directory of the repository (~/.julia/dev/ITensors) to format your code according to our style guidelines.
• I have performed a self-review of my own code.
• I have commented my code, particularly in hard-to-understand areas.
• I have added tests that verify the behavior of the changes I made.
• I have made corresponding changes to the documentation.
• My changes generate no new warnings.
• Any dependent changes have been merged and published in downstream modules.

[kmp5VT]

@emstoudenmire I looked into the error and its one I helped Ryan with last week. A recent update to the Julia version 1.9 causes an issue in the Zygote package. As a result one of our unit tests fails. So the failing test is not related to the work in this PR. (Ryan made a PR with only a change to documentation and was worried about why the unittest were failing)

[emstoudenmire]

To keep you in the loop @mtfishman , Karl and I discussed this in detail today, and agreed it's more of an experimental PR to see what kinds of compilation speedups might be possible to get and how to get them. Karl is going to add some more examples that put even more pressure on the compiler to see how big the speedups can be. A final PR implementing these ideas would need to be designed differently (some parts of the new code here being absorbed into functions above, and other parts being moved to NDTensors.jl).

[mtfishman]

Sounds good, glad this is getting explored and it's good timing with the release of Julia 1.9 which should make precompilation more effective once we refactor the contraction code properly.

This was a branch I was working on which refactored the contraction code logic to help with precompilation: https://github.com/ITensor/ITensors.jl/tree/ITensors_precompile. It's similar in spirit, where I was creating function barriers to split apart the logic of converting the tensors to matrices and then performing the contraction, but at a lower level in the NDTensors code. The goal was to make it so that less of the code had to get recompiled for tensors with new tensor orders that get input into the contract function, and therefore make precompilation more effective. I think there is a way to do that and make the contract code logic simpler, I would even hope to get rid of, or at least simplify, the ContractionProperties struct.

In our next meeting, we can try to go through the logic of the contract function in NDTensors and go through the changes I had been making in the ITensors_precompile branch.

[mtfishman]

I'll close this since it will be easier to refactor to improve compile times with changes in #1206, #1222, etc.