Problem
The GPU path never keeps data on the GPU. There are two separate host
round-trips, and both are on the hot path.
1. Every result is copied back to the host
All four implementations end with:
return qtn.MatrixProductOperator([to_numpy(t) for t in eta])
So a GPU call is always H2D → compute → D2H. In an iterative workload — the
main use case for this library, where the output of one apply is the input to
the next — every single iteration pays a full device↔host round trip for the
entire tensor network, even though the data never needed to leave the device.
quimb handles CuPy-backed tensors natively via autoray, so this copy is not
required.
2. The truncation SVD runs on the CPU, once per site
src/src_method/utils/linalg.py:
Q, R = xp.linalg.qr(matrix.T if transpose else matrix)
R_np = R.get() if hasattr(R, "get") else np.asarray(R) # <-- blocking D2H
U, S, _ = np.linalg.svd(R_np.T if transpose else R_np, full_matrices=False)
R.get() is a blocking transfer: it forces a full device synchronisation at
every site, serialising the GPU pipeline n_sites times per sweep, and then runs
the decomposition on the CPU while the GPU idles. For large chi_out the SVD
itself is substantial work (measured ~47 ms for a 256×256 complex128 R on CPU)
being done on the slower device.
Proposed fix
- Run the SVD with
xp.linalg.svd so it stays on whichever device the data is
already on. Only the scalar rank needs to reach the host, and the rank can be
computed on device (int((cutoff * S[0] <= S).sum())) so at most one small
sync per site remains — ideally batched or avoided.
- Stop unconditionally calling
to_numpy on the result. Return tensors backed
by the array module the computation ran on, so device="gpu" results stay on
the GPU and can be chained without copies.
- To keep this from breaking existing callers, add an explicit knob controlling
where the result lands, e.g. to_host: bool = True initially (preserving
today's behaviour) with a documented path to False, or a
device="gpu" result that the caller converts explicitly.
- Benchmark before/after on the GPU runner for a representative chained
workload and put the numbers in the PR.
Acceptance criteria
Problem
The GPU path never keeps data on the GPU. There are two separate host
round-trips, and both are on the hot path.
1. Every result is copied back to the host
All four implementations end with:
So a GPU call is always
H2D → compute → D2H. In an iterative workload — themain use case for this library, where the output of one
applyis the input tothe next — every single iteration pays a full device↔host round trip for the
entire tensor network, even though the data never needed to leave the device.
quimbhandles CuPy-backed tensors natively viaautoray, so this copy is notrequired.
2. The truncation SVD runs on the CPU, once per site
src/src_method/utils/linalg.py:R.get()is a blocking transfer: it forces a full device synchronisation atevery site, serialising the GPU pipeline n_sites times per sweep, and then runs
the decomposition on the CPU while the GPU idles. For large
chi_outthe SVDitself is substantial work (measured ~47 ms for a 256×256 complex128 R on CPU)
being done on the slower device.
Proposed fix
xp.linalg.svdso it stays on whichever device the data isalready on. Only the scalar rank needs to reach the host, and the rank can be
computed on device (
int((cutoff * S[0] <= S).sum())) so at most one smallsync per site remains — ideally batched or avoided.
to_numpyon the result. Return tensors backedby the array module the computation ran on, so
device="gpu"results stay onthe GPU and can be chained without copies.
where the result lands, e.g.
to_host: bool = Trueinitially (preservingtoday's behaviour) with a documented path to
False, or adevice="gpu"result that the caller converts explicitly.workload and put the numbers in the PR.
Acceptance criteria
.get()/ D2H transfer inside the per-site loop.applycalls do not round-trip through host memory.