You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@mxnet.apache.org by GitBox <gi...@apache.org> on 2018/07/26 01:05:53 UTC
[GitHub] eric-haibin-lin closed pull request #11825: Fix non-determinism of
dot(csr.T, dns) = dns with tests
eric-haibin-lin closed pull request #11825: Fix non-determinism of dot(csr.T, dns) = dns with tests
URL: https://github.com/apache/incubator-mxnet/pull/11825
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/common/utils.h b/src/common/utils.h
index d7ed4ddf04c..96949a047fb 100644
--- a/src/common/utils.h
+++ b/src/common/utils.h
@@ -663,6 +663,18 @@ constexpr size_t MaxIntegerValue<mshadow::half::half_t>() {
return size_t(2) << 10;
}
+MSHADOW_XINLINE int ilog2ul(size_t a) {
+ int k = 1;
+ while (a >>= 1) ++k;
+ return k;
+}
+
+MSHADOW_XINLINE int ilog2ui(unsigned int a) {
+ int k = 1;
+ while (a >>= 1) ++k;
+ return k;
+}
+
} // namespace common
} // namespace mxnet
#endif // MXNET_COMMON_UTILS_H_
diff --git a/src/operator/tensor/dot-inl.cuh b/src/operator/tensor/dot-inl.cuh
index 8aedec066ba..6c94ad1eca1 100644
--- a/src/operator/tensor/dot-inl.cuh
+++ b/src/operator/tensor/dot-inl.cuh
@@ -31,6 +31,7 @@
#include "./sort_op.h"
#include "./util/tensor_util-inl.h"
#include "./util/tensor_util-inl.cuh"
+#include "../../common/utils.h"
namespace mxnet {
namespace op {
@@ -176,119 +177,6 @@ struct DotCsrTransDnsDnsScalarKernel {
}
};
-/*!
- * \brief GPU warp kernel of dot(csr.T, dns1) = dns2
- * Parallelization by columns: 1 warp computes one lhs column for one rhs column
- */
-struct DotCsrTransDnsDnsWarpKernel {
- /*!
- * \brief see DotCsrTransDnsDnsScalarKernel Map for documentation.
- */
- template<typename DType, typename IType, typename CType>
- __device__ __forceinline__ static void Map(int tid,
- DType* out,
- const DType* data_l,
- const IType* indptr_l,
- const CType* col_idx_l,
- const DType* data_r,
- const nnvm::dim_t num_cols_r) {
- using nnvm::dim_t;
- const dim_t warp_id = tid / 32; // global warp id
- const dim_t lane = tid & (32-1); // local thread id within warp
- const dim_t icol = warp_id / num_cols_r; // lhs column that this warp computes
- const dim_t kcol = warp_id % num_cols_r; // rhs column that this warp computes
-
- // Compute range of nnz elements in this column
- const dim_t low = static_cast<dim_t>(indptr_l[icol]);
- const dim_t high = static_cast<dim_t>(indptr_l[icol+1]);
-
- // Iterate through the nnz elements in this column
- for (dim_t j = low+lane; j < high; j+=32) {
- const dim_t irow = static_cast<dim_t>(col_idx_l[j]);
- const DType val = data_l[j]*data_r[icol*num_cols_r+kcol];
- atomicAdd(static_cast<DType *>(&(out[irow*num_cols_r+kcol])), val);
- }
- }
-};
-
-/*!
- * \brief GPU thread block kernel of dot(csr.T, dns1) = dns2
- * Parallelization by columns: 1 thread block computes one lhs column for all rhs columns
- */
-struct DotCsrTransDnsDnsThreadBlockKernel {
- /*!
- * \brief see DotCsrTransDnsDnsScalarKernel Map for documentation.
- */
- template<typename DType, typename IType, typename CType>
- __device__ __forceinline__ static void Map(int tid,
- DType* out,
- const DType* data_l,
- const IType* indptr_l,
- const CType* col_idx_l,
- const DType* data_r,
- const nnvm::dim_t num_cols_r) {
- using nnvm::dim_t;
- const dim_t warps_per_block = blockDim.x / 32; // number of warps in this thread block
- const dim_t warp_id = tid / 32; // global warp id
- const dim_t lane = tid & (32-1); // local thread id within warp
- const dim_t icol = blockIdx.x; // lhs column that this thread block computes
- const dim_t kcol = warp_id % warps_per_block; // rhs column where warp starts computing (offset)
-
- // Compute range of nnz elements in this lhs column
- const dim_t low = static_cast<dim_t>(indptr_l[icol]);
- const dim_t high = static_cast<dim_t>(indptr_l[icol+1]);
-
- // Iterate through the nnz elements in this lhs column
- for (dim_t j = low+lane; j < high; j+=32) {
- const dim_t irow = static_cast<dim_t>(col_idx_l[j]);
- const DType datum_l = data_l[j];
- // Iterate over rhs columns that this warp computes
- for (dim_t k = kcol; k < num_cols_r; k+=warps_per_block) {
- const DType val = datum_l*data_r[icol*num_cols_r+k];
- atomicAdd(static_cast<DType *>(&(out[irow*num_cols_r+k])), val);
- }
- }
- }
-};
-
-/*!
- * \brief GPU warp block kernel of dot(csr.T, dns1) = dns2
- * Parallelization by columns: 1 warp computes one lhs column for all rhs columns
- */
-struct DotCsrTransDnsDnsWarpBlockKernel {
- /*!
- * \brief see DotCsrTransDnsDnsScalarKernel Map for documentation.
- */
- template<typename DType, typename IType, typename CType>
- __device__ __forceinline__ static void Map(int tid,
- DType* out,
- const DType* data_l,
- const IType* indptr_l,
- const CType* col_idx_l,
- const DType* data_r,
- const nnvm::dim_t num_cols_r) {
- using nnvm::dim_t;
- const dim_t warp_id = tid / 32; // global warp id
- const dim_t lane = tid & (32-1); // local thread id within warp
- const dim_t icol = warp_id; // lhs column that this warp computes
-
- // Compute range of nnz elements in this column
- const dim_t low = static_cast<dim_t>(indptr_l[icol]);
- const dim_t high = static_cast<dim_t>(indptr_l[icol+1]);
-
- // Iterate through the nnz elements in lhs column
- for (dim_t j = low+lane; j < high; j+=32) {
- const dim_t irow = static_cast<dim_t>(col_idx_l[j]);
- const DType datum_l = data_l[j];
- // Iterate over all rhs columns
- for (dim_t k = 0; k < num_cols_r; k++) {
- const DType val = datum_l*data_r[icol*num_cols_r+k];
- atomicAdd(static_cast<DType *>(&(out[irow*num_cols_r+k])), val);
- }
- }
- }
-};
-
/*!
* \brief GPU Kernel of dot(csr.T, rsp1) = rsp2
* Parallelization by rows: 1 thread/row
@@ -510,7 +398,7 @@ inline void DotCsrDnsDnsImpl(const OpContext& ctx,
return;
}
- using mshadow::cuda::kBaseThreadNum;
+ using namespace mshadow;
using mxnet_op::Kernel;
using mxnet_op::set_zero;
using nnvm::dim_t;
@@ -518,7 +406,6 @@ inline void DotCsrDnsDnsImpl(const OpContext& ctx,
const dim_t num_rows_l = lhs.shape()[0];
const dim_t num_cols_r = rhs.shape_[1];
const dim_t threads_per_warp = mxnet_op::cuda_get_device_prop().warpSize;
- const dim_t threads_per_block = kBaseThreadNum;
dim_t num_threads;
// TODO: remove kernel dependency on warpSize=32
if (threads_per_warp != 32) {
@@ -539,86 +426,120 @@ inline void DotCsrDnsDnsImpl(const OpContext& ctx,
Kernel<set_zero, gpu>::Launch(s, num_threads, data_out.dptr<DType>());
}
if (trans_lhs) {
- // Different kernel versions are optimized for different matrix instances
- // TODO: switch between kernel versions depending on input
- // (1) 'Scalar kernel' (one thread computing one output element )
- // (2) 'Warp kernel' (one warp computing one lhs column for one rhs column )
- // (3) 'Thread block kernel' (one thread block computing one lhs column for all rhs columns)
- // (4) 'Warp block kernel' (one warp computing one lhs column for all rhs columns)
- const int kernel_version = 0;
- switch (kernel_version) {
- case 1:
- num_threads = data_out.Size();
- MXNET_ASSIGN_REQ_SWITCH(req, ReqType, {
- Kernel<DotCsrTransDnsDnsScalarKernel<ReqType>, gpu>::Launch(s, num_threads,
- data_out.dptr<DType>(), data_l.dptr<DType>(), indptr_l.dptr<IType>(),
- col_idx_l.dptr<CType>(), data_r.dptr<DType>(), num_rows_l, num_cols_r);
- });
- break;
- case 2:
- num_threads = threads_per_warp * num_rows_l * num_cols_r;
- Kernel<DotCsrTransDnsDnsWarpKernel, gpu>::Launch(s, num_threads,
- data_out.dptr<DType>(), data_l.dptr<DType>(), indptr_l.dptr<IType>(),
- col_idx_l.dptr<CType>(), data_r.dptr<DType>(), num_cols_r);
- break;
- case 3:
- num_threads = threads_per_block * num_rows_l;
- Kernel<DotCsrTransDnsDnsThreadBlockKernel, gpu>::Launch(s, num_threads,
- data_out.dptr<DType>(), data_l.dptr<DType>(), indptr_l.dptr<IType>(),
- col_idx_l.dptr<CType>(), data_r.dptr<DType>(), num_cols_r);
- break;
- case 4:
- num_threads = threads_per_warp * num_rows_l;
- Kernel<DotCsrTransDnsDnsWarpBlockKernel, gpu>::Launch(s, num_threads,
+ // TODO(haojin2): Switching to deterministic algorithm for now.
+ // Further optimizations to come later.
+ const nnvm::dim_t num_csr_rows = lhs.shape()[0];
+ const nnvm::dim_t num_csr_cols = lhs.shape()[1];
+ const nnvm::dim_t num_dns_rows = rhs.shape_[0];
+ const nnvm::dim_t nnz = lhs.storage_shape().Size();
+
+ IType* original_idx_ptr = nullptr;
+ IType* csc_indices_ptr = nullptr;
+ IType* csc_cols_ptr = nullptr;
+ CType* csr_rows_ptr = nullptr;
+ CType* csc_indptr_ptr = nullptr;
+ DType* csc_data_ptr = nullptr;
+ char* temp_storage_ptr = nullptr;
+ size_t original_idx_bytes = nnz*sizeof(IType);
+ size_t csc_indices_bytes = nnz*sizeof(IType);
+ size_t csc_cols_bytes = nnz*sizeof(IType);
+ size_t csr_rows_bytes = nnz*sizeof(CType);
+ size_t csc_indptr_bytes = (num_csr_cols+1)*sizeof(CType);
+ size_t csc_data_bytes = nnz*sizeof(DType);
+ size_t scan_temp_storage_bytes = 0;
+ size_t temp_storage_bytes = SortByKeyWorkspaceSize<IType, IType, gpu>(nnz);
+ IType* csr_indices_ptr = col_idx_l.dptr<IType>();
+ cub::DeviceScan::ExclusiveSum(temp_storage_ptr,
+ scan_temp_storage_bytes,
+ csc_indptr_ptr,
+ csc_indptr_ptr,
+ num_csr_cols+1,
+ mshadow::Stream<gpu>::GetStream(s));
+ temp_storage_bytes = std::max(temp_storage_bytes, scan_temp_storage_bytes);
+ temp_storage_bytes += (sizeof(dim_t) - temp_storage_bytes % sizeof(dim_t));
+ size_t total_workspace_bytes =
+ original_idx_bytes + csc_indices_bytes + csc_cols_bytes + csr_rows_bytes +
+ csc_indptr_bytes + csc_data_bytes + temp_storage_bytes;
+ total_workspace_bytes += (sizeof(IType) - total_workspace_bytes % sizeof(IType));
+ Tensor<gpu, 1, char> workspace = ctx.requested[0]
+ .get_space_typed<gpu, 1, char>(Shape1(total_workspace_bytes), s);
+ original_idx_ptr = reinterpret_cast<IType*>(workspace.dptr_);
+ csc_indices_ptr = reinterpret_cast<IType*>(workspace.dptr_ + original_idx_bytes);
+ csc_cols_ptr = reinterpret_cast<IType*>(workspace.dptr_ + original_idx_bytes +
+ csc_indices_bytes);
+ csr_rows_ptr = reinterpret_cast<CType*>(workspace.dptr_ + original_idx_bytes +
+ csc_indices_bytes + csc_cols_bytes);
+ csc_indptr_ptr = reinterpret_cast<CType*>(workspace.dptr_ + original_idx_bytes +
+ csc_indices_bytes + csc_cols_bytes +
+ csr_rows_bytes);
+ temp_storage_ptr = workspace.dptr_ + original_idx_bytes + csc_indices_bytes +
+ csc_cols_bytes + csr_rows_bytes + csc_indptr_bytes;
+ csc_data_ptr = reinterpret_cast<DType*>(
+ workspace.dptr_ + total_workspace_bytes - csc_data_bytes);
+
+ // Fill original_idx
+ mxnet_op::Kernel<range_fwd, gpu>::Launch(
+ s, nnz, 1, IType(0), IType(1), kWriteTo, original_idx_ptr);
+ // Fill csc_cols with copy of csr_indices
+ mxnet_op::Kernel<mxnet_op::op_with_req<mshadow_op::identity, kWriteTo>, gpu>::Launch(
+ s, nnz, csc_cols_ptr, csr_indices_ptr);
+ // Allocate the tensors needed for SortByKey
+ Tensor<gpu, 1, IType> original_idx(original_idx_ptr, Shape1(nnz), s);
+ Tensor<gpu, 1, IType> csc_cols(csc_cols_ptr, Shape1(nnz), s);
+ Tensor<gpu, 1, char> temp_storage(temp_storage_ptr, Shape1(temp_storage_bytes), s);
+
+ int num_bits = common::ilog2ul(num_csr_cols - 1);
+ SortByKey(csc_cols, original_idx, true, &temp_storage, 0, num_bits);
+
+ // Scatter csr indptr to row id
+ mxnet_op::Kernel<CsrRowScatterKernel, gpu>::Launch(
+ s, num_csr_rows, indptr_l.dptr<CType>(), csr_rows_ptr, num_csr_rows);
+ // Reset indptr to zero
+ mxnet_op::Kernel<mxnet_op::set_zero, gpu>::Launch(s, num_csr_cols+1, csc_indptr_ptr);
+ // Histogram on the sorted cols
+ mxnet_op::Kernel<HistogramKernel, gpu>::Launch(
+ s, nnz, csc_indptr_ptr, csc_cols_ptr, nnz);
+ // Scan the bin counts for every column to get csc_indptr
+ cub::DeviceScan::ExclusiveSum(temp_storage_ptr,
+ temp_storage_bytes,
+ csc_indptr_ptr,
+ csc_indptr_ptr,
+ num_csr_cols+1,
+ mshadow::Stream<gpu>::GetStream(s));
+ // Assign data to csc matrix arrays
+ mxnet_op::Kernel<CscDataIndicesKernel, gpu>::Launch(
+ s, nnz, original_idx_ptr, data_l.dptr<DType>(), csr_rows_ptr, csc_data_ptr,
+ csc_indices_ptr, nnz);
+ if (num_cols_r > 4) {
+ num_threads = data_out.Size();
+ MXNET_ASSIGN_REQ_SWITCH(req, ReqType, {
+ Kernel<DotCsrDnsDnsScalarKernel<ReqType>, gpu>::Launch(s, num_threads,
+ data_out.dptr<DType>(), csc_data_ptr, csc_indptr_ptr,
+ csc_indices_ptr, data_r.dptr<DType>(), num_cols_r);
+ });
+ } else {
+ num_threads = threads_per_warp * num_rows_l * num_cols_r;
+ MXNET_ASSIGN_REQ_SWITCH(req, ReqType, {
+ Kernel<DotCsrDnsDnsVectorKernel<ReqType>, gpu>::Launch(s, num_threads,
+ data_out.dptr<DType>(), csc_data_ptr, csc_indptr_ptr,
+ csc_indices_ptr, data_r.dptr<DType>(), num_cols_r);
+ });
+ }
+ } else {
+ if (num_cols_r > 4) {
+ num_threads = data_out.Size();
+ MXNET_ASSIGN_REQ_SWITCH(req, ReqType, {
+ Kernel<DotCsrDnsDnsScalarKernel<ReqType>, gpu>::Launch(s, num_threads,
data_out.dptr<DType>(), data_l.dptr<DType>(), indptr_l.dptr<IType>(),
col_idx_l.dptr<CType>(), data_r.dptr<DType>(), num_cols_r);
- break;
- default:
- num_threads = threads_per_warp * num_rows_l * num_cols_r;
- Kernel<DotCsrTransDnsDnsWarpKernel, gpu>::Launch(s, num_threads,
+ });
+ } else {
+ num_threads = threads_per_warp * num_rows_l * num_cols_r;
+ MXNET_ASSIGN_REQ_SWITCH(req, ReqType, {
+ Kernel<DotCsrDnsDnsVectorKernel<ReqType>, gpu>::Launch(s, num_threads,
data_out.dptr<DType>(), data_l.dptr<DType>(), indptr_l.dptr<IType>(),
col_idx_l.dptr<CType>(), data_r.dptr<DType>(), num_cols_r);
- break;
- }
- } else {
- // Different kernel versions are optimized for different matrix instances
- // (1) 'Scalar kernel' (one thread computing one output element)
- // (2) 'Vector kernel' (one warp computing one output element)
- const int kernel_version = 0;
- switch (kernel_version) {
- case 1:
- num_threads = data_out.Size();
- MXNET_ASSIGN_REQ_SWITCH(req, ReqType, {
- Kernel<DotCsrDnsDnsScalarKernel<ReqType>, gpu>::Launch(s, num_threads,
- data_out.dptr<DType>(), data_l.dptr<DType>(), indptr_l.dptr<IType>(),
- col_idx_l.dptr<CType>(), data_r.dptr<DType>(), num_cols_r);
- });
- break;
- case 2:
- num_threads = threads_per_warp * num_rows_l * num_cols_r;
- MXNET_ASSIGN_REQ_SWITCH(req, ReqType, {
- Kernel<DotCsrDnsDnsVectorKernel<ReqType>, gpu>::Launch(s, num_threads,
- data_out.dptr<DType>(), data_l.dptr<DType>(), indptr_l.dptr<IType>(),
- col_idx_l.dptr<CType>(), data_r.dptr<DType>(), num_cols_r);
- });
- break;
- default:
- if (num_cols_r > 4) {
- num_threads = data_out.Size();
- MXNET_ASSIGN_REQ_SWITCH(req, ReqType, {
- Kernel<DotCsrDnsDnsScalarKernel<ReqType>, gpu>::Launch(s, num_threads,
- data_out.dptr<DType>(), data_l.dptr<DType>(), indptr_l.dptr<IType>(),
- col_idx_l.dptr<CType>(), data_r.dptr<DType>(), num_cols_r);
- });
- } else {
- num_threads = threads_per_warp * num_rows_l * num_cols_r;
- MXNET_ASSIGN_REQ_SWITCH(req, ReqType, {
- Kernel<DotCsrDnsDnsVectorKernel<ReqType>, gpu>::Launch(s, num_threads,
- data_out.dptr<DType>(), data_l.dptr<DType>(), indptr_l.dptr<IType>(),
- col_idx_l.dptr<CType>(), data_r.dptr<DType>(), num_cols_r);
- });
- }
- break;
+ });
}
}
});
@@ -671,13 +592,6 @@ struct DotCsrTransDnsRspKernel {
}
};
-// Returns integer log2(a) rounded up
-inline int log2i(size_t a) {
- int k = 1;
- while (a >>= 1) k++;
- return k;
-}
-
/*!
* \brief GPU Impl of dot(csr.T, dns) = rsp
*/
@@ -777,7 +691,7 @@ inline void DotCsrDnsRspImpl(const OpContext& ctx,
Tensor<gpu, 1, IType> original_idx(original_idx_ptr, Shape1(nnz), s);
Tensor<gpu, 1, char> temp_storage(temp_storage_ptr, Shape1(total_temp_bytes), s);
- int num_bits = log2i(num_cols_l - 1);
+ int num_bits = common::ilog2ul(num_cols_l - 1);
SortByKey(col_idx_copy, original_idx, true, &temp_storage, 0, num_bits);
// over-allocate aux indices
@@ -1124,7 +1038,7 @@ inline void DotDnsCsrDnsImpl(const OpContext& ctx, const gpu& gpu_dev,
Tensor<gpu, 1, IType> csc_cols(csc_cols_ptr, Shape1(nnz), s);
Tensor<gpu, 1, char> temp_storage(temp_storage_ptr, Shape1(temp_storage_bytes), s);
- int num_bits = log2i(num_csr_cols - 1);
+ int num_bits = common::ilog2ul(num_csr_cols - 1);
SortByKey(csc_cols, original_idx, true, &temp_storage, 0, num_bits);
// Scatter csr indptr to row id
diff --git a/src/operator/tensor/indexing_op.cu b/src/operator/tensor/indexing_op.cu
index b0ee05ea139..39fd81ef200 100644
--- a/src/operator/tensor/indexing_op.cu
+++ b/src/operator/tensor/indexing_op.cu
@@ -217,7 +217,7 @@ void SparseEmbeddingDeterministicKernelLaunch(const OpContext& ctx,
Kernel<range_fwd, gpu>::Launch(s, data_size, 1, static_cast<dim_t>(0),
static_cast<dim_t>(1), kWriteTo, original_idx);
// sort data with its original idx
- int num_bits = ilog2(num_rows - 1);
+ int num_bits = common::ilog2ui(num_rows - 1);
char* temp_storage_ptr = reinterpret_cast<char*>(temp_storage);
Tensor<gpu, 1, char> temp_storage_tensor(temp_storage_ptr,
Shape1(sort_workspace_size), s);
diff --git a/src/operator/tensor/indexing_op.h b/src/operator/tensor/indexing_op.h
index 07d96dcbdb9..edaf9397303 100644
--- a/src/operator/tensor/indexing_op.h
+++ b/src/operator/tensor/indexing_op.h
@@ -44,6 +44,7 @@
#include "./sort_op.h"
#include "./init_op.h"
#include "../../engine/openmp.h"
+#include "../../common/utils.h"
#ifdef __CUDACC__
#include "./indexing_op-inl.cuh"
#endif
@@ -528,13 +529,6 @@ void SparseEmbeddingOpForwardEx(const nnvm::NodeAttrs& attrs,
}
}
-// Returns integer log2(a) rounded up
-inline int ilog2(unsigned int a) {
- int k = 1;
- while (a >>= 1) k++;
- return k;
-}
-
/*! \brief cast to type and clip to range [0, K - 1]
*/
struct tcast_clip {
@@ -906,7 +900,7 @@ void TakeOpBackwardImpl(mshadow::Stream<cpu>* s,
s, idxshape.Size(), sorted_idx_ptr, sorted_idx_ptr, static_cast<int>(arrshape[axis]));
}
Tensor<cpu, 1, int> original_idx(original_idx_ptr, Shape1(idxshape.Size()), s);
- int num_bits = ilog2(static_cast<unsigned int>(idxshape.Size()) - 1);
+ int num_bits = common::ilog2ui(static_cast<unsigned int>(idxshape.Size()) - 1);
Tensor<cpu, 1, int> sorted_idx(sorted_idx_ptr, Shape1(idxshape.Size()), s);
SortByKey(sorted_idx, original_idx, true, &temp_storage, 0, num_bits);
for (size_t i = 0; i < idxshape.Size(); ++i) {
@@ -1000,7 +994,7 @@ void TakeOpBackwardImpl(mshadow::Stream<gpu>* s,
}
Tensor<gpu, 1, int> original_idx(original_idx_ptr, Shape1(idxshape.Size()), s);
Tensor<gpu, 1, char> temp_storage(temp_storage_ptr, Shape1(temp_storage_bytes), s);
- int num_bits = ilog2(static_cast<unsigned int>(idxshape.Size()) - 1);
+ int num_bits = common::ilog2ui(static_cast<unsigned int>(idxshape.Size()) - 1);
Tensor<gpu, 1, int> sorted_idx(sorted_idx_ptr, Shape1(idxshape.Size()), s);
SortByKey(sorted_idx, original_idx, true, &temp_storage, 0, num_bits);
cub::DeviceScan::ExclusiveSum(temp_storage_ptr,
diff --git a/src/storage/pooled_storage_manager.h b/src/storage/pooled_storage_manager.h
index bed97301fd1..f3a9b16cdd8 100644
--- a/src/storage/pooled_storage_manager.h
+++ b/src/storage/pooled_storage_manager.h
@@ -38,6 +38,7 @@
#include <new>
#include "./storage_manager.h"
#include "../common/cuda_utils.h"
+#include "../common/utils.h"
namespace mxnet {
@@ -173,10 +174,10 @@ class GPUPooledRoundedStorageManager final : public StorageManager {
LOG(FATAL) << "MXNET_GPU_MEM_POOL_PAGE_SIZE cannot be set to a value smaller than 32. " \
<< "Got: " << page_size_ << ".";
}
- if (page_size_ != 1ul << log2_round_up(page_size_)) {
+ if (page_size_ != 1ul << common::ilog2ul(page_size_ - 1)) {
LOG(FATAL) << "MXNET_GPU_MEM_POOL_PAGE_SIZE must be a power of 2. Got: " << page_size_ << ".";
}
- page_size_ = log2_round_up(page_size_);
+ page_size_ = common::ilog2ul(page_size_ - 1);
if (cut_off_ < 20 || cut_off_ > LOG2_MAX_MEM) {
LOG(FATAL) << "MXNET_GPU_MEM_POOL_ROUND_LINEAR_CUTOFF cannot be set to a value " \
<< "smaller than 20 or greater than " << LOG2_MAX_MEM << ". Got: " \
@@ -205,9 +206,6 @@ class GPUPooledRoundedStorageManager final : public StorageManager {
}
private:
- inline int log2_round_up(size_t s) {
- return static_cast<int>(std::ceil(std::log2(s)));
- }
inline int div_pow2_round_up(size_t s, int divisor_log2) {
// (1025, 10) -> 2
// (2048, 10) -> 2
@@ -216,7 +214,7 @@ class GPUPooledRoundedStorageManager final : public StorageManager {
return static_cast<int>(result + (s > (result << divisor_log2) ? 1 : 0));
}
inline int get_bucket(size_t s) {
- int log_size = log2_round_up(s);
+ int log_size = common::ilog2ul(s - 1);
if (log_size > static_cast<int>(cut_off_))
return div_pow2_round_up(s, cut_off_) - 1 + cut_off_;
else
diff --git a/tests/python/unittest/test_sparse_operator.py b/tests/python/unittest/test_sparse_operator.py
index 09b0ae3a9a6..b4783bd2382 100644
--- a/tests/python/unittest/test_sparse_operator.py
+++ b/tests/python/unittest/test_sparse_operator.py
@@ -1424,7 +1424,7 @@ def test_sparse_dot_zero_output(lhs_shape, trans_lhs, rhs_num_cols):
@with_seed()
def test_sparse_dot_determinism():
- def test_dot_determinism(lhs_stype, rhs_stype, lhs_density, rhs_density, transpose_a, transpose_b, forward_stype):
+ def check_dot_determinism(lhs_stype, rhs_stype, lhs_density, rhs_density, transpose_a, transpose_b, forward_stype):
lhs_row = rnd.randint(50, 100)
lhs_col = rnd.randint(50, 100)
if transpose_a:
@@ -1444,10 +1444,11 @@ def test_dot_determinism(lhs_stype, rhs_stype, lhs_density, rhs_density, transpo
res2 = mx.nd.sparse.dot(lhs, rhs, transpose_a=transpose_a, transpose_b=transpose_b, forward_stype=forward_stype)
assert_almost_equal(res1.asnumpy(), res2.asnumpy(), rtol=0.0, atol=0.0)
- test_dot_determinism('csr', 'default', 0.1, 1.0, True, False, 'row_sparse')
+ check_dot_determinism('csr', 'default', 0.1, 1.0, True, False, 'row_sparse')
forward_stype = 'csr' if default_context() == mx.cpu() else 'default'
- test_dot_determinism('default', 'csr', 1.0, 0.1, False, False, forward_stype)
- test_dot_determinism('default', 'csr', 1.0, 0.1, False, True, forward_stype)
+ check_dot_determinism('default', 'csr', 1.0, 0.1, False, False, forward_stype)
+ check_dot_determinism('default', 'csr', 1.0, 0.1, False, True, forward_stype)
+ check_dot_determinism('csr', 'default', 0.1, 1.0, True, False, 'default')
@with_seed()
----------------------------------------------------------------
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:
users@infra.apache.org
With regards,
Apache Git Services