/******************************************************************************
 * Copyright 1998-2019 Lawrence Livermore National Security, LLC and other
 * HYPRE Project Developers. See the top-level COPYRIGHT file for details.
 *
 * SPDX-License-Identifier: (Apache-2.0 OR MIT)
 ******************************************************************************/

/******************************************************************************
 *
 * Member functions for hypre_CSRMatrix class.
 *
 *****************************************************************************/

#include "seq_mv.h"

#ifdef HYPRE_PROFILE
HYPRE_Real hypre_profile_times[HYPRE_TIMER_ID_COUNT] = { 0 };
#endif

/*--------------------------------------------------------------------------
 * hypre_CSRMatrixCreate
 *--------------------------------------------------------------------------*/

hypre_CSRMatrix *
hypre_CSRMatrixCreate( HYPRE_Int num_rows,
                       HYPRE_Int num_cols,
                       HYPRE_Int num_nonzeros )
{
   hypre_CSRMatrix  *matrix;

   matrix = hypre_CTAlloc(hypre_CSRMatrix, 1, HYPRE_MEMORY_HOST);

   hypre_CSRMatrixData(matrix)           = NULL;
   hypre_CSRMatrixI(matrix)              = NULL;
   hypre_CSRMatrixJ(matrix)              = NULL;
   hypre_CSRMatrixBigJ(matrix)           = NULL;
   hypre_CSRMatrixRownnz(matrix)         = NULL;
   hypre_CSRMatrixNumRows(matrix)        = num_rows;
   hypre_CSRMatrixNumCols(matrix)        = num_cols;
   hypre_CSRMatrixNumNonzeros(matrix)    = num_nonzeros;
   hypre_CSRMatrixMemoryLocation(matrix) = HYPRE_MEMORY_SHARED; /* HYPRE_MEMORY_UNSET; */ /* TODO !!! */

   /* set defaults */
   hypre_CSRMatrixOwnsData(matrix)  = 1;
   hypre_CSRMatrixNumRownnz(matrix) = num_rows;

#ifdef HYPRE_BIGINT
   matrix->i_short=NULL;
   matrix->j_short=NULL;
#endif
   return matrix;
}

/*--------------------------------------------------------------------------
 * hypre_CSRMatrixDestroy
 *--------------------------------------------------------------------------*/

HYPRE_Int
hypre_CSRMatrixDestroy( hypre_CSRMatrix *matrix )
{
   HYPRE_Int ierr = 0;

   if (matrix)
   {
      HYPRE_Int memory_location = hypre_CSRMatrixMemoryLocation(matrix);

      hypre_TFree(hypre_CSRMatrixI(matrix),      memory_location);
      hypre_TFree(hypre_CSRMatrixRownnz(matrix), memory_location);

      if ( hypre_CSRMatrixOwnsData(matrix) )
      {
         hypre_TFree(hypre_CSRMatrixData(matrix), memory_location);
         hypre_TFree(hypre_CSRMatrixJ(matrix),    memory_location);
         hypre_TFree(hypre_CSRMatrixBigJ(matrix), memory_location);
      }

      hypre_TFree(matrix, HYPRE_MEMORY_HOST);
   }

   return ierr;
}

/*--------------------------------------------------------------------------
 * hypre_CSRMatrixInitialize
 *--------------------------------------------------------------------------*/

