| /* |
| * Licensed to the Apache Software Foundation (ASF) under one |
| * or more contributor license agreements. See the NOTICE file |
| * distributed with this work for additional information |
| * regarding copyright ownership. The ASF licenses this file |
| * to you under the Apache License, Version 2.0 (the |
| * "License"); you may not use this file except in compliance |
| * with the License. You may obtain a copy of the License at |
| * |
| * http://www.apache.org/licenses/LICENSE-2.0 |
| * |
| * Unless required by applicable law or agreed to in writing, |
| * software distributed under the License is distributed on an |
| * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| * KIND, either express or implied. See the License for the |
| * specific language governing permissions and limitations |
| * under the License. |
| */ |
| |
| /*! |
| * Copyright (c) 2015 by Contributors |
| * \file ndarray_function-inl.h |
| * \brief The real implementation of NDArray functions. |
| */ |
| #ifndef MXNET_NDARRAY_NDARRAY_FUNCTION_INL_H_ |
| #define MXNET_NDARRAY_NDARRAY_FUNCTION_INL_H_ |
| |
| #include <vector> |
| #include "./ndarray_function.h" |
| // this file will be included twice by CPU and GPU |
| // macro to help specialize evaluation function |
| |
| #ifndef DECL_TERNARY |
| #define DECL_TERNARY(XPU, OP, FUN) \ |
| template<> \ |
| void Eval<XPU, OP>(const TBlob &lhs, const TBlob &mhs, \ |
| const TBlob &rhs, TBlob *ret, RunContext ctx) { \ |
| FUN<XPU, OP>(lhs, mhs, rhs, ret, ctx); \ |
| } |
| #endif |
| |
| #ifndef DECL_BINARY |
| #define DECL_BINARY(XPU, OP, FUN) \ |
| template<> \ |
| void Eval<XPU, OP>(const TBlob &lhs, const TBlob &rhs, TBlob *ret, RunContext ctx) { \ |
| FUN<XPU, OP>(lhs, rhs, ret, ctx); \ |
| } |
| #endif |
| |
| #ifndef DECL_BINARY_LAUNCH |
| #define DECL_BINARY_LAUNCH(XPU, OP) \ |
| template <> \ |
| void BinaryOpKernelImpl<OP, XPU>(mshadow::Stream<XPU> *s, \ |
| const TBlob& lhs, const TBlob& rhs, TBlob *out) { \ |
| BinaryOpKernelLaunch<OP>(s, lhs, rhs, out); \ |
| } |
| #endif |
| |
| #ifndef DECL_SCALAR |
| #define DECL_SCALAR(XPU, OP, FUN, REVERSE) \ |
| template<> \ |
| void Eval<XPU, OP, REVERSE>(const TBlob &lhs, const real_t &rhs, \ |
| TBlob *ret, RunContext ctx) { \ |
| FUN<XPU, OP, REVERSE>(lhs, rhs, ret, ctx); \ |
| } |
| #endif |
| |
| #if defined(__CUDACC__) |
| #define DEVICE gpu |
| #else |
| #define DEVICE cpu |
| #endif |
| |
| namespace mxnet { |
| namespace ndarray { |
| |
| // true implementation |
| template<typename xpu, typename OP> |
| void EvalBinary_(const TBlob &lhs, const TBlob &rhs, |
| TBlob *ret, RunContext ctx) { |
| using namespace mshadow::expr; |
| mshadow::Stream<xpu> *s = ctx.get_stream<xpu>(); |
| CHECK_EQ(ret->type_flag_, lhs.type_flag_) |
| << "Only support input/output with the same data type"; |
| CHECK_EQ(ret->type_flag_, rhs.type_flag_) |
| << "Only support input/output with the same data type"; |
| MSHADOW_TYPE_SWITCH(ret->type_flag_, DType, { |
| ret->FlatTo2D<xpu, DType>(s) |
| = F<typename OP::mshadow_op>(lhs.FlatTo2D<xpu, DType>(s), |
| rhs.FlatTo2D<xpu, DType>(s)); |
| }); |
| } |
| |
| template<typename xpu, typename OP> |
| void EvalOneHot_(const TBlob &index, const TBlob &rhs, |
| TBlob *ret, RunContext ctx) { |
| LOG(INFO) << "The operator onehot_encode is deprecated; use one_hot instead."; |
| using namespace mshadow::expr; |
| mshadow::Stream<xpu> *s = ctx.get_stream<xpu>(); |
| // TODO(eric): support mixed type encoding, i.e. int index and float rhs. |
| CHECK_EQ(ret->type_flag_, mshadow::default_type_flag) |
| << "one_hot_encode only support float32 as input/output"; |
| CHECK_EQ(rhs.type_flag_, mshadow::default_type_flag) |
| << "one_hot_encode only support float32 as input/output"; |
| CHECK_EQ(index.type_flag_, mshadow::default_type_flag) |
| << "one_hot_encode only support float32 as input/output"; |
| ret->get<xpu, 2, real_t>(s) = |
| one_hot_encode(index.get<xpu, 1, real_t>(s), |
| rhs.shape_[1]); |
| } |
| |
| template<typename xpu, typename OP> |
| void EvalMatChooseRowElem_(const TBlob &lhs, const TBlob &rhs, |
| TBlob *ret, RunContext ctx) { |
| using namespace mshadow::expr; |
| mshadow::Stream<xpu> *s = ctx.get_stream<xpu>(); |
| // TODO(eric): support mixed type choose, i.e. int index and float rhs. |
| CHECK_EQ(ret->type_flag_, mshadow::default_type_flag) |
| << "mat_choose_row_element only support float32 as input/output"; |
| CHECK_EQ(rhs.type_flag_, mshadow::default_type_flag) |
| << "mat_choose_row_element only support float32 as input/output"; |
| CHECK_EQ(lhs.type_flag_, mshadow::default_type_flag) |
| << "mat_choose_row_element only support float32 as input/output"; |
| ret->get<xpu, 1, real_t>(s) |
| = mat_choose_row_element(lhs.get<xpu, 2, real_t>(s), |
| rhs.get<xpu, 1, real_t>(s)); |
| } |
| |
| template<typename xpu, typename OP> |
| void EvalMatFillRowElem_(const TBlob &lhs, const TBlob &mhs, const TBlob &rhs, |
| TBlob *ret, RunContext ctx) { |
| using namespace mshadow::expr; |
| mshadow::Stream<xpu> *s = ctx.get_stream<xpu>(); |
| ret->get<xpu, 2, real_t>(s) |
| = mat_fill_row_element(lhs.get<xpu, 2, real_t>(s), |
| mhs.get<xpu, 1, real_t>(s), |
| rhs.get<xpu, 1, real_t>(s)); |
| } |
| |
| template<typename xpu, typename OP, bool reverse> |
| void EvalScalar_(const TBlob &lhs, const real_t &rhs, |
| TBlob *ret, RunContext ctx) { |
| using namespace mshadow::expr; |
| mshadow::Stream<xpu> *s = ctx.get_stream<xpu>(); |
| CHECK_EQ(ret->type_flag_, lhs.type_flag_) |
| << "Only support input/output with the same data type"; |
| if (reverse) { |
| MSHADOW_TYPE_SWITCH(ret->type_flag_, DType, { |
| ret->FlatTo2D<xpu, DType>(s) |
| = F<typename OP::mshadow_op>(scalar(DType(rhs)), lhs.FlatTo2D<xpu, DType>(s)); |
| }); |
| } else { |
| MSHADOW_TYPE_SWITCH(ret->type_flag_, DType, { |
| ret->FlatTo2D<xpu, DType>(s) |
| = F<typename OP::mshadow_op>(lhs.FlatTo2D<xpu, DType>(s), scalar(DType(rhs))); |
| }); |
| } |
| } |
| |
| template<> |
| void EvalClip<DEVICE>(const TBlob &src, const real_t &a_min, const real_t &a_max, |
| TBlob *ret, RunContext ctx) { |
| typedef DEVICE xpu; |
| using namespace mshadow::expr; |
| mshadow::Stream<xpu> *s = ctx.get_stream<xpu>(); |
| CHECK_EQ(ret->type_flag_, src.type_flag_) |
| << "Only support input/output with the same data type"; |
| MSHADOW_TYPE_SWITCH(ret->type_flag_, DType, { |
| ret->FlatTo2D<xpu, DType>(s) |
| = F<ClipMax::mshadow_op>( |
| F<ClipMin::mshadow_op>(src.FlatTo2D<xpu, DType>(s), scalar(DType(a_min))), |
| scalar(DType(a_max))); |
| }); |
| } |
| |
| template<> |
| void EvalRandom<DEVICE, UniformDistribution>(const real_t &a, |
| const real_t &b, |
| const Resource &resource, |
| TBlob *ret, |
| RunContext ctx) { |
| typedef DEVICE xpu; |
| mshadow::Stream<xpu> *s = ctx.get_stream<xpu>(); |
| switch (ret->type_flag_) { |
| case mshadow::kFloat32: |
| { |
| mshadow::Random<xpu, float> *prnd = resource.get_random<xpu, float>(s); |
| mshadow::Tensor<xpu, 2, float> tmp = ret->FlatTo2D<xpu, float>(s); |
| prnd->SampleUniform(&tmp, float(a), float(b)); // NOLINT(*) |
| break; |
| } |
| case mshadow::kFloat64: |
| { |
| mshadow::Random<xpu, double> *prnd = resource.get_random<xpu, double>(s); |
| mshadow::Tensor<xpu, 2, double> tmp = ret->FlatTo2D<xpu, double>(s); |
| prnd->SampleUniform(&tmp, double(a), double(b)); // NOLINT(*) |
| break; |
| } |
| default: |
| LOG(FATAL) << "Random only support float32 and float64"; |
| } |
| } |
| |
| template<> |
| void EvalRandom<DEVICE, GaussianDistribution>( |
| const real_t &mu, |
| const real_t &sigma, |
| const Resource &resource, |
| TBlob *ret, |
| RunContext ctx) { |
| typedef DEVICE xpu; |
| mshadow::Stream<xpu> *s = ctx.get_stream<xpu>(); |
| switch (ret->type_flag_) { |
| case mshadow::kFloat32: |
| { |
| mshadow::Random<xpu, float> *prnd = resource.get_random<xpu, float>(s); |
| mshadow::Tensor<xpu, 2, float> tmp = ret->FlatTo2D<xpu, float>(s); |
| prnd->SampleGaussian(&tmp, float(mu), float(sigma)); // NOLINT(*) |
| break; |
| } |
| case mshadow::kFloat64: |
| { |
| mshadow::Random<xpu, double> *prnd = resource.get_random<xpu, double>(s); |
| mshadow::Tensor<xpu, 2, double> tmp = ret->FlatTo2D<xpu, double>(s); |
| prnd->SampleGaussian(&tmp, double(mu), double(sigma)); // NOLINT(*) |
| break; |
| } |
| default: |
| LOG(FATAL) << "Random only support float32 and float64"; |
| } |
| } |
| |
| template<> |
| void EvalRandom<DEVICE, GammaDistribution>( |
| const real_t &alpha, |
| const real_t &beta, |
| const Resource &resource, |
| TBlob *ret, |
| RunContext ctx) { |
| typedef cpu xpu; // No support for gpu for this distribution. |
| mshadow::Stream<xpu> *s = ctx.get_stream<xpu>(); |
| switch (ret->type_flag_) { |
| case mshadow::kFloat32: |
| { |
| mshadow::Random<xpu, float> *prnd = resource.get_random<xpu, float>(s); |
| mshadow::Tensor<xpu, 2, float> tmp = ret->FlatTo2D<xpu, float>(s); |
| prnd->SampleGamma(&tmp, float(alpha), float(beta)); // NOLINT(*) |
| break; |
| } |
| case mshadow::kFloat64: |
| { |
| mshadow::Random<xpu, double> *prnd = resource.get_random<xpu, double>(s); |
| mshadow::Tensor<xpu, 2, double> tmp = ret->FlatTo2D<xpu, double>(s); |
| prnd->SampleGamma(&tmp, double(alpha), double(beta)); // NOLINT(*) |
| break; |
| } |
| default: |
| LOG(FATAL) << "Random only support float32 and float64"; |
| } |
| } |
| |
| |
| template<> |
| void EvalRandom<DEVICE, ExponentialDistribution>( |
| const real_t &lambda, |
| const real_t &dummy, // this is to satisfy the SampleOp lambda signature |
| const Resource &resource, |
| TBlob *ret, |
| RunContext ctx) { |
| typedef cpu xpu; // No support for gpu for this distribution. |
| mshadow::Stream<xpu> *s = ctx.get_stream<xpu>(); |
| switch (ret->type_flag_) { |
| case mshadow::kFloat32: |
| { |
| mshadow::Random<xpu, float> *prnd = resource.get_random<xpu, float>(s); |
| mshadow::Tensor<xpu, 2, float> tmp = ret->FlatTo2D<xpu, float>(s); |
| prnd->SampleExponential(&tmp, float(lambda)); // NOLINT(*) |
| break; |
| } |
| case mshadow::kFloat64: |
| { |
| mshadow::Random<xpu, double> *prnd = resource.get_random<xpu, double>(s); |
| mshadow::Tensor<xpu, 2, double> tmp = ret->FlatTo2D<xpu, double>(s); |
| prnd->SampleExponential(&tmp, double(lambda)); // NOLINT(*) |
| break; |
| } |
| default: |
| LOG(FATAL) << "Random only support float32 and float64"; |
| } |
| } |
| |
| template<> |
| void EvalRandom<DEVICE, PoissonDistribution>( |
| const real_t &lambda, |
| const real_t &dummy, // this is to satisfy the SampleOp lambda signature |
| const Resource &resource, |
| TBlob *ret, |
| RunContext ctx) { |
| typedef cpu xpu; // No support for gpu for this distribution. |
| mshadow::Stream<xpu> *s = ctx.get_stream<xpu>(); |
| switch (ret->type_flag_) { |
| case mshadow::kFloat32: |
| { |
| mshadow::Random<xpu, float> *prnd = resource.get_random<xpu, float>(s); |
| mshadow::Tensor<xpu, 2, float> tmp = ret->FlatTo2D<xpu, float>(s); |
| prnd->SamplePoisson(&tmp, float(lambda)); // NOLINT(*) |
| break; |
| } |
| case mshadow::kFloat64: |
| { |
| mshadow::Random<xpu, double> *prnd = resource.get_random<xpu, double>(s); |
| mshadow::Tensor<xpu, 2, double> tmp = ret->FlatTo2D<xpu, double>(s); |
| prnd->SamplePoisson(&tmp, double(lambda)); // NOLINT(*) |
| break; |
| } |
| default: |
| LOG(FATAL) << "Random only support float32 and float64"; |
| } |
| } |
| |
| template<> |
| void EvalRandom<DEVICE, NegBinomialDistribution>( |
| const real_t &k, |
| const real_t &p, |
| const Resource &resource, |
| TBlob *ret, |
| RunContext ctx) { |
| typedef cpu xpu; // No support for gpu for this distribution. |
| mshadow::Stream<xpu> *s = ctx.get_stream<xpu>(); |
| switch (ret->type_flag_) { |
| case mshadow::kFloat32: |
| { |
| mshadow::Random<xpu, float> *prnd = resource.get_random<xpu, float>(s); |
| mshadow::Tensor<xpu, 2, float> tmp = ret->FlatTo2D<xpu, float>(s); |
| prnd->SampleNegativeBinomial(&tmp, float(k), float(p)); // NOLINT(*) |
| break; |
| } |
| case mshadow::kFloat64: |
| { |
| mshadow::Random<xpu, double> *prnd = resource.get_random<xpu, double>(s); |
| mshadow::Tensor<xpu, 2, double> tmp = ret->FlatTo2D<xpu, double>(s); |
| prnd->SampleNegativeBinomial(&tmp, double(k), double(p)); // NOLINT(*) |
| break; |
| } |
| default: |
| LOG(FATAL) << "Random only support float32 and float64"; |
| } |
| } |
| |
| template<> |
| void EvalRandom<DEVICE, GenNegBinomialDistribution>( |
| const real_t &mu, |
| const real_t &alpha, |
| const Resource &resource, |
| TBlob *ret, |
| RunContext ctx) { |
| typedef cpu xpu; // No support for gpu for this distribution. |
| mshadow::Stream<xpu> *s = ctx.get_stream<xpu>(); |
| switch (ret->type_flag_) { |
| case mshadow::kFloat32: |
| { |
| mshadow::Random<xpu, float> *prnd = resource.get_random<xpu, float>(s); |
| mshadow::Tensor<xpu, 2, float> tmp = ret->FlatTo2D<xpu, float>(s); |
| prnd->SampleGeneralizedNegativeBinomial(&tmp, float(mu), float(alpha)); // NOLINT(*) |
| break; |
| } |
| case mshadow::kFloat64: |
| { |
| mshadow::Random<xpu, double> *prnd = resource.get_random<xpu, double>(s); |
| mshadow::Tensor<xpu, 2, double> tmp = ret->FlatTo2D<xpu, double>(s); |
| prnd->SampleGeneralizedNegativeBinomial(&tmp, double(mu), double(alpha)); // NOLINT(*) |
| break; |
| } |
| default: |
| LOG(FATAL) << "Random only support float32 and float64"; |
| } |
| } |
| |
| template<> |
| void Eval<DEVICE>(const real_t &rhs, TBlob *ret, RunContext ctx) { |
| mshadow::Stream<DEVICE> *s = ctx.get_stream<DEVICE>(); |
| MSHADOW_TYPE_SWITCH(ret->type_flag_, DType, { |
| ret->FlatTo2D<DEVICE, DType>(s) = DType(rhs); |
| }); |
| } |
| |
| template<> |
| void ElementwiseSum<DEVICE>(const std::vector<TBlob> source, |
| TBlob *dst, |
| RunContext ctx) { |
| typedef DEVICE xpu; |
| using namespace mshadow; |
| using namespace mshadow::expr; |
| Stream<xpu> *s = ctx.get_stream<xpu>(); |
| for (size_t i = 1; i < source.size(); ++i) { |
| CHECK_EQ(source[i].type_flag_, dst->type_flag_) |
| << "Only support input/output with the same data type"; |
| } |
| MSHADOW_TYPE_SWITCH(dst->type_flag_, DType, { |
| Tensor<xpu, 2, DType> out = dst->FlatTo2D<xpu, DType>(s); |
| |
| switch (source.size()) { |
| case 2: { |
| Tensor<xpu, 2, DType> in_0 = source[0].FlatTo2D<xpu, DType>(s); |
| Tensor<xpu, 2, DType> in_1 = source[1].FlatTo2D<xpu, DType>(s); |
| out = in_0 + in_1; |
| break; |
| } |
| case 3: { |
| Tensor<xpu, 2, DType> in_0 = source[0].FlatTo2D<xpu, DType>(s); |
| Tensor<xpu, 2, DType> in_1 = source[1].FlatTo2D<xpu, DType>(s); |
| Tensor<xpu, 2, DType> in_2 = source[2].FlatTo2D<xpu, DType>(s); |
| out = in_0 + in_1 + in_2; |
| break; |
| } |
| case 4: { |
| Tensor<xpu, 2, DType> in_0 = source[0].FlatTo2D<xpu, DType>(s); |
| Tensor<xpu, 2, DType> in_1 = source[1].FlatTo2D<xpu, DType>(s); |
| Tensor<xpu, 2, DType> in_2 = source[2].FlatTo2D<xpu, DType>(s); |
| Tensor<xpu, 2, DType> in_3 = source[3].FlatTo2D<xpu, DType>(s); |
| out = in_0 + in_1 + in_2 + in_3; |
| break; |
| } |
| default: { |
| Tensor<xpu, 2, DType> in_0 = source[0].FlatTo2D<xpu, DType>(s); |
| out = F<op::mshadow_op::identity>(in_0); |
| for (size_t i = 1; i < source.size(); ++i) { |
| out += source[i].FlatTo2D<xpu, DType>(s); |
| } |
| break; |
| } |
| } |
| }); |
| } |
| |
| template <> |
| void EvalBroadcast<DEVICE>(TBlob const& src, TBlob* ret, int size, RunContext ctx) { |
| typedef DEVICE xpu; |
| mshadow::Stream<xpu>* s = ctx.get_stream<xpu>(); |
| mshadow::Tensor<xpu, 3> out = ret->get<xpu, 3, real_t>(s); |
| mshadow::Tensor<xpu, 2> in = src.get<xpu, 2, real_t>(s); |
| out = mshadow::expr::broadcast_with_axis(in, 0, size); |
| } |
| |
| template<typename OP, typename xpu> |
| void BinaryOpKernelLaunch(mshadow::Stream<xpu>* s, const TBlob& lhs, const TBlob& rhs, TBlob *out) { |
| using namespace op::mxnet_op; |
| using namespace mshadow; |
| MSHADOW_TYPE_SWITCH(out->type_flag_, DType, { |
| Kernel<op_with_req<OP, kWriteInplace>, xpu >:: |
| Launch(s, |
| lhs.Size(), |
| out->dptr<DType>(), |
| lhs.dptr<DType>(), |
| rhs.dptr<DType>()); |
| }); |
| } |
| // declarations |
| DECL_BINARY(DEVICE, MatChooseRowElem, EvalMatChooseRowElem_) |
| DECL_TERNARY(DEVICE, MatFillRowElem, EvalMatFillRowElem_) |
| DECL_BINARY(DEVICE, OneHotEncode, EvalOneHot_) |
| DECL_SCALAR(DEVICE, Plus, EvalScalar_, true) |
| DECL_SCALAR(DEVICE, Minus, EvalScalar_, true) |
| DECL_SCALAR(DEVICE, Mul, EvalScalar_, true) |
| DECL_SCALAR(DEVICE, Div, EvalScalar_, true) |
| DECL_BINARY_LAUNCH(DEVICE, Plus) |
| DECL_BINARY_LAUNCH(DEVICE, Minus) |
| DECL_BINARY_LAUNCH(DEVICE, Mul) |
| DECL_BINARY_LAUNCH(DEVICE, Div) |
| |
| // for reverse seq |
| DECL_SCALAR(DEVICE, Plus, EvalScalar_, false) |
| DECL_SCALAR(DEVICE, Minus, EvalScalar_, false) |
| DECL_SCALAR(DEVICE, Mul, EvalScalar_, false) |
| DECL_SCALAR(DEVICE, Div, EvalScalar_, false) |
| } // namespace ndarray |
| } // namespace mxnet |
| |
| #endif // MXNET_NDARRAY_NDARRAY_FUNCTION_INL_H_ |
