加快牛顿寻找第n根的方法

时间:2013-06-12 21:01:35

标签: c simd intrinsics avx

让我用一个声明来预测这个问题;这段代码按预期工作,但它的速度很慢,非常慢。有没有办法让牛顿方法收敛得更快,或者设置一个__m256 var等于单个浮点数的方法而不会弄乱浮点数组等?

__m256 nthRoot(__m256 a, int root){

#define aligned __declspec(align(16)) float

// uses the calculation
// n_x+1 = (1/root)*(root * x + a / pow(x,root))


//initial numbers
aligned r[8];
aligned iN[8];
aligned mN[8];

//Function I made to fill arrays
/*
template<class T>
void FillArray(T a[],T b)
{
int n = sizeof(a)/sizeof(T);
for(int i = 0; i < n; a[i++] = b);
}*/

//fills the arrays
FillArray(iN,(1.0f/(float)root));
FillArray(mN,(float)(root-1));
FillArray(r,(float)root);

//loads the arrays into the sse componenets
__m256 R = _mm256_load_ps(r);
__m256 Ni = _mm256_load_ps(iN);
__m256 Nm = _mm256_load_ps(mN);

    //sets initaial guess to 1 / (a * root)
__m256 x = _mm256_rcp_ps(_mm256_mul_ps(R,a));

for(int i = 0; i < 20 ; i ++){
    __m256 tmpx = x;
    for(int k = 0 ; k < root -2 ; k++){
        tmpx = _mm256_mul_ps(x,tmpx);
    }
    //f over f'
    __m256 tar = _mm256_mul_ps(a,_mm256_rcp_ps(tmpx)); 
    //fmac with Ni*X+tar
    tar = _mm256_fmadd_ps(Nm,x,tar);
    //Multipled by Ni
    x = _mm256_mul_ps(Ni,tar);
}
return x;
}

编辑#1

__m256 SSEnthRoot(__m256 a, int root){

__m256 R = _mm256_set1_ps((float)root);
__m256 Ni = _mm256_set1_ps((1.0f)/((float)root));
__m256 Nm = _mm256_set1_ps((float)(root -1));

__m256 x = _mm256_mul_ps(a,_mm256_rcp_ps(R));

for(int i = 0; i < 10 ; i ++){
    __m256 tmpx = x;
    for(int k = 0 ; k < root -2 ; k++){
        tmpx = _mm256_mul_ps(x,tmpx);
    }
    //f over f'
    __m256 tar = _mm256_mul_ps(a,_mm256_rcp_ps(tmpx)); 
    //mult nm x then add tar because my compiler stoped thinking that fmadd is a valid instruction
    tar = _mm256_add_ps(_mm256_mul_ps(Nm,x),tar);
    //Multiplied by the inverse of power
    x = _mm256_mul_ps(Ni,tar);
}

return x;
}

任何使牛顿方法更快收敛的提示或指针(不是内存种类)都会受到赞赏。

使用_mm256_rcp_ps()在_mm256_set1_ps()函数调用中删除编辑#2,因为我已经将我需要的倒数加载到R

__m256 SSEnthRoot(__m256 a, int root){
__m256 R = _mm256_set1_ps((float)root);
__m256 Ni = _mm256_rcp_ps(R);
__m256 Nm = _mm256_set1_ps((float)(root -1));

__m256 x = _mm256_mul_ps(a,Ni);
for(int i = 0; i < 20 ; i ++){
    __m256 tmpx = x;
    for(int k = 0 ; k < root -2 ; k++)
        tmpx = _mm256_mul_ps(x,tmpx);
    //f over f'
    __m256 tar = _mm256_mul_ps(a,_mm256_rcp_ps(tmpx)); 
    //fmac with Ni*X+tar
            //my compiler believes in fmac again
    tar = _mm256_fmadd_ps(Nm,x,tar);
    //Multiplied by the inverse of power
    x = _mm256_mul_ps(Ni,tar);
}
return x;
}

编辑#3

