Parallel assembly of matrix

Question

I have a matrix which I want to assembly quickly, which is in block form: $$ A = \pmatrix{ A_{11} & A_{12} & A_{13} \\ A_{21} & A_{22} & A_{23} \\ A_{31} & A_{32} & A_{33}} $$ with: $$ A_{ij} = \sum_{k=1}^N \alpha_{ijk} v_k v_k^T $$ and my current implementation is quite slow, so I am looking for advice on how to speed it up. For a given $k$, the matrix $v_k v_k^T$ appears in each block $A_{ij}$ so it can be precomputed and then copied into each block. I use OpenMP to distribute the calculation of $v_k v_k^T$ across multiple threads, which is quite fast. The bottleneck of the code is when I need to update the matrix $A$. I give each thread its own matrix A to update and then sum them all together at the end for the final A. My pseudo code is below:

# pragma omp parallel for
loop from k = 1 to N
  v_k = buildv(k);   // v_k and G_k are stored in thread-specific storage
  G_k = v_k v_k^T;   // using DSYR - this step is fast
  loop i = 1 to 3    // this part is really SLOW
    loop j = 1 to 3
      alpha_ijk = ...; // this is fast
      A[thread_id](i,j) += alpha_ijk * G_k; // I wrote a matrix DAXPY routine for this

sum A[thread_id] over all threads to build final A

When I wrote the code, I had thought that the calculation of $v_k$ and $G_k$ would be the slow parts, since they involve computations (BLAS calls). However I was suprised that the slow part is just copying already computed results to the matrix $A$ (even if each thread has its own $A$). So I guess I need to focus on parallizing the copy operation somehow. Does anyone have any suggestions?

Note that $N$ is quite large (about 5 million), which is why I focused on parallezing the computation of $v_k v_k^T$.

EDIT

As requested, I have coded up a minimal working example (available here). I compiled the code below with:

gcc -g -Wall -O2 -fopenmp -o matrix matrix.c -lm -lgsl -lcblas -latlas

Note that the prerequesites to run this are the GSL library and ATLAS BLAS - though you could replace ATLAS with any cblas library (-lgslcblas for example).

I added a #if flag to enable/disable the matrix copy operation which is the slow part. When the matrices are copied into the larger thread-specific matrices, the total runtime on my machine is 81 seconds. When this code is disabled, the total run time is 4.1 seconds. I have 24 cores on the machine.

For simplicity, I use a random number generator to specify the vectors $v_k$.

#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <omp.h>
#include <sys/time.h>

#include <gsl/gsl_math.h>
#include <gsl/gsl_matrix.h>
#include <gsl/gsl_blas.h>
#include <gsl/gsl_rng.h>

static void
random_vector(const double a, const double b, gsl_vector * v, gsl_rng * r)
{
  size_t i;

  for (i = 0; i < v->size; ++i)
    {
      double vi = (b - a) * gsl_rng_uniform(r) + a;
      gsl_vector_set(v, i, vi);
    }
}

static int
matrix_daxpy(const double alpha, const gsl_matrix * X, gsl_matrix * Y)
{
  const size_t M = Y->size1;
  const size_t N = Y->size2;

  if (X->size1 != M || X->size2 != N)
    {
      GSL_ERROR ("matrices must have same dimensions", GSL_EBADLEN);
    }
  else
    {
      if (alpha != 0.0)
        {
          const size_t tda_X = X->tda;
          const size_t tda_Y = Y->tda;
          size_t i, j;

          for (i = 0; i < M; ++i)
            {
              for (j = 0; j < N; ++j)
                {
                  Y->data[tda_Y * i + j] += alpha * X->data[tda_X * i + j];
                }
            }
        }

      return GSL_SUCCESS;
    }
}

int
main(int argc, char *argv[])
{
  const size_t max_threads = (size_t) omp_get_max_threads();
  const size_t N = 500000;  /* my real N is around 5 million */
  const size_t m = 300;
  const size_t nblocks = 3; /* number of block matrices in A */
  const double a = -10.0;
  const double b = 10.0;
  gsl_vector **v = malloc(max_threads * sizeof(gsl_vector *));
  gsl_matrix **GTG = malloc(max_threads * sizeof(gsl_matrix *));
  gsl_rng **r = malloc(max_threads * sizeof(gsl_rng *));
  gsl_matrix **A = malloc(max_threads * sizeof(gsl_matrix *));
  size_t k;
  struct timeval tv0, tv1;

  for (k = 0; k < max_threads; ++k)
    {
      v[k] = gsl_vector_alloc(m);
      GTG[k] = gsl_matrix_calloc(m, m);
      A[k] = gsl_matrix_calloc(nblocks * m, nblocks * m);
      r[k] = gsl_rng_alloc(gsl_rng_default);
    }

  gettimeofday(&tv0, NULL);

#pragma omp parallel for
  for (k = 0; k < N; ++k)
    {
      int thread_id = omp_get_thread_num();
      size_t i, j;

      /* construct v_k */
      random_vector(a, b, v[thread_id], r[thread_id]);

      /* GTG += v_k v_k^T */
      gsl_blas_dsyr(CblasLower, 1.0, v[thread_id], GTG[thread_id]);

      /* copy lower triangle to upper triangle */
      gsl_matrix_transpose_tricpy('L', 0, GTG[thread_id], GTG[thread_id]);

#if 1
      /* SLOW PART: now copy alpha_{ijk} v_k v_k^T into thread-specific final A */
      for (i = 0; i < nblocks; ++i)
        {
          for (j = 0; j <= i; ++j)
            {
              double alpha_ijk = 1.0; /* for simplicity */
              gsl_matrix_view Av = gsl_matrix_submatrix(A[thread_id], i * m, j * m, m, m); /* A_{ij} m-by-m block */

              /* copy alpha_ijk v_k v_k^T into appropriate spot */
              matrix_daxpy(alpha_ijk, GTG[thread_id], &Av.matrix);
            }
        }
#endif
    }

  gettimeofday(&tv1, NULL);
  fprintf(stderr, "elapsed time = %.4f [sec]\n",
          (tv1.tv_sec + tv1.tv_usec * 1.0e-6) - (tv0.tv_sec + tv0.tv_usec * 1.0e-6));

  return 0;
}

Answer 1

You are updating $m^2\times\mathit{nblocks}^2$ memory locations at each iteration, 8 bytes each, a total of $N$ times, so you're getting memory bandwidth $N\times m^2\times\mathit{nblocks}^2\times 8\mathrm{B}\,/\,81\mathrm{s} = 40\mathrm{GB}/\mathrm{s}$. (When you disable the slow part, it's the same formula with $\mathit{nblocks}=1$, $t=4.1\mathrm{s}$, giving $82\mathrm{GB}/\mathrm{s}$.) This doesn't count updates to $G$, so might be off by a small factor, but that doesn't matter. It seems likely to me that your code is memory-limited and that you are already within a small integer factor of max possible memory bandwidth (whatever that is on your machine).

One possible solution in such cases is to reduce the matrix blocks further until they fit in cache, and pay extra attention to only manipulate values that are already in cache. This is very similar to what a typical gemm routine does with matrix blocks. So make the blocks smaller so that $8\mathrm{B}\times m^2$ fits into your per-thread cache: if it's 256KB, that's $\sqrt{256\times 2^{10}/8}\sim 180$ give or take (it's often easier to just experiment with setting this value directly by hand instead of trying to compute it on your own). For each such block, you would still accumulate $\alpha_{ijk}v_k v_k^\top$ but using only a part of each $v_k$, depending on the block, so it's the ger non-symmetric rank-1 update. Note: I haven't tested this idea in code.

P.S. I think you can also save a factor of two by only doing the triangular-transpose-copy at the very end, once per block, instead of at every iteration.

Answer 2

The first suggestion (without seeing the code) is to check if your memory access pattern is reasonable or, preferably, optimal for the critical parts of the code.

One needs to keep in mind how the data is stored in RAM and minimize cache misses during your operations. In your case, you should check

• whether your large matrix $A$, its sub-blocks $A_{ij}$ are stored column-wise or row-wise. Whichever way they are stored, it should be reflected in the way they are accessed for read/write operations.
• try to avoid DSYR and DAXPY calls to haveINCX and INCY different from 1. If the storage spacing increments are different from 1 (non-contiguous memory access pattern), it will probably be beneficial to reorganize the data. This is unlikely the problem for DSYR since you mentioned it to be fast.
• you mentioned that you've written your own DAXPY for the matrix. I am not sure if you need one of your own since you can call regular DAXPY, just with N now being not the length of the vector, but the size of the matrix. Your custom implementation of DAXPY is highly likely to slow down your computations.

• changing the order of operations might also help. Right now, you have the loop over $k$ on top. I would try looping over the blocks first and then have an OpenMP-parallelized loop for $k$ inside. The number of blocks is very small, and precomputing $G_k$ and $v_k$ (shared by 9 blocks) can be traded for data locality.

• you may try to manually unroll the "block-loop" for $i,j$. It's hard to say how you are storing the individual blocks and what compiler options are used, but changing from 2-D indexing $i,j = 1,\ldots,3$ to 1-D indexing $b=1,\ldots,9$ can be beneficial. It will certainly remove one level of folded loops without too much hassle.