# AVX assembly for fast atan2 approximation

I'm writing a fast atan2 approximation, and would like some feedback on my assembly in particular. I know one of the first things is that people will question why I'm using inline assembly instead of instrinsics. This comes down to the fact my code is generally my deliverable, so the less I impose on people's build process the better. With intrinsics, you have to compile them separately with the appropriate compiler flag to even get access to them, and then link in the result. I like to write header libraries that can just be included and used. Inline assembly satisfies this for me, and these are generally simple fixed functions, so maintenance concerns don't worry me.

The code is created by allowing a good compiler (gcc 7.1) to vectorize, and then taking the output, inlining it and cleaning it up and optimizing from there. This will compile and run all the way back to gcc4.4 (which I have to support, ugh).

Thoughts and feedback welcome.

//-*-c++-*-
#pragma once

#include <cmath>
#include <algorithm>

// older versions of gcc didn't know about ymm registers in clobber list
#if (__GNUC__ == 4 && __GNUC_MINOR__ < 9) || (__GNUC__ < 4)
#define MMREG(n) "xmm"#n
#else
#define MMREG(n) "ymm"#n
#endif

// approximation of the atan2(y,x) function. This is approximately 5x faster
// than atan2() with -ffast-math on. Absolute error measured is ~2e-4, or < .1 degree.
static inline float fast_atan2(float y, float x) {
using namespace std;

if (x == 0 && y== 0) {
return 0;
}

// 7th order polynomial approximation of atan(z) on [-1,1], slightly
// tweaked to remove a multiply at the cost of very slightly higher
// error.
float a = min(abs(x),abs(y))/max(abs(x),abs(y));
float s = a*a;
float r = ((-0.0464964749f*s + 0.15931422f)*s - 0.327622764f)*s*a + a;

if (abs(y) > abs(x)) r = (float)M_PI_2 - r;
if (x < 0)           r = (float)M_PI   - r;
if (y < 0)           r =               - r;

return r;
}

