Performance: 1D Transport

The 1D plans of 1D Transport, on batches of 256 independent problems: the threads of the host, or the card, share them. Two kinds of measures only: diracs, and images of affine pixels ( Image( degree = 1 ), a piecewise linear density ). Three cases:

1. semi-discrete transports: diracs against an image;
2. transports between two images, two continuous densities;
3. against POT, on the transport between two images, discretized for POT.

The images sample sums of two thin gaussians ( standard deviation 0.05 ), shifted from one problem to the next. The diracs are drawn uniformly. Everything is in double; the inputs are already where the computation runs, and the time stops when the costs are on the host.

Semi-discrete transports

SdotPlan1d on a batch: each problem's diracs sorted, then one sweep along its target, an image of 1e4 pixels. The costs are exact.

sdot, EPYC 7763, 128 coressdot, H100

Total time for 256 transports against images of 1e4 pixels: a slope of one decade per decade is linear. Log-log; hover a point for its time.

256 transports of 1e6 diracs take 0.82 s on 128 cores, 0.09 s on an H100.

Between two images

ContinuousOtPlan1d between two images: no sort, the merge of the pixels of both, O( pixels ).

sdot, EPYC 7763, 128 coressdot, H100

Total time for 256 pairs of images: a slope of one decade per decade is linear. Log-log; hover a point for its time.

256 pairs of images of 1e6 pixels take 1.6 s on 128 cores, 1.7 s on the H100.

The cost is exact up to the quadrature of each merged piece: against an affine target, the transport map involves a square root, and so does its integral ( a closed form with an arcsinh, unstable when a slope vanishes ). Seven Gauss-Legendre points per piece agree with a 64-point reference to 7e-15 on images of 1024 pixels.

Over the threads

The problems are independent: the threads share them.

sdotsdot, H100 ( 132 SMs )ideal scaling

Cycles for 256 problems of 1e5 diracs ( or pixels ): the time times the clock, the median of the clocks measured on the machine at that thread count ( the cores' on the CPU, the SMs' on a card ): one run's sample is too noisy. Log-log; hover a point for its time and clock.

The images gain 13 times on 128 cores, the diracs 11 times; past 32 threads both gain little, bound by the memory.

Against POT

POT has no continuous transport. The fair comparison is therefore on two continuous densities, both discretized alike for POT: each image becomes m diracs, the centres of a regular grid of its support weighted by the masses of their cells, the same number on both sides, then ot.wasserstein_1d( p = 2 ), one column per pair. The more diracs, the closer POT gets to the continuous $W_2$, which sdot computes exactly: the chart puts the time of each method against the error on $W_2$. The two do not solve the same problem; the chart gives the time POT needs for a given accuracy, at as many threads: sdot on one thread against POT on numpy, on all the cores against POT on torch, and on the card against POT on its jax backend.

sdot, one threadsdot, all the cores ( the card )POT, numpy ( one thread )POT, torch on all the cores ( jax on the card )

Time against the relative error, 256 pairs of images of 1e4 pixels, in double; a curve joins the runs of a method in the order of their size. Log-log; hover a point for its time.

POT's error on $W_2$ decreases as 1 / m² ( two digits per decade of diracs ). sdot, exact, takes 0.047 s on 128 cores: POT on torch, on the same cores, takes as long for an error of about 1e-6, 0.31 s for 2.4e-8 and 25 s for 2.7e-12. On one thread, 0.32 s for sdot against 0.6 s for POT on numpy at 2.4e-8. On the H100, POT is faster down to about 1e-9 ( 7 ms for 2.4e-8, 34 ms for 2.6e-9, against 40 ms for sdot, of which 18 ms of kernels: the rest is the cost of the calls ), and slower beyond ( 1.6 s for 2.6e-11 ).

The tables

Semi-discrete transports

	sdot
	all the cores	H100
1e3	0.0362 s	0.0318 s
1e4	0.0474 s	0.0303 s
1e5	0.159 s	0.0370 s
1e6	0.818 s	0.0929 s

Seconds, 256 semi-discrete transports: n diracs each against an image of 1e4 affine pixels, the inputs already where the computation runs, until the costs are on the host.

Between two images

	sdot
	all the cores	H100
1e3	0.0213 s	0.0247 s
1e4	0.0501 s	0.00749 s
1e5	0.247 s	0.193 s
1e6	1.59 s	1.73 s

Seconds, 256 transports between two images of n affine pixels each.

Over the threads

	threads
	1 / 8 / 32 / 128 threads
diracs	1.33 s / 0.233 s / 0.127 s / 0.159 s
images	3.06 s / 0.516 s / 0.252 s / 0.247 s

Seconds, over the threads: 256 transports of 1e5 diracs against 1e4 pixels, of two images of 1e5 pixels.

Against POT

	POT, numpy, one thread		POT, torch, all the cores	POT, jax
	time	error	time	H100
m = 1e3	0.0489 s	2.4e-6	0.0561 s	0.00336 s
m = 1e4	0.603 s	2.4e-8	0.306 s	0.00749 s
m = 1e5	7.77 s	2.3e-10	2.74 s	0.317 s
m = 1e6	83.6 s	2.7e-12	24.8 s	–
m = 3e6	257 s	2.9e-13	67.2 s	–

Seconds and median relative error on W2 ( against sdot ), 256 pairs of images of 1e4 affine pixels, in double. POT: each image replaced by m diracs ( the same number on both sides, made beforehand ), one column per pair; on a card, with its jax backend.

	one thread	all the cores	sdot
	time	time	H100
sdot	0.320 s	0.0471 s	0.0398 s

sdot on the same 256 pairs, exact: seconds.

The numbers come from bench/bench_1d.py, jobs one_d* of bench/page.py. The machines and the protocol are on the overview.