You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@mxnet.apache.org by jx...@apache.org on 2017/12/09 00:55:39 UTC

[incubator-mxnet] branch master updated: change workspace allocation in random samplers (#8935)

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 <as...@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
['"commits@mxnet.apache.org" <co...@mxnet.apache.org>'].