HYPRE_Int
hypre_CSRMatrixInitialize_v2( hypre_CSRMatrix *matrix, HYPRE_Int bigInit, HYPRE_Int memory_location )
{
   HYPRE_Int  num_rows     = hypre_CSRMatrixNumRows(matrix);
   HYPRE_Int  num_nonzeros = hypre_CSRMatrixNumNonzeros(matrix);
   /* HYPRE_Int  num_rownnz = hypre_CSRMatrixNumRownnz(matrix); */

   HYPRE_Int ierr = 0;

   hypre_CSRMatrixMemoryLocation(matrix) = memory_location;

   /* Caveat: for pre-existing i, j, data, their memory location must be guaranteed to be consistent with `memory_location'
    * Otherwise, mismatches will exist and problems will be encountered when being used, and freed */

   if ( !hypre_CSRMatrixData(matrix) && num_nonzeros )
   {
      hypre_CSRMatrixData(matrix) = hypre_CTAlloc(HYPRE_Complex, num_nonzeros, memory_location);
   }
   /*
   else
   {
     //if (PointerAttributes(hypre_CSRMatrixData(matrix))==HYPRE_HOST_POINTER) printf("MATREIX INITIAL WITH JHOST DATA\n");
   }
   */

   if ( !hypre_CSRMatrixI(matrix) )
   {
      hypre_CSRMatrixI(matrix) = hypre_CTAlloc(HYPRE_Int, num_rows + 1, memory_location);
   }

   /*
   if (!hypre_CSRMatrixRownnz(matrix))
   {
      hypre_CSRMatrixRownnz(matrix) = hypre_CTAlloc(HYPRE_Int,  num_rownnz, memory_location);
   }
   */

   if (bigInit)
   {
      if ( !hypre_CSRMatrixBigJ(matrix) && num_nonzeros )
      {
         hypre_CSRMatrixBigJ(matrix) = hypre_CTAlloc(HYPRE_BigInt, num_nonzeros, memory_location);
      }
   }
   else
   {
      if ( !hypre_CSRMatrixJ(matrix) && num_nonzeros )
      {
         hypre_CSRMatrixJ(matrix) = hypre_CTAlloc(HYPRE_Int, num_nonzeros, memory_location);
      }
   }

   return ierr;
}

HYPRE_Int
hypre_CSRMatrixInitialize( hypre_CSRMatrix *matrix )
{
   HYPRE_Int ierr;

   ierr = hypre_CSRMatrixInitialize_v2( matrix, 0, hypre_CSRMatrixMemoryLocation(matrix) );

   return ierr;
}

/*--------------------------------------------------------------------------
 * hypre_CSRMatrixBigInitialize
 *--------------------------------------------------------------------------*/

HYPRE_Int
hypre_CSRMatrixBigInitialize( hypre_CSRMatrix *matrix )
{
   HYPRE_Int ierr;

   ierr = hypre_CSRMatrixInitialize_v2( matrix, 1, hypre_CSRMatrixMemoryLocation(matrix) );

   return ierr;
}

/*--------------------------------------------------------------------------
 * hypre_CSRMatrixBigJtoJ
 *--------------------------------------------------------------------------*/

HYPRE_Int
hypre_CSRMatrixBigJtoJ( hypre_CSRMatrix *matrix )
{
   HYPRE_Int     num_nonzeros = hypre_CSRMatrixNumNonzeros(matrix);
   HYPRE_BigInt *matrix_big_j = hypre_CSRMatrixBigJ(matrix);
   HYPRE_Int    *matrix_j = NULL;
   HYPRE_Int     i;
/*   HYPRE_Int  num_rownnz = hypre_CSRMatrixNumRownnz(matrix); */

   HYPRE_Int  ierr=0;

   if (num_nonzeros && matrix_big_j )
   {
      matrix_j = hypre_CTAlloc(HYPRE_Int, num_nonzeros, HYPRE_MEMORY_SHARED);
      for (i=0; i < num_nonzeros; i++)
      {
         matrix_j[i] = (HYPRE_Int) matrix_big_j[i];
      }
      hypre_CSRMatrixJ(matrix) = matrix_j;
      hypre_TFree(matrix_big_j, HYPRE_MEMORY_SHARED);
      hypre_CSRMatrixBigJ(matrix) = NULL;
   }

   return ierr;
}

/*--------------------------------------------------------------------------
 * hypre_CSRMatrixJtoBigJ
 *--------------------------------------------------------------------------*/

HYPRE_Int
hypre_CSRMatrixJtoBigJ( hypre_CSRMatrix *matrix )
{
   HYPRE_Int     num_nonzeros = hypre_CSRMatrixNumNonzeros(matrix);
   HYPRE_Int    *matrix_j = hypre_CSRMatrixJ(matrix);
   HYPRE_BigInt *matrix_big_j = NULL;
   HYPRE_Int  i;
/*   HYPRE_Int  num_rownnz = hypre_CSRMatrixNumRownnz(matrix); */

   HYPRE_Int  ierr=0;

   if (num_nonzeros && matrix_j )
   {
      matrix_big_j = hypre_CTAlloc(HYPRE_BigInt,  num_nonzeros, HYPRE_MEMORY_SHARED);
      for (i=0; i < num_nonzeros; i++)
      {
         matrix_big_j[i] = (HYPRE_BigInt) matrix_j[i];
      }
      hypre_CSRMatrixBigJ(matrix) = matrix_big_j;
      hypre_TFree(matrix_j, HYPRE_MEMORY_SHARED);
      hypre_CSRMatrixJ(matrix) = NULL;
   }

   return ierr;
}

