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>'].