linear_system_banded_avx.c

#include <math.h>
#include "mesh.h"
#include "linear_system.h"
#include <stdlib.h>
#include <string.h>

#define CONMAX 12
void connectivity_insert(int *connectivity, int i, int j)
{
  for (int k = 0; k < CONMAX; ++k) {
    int *p = connectivity + CONMAX*i + k;
    if (*p == -1)
      *p = j;
    if (*p == j)
      return;
  }
  printf("ERROR : node %i has more than %i neighbours\n", i, CONMAX);
}

void reverse_cuthill_mckee(Mesh *mesh, int *ordering)
{
  int *node_connectivity = malloc(sizeof(int)*mesh->n_nodes*CONMAX);
  for (int i = 0; i < mesh->n_nodes*CONMAX; ++i) {
    node_connectivity[i] = -1;
  }
  for (int i = 0; i < mesh->n_triangles; ++i) {
    int *tri = mesh->triangles + i*3;
    connectivity_insert(node_connectivity, tri[0], tri[1]);
    connectivity_insert(node_connectivity, tri[1], tri[0]);
    connectivity_insert(node_connectivity, tri[0], tri[2]);
    connectivity_insert(node_connectivity, tri[2], tri[0]);
    connectivity_insert(node_connectivity, tri[1], tri[2]);
    connectivity_insert(node_connectivity, tri[2], tri[1]);
  }
  int *node_degree = malloc(sizeof(int)*mesh->n_nodes);
  for (int i = 0; i < mesh->n_nodes; ++i) {
    node_degree[i] = 0;
    for (int j = 0; j < CONMAX; ++j) {
      if (node_connectivity[CONMAX*i+j] == -1)
        break;
      node_degree[i] += 1;
    }
  }
  int *queue = malloc(sizeof(int)*mesh->n_nodes);
  queue[0] = 0;
  for (int i = 0; i < mesh->n_nodes; ++i){
    ordering[i] = -1;
    if (node_degree[queue[0]] > node_degree[i] && node_degree[i] !=0)
      queue[0] = i;
  }
  int stage_start = 0;
  int stage_end = 1;
  int queue_end = 1;
  int id = 0;
  while(stage_start != stage_end) {
    for (int i = stage_start; i < stage_end; ++i) {
      int c = queue[i];
      ordering[c] = id++;
      for(int j = 0; j < node_degree[c]; ++j) {
        int o = node_connectivity[c*CONMAX+j];
        if (ordering[o] == -1) {
          ordering[o] = -2;
#if 1
          queue[queue_end++] = o;
#else
          int k = stage_end;
          while(k < queue_end && node_degree[queue[k]] < node_degree[o])
            k++;
          for (int l = queue_end; l > k; --l)
            queue[l] = queue[l-1];
          queue[k] = o;
          queue_end++;
#endif
        }
      }
    }
    stage_start = stage_end;
    stage_end = queue_end;
  }
  for (int i = 0; i < mesh->n_nodes; ++i)
    if(ordering[i] >= 0)
      ordering[i] = id-1-ordering[i];
  printf("id = %i\n", id);
  free(queue);
  free(node_degree);
  free(node_connectivity);
}

#if defined(__AVX2__)
#define PADDING 4
#elif defined(__AVX__)
#define PADDING 2
#else
#define PADDING 1
#endif

#include <immintrin.h>
// y[from:to] += a*x[from:to]
// y and x must be 256-bit aligned
// memory should be allocated in consequence
static void avx_row_axpy(double a, double *__restrict__ x, double *__restrict__ y, int from, int to)
{
#if defined(__AVX2__)
  int pto = (to+3)/4*4;
  int pfrom = (from+3)/4*4;
  for (int i = from; i < pfrom; ++i)
    y[i] += a*x[i];
  __m256d ma = _mm256_set1_pd(a);
  for (int i = pfrom; i < pto; i+=4) {
    __m256d mx = _mm256_load_pd(x+i);
    __m256d my = _mm256_load_pd(y+i);
    _mm256_store_pd(y+i,_mm256_fmadd_pd(ma,mx,my));
  }
#elif defined(__AVX__)
  int pto = (to+1)/2*2;
  int pfrom = (from+1)/2*2;
  for (int i = from; i < pfrom; ++i)
    y[i] += a*x[i];
  __m128d ma = _mm_set1_pd(a);
  for (int i = pfrom; i < pto; i+=2) {
    __m128d mx = _mm_load_pd(x+i);
    __m128d my = _mm_load_pd(y+i);
#if defined(__FMA__)
    _mm_store_pd(y+i, _mm_fmadd_pd(ma,mx,my));
#else
    __m128d mz = _mm_mul_pd(ma, mx);
    _mm_store_pd(y+i, _mm_add_pd(my,mz));
#endif
  }
#else
  for (int i = from; i < to; ++i)
    y[i] += a*x[i];
#endif
}

struct LinearSystemStruct{
  double *M;
  int *row_start;
  int *row_end;
  int *row_lu_end;
  double **rows;
  double *b;
  double *x;
  int *node_map;
  int n_fields;
  int n;
  Mesh *mesh;
  int *isfixed;
};

static int imin(int a, int b) {
  return a < b ? a : b;
}

static int imax(int a, int b) {
  return a > b ? a : b;
}

static int LUPDecompose(LinearSystem *system){
  int i,j;
  int N = system->n;
  double **A = system->rows;
  for(i=0;i<N;i++){
    for(j=i+1;j<system->row_lu_end[i];j++){
      if(i < system->row_start[j] || A[j][i] == 0.)
        continue;
      A[j][i]/=A[i][i];
      avx_row_axpy(-A[j][i],A[i],A[j],i+1,system->row_end[i]);
    }
  }
  return(1);
}

static void LUPSolve(LinearSystem *system){
  double *x = system->x;
  double *b= system->b;
  int N = system->n;
  double **A = system->rows;
  for(int i=0;i<N;i++){
    //x[i]=b[P[i]];
    x[i]=b[i];
    for(int k=system->row_start[i];k<i;k++)
      x[i]-=A[i][k]*x[k];
  }
  for(int i=N-1;i>=0;i--){
    for(int k=i+1;k<system->row_end[i] && k < N;k++){
      x[i]-=A[i][k]*x[k];
    }
    x[i]=x[i]/A[i][i];
  }
}

static void linear_system_zero_row(LinearSystem *system, int i)
{
  for (int j = system->row_start[i]; j < system->row_end[i]; ++j) {
    system->rows[i][j] = 0.;
  }
}

LinearSystem *linear_system_new(Mesh *mesh, int n_fields, int n_boundaries, const StrongBoundary *boundaries)
{
  LinearSystem *system = malloc(sizeof(LinearSystem));
  system->n_fields = n_fields;
  system->mesh = mesh;
  system->node_map = malloc(sizeof(int)*mesh->n_nodes);
  reverse_cuthill_mckee(mesh, system->node_map);
  int n_used_nodes = 0;
  for (int i = 0; i < mesh->n_nodes; ++i)
    if (system->node_map[i] >= 0)
      n_used_nodes ++;

  system->n = n_used_nodes*system->n_fields;
  int *node_row_start = malloc(sizeof(int)*n_used_nodes);
  int *node_row_end = malloc(sizeof(int)*n_used_nodes);
  for (int i = 0; i < n_used_nodes; ++i) {
    node_row_end[i] = 0;
    node_row_start[i] = n_used_nodes;
  }
  for (int i = 0; i < mesh->n_triangles; ++i) {
    int *tri = mesh->triangles + i*3;
    for (int j = 0; j < 3; ++j) {
      int n0 = system->node_map[tri[j]];
      for (int k = 0; k < 3; ++k) {
        int n1 = system->node_map[tri[k]];
        if (node_row_start[n0] > n1)
          node_row_start[n0] = n1;
        if (node_row_end[n0] < n1)
          node_row_end[n0] = n1;
      }
    }
  }
  int total_size = 0;
  system->row_start = malloc(sizeof(int)*system->n);
  system->row_end = malloc(sizeof(int)*system->n);
  system->row_lu_end = malloc(sizeof(int)*system->n);
  for (int i = 0; i < n_used_nodes; ++i) {
    for (int j = 0; j < n_fields; ++j) {
      int k = i*n_fields +j;
      system->row_start[k] = (node_row_start[i]*n_fields)/PADDING*PADDING;
      system->row_end[k] = (node_row_end[i]*n_fields+n_fields+(PADDING-1))/PADDING*PADDING;
      system->row_lu_end[k] = k;
    }
  }
  for (int i = 0; i < system->n; ++i) {
    for (int j = system->row_start[i]; j < i; ++j) {
      if (system->row_lu_end[j] < i+1) system->row_lu_end[j] = i+1;
      if (system->row_end[i] < system->row_end[j]) system->row_end[i] = system->row_end[j];
    }
    total_size += system->row_end[i]-system->row_start[i];
  }
  free(node_row_start);
  free(node_row_end);
  system->M = _mm_malloc(sizeof(double)*total_size, 32);
  system->rows = malloc(sizeof(double*)*system->n);
  for (int i = 0; i < total_size; ++i)
    system->M[i] = 0;
  system->rows[0] = system->M;
  for (int i = 1; i < system->n; ++i){
    system->rows[i] = system->rows[i-1]+system->row_end[i-1]-system->row_start[i];
  }
  system->b = malloc(sizeof(double)*system->n);
  system->x = malloc(sizeof(double)*system->n);
  system->isfixed = malloc(sizeof(int)*system->n);
  for (int i = 0; i < system->n; ++i) {
    system->isfixed[i] = 0;
  }
  for (int ibnd = 0; ibnd < n_boundaries; ++ibnd) {
    const StrongBoundary *bnd = boundaries + ibnd;
    int iphys;
    for (iphys = 0; iphys < mesh->n_phys; ++iphys) {
      if (strcmp(bnd->tag, mesh->phys_name[iphys]) == 0)
        break;
    }
    if (iphys == mesh->n_phys){
      printf("Boundary tag \"%s\" not found.\n", bnd->tag);
      continue;
    }
    for (int i = 0; i < mesh->phys_n_nodes[iphys]; ++i){
      if (system->node_map[mesh->phys_nodes[iphys][i]] >= 0)
        system->isfixed[system->node_map[mesh->phys_nodes[iphys][i]]*system->n_fields+bnd->field] = 1;
    }
  }
  return system;
}

void linear_system_add_to_matrix(LinearSystem *system, int el0, int el1, const double *local_matrix){
  int *tri0 = &system->mesh->triangles[el0*3];
  int *tri1 = &system->mesh->triangles[el1*3];
  int nf = system->n_fields;
  for (int i = 0; i < 3; ++i) {
    for (int inf = 0; inf < 3; ++inf) {
      int ii = system->node_map[tri0[i]]*nf + inf;
      for (int j = 0; j < 3; ++j) {
        for (int jnf = 0; jnf < nf; ++jnf) {
          int jj = system->node_map[tri1[j]]*nf + jnf;
          system->rows[ii][jj] += local_matrix[(inf*3+i)*nf*3+jnf*3+j];
        }
      }
    }
  }
}

void linear_system_add_to_rhs(LinearSystem *system, int el0, const double *local_vector)
{
  const int *tri = system->mesh->triangles + el0*3;
  int n_fields = system->n_fields;
  for (int i = 0; i < n_fields; ++i) {
    for (int j = 0; j < 3; ++j) {
      int m = system->node_map[tri[j]]*n_fields+i;
      system->b[m] += local_vector[i*3+j];
    }
  }
}

void linear_system_zero_matrix_and_rhs(LinearSystem *system)
{
  for (int i = 0; i < system->n; ++i){
    system->b[i] = 0.;
    linear_system_zero_row(system, i);
  }
}

void linear_system_solve(LinearSystem *system, double *solution){
  double **rows = system->rows;
  for (int i = 0; i < system->n; ++i){
    if (system->isfixed[i]) {
      linear_system_zero_row(system, i);
      rows[i][i] = 1.;
      system->b[i] = 0;
    }
  }
  LUPDecompose(system);
  LUPSolve(system);
  for (int i = 0; i < system->mesh->n_nodes; ++i){
    int ii = system->node_map[i];
    if (ii < 0)
      continue;
    for (int j = 0; j < system->n_fields; ++j){
      solution[i*system->n_fields+j] = system->x[ii*system->n_fields+j];
    }
  }
}

void linear_system_free(LinearSystem *system)
{
  free(system->b);
  free(system->x);
  free(system->M);
  free(system->rows);
  free(system->row_start);
  free(system->row_end);
  free(system->row_lu_end);
  free(system->isfixed);
  free(system->node_map);
  free(system);
}

double linear_system_get_rhs_norm(LinearSystem *system)
{
  double norm = 0;
  for (int i = 0; i < system->n;++i)
    if (!system->isfixed[i])
      norm += system->b[i]*system->b[i];
  return sqrt(norm);
}

void initialize_linear_solver(int argc, char **argv){}