piiswrong commented on a change in pull request #10351: [MXNET-259] Performance improvement
of random.shuffle
URL: https://github.com/apache/incubator-mxnet/pull/10351#discussion_r178471176
##########
File path: src/operator/random/shuffle_op.cc
##########
@@ -55,18 +56,24 @@ void Shuffle1D(DType* const out, const index_t size, Rand* const prnd)
{
template<typename DType, typename Rand>
void ShuffleND(DType* const out, const index_t size, const index_t first_axis_len,
- Rand* const prnd) {
+ Rand* const prnd, const OpContext& ctx) {
// Fisher-Yates shuffling
+ using namespace mxnet_op;
const index_t stride = size / first_axis_len;
auto rand_n = [prnd](index_t n) {
std::uniform_int_distribution<index_t> dist(0, n - 1);
return dist(*prnd);
};
CHECK_GT(first_axis_len, 0U);
+ const size_t stride_bytes = sizeof(DType) * stride;
+ Tensor<cpu, 1, char> buf =
+ ctx.requested[1].get_space_typed<cpu, 1, char>(Shape1(stride_bytes), ctx.get_stream<cpu>());
Review comment:
I think you need to add Fresources to operator registration?
----------------------------------------------------------------
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
|