/*--------------------------------------------------------------------------
 * hypre_CSRMatrixSetDataOwner
 *--------------------------------------------------------------------------*/

HYPRE_Int
hypre_CSRMatrixSetDataOwner( hypre_CSRMatrix *matrix,
                             HYPRE_Int        owns_data )
{
   HYPRE_Int ierr = 0;

   hypre_CSRMatrixOwnsData(matrix) = owns_data;

   return ierr;
}

/*--------------------------------------------------------------------------
 * hypre_CSRMatrixSetRownnz
 *
 * function to set the substructure rownnz and num_rowsnnz inside the CSRMatrix
 * it needs the A_i substructure of CSRMatrix to find the nonzero rows.
 * It runs after the create CSR and when A_i is known..It does not check for
 * the existence of A_i or of the CSR matrix.
 *--------------------------------------------------------------------------*/

HYPRE_Int
hypre_CSRMatrixSetRownnz( hypre_CSRMatrix *matrix )
{
   HYPRE_Int  ierr = 0;
   HYPRE_Int  num_rows = hypre_CSRMatrixNumRows(matrix);
   HYPRE_Int  *A_i = hypre_CSRMatrixI(matrix);
   HYPRE_Int  *Arownnz;

   HYPRE_Int i, adiag;
   HYPRE_Int irownnz = 0;

   for (i = 0; i < num_rows; i++)
   {
      adiag = A_i[i+1] - A_i[i];
      if(adiag > 0)
      {
         irownnz++;
      }
   }

   hypre_CSRMatrixNumRownnz(matrix) = irownnz;

   if (irownnz == 0 || irownnz == num_rows)
   {
      hypre_CSRMatrixRownnz(matrix) = NULL;
   }
   else
   {
      Arownnz = hypre_CTAlloc(HYPRE_Int,  irownnz, HYPRE_MEMORY_SHARED);
      irownnz = 0;
      for (i = 0; i < num_rows; i++)
      {
         adiag = A_i[i+1] - A_i[i];
         if(adiag > 0)
         {
            Arownnz[irownnz++] = i;
         }
      }
      hypre_CSRMatrixRownnz(matrix) = Arownnz;
   }

   return ierr;
}

/*--------------------------------------------------------------------------
 * hypre_CSRMatrixRead
 *--------------------------------------------------------------------------*/

hypre_CSRMatrix *
hypre_CSRMatrixRead( char *file_name )
{
   hypre_CSRMatrix  *matrix;

   FILE    *fp;

   HYPRE_Complex *matrix_data;
   HYPRE_Int     *matrix_i;
   HYPRE_Int     *matrix_j;
   HYPRE_Int      num_rows;
   HYPRE_Int      num_nonzeros;
   HYPRE_Int      max_col = 0;

   HYPRE_Int      file_base = 1;

   HYPRE_Int      j;

   /*----------------------------------------------------------
    * Read in the data
    *----------------------------------------------------------*/
   fp = fopen(file_name, "r");

   hypre_fscanf(fp, "%d", &num_rows);

   matrix_i = hypre_CTAlloc(HYPRE_Int,  num_rows + 1, HYPRE_MEMORY_SHARED);
   for (j = 0; j < num_rows+1; j++)
   {
      hypre_fscanf(fp, "%d", &matrix_i[j]);
      matrix_i[j] -= file_base;
   }

   num_nonzeros = matrix_i[num_rows];

   matrix = hypre_CSRMatrixCreate(num_rows, num_rows, matrix_i[num_rows]);
   hypre_CSRMatrixI(matrix) = matrix_i;
   hypre_CSRMatrixInitialize(matrix);

   matrix_j = hypre_CSRMatrixJ(matrix);
   for (j = 0; j < num_nonzeros; j++)
   {
      hypre_fscanf(fp, "%d", &matrix_j[j]);
      matrix_j[j] -= file_base;

      if (matrix_j[j] > max_col)
      {
         max_col = matrix_j[j];
      }
   }

   matrix_data = hypre_CSRMatrixData(matrix);
   for (j = 0; j < matrix_i[num_rows]; j++)
   {
      hypre_fscanf(fp, "%le", &matrix_data[j]);
   }

   fclose(fp);

   hypre_CSRMatrixNumNonzeros(matrix) = num_nonzeros;
   hypre_CSRMatrixNumCols(matrix) = ++max_col;

   return matrix;
}