// take an array of interleaved (x,y) pairs and computes fast_atan2(y,x) estimate on them.
// approximately 15-40x faster than a simple loop with atan2, depending on
// input buffer size.  Generated from vectorized output of fast_atan function.
//
// @param out    array to write output to
// @param in     input array containing interleaved pairs
// @param npair  number of input pairs to process
//
static inline void vatan2_avx(float* __restrict__ out, const float* __restrict__ in, ssize_t npair) {
// compute how many iterations to do and remainder of pairs left to do manually
size_t iters = npair/8;
size_t rem   = npair-iters*8;

// constant vectors
static const uint32_t posnan[8]  = {  0x7fffffff, 0x7fffffff, 0x7fffffff, 0x7fffffff, 0x7fffffff, 0x7fffffff, 0x7fffffff, 0x7fffffff };
static const uint32_t negnan[8]  = {  0xffffffff, 0xffffffff, 0xffffffff, 0xffffffff, 0xffffffff, 0xffffffff, 0xffffffff, 0xffffffff };
static const uint32_t signbit[8] = {  0x80000000, 0x80000000, 0x80000000, 0x80000000, 0x80000000, 0x80000000, 0x80000000, 0x80000000 };
static const float    ones[8]    = {  1,1,1,1,1,1,1,1 };
static const float    mpi_2[8]   = {  1.57079637,    1.57079637,    1.57079637,    1.57079637,    1.57079637,    1.57079637,    1.57079637,    1.57079637   };
static const float    mpi[8]     = {  3.14159274,    3.14159274,    3.14159274,    3.14159274,    3.14159274,    3.14159274,    3.14159274,    3.14159274   };
static const float    coefa[8]   = { -0.0464964733, -0.0464964733, -0.0464964733, -0.0464964733, -0.0464964733, -0.0464964733, -0.0464964733, -0.0464964733 };
static const float    coefb[8]   = {  0.159314215,   0.159314215,   0.159314215,   0.159314215,   0.159314215,   0.159314215,   0.159314215,   0.159314215  };
static const float    coefc[8]   = { -0.327622771,  -0.327622771,  -0.327622771,  -0.327622771,  -0.327622771,  -0.327622771,  -0.327622771,  -0.327622771  };

__asm__(
"    vxorps  %%ymm8, %%ymm8, %%ymm8\n\t" // ymm8 = 0
"    vmovups %[posnan], %%ymm9 \n\t"     // abs() mask
"    vmovups %[coefa],  %%ymm15\n\t"
"    vmovups %[coefb],  %%ymm14\n\t"
"    vmovups %[coefc],  %%ymm13\n\t"
"    vmovups %[ones],   %%ymm12\n\t"
"    vmovups %[mpi_2],  %%ymm11\n\t"
"    vmovups %[mpi],    %%ymm10\n\t"

// setup indices, pointers
"    mov %[in],  %%rax\n\t" // input pointer
"    mov %[out], %%rcx\n\t" // output pointer
"    xor %%r8d,  %%r8d\n\t" // r8 = 0

".p2align 4\n\t"
"LOOP%=:\n\t"
// load bottom part of ymm0 and ymm1
"    vmovups     (%%rax), %%xmm0\n\t"
"    vmovups 0x20(%%rax), %%xmm1\n\t"
"    add     $0x01, %%r8\n\t" // r8 += 1 " add$0x40,  %%rax\n\t" // in  += 16
"    add     $0x20, %%rcx\n\t" // out += 8 // load top part " vinsertf128$0x1,-0x30(%%rax), %%ymm0, %%ymm0\n\t"
"    vinsertf128 $0x1,-0x10(%%rax), %%ymm1, %%ymm1\n\t" // de-interleave x,y pairs into separate registers " vshufps$0x88, %%ymm1, %%ymm0, %%ymm3\n\t"
"    vshufps     $0xdd, %%ymm1, %%ymm0, %%ymm0\n\t" " vperm2f128$0x03, %%ymm3, %%ymm3, %%ymm2\n\t"
"    vperm2f128  $0x03, %%ymm0, %%ymm0, %%ymm1\n\t" " vshufps$0x44, %%ymm2, %%ymm3, %%ymm4\n\t"
"    vshufps     $0xee, %%ymm2, %%ymm3, %%ymm2\n\t" " vshufps$0x44, %%ymm1, %%ymm0, %%ymm3\n\t"
"    vshufps     $0xee, %%ymm1, %%ymm0, %%ymm1\n\t" " vinsertf128$0x01, %%xmm2, %%ymm4, %%ymm2\n\t"
"    vinsertf128 $0x01, %%xmm1, %%ymm3, %%ymm3\n\t" // absolute values and zero check " vandps %%ymm9, %%ymm2, %%ymm4\n\t" // abs(x) " vcmpeqps %%ymm8, %%ymm2, %%ymm0\n\t" // x == 0? " vandps %%ymm9, %%ymm3, %%ymm6\n\t" // abs(y) " vcmpeqps %%ymm8, %%ymm3, %%ymm1\n\t" // y == 0? // compute argument a to polynomial " vmaxps %%ymm4, %%ymm6, %%ymm5\n\t" // max(abs(x), abs(y)) " vandps %%ymm0, %%ymm1, %%ymm1\n\t" // x == 0 && y == 0 " vminps %%ymm4, %%ymm6, %%ymm0\n\t" // min(abs(x), abs(y)) " vcmpltps %%ymm6, %%ymm4, %%ymm4\n\t" // abs(x) < abs(y) " vrcpps %%ymm5, %%ymm7 \n\t" // compute 1/max(abs(x), abs(y)) " vmulps %%ymm5, %%ymm7, %%ymm5\n\t" " vcmpltps %%ymm8, %%ymm2, %%ymm2\n\t" // x < 0 // compute polynomial " vmulps %%ymm5, %%ymm7, %%ymm5\n\t" " vaddps %%ymm7, %%ymm7, %%ymm7\n\t" " vsubps %%ymm5, %%ymm7, %%ymm7\n\t" " vmulps %%ymm7, %%ymm0, %%ymm5\n\t" " vmulps %%ymm5, %%ymm5, %%ymm7\n\t" " vmulps %%ymm15,%%ymm7, %%ymm0\n\t" " vaddps %%ymm14,%%ymm0, %%ymm0\n\t" " vmulps %%ymm7, %%ymm0, %%ymm0\n\t" " vaddps %%ymm13,%%ymm0, %%ymm0\n\t" " vmulps %%ymm7, %%ymm0, %%ymm0\n\t" // finish up " vxorps %[negnan],%%ymm1,%%ymm7\n\t" " vaddps %%ymm12,%%ymm0, %%ymm0\n\t" " vandps %%ymm4, %%ymm7, %%ymm4\n\t" " vandps %%ymm2, %%ymm7, %%ymm2\n\t" " vmulps %%ymm5, %%ymm0, %%ymm0\n\t" " vsubps %%ymm0, %%ymm11,%%ymm5\n\t" " vblendvps %%ymm4, %%ymm5, %%ymm0, %%ymm0\n\t" " vsubps %%ymm0, %%ymm10,%%ymm5\n\t" " vblendvps %%ymm2, %%ymm5, %%ymm0, %%ymm0\n\t" " vcmpleps %%ymm3, %%ymm8, %%ymm2\n\t" " vxorps %[signbit], %%ymm0, %%ymm4\n\t" " vcmpltps %%ymm8, %%ymm3, %%ymm3\n\t" " vandps %%ymm2, %%ymm7, %%ymm2\n\t" " vandps %%ymm3, %%ymm7, %%ymm7\n\t" " vblendvps %%ymm1, %%ymm8, %%ymm4, %%ymm1\n\t" " vblendvps %%ymm7, %%ymm4, %%ymm1, %%ymm1\n\t" " vblendvps %%ymm2, %%ymm0, %%ymm1, %%ymm1\n\t" // store to result " vmovups %%xmm1,-0x20(%%rcx)\n\t" " vextractf128$0x1,%%ymm1,-0x10(%%rcx)\n\t"

// are we done?
"    cmp    %[iters],%%r8\n\t"
"    jb     LOOP%=\n\t"
"    vzeroupper\n\t"
:
: [posnan]  "m" (*posnan),  [negnan] "m" (*negnan), [coefa] "m" (*coefa), [coefb] "m"  (*coefb),
[coefc]   "m" (*coefc),   [ones]   "m" (*ones),   [mpi_2] "m" (*mpi_2), [mpi]   "m"  (*mpi),
[signbit] "m" (*signbit), [in]     "r" (in),      [out]   "r" (out),    [iters] "er" (iters)
: MMREG(0), MMREG(1), MMREG(2),  MMREG(3),  MMREG(4),  MMREG(5),  MMREG(6),  MMREG(7),
MMREG(8), MMREG(9), MMREG(10), MMREG(11), MMREG(12), MMREG(13), MMREG(14), MMREG(15),
"rax", "rcx", "r8", "memory"
);

// finish remainder
if (rem > 0) {
in  += iters*16;
out += iters*8;

for (size_t ii=0; ii < rem; ii++) {
out[ii] = fast_atan2(in[2*ii+1], in[2*ii+0]);
}
}
}

