method to make this for-loop faster in C++ possibly with NVidia?

I want to make a C++ function faster. I am asking you about potential ways to do it.

I can use up to 32 OMP threads.

I can use an NVidia GPU.

A MWE for the function is:

#include <iostream>
#include <complex>
#include <cmath>

typedef std::numeric_limits<double> dbl;
#define _USE_MATH_DEFINES

#include <omp.h>

const std::complex<double> I(0.0, 1.0); // imaginary unit, I*I = -1
std::complex<double> zero_imag (0.0, 0.0);

const int N_rs = 1500;
const int l_max = 70;
const int lmax = 70;
const int N_thetas = l_max + 1;
const int N_phis = 2 * l_max + 2;
const int N_ps = 600;
const int nphi = 2 * l_max + 2;
const double sqrt_of_2_over_pi = sqrt( 2.0 / M_PI );


void rtop(std::complex<double> * Psi_outer_spec,
          std::complex<double> * Psi_outer_spec_plm,
          double * BJ,
          double * wrk,
          std::complex<double> * wrk2,
          double * ris_without_ends,
          double * r_primes_without_ends,
          double * weights_Lobatto_without_ends
         )
{


    int l, kk, kkk, m;
    long int idx, idxx, idxxx;

    // #pragma omp parallel for firstprivate (wrk2) private(l, kkk, idx, m, kk, idxx, idxxx) schedule(static)
    // #pragma omp target teams distribute parallel for firstprivate(wrk2) private(l, kkk, idx, m, kk, idxx, idxxx)
    for (int i = 0; i <= (N_ps - 1); i++) { // THIS IS THE BOTTLENECK !!!
       
        std::complex<double> sum1 = std::complex<double> (0.0, 0.0); // each thread creates a sum1 on its own

        for (l = 0; l <= lmax; l++) {

            for (kkk = 0; kkk <= (N_rs-1); kkk++) {
                idx = i * (N_rs*(l_max+1)) + kkk * (l_max+1) + l;
                wrk2[kkk] = pow(-I, l) * BJ[idx] * wrk[kkk];
            }

            for (m = 0; m <= (nphi-1); m++) {

                sum1 = zero_imag;
                for (kk = 0; kk <= (N_rs-1); kk++) {
                    idxx = kk * (N_thetas*N_phis) + l * N_phis + m;
                    sum1 += Psi_outer_spec[idxx] * wrk2[kk];

                }

                idxxx = i * (N_thetas*N_phis) + l * N_phis + m;
                Psi_outer_spec_plm[idxxx] = sum1 * sqrt_of_2_over_pi;
                                       
            }
            // END for m loop
        }
        // END for l loop
    }    
    // END for i loop
}

int main() {

    double * wrk = new double [N_rs];
    std::complex<double> * wrk2 = new std::complex<double> [N_rs];

    double * ris_without_ends = new double [N_rs];
    double * r_primes_without_ends = new double [N_rs];
    double * weights_Lobatto_without_ends = new double [N_rs];

    double * BJ = new double [N_ps * N_rs * (l_max+1)];

    std::complex<double> * Psi_outer_spec = new std::complex<double> [N_rs * N_thetas * N_phis];
    std::complex<double> * Psi_outer_spec_plm = new std::complex<double> [N_ps * N_thetas * N_phis];

    rtop(Psi_outer_spec, Psi_outer_spec_plm, BJ, wrk, wrk2, ris_without_ends, r_primes_without_ends, weights_Lobatto_without_ends);
   
    return 0;
}

The associated CMakeLists.txt is:

cmake_minimum_required(VERSION 3.0 FATAL_ERROR)

set(CMAKE_VERBOSE_MAKEFILE ON)

set(CMAKE_C_COMPILER "gcc")
set(CMAKE_CXX_COMPILER "g++")

project(trial)

set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -O3 -pedantic -Wall")

find_package(OpenMP)

add_executable(trial trial.cpp)

if(OpenMP_CXX_FOUND)
target_link_libraries(trial PUBLIC OpenMP::OpenMP_CXX)
endif()

set_property(TARGET trial PROPERTY CXX_STANDARD 17)

Compile as: $ cmake .. then $ cmake --build . --config Release.

My output is:

-- The C compiler identification is GNU 11.3.0
-- The CXX compiler identification is GNU 11.3.0
-- Detecting C compiler ABI info
-- Detecting C compiler ABI info - done
-- Check for working C compiler: /apps20/sw/eb/software/GCCcore/11.3.0/bin/gcc - skipped
-- Detecting C compile features
-- Detecting C compile features - done
-- Detecting CXX compiler ABI info
-- Detecting CXX compiler ABI info - done
-- Check for working CXX compiler: /apps20/sw/eb/software/GCCcore/11.3.0/bin/g++ - skipped
-- Detecting CXX compile features
-- Detecting CXX compile features - done
-- Found OpenMP_C: -fopenmp (found version "4.5")
-- Found OpenMP_CXX: -fopenmp (found version "4.5")
-- Found OpenMP: TRUE (found version "4.5")
-- Configuring done
-- Generating done
-- Build files have been written to: /work4/clf/ouatu/trial_for_SO/build 

Then for the build:

[ 50%] Building CXX object CMakeFiles/trial.dir/trial.cpp.o
[100%] Linking CXX executable trial
[100%] Built target trial

What I have tried:

• With OpenMP parallel for, I do get speedup.

• I fail with OpenMP GPU-offloading (it seems my compiler flags do not
  make the offloading possible). (these flags are hidden from the shown CMakeLists.txt for this MWE)

• I am open to any other suggestions.

For example, would rtop benefit from being a CUDA kernel? Is it hard to make it that way?

Thank you!

I suggest an OpenMP version with some optimizations and tweaks. A quick review of some changes and what to look out for:

The whole business with wrk2[kkk] = pow(-I, l) * ... is doubly redundant. For one, pow(-I, l) is an elegant but expensive way of expressing just 4 different values. Second, it's only used as a factor in a dot product. You can fold the whole thing into the final multiplication sum1 * sqrt_of_2_over_pi. That also allows wrk2 to be real-valued, which also turns the innermost loop from a complex-complex dot product to a complex-real dot product.

The multi-dimensional index calculations like idx = i * (N_rs*(l_max+1)) + kkk * (l_max+1) + l should be done following the Horner method to avoid redundant multiplications. More of a nitpick but also clearner. For example here idx = (i * N_rs + kkk) * (l_max+1) + l. While we're at it, be careful with your index variables. They are all int. Especially the 3-dimensional arrays can quickly grow to multiple GiB in size at which point you will experience integer overflows. Switch to std::ptrdiff_t if you fear that this may become a problem.

The order of iterations over BJ and Psi_outer_spec_plm aren't ideal. If possible, BJ should swap the two inner dimensions for better locality of data, which would also allow for vectorization of the loop initializing wrk2. Psi_outer_spec is even worse because you iterate along the outer dimension in your innermost loop. However, I assume this order was picked so that it is the same as with Psi_outer_spec_plm and for that it is good. In any case, this higher stride prevents vectorization.

I don't see a reason why you declare counter and index variables outside the scope where they are used. Even modern C standards allow declaring them inside for loops, let alone C++. For parallelization, you want to limit the number of shared or accidentally shared variables.

Speaking of shared data, as far as I can see, the only shared memory in which threads may overlap is the wrk2 array. That can simply be allocated per thread, which brings us to the final implementation.

#   pragma omp parallel
    {
        auto wrk2 = std::make_unique<double[]>(N_rs);
#       pragma omp for collapse(2) nowait
        for (int i = 0; i <= (N_ps - 1); i++) {
            for (int l = 0; l <= lmax; l++) {
                for (int kkk = 0; kkk <= (N_rs-1); kkk++) {
                    int idx = (i * N_rs + kkk) * (lmax + 1) + l;
                    wrk2[kkk] = BJ[idx] * wrk[kkk];
                }
                constexpr std::complex<double> I(0., 1.);
                std::complex<double> factor(-sqrt_of_2_over_pi);
                if(l & 1)
                    factor *= I;
                if(l & 2)
                    factor = -factor;
                for (int m = 0; m <= (N_phis-1); m++) {
                    std::complex<double> sum1;
                    for (int kk = 0; kk <= (N_rs-1); kk++) {
                        int idx = (kk * N_thetas + l) * N_phis + m;
                        sum1 += Psi_outer_spec[idx] * wrk2[kk];
                    }
                    int idx = (i * N_thetas + l) * N_phis + m;
                    Psi_outer_spec_plm[idx] = sum1 * factor;
                }
            }
        }
    }

