Skip to content

GPU path round-trips to host on every call and every site #14

Description

@Panadestein

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

  • No .get() / D2H transfer inside the per-site loop.
  • Chained GPU apply calls 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

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions