<commit_before>
#include "backend_impl.h"

#undef  __FUNCT__
#define __FUNCT__ "BuildMat_Full"
PetscErrorCode BuildMat_Full(PetscInt L,PetscInt nterms,
                             const PetscInt* masks,
                             const PetscInt* signs,
                             const PetscScalar* coeffs,Mat *A)
{
  PetscErrorCode ierr;
  PetscInt N,i,state,Istart,Iend,nrows,local_bits,nonlocal_mask;
  int mpi_rank,mpi_size;
  PetscInt d_nz,o_nz;

  PetscInt lstate,sign;
  PetscScalar tmp_val;

  N = 1<<L;

  MPI_Comm_rank(PETSC_COMM_WORLD,&mpi_rank);
  MPI_Comm_size(PETSC_COMM_WORLD,&mpi_size);

  ierr = MatCreate(PETSC_COMM_WORLD,A);CHKERRQ(ierr);
  ierr = MatSetSizes(*A,PETSC_DECIDE,PETSC_DECIDE,N,N);CHKERRQ(ierr);
  ierr = MatSetFromOptions(*A);CHKERRQ(ierr);

  nrows = PETSC_DECIDE;
  PetscSplitOwnership(PETSC_COMM_WORLD,&nrows,&N);

  /* find how many spins are local, for allocation purposes */
  local_bits = 0;
  while(nrows >>= 1) ++local_bits;

  /* now count how many terms are not on our processor */

  /* this creates something that looks like 111110000000 */
  nonlocal_mask = (-1) << local_bits;

  d_nz = o_nz = 0;
  for (i=0;i<nterms;++i) {
    /* only count each element once, even though
       there might be a few terms that contribute to it */
    if (i>0 && masks[i-1] == masks[i]) continue;

    if (masks[i] & nonlocal_mask) ++o_nz;
    else ++d_nz;
  }

  /* it seems from the documentation you can just call both of these and it will use the right one */
  ierr = MatSeqAIJSetPreallocation(*A,d_nz,NULL);CHKERRQ(ierr);
  ierr = MatMPIAIJSetPreallocation(*A,d_nz,NULL,o_nz,NULL);CHKERRQ(ierr);

  ierr = MatGetOwnershipRange(*A,&Istart,&Iend);CHKERRQ(ierr);

  /* this is where the magic happens */
  for (state=Istart;state<Iend;++state) {
    for (i=0;i<nterms;) {
      lstate = state ^ masks[i];
      tmp_val = 0;
      /* sum all terms for this matrix element */
      do {
        /* this requires gcc builtins */
        sign = 1 - 2*(__builtin_popcount(state & signs[i]) % 2);
        tmp_val += sign * coeffs[i];
        ++i;
      } while (i<nterms && masks[i-1] == masks[i]);

      /* the elements must not be repeated or else INSERT_VALUES is wrong! */
      /* I could just use ADD_VALUES but if they are repeated there is a bug somewhere else */
      ierr = MatSetValues(*A,1,&lstate,1,&state,&tmp_val,INSERT_VALUES);CHKERRQ(ierr);
    }
  }

  ierr = MatAssemblyBegin(*A,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
  ierr = MatAssemblyEnd(*A,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);

  return ierr;

}

#undef  __FUNCT__
#define __FUNCT__ "BuildMat_Shell"
PetscErrorCode BuildMat_Shell(PetscInt L,PetscInt nterms,
                              const PetscInt* masks,
                              const PetscInt* signs,
                              const PetscScalar* coeffs,Mat *A)
{
  PetscErrorCode ierr;
  PetscInt N,n;
  shell_context *ctx;

  N = 1<<L;

  n = PETSC_DECIDE;
  PetscSplitOwnership(PETSC_COMM_WORLD,&n,&N);

  ierr = BuildContext(L,nterms,masks,signs,coeffs,&ctx);CHKERRQ(ierr);

  ierr = MatCreateShell(PETSC_COMM_WORLD,n,n,N,N,ctx,A);CHKERRQ(ierr);
  ierr = MatShellSetOperation(*A,MATOP_MULT,(void(*)(void))MatMult_Shell);
  ierr = MatShellSetOperation(*A,MATOP_NORM,(void(*)(void))MatNorm_Shell);

  return ierr;

}

#undef  __FUNCT__
#define __FUNCT__ "MatMult_Shell"
PetscErrorCode MatMult_Shell(Mat A,Vec x,Vec b)
{
  PetscErrorCode ierr;
  PetscInt state,i,Istart,Iend,start,block_start,sign,m,s;
  PetscInt proc,proc_idx,proc_mask,proc_size,log2size;
  PetscBool assembling;
  const PetscScalar *x_array;
  PetscScalar c;
  shell_context *ctx;

  /* cache */
  PetscInt *lidx;
  PetscInt cache_idx,cache_idx_max;
  PetscScalar *values;

  PetscInt mpi_rank,mpi_size;
  MPI_Comm_rank(PETSC_COMM_WORLD,&mpi_rank);
  MPI_Comm_size(PETSC_COMM_WORLD,&mpi_size);

  ierr = MatShellGetContext(A,&ctx);CHKERRQ(ierr);

  ierr = PetscMalloc1(VECSET_CACHE_SIZE,&lidx);CHKERRQ(ierr);
  ierr = PetscMalloc1(VECSET_CACHE_SIZE,&values);CHKERRQ(ierr);

  /* clear out the b vector */
  ierr = VecSet(b,0);CHKERRQ(ierr);

  ierr = VecGetOwnershipRange(x,&Istart,&Iend);CHKERRQ(ierr);
  ierr = VecGetArrayRead(x,&x_array);CHKERRQ(ierr);

  /* first need to compute some constants etc */
  log2size = __builtin_ctz(mpi_size);

  /* this computes a mask over the indices that define the processor */
  proc_mask = ((1 << log2size)-1) << (ctx->L - log2size);
  proc_idx = mpi_rank << (ctx->L - log2size);
  proc_size = Iend-Istart;

  /* we are not already sending values to another processor */
  assembling = PETSC_FALSE;

  for (proc=0;proc < mpi_size;++proc) {

    /* NOTE: these breaks and continues are dangerous if
     * the data doesn't divide perfectly across processors.
     * When switching to a subspace need to be careful!
     */

    /* if there are none for this process, skip it */
    if (ctx->mask_starts[proc] == ctx->mask_starts[proc+1]) continue;

    /* if we've hit the end of the masks, stop */
    if (ctx->mask_starts[proc] == ctx->nterms) break;

    start = proc_mask & (proc_idx ^ ctx->masks[ctx->mask_starts[proc]]);
    for (block_start=start;
         block_start<start+proc_size;
         block_start+=VECSET_CACHE_SIZE) {

      cache_idx_max = PetscMin((start+proc_size)-block_start,VECSET_CACHE_SIZE);

      PetscMemzero(values,sizeof(PetscScalar)*VECSET_CACHE_SIZE);
      for (cache_idx=0;cache_idx<cache_idx_max;++cache_idx) {
        lidx[cache_idx] = block_start + cache_idx;
      }

      for (i=ctx->mask_starts[proc];i<ctx->mask_starts[proc+1];++i) {

        m = ctx->masks[i];
        s = ctx->signs[i];
        c = ctx->coeffs[i];

        for (cache_idx=0;cache_idx<cache_idx_max;++cache_idx) {
          state = (block_start + cache_idx) ^ m;
          sign = 1 - 2*(__builtin_popcount(state & s)&1);
          values[cache_idx] += x_array[state - Istart] * sign * c;
        }
      }

      /* set the last values that weren't set at the end there */
      if (assembling) {
        ierr = VecAssemblyEnd(b);CHKERRQ(ierr);
        assembling = PETSC_FALSE;
      }
      ierr = VecSetValues(b,cache_idx_max,lidx,values,ADD_VALUES);CHKERRQ(ierr);

      ierr = VecAssemblyBegin(b);CHKERRQ(ierr);
      assembling = PETSC_TRUE;
    }
  }

  if (assembling) {
    ierr = VecAssemblyEnd(b);CHKERRQ(ierr);
  }<commit_msg>rework, solid speedup<commit_after>
#include "backend_impl.h"

#undef  __FUNCT__
#define __FUNCT__ "BuildMat_Full"
PetscErrorCode BuildMat_Full(PetscInt L,PetscInt nterms,
                             const PetscInt* masks,
                             const PetscInt* signs,
                             const PetscScalar* coeffs,Mat *A)
{
  PetscErrorCode ierr;
  PetscInt N,i,state,Istart,Iend,nrows,local_bits,nonlocal_mask;
  int mpi_rank,mpi_size;
  PetscInt d_nz,o_nz;

  PetscInt lstate,sign;
  PetscScalar tmp_val;

  N = 1<<L;

  MPI_Comm_rank(PETSC_COMM_WORLD,&mpi_rank);
  MPI_Comm_size(PETSC_COMM_WORLD,&mpi_size);

  ierr = MatCreate(PETSC_COMM_WORLD,A);CHKERRQ(ierr);
  ierr = MatSetSizes(*A,PETSC_DECIDE,PETSC_DECIDE,N,N);CHKERRQ(ierr);
  ierr = MatSetFromOptions(*A);CHKERRQ(ierr);

  nrows = PETSC_DECIDE;
  PetscSplitOwnership(PETSC_COMM_WORLD,&nrows,&N);

  /* find how many spins are local, for allocation purposes */
  local_bits = 0;
  while(nrows >>= 1) ++local_bits;

  /* now count how many terms are not on our processor */

  /* this creates something that looks like 111110000000 */
  nonlocal_mask = (-1) << local_bits;

  d_nz = o_nz = 0;
  for (i=0;i<nterms;++i) {
    /* only count each element once, even though
       there might be a few terms that contribute to it */
    if (i>0 && masks[i-1] == masks[i]) continue;

    if (masks[i] & nonlocal_mask) ++o_nz;
    else ++d_nz;
  }

  /* it seems from the documentation you can just call both of these and it will use the right one */
  ierr = MatSeqAIJSetPreallocation(*A,d_nz,NULL);CHKERRQ(ierr);
  ierr = MatMPIAIJSetPreallocation(*A,d_nz,NULL,o_nz,NULL);CHKERRQ(ierr);

  ierr = MatGetOwnershipRange(*A,&Istart,&Iend);CHKERRQ(ierr);

  /* this is where the magic happens */
  for (state=Istart;state<Iend;++state) {
    for (i=0;i<nterms;) {
      lstate = state ^ masks[i];
      tmp_val = 0;
      /* sum all terms for this matrix element */
      do {
        /* this requires gcc builtins */
        sign = 1 - 2*(__builtin_popcount(state & signs[i]) % 2);
        tmp_val += sign * coeffs[i];
        ++i;
      } while (i<nterms && masks[i-1] == masks[i]);

      /* the elements must not be repeated or else INSERT_VALUES is wrong! */
      /* I could just use ADD_VALUES but if they are repeated there is a bug somewhere else */
      ierr = MatSetValues(*A,1,&lstate,1,&state,&tmp_val,INSERT_VALUES);CHKERRQ(ierr);
    }
  }

  ierr = MatAssemblyBegin(*A,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
  ierr = MatAssemblyEnd(*A,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);

  return ierr;

}

#undef  __FUNCT__
#define __FUNCT__ "BuildMat_Shell"
PetscErrorCode BuildMat_Shell(PetscInt L,PetscInt nterms,
                              const PetscInt* masks,
                              const PetscInt* signs,
                              const PetscScalar* coeffs,Mat *A)
{
  PetscErrorCode ierr;
  PetscInt N,n;
  shell_context *ctx;

  N = 1<<L;

  n = PETSC_DECIDE;
  PetscSplitOwnership(PETSC_COMM_WORLD,&n,&N);

  ierr = BuildContext(L,nterms,masks,signs,coeffs,&ctx);CHKERRQ(ierr);

  ierr = MatCreateShell(PETSC_COMM_WORLD,n,n,N,N,ctx,A);CHKERRQ(ierr);
  ierr = MatShellSetOperation(*A,MATOP_MULT,(void(*)(void))MatMult_Shell);
  ierr = MatShellSetOperation(*A,MATOP_NORM,(void(*)(void))MatNorm_Shell);

  return ierr;

}

static inline PetscScalar cix(PetscReal c,const PetscScalar *x)
{
  PetscScalar y;
  PetscReal* yp;
  yp = (PetscReal*)&y;
  yp[0] = -c*((PetscReal*)x)[1];
  yp[1] = c*((PetscReal*)x)[0];
  return y;
}

#undef  __FUNCT__
#define __FUNCT__ "MatMult_Shell"
PetscErrorCode MatMult_Shell(Mat A,Vec x,Vec b)
{
  PetscErrorCode ierr;
  PetscInt state,i,Istart,Iend,start,block_start,sign,m,s;
  PetscInt proc,proc_idx,proc_mask,proc_size,log2size;
  PetscBool assembling,real;
  const PetscScalar *x_array, *x_begin, *xp;
  PetscScalar c;
  PetscReal cr;
  shell_context *ctx;

  /* cache */
  PetscInt *lidx;
  PetscInt cache_idx,cache_idx_max,cache_idx_submax,iterate_max;
  PetscScalar *summed_c,*values,*vp;

  PetscInt mpi_rank,mpi_size;
  MPI_Comm_rank(PETSC_COMM_WORLD,&mpi_rank);
  MPI_Comm_size(PETSC_COMM_WORLD,&mpi_size);

  ierr = MatShellGetContext(A,&ctx);CHKERRQ(ierr);

  ierr = PetscMalloc1(VECSET_CACHE_SIZE,&lidx);CHKERRQ(ierr);
  ierr = PetscMalloc1(VECSET_CACHE_SIZE,&values);CHKERRQ(ierr);

  /* clear out the b vector */
  ierr = VecSet(b,0);CHKERRQ(ierr);

  ierr = VecGetOwnershipRange(x,&Istart,&Iend);CHKERRQ(ierr);
  ierr = VecGetArrayRead(x,&x_array);CHKERRQ(ierr);

  /* first need to compute some constants etc */
  log2size = __builtin_ctz(mpi_size);

  /* this computes a mask over the indices that define the processor */
  proc_mask = ((1 << log2size)-1) << (ctx->L - log2size);
  proc_idx = mpi_rank << (ctx->L - log2size);
  proc_size = Iend-Istart;
  x_begin = x_array - Istart;

  /* we are not already sending values to another processor */
  assembling = PETSC_FALSE;

  for (proc=0;proc < mpi_size;++proc) {

    /* NOTE: these breaks and continues are dangerous if
     * the data doesn't divide perfectly across processors.
     * When switching to a subspace need to be careful!
     */

    /* if there are none for this process, skip it */
    if (ctx->mask_starts[proc] == ctx->mask_starts[proc+1]) continue;

    /* if we've hit the end of the masks, stop */
    if (ctx->mask_starts[proc] == ctx->nterms) break;

    start = proc_mask & (proc_idx ^ ctx->masks[ctx->mask_starts[proc]]);
    for (block_start=start;
         block_start<start+proc_size;
         block_start+=VECSET_CACHE_SIZE) {

      cache_idx_max = PetscMin((start+proc_size)-block_start,VECSET_CACHE_SIZE);

      PetscMemzero(values,sizeof(PetscScalar)*VECSET_CACHE_SIZE);
      for (cache_idx=0;cache_idx<cache_idx_max;++cache_idx) {
        lidx[cache_idx] = block_start + cache_idx;
      }

      for (i=ctx->mask_starts[proc];i<ctx->mask_starts[proc+1];++i) {

        m = ctx->masks[i];
        s = ctx->signs[i];
        c = ctx->coeffs[i];

        iterate_max = 1<<PetscMin(__builtin_ctz(m),__builtin_ctz(s));

        cr = PetscRealPart(c);
        real = (cr != 0);
        if (!real) cr = PetscImaginaryPart(c);

        if (real) {
          if (iterate_max < 1<<3) {
            for (cache_idx=0;cache_idx<cache_idx_max;++cache_idx) {
              state = (block_start+cache_idx) ^ m;
              sign = __builtin_popcount(state & s)&1;
              values[cache_idx] -= ((sign^(sign-1))*cr)*x_begin[state];
            }
          }
          else {
            for (cache_idx=0;cache_idx<cache_idx_max;++cache_idx) {
              state = (block_start+cache_idx) ^ m;
              sign = __builtin_popcount(state & s)&1;

              vp = values + cache_idx;
              xp = x_begin + state;
              (*vp) -= ((sign^(sign-1))*cr)*(*xp);

              cache_idx_submax = PetscMin(cache_idx+(iterate_max-(state%iterate_max)),cache_idx_max);
              for (++cache_idx,++xp,++vp;cache_idx<cache_idx_submax;++cache_idx,++xp,++vp) {
                (*vp) -= ((sign^(sign-1))*cr)*(*xp);
              }
              --cache_idx;
            }
          }
        }
        else {
          if (iterate_max < 1<<3) {
            for (cache_idx=0;cache_idx<cache_idx_max;++cache_idx) {
              state = (block_start+cache_idx) ^ m;
              sign = __builtin_popcount(state & s)&1;
              values[cache_idx] -= (sign^(sign-1))*cix(cr,x_begin + state);
            }
          }
          else {
            for (cache_idx=0;cache_idx<cache_idx_max;++cache_idx) {
              state = (block_start+cache_idx) ^ m;
              sign = __builtin_popcount(state & s)&1;

              vp = values + cache_idx;
              xp = x_begin + state;
              (*vp) -= (sign^(sign-1))*cix(cr,xp);

              cache_idx_submax = PetscMin(cache_idx+(iterate_max-(state%iterate_max)),cache_idx_max);
              for (++cache_idx,++xp,++vp;cache_idx<cache_idx_submax;++cache_idx,++xp,++vp) {
                (*vp) -= (sign^(sign-1))*cix(cr,xp);
              }
              --cache_idx;
            }
          }
        }
      }

      /* set the last values that weren't set at the end there */
      if (assembling) {
        ierr = VecAssemblyEnd(b);CHKERRQ(ierr);
        assembling = PETSC_FALSE;
      }
      ierr = VecSetValues(b,cache_idx_max,lidx,values,ADD_VALUES);CHKERRQ(ierr);

      ierr = VecAssemblyBegin(b);CHKERRQ(ierr);
      assembling = PETSC_TRUE;
    }
  }

  if (assembling) {
    ierr = VecAssemblyEnd(b);CHKERRQ(ierr);
  }
{'fraction_non_alphanumeric': 0.13143896033937108, 'fraction_numerical': 0.007676250757524746, 'mean_word_length': 3.1497625034925956, 'pattern_counts': {'":': 0, '<': 58, '<?xml version=': 0, '>': 35, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 2, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '221800', 'n_tokens_mistral': 5737, 'n_tokens_neox': 5467, 'n_words': 1038}