hand_coded_abb_to_abb_beta_fock_scheme6.cpp

#include <vector>
#include <memory>
#include <iostream>

#include <essqc/exception.h>
#include <essqc/lapack_blas_interface.h>
#include <essqc/threading.h>
#include <essqc/esmp2/tensor_indexing_classes.h>
#include <essqc/esmp2/global_work_space.h>
#include <essqc/esmp2/hand_coded_abb_to_abb_beta_fock.h>

void essqc::esmp2::hc_abb_to_abb_beta_fock_s6_batched(const int nocc,
						      const int nvir,
						      const int nnto,
						      const std::vector<double> & fmat,
						      Tensor_abb_ijkabc & ten_in,
						      Tensor_abb_ijkabc & ten_out)
{
  // setting up some data structures and saving some values useful in the batched dgemm calls
  const int norb = nocc+nvir;

  //number of gemms that will be grouped into a batch
  const size_t batch_size = std::max(16, 2*essqc::get_number_of_omp_threads()); // somewhat arbitrary choice

  // buffer size between each dgemm call
  const size_t buffer_size = 256;

  // pointer to the beginning of the  block of memory for holding the matricies within a batch
  // need to be careful to not overwrite this memory space
  // this is slightly larger than actually needed
  const size_t work_size = size_t(nocc*nocc) + size_t((batch_size+1) * 2*nocc*nocc) + buffer_size*batch_size;
  double * const work_space = essqc::esmp2::get_global_work_space(work_size);

  // Allocate vectors to hold the loop index values corresponding to a batch's gemms.
  //  -- Note that we make room for one more than the batch size.
  //     The extra spot will be used to remember the index values at
  //     which the next batch will start.
  std::vector<int> outer_loop_index_values(batch_size + 1, 0);
  std::vector<int> middl_loop_index_values(batch_size + 1, 0);
  std::vector<int> inner_loop_index_values(batch_size + 1, 0);

  // Allocate vectors to hold pointers to the a, b, and c matrix inputs to dgemm.
  //  -- Note that we again make room for one more than the batch size.
  //     This is to handle the possible special case in which the last group
  //     of gemms fits evenly into a batch of size batch_size+1.
  std::vector<double *> gemm_a(batch_size + 1, NULL);
  std::vector<double *> gemm_b(batch_size + 1, NULL);
  std::vector<double *> gemm_c(batch_size + 1, NULL);

  // Store the occupied-occupied slice of the fock matrix near the beginning of the work space.
  double * const oofock = work_space;
  for ( int jp = 0; jp < nocc; jp++ )
    for ( int j = 0; j < nocc; j++ )
      oofock[ j + nocc*jp ] = fmat.at(j + norb*jp);
  
  //original autogenerated code
  //  for ( int i = 0; i < nocc; i++ ) {
  //  for ( int a = nocc; a < norb; a++ ) {
  //  for ( int j = 0; j < nocc; j++ ) {
  //    const int kstart = ( a == i + nocc ?  j+1 : 0        );
  //  for ( int k = kstart; k < nocc; k++ ) {
  //    const int bstart = ( a == i + nocc ? nocc : j+nocc   );
  //    const int bend   = ( a == i + nocc ? norb : j+nocc+1 );
  //  for ( int b = bstart; b < bend; b++ ) {
  //    if ( i + nocc == a && j + nocc == b ) {
  //      if ( i >= nnto && j >= nnto ) continue;
  //    } else { 
  //      if ( i + nocc == a && i >= nnto ) continue;
  //      if ( j + nocc == b && j >= nnto ) continue;
  //    }
  //    const int cstart = ( a == i + nocc ?  b+1 : nocc     );
  //  for ( int c = cstart; c < norb ; c++ ) {
  //    if ( j == k ) continue;
  //    if ( b == c ) continue;
  //    if ( b == j + nocc && c == k + nocc && k <= j ) continue;
  //  for ( int jp = 0; jp < nocc; jp++ ) {
  //    const int kpstart = ( a == i + nocc ?  jp+1 : 0         );
  //    if ( k <  kpstart ) continue;
  //    const int bpstart = ( a == i + nocc ?  nocc : jp+nocc   );
  //    const int bpend   = ( a == i + nocc ?  norb : jp+nocc+1 );
  //    if ( b <  bpstart ) continue;
  //    if ( b >= bpend ) continue;
  //    if ( i + nocc == a && jp + nocc == b ) {
  //      if ( i >= nnto && jp >= nnto ) continue;
  //    } else { 
  //      if ( i + nocc == a && i >= nnto ) continue;
  //      if ( jp + nocc == b && jp >= nnto ) continue;
  //    }
  //    if ( jp == k ) continue;
  //    if ( b == c ) continue;
  //    if ( b == jp + nocc && c == k + nocc && k <= jp ) continue;
  //    if ( a == i + nocc && k <= jp ) continue;
  //    if ( a == i + nocc && c <= b ) continue;
  //    if ( a != i + nocc && b != jp + nocc ) continue;
  //    out_tensor(i,j,k,a,b,c) -= fock_matrix.at(jp*norb+j) * in_tensor(i,jp,k,a,b,c);
  //  }  }  }  }  }  }  }

  if (true){
    // case 1: only a and i are paired, scales as O^3V^2
    //  for ( int i = 0; i < nnto; i++ ) {
    //    if (i+nocc>=norb) continue;
    //    const int a = i+nocc;
    //  for ( int j = 0; j < nocc; j++ ) {
    //  for ( int k = j+1; k < nocc; k++ ) {
    //  for ( int b = nocc; b < norb; b++ ) {
    //    if (b==j+nocc) continue;
    //  for ( int c = b+1; c < norb ; c++ ) {
    //  for ( int jp = 0; jp < nocc; jp++ ) {
    //    if ( k <  jp+1 ) continue;
    //    if (b==jp+nocc) continue;
    //    out_tensor(i,j,k,a,b,c) -= fock_matrix.at(jp*norb+j) * in_tensor(i,jp,k,a,b,c);
    //  }  }  }  }  }  }  }

    // determine how many gemms there are for these terms
    size_t n_gems = 0;
    for (int i=0; i<nnto; i++){
      if (i+nocc >= norb) continue;
      for (int b=nocc; b<norb-1; b++){
	for (int c=b+1; c<norb; c++){
	  n_gems++;
	}
      }
    }

    // initialize flag to determine if the current batch is the last batch
    bool last_batch = false;

    // store initial values of the indicies
    *outer_loop_index_values.rbegin() = 0;      // i start
    *middl_loop_index_values.rbegin() = nocc;   // b start
    *inner_loop_index_values.rbegin() = nocc+1; // c start

    //loop over batches
    for (size_t total_gemm_count = 0; total_gemm_count < n_gems;){
      bool keep_going = true; // flag for if more gemms should be added to current batch
      bool starter_iter = true; // flag for if this is the starting iteration of the current batch
      size_t batch_gemm_count = 0; // counter for how many gemms are in the current batch

      // loop through indicies
      for (int i = (starter_iter ? *outer_loop_index_values.rbegin() : 0); i<nnto && keep_going; i++){
	if (i+nocc >= norb) continue;
	for (int b = (starter_iter ? *middl_loop_index_values.rbegin() : nocc); b<norb-1 && keep_going; b++){
	  for(int c = (starter_iter ? *inner_loop_index_values.rbegin() : b+1); c<norb && keep_going; c++){


	    // record current index values (add set of indicies to batch)
	    outer_loop_index_values[batch_gemm_count] = i;
            middl_loop_index_values[batch_gemm_count] = b;
            inner_loop_index_values[batch_gemm_count] = c;

	    // increment total gemm counter and number of gemms in current batch
	    total_gemm_count++;
	    batch_gemm_count++;

	    // determine if this is the last gemm and set booleans accordingly
	    if (total_gemm_count == n_gems){
	      last_batch = true;
	      keep_going = false;

	    // If instead we have filled the batch with its quota of gemms and recorded the index
            // values at which the next batch should start, set the flag indicating we should stop looping
            // and make sure to "un-count" the gemm corresponding to the saved index values for the
            // next batch, as that gemm will not be processed in the current batch;
            } else if ( batch_gemm_count == batch_size + 1 ) {
              total_gemm_count--; // "un-counting" the gemm corresponding to the next batch's starting index values
              keep_going = false;
            }

	    // Set the flag to tell that we are no longer on the first iteration so that the loop indices will start
            // from their normal starting values rather than the initiation values for this batch.
            starter_iter = false;
	    
	  }//end loop over i
	}//end loop over b
      }//end loop over c

      // Get an end value for the current batch.  Note this is one larger for the last batch, as in that case there are
      // no saved index values for the next batch, and so all saved index values correspond to gemms in the current batch.
      // The gemms in the batch will be numbered 0, 1, 2, ..., batch_end-1
      const size_t batch_end = batch_gemm_count - ( last_batch ? 0 : 1 );

      // some useful constants for the tensor indexation, see tensor_indexing_classes.h
      const int nop = essqc::esmp2::n_pair(nocc);
      const int nvp = essqc::esmp2::n_pair(nvir);

      #pragma omp parallel for
      for (size_t x=0; x<batch_end; x++){
	gemm_a.at(x) = oofock;
	//                             for the oofock   for the other gemm_b's and gemm_c's
	gemm_b.at(x) = work_space + size_t(nocc*nocc) + size_t(2*x*nocc*nocc) + x*buffer_size;
	gemm_c.at(x) = gemm_b.at(x) + nocc*nocc;

	const int i = outer_loop_index_values.at(x);
	const int b = middl_loop_index_values.at(x);
	const int c = inner_loop_index_values.at(x);

	const int bc = essqc::esmp2::cmpd_index_p_lt_q(nvir, b-nocc, c-nocc);

	std::fill(gemm_b[x], gemm_b[x] + 2*nocc*nocc, 0.0); // filling both gemm_b and gemm_c with zeros

	const double * tin = ten_in.raw_data_ptr();

	for (int jp=0; jp<nocc-1; jp++){
	  if(b==jp+nocc) continue;
	  for (int k=jp+1; k<nocc; k++){
	    const int jk = essqc::esmp2::cmpd_index_p_lt_q(nocc, jp, k);
	    gemm_b[x][jp + nocc*(k)] = tin[nop*nvp*i + nvp*jk + bc];
	  }
	}
      }

      // do matrix multiplications
      {
	int gemm_gc = 1; // number of groups
	int gemm_gs = batch_end; // group sizes
	char gemm_transa = 'N';
	char gemm_transb = 'N';
	int gemm_m = nocc;
	int gemm_n = nocc;
	int gemm_k = nocc;
	double gemm_alpha = 1.0;
	double gemm_beta  = 0.0;
	int gemm_lda = nocc;
	int gemm_ldb = nocc;
	int gemm_ldc = nocc;
	essqc::dgemm_batch(&gemm_transa, &gemm_transb, &gemm_m, &gemm_n, &gemm_k, &gemm_alpha, &gemm_a.at(0), &gemm_lda, &gemm_b.at(0), &gemm_ldb, &gemm_beta, &gemm_c.at(0), &gemm_ldc, gemm_gc, &gemm_gs);
      }
      
      // add output of matrix multiplications to output tensor
      #pragma omp parallel for
      for ( int x = 0; x < batch_end; x++ ) {
	
        const int i = outer_loop_index_values.at(x);
        const int b = middl_loop_index_values.at(x);
        const int c = inner_loop_index_values.at(x);

	const int bc = essqc::esmp2::cmpd_index_p_lt_q(nvir, b-nocc, c-nocc);
	
	// put output in ten_out
        double * const tout = ten_out.raw_data_ptr();
        for ( int j = 0; j < nocc-1; j++ ) {
	  if (b==j+nocc) continue;
          for ( int k = j+1; k < nocc; k++ ) {
	    const int jk = essqc::esmp2::cmpd_index_p_lt_q(nocc, j, k);
            tout[nop*nvp*i + nvp*jk + bc] -= gemm_c[x][ j + nocc*(k) ];
          }
        }

      }

      if ( last_batch && total_gemm_count != n_gems )
        throw essqc::Exception("( a == i + nocc ) part of hc_abb_to_abb_beta_fock_s6_batched saw ( last_batch == true && total_count != n_gems )");

    }

    if ( ! last_batch )
      throw essqc::Exception("( a == i + nocc ) part of hc_abb_to_abb_beta_fock_s6_batched failed to set last_batch = true");
	
  }

  if (true){
    // case2: a and i, b and j are paired, scales as O^3V
    //  for ( int i = 0; i < nnto; i++ ) {
    //    if (i+nocc>=norb) continue;
    //    const int a = i+nocc;
    //  for ( int j = 0; j < nocc; j++ ) {
    //    if (j+nocc>=norb) continue;
    //    const int b = j+nocc;
    //  for ( int k = j+1; k < nocc; k++ ) {
    //  for ( int c = b+1; c < norb ; c++ ) {
    //  for ( int jp = 0; jp < nocc; jp++ ) {
    //    if (b==jp+nocc) continue;
    //    if ( k <  jp+1 ) continue;
    //    out_tensor(i,j,k,a,b,c) -= fock_matrix.at(jp*norb+j) * in_tensor(i,jp,k,a,b,c);
    //  }  }  }  }  }  }  }

    std::vector<int> ivals;
    std::vector<int> jvals;
    std::vector<int> kvals;
    
    for (int i=0; i<nnto; i++){
      if (i+nocc >= norb) continue;
      for (int j=0; j<nocc-1; j++){
        if (j+nocc >= norb) continue;
        for (int k=j+1; k<nocc; k++){
	  ivals.push_back(i);
	  jvals.push_back(j);
	  kvals.push_back(k);
        }
      }
    }
    
    #pragma omp parallel for
    for (int y=0; y<ivals.size(); y++){
      const int i=ivals.at(y);
      const int j=jvals.at(y);
      const int k=kvals.at(y);
      for (int c=j+nocc+1; c<norb; c++){
        for (int jp=0; jp<k; jp++){
	  if (j==jp) continue;
	  ten_out(i,j,k,i+nocc,j+nocc,c) -= fmat.at(jp*norb+j) * ten_in(i,jp,k,i+nocc,j+nocc,c);
        }
      }
    }
    
  }

  if (true){
    // case 3: a and i, b and jp are paired, scales as O^3V
    //  for ( int i = 0; i < nnto; i++ ) {
    //    if (i+nocc>=norb) continue;
    //    const int a = i+nocc;
    //  for ( int j = 0; j < nocc; j++ ) {
    //  for ( int k = j+1; k < nocc; k++ ) {
    //  for ( int jp = 0; jp < nocc; jp++ ) {
    //    if (jp+nocc>=norb) continue;
    //    const int b = jp+nocc;
    //    if (jp==j) continue;
    //  for ( int c = jp+nocc+1; c < norb ; c++ ) {
    //    if ( k <  jp+1 ) continue;
    //    out_tensor(i,j,k,a,b,c) -= fock_matrix.at(jp*norb+j) * in_tensor(i,jp,k,a,b,c);
    //  }  }  }  }  }  }  }
    std::vector<int> ivals;
    std::vector<int> jvals;
    std::vector<int> kvals;
    
    for (int i=0; i<nnto; i++){
      if (i+nocc >= norb) continue;
      for (int j=0; j<nocc-1; j++){
        for (int k=j+1; k<nocc; k++){
	  ivals.push_back(i);
	  jvals.push_back(j);
	  kvals.push_back(k);
        }
      }
    }
    
    #pragma omp parallel for
    for (int y=0; y<ivals.size(); y++){
      const int i=ivals.at(y);
      const int j=jvals.at(y);
      const int k=kvals.at(y);
      for (int jp=0; jp<k; jp++){
	if (jp+nocc>=norb) continue;
	if (jp==j) continue;
        for (int c=jp+nocc+1; c<norb; c++){
	  ten_out(i,j,k,i+nocc,jp+nocc,c) -= fmat.at(jp*norb+j) * ten_in(i,jp,k,i+nocc,jp+nocc,c);
        }
      }
    }
    
  }
  
  if (true){
    // case 4: all pairings, scales as O^2V
    // for ( int i = 0; i < nocc; i++ ) {
    //   if (i+nocc>=norb) continue;
    //   const int a = i+nocc;
    // for ( int j = 0; j < nocc; j++ ) {
    //   if (j+nocc>=norb) continue;
    //   const int b = j+nocc;
    //   consti int jp=j;
    // for ( int k = j+1; k < nocc; k++ ) {
    //   if ( i >= nnto && j >= nnto ) continue;
    // for ( int c = b+1; c < norb ; c++ ) {
    //   if ( i >= nnto && jp >= nnto ) continue;
    //   out_tensor(i,j,k,a,b,c) -= fock_matrix.at(jp*norb+j) * in_tensor(i,jp,k,a,b,c);
    // }  }  }  }  }  }  }

    std::vector<int> ivals;
    std::vector<int> jvals;
    std::vector<int> kvals;
    
    for (int i=0; i<nocc; i++){
      for (int j=0; j<nocc-1; j++){
        if (i>=nnto && j>=nnto) continue;
        for (int k=j+1; k<nocc; k++){
	  ivals.push_back(i);
	  jvals.push_back(j);
	  kvals.push_back(k);
        }
      }
    }
    
    #pragma omp parallel for
    for (int y=0; y<ivals.size(); y++){
      const int i=ivals.at(y);
      const int j=jvals.at(y);
      const int k=kvals.at(y);
      for (int c=j+nocc+1; c<norb; c++){
	ten_out(i,j,k,i+nocc,j+nocc,c) -= fmat.at(j*norb+j) * ten_in(i,j,k,i+nocc,j+nocc,c);
      }
    }
    
  }

  if (true){
    // case 5: b and j are paired, meaning jp and b must also be paired, scales as O^2V^2
    //  for ( int i = 0; i < nocc; i++ ) {
    //  for ( int a = nocc; a < norb; a++ ) {
    //    if (a==i+nocc) continue;
    //  for ( int j = 0; j < nnto; j++ ) {
    //    if (j+nocc>=norb) continue;
    //    const int b=j+nocc;
    //    const int jp=j;
    //  for ( int k = 0; k < nocc; k++ ) {
    //  for ( int c = nocc; c < norb ; c++ ) {
    //    if ( j == k ) continue;
    //    if ( b == c ) continue;
    //    if ( b == j + nocc && c == k + nocc && k <= j ) continue;
    //    out_tensor(i,j,k,a,b,c) -= fock_matrix.at(jp*norb+j) * in_tensor(i,jp,k,a,b,c);
    //  }  }  }  }  }  }  }

    std::vector<int> ivals;
    std::vector<int> avals;
    std::vector<int> jvals;
    
    for (int i=0; i<nocc; i++){
      for (int a=nocc; a<norb; a++){
	if (a==i+nocc) continue;
        for (int j=0; j<nnto; j++){
	  if (j+nocc >= norb) continue;
	  ivals.push_back(i);
	  avals.push_back(a);
	  jvals.push_back(j);
        }
      }
    }
    
    #pragma omp parallel for
    for (int y=0; y<ivals.size(); y++){
      const int i=ivals.at(y);
      const int a=avals.at(y);
      const int j=jvals.at(y);
      for (int k=0; k<nocc; k++){
	if (j==k) continue;
        for (int c=nocc; c<norb; c++){
	  if (j+nocc==c) continue;
	  if (c==k+nocc && k<=j) continue;
	  ten_out(i,j,k,a,j+nocc,c) -= fmat.at(j*norb+j) * ten_in(i,j,k,a,j+nocc,c);
        }
      }
    }
    
  }

}