Zero-copy interoperability between PETSc device sparse operators and PyTorch autograd: paper, code, and measurement data
Abstract
Everything behind the paper Zero-copy interoperability between PETSc device sparse operators and PyTorch autograd: the bridge library, its correctness gates, the benchmark scripts, the measurement data behind every table and figure, and the paper source and PDF. The bridge exposes a PETSc Mat living in GPU memory as a differentiable operator inside a PyTorch graph, with no host round trip on the hot path. It covers the CUDA scalar and native block-sparse classes, trainable nonzero values over a fixed sparsity pattern, batched right-hand sides, and distributed (MPI) operators. Contents. code/ — the bridge library, its correctness gates (including committed real-elasticity gate matrices) and the benchmarks that produced the data, BSD-2-Clause. data/ — the measurement CSVs behind every table and figure, plus the memory-probe log behind the host-staging claim. paper/ — LaTeX source, figures, and the built PDF. scripts/emit-paper-tables.py — regenerates every table body in the paper from data/ and, with --check, verifies every number in the paper's tables against the data; it reports 208 numbers across six tables, all reproduced. code/petsc-blocked-amg.patch — the PETSc side. Provenance. The measurements ran on NVIDIA A100s: the single-GPU sections on one A100-SXM4-40GB, the distributed section on Perlmutter A100-SXM4-80GB nodes. The elasticity operators are produced by stock PETSc snes/tutorials/ex56 (release 3.25.4 base); the exact command, including the RCM ordering, is in the paper's Availability section. code/petsc-blocked-amg.patch applies to upstream petsc/petsc commit d2a0018167b and reconstructs the tree the measurements ran on, carrying the BAIJCUDA block classes, the device dense-operand block products, and the petsc4py binding extensions. The annotated tag petsc-pytorch-bridge-v1 in the source repository marks the state the numbers come from; this archive is tagged petsc-pytorch-bridge-v2. Licensing. code/ is BSD-2-Clause (see code/LICENSE). The paper source, figures, and measurement data under paper/ and data/ are CC-BY-4.0.
// Source
Authors: Mark F. Adams
Institutions: Lawrence Berkeley National Laboratory