[Python] Unstable timings of C++/CUDA function depending on how it is called from python...

I have a C++ library relying on CUDA and OpenMP to perform a matrix-vector operation on the GPU, on top which I have a python wrapper made with pybind11. I finally define a pure python module that imports this pybind11-made module and defines a class that inherits from scipy.sparse.linalg.LinearOperator, the constructor of this class calls the constructor of my underlying C++ class and a _matvec method is defined that calls for the CUDA C++ matrix-vector operation.

When using it as in :

# First call to "warm-up"
res = custom_operator @ x
# Then subsequent calls to the operation to time it
for _ in range(10):
res = custom_operator @ x
# Finally solving a linear system using scipy's GMRES
x,info = gmres(custom_operator, b)


And profiling using Nsight, with trace enabled on CUDA, OpenMP, OS Runtine and NVTX (with a NVTX3_FUNC_RANGE() at the entry of the C++ code of this matvec function) I get the following results :

I observe :

• (in red) the first call to the matvec function (that takes longer due to CUDA "warm-up")
• (in blue) the 10 calls made directly by myself
• (in green) the first call, of many, to the function made by scipy's solver.

I observe that the call made from scipy's solver takes longer (as seen on the span of the NVTX event) but that this time is not associated to any CUDA API calls.

Then, if I drop the OpenMP part (that I normally use to manage multiple GPUs and to send some computation that can't be done on the GPU on the CPU) then I obtain the following :

Here I don't observe this slow down : the matvec calls made from scipy's solver take as much time as the 10 first ones. The culptit would be some threads interactions between the OpenMP threads of my C++ library and the threads of scipy (or other underlying python libraries) ?

If it is, what could I do prevent such slow-down due to threads interactions ?

For completeness, here is the full code of this matvec function :

template <>
il::Array<double> HmatCuda<double>::matvec(il::ArrayView<double> x) {

NVTX3_FUNC_RANGE();

il::Array<double> y(vector_size_, 0.0);

#pragma omp parallel num_threads(this->n_openMP_threads_)
{

int thread_id = omp_get_thread_num();

il::Array<double> y_private(vector_size_, 0.0);

// The first num_gpus_ threads manage GPUs
if (thread_id < num_gpus_) {

int gpu_id = thread_id;

CHECK_CUDA_ERROR(cudaSetDevice(gpu_id));

// Copying lhs vector host > device
CHECK_CUDA_ERROR(cudaMemcpy(d_x_[gpu_id], x.data(), vector_size_bytes_, cudaMemcpyHostToDevice));

double alpha = 1.0;
double beta = 0.0;

// FR blocks computation
int fr_block_size = this->hr_->leaf_size * this->dof_dimension_;

// part of the computation assigned to a first cuda stream
cusparseSetStream(cusparse_handle_[gpu_id], cuda_streams_[gpu_id][0]);

cusparseDbsrmv(
cusparse_handle_[gpu_id],
CUSPARSE_DIRECTION_COLUMN,
CUSPARSE_OPERATION_NON_TRANSPOSE,
total_block_size_,
total_block_size_,
num_FR_per_gpu_[gpu_id],
&alpha,
FR_bsr_descr_[gpu_id],
d_FR_data_[gpu_id],
d_FR_bsrRowPtr_[gpu_id],
d_FR_bsrColInd_[gpu_id],
fr_block_size,
d_x_[gpu_id],
&beta,
d_y_[gpu_id]
);


// part of the computation assigned to a second cuda stream
CHECK_CUBLAS_ERROR(cublasSetStream(cublas_handle_[gpu_id], cuda_streams_[gpu_id][1]));

for (int block_size : LR_std_sizes_per_gpu_[gpu_id]){

int num_blocks = LR_std_indices_[block_size].size();

// Compute tmp = B*x
CHECK_CUBLAS_ERROR(cublasDgemmBatched(
cublas_handle_[gpu_id], CUBLAS_OP_N, CUBLAS_OP_N,
this->fixed_rank_*this->dof_dimension_,
1, // 1 column matrix
block_size*this->dof_dimension_,
&alpha,
(const double**)d_LR_B_data_pointers_[gpu_id][block_size],
this->fixed_rank_*this->dof_dimension_,
(const double**)d_LR_x_pointers_[gpu_id][block_size],
block_size*this->dof_dimension_,
&beta,
d_LR_tmp_pointers_[gpu_id][block_size],
this->fixed_rank_*this->dof_dimension_,
num_blocks
));

// Compute y = A*tmp
CHECK_CUBLAS_ERROR(cublasDgemmBatched(
cublas_handle_[gpu_id], CUBLAS_OP_N, CUBLAS_OP_N,
block_size*this->dof_dimension_,
1, // 1 column matrix
this->fixed_rank_*this->dof_dimension_,
&alpha,
(const double**)d_LR_A_data_pointers_[gpu_id][block_size],
block_size*this->dof_dimension_,
(const double**)d_LR_tmp_pointers_[gpu_id][block_size],
this->fixed_rank_*this->dof_dimension_,
&beta,
d_LR_y_pointers_[gpu_id][block_size],
block_size*this->dof_dimension_,
num_blocks
));
}


// We sync the streams
cudaDeviceSynchronize();

// Here we gather the partial results
scatter_add(
d_y_[gpu_id],
d_y_partial_LR_[gpu_id],
d_LR_y_partial_src_indices_[gpu_id],
d_LR_y_partial_dest_indices_[gpu_id],
d_LR_y_partial_lengths_[gpu_id],
num_LR_per_gpu_[gpu_id]
);

cudaDeviceSynchronize();

// Finnally, copy it back to cpu
double* y_data = y_private.Data();
CHECK_CUDA_ERROR(cudaMemcpy(y_data, d_y_[gpu_id], vector_size_bytes_, cudaMemcpyDeviceToHost));


} // GPU - assigned threads

else { // remaining threads > computation that cannot be performed on the GPU

#pragma omp for schedule(guided) nowait
for (int i : FR_non_std_indices) {
auto i0 = hr_->pattern_.FRB_pattern(1, i);
auto j0 = hr_->pattern_.FRB_pattern(2, i);
auto iend = hr_->pattern_.FRB_pattern(3, i);
auto jend = hr_->pattern_.FRB_pattern(4, i);

auto a = (*full_rank_blocks_).view();

auto xs = x.view(il::Range{j0 * dof_dimension_, jend * dof_dimension_});
auto ys = y_private.Edit(il::Range{i0 * dof_dimension_, iend * dof_dimension_});

il::blas(1.0, a, xs, 1.0, il::io, ys);
}

#pragma omp for schedule(guided) nowait
for (int i : LR_non_std_indices_) {

auto i0 = hr_->pattern_.LRB_pattern(1, i);
auto j0 = hr_->pattern_.LRB_pattern(2, i);
auto iend = hr_->pattern_.LRB_pattern(3, i);
auto jend = hr_->pattern_.LRB_pattern(4, i);

auto a = low_rank_blocks_->A.view();
auto b = low_rank_blocks_->B.view();

auto xs = x.view(il::Range{j0 * dof_dimension_, jend * dof_dimension_});
auto ys =y_private.Edit(il::Range{i0 * dof_dimension_, iend * dof_dimension_});
auto r = a.size(1);
il::Array<double> tmp{r, 0.0};

il::blas(1.0, b, il:ot::None, xs, 0.0, il::io, tmp.Edit()); // Note here we have stored b (not b^T)
il::blas(1.0, a, tmp.view(), 1.0, il::io, ys);
}
}

// We sum the partial results of each thread
#pragma omp critical
{
il::blas(1., y_private.view(), il::io_t{}, y.Edit());
}

} // pragma omp parallel

return y;
}

Continue reading...