/*--------------------------------------------------------------------------
 * hypre_CSRMatrixPrint
 *--------------------------------------------------------------------------*/

HYPRE_Int
hypre_CSRMatrixPrint( hypre_CSRMatrix *matrix,
                      const char      *file_name )
{
   FILE    *fp;

   HYPRE_Complex *matrix_data;
   HYPRE_Int     *matrix_i;
   HYPRE_Int     *matrix_j;
   HYPRE_Int      num_rows;

   HYPRE_Int      file_base = 1;

   HYPRE_Int      j;

   HYPRE_Int      ierr = 0;

   /*----------------------------------------------------------
    * Print the matrix data
    *----------------------------------------------------------*/

   matrix_data = hypre_CSRMatrixData(matrix);
   matrix_i    = hypre_CSRMatrixI(matrix);
   matrix_j    = hypre_CSRMatrixJ(matrix);
   num_rows    = hypre_CSRMatrixNumRows(matrix);

   fp = fopen(file_name, "w");

   hypre_fprintf(fp, "%d\n", num_rows);

   for (j = 0; j <= num_rows; j++)
   {
      hypre_fprintf(fp, "%d\n", matrix_i[j] + file_base);
   }

   for (j = 0; j < matrix_i[num_rows]; j++)
   {
      hypre_fprintf(fp, "%d\n", matrix_j[j] + file_base);
   }

   if (matrix_data)
   {
      for (j = 0; j < matrix_i[num_rows]; j++)
      {
#ifdef HYPRE_COMPLEX
         hypre_fprintf(fp, "%.14e , %.14e\n",
                       hypre_creal(matrix_data[j]), hypre_cimag(matrix_data[j]));
#else
         hypre_fprintf(fp, "%.14e\n", matrix_data[j]);
#endif
      }
   }
   else
   {
      hypre_fprintf(fp, "Warning: No matrix data!\n");
   }

   fclose(fp);

   return ierr;
}

HYPRE_Int
hypre_CSRMatrixPrintMM( hypre_CSRMatrix *matrix,
                        HYPRE_Int        basei,
                        HYPRE_Int        basej,
                        HYPRE_Int        trans,
                        const char      *file_name )
{
   FILE    *fp;

   HYPRE_Complex *matrix_data;
   HYPRE_Int     *matrix_i;
   HYPRE_Int     *matrix_j;
   HYPRE_Int      num_rows, num_cols;

   /* HYPRE_Int      file_base = 1; */

   HYPRE_Int      j,k;

   HYPRE_Int      ierr = 0;

   /*----------------------------------------------------------
    * Print the matrix data
    *----------------------------------------------------------*/

   matrix_data = hypre_CSRMatrixData(matrix);
   matrix_i    = hypre_CSRMatrixI(matrix);
   matrix_j    = hypre_CSRMatrixJ(matrix);
   num_rows    = hypre_CSRMatrixNumRows(matrix);
   num_cols    = hypre_CSRMatrixNumCols(matrix);

   if (file_name)
   {
      fp = fopen(file_name, "w");
   }
   else
   {
      fp = stdout;
   }

   hypre_fprintf(fp, "%%%%MatrixMarket matrix coordinate real general\n");

   hypre_assert(matrix_i[num_rows] == hypre_CSRMatrixNumNonzeros(matrix));

   if (!trans)
   {
      hypre_fprintf(fp, "%d %d %d\n", num_rows, num_cols, hypre_CSRMatrixNumNonzeros(matrix));
   }
   else
   {
      hypre_fprintf(fp, "%d %d %d\n", num_cols, num_rows, hypre_CSRMatrixNumNonzeros(matrix));
   }

   for (j = 0; j < num_rows; j++)
   {
      for (k = matrix_i[j]; k < matrix_i[j+1]; k++)
      {
         if (!trans)
         {
            hypre_fprintf(fp, "%d %d %.15e\n", j + basei, matrix_j[k] + basej, matrix_data[k]);
         }
         else
         {
            hypre_fprintf(fp, "%d %d %.15e\n", matrix_j[k] + basej, j + basei, matrix_data[k]);
         }
      }
   }

   if (file_name)
   {
      fclose(fp);
   }

   return ierr;
}

