This is an automated email from the ASF dual-hosted git repository. jxie pushed a commit to branch master in repository https://gitbox.apache.org/repos/asf/incubator-mxnet.git
The following commit(s) were added to refs/heads/master by this push: new 01b7ae3 change workspace allocation in random samplers (#8935) 01b7ae3 is described below commit 01b7ae3a4db14a47f426882b0fbf2e32192a3954 Author: moin <asmushet...@yahoo.de> AuthorDate: Sat Dec 9 01:55:36 2017 +0100 change workspace allocation in random samplers (#8935) --- src/operator/random/sample_op.h | 105 +++++++++++++++++++--------------------- 1 file changed, 51 insertions(+), 54 deletions(-) diff --git a/src/operator/random/sample_op.h b/src/operator/random/sample_op.h index 84bb851..9fdff03 100644 --- a/src/operator/random/sample_op.h +++ b/src/operator/random/sample_op.h @@ -242,44 +242,30 @@ using FSampleCompute = std::function<void (const nnvm::NodeAttrs& attrs, using mxnet::TBlob; -// Convenience class that transfers a host based scalar into an -// array on either the host or the device. Needed as -// the core samplers expect parameters to be tensors located on the -// appropriate device. -template<typename xpu> -Context AllocContext(); -template<> -MSHADOW_FORCE_INLINE Context AllocContext<cpu>() { return Context::CPU(); } -template<> -MSHADOW_FORCE_INLINE Context AllocContext<gpu>() { return Context::GPU(); } - +// Allocates a single chunk of workspace memory and partitions it into three +// workspace tensors that hold the seeds as well as the distribution parameters. template<typename xpu, typename DType> -struct Scalar2Array { - Storage::Handle array; - Scalar2Array(DType scalar, const OpContext& ctx) { - Stream<xpu> *s = ctx.get_stream<xpu>(); - array = Storage::Get()->Alloc(sizeof(DType), AllocContext<xpu>()); - Tensor<xpu, 1, DType> src(Ref(), Shape1(1), s); - Copy(src, Tensor<cpu, 1, DType>(&scalar, Shape1(1)), s); - } - ~Scalar2Array() { - Storage::Get()->Free(array); - } - DType *Ref() { return static_cast<DType*>(array.dptr); } - Tensor<xpu, 1, DType> GetTensor() { return Tensor<xpu, 1, DType>(Ref(), Shape1(1)); } -}; - -// Convienience function to generate the required number of seeds for sampling -template<typename xpu> -MSHADOW_FORCE_INLINE Tensor<xpu, 1, unsigned int> GetSeeds(index_t N, const OpContext& ctx) { +MSHADOW_FORCE_INLINE void GetSamplingTempData(index_t N, DType p1, DType p2, const OpContext& ctx, + Tensor<xpu, 1, unsigned int>* seeds, + Tensor<xpu, 1, DType>* parm1, + Tensor<xpu, 1, DType>* parm2) { Stream<xpu> *s = ctx.get_stream<xpu>(); const index_t nSeeds(OptSampleSeedNum<xpu>(N)); - Tensor<xpu, 1, unsigned int> seeds - = ctx.requested[1].get_space_typed<xpu, 1, unsigned int>(Shape1(nSeeds), ctx.get_stream<xpu>()); - ctx.requested[0].get_random<xpu, float>(s)->GetRandInt(seeds); - return seeds; + // Combined memory requirement for the workspace data. + const index_t nInt(nSeeds + (2 * sizeof(DType) + sizeof(unsigned) - 1) / sizeof(unsigned)); + Tensor<xpu, 1, unsigned> wspace + = ctx.requested[1].get_space_typed<xpu, 1, unsigned>(Shape1(nInt), s); + // Partition workspace into three chunks and initialize them. + *seeds = Tensor<xpu, 1, unsigned>(wspace.dptr_, Shape1(nSeeds), s); + ctx.requested[0].get_random<xpu, float>(s)->GetRandInt(*seeds); + DType *pspace = static_cast<DType*>(static_cast<void*>(wspace.dptr_+nSeeds)); + *parm1 = Tensor<xpu, 1, DType>(pspace, Shape1(1), s); + Copy(*parm1, Tensor<cpu, 1, DType>(&p1, Shape1(1)), s); + *parm2 = Tensor<xpu, 1, DType>(pspace+1, Shape1(1), s); + Copy(*parm2, Tensor<cpu, 1, DType>(&p2, Shape1(1)), s); } + template<typename xpu, typename Sampler> struct SampleMaster; @@ -292,12 +278,14 @@ struct SampleMaster<xpu, UniformSampler<xpu>> { Stream<xpu> *s = ctx.get_stream<xpu>(); const SampleUniformParam& param = nnvm::get<SampleUniformParam>(attrs.parsed); CHECK_GE(param.high, param.low) << "low must be less or equal to high in uniform distribution"; - Scalar2Array<xpu, float> low(param.low, ctx), high(param.high, ctx); - Tensor<xpu, 1, unsigned int> seeds(GetSeeds<xpu>(outputs->Size(), ctx)); + Tensor<xpu, 1, unsigned int> seeds; + Tensor<xpu, 1, float> low, high; + GetSamplingTempData<xpu, float>(outputs->Size(), param.low, param.high, ctx, + &seeds, &low, &high); UniformSampler<xpu> sampler; MSHADOW_REAL_TYPE_SWITCH(outputs[0].type_flag_, OType, { Tensor<xpu, 1, OType> out = outputs->FlatTo1D<xpu, OType>(s); - sampler.Sample(low.GetTensor(), high.GetTensor(), out, seeds, s); + sampler.Sample(low, high, out, seeds, s); }); } }; @@ -311,12 +299,14 @@ struct SampleMaster<xpu, NormalSampler<xpu>> { Stream<xpu> *s = ctx.get_stream<xpu>(); const SampleNormalParam& param = nnvm::get<SampleNormalParam>(attrs.parsed); CHECK_GT(param.scale, 0) << "scale parameter in gaussian has to be positive"; - Scalar2Array<xpu, float> loc(param.loc, ctx), scale(param.scale, ctx); - Tensor<xpu, 1, unsigned int> seeds(GetSeeds<xpu>(outputs->Size(), ctx)); + Tensor<xpu, 1, unsigned int> seeds; + Tensor<xpu, 1, float> loc, scale; + GetSamplingTempData<xpu, float>(outputs->Size(), param.loc, param.scale, ctx, + &seeds, &loc, &scale); NormalSampler<xpu> sampler; MSHADOW_REAL_TYPE_SWITCH(outputs[0].type_flag_, OType, { Tensor<xpu, 1, OType> out = outputs->FlatTo1D<xpu, OType>(s); - sampler.Sample(loc.GetTensor(), scale.GetTensor(), out, seeds, s); + sampler.Sample(loc, scale, out, seeds, s); }); } }; @@ -331,12 +321,14 @@ struct SampleMaster<xpu, GammaSampler<xpu>> { const SampleGammaParam& param = nnvm::get<SampleGammaParam>(attrs.parsed); CHECK_GT(param.alpha, 0) << "alpha parameter in gamma distribution has to be positive"; CHECK_GT(param.beta, 0) << "beta parameter in gamma distribution has to be positive"; - Scalar2Array<xpu, float> alpha(param.alpha, ctx), beta(param.beta, ctx); - Tensor<xpu, 1, unsigned int> seeds(GetSeeds<xpu>(outputs->Size(), ctx)); + Tensor<xpu, 1, unsigned int> seeds; + Tensor<xpu, 1, float> alpha, beta; + GetSamplingTempData<xpu, float>(outputs->Size(), param.alpha, param.beta, ctx, + &seeds, &alpha, &beta); GammaSampler<xpu> sampler; MSHADOW_REAL_TYPE_SWITCH(outputs[0].type_flag_, OType, { Tensor<xpu, 1, OType> out = outputs->FlatTo1D<xpu, OType>(s); - sampler.Sample(alpha.GetTensor(), beta.GetTensor(), out, seeds, s); + sampler.Sample(alpha, beta, out, seeds, s); }); } }; @@ -350,12 +342,13 @@ struct SampleMaster<xpu, ExponentialSampler<xpu>> { Stream<xpu> *s = ctx.get_stream<xpu>(); const SampleExponentialParam& param = nnvm::get<SampleExponentialParam>(attrs.parsed); CHECK_GT(param.lam, 0) << "lambda parameter in exponential distribution has to be positive"; - Scalar2Array<xpu, float> lam(param.lam, ctx); - Tensor<xpu, 1, unsigned int> seeds(GetSeeds<xpu>(outputs->Size(), ctx)); + Tensor<xpu, 1, unsigned int> seeds; + Tensor<xpu, 1, float> lam, dummy; + GetSamplingTempData<xpu, float>(outputs->Size(), param.lam, 0, ctx, &seeds, &lam, &dummy); ExponentialSampler<xpu> sampler; MSHADOW_REAL_TYPE_SWITCH(outputs[0].type_flag_, OType, { Tensor<xpu, 1, OType> out = outputs->FlatTo1D<xpu, OType>(s); - sampler.Sample(lam.GetTensor(), out, seeds, s); + sampler.Sample(lam, out, seeds, s); }); } }; @@ -369,12 +362,13 @@ struct SampleMaster<xpu, PoissonSampler<xpu>> { Stream<xpu> *s = ctx.get_stream<xpu>(); const SamplePoissonParam& param = nnvm::get<SamplePoissonParam>(attrs.parsed); CHECK_GE(param.lam, 0) << "lambda parameter in poisson distribution has to be non-negative"; - Scalar2Array<xpu, float> lam(param.lam, ctx); - Tensor<xpu, 1, unsigned int> seeds(GetSeeds<xpu>(outputs->Size(), ctx)); + Tensor<xpu, 1, unsigned int> seeds; + Tensor<xpu, 1, float> lam, dummy; + GetSamplingTempData<xpu, float>(outputs->Size(), param.lam, 0, ctx, &seeds, &lam, &dummy); PoissonSampler<xpu> sampler; MSHADOW_REAL_TYPE_SWITCH(outputs[0].type_flag_, OType, { Tensor<xpu, 1, OType> out = outputs->FlatTo1D<xpu, OType>(s); - sampler.Sample(lam.GetTensor(), out, seeds, s); + sampler.Sample(lam, out, seeds, s); }); } }; @@ -389,12 +383,13 @@ struct SampleMaster<xpu, NegativeBinomialSampler<xpu>> { const SampleNegBinomialParam& param = nnvm::get<SampleNegBinomialParam>(attrs.parsed); CHECK_GE(param.k, 0) << "k parameter in negative binomial distribution has to be non-negative"; CHECK_GE(param.p, 0) << "p parameter in negative binomial distribution has to be non-negative"; - Scalar2Array<xpu, float> k(param.k, ctx), p(param.p, ctx); - Tensor<xpu, 1, unsigned int> seeds(GetSeeds<xpu>(outputs->Size(), ctx)); + Tensor<xpu, 1, unsigned int> seeds; + Tensor<xpu, 1, float> k, p; + GetSamplingTempData<xpu, float>(outputs->Size(), param.k, param.p, ctx, &seeds, &k, &p); NegativeBinomialSampler<xpu> sampler; MSHADOW_REAL_TYPE_SWITCH(outputs[0].type_flag_, OType, { Tensor<xpu, 1, OType> out = outputs->FlatTo1D<xpu, OType>(s); - sampler.Sample(k.GetTensor(), p.GetTensor(), out, seeds, s); + sampler.Sample(k, p, out, seeds, s); }); } }; @@ -411,12 +406,14 @@ struct SampleMaster<xpu, GeneralizedNegativeBinomialSampler<xpu>> { << "mu parameter in generalized negative binomial distribution has to be non-negative"; CHECK_GE(param.alpha, 0) << "alpha parameter in generalized negative binomial distribution has to be non-negative"; - Scalar2Array<xpu, float> mu(param.mu, ctx), alpha(param.alpha, ctx); - Tensor<xpu, 1, unsigned int> seeds(GetSeeds<xpu>(outputs->Size(), ctx)); + Tensor<xpu, 1, unsigned int> seeds; + Tensor<xpu, 1, float> mu, alpha; + GetSamplingTempData<xpu, float>(outputs->Size(), param.mu, param.alpha, ctx, + &seeds, &mu, &alpha); GeneralizedNegativeBinomialSampler<xpu> sampler; MSHADOW_REAL_TYPE_SWITCH(outputs[0].type_flag_, OType, { Tensor<xpu, 1, OType> out = outputs->FlatTo1D<xpu, OType>(s); - sampler.Sample(mu.GetTensor(), alpha.GetTensor(), out, seeds, s); + sampler.Sample(mu, alpha, out, seeds, s); }); } }; -- To stop receiving notification emails like this one, please contact ['"comm...@mxnet.apache.org" <comm...@mxnet.apache.org>'].