void vatan2_reg(float* out, const float* in, ssize_t npair) {
for (ssize_t ii=0; ii < npair; ii++) {
out[ii] = fast_atan2(in[2*ii+1], in[2*ii+0]);
}
}


and a test harness:

#include <math.h>
#include <stdio.h>
#include <stdint.h>
#include <stdlib.h>
#include <string.h>

#include <vector>

#include <simd.h>

inline uint64_t rdtscp() {
uint32_t eax,edx;
asm volatile ("rdtscp\n" : "=a" (eax), "=d" (edx) :: "ecx");
return ((uint64_t)edx << 32) | eax;
}

int main(int argc, const char* argv[]) {
ssize_t npoint = 100*1024*1024ull + 5;
if (argc > 1) { npoint = atoi(argv[1]); }

float *data   = (float*)malloc(npoint*2*sizeof(float));
float *odata  = (float*)malloc(npoint*sizeof(float));
float *odata2 = (float*)malloc(npoint*sizeof(float));
for (ssize_t ii=0; ii < npoint; ii++) {
data[2*ii+1] = rand();
data[2*ii+0] = rand();
}

printf("done generating\n");

uint64_t fasttsc = rdtscp();
vatan2_avx(odata, data, npoint);
fasttsc = rdtscp() - fasttsc;

uint64_t normtsc = rdtscp();
for (ssize_t ii=0; ii < npoint; ii++) {
odata2[ii] = atan2f(data[2*ii+1], data[2*ii+0]);
}
normtsc = rdtscp() - normtsc;

double err=0;
for (ssize_t ii=0; ii < npoint; ii++) {
double e = std::abs(odata[ii] - odata2[ii]);
if (e > err) {
err = e;
printf("max error so far: %.4e\n", err);
}
}

printf("\n");
printf("%6zd  %15zd  %12.3f  %15zd  %12.3f  %.3f\n", npoint,
fasttsc, (double)fasttsc/npoint, normtsc, (double)normtsc/npoint, (double)normtsc/(double)fasttsc
);

FILE *out = fopen("/dev/null", "w");
fwrite(odata,  npoint, sizeof(float), out);
fwrite(odata2, npoint, sizeof(float), out);
fclose(out);
}