Note how the usual pragma omp parallel for is split into an omp parallel and a separate omp for to allow for allocating the temporary memory. The collapse(2) means that both outer loops are parallelized.

Other things to consider:

• The inner dot product may be computed faster by an accelerated BLAS library or something similar. I think Eigen should work well here but one might need to coerce it a bit into working with this memory layout
• It looks somewhat like we could change the m loop into a matrix-vector product, which might solve some of our vectorization / memory access issues via a BLAS library
• Since you asked about compile options, -march=native or whatever baseline architecture you want should be worthwhile here. -mavx2 -mfma may be a good compromise to handle all relatively recent CPUs without specializing the binary too much

Edit: Matrix-vector product

Coming back to the idea to offload the loop over m into a matrix vector product, we have to reinterpret the slice of Psi_outer_spec that we use as a matrix. I choose a column-major matrix since I want to use Eigen3 for this step.

• The number of rows are N_phi (loop counter m)
• The number of columns are N_rs (loop counter kk)
• From one column to the next we have a stride / a.k.a. leading dimension of N_phi * N_theta
• The offset of the top left corner is l * N_phis

Assuming this is correct, we can map our arrays into Eigen vectors and matrices and let it handle the transposed accesses. This turns everything below the wrk2 initialization into this code

using MatrixMap = Eigen::Map<const Eigen::MatrixXcd,
        Eigen::Unaligned, Eigen::OuterStride<>>;
MatrixMap Psi_slice(
        Psi_outer_spec + l * N_phis /*top left corner*/,
        N_phis /*rows*/, N_rs /*cols*/,
        Eigen::OuterStride<>(N_phis * N_thetas));
const auto wrk2_mapped = Eigen::VectorXd::Map(wrk2.get(), N_rs);
auto Psi_plm_mapped = Eigen::VectorXcd::Map(
        Psi_outer_spec_plm + (i * N_thetas + l) * N_phis, N_phis);
Psi_plm_mapped.noalias() = Psi_slice * wrk2_mapped * factor;

And now this step obviously raises the question of whether we can turn the whole thing into a matrix-matrix product with some pre- or post-processing, which might take care of the entire parallelization and potential offloading to GPUs. And this is why I asked for a mathematical description instead of doing this wild goose chase through the code

Edit 2: Matrix-matrix product

It is indeed possible to rewrite it as a matrix-matrix product. The trick is the observation that Psi_outer_spec is independent from i. Therefore if we switch the two outer loops, we can compute all values for one l over all i in one operation.

While doing so, I switch back to wrk2 being complex and including the factor. This technically requires more compute time and memory but with a matrix-matrix product, you may want to dispatch to a BLAS backend, either directly for example via OpenBLAS, via Eigen's backends or even GPU acceleration such as CuBLAS. And for that you need a complex-complex multiplication.

Eigen::MatrixXcd wrk2mat(N_rs, N_ps);
for (int l = 0; l <= lmax; l++) {
    std::complex<double> factor(-sqrt_of_2_over_pi);
    if(l & 1)
        factor *= I;
    if(l & 2)
        factor = -factor;
#   pragma omp parallel for
    for (int i = 0; i <= N_ps - 1; i++) {
        for (int k = 0; k <= N_rs - 1; ++k) {
            int idx = (i * N_rs + k) * (lmax + 1) + l;
            wrk2mat(k, i) = BJ[idx] * wrk[k] * factor;
        }
    }
    using ConstMatrixMap = Eigen::Map<const Eigen::MatrixXcd,
            Eigen::Unaligned, Eigen::OuterStride<>>;
    ConstMatrixMap Psi_slice(
            Psi_outer_spec + l * N_phis /*top left corner*/,
            N_phis /*rows*/, N_rs /*cols*/,
            Eigen::OuterStride<>(N_phis * N_thetas));
    using MatrixMap = Eigen::Map<Eigen::MatrixXcd,
            Eigen::Unaligned, Eigen::OuterStride<>>;
    MatrixMap Psi_plm_mapped(
            Psi_outer_spec_plm + l * N_phis,
            N_phis, N_ps,
            Eigen::OuterStride<>((lmax + 1) * N_phis));
    Psi_plm_mapped.noalias() = Psi_slice * wrk2mat;
}

The matrix-matrix product should be internally parallelized as long as the matrices are large enough. If this is not always the case, you can wrap the whole thing into an runtime-optional parallel block. Roughly like this:

