Repository navigation
Description
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.svdso 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_numpyon the result. Return tensors backed
by the array module the computation ran on, sodevice="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 = Trueinitially (preserving
today's behaviour) with a documented path toFalse, 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
- No
.get()/ D2H transfer inside the per-site loop. - Chained GPU
applycalls do not round-trip through host memory. - CPU results and numerics are unchanged.
- GPU tests still pass; before/after timings recorded in the PR description.
Activity
LucaLM02 commentedon Oct 3, 2026
hi @robertodr, i would choose this for the technical assignment