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);
}
_mm256_permutevar8x32_ps
is a bigger problem for AMD \$\endgroup\$