Умножение разреженных матриц

/// <summary>
/// Умножение матрицы на матрицу
/// </summary>
/// <typeparam name="T"> тип элементов матриц </typeparam>
/// <param name="A"> 1-ая матрица </param>
/// <param name="B"> 2-ая матрица </param>
/// <returns> матрица, являющаяся результатом произведения обеих матриц </returns>
template <class T>
CSR_Matrix<T> operator * (CSR_Matrix<T> A, CSR_Matrix<T> B)
{
    assert(A.m() == B.n());

    // Число строк, столбцов, ненулевых элементов матрицы-результата

    int _N = A.n(), _M = B.m(), _NNZ = 0;

    int p = GridThreadsNum;

    // Распределяем разреженные вектор-строки матрицы по потокам

    vector <int> thread_pos(p + 1);

    for (int i = p - 1; i >= 0; --i)
        thread_pos[p - i] = thread_pos[p - i - 1] + (_N + i) / p;

    vector <T> values_A = A.values(), values_B = B.values();
    vector <int> cols_A = A.cols(), cols_B = B.cols();
    vector <int> rowptr_A = A.rowptr(), rowptr_B = B.rowptr();

    vector <int> rowptr(_N + 1);
    vector <vector <T>> rows(p, vector <T>(_M));

    /*
    // Сжатый массив строк для матрицы-результата
    vector <T> values;
    vector <int> cols, rowptr(1);

    // Вспомогательный массив, хранящий текущую рассматриваемую потоком
    // строку (в плотном формате, то есть не в разреженном)
    vector <vector <T>> rows(p, vector <T> (_M));

    // Преподсчет элементов массива rowptr
    // rowptr[i + 2] = k <=> в i-ой строке
    // матрицы-результата k ненулевых элементов

    omp_set_num_threads(p);
    #pragma omp parallel
    {
        int ThreadID = omp_get_thread_num();

        for (int i = thread_pos[ThreadID]; i < thread_pos[ThreadID + 1]; ++i)
        {
            for (int j = rowptr_A[i]; j < rowptr_A[i + 1]; ++j)
            {
                int col_A = cols_A[j]; T value_A = values_A[j];

                for (int k = rowptr_B[col_A]; k < rowptr_B[col_A + 1]; ++k)
                {
                    int col_B = cols_B[k]; T value_B = values_B[k];
                    rows[ThreadID][col_B] += value_A * value_B;
                }
            }

            for (int j = 0; j < _M; ++j)
                if (rows[ThreadID][j])
                {
                    values.push_back(rows[ThreadID][j]);
                    cols.push_back(j); ++_NNZ;

                    rows[ThreadID][j] = 0;
                }

            rowptr.push_back(_NNZ);
        }
    } */

    omp_set_num_threads(p);
    #pragma omp parallel
    {
        int ThreadID = omp_get_thread_num();

        for (int i = thread_pos[ThreadID]; i < thread_pos[ThreadID + 1]; ++i)
        {
            for (int j = rowptr_A[i]; j < rowptr_A[i + 1]; ++j)
            {
                int col_A = cols_A[j]; T value_A = values_A[j];

                for (int k = rowptr_B[col_A]; k < rowptr_B[col_A + 1]; ++k)
                {
                    int col_B = cols_B[k]; T value_B = values_B[k];
                    rows[ThreadID][col_B] += value_A * value_B;
                }
            }

            for (int j = 0; j < _M; ++j)
                if (rows[ThreadID][j])
                {
                    if (i == _N - 1) ++_NNZ;
                    else ++rowptr[i + 2];

                    rows[ThreadID][j] = 0;
                }
        }
    }

    for (int i = 3; i <= _N; ++i)
        rowptr[i] += rowptr[i - 1];

    _NNZ += rowptr[_N];

    vector <T> values(_NNZ);
    vector <int> cols(_NNZ);

    // Основной алгоритм: заполнение массивов values и
    // cols соответствующими значениями

    #pragma omp parallel
    {
        int ThreadID = omp_get_thread_num();

        for (int i = thread_pos[ThreadID]; i < thread_pos[ThreadID + 1]; ++i)
        {
            for (int j = rowptr_A[i]; j < rowptr_A[i + 1]; ++j)
            {
                int col_A = cols_A[j]; T value_A = values_A[j];

                for (int k = rowptr_B[col_A]; k < rowptr_B[col_A + 1]; ++k)
                {
                    int col_B = cols_B[k]; T value_B = values_B[k];
                    rows[ThreadID][col_B] += value_A * value_B;
                }
            }

            for (int j = 0; j < _M; ++j)
                if (rows[ThreadID][j])
                {
                    int pos = rowptr[i + 1];

                    values[pos] = rows[ThreadID][j];
                    cols[pos] = j;

                    rows[ThreadID][j] = 0;
                    ++rowptr[i + 1];
                }
        }
    }

    CSR_Matrix<T>* res = new CSR_Matrix<T>(_N, _M, _NNZ, values, cols, rowptr);
    return *res;
}