Running decomp.py operations faster

[ACE07-Sev]

Greetings there dear Johnnie,

I hope this letter finds you well. I have a few questions regarding how I can improve runtime to be faster when running the below operations. I have done a quick profiler to show which ones are taking the most time and need to be done faster:

         463117405 function calls (447939147 primitive calls) in 4968.795 seconds

   Ordered by: internal time

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
    43050 2097.389    0.049 2097.773    0.049 /home/ace07/Documents/GitHub/QMPRS/.venv/lib/python3.11/site-packages/quimb/tensor/decomp.py:324(svd_truncated_numba)
    43036 1825.831    0.042 1826.118    0.042 /home/ace07/Documents/GitHub/QMPRS/.venv/lib/python3.11/site-packages/quimb/tensor/decomp.py:682(qr_stabilized_numba)
    71722  503.232    0.007  510.347    0.007 /home/ace07/Documents/GitHub/QMPRS/.venv/lib/python3.11/site-packages/numpy/core/numeric.py:932(tensordot)
   461354  157.020    0.000  679.937    0.001 /home/ace07/Documents/GitHub/QMPRS/.venv/lib/python3.11/site-packages/autoray/autoray.py:30(do)
8812213/32769   12.069    0.000   34.827    0.001 /usr/local/lib/python3.11/copy.py:128(deepcopy)
    86100    8.966    0.000 3973.492    0.046 /home/ace07/Documents/GitHub/QMPRS/.venv/lib/python3.11/site-packages/quimb/tensor/tensor_core.py:415(tensor_split)

As you can see, the majority of the runtime is dominated by decomp.py methods svd_truncated_numba and qr_stabilized_numba (taking 3922 seconds out of 4968 seconds), and with addition of do and tensordot it increases even more (taking 4582 seconds overall). I would like to decrease this significantly and understand there are options implemented in quimb for doing so.

Could I have some suggestions regarding this? Preferably those which would significantly reduce this runtime.

[ACE07-Sev]

I imagine seeing the code would help, if you like I can add you to my repo to see the code (it's two .py files so it wouldn't take much of your time).

[ACE07-Sev]

I was able to generate a call graph to see what methods/functions call the above methods/functions. I can see that .tensor_split() is the main source of the bottleneck as it calls .do() which itself calls .tensordot(), as well as calling both .svd_truncated_numpy() and qr_stabilized_numba().

Here is the call graph (kindly look at the yellow box):

[jcmgray]

As you have found, the runtime is being dominated by the low level QR and SVD routines, meaning that improvements will have to come from hardware or algorithm changes.

1. Probably the code is very amenable to GPU acceleration if you have access to one.
2. It looks like a 1D compress is the dominating sub routine. You could try using the tensor_network_1d_compress function which has different methods of performing the compression with different performance and accuracy tradeoffs.

[ACE07-Sev]

Thank you very much for the prompt response! Trying the tensor_network_1d_compress now.

[ACE07-Sev]

Apologies for the massive delay. I got sidetracked with a few other tasks for the package. Here's what I have for my compress wrapper.

def compress(
            self,
            max_bond_dimension: int | None = None,
            mode: Literal["left", "right"] | None = None
        ) -> None:
        """ SVD Compress the bond dimension of the MPS.

        ```
         a)│   │        b)│        │        c)│       │
         ━━●━━━●━━  ->  ━━>━━○━━○━━<━━  ->  ━━>━━━M━━━<━━
           │   │          │  ....  │          │       │
          <*> <*>          contract              <*>
          QR   LQ            -><-                SVD

            d)│            │        e)│   │
        ->  ━━>━━━ML──MR━━━<━━  ->  ━━●───●━━
              │....    ....│          │   │
            contract  contract          ^compressed bond
               -><-      -><-
        ```

        Parameters
        ----------
        `max_bond_dimension` : int
            The maximum bond dimension of the MPS.
        `mode` : Literal["left", "right"], optional
            The mode of compression.

        Raises
        ------
        ValueError
            - If `mode` is not "left", or "right".

        Usage
        -----
        >>> mps.compress(max_bond_dimension=16)
        >>> mps.compress(mode="left")
        >>> mps.compress(max_bond_dimension=32, mode="right")
        """
        if not (max_bond_dimension or mode):
            self.mps.compress()

        elif not mode:
            self.mps = tensor_network_1d_compress(tn=self.mps, max_bond=max_bond_dimension)

        else:
            if mode in ["left", "right", "flat"]:
                if not max_bond_dimension:
                    self.mps.compress(form=mode)
                else:
                    self.mps.compress(form=mode, max_bond=max_bond_dimension)
            else:
                raise ValueError(f"`mode` must be either 'left', or 'right'. Received {mode}.")

I checked out tensor_network_1d_compress() and from what I saw it does not support form, correct? I was wondering if you have any recommendations for improving how I do compression.