Papers on CUDA ports often compare against CPU implementations that are not optimized for highest performance. Although CUDA hardware generally is faster than CPUs at the workloads described in these papers, the reported speedup is often higher than it would be if the CPU implementation had been optimized properly.
To gain some insight into the tradeoffs between CUDA and modern CPU optimizations, we optimized the N-body computation using two key strategies that are necessary for multicore CPUs to achieve peak performance:
SIMD (“single instruction multiple data”) instructions can perform multiple single-precision floating point operations in a single instruction, and
multithreading achieves near-linear speedups in the number of execution cores available in the CPU. Multicore CPUs have been widely available since 2006, servers with dozens of cores are now available, and N-Body computations are expected to scale almost linearly in the number of cores.
Since N-Body computations have such high computational density, we will not concern ourselves with affinity (for example, trying to use NUMA APIs to associate memory buffers with certain CPUs) – there is so much reuse in this computation that caches in the CPU keep external memory traffic to a trickle.
The Advanced Vector Extensions (AVX) were added to the x86
architecture in 2011, with Intel’s Sandy Bridge and AMD’s Bulldozer. AVX
widened the earlier SSE register file to a set of sixteen 256-bit YMM
registers that operate on eight packed 32-bit floating point values6; for example, the
VADDPS instruction performs eight floating point additions
in parallel, on corresponding packed floats in YMM registers.
When porting N-Body to the AVX instruction set, the AOS (array of
structures) memory layout that we have been using becomes problematic:
the instruction set operates on 8 bodies at a time, with their X, Y, Z
and Mass components separated and packed into registers. Rather than
shuffle the data into that form when computing the body-body
interactions, we rearrange the memory layout as structure of arrays:
instead of a single array of float4 (each element being the
X, Y, Z and Mass values for a given body), we use four arrays of float,
with an array of X values, an array of Y values, and so on. With the
data rearranged in this way, eight bodies’ descriptions can be loaded
into YMM registers with a single instruction each; the difference
vectors between eight bodies’ positions can be computed with just 3
machine instructions; and so on.
To simplify coding, Intel has worked with compiler vendors to add
cross-platform support for these instructions: a special data type
__m256 corresponds to the 256-bit register and operand
size, and intrinsic functions such as _mm256_sub_ps() that
correspond to the VSUBPS instruction.
For purposes of our N-Body implementation, we also need a
full-precision reciprocal square root implementation; the AVX
instruction set has an instruction VRSQRTPS that computes
an approximation of the reciprocal square root, but its 12-bit estimate
must be refined by a Newton-Raphson iteration to achieve full float
precision7:
\[x_{0} = VRSQRTPS(a) \]
\[x_{1} = \frac{x_{0}\left( 3 - a{x_{0}}^{2} \right)}{2}\]
Listing 14-8 gives an AVX implementation of the body-body computation
that takes the two bodies’ descriptions as __m256
variables, computes the eight body-body forces in parallel, and passes
back the 3 resulting force vectors. Listing 14-8 is functionally
equivalent to Listings 14-1 and 14-2, though markedly less readable.
Note that the x0, y0, and z0
variables contain descriptions of the same body, replicated across the
__m256 variable eight times.
static inline __m256rcp_sqrt_nr_ps(const __m256 x){ const __m256 nr = _mm256_rsqrt_ps(x), muls = _mm256_mul_ps(_mm256_mul_ps(nr, nr), x), beta = _mm256_mul_ps(_mm256_set1_ps(0.5f), nr), gamma = _mm256_sub_ps(_mm256_set1_ps(3.0f), muls); return _mm256_mul_ps(beta, gamma);} //// Sum the eight lanes of a YMM register: fold the high 128 bits onto the low// 128 and reduce those four, leaving the total in the low lane.//static inline __m128horizontal_sum_ps( const __m256 x ){ __m128 s = _mm_add_ps( _mm256_castps256_ps128(x), _mm256_extractf128_ps(x, 1) ); s = _mm_add_ps( s, _mm_movehl_ps( s, s ) ); return _mm_add_ss( s, _mm_shuffle_ps( s, s, 1 ) );} inline voidbodyBodyInteraction( __m256& fx, __m256& fy, __m256& fz, const __m256& x0, const __m256& y0, const __m256& z0, const __m256& x1, const __m256& y1, const __m256& z1, const __m256& mass1, const __m256& softeningSquared ){ // r_01 [3 FLOPS] __m256 dx = _mm256_sub_ps( x1, x0 ); __m256 dy = _mm256_sub_ps( y1, y0 ); __m256 dz = _mm256_sub_ps( z1, z0 ); // d^2 + e^2 [6 FLOPS] __m256 distSq = _mm256_add_ps( _mm256_add_ps( _mm256_mul_ps( dx, dx ), _mm256_mul_ps( dy, dy ) ), _mm256_mul_ps( dz, dz ) ); distSq = _mm256_add_ps( distSq, softeningSquared ); // invDistCube = 1/distSq^(3/2) [4 FLOPS (2 mul, 1 sqrt, 1 inv)] __m256 invDist = rcp_sqrt_nr_ps( distSq ); __m256 invDistCube = _mm256_mul_ps( invDist, _mm256_mul_ps( invDist, invDist ) ); // s = m_j * invDistCube [1 FLOP] __m256 s = _mm256_mul_ps( mass1, invDistCube ); // (m_1 * r_01) / (d^2 + e^2)^(3/2) [6 FLOPS] fx = _mm256_add_ps( fx, _mm256_mul_ps( dx, s ) ); fy = _mm256_add_ps( fy, _mm256_mul_ps( dy, s ) ); fz = _mm256_add_ps( fz, _mm256_mul_ps( dz, s ) );}To take advantage of multiple cores, we must spawn multiple threads and have each thread perform a subset of the computation.
The same strategy is used for multiple CPU cores as for multiple GPUs8: just evenly divide the output rows among threads (one per CPU core) and, for each timestep, have the “parent” thread spawn the worker threads to perform their work and then wait for them to finish.
The multithreaded samples use C++ std::thread (Section
A.2). The number of CPU cores comes from
std::thread::hardware_concurrency(); the parent thread
fills a std::vector<std::thread>, one worker per
core, then joins them to wait until the workers are finished.
Listing 14-9 gives the code that dispatches the N-Body calculation to
worker CPU threads. Because the AVX kernels consume eight bodies at a
time, it first calls requireBodyCountForAVX() to reject a
body count that is not a multiple of eight—refusing the computation
rather than silently dropping the leftover bodies. The
avxDelegation structures communicate the work to each
thread; each std::thread is constructed with the worker
function and a pointer to that thread’s avxDelegation
structure.
floatComputeGravitation_SIMD_threaded( float *force[3], float *pos[4], float *mass, float softeningSquared, size_t N){ // AVX processes eight bodies at a time; refuse a body count that is not // a multiple of eight rather than silently dropping the remainder. requireBodyCountForAVX( N ); std::chrono::steady_clock::time_point start, end; start = std::chrono::steady_clock::now(); { avxDelegation *pavx = new avxDelegation[g_numCPUCores]; std::vector<std::thread> threads; for ( size_t i = 0; i < g_numCPUCores; i++ ) { pavx[i].hostPosSOA[0] = pos[0]; pavx[i].hostPosSOA[1] = pos[1]; pavx[i].hostPosSOA[2] = pos[2]; pavx[i].hostMassSOA = mass; pavx[i].hostForceSOA[0] = force[0]; pavx[i].hostForceSOA[1] = force[1]; pavx[i].hostForceSOA[2] = force[2]; pavx[i].softeningSquared = softeningSquared; // Divide the bodies as evenly as possible among the cores. This // split is independent of the AVX width: a core may be handed any // number of bodies, so an uneven division just gives some cores one // more body than others rather than dropping the remainder. size_t begin = N * i / g_numCPUCores; size_t end = N * (i + 1) / g_numCPUCores; pavx[i].i = begin; pavx[i].n = end - begin; pavx[i].N = N; threads.emplace_back( avxWorkerThread, &pavx[i] ); } for ( auto &t : threads ) t.join(); delete[] pavx; } end = std::chrono::steady_clock::now(); return (float) std::chrono::duration<double>(end - start).count() * 1000.0f;}Finally, Listing 14-10 gives the avxDelegation structure
and the delegation function invoked by
ComputeGravitation_SIMD_threaded in Listing 14-9. It
performs the body-body calculations eight at a time, accumulating eight
partial sums that are added together with horizontal_sum_ps() before
storing the final output forces. This function, along with all the
functions that it calls, uses the SOA memory layout for all inputs and
outputs.
struct avxDelegation { size_t i; // base offset for this thread to process size_t n; // size of this thread's problem size_t N; // total number of bodies float *hostPosSOA[3]; float *hostMassSOA; float *hostForceSOA[3]; float softeningSquared; }; static voidavxWorkerThread( void *_p ){ avxDelegation *p = (avxDelegation *) _p; const __m256 softening = _mm256_set1_ps( p->softeningSquared ); for (int k = 0; k < p->n; k++) { int i = p->i + k; __m256 ax = _mm256_setzero_ps(); __m256 ay = _mm256_setzero_ps(); __m256 az = _mm256_setzero_ps(); __m256 x0 = _mm256_set1_ps( p->hostPosSOA[0][i] ); __m256 y0 = _mm256_set1_ps( p->hostPosSOA[1][i] ); __m256 z0 = _mm256_set1_ps( p->hostPosSOA[2][i] ); for ( int j = 0; j < p->N/8; j++ ) { bodyBodyInteraction( ax, ay, az, x0, y0, z0, _mm256_loadu_ps( p->hostPosSOA[0] + 8*j ), _mm256_loadu_ps( p->hostPosSOA[1] + 8*j ), _mm256_loadu_ps( p->hostPosSOA[2] + 8*j ), _mm256_loadu_ps( p->hostMassSOA + 8*j ), softening ); } // Sum the eight partial forces accumulated in each YMM register _mm_store_ss( &p->hostForceSOA[0][i], horizontal_sum_ps( ax ) ); _mm_store_ss( &p->hostForceSOA[1][i], horizontal_sum_ps( ay ) ); _mm_store_ss( &p->hostForceSOA[2][i], horizontal_sum_ps( az ) ); }}avxWorkerThread (source on GitHub)AVX can also treat the YMM registers as packed integers (with the AVX2 extensions) or as four packed double-precision floating point values, but we do not use any of those features. The later AVX-512 extensions double the width again, to 512-bit ZMM registers holding sixteen floats; widening this code to AVX-512 would be straightforward on the CPUs that support it.↩︎
This code is not present in the compiler’s intrinsics support, and is surprisingly difficult to find. Our implementation adapts the widely circulated SSE approximation code to 256-bit registers.↩︎
In fact, both the multithreaded CPU implementation and
the multi-GPU support of Chapter 9 spawn their worker threads the same
way, with the C++ standard library’s std::thread.↩︎