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

Add Grid stride pairwise dist and fused L2 NN kernels #232

Merged
merged 11 commits into from
Jun 2, 2021
Merged

Add Grid stride pairwise dist and fused L2 NN kernels #232

Show file tree
Hide file tree
Changes from 9 commits
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
3a4ec66
Refactor fusedL2NN to use pairwiseDistance class. invert block y/x di…
mdoijade May 12, 2021
76f9a72
-- add grid stride support to pairwise distance based cosine, l2, l1 …
mdoijade May 17, 2021
af89085
--Add grid stride based fusedL2NN kernel, this gives approx 1.67x spe…
mdoijade May 19, 2021
9c71c4a
Add note on reason to use thread 0 from each warp to write final redu…
mdoijade May 19, 2021
4d76b57
fix clangformat and copyright year
mdoijade May 19, 2021
da2d768
Merge branch 'branch-21.06' into gridStridedDist
mdoijade May 19, 2021
4ada29e
--Add additional Mblk + Nblk shmem for storing norms, and reuse xNorm…
mdoijade May 20, 2021
2e804c2
Use cudaOccupancyMaxActiveBlocksPerSM instead of hard-coded launch bo…
mdoijade May 24, 2021
3408a40
Merge branch 'branch-21.06' into gridStridedDist
mdoijade May 27, 2021
69b316d
initialize regx and regy during each prolog call
mdoijade Jun 1, 2021
9a30a87
initialize ldgX, ldgY in prolog
mdoijade Jun 1, 2021
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
35 changes: 20 additions & 15 deletions cpp/include/raft/distance/cosine.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -61,8 +61,6 @@ void cosineImpl(const DataT *x, const DataT *y, const DataT *xn,
typedef
typename std::conditional<isRowMajor, RowPolicy, ColPolicy>::type KPolicy;

dim3 grid(raft::ceildiv<int>(m, KPolicy::Mblk),
raft::ceildiv<int>(n, KPolicy::Nblk));
dim3 blk(KPolicy::Nthreads);

// Accumulation operation lambda
Expand All @@ -73,7 +71,8 @@ void cosineImpl(const DataT *x, const DataT *y, const DataT *xn,
// epilogue operation lambda for final value calculation
auto epilog_lambda = [] __device__(
AccT acc[KPolicy::AccRowsPerTh][KPolicy::AccColsPerTh],
DataT * regxn, DataT * regyn) {
DataT * regxn, DataT * regyn, IdxT gridStrideX,
IdxT gridStrideY) {
#pragma unroll
for (int i = 0; i < KPolicy::AccRowsPerTh; ++i) {
#pragma unroll
Expand All @@ -83,20 +82,26 @@ void cosineImpl(const DataT *x, const DataT *y, const DataT *xn,
}
};

constexpr size_t shmemSize =
KPolicy::SmemSize + ((KPolicy::Mblk + KPolicy::Nblk) * sizeof(DataT));
if (isRowMajor) {
pairwiseDistanceMatKernel<true, DataT, AccT, OutT, IdxT, KPolicy,
decltype(core_lambda), decltype(epilog_lambda),
FinalLambda, true>
<<<grid, blk, KPolicy::SmemSize, stream>>>(x, y, xn, yn, m, n, k, lda,
ldb, ldd, dOutput, core_lambda,
epilog_lambda, fin_op);
auto cosineRowMajor =
pairwiseDistanceMatKernel<true, DataT, AccT, OutT, IdxT, KPolicy,
decltype(core_lambda), decltype(epilog_lambda),
FinalLambda, true>;
dim3 grid = launchConfigGenerator<KPolicy>(m, n, shmemSize, cosineRowMajor);
cosineRowMajor<<<grid, blk, shmemSize, stream>>>(
x, y, xn, yn, m, n, k, lda, ldb, ldd, dOutput, core_lambda, epilog_lambda,
fin_op);
} else {
pairwiseDistanceMatKernel<true, DataT, AccT, OutT, IdxT, KPolicy,
decltype(core_lambda), decltype(epilog_lambda),
FinalLambda, false>
<<<grid, blk, KPolicy::SmemSize, stream>>>(x, y, xn, yn, m, n, k, lda,
ldb, ldd, dOutput, core_lambda,
epilog_lambda, fin_op);
auto cosineColMajor =
pairwiseDistanceMatKernel<true, DataT, AccT, OutT, IdxT, KPolicy,
decltype(core_lambda), decltype(epilog_lambda),
FinalLambda, false>;
dim3 grid = launchConfigGenerator<KPolicy>(m, n, shmemSize, cosineColMajor);
cosineColMajor<<<grid, blk, shmemSize, stream>>>(
x, y, xn, yn, m, n, k, lda, ldb, ldd, dOutput, core_lambda, epilog_lambda,
fin_op);
}

CUDA_CHECK(cudaGetLastError());
Expand Down
77 changes: 47 additions & 30 deletions cpp/include/raft/distance/euclidean.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -60,8 +60,6 @@ void euclideanExpImpl(const DataT *x, const DataT *y, const DataT *xn,
typedef
typename std::conditional<isRowMajor, RowPolicy, ColPolicy>::type KPolicy;

dim3 grid(raft::ceildiv<int>(m, KPolicy::Mblk),
raft::ceildiv<int>(n, KPolicy::Nblk));
dim3 blk(KPolicy::Nthreads);

// Accumulation operation lambda
Expand All @@ -72,7 +70,8 @@ void euclideanExpImpl(const DataT *x, const DataT *y, const DataT *xn,
// epilogue operation lambda for final value calculation
auto epilog_lambda = [sqrt] __device__(
AccT acc[KPolicy::AccRowsPerTh][KPolicy::AccColsPerTh],
DataT * regxn, DataT * regyn) {
DataT * regxn, DataT * regyn, IdxT gridStrideX,
IdxT gridStrideY) {
#pragma unroll
for (int i = 0; i < KPolicy::AccRowsPerTh; ++i) {
#pragma unroll
Expand All @@ -91,20 +90,29 @@ void euclideanExpImpl(const DataT *x, const DataT *y, const DataT *xn,
}
};

constexpr size_t shmemSize =
KPolicy::SmemSize + ((KPolicy::Mblk + KPolicy::Nblk) * sizeof(DataT));
if (isRowMajor) {
pairwiseDistanceMatKernel<true, DataT, AccT, OutT, IdxT, KPolicy,
decltype(core_lambda), decltype(epilog_lambda),
FinalLambda, true>
<<<grid, blk, KPolicy::SmemSize, stream>>>(x, y, xn, yn, m, n, k, lda,
ldb, ldd, dOutput, core_lambda,
epilog_lambda, fin_op);
auto euclideanExpRowMajor =
pairwiseDistanceMatKernel<true, DataT, AccT, OutT, IdxT, KPolicy,
decltype(core_lambda), decltype(epilog_lambda),
FinalLambda, true>;
dim3 grid =
launchConfigGenerator<KPolicy>(m, n, shmemSize, euclideanExpRowMajor);

euclideanExpRowMajor<<<grid, blk, shmemSize, stream>>>(
x, y, xn, yn, m, n, k, lda, ldb, ldd, dOutput, core_lambda, epilog_lambda,
fin_op);
} else {
pairwiseDistanceMatKernel<true, DataT, AccT, OutT, IdxT, KPolicy,
decltype(core_lambda), decltype(epilog_lambda),
FinalLambda, false>
<<<grid, blk, KPolicy::SmemSize, stream>>>(x, y, xn, yn, m, n, k, lda,
ldb, ldd, dOutput, core_lambda,
epilog_lambda, fin_op);
auto euclideanExpColMajor =
pairwiseDistanceMatKernel<true, DataT, AccT, OutT, IdxT, KPolicy,
decltype(core_lambda), decltype(epilog_lambda),
FinalLambda, false>;
dim3 grid =
launchConfigGenerator<KPolicy>(m, n, shmemSize, euclideanExpColMajor);
euclideanExpColMajor<<<grid, blk, shmemSize, stream>>>(
x, y, xn, yn, m, n, k, lda, ldb, ldd, dOutput, core_lambda, epilog_lambda,
fin_op);
}

CUDA_CHECK(cudaGetLastError());
Expand Down Expand Up @@ -229,8 +237,7 @@ void euclideanUnExpImpl(const DataT *x, const DataT *y, IdxT m, IdxT n, IdxT k,

typedef
typename std::conditional<isRowMajor, RowPolicy, ColPolicy>::type KPolicy;
dim3 grid(raft::ceildiv<int>(m, KPolicy::Mblk),
raft::ceildiv<int>(n, KPolicy::Nblk));

dim3 blk(KPolicy::Nthreads);

// Accumulation operation lambda
Expand All @@ -242,7 +249,8 @@ void euclideanUnExpImpl(const DataT *x, const DataT *y, IdxT m, IdxT n, IdxT k,
// epilogue operation lambda for final value calculation
auto epilog_lambda = [sqrt] __device__(
AccT acc[KPolicy::AccRowsPerTh][KPolicy::AccColsPerTh],
DataT * regxn, DataT * regyn) {
DataT * regxn, DataT * regyn, IdxT gridStrideX,
IdxT gridStrideY) {
if (sqrt) {
#pragma unroll
for (int i = 0; i < KPolicy::AccRowsPerTh; ++i) {
Expand All @@ -255,19 +263,28 @@ void euclideanUnExpImpl(const DataT *x, const DataT *y, IdxT m, IdxT n, IdxT k,
};

if (isRowMajor) {
pairwiseDistanceMatKernel<false, DataT, AccT, OutT, IdxT, KPolicy,
decltype(core_lambda), decltype(epilog_lambda),
FinalLambda>
<<<grid, blk, KPolicy::SmemSize, stream>>>(
x, y, nullptr, nullptr, m, n, k, lda, ldb, ldd, dOutput, core_lambda,
epilog_lambda, fin_op);
auto euclideanUnExpRowMajor =
pairwiseDistanceMatKernel<false, DataT, AccT, OutT, IdxT, KPolicy,
decltype(core_lambda), decltype(epilog_lambda),
FinalLambda, true>;
dim3 grid = launchConfigGenerator<KPolicy>(m, n, KPolicy::SmemSize,
euclideanUnExpRowMajor);

euclideanUnExpRowMajor<<<grid, blk, KPolicy::SmemSize, stream>>>(
x, y, nullptr, nullptr, m, n, k, lda, ldb, ldd, dOutput, core_lambda,
epilog_lambda, fin_op);

} else {
pairwiseDistanceMatKernel<false, DataT, AccT, OutT, IdxT, KPolicy,
decltype(core_lambda), decltype(epilog_lambda),
FinalLambda, isRowMajor>
<<<grid, blk, KPolicy::SmemSize, stream>>>(
x, y, nullptr, nullptr, m, n, k, lda, ldb, ldd, dOutput, core_lambda,
epilog_lambda, fin_op);
auto euclideanUnExpColMajor =
pairwiseDistanceMatKernel<false, DataT, AccT, OutT, IdxT, KPolicy,
decltype(core_lambda), decltype(epilog_lambda),
FinalLambda, false>;
dim3 grid = launchConfigGenerator<KPolicy>(m, n, KPolicy::SmemSize,
euclideanUnExpColMajor);

euclideanUnExpColMajor<<<grid, blk, KPolicy::SmemSize, stream>>>(
x, y, nullptr, nullptr, m, n, k, lda, ldb, ldd, dOutput, core_lambda,
epilog_lambda, fin_op);
}

CUDA_CHECK(cudaGetLastError());
Expand Down
Loading