bool small_matrices = ...;
#pragma omp parallel if(small_matrices)
{
    Eigen::MatrixXcd wrk2mat(N_rs, N_ps);
#   pragma omp for nowait
    for (int l = 0; l <= lmax; l++) {
        ...
    }
}

Since OpenMP normally deactivates nested parallelization, this will automatically deactivate all the inner parallel sections and run them sequentially.

10x faster solution using GPU and CuBLAS. Also fix wrong result bug from current accepted answer

Test computer specs:

• Ryzen 2950X, quad channel RAM 2133MHz
• 2080ti (terrible at FP64)
• PCIE 3.0 x16 lanes
• Ubuntu 20.04, CUDA 12.1, Eigen 3.3.9
• IDK what backend Eigen is using because it's not my PC, I think it's MKL

The time cost for each implementation using this PC:

• Original code: ~75s
• Eigen code OMP_NUM_THREADS=1: 8s
• Eigen code OMP_NUM_THREADS=16: 1.1s
• CUDA + CuBLAS: 0.26s.

I measure each segment by using MyTimer and cudaDeviceSynchronize() after each CUDA call (not shown in the code below). The total run time when not measuring is slightly lower, because some transfer/compute is being overlapped (I think).

• BJ, d_wrk, Psi_outer_spec Host to Device transfer: ~50ms
• compute_wrk2mat: ~15ms
• cublasZgemmStridedBatched: ~180ms
• Psi_outer_spec_plm Device to Host transfer: ~10ms

Predicting 2080ti vs A100 performance:

• FP32: 13.45 TFLOP vs 19.75 TFLOP
• FP64: 0.42 TFLOP vs 9.75 TFLOP
• Memory bandwidth: 616 GB/s vs 1555 GB/s
• PCIE 4.0 is 2x faster than 3.0

