Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Rewrite CD solver using more BLAS #4446

Merged
merged 17 commits into from
Feb 9, 2022
Merged

Rewrite CD solver using more BLAS #4446

Changes from 1 commit
Commits
Show all changes
17 commits
Select commit Hold shift + click to select a range
1a4fd90
Rewrite CD solver using more BLAS
achirkin Dec 13, 2021
e579757
Merge branch 'branch-22.02' into enh-coordinate-descent-via-blas
achirkin Dec 20, 2021
45dc709
Refactoring for the new raft
achirkin Dec 20, 2021
58fd3e3
Use RAFT_CUBLAS_TRY in place of CUBLAS_CHECK
achirkin Dec 21, 2021
b4040ea
Merge branch 'branch-22.02' into enh-coordinate-descent-via-blas
achirkin Jan 14, 2022
ab94f30
Merge branch 'branch-22.02' into enh-coordinate-descent-via-blas
achirkin Jan 17, 2022
7b246e5
Don't use cublas_wrappers
achirkin Jan 20, 2022
e719dc2
Merge branch 'branch-22.04' into enh-coordinate-descent-via-blas
achirkin Jan 25, 2022
fc8c405
Merge branch 'branch-22.04' into enh-coordinate-descent-via-blas
achirkin Feb 2, 2022
11685e1
unpin my raft fork
achirkin Feb 5, 2022
23c9a5e
Add more comments
achirkin Feb 7, 2022
2dc39bc
Merge branch 'branch-22.04' into enh-coordinate-descent-via-blas
achirkin Feb 8, 2022
2a2dd45
Merge remote-tracking branch 'rapidsai/branch-22.04' into enh-coordin…
achirkin Feb 8, 2022
e8c1358
Fix imports (pending raft PR)
achirkin Feb 8, 2022
92817ce
unpin raft branch
achirkin Feb 8, 2022
8df16b0
Merge branch 'branch-22.04' into enh-coordinate-descent-via-blas
achirkin Feb 9, 2022
35eaef4
Merge branch 'branch-22.04' into enh-coordinate-descent-via-blas
achirkin Feb 9, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
124 changes: 81 additions & 43 deletions cpp/src/solver/cd.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,11 +18,13 @@

