jiajinyu closed pull request #12958: Improve dot(csr, rsp) on CPU by 10x
URL: https://github.com/apache/incubator-mxnet/pull/12958
This is a PR merged from a forked repository.
As GitHub hides the original diff on merge, it is displayed below for
the sake of provenance:
As this is a foreign pull request (from a fork), the diff is supplied
below (as it won't show otherwise due to GitHub magic):
diff --git a/src/operator/tensor/dot-inl.h b/src/operator/tensor/dot-inl.h
index 5e469108eda..9d927d819a7 100644
--- a/src/operator/tensor/dot-inl.h
+++ b/src/operator/tensor/dot-inl.h
@@ -521,18 +521,22 @@ struct DotCsrRspDnsByRowBlocks {
const RType* row_idx_ptr = first;
// end of binary search
if (row_idx_ptr == row_idx_r+nnr_r || *row_idx_ptr >
col_idx_l[indptr_l[j+1]-1]) continue;
- for (IType k = indptr_l[j]; k < indptr_l[j+1] && row_idx_ptr !=
row_idx_r+nnr_r;) {
- if (col_idx_l[k] == *row_idx_ptr) {
- const dim_t offset_r = (row_idx_ptr - row_idx_r) * num_cols;
- for (dim_t l = 0; l < num_cols; ++l) {
- out[offset_out+l] += data_l[k] * data_r[offset_r+l];
- }
- ++k;
- ++row_idx_ptr;
- } else if (col_idx_l[k] < *row_idx_ptr) {
- ++k;
- } else {
- ++row_idx_ptr;
+ const RType* end = row_idx_r + nnr_r;
+ const RType* start = row_idx_ptr;
+ for (IType k = indptr_l[j]; k < indptr_l[j+1] && start < end; ++k) {
+ const CType v = col_idx_l[k];
+ if (v < *start) {
+ continue;
+ }
+ const RType* p = std::lower_bound(start, end, v);
+ start = p;
+ if (p >= end || v < *p) {
+ continue;
+ }
+ start += 1;
+ const dim_t offset_r = (p - row_idx_r) * num_cols;
+ for (dim_t l = 0; l < num_cols; ++l) {
+ out[offset_out+l] += data_l[k] * data_r[offset_r+l];
}
}
}
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
With regards,
Apache Git Services