__m256 SSEnthRoot(__m256 a, int root){
__m256 Ni = _mm256_set1_ps(1.0f/(float)root);
__m256 Nm = _mm256_set1_ps((float)(root -1));
__m256 x = _mm256_mul_ps(a,Ni);
for(int i = 0; i < 20 ; i ++){
    __m256 tmpx = x;
    for(int k = 0 ; k < root -2 ; k++)
        tmpx = _mm256_mul_ps(x,tmpx);
    __m256 tar = _mm256_mul_ps(a,_mm256_rcp_ps(tmpx)); 
    tar = _mm256_fmadd_ps(Nm,x,tar);
    x = _mm256_mul_ps(Ni,tar);
}
return x;
}

3 个答案:

答案 0 :(得分:2)

__m256向量的所有元素设置为单个值:

__m256 v = _mm256_set1_ps(1.0f);

或在您的具体案例中:

__m256 R =  _mm256_set1_ps((float)root);
__m256 Ni = _mm256_set1_ps((1.0f/(float)root));
__m256 Nm = _mm256_set1_ps((float)(root-1));

显然,一旦你做出这个改变,你就可以摆脱FillArray的东西。

答案 1 :(得分:2)

您的pow功能效率低下。

for(int k = 0 ; k < root -2 ; k++)
        tmpx = _mm256_mul_ps(x,tmpx);

在你的例子中,你是第29根。您需要pow(x, 29-1) = x^28。目前你使用27次乘法,但只有6次乘法就可以做到这一点。

x^28 = (x^4)*(x^8)*(x^16)
x^4 = y -> 2 multiplications
x^8 = y*y = z -> 1 multiplication
x^16 = z^2 = w-> 1 multiplications
y*z*w -> 2 multiplications
6 multiplications in total

以下是您的代码的改进版本,其速度大约是我系统的两倍。它使用我创建的新pow_avx_fast函数,它使用AVX一次为8个浮点数执行x ^ n。它确实例如x ^ 28乘以6次乘法而不是27次。请进一步查看我的答案。我发现了一个版本,可以在一定的容差xacc内找到结果。如果收敛很快,这可能会快得多。

inline __m256 pow_avx_fast(__m256 x, const int n) {
    //n must be greater than zero
    if(n%2 == 0) {
        return pow_avx_fast(_mm256_mul_ps(x, x), n/2);
    }
    else {
        if(n>1) return _mm256_mul_ps(x,pow_avx_fast(_mm256_mul_ps(x, x), (n-1)/2));
        return x;
    }
}

inline __m256 SSEnthRoot_fast(__m256 a, int root) {
    // n_x+1 = (1/root)*((root-1) * x + a / pow(x,root-1))
    __m256 R = _mm256_set1_ps((float)root);
    __m256 Ni = _mm256_rcp_ps(R);
    __m256 Nm = _mm256_set1_ps((float)(root -1));

    __m256 x = _mm256_mul_ps(a,Ni);
    for(int i = 0; i < 20 ; i ++) {
        __m256 tmpx = pow_avx_fast(x, root-1);
        //f over f'
        __m256 tar = _mm256_mul_ps(a,_mm256_rcp_ps(tmpx)); 
        //fmac with Ni*X+tar
        //tar = _mm256_fmadd_ps(Nm,x,tar);
        tar = _mm256_add_ps(_mm256_mul_ps(Nm,x),tar);
        //Multiplied by the inverse of power
        x = _mm256_mul_ps(Ni,tar);
    }
    return x;
}

有关如何编写高效pow函数的详细信息,请参阅这些链接http://en.wikipedia.org/wiki/Addition-chain_exponentiationhttp://en.wikipedia.org/wiki/Exponentiation_by_squaring

另外,你最初的猜测可能不太好。这是基于您的方法找到第n个根的标量代码(但使用的数学pow函数可能比您的快)。解决16的第4个根(即2)需要大约50次迭代。对于你使用的20次迭代,它返回超过4000,这不是接近2.0的地方。因此,您需要调整方法以进行足够的迭代,以确保在一定容差范围内得到合理的答案。