compile with:

g++ -O3 -ffast-math -mavx -Wall -Wextra -I. test.cc -o test

• Looking at your code, it's highly structured and logically organized so humans can read it. But that isn't always optimal for execution. Writing this with intrinsics and examining optimized output might be informative. FWIW, gcc has a pragma for setting compiler optimizations. Not trying to push you to intrinsics, but you might want to give it a second thought. However if you ultimately decide against intrinsics, how about using an actual asm routine? Getting constraints correct can be really hard. – David Wohlferd Sep 2 '17 at 4:54
• You should really use intrinsics because with intrinsics, the compiler can interleave your code with other code in the function this is inlined into and improve registration choices. This generally improves performance. With inline assembly, your function is a black box the compiler cannot look inside and optimize. That's really not something you want. – FUZxxl Sep 2 '17 at 14:11
• @DavidWohlferd The pragma won't allow me to #include the appropriate instrinsics, I have to have the flag on for the whole compilation unit which is unfriendly for my use. – gct Sep 2 '17 at 14:30
• @FUZxxl I agree, but in general I'm writing these functions to operate on, at minimum, vectors with thousands of values of in them, so I'm not terribly worried if there's a little inefficiency with the register allocation across the loop. – gct Sep 2 '17 at 14:31
• @SeanMcAllister And you are also not worried that anybody might have to maintain that? – FUZxxl Sep 2 '17 at 14:37

So, just for fun, I've converted your code to intrinsics (and probably made a bunch of mistakes in the process):