So I predict the A100 will run this in (50ms + 10ms) * 0.5 (pcie_transfer) + 15ms * 616 / 1555 (memory_bound_kernel) + 180ms / 10 (fp64_compute_bound_kernel ~~50ms . If anyone has free time, please run a benchmark because I'm also curious.

Finally, the code. First, the accepted answer (super interesting and well analyzed, btw) has a small mistake that causes it to output different results compared to your original code in the question.

    std::complex<double> factor(-sqrt_of_2_over_pi);
    if(l & 1)
        factor *= I;
    if(l & 2)
        factor = -factor;

it should be

    std::complex<double> factor(sqrt_of_2_over_pi);
    if(l & 1)
        factor *= -I;
    if(l & 2)
        factor = -factor;

The code below is a full runnable program that benchmark + check correctness of all 3 implementations in this post. The GPU version uses a class RtopCalculator object instead of a function, so that it can reuse resources (device arrays) instead of allocate/free each time the function is called. The class is RAII, and will safely free resources when destructed.

I dont use pinned memory because the expected use case is passing array from Numpy (OP says so in comment). So using normal memory is a better practical benchmark.

#include <iostream>
#include <complex>
#include <cmath>
#include <random>
#include <chrono>
#include <fstream>
#include <Eigen/Dense>
#include <omp.h>
#include <cuda_runtime.h>
#include <cublas_v2.h>
#include <thrust/execution_policy.h>
#include <thrust/device_vector.h>
#include <thrust/complex.h>
//------------
// Utility stuffs to test
// Helper function to check cuBLAS status
#define CUBLAS_CHECK(err)                                                                          \
do {                                                                                           \
cublasStatus_t err_ = (err);                                                               \
if (err_ != CUBLAS_STATUS_SUCCESS) {                                                       \
printf("cublas error %d at %s:%d\n", err_, __FILE__, __LINE__);                        \
throw std::runtime_error("cublas error");                                              \
}                                                                                          \
} while (0)
inline void gpuAssert(cudaError_t code, const char *file, int line, bool printing = false)
{   
if (code != cudaSuccess)
{
std::string mess = std::string("GPUassert: ") + std::string(cudaGetErrorString(code)) 
+ " " + std::string(file) + " " + std::to_string(line);
if (printing) std::cout << mess << std::endl;
throw std::runtime_error(mess.c_str());
}
auto lastError = cudaGetLastError();
if (lastError != cudaSuccess)
{
std::string mess = std::string("GPUassert: ") + std::string(cudaGetErrorString(lastError)) 
+ " " + std::string(file) + " " + std::to_string(line);
std::cout << "UNDETECTED_ERROR " << mess << std::endl;
throw std::runtime_error(mess.c_str());
}
}
// CUDA API error checking
#define CUDA_CHECK(ans) { gpuAssert((ans), __FILE__, __LINE__, true); }
#define CUDA_CHECK_NOLOG(ans) { gpuAssert((ans), __FILE__, __LINE__); }
double eps_ = 1e-6;
template <typename num_t>
num_t rcmp(num_t a, num_t b, num_t eps = eps_) {
if (std::isnan(a) && std::isnan(b)) return 0;
if (std::isnan(a + b)) return NAN;
num_t t = (a - b) / (std::max(std::abs(a), std::abs(b)) + 1e-18);
return t < -eps ? -1 : +eps < t;
}
class MyTimer {
std::chrono::time_point<std::chrono::system_clock> start;
public:
void startCounter() {
start = std::chrono::system_clock::now();
}
int64_t getCounterNs() {
return std::chrono::duration_cast<std::chrono::nanoseconds>(std::chrono::system_clock::now() - start).count();
}
int64_t getCounterMs() {
return std::chrono::duration_cast<std::chrono::milliseconds>(std::chrono::system_clock::now() - start).count();
}
double getCounterMsPrecise() {
return std::chrono::duration_cast<std::chrono::nanoseconds>(std::chrono::system_clock::now() - start).count()
/ 1000000.0;
}
};
//----------------
//----------------
//----------------
std::mt19937 rander(42);
double myrand() {
return double(rander() % 10000) / (rander() % 10000 + 1);
}
const std::complex<double> I(0.0, 1.0); // imaginary unit, I*I = -1
std::complex<double> zero_imag (0.0, 0.0);
const double sqrt_of_2_over_pi = sqrt( 2.0 / M_PI   );
void GenData(
int64_t N_rs, int64_t l_max, int64_t N_ps,
std::complex<double>* Psi_outer_spec, double* BJ, double* wrk
)
{
int64_t N_thetas = l_max + 1;
int64_t N_phis = 2 * l_max + 2;
for (int64_t i = 0; i < N_rs; i++) wrk[i] = myrand();
for (int64_t i = 0; i < N_ps * N_rs * N_thetas; i++) BJ[i] = myrand();
for (int64_t i = 0; i < N_rs * N_thetas * N_phis; i++) {
Psi_outer_spec[i].real(myrand());
Psi_outer_spec[i].imag(myrand());
}
}
void rtop(
int64_t N_rs, int64_t l_max, int64_t N_ps,
const std::complex<double>* __restrict__ Psi_outer_spec,
std::complex<double>* __restrict__ Psi_outer_spec_plm,
const double* __restrict__ BJ,
const double* __restrict__ wrk,
std::complex<double>* __restrict__ wrk2,
double* __restrict__ ris_without_ends = nullptr,
double* __restrict__ r_primes_without_ends = nullptr,
double* __restrict__ weights_Lobatto_without_ends = nullptr
)
{        
int64_t N_thetas = l_max + 1;
int64_t N_phis = 2 * l_max + 2;
int64_t l, kk, kkk, m;
int64_t idx, idxx, idxxx;
for (int64_t i = 0; i <= (N_ps - 1); i++) { // THIS IS THE BOTTLENECK !!!
std::complex<double> sum1 = std::complex<double> (0.0, 0.0); // each thread creates a sum1 on its own
for (l = 0; l <= l_max; l++) {
for (kkk = 0; kkk <= (N_rs-1); kkk++) {
idx = i * (N_rs*(l_max+1)) + kkk * (l_max+1) + l;
wrk2[kkk] = pow(-I, l) * BJ[idx] * wrk[kkk];
}
for (m = 0; m <= (N_phis-1); m++) {
sum1 = zero_imag;
for (kk = 0; kk <= (N_rs-1); kk++) {
idxx = kk * (N_thetas*N_phis) + l * N_phis + m;
sum1 += Psi_outer_spec[idxx] * wrk2[kk];
}
idxxx = i * (N_thetas*N_phis) + l * N_phis + m;
Psi_outer_spec_plm[idxxx] = sum1 * sqrt_of_2_over_pi;
}
// END for m loop
}
// END for l loop
}    
// END for i loop
}
void rtop_eigen(
int64_t N_rs, int64_t l_max, int64_t N_ps,
const std::complex<double>* __restrict__ Psi_outer_spec,
std::complex<double>* __restrict__ Psi_outer_spec_plm,
const double* __restrict__ BJ,
const double* __restrict__ wrk,
std::complex<double>* __restrict__ wrk2,
double* __restrict__ ris_without_ends = nullptr,
double* __restrict__ r_primes_without_ends = nullptr,
double* __restrict__ weights_Lobatto_without_ends = nullptr
)
{
int64_t N_thetas = l_max + 1;
int64_t N_phis = 2 * l_max + 2;
Eigen::MatrixXcd wrk2mat(N_rs, N_ps);
for (int64_t l = 0; l <= l_max; l++) {
std::complex<double> factor(sqrt_of_2_over_pi);
if(l & 1)
factor *= -I;
if(l & 2)
factor = -factor;
#   pragma omp parallel for
for (int64_t i = 0; i <= N_ps - 1; i++) {
for (int64_t k = 0; k <= N_rs - 1; ++k) {
int64_t idx = (i * N_rs + k) * (l_max + 1) + l;
wrk2mat(k, i) = BJ[idx] * wrk[k] * factor;
}
}
using ConstMatrixMap = Eigen::Map<const Eigen::MatrixXcd,
Eigen::Unaligned, Eigen::OuterStride<>>;
ConstMatrixMap Psi_slice(
Psi_outer_spec + l * N_phis /*top left corner*/,
N_phis /*rows*/, N_rs /*cols*/,
Eigen::OuterStride<>(N_phis * N_thetas));
using MatrixMap = Eigen::Map<Eigen::MatrixXcd,
Eigen::Unaligned, Eigen::OuterStride<>>;
MatrixMap Psi_plm_mapped(
Psi_outer_spec_plm + l * N_phis,
N_phis, N_ps,
Eigen::OuterStride<>((l_max + 1) * N_phis));
Psi_plm_mapped.noalias() = Psi_slice * wrk2mat;
}
}
namespace {
__global__
void compute_wrk2mat(
int64_t N_ps, int64_t N_rs, int64_t l_max,
const double* __restrict__ BJ,
const double* __restrict__ wrk,
cuDoubleComplex* __restrict__ wrk2mat
)
{
constexpr double sqrt_of_2_over_pi = 0.79788456080286535587989;
for (int l = 0; l <= (int)l_max; l++) {
cuDoubleComplex* wrk2mat_offset = wrk2mat + l * N_rs * N_ps;
double factor_real = sqrt_of_2_over_pi;
double factor_imag = 0;
if (l & 1) {
double temp_real = factor_real;
factor_real = factor_imag;
factor_imag = -temp_real;
}
if (l & 2) {
factor_real = -factor_real;
factor_imag = -factor_imag;
}
for (int i = blockIdx.x; i < N_ps; i += gridDim.x)
for (int k = threadIdx.x; k < N_rs; k += blockDim.x) {
int64_t idx = (i * N_rs + k) * (l_max + 1) + l;
wrk2mat_offset[k + i * N_rs].x = BJ[idx] * wrk[k] * factor_real;
wrk2mat_offset[k + i * N_rs].y = BJ[idx] * wrk[k] * factor_imag;
}
}
}
}
class RtopCalculator {
private:
cudaStream_t main_stream_;
cudaStream_t side_stream_;
cublasHandle_t cublas_handle_;
int64_t N_rs_;
int64_t l_max_;
int64_t N_thetas_;
int64_t N_phis_;
int64_t N_ps_;
thrust::device_vector<double> d_BJ_;
thrust::device_vector<double> d_wrk_;
thrust::device_vector<thrust::complex<double>> d_wrk2mat_;
thrust::device_vector<thrust::complex<double>> d_Psi_outer_spec_;
thrust::device_vector<thrust::complex<double>> d_Psi_outer_spec_plm_;
void allocate_internal() {
d_BJ_.resize(N_rs_ * N_ps_ * N_thetas_);
d_wrk_.resize(N_rs_);
d_wrk2mat_.resize(N_thetas_ * N_rs_ * N_ps_);
d_Psi_outer_spec_.resize(N_rs_ * N_thetas_ * N_phis_);
d_Psi_outer_spec_plm_.resize(N_ps_ * N_thetas_ * N_phis_);
}
public:
RtopCalculator() {
cudaStreamCreate(&main_stream_);
cudaStreamCreate(&side_stream_);
cublasCreate(&cublas_handle_);
cublasSetStream(cublas_handle_, main_stream_);
}
~RtopCalculator() {
cudaStreamDestroy(main_stream_);
cudaStreamDestroy(side_stream_);
cublasDestroy(cublas_handle_);
}
void allocate(int64_t N_rs, int64_t l_max, int64_t N_ps) {
N_rs_ = N_rs;
l_max_ = l_max;
N_thetas_ = l_max + 1;
N_phis_ = 2 * l_max + 2;
N_ps_ = N_ps;
allocate_internal();
}
void compute(
int64_t N_rs, int64_t l_max, int64_t N_ps,
const double* __restrict__ BJ,
const double* __restrict__ wrk,
const double* __restrict__ Psi_outer_spec, // std::complex<double>
double* __restrict__ Psi_outer_spec_plm,
double* __restrict__ ris_without_ends = nullptr,
double* __restrict__ r_primes_without_ends = nullptr,
double* __restrict__ weights_Lobatto_without_ends = nullptr
)
{
allocate(N_rs, l_max, N_ps);        
int64_t N_phis = N_phis_;
int64_t N_thetas = N_thetas_;
double* d_BJ = thrust::raw_pointer_cast(d_BJ_.data());
double* d_wrk = thrust::raw_pointer_cast(d_wrk_.data());
thrust::complex<double>* d_wrk2mat = thrust::raw_pointer_cast(d_wrk2mat_.data());
thrust::complex<double>* d_Psi_outer_spec = thrust::raw_pointer_cast(d_Psi_outer_spec_.data());
thrust::complex<double>* d_Psi_outer_spec_plm = thrust::raw_pointer_cast(d_Psi_outer_spec_plm_.data());
// the ordering of the next 4 statements are intended to interleave data transfer and compute
// Cost 1
cudaMemcpyAsync(d_BJ, BJ, N_rs * N_ps * (l_max + 1) * sizeof(double), cudaMemcpyHostToDevice, main_stream_);
cudaMemcpyAsync(d_wrk, wrk, N_rs * sizeof(double), cudaMemcpyHostToDevice, main_stream_);                
compute_wrk2mat<<<512, 256, 0, main_stream_>>>(
N_ps, N_rs, l_max, d_BJ, d_wrk, reinterpret_cast<cuDoubleComplex*>(d_wrk2mat)
);
cudaMemcpyAsync(
d_Psi_outer_spec, Psi_outer_spec,
N_rs * N_thetas * N_phis * sizeof(std::complex<double>),
cudaMemcpyHostToDevice, side_stream_
);
CUDA_CHECK(cudaStreamSynchronize(main_stream_));
CUDA_CHECK(cudaStreamSynchronize(side_stream_));
// Cost 2
int64_t M = N_phis;
int64_t K = N_rs;
int64_t N = N_ps;
int64_t lda = N_phis * N_thetas;
int64_t ldb = N_rs;
int64_t ldc = (l_max + 1) * N_phis;
cuDoubleComplex* d_A = reinterpret_cast<cuDoubleComplex*>(d_Psi_outer_spec);
cuDoubleComplex* d_B = reinterpret_cast<cuDoubleComplex*>(d_wrk2mat);
cuDoubleComplex* d_C = reinterpret_cast<cuDoubleComplex*>(d_Psi_outer_spec_plm);
int64_t strideA = N_phis;
int64_t strideB = N_rs * N_ps;
int64_t strideC = N_phis;
std::complex<double> alpha(1.0, 0.0);
std::complex<double> beta(0.0, 0.0);
CUBLAS_CHECK(cublasZgemmStridedBatched(
cublas_handle_,
CUBLAS_OP_N, CUBLAS_OP_N,
M, N, K,
(cuDoubleComplex*)&alpha,
d_A, lda, strideA,
d_B, ldb, strideB,
(cuDoubleComplex*)&beta,
d_C, ldc, strideC,
l_max + 1
));
// Cost 3
cudaMemcpyAsync(
Psi_outer_spec_plm, reinterpret_cast<double*>(d_Psi_outer_spec_plm),
N_ps * N_thetas * N_phis * sizeof(std::complex<double>),
cudaMemcpyDeviceToHost, main_stream_
);
CUDA_CHECK(cudaStreamSynchronize(main_stream_));
}
};
int main(int argc, char* argv[]) {
bool full_test = 0;
if (argc > 1) full_test = 1;
std::cout << "Full test = " << full_test << "\n";
const int64_t N_rs = 1500;
const int64_t l_max = 70;
const int64_t N_thetas = l_max + 1;
const int64_t N_phis = 2 * l_max + 2;
const int64_t N_ps = 600;
MyTimer timer;
double total_cost[3] = {0};
double* wrk = new double [N_rs];
std::complex<double>* wrk2 = new std::complex<double> [N_rs * N_ps];
double* BJ = new double [N_ps * N_rs * (l_max+1)];
std::complex<double>* Psi_outer_spec = new std::complex<double> [N_rs * N_thetas * N_phis];
std::complex<double>* Psi_outer_spec_plm_0 = new std::complex<double> [N_ps * N_thetas * N_phis];
std::complex<double>* Psi_outer_spec_plm_1 = new std::complex<double> [N_ps * N_thetas * N_phis];
std::complex<double>* Psi_outer_spec_plm_2 = new std::complex<double> [N_ps * N_thetas * N_phis];
RtopCalculator calculator;
calculator.allocate(N_rs, l_max, N_ps);
int ntest = 5;
int wrong = 0;
for (int t = 1; t <= ntest; t++) {
std::cout << "Start test " << t << "\n";
GenData(N_rs, l_max, N_ps, Psi_outer_spec, BJ, wrk);
if (full_test) {
timer.startCounter();        
rtop(N_rs, l_max, N_ps, Psi_outer_spec, Psi_outer_spec_plm_0, BJ, wrk, wrk2);
total_cost[0] += timer.getCounterMsPrecise();
}
timer.startCounter();
rtop_eigen(N_rs, l_max, N_ps, Psi_outer_spec, Psi_outer_spec_plm_1, BJ, wrk, wrk2);
total_cost[1] += timer.getCounterMsPrecise();
timer.startCounter();
calculator.compute(
N_rs, l_max, N_ps,
BJ,
wrk,
reinterpret_cast<double*>(Psi_outer_spec),
reinterpret_cast<double*>(Psi_outer_spec_plm_2)
);
total_cost[2] += timer.getCounterMsPrecise();
std::cout << "cost = " << total_cost[0] << " " << total_cost[1] << " " << total_cost[2] << "\n";
for (int i = 0; i < N_ps; i++) 
for (int l = 0; l < N_thetas; l++)
for (int m = 0; m < N_phis; m++) {
int64_t idx = i * (N_thetas * N_phis) + l * N_phis + m;
auto res0 = Psi_outer_spec_plm_0[idx];
auto res1 = Psi_outer_spec_plm_1[idx];
auto res2 = Psi_outer_spec_plm_2[idx];
if (full_test) {
if (rcmp(res0.real(), res1.real()) || rcmp(res0.imag(), res1.imag()) ||
rcmp(res0.real(), res2.real()) || rcmp(res0.imag(), res2.imag())
) {
std::cout << "Error at (i=" << i << ",l=" << l << ",m=" << m << "): " << res0 << "; " << res1 << " " << res2 << "\n";
wrong++;
if (wrong == 20) exit(1);
}
} else {
if (rcmp(res1.real(), res2.real()) || rcmp(res1.imag(), res2.imag())) {
std::cout << "Error at (i=" << i << ",l=" << l << ",m=" << m << "): " << res1 << "; " << res2 << "\n";                
wrong++;
if (wrong == 20) exit(1);
}
}
}
}
return 0;
}

To run and test, use run.sh.

threads=$1  # The custom variable (number of threads) passed as the first argument
if [ -z "$threads" ]; then
threads=8
fi
export OMP_NUM_THREADS=$threads
export OPENBLAS_NUM_THREADS=$threads
nvcc -o main quantum.cu -O3 -std=c++17 -lcudart -lcublas -Xcompiler -march=native -Xcompiler -fopenmp -arch=native
echo "Running with $1 threads"
time ./main $2

For example:

./run.sh 1 # test Eigen 1 thread vs GPU
./run.sh 16 # test Eigen 16 thread vs GPU
./run.sh 16 1 # test all 3 versions. Remember the original code is very slow.

It's possible to turn this into a shared library which can be imported into Python and use with Tensorflow/Torch. But it's out of the scope of the question, so you can send me an email.