float fx(float a, int n, float x) {
    return 1.0f/n * ((n-1)*x + a/pow(x, n-1));
}
float scalar_nthRoot_v2(float a, int root) {
    //sets initaial guess to 1 / (a * root)
    float x = 1.0f/(a*root);
    printf("x0 %f\n", x);
    for(int i = 0; i<50; i++) {
        x = fx(a, root, x);
        printf("x %f\n", x);
    }
    return x;
}

我从这里得到了牛顿方法的公式。 http://en.wikipedia.org/wiki/Nth_root_algorithm

以下是您的函数的一个版本,它在一定的容差xacc内提供结果,或者在nmax次迭代后没有收敛时退出。如果在少于20次迭代中发生收敛,则此函数可能比您的方法快得多。它要求所有八个浮子一次收敛。换句话说,如果七个收敛而一个不收敛,则其他七个必须等待不收敛的那个。这就是SIMD的问题(在GPU上也是如此),但总的来说它比没有SIMD的情况更快。

int get_mask(const __m256 dx, const float xacc) {
    __m256i mask = _mm256_castps_si256(_mm256_cmp_ps(dx, _mm256_set1_ps(xacc), _CMP_GT_OQ));
    return _mm_movemask_epi8(_mm256_castsi256_si128(mask)) + _mm_movemask_epi8(_mm256_extractf128_si256(mask,1));
}

inline __m256 SSEnthRoot_fast_xacc(const __m256 a, const int root, const int nmax, float xacc) {
    // n_x+1 = (1/root)*(root * x + a / pow(x,root))
    __m256 R = _mm256_set1_ps((float)root);
    __m256 Ni = _mm256_rcp_ps(R);
    //__m256 Ni = _mm256_set1_ps(1.0f/root);
    __m256 Nm = _mm256_set1_ps((float)(root -1));

    __m256 x = _mm256_mul_ps(a,Ni);

    for(int i = 0; i <nmax ; i ++) {
        __m256 tmpx = pow_avx_fast(x, root-1);
        __m256 tar = _mm256_mul_ps(a,_mm256_rcp_ps(tmpx)); 
        //tar = _mm256_fmadd_ps(Nm,x,tar);
        tar = _mm256_add_ps(_mm256_mul_ps(Nm,x),tar);
        tmpx = _mm256_mul_ps(Ni,tar);
        __m256 dx = _mm256_sub_ps(tmpx,x);
        dx = _mm256_max_ps(_mm256_sub_ps(_mm256_setzero_ps(), dx), dx); //fabs(dx)
        int cnt = get_mask(dx, xacc);
        if(cnt == 0) return x;
        x = tmpx;
    }
    return x; //at least one value out of eight did not converge by nmax.
}

这是avx的pow函数的更一般版本,也适用于n <= 0。

__m256 pow_avx(__m256 x, const int n) {
    if(n<0) {
        return pow_avx(_mm256_rcp_ps(x), -n);
    }
    else if(n == 0) { 
        return _mm256_set1_ps(1.0f);
    }
    else if(n == 1) {
        return x;
    }
    else if(n%2 ==0) {
        return pow_avx(_mm256_mul_ps(x, x), n/2);
    }
    else {
        return _mm256_mul_ps(x,pow_avx(_mm256_mul_ps(x, x), (n-1)/2));
    }
}

其他一些建议

您可以使用找到第n个根的SIMD数学库。 SIMD math libraries for SSE and AVX

对于英特尔,您可以使用昂贵且封闭源代码的SVML(英特尔的OpenCL驱动程序使用SVML,因此您可以免费获得它)。对于AMD,您可以使用免费但封闭源的LIBM。有几个开源SIMD数学库,如http://software-lisc.fbk.eu/avx_mathfun/https://bitbucket.org/eschnett/vecmathlib/wiki/Home

答案 2 :(得分:0)

也许您应该在日志域中执行此操作。

pow(a,1/root) == exp( log(x) /root) 

Julien Pommier有sse_mathfun.h具有SSE,SSE2日志和exp功能,但我不能说我特别使用了这些功能。这些技术可以扩展到avx。