HYPRE_Int
hypre_CSRMatrixPrint2( hypre_CSRMatrix *matrix,
                       const char      *file_name )
{
   return hypre_CSRMatrixPrintMM(matrix, 0, 0, 0, file_name);
}

/*--------------------------------------------------------------------------
 * hypre_CSRMatrixPrintHB: print a CSRMatrix in Harwell-Boeing format
 *--------------------------------------------------------------------------*/

HYPRE_Int
hypre_CSRMatrixPrintHB( hypre_CSRMatrix *matrix_input,
                        char            *file_name )
{
   FILE            *fp;
   hypre_CSRMatrix *matrix;
   HYPRE_Complex   *matrix_data;
   HYPRE_Int       *matrix_i;
   HYPRE_Int       *matrix_j;
   HYPRE_Int        num_rows;
   HYPRE_Int        file_base = 1;
   HYPRE_Int        j, totcrd, ptrcrd, indcrd, valcrd, rhscrd;
   HYPRE_Int        ierr = 0;

   /*----------------------------------------------------------
    * Print the matrix data
    *----------------------------------------------------------*/

   /* First transpose the input matrix, since HB is in CSC format */
   hypre_CSRMatrixTranspose(matrix_input, &matrix, 1);

   matrix_data = hypre_CSRMatrixData(matrix);
   matrix_i    = hypre_CSRMatrixI(matrix);
   matrix_j    = hypre_CSRMatrixJ(matrix);
   num_rows    = hypre_CSRMatrixNumRows(matrix);

   fp = fopen(file_name, "w");

   hypre_fprintf(fp, "%-70s  Key     \n", "Title");
   ptrcrd = num_rows;
   indcrd = matrix_i[num_rows];
   valcrd = matrix_i[num_rows];
   rhscrd = 0;
   totcrd = ptrcrd + indcrd + valcrd + rhscrd;
   hypre_fprintf (fp, "%14d%14d%14d%14d%14d\n",
                  totcrd, ptrcrd, indcrd, valcrd, rhscrd);
   hypre_fprintf (fp, "%-14s%14i%14i%14i%14i\n", "RUA",
                  num_rows, num_rows, valcrd, 0);
   hypre_fprintf (fp, "%-16s%-16s%-16s%26s\n", "(1I8)", "(1I8)", "(1E16.8)", "");

   for (j = 0; j <= num_rows; j++)
   {
      hypre_fprintf(fp, "%8d\n", matrix_i[j] + file_base);
   }

   for (j = 0; j < matrix_i[num_rows]; j++)
   {
      hypre_fprintf(fp, "%8d\n", matrix_j[j] + file_base);
   }

   if (matrix_data)
   {
      for (j = 0; j < matrix_i[num_rows]; j++)
      {
#ifdef HYPRE_COMPLEX
         hypre_fprintf(fp, "%16.8e , %16.8e\n",
                       hypre_creal(matrix_data[j]), hypre_cimag(matrix_data[j]));
#else
         hypre_fprintf(fp, "%16.8e\n", matrix_data[j]);
#endif
      }
   }
   else
   {
      hypre_fprintf(fp, "Warning: No matrix data!\n");
   }

   fclose(fp);

   hypre_CSRMatrixDestroy(matrix);

   return ierr;
}

/*--------------------------------------------------------------------------
 * hypre_CSRMatrixCopy: copy A to B,
 * if copy_data = 0 only the structure of A is copied to B.
 * the routine does not check if the dimensions/sparsity of A and B match !!!
 *--------------------------------------------------------------------------*/