#include <raft/cudart_utils.h>
#include <raft/linalg/cublas_wrappers.h>
#include <raft/linalg/gemv.h>
#include <cuml/solvers/params.hpp>
#include <functions/linearReg.cuh>
#include <functions/penalty.cuh>
#include <functions/softThres.cuh>
#include <glm/preprocess.cuh>
// #include <raft/common/nvtx.hpp>
achirkin marked this conversation as resolved.
Show resolved Hide resolved
#include <raft/cuda_utils.cuh>
#include <raft/linalg/add.cuh>
#include <raft/linalg/eltwise.cuh>
Expand All @@ -39,6 +41,36 @@ namespace Solver {

using namespace MLCommon;

namespace {

/** Epoch and iteration -related state. */
template <typename math_t>
struct ConvState {
math_t coef;
math_t coefMax;
math_t diffMax;
};

template <typename math_t>
achirkin marked this conversation as resolved.
Show resolved Hide resolved
__global__ void __launch_bounds__(1, 1) updateCoefKernel(math_t* coefLoc,
const math_t* squaredLoc,
ConvState<math_t>* convStateLoc,
const math_t l1_alpha)
{
auto coef = *coefLoc;
auto r = coef > l1_alpha ? coef - l1_alpha : (coef < -l1_alpha ? coef + l1_alpha : 0);
auto squared = *squaredLoc;
r = squared > math_t(1e-5) ? r / squared : math_t(0);
auto diff = raft::myAbs(convStateLoc->coef - r);
if (convStateLoc->diffMax < diff) convStateLoc->diffMax = diff;
auto absv = raft::myAbs(r);
if (convStateLoc->coefMax < absv) convStateLoc->coefMax = absv;
convStateLoc->coef = -r;
*coefLoc = r;
}

} // namespace

/**
* Fits a linear, lasso, and elastic-net regression model using Coordinate Descent solver
* @param handle
Expand Down Expand Up @@ -95,22 +127,20 @@ void cdFit(const raft::handle_t& handle,
math_t tol,
cudaStream_t stream)
{
// RAFT_USING_NVTX_RANGE("ML::Solver::cdFit-%d-%d", n_rows, n_cols);
ASSERT(n_cols > 0, "Parameter n_cols: number of columns cannot be less than one");
ASSERT(n_rows > 1, "Parameter n_rows: number of rows cannot be less than two");
ASSERT(loss == ML::loss_funct::SQRD_LOSS,
"Parameter loss: Only SQRT_LOSS function is supported for now");

cublasHandle_t cublas_handle = handle.get_cublas_handle();

rmm::device_uvector<math_t> pred(n_rows, stream);
rmm::device_uvector<math_t> residual(n_rows, stream);
rmm::device_uvector<math_t> squared(n_cols, stream);
rmm::device_uvector<math_t> mu_input(0, stream);
rmm::device_uvector<math_t> mu_labels(0, stream);
rmm::device_uvector<math_t> norm2_input(0, stream);

std::vector<math_t> h_coef(n_cols, math_t(0));

if (fit_intercept) {
mu_input.resize(n_cols, stream);
mu_labels.resize(1, stream);
Expand All @@ -135,7 +165,7 @@ void cdFit(const raft::handle_t& handle,
initShuffle(ri, g);

math_t l2_alpha = (1 - l1_ratio) * alpha * n_rows;
alpha = l1_ratio * alpha * n_rows;
math_t l1_alpha = l1_ratio * alpha * n_rows;

if (normalize) {
math_t scalar = math_t(1.0) + l2_alpha;
Expand All @@ -148,59 +178,67 @@ void cdFit(const raft::handle_t& handle,

raft::copy(residual.data(), labels, n_rows, stream);

ConvState<math_t> h_convState;
rmm::device_uvector<ConvState<math_t>> convStateBuf(1, stream);
auto convStateLoc = convStateBuf.data();

CUBLAS_CHECK(
raft::linalg::cublassetpointermode(cublas_handle, CUBLAS_POINTER_MODE_DEVICE, stream));
rmm::device_scalar<math_t> cublas_alpha(1.0, stream);
rmm::device_scalar<math_t> cublas_beta(0.0, stream);

achirkin marked this conversation as resolved.
Show resolved Hide resolved
for (int i = 0; i < epochs; i++) {
// RAFT_USING_NVTX_RANGE(stream, "ML::Solver::cdFit::epoch-%d", i);
achirkin marked this conversation as resolved.
Show resolved Hide resolved
if (i > 0 && shuffle) { Solver::shuffle(ri, g); }

math_t coef_max = 0.0;
math_t d_coef_max = 0.0;
math_t coef_prev = 0.0;
CUDA_CHECK(cudaMemsetAsync(convStateLoc, 0, sizeof(ConvState<math_t>), stream));
achirkin marked this conversation as resolved.
Show resolved Hide resolved

for (int j = 0; j < n_cols; j++) {
// RAFT_USING_NVTX_RANGE(stream, "ML::Solver::cdFit::col-%d", j);
int ci = ri[j];
math_t* coef_loc = coef + ci;
math_t* squared_loc = squared.data() + ci;
math_t* input_col_loc = input + (ci * n_rows);

raft::linalg::multiplyScalar(pred.data(), input_col_loc, h_coef[ci], n_rows, stream);
raft::linalg::add(residual.data(), residual.data(), pred.data(), n_rows, stream);
raft::linalg::gemm(handle,
input_col_loc,
n_rows,
1,
residual.data(),
coef_loc,
1,
1,
CUBLAS_OP_T,
CUBLAS_OP_N,
stream);

if (l1_ratio > math_t(0.0)) Functions::softThres(coef_loc, coef_loc, alpha, 1, stream);

raft::linalg::eltwiseDivideCheckZero(coef_loc, coef_loc, squared_loc, 1, stream);

coef_prev = h_coef[ci];
raft::update_host(&(h_coef[ci]), coef_loc, 1, stream);
CUDA_CHECK(cudaStreamSynchronize(stream));

math_t diff = abs(coef_prev - h_coef[ci]);

if (diff > d_coef_max) d_coef_max = diff;

if (abs(h_coef[ci]) > coef_max) coef_max = abs(h_coef[ci]);

raft::linalg::multiplyScalar(pred.data(), input_col_loc, h_coef[ci], n_rows, stream);
raft::linalg::subtract(residual.data(), residual.data(), pred.data(), n_rows, stream);
raft::copy(&(convStateLoc->coef), coef_loc, 1, stream);
achirkin marked this conversation as resolved.
Show resolved Hide resolved
CUBLAS_CHECK(raft::linalg::cublasaxpy(
cublas_handle, n_rows, coef_loc, input_col_loc, 1, residual.data(), 1, stream));

raft::linalg::cublasgemv(cublas_handle,
achirkin marked this conversation as resolved.
Show resolved Hide resolved
CUBLAS_OP_N,
1,
n_rows,
cublas_alpha.data(),
input_col_loc,
1,
residual.data(),
1,
cublas_beta.data(),
coef_loc,
1,
stream);

updateCoefKernel<math_t><<<dim3(1, 1, 1), dim3(1, 1, 1), 0, stream>>>(
coef_loc, squared_loc, convStateLoc, l1_alpha);
CUDA_CHECK(cudaGetLastError());
achirkin marked this conversation as resolved.
Show resolved Hide resolved

CUBLAS_CHECK(raft::linalg::cublasaxpy(cublas_handle,
achirkin marked this conversation as resolved.
Show resolved Hide resolved
n_rows,
&(convStateLoc->coef),
input_col_loc,
1,
residual.data(),
1,
stream));
}
raft::update_host(&h_convState, convStateLoc, 1, stream);
CUDA_CHECK(cudaStreamSynchronize(stream));
achirkin marked this conversation as resolved.
Show resolved Hide resolved

bool flag_continue = true;
if (coef_max == math_t(0)) { flag_continue = false; }

if ((d_coef_max / coef_max) < tol) { flag_continue = false; }

if (!flag_continue) { break; }
if (h_convState.coefMax < tol || (h_convState.diffMax / h_convState.coefMax) < tol) break;
}

CUBLAS_CHECK(raft::linalg::cublassetpointermode(cublas_handle, CUBLAS_POINTER_MODE_HOST, stream));
achirkin marked this conversation as resolved.
Show resolved Hide resolved

if (fit_intercept) {
GLM::postProcessData(handle,
input,
Expand Down