SDOT

Appearance

Performance: Semi-discrete OT ​

Geogram ( Bruno Lévy's exploragram, OptimalTransportMap2d / 3d ) solves the same problem with the damped Newton method of Kitagawa, Mérigot and Thibert ( KMT ). sdot is a damped Newton method too, but not KMT. In short:

  • a log phase: the first iterations write the Newton system and the merit of the step on log( mass / target ), so a nearly empty cell asks for a fraction of its gap only ( half the diagrams on the hard 2D cases ); once max |mass - target| / target <= 2, the solver switches to the plain difference, the true Newton, which converges quadratically;
  • the step is not halved from 1 at each iteration: the trials restart from the last accepted step, and the step comes from the limits at which the cells along the direction would be crushed;
  • continuation: when cells start empty, the target density is first smoothed, then sharpened stage by stage ( continuation ); the Lebesgue cases below do not need it;
  • seeds too close for the weights to separate are merged, then split again.

Here it gets the same clouds, the same box, the same stopping criterion and the same cores:

sdotgeogramsdot, RTX 2080 Ti ( 68 SMs )sdot, H100 ( 132 SMs )ideal scaling
2D uniform, 1e5100M300M1G3G10G30G1248163264128threads · SMssdot, 1 thread: 11G cycles, 3.3 s at 3.42 GHzsdot, 2 threads: 7.1G cycles, 2.08 s at 3.39 GHzsdot, 4 threads: 3.9G cycles, 1.18 s at 3.33 GHzsdot, 8 threads: 2.4G cycles, 741 ms at 3.27 GHzsdot, 16 threads: 2G cycles, 661 ms at 3.04 GHzsdot, 32 threads: 1.6G cycles, 530 ms at 2.99 GHzsdot, 64 threads: 1.5G cycles, 520 ms at 2.87 GHzsdot, 96 threads: 1.8G cycles, 680 ms at 2.71 GHzsdot, 128 threads: 1.8G cycles, 734 ms at 2.48 GHzsdot, RTX 2080 Ti ( 68 SMs ): 210M cycles, 113 ms at 1.86 GHzsdot, H100 ( 132 SMs ): 130M cycles, 66.9 ms at 1.98 GHzgeogram, 1 thread: 38G cycles, 11 s at 3.42 GHzgeogram, 2 threads: 32G cycles, 9.31 s at 3.39 GHzgeogram, 4 threads: 40G cycles, 12.1 s at 3.33 GHzgeogram, 8 threads: 38G cycles, 11.6 s at 3.27 GHzgeogram, 16 threads: 38G cycles, 12.6 s at 3.04 GHzgeogram, 32 threads: 36G cycles, 12.1 s at 2.99 GHzgeogram, 64 threads: 35G cycles, 12.2 s at 2.87 GHzgeogram, 96 threads: 33G cycles, 12 s at 2.71 GHzgeogram, 128 threads: 37G cycles, 15 s at 2.48 GHz
2D clustered ( 5 lines ), 1e51G10G100G1T1248163264128threads · SMssdot, 1 thread: 31G cycles, 9.2 s at 3.42 GHzsdot, 2 threads: 26G cycles, 7.68 s at 3.39 GHzsdot, 4 threads: 14G cycles, 4.29 s at 3.33 GHzsdot, 8 threads: 7.7G cycles, 2.35 s at 3.27 GHzsdot, 16 threads: 5.8G cycles, 1.89 s at 3.04 GHzsdot, 32 threads: 4.5G cycles, 1.5 s at 2.99 GHzsdot, 64 threads: 3.8G cycles, 1.33 s at 2.87 GHzsdot, 96 threads: 4.9G cycles, 1.79 s at 2.71 GHzsdot, 128 threads: 4.9G cycles, 1.99 s at 2.48 GHzsdot, RTX 2080 Ti ( 68 SMs ): 680M cycles, 365 ms at 1.86 GHzsdot, H100 ( 132 SMs ): 460M cycles, 234 ms at 1.98 GHzgeogram, 32 threads: 750G cycles, 252 s at 2.99 GHzgeogram, 96 threads: 700G cycles, 260 s at 2.71 GHz
3D uniform, 1e5300M1G3G10G30G1248163264128threads · SMssdot, 1 thread: 38G cycles, 11.2 s at 3.42 GHzsdot, 2 threads: 20G cycles, 5.83 s at 3.39 GHzsdot, 4 threads: 10G cycles, 3.11 s at 3.33 GHzsdot, 8 threads: 5.8G cycles, 1.78 s at 3.27 GHzsdot, 16 threads: 3.6G cycles, 1.18 s at 3.04 GHzsdot, 32 threads: 2.6G cycles, 863 ms at 2.99 GHzsdot, 64 threads: 2.1G cycles, 749 ms at 2.87 GHzsdot, 96 threads: 2.4G cycles, 885 ms at 2.71 GHzsdot, 128 threads: 2.6G cycles, 1.04 s at 2.48 GHzsdot, RTX 2080 Ti ( 68 SMs ): 660M cycles, 356 ms at 1.86 GHzsdot, H100 ( 132 SMs ): 220M cycles, 111 ms at 1.98 GHzgeogram, 1 thread: 54G cycles, 15.7 s at 3.42 GHzgeogram, 2 threads: 40G cycles, 11.9 s at 3.39 GHzgeogram, 4 threads: 34G cycles, 10.1 s at 3.33 GHzgeogram, 8 threads: 28G cycles, 8.52 s at 3.27 GHzgeogram, 16 threads: 24G cycles, 7.88 s at 3.04 GHzgeogram, 32 threads: 21G cycles, 6.9 s at 2.99 GHzgeogram, 64 threads: 19G cycles, 6.69 s at 2.87 GHzgeogram, 96 threads: 18G cycles, 6.57 s at 2.71 GHzgeogram, 128 threads: 24G cycles, 9.6 s at 2.48 GHz
3D clustered ( 4 planes ), 1e51G3G10G30G100G1248163264128threads · SMssdot, 1 thread: 140G cycles, 39.9 s at 3.42 GHzsdot, 2 threads: 70G cycles, 20.5 s at 3.39 GHzsdot, 4 threads: 35G cycles, 10.6 s at 3.33 GHzsdot, 8 threads: 19G cycles, 5.66 s at 3.27 GHzsdot, 16 threads: 11G cycles, 3.5 s at 3.04 GHzsdot, 32 threads: 6.3G cycles, 2.1 s at 2.99 GHzsdot, 64 threads: 4.6G cycles, 1.59 s at 2.87 GHzsdot, 96 threads: 4.6G cycles, 1.71 s at 2.71 GHzsdot, 128 threads: 5.2G cycles, 2.1 s at 2.48 GHzsdot, RTX 2080 Ti ( 68 SMs ): 2G cycles, 1.1 s at 1.86 GHzsdot, H100 ( 132 SMs ): 680M cycles, 343 ms at 1.98 GHzgeogram, 1 thread: 170G cycles, 50.2 s at 3.42 GHzgeogram, 2 threads: 130G cycles, 38.4 s at 3.39 GHzgeogram, 4 threads: 95G cycles, 28.4 s at 3.33 GHzgeogram, 8 threads: 75G cycles, 22.8 s at 3.27 GHzgeogram, 16 threads: 61G cycles, 19.9 s at 3.04 GHzgeogram, 32 threads: 55G cycles, 18.4 s at 2.99 GHzgeogram, 64 threads: 55G cycles, 19.2 s at 2.87 GHzgeogram, 96 threads: 59G cycles, 21.6 s at 2.71 GHzgeogram, 128 threads: 83G cycles, 33.3 s at 2.48 GHz

Cycles per solve: 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.

On all the cores, sdot solves 1e5 uniform diracs in 0.73 s in 2D and 1.0 s in 3D; its best times are on one socket, 64 threads: 0.52 s and 0.75 s ( see below ). On the clustered 2D cloud, geogram does not reach the tolerance ( see below ). An H100 solves the 2D case in 67 ms, a 2080 Ti in 0.11 s; the 3D case in 0.11 s and 0.36 s. The image and gaussian targets of the tables include a width continuation ( 16 to 22 smoothing stages ), which explains why they cost more than the Lebesgue cases.

How geogram was run:

  • geogram's own Newton loop ( optimize_full_Newton ) was kept as it is, with one change: it stops on sdot's max |mass - target| / target <= 1e-6 instead of its own criterion, a relative RMS error of 1 %. Its default use, a multilevel L-BFGS to 1 %, solves an easier problem and is not timed here.
  • Its linear solver is the default one: a conjugate gradient with a Jacobi preconditioner, tolerance 1e-3, at most 1000 iterations. It reaches the cap at every step: in 2D the linear solves take most of the time, in 3D the diagrams do. A tighter tolerance does not change the time or the number of steps.
  • The box is integrated as a mesh of simplices. Its cells match sdot's to 1e-10.
  • On the 2D clustered cloud, geogram's step control stalls: it reaches the tolerance in 2 runs of 9 ( one per thread count ), and after 100 steps the residual of the others is between 4e-6 and 5e-3. Lévy's multilevel scheme, run with a Newton solver, stops after the first level: interpolating the weights to the next level leaves cells empty, and exploragram's Newton stops on an empty cell. The equal clouds have coincident seeds, which geogram does not merge.

Scaling with the threads ​

A Newton solve alternates diagrams and linear solves. In 3D the diagrams take most of the time on one thread, and the solve gains 15× on 64 threads ( 25× on the clustered cloud ); what is left is in the linear solves and the tree. In 2D the cells are cheap and the linear solves take a larger share: on several threads they go through a multigrid, which is limited by the memory bandwidth, and the solve gains about 6× ( on one thread, a sparse Cholesky is used instead ). Past 64 threads the second socket adds memory traffic between the sockets and the solve gets slower again; on this node the time of the linear solves also varies by about 2× from one run to the next, whatever the number of threads. geogram gains about 2.4× in 3D and nothing in 2D.

Scaling with the number of diracs ​

sdot, EPYC 7763, 128 coressdot, RTX 2080 Tisdot, H100geogram, EPYC 7763, 128 cores
2D uniform1 µs10 µs100 µs1e31e41e51e61e7diracssdot, RTX 2080 Ti, n = 1e3: 26.8 µs at 1.86 GHzsdot, RTX 2080 Ti, n = 1e4: 5.52 µs at 1.86 GHzsdot, RTX 2080 Ti, n = 1e5: 1.13 µs at 1.86 GHzsdot, RTX 2080 Ti, n = 1e6: 705 ns at 1.86 GHzsdot, H100, n = 1e3: 22 µs at 1.98 GHzsdot, H100, n = 1e4: 4.26 µs at 1.98 GHzsdot, H100, n = 1e5: 669 ns at 1.98 GHzsdot, H100, n = 1e6: 215 ns at 1.98 GHzsdot, H100, n = 1e7: 245 ns at 1.98 GHzsdot, 128 threads, n = 1e3: 83.5 µs at 2.48 GHzsdot, 128 threads, n = 1e4: 21.8 µs at 2.48 GHzsdot, 128 threads, n = 1e5: 8.06 µs at 2.48 GHzsdot, 128 threads, n = 1e6: 6.46 µs at 2.48 GHzsdot, 128 threads, n = 1e7: 18.3 µs at 2.48 GHzgeogram, 128 threads, n = 1e3: 428 µs at 2.48 GHzgeogram, 128 threads, n = 1e4: 86.8 µs at 2.48 GHzgeogram, 128 threads, n = 1e5: 183 µs at 2.48 GHzgeogram, 128 threads, n = 1e6: 416 µs at 2.48 GHz
3D uniform1 µs10 µs100 µs1 ms1e31e41e51e61e7diracssdot, RTX 2080 Ti, n = 1e3: 17.5 µs at 1.86 GHzsdot, RTX 2080 Ti, n = 1e4: 5.25 µs at 1.86 GHzsdot, RTX 2080 Ti, n = 1e5: 3.56 µs at 1.86 GHzsdot, RTX 2080 Ti, n = 1e6: 3.59 µs at 1.86 GHzsdot, H100, n = 1e3: 18.7 µs at 1.98 GHzsdot, H100, n = 1e4: 3.18 µs at 1.98 GHzsdot, H100, n = 1e5: 1.11 µs at 1.98 GHzsdot, H100, n = 1e6: 799 ns at 1.98 GHzsdot, H100, n = 1e7: 1.64 µs at 1.98 GHzsdot, 128 threads, n = 1e3: 106 µs at 2.48 GHzsdot, 128 threads, n = 1e4: 23.7 µs at 2.48 GHzsdot, 128 threads, n = 1e5: 11.1 µs at 2.48 GHzsdot, 128 threads, n = 1e6: 6.05 µs at 2.48 GHzsdot, 128 threads, n = 1e7: 8.92 µs at 2.48 GHzgeogram, 128 threads, n = 1e3: 3.21 ms at 2.48 GHzgeogram, 128 threads, n = 1e4: 376 µs at 2.48 GHzgeogram, 128 threads, n = 1e5: 156 µs at 2.48 GHzgeogram, 128 threads, n = 1e6: 114 µs at 2.48 GHzgeogram, 128 threads, n = 1e7: 167 µs at 2.48 GHz

Time per dirac, a whole solve: flat is linear in the number of diracs. Log-log; hover a point for its time.

The number of Newton iterations barely depends on n. The total cost therefore follows the cost of one diagram plus one linear solve, which are both close to linear in n. On the CPU, 1e7 diracs take 183 s in 2D and 89 s in 3D on 128 cores; in 2D the linear solves are two thirds of it ( a uniform cloud of 1e7 seeds has close pairs, solved as groups, Iterative( duplicates = "split" ) ). A solve of 1e7 diracs takes 2.4 s in 2D and 16 s in 3D on an H100 ( an eager call, timed by its kernels: its wall time also counts python's loop ). It does not fit on a 2080 Ti ( 11 GB ).

geogram solves 1e7 diracs in 3D in 1674 s ( 6 steps ). In 2D its number of steps grows with n ( 7 at 1e5, 26 at 1e6 ), as its conjugate gradient stops at the cap before solving the system; the 1e7 solve was stopped after 5 hours without having converged, so the chart has no point there.

At 1e7 diracs in 3D, the float kernel can no longer resolve the cells ( each cell is about 4.6e-3 wide in a unit box ), so the solve switches to the double kernel. On cards with slow FP64 this switch costs half of the total time. GPU memory is about 900 bytes per dirac in 2D and 1300 in 3D at 1e6. A solve that does not fit raises a MemoryError before it starts.

Scaling with the dimension ​

sdot
A whole solve1 s10 s100 s1000 s2D3D4D5D6Ddimensionsdot, 2D: 214 ms at 2.48 GHzsdot, 3D: 292 ms at 2.48 GHzsdot, 4D: 1.04 s at 2.48 GHzsdot, 5D: 20.6 s at 2.48 GHzsdot, 6D: 754 s at 2.48 GHz
One diagram, per dirac1 µs10 µs100 µs1 ms10 ms2D3D4D5D6Ddimensionsdot, 2D: 180 ns at 2.48 GHzsdot, 3D: 1.14 µs at 2.48 GHzsdot, 4D: 13.3 µs at 2.48 GHzsdot, 5D: 336 µs at 2.48 GHzsdot, 6D: 15.1 ms at 2.48 GHz

Time, n = 1e4, all the cores: semi-log, the dimension along a linear axis. Hover a point for its time.

Newton is not affected by the dimension: the iteration count and the cost of the linear solves stay the same up to 6D ( the table ). The whole cost is in the cells. A Laguerre cell of a uniform cloud has a number of vertices that grows very fast with d ( about 6 in 2D, 25 in 3D, 120 in 4D, 670 in 5D, 3300 in 6D ), and the cell is cut, measured and triangulated vertex by vertex. At 1e4 diracs, a diagram costs about 6 times more in 3D than in 2D, then each added dimension multiplies it by 12 ( 4D ), 25 ( 5D ) and 45 ( 6D ).

In practice, 4D is comfortable up to a few 10⁵ diracs and 5D up to about 10⁵. 6D is limited to about 10⁴ diracs. Beyond that, the cost comes from the geometry of the cells, not from the solver.

The tables ​

Newton ​

CPUGPUspeed-up
128 threadsiterations / diagramsRTX 2080 Ti / H100RTX 2080 Ti / H100
2D uniform, 1e50.734 s6 / 70.113 s / 0.0669 s6.5× / 11×
2D uniform, 1e6––0.705 s / 0.215 s– / –
2D clustered ( 5 lines ), 1e51.99 s12 / 190.365 s / 0.234 s5.4× / 8.5×
2D 512² image, 2e47.25 s95 / 1170.906 s / 0.751 s8.0× / 9.7×
2D 4 gaussians ( σ = 0.02 ), 1e520.7 s133 / 2087.37 s / 2.93 s2.8× / 7.1×
3D uniform, 1e51.04 s5 / 60.356 s / 0.111 s2.9× / 9.4×
3D uniform, 1e6––3.59 s / 0.799 s– / –
3D clustered ( 4 planes ), 1e52.10 s9 / 171.10 s / 0.343 s1.9× / 6.1×

A whole solve, cold start, tree included, `max |mass - target| / target <= 1e-6`. CPU: all the physical cores, double. GPU: jitted, `mixed`; speed-up: against the CPU.

totalper dirac
RTX 2080 Ti / H100RTX 2080 Ti / H100
2D, 1e50.113 s / 0.0669 s1.1 µs / 0.67 µs
2D, 1e60.705 s / 0.215 s0.71 µs / 0.22 µs
2D, 1e7– / 2.45 s– / 0.24 µs
3D, 1e50.356 s / 0.111 s3.6 µs / 1.1 µs
3D, 1e63.59 s / 0.799 s3.6 µs / 0.80 µs
3D, 1e7– / 16.4 s– / 1.6 µs

A whole solve on the card, `mixed`, uniform diracs, Lebesgue. 1e7 ( `--full` ): an eager call, timed by its kernels.

sdot and geogram ​

sdotgeogram
1 / 32 / 128 threadsRTX 2080 Ti / H100iterations / diagrams1 / 32 / 128 threadsiterations / diagramslinear solves
2D uniform3.30 s / 0.530 s / 0.734 s0.113 s / 0.0669 s6 / 711.0 s / 12.1 s / 15.0 s6 / 813.7 s
2D clustered ( 5 lines )9.20 s / 1.50 s / 1.99 s0.365 s / 0.234 s12 / 19no conv. ( 271 s ) / 252 s / no conv. ( 298 s )100 / 719215 s
3D uniform11.2 s / 0.863 s / 1.04 s0.356 s / 0.111 s5 / 615.7 s / 6.90 s / 9.60 s5 / 83.00 s
3D clustered ( 4 planes )39.9 s / 2.10 s / 2.10 s1.10 s / 0.343 s9 / 1750.2 s / 18.4 s / 33.3 s13 / 2712.1 s

A whole semi-discrete solve, Lebesgue on the unit box, n = 1e5, cold start ( zero weights ), double ( GPU: jitted, `mixed` ), stopped on `max |mass - target| / target <= 1e-6` ( geogram's residual measured again by sdot ). Shaded: the fastest on the CPU, and on the GPU.

Dimension ​

totaliterations / diagramsone diagram, per dirac
2D0.214 s5 / 60.18 µs
3D0.292 s5 / 61.1 µs
4D1.04 s5 / 613 µs
5D20.6 s5 / 6336 µs
6D754 s4 / 515 ms

Lebesgue on the unit hypercube, n = 1e4, cold start, all the physical cores, double, one run. 6D: `--full`.

The machines and the protocol ( cold starts, repetitions, clocks ) are on the overview.

SDOT — H. Leclerc, Q. Mérigot, T. Gallouët · LMO / INRIA PARMA