HYPRE_Int
hypre_CSRMatrixCopy( hypre_CSRMatrix *A, hypre_CSRMatrix *B, HYPRE_Int copy_data )
{
   HYPRE_Int ierr = 0;
   HYPRE_Int num_rows = hypre_CSRMatrixNumRows(A);
   HYPRE_Int num_nonzeros = hypre_CSRMatrixNumNonzeros(A);

   HYPRE_Int     *A_i    = hypre_CSRMatrixI(A);
   HYPRE_Int     *A_j    = hypre_CSRMatrixJ(A);
   HYPRE_BigInt  *A_bigj = hypre_CSRMatrixBigJ(A);
   HYPRE_Complex *A_data;

   HYPRE_Int     *B_i    = hypre_CSRMatrixI(B);
   HYPRE_Int     *B_j    = hypre_CSRMatrixJ(B);
   HYPRE_BigInt  *B_bigj = hypre_CSRMatrixBigJ(B);
   HYPRE_Complex *B_data;

   HYPRE_Int memlocA = hypre_CSRMatrixMemoryLocation(A);
   HYPRE_Int memlocB = hypre_CSRMatrixMemoryLocation(B);

   hypre_TMemcpy(B_i, A_i, HYPRE_Int, num_rows + 1, memlocB, memlocA);

   if (A_j && B_j)
   {
      hypre_TMemcpy(B_j, A_j, HYPRE_Int, num_nonzeros, memlocB, memlocA);
   }

   if (A_bigj && B_bigj)
   {
      hypre_TMemcpy(B_bigj, A_bigj, HYPRE_BigInt, num_nonzeros, memlocB, memlocA);
   }

   if (copy_data)
   {
      A_data = hypre_CSRMatrixData(A);
      B_data = hypre_CSRMatrixData(B);
      hypre_TMemcpy(B_data, A_data, HYPRE_Complex, num_nonzeros, memlocB, memlocA);
   }

   return ierr;
}

/*--------------------------------------------------------------------------
 * hypre_CSRMatrixClone
 * Creates and returns a new copy of the argument, A.
 * Copying is a deep copy in that no pointers are copied; new arrays are
 * created where necessary.
 *--------------------------------------------------------------------------*/

hypre_CSRMatrix*
hypre_CSRMatrixClone_v2( hypre_CSRMatrix *A, HYPRE_Int copy_data, HYPRE_Int memory_location )
{
   HYPRE_Int num_rows = hypre_CSRMatrixNumRows(A);
   HYPRE_Int num_cols = hypre_CSRMatrixNumCols(A);
   HYPRE_Int num_nonzeros = hypre_CSRMatrixNumNonzeros( A );

   hypre_CSRMatrix *B = hypre_CSRMatrixCreate( num_rows, num_cols, num_nonzeros );

   HYPRE_Int bigInit = hypre_CSRMatrixBigJ(A) != NULL;

   hypre_CSRMatrixInitialize_v2(B, bigInit, memory_location);

   hypre_CSRMatrixCopy(A, B, copy_data);

   return B;
}

hypre_CSRMatrix*
hypre_CSRMatrixClone( hypre_CSRMatrix *A, HYPRE_Int copy_data )
{
   return hypre_CSRMatrixClone_v2(A, copy_data, HYPRE_MEMORY_SHARED);
}

/*--------------------------------------------------------------------------
 * hypre_CSRMatrixUnion
 * Creates and returns a matrix whose elements are the union of those of A and B.
 * Data is not computed, only structural information is created.
 * A and B must have the same numbers of rows.
 * Nothing is done about Rownnz.
 *
 * If col_map_offd_A and col_map_offd_B are zero, A and B are expected to have
 * the same column indexing.  Otherwise, col_map_offd_A, col_map_offd_B should
 * be the arrays of that name from two ParCSRMatrices of which A and B are the
 * offd blocks.
 *
 * The algorithm can be expected to have reasonable efficiency only for very
 * sparse matrices (many rows, few nonzeros per row).
 * The nonzeros of a computed row are NOT necessarily in any particular order.
 *--------------------------------------------------------------------------*/

hypre_CSRMatrix*
hypre_CSRMatrixUnion( hypre_CSRMatrix *A,
                      hypre_CSRMatrix *B,
                      HYPRE_BigInt *col_map_offd_A,
                      HYPRE_BigInt *col_map_offd_B,
                      HYPRE_BigInt **col_map_offd_C )
{
   HYPRE_Int num_rows = hypre_CSRMatrixNumRows( A );
   HYPRE_Int num_cols_A = hypre_CSRMatrixNumCols( A );
   HYPRE_Int num_cols_B = hypre_CSRMatrixNumCols( B );
   HYPRE_Int num_cols;
   HYPRE_Int num_nonzeros;
   HYPRE_Int *A_i = hypre_CSRMatrixI(A);
   HYPRE_Int *A_j = hypre_CSRMatrixJ(A);
   HYPRE_Int *B_i = hypre_CSRMatrixI(B);
   HYPRE_Int *B_j = hypre_CSRMatrixJ(B);
   HYPRE_Int *C_i;
   HYPRE_Int *C_j;
   HYPRE_Int *jC = NULL;
   HYPRE_BigInt jBg, big_jA, big_jB;
   HYPRE_Int i, jA, jB;
   HYPRE_Int ma, mb, mc, ma_min, ma_max, match;
   hypre_CSRMatrix * C;

   hypre_assert( num_rows == hypre_CSRMatrixNumRows(B) );
   if ( col_map_offd_B ) hypre_assert( col_map_offd_A );
   if ( col_map_offd_A ) hypre_assert( col_map_offd_B );

   /* ==== First, go through the columns of A and B to count the columns of C. */
   if ( col_map_offd_A==0 )
   {  /* The matrices are diagonal blocks.
         Normally num_cols_A==num_cols_B, col_starts is the same, etc.
      */
      num_cols = hypre_max( num_cols_A, num_cols_B );
   }
   else
   {  /* The matrices are offdiagonal blocks. */
      jC = hypre_CTAlloc( HYPRE_Int,  num_cols_B , HYPRE_MEMORY_SHARED);
      num_cols = num_cols_A;  /* initialization; we'll compute the actual value */
      for ( jB=0; jB<num_cols_B; ++jB )
      {
         match = 0;
         jBg = col_map_offd_B[jB];
         for ( ma=0; ma<num_cols_A; ++ma )
         {
            if ( col_map_offd_A[ma]==jBg )
               match = 1;
         }
         if ( match==0 )
         {
            jC[jB] = num_cols;
            ++num_cols;
         }
      }
   }

   /* ==== If we're working on a ParCSRMatrix's offd block,
      make and load col_map_offd_C */
   if ( col_map_offd_A )
   {
      *col_map_offd_C = hypre_CTAlloc( HYPRE_BigInt,  num_cols , HYPRE_MEMORY_SHARED);
      for ( jA=0; jA<num_cols_A; ++jA )
         (*col_map_offd_C)[jA] = col_map_offd_A[jA];
      for ( jB=0; jB<num_cols_B; ++jB )
      {
         match = 0;
         jBg = col_map_offd_B[jB];
         for ( ma=0; ma<num_cols_A; ++ma )
         {
            if ( col_map_offd_A[ma]==jBg )
               match = 1;
         }
         if ( match==0 )
            (*col_map_offd_C)[ jC[jB] ] = jBg;
      }
   }


   /* ==== The first run through A and B is to count the number of nonzero elements,
      without HYPRE_Complex-counting duplicates.  Then we can create C. */
   num_nonzeros = hypre_CSRMatrixNumNonzeros(A);
   for ( i=0; i<num_rows; ++i )
   {
      ma_min = A_i[i];  ma_max = A_i[i+1];
      for ( mb=B_i[i]; mb<B_i[i+1]; ++mb )
      {
         jB = B_j[mb];
         if ( col_map_offd_B ) big_jB = col_map_offd_B[jB];
         match = 0;
         for ( ma=ma_min; ma<ma_max; ++ma )
         {
            jA = A_j[ma];
            if ( col_map_offd_A ) big_jA = col_map_offd_A[jA];
            if ( big_jB == big_jA )
            {
               match = 1;
               if( ma==ma_min ) ++ma_min;
               break;
            }
         }
         if ( match==0 )
            ++num_nonzeros;
      }
   }

   C = hypre_CSRMatrixCreate( num_rows, num_cols, num_nonzeros );
   hypre_CSRMatrixInitialize( C );


   /* ==== The second run through A and B is to pick out the column numbers
      for each row, and put them in C. */
   C_i = hypre_CSRMatrixI(C);
   C_i[0] = 0;
   C_j = hypre_CSRMatrixJ(C);
   mc = 0;
   for ( i=0; i<num_rows; ++i )
   {
      ma_min = A_i[i];  ma_max = A_i[i+1];
      for ( ma=ma_min; ma<ma_max; ++ma )
      {
         C_j[mc] = A_j[ma];
         ++mc;
      }
      for ( mb=B_i[i]; mb<B_i[i+1]; ++mb )
      {
         jB = B_j[mb];
         if ( col_map_offd_B ) big_jB = col_map_offd_B[jB];
         match = 0;
         for ( ma=ma_min; ma<ma_max; ++ma )
         {
            jA = A_j[ma];
            if ( col_map_offd_A ) big_jA = col_map_offd_A[jA];
            if ( big_jB == big_jA )
            {
               match = 1;
               if( ma==ma_min ) ++ma_min;
               break;
            }
         }
         if ( match==0 )
         {
            if ( col_map_offd_A )
               C_j[mc] = jC[ B_j[mb] ];
            else
               C_j[mc] = B_j[mb];
            /* ... I don't know whether column indices are required to be in any
               particular order.  If so, we'll need to sort. */
            ++mc;
         }
      }
      C_i[i+1] = mc;
   }

   hypre_assert( mc == num_nonzeros );
   if (jC) hypre_TFree( jC , HYPRE_MEMORY_SHARED);

   return C;
}

static HYPRE_Int hypre_CSRMatrixGetLoadBalancedPartitionBoundary(hypre_CSRMatrix *A, HYPRE_Int idx)
{
   HYPRE_Int num_nonzerosA = hypre_CSRMatrixNumNonzeros(A);
   HYPRE_Int num_rowsA = hypre_CSRMatrixNumRows(A);
   HYPRE_Int *A_i = hypre_CSRMatrixI(A);

   HYPRE_Int num_threads = hypre_NumActiveThreads();

   HYPRE_Int nonzeros_per_thread = (num_nonzerosA + num_threads - 1)/num_threads;

   if (idx <= 0)
   {
      return 0;
   }
   else if (idx >= num_threads)
   {
      return num_rowsA;
   }
   else
   {
      return (HYPRE_Int)(hypre_LowerBound(A_i, A_i + num_rowsA, nonzeros_per_thread*idx) - A_i);
   }
}

HYPRE_Int hypre_CSRMatrixGetLoadBalancedPartitionBegin(hypre_CSRMatrix *A)
{
   return hypre_CSRMatrixGetLoadBalancedPartitionBoundary(A, hypre_GetThreadNum());
}

HYPRE_Int hypre_CSRMatrixGetLoadBalancedPartitionEnd(hypre_CSRMatrix *A)
{
   return hypre_CSRMatrixGetLoadBalancedPartitionBoundary(A, hypre_GetThreadNum() + 1);
}

HYPRE_Int
hypre_CSRMatrixPrefetch( hypre_CSRMatrix *A, HYPRE_Int to_location )
{
   HYPRE_Int      ierr = 0;

#ifdef HYPRE_USING_UNIFIED_MEMORY
   if (hypre_GetActualMemLocation(hypre_CSRMatrixMemoryLocation(A)) != HYPRE_MEMORY_SHARED)
   {
      return 1;
   }

   HYPRE_Complex *data = hypre_CSRMatrixData(A);
   HYPRE_Int     *ia   = hypre_CSRMatrixI(A);
   HYPRE_Int     *ja   = hypre_CSRMatrixJ(A);
   HYPRE_Int      nrow = hypre_CSRMatrixNumRows(A);
   HYPRE_Int      nnzA = hypre_CSRMatrixNumNonzeros(A);

   /* speical use of TMemcpy for prefetch */
   hypre_TMemcpy(data, data, HYPRE_Complex, nnzA, to_location, HYPRE_MEMORY_SHARED);
   //hypre_HandleCudaPrefetchStreamNum(hypre_handle) ++;
   hypre_TMemcpy(ia, ia, HYPRE_Int, nrow+1, to_location, HYPRE_MEMORY_SHARED);
   //hypre_HandleCudaPrefetchStreamNum(hypre_handle) ++;
   hypre_TMemcpy(ja, ja, HYPRE_Int, nnzA, to_location, HYPRE_MEMORY_SHARED);
#endif

   return ierr;
}

/*
HYPRE_Int hypre_CSRMatrixIsManaged(hypre_CSRMatrix *a)
{
}
*/

