1
\$\begingroup\$

I have a function that I've been trying to make faster that computes \$XA + B\$, where \$A \in \mathbb{R}^{n \times n}\$ is a tridiagonal matrix, and \$X, B \in \mathbb{R}^{m \times n}\$. In the code, the lower, main, and upper diagonals are represented by dl, dm, du. The naive implementation of this is,

/* Compute BA + X where A is tridiagonal shape (n, n), B, X are matrices shape (m, n) */
void dgtaxmt(const double *dl, const double *dm, const double *du,
             const double *x, const double *b, double *out,
             const unsigned int n, const unsigned int m, double *temp)
{
    int i, j, ii;

    for (i = 0; i < m; i++){
        ii = i * n;
        temp[ii] = (x[ii] * dm[0]) + (x[ii + 1] * dl[0]) + b[ii];
        for (j = 1; j < n - 1; j++){
            temp[ii + j] = ((x[ii + j - 1] * du[j - 1]) + ((x[ii + j] * dm[j]) + ((x[ii + j + 1] * dl[j]) + b[ii + j])));
        }
        temp[ii + j] = (x[ii + j] * dm[j]) + (x[ii + j - 1] * du[j - 1]) + b[ii + j];
    }
    // store the transpose of temp into out
    // otrans(temp, out, m, n, 16);

}

Note that all of the pointers use the restrict keyword, but I left them out for readability. My attempts at increasing performance include tiling and using AVX intrinsics in hopes that not having to reload dl, dm, du would speed up the code, but it actually made it a lot slower. Is there anyway to vectorize this function in an efficient manner? Clang doesn't report any vectorization, and the assembly code it generates is using the xmm registers, but only the first value. Could there be a different order of operations that would allow better re-usage of data that's been previously loaded?

EDIT Here is the attempt I made with AVX intrinsics

void dgtaxmt_avx(const double *dl, const double *dm, const double *du,
                 const double *x, const double *b, double *out,
                 const unsigned int n, const unsigned int m, double *temp)
{
    // idx is for "rolling" values in vector to left by one index
    const __m256i idx = _mm256_set_epi32(1, 0, 7, 6, 5, 4, 3, 2);
    __m256d dl_vec, dm_vec, du_vec, bn, xnm1, xn, xn1, tmp4;
    __m128d tmp2;

    unsigned int i, j, ii;
    const unsigned int r = ((n - 2) & (-4)) + 1;

    for (i = 0; i < m; i++){
        ii = i * n;
        temp[ii] = (x[ii] * dm[0]) + (x[ii + 1] * dl[0]) + b[ii];
        for (j = 1; j < r; j += 4){
            dl_vec = _mm256_loadu_pd(&dl[j - 1]);
            dm_vec = _mm256_loadu_pd(&dm[j]);
            du_vec = _mm256_loadu_pd(&du[j]);
            bn     = _mm256_loadu_pd(&b[ii + j]);

            xnm1  = _mm256_loadu_pd(&x[ii + j - 1]);
            tmp2  = _mm_loadu_pd(&x[ii + j + 3]);
            xn1   = _mm256_set_m128d(tmp2, _mm256_extractf128_pd(xnm1, 1));

            // use permutations to avoid doing extra loads
            xn   = (__m256d)_mm256_permutevar8x32_ps((__m256)xnm1, idx);
            tmp4 = (__m256d)_mm256_permutevar8x32_ps((__m256)xn1,  idx);
            tmp2 = _mm256_extractf128_pd(tmp4, 0);
            xn   = _mm256_insertf128_pd(xn, tmp2, 1);

            tmp4 = _mm256_fmadd_pd(du_vec, xn1, bn);
            tmp4 = _mm256_fmadd_pd(dm_vec, xn, tmp4);
            tmp4 = _mm256_fmadd_pd(dl_vec, xnm1, tmp4);
            
            _mm256_storeu_pd(&out[ii + j], tmp4);
        }
        for (j = r; j < n - 1; j++){
            temp[ii + j] = ((x[ii + j - 1] * du[j - 1]) + ((x[ii + j] * dm[j]) + ((x[ii + j + 1] * dl[j]) + b[ii + j])));
        }
        temp[ii + j] = (x[ii + j] * dm[j]) + (x[ii + j - 1] * du[j - 1]) + b[ii + j];
    }
    
    otrans(temp, out, m, n, 16);
}
\$\endgroup\$
10
  • \$\begingroup\$ Can you show the attempted AVX code as well? \$\endgroup\$
    – harold
    Mar 11, 2021 at 18:10
  • \$\begingroup\$ @harold I edited the question to include the AVX code. Sorry it's a littles messy since it was "experimental" code \$\endgroup\$
    – vlovero
    Mar 11, 2021 at 18:44
  • \$\begingroup\$ I actually did not measure a serious performance issue in the AVX code. Are you running this on an AMD processor? _mm256_permutevar8x32_ps is a bigger problem for AMD \$\endgroup\$
    – harold
    Mar 11, 2021 at 22:15
  • \$\begingroup\$ 1. How large are \$n\$ and \$m\$, and 2. Did you try computing the result column-wise (rather than row-wise)? \$\endgroup\$
    – vnp
    Mar 11, 2021 at 22:16
  • \$\begingroup\$ @harold I'm on a macbook running an intel i5-5257U CPU @ 2.70GHz and I'm not entirely sure why I'm seeing the difference. In the intel documentation for broadwell, it says the efficiency of these instructions aren't too bad \$\endgroup\$
    – vlovero
    Mar 11, 2021 at 22:26

0

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge that you have read and understand our privacy policy and code of conduct.

Browse other questions tagged or ask your own question.