extern void vatan2i(float* __restrict__ out, const float* __restrict__ in, ssize_t npair) {
// compute how many iterations to do and remainder of pairs left to do manually
size_t iters = npair/8;
size_t rem   = npair-iters*8;

size_t r8;
const float *__restrict__ rax;
float *__restrict__ rcx;
__m128 xmm0, xmm1, xmm2;
__m256 ymm0, ymm1, ymm2, ymm3, ymm4, ymm5, ymm6, ymm7;
__m256 ymm8, ymm9, ymm10, ymm11, ymm12, ymm13, ymm14, ymm15;

ymm8 = _mm256_setzero_ps();
ymm9 = _mm256_castsi256_ps(_mm256_set1_epi32(0x7fffffff)); // posnan
ymm15 = _mm256_set1_ps(-0.0464964733); // coefa
ymm14 = _mm256_set1_ps(0.159314215); // coefb
ymm13 = _mm256_set1_ps(-0.327622771); // coefc
ymm12 = _mm256_set1_ps(1); // ones
ymm11 = _mm256_set1_ps(1.57079637); // mpi_2
ymm10 = _mm256_set1_ps(3.14159274); // mpi

// setup indices, pointers
rax = in;
rcx = out;
r8 = 0;

do {
// load bottom part of ymm0 and ymm1
r8 += 1;
rax += 16;
rcx += 8;

ymm0 = _mm256_castps128_ps256(xmm0);
ymm1 = _mm256_castps128_ps256(xmm1);
ymm0 = _mm256_insertf128_ps(ymm0, _mm_loadu_ps(rax - 12), 1);
ymm1 = _mm256_insertf128_ps(ymm1, _mm_loadu_ps(rax - 4), 1);

// de-interleave x,y pairs into separate registers
ymm3 = _mm256_shuffle_ps(ymm0, ymm1, 0x88);
ymm0 = _mm256_shuffle_ps(ymm0, ymm1, 0xdd);
ymm2 = _mm256_permute2f128_ps(ymm3, ymm3, 0x03);
ymm1 = _mm256_permute2f128_ps(ymm0, ymm0, 0x03);
ymm4 = _mm256_shuffle_ps(ymm3, ymm2, 0x44);
ymm2 = _mm256_shuffle_ps(ymm3, ymm2, 0xee);
ymm3 = _mm256_shuffle_ps(ymm0, ymm1, 0x44);
ymm1 = _mm256_shuffle_ps(ymm0, ymm1, 0xee);
xmm1 = _mm256_castps256_ps128(ymm1);
xmm2 = _mm256_castps256_ps128(ymm2);
ymm2 = _mm256_insertf128_ps(ymm4, xmm2, 1);
ymm3 = _mm256_insertf128_ps(ymm3, xmm1, 1);

// absolute values and zero check
ymm4 = _mm256_and_ps(ymm2, ymm9);
ymm0 = _mm256_cmp_ps(ymm2, ymm8, 0); // eq
ymm6 = _mm256_and_ps(ymm3, ymm9);
ymm1 = _mm256_cmp_ps(ymm3, ymm8, 0); // eq

// compute argument a to polynomial
ymm5 = _mm256_max_ps(ymm6, ymm4);
ymm1 = _mm256_and_ps(ymm1, ymm0);
ymm0 = _mm256_min_ps(ymm6, ymm4);
ymm4 = _mm256_cmp_ps(ymm4, ymm6, 1); // lt
ymm7 = _mm256_rcp_ps(ymm5);
ymm5 = _mm256_mul_ps(ymm7, ymm5);
ymm2 = _mm256_cmp_ps(ymm2, ymm8, 1); // lt

// compute polynomial
ymm5 = _mm256_mul_ps(ymm7, ymm5);
ymm7 = _mm256_sub_ps(ymm7, ymm5);
ymm5 = _mm256_mul_ps(ymm0, ymm7);
ymm7 = _mm256_mul_ps(ymm5, ymm5);
ymm0 = _mm256_mul_ps(ymm7, ymm15);
ymm0 = _mm256_mul_ps(ymm0, ymm7);
ymm0 = _mm256_mul_ps(ymm0, ymm7);

// finish up
ymm7 = _mm256_xor_ps(ymm1, _mm256_castsi256_ps(_mm256_set1_epi32(0xffffffff))); // negnan
ymm4 = _mm256_and_ps(ymm7, ymm4);
ymm2 = _mm256_and_ps(ymm7, ymm2);
ymm0 = _mm256_mul_ps(ymm0, ymm5);
ymm5 = _mm256_sub_ps(ymm11, ymm0);
ymm0 = _mm256_blendv_ps(ymm0, ymm5, ymm4);
ymm5 = _mm256_sub_ps(ymm10, ymm0);
ymm0 = _mm256_blendv_ps(ymm0, ymm5, ymm2);
ymm2 = _mm256_cmp_ps(ymm8, ymm3, 2); // le
ymm4 = _mm256_xor_ps(ymm0, _mm256_castsi256_ps(_mm256_set1_epi32(0x80000000))); // signbit
ymm3 = _mm256_cmp_ps(ymm3, ymm8, 1); // lt
ymm2 = _mm256_and_ps(ymm7, ymm2);
ymm7 = _mm256_and_ps(ymm7, ymm3);
ymm1 = _mm256_blendv_ps(ymm4, ymm8, ymm1);
ymm1 = _mm256_blendv_ps(ymm1, ymm4, ymm7);
ymm1 = _mm256_blendv_ps(ymm1, ymm0, ymm2);

// store to result
xmm1 = _mm256_castps256_ps128(ymm1);
_mm_store_ps(rcx - 8, xmm1);
_mm_store_ps(rcx - 4, _mm256_extractf128_ps(ymm1, 1));
} while (r8 < iters);

// finish remainder
if (rem > 0) {
in  += iters*16;
out += iters*8;

for (size_t ii=0; ii < rem; ii++) {
out[ii] = fast_atan2(in[2*ii+1], in[2*ii+0]);
}
}
}


When I compile this with a reasonably modern compiler (here: clang 4.0) I immediately see a number of different choices the compiler makes:

• Instead of allocating arrays for your constants, the compiler uses vbroadcastss to broadcast single floats into each vector entry. This reduces the size of the auxiliary arrays needed, improving cache usage.
• where arrays are needed to represent constants, the compiler can place them in a special section that merges all constant arrays with the same contents. This is the case for example when your function is inlined multiple times. This does not work with static const arrays because the C standard makes some guarantees about their addresses being distinct.
• The compiler emits slightly better shuffling code using vunpcklpd and vunpckhpd.
• The compiler does some funny things with your “compute argument a to polynomial” section I do not really understand.
• The compiler turns your two stores into one by simply emitting vmovups %ymm1, -0x20(%rcx)

Also, my own 50 cents:

vxorps (negnan),%ymm1,%ymm7


why don't you do

vcmpneqps %ymm8,%ymm3,%ymm7


to avoid a costly load from memory?

You should begin your own assembly label names with .L indicating a local label. This is useful because otherwise, LOOP looks just like an ordinary function name and could cause collisions.

About your comment that you don't want special compiler flags: The whole point of these is allowing the compiler to select the best instructions for a given CPU. So rather, you should rewrite your code with #ifdef __AVX__ and the like to detect what CPU features the compiler has enabled and then select an optimized implementation depending on that. This way, your code does the right thing and performs well regardless of whether the user enabled AVX or not. And if the user did not enable AVX, you don't get an invalid instruction at runtime.

• Definitely open to using compiler output to improve the assembly. What architecture did you compile for? I generated the original assembly using the auto vectorizer in gcc 7.1 and it decided breaking up the loads was a good call for some reason. I plan on detecting AVX directly using cpuid so I can set a function pointer to the correct implementation. My users will often compile on one machine and run on another over NFS... – gct Sep 2 '17 at 15:04
• @SeanMcAllister I compiled for Haswell with clang 4.0 on FreeBSD. Interestingly, clang doesn't prefer vbroadcastss for all architectures. – FUZxxl Sep 2 '17 at 15:29
• I'm compiling for a generic architecture I guess, so it's probably being conservative with the loads/stores, I might target something like ivy bridge that's not super recent but not super old. I never know what my code will run on intel wise, so getting too far into the weeds with optimization probably isn't profitable. – gct Sep 2 '17 at 16:42
• @PeterCordes Perhaps it might make sense to look at the unvectorized code and start anew with the implementation beginning with that. – FUZxxl Sep 15 '17 at 10:31
• @FUZxxl: I think so. Also, probably using a _mm256_div_ps would be better than a vrcpps + Newton, especially on Skylake but there's enough non-divider-port stuff going on that a single vdivps is a throughput win over vrcpps + multiple other instructions on Haswell or earlier, too. – Peter Cordes Sep 15 '17 at 10:39