/
githubmirror
/
incubator-mxnet
Обзор
Документация
Войти
/
githubmirror
/
incubator-mxnet
Код
Запросы
0
Пакеты
0
Релизы
0
Аналитика
Безопасность
master
src/operator/operator_common.h
717 строк
25 KB
bgawrych
oneDNN FullyConnected weight caching & refactor (#21047)
06 июл 2022, 13:02
Не верифицирован
06 июл 2022, 13:02
84b1626
Код
Авторство
О чём код?
/* * 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. */ /*! * \file operator_common.h * \brief common internal header of most operators * this header includes utility functions operator can use * \author Bing Xu */ #ifndef MXNET_OPERATOR_OPERATOR_COMMON_H_ #define MXNET_OPERATOR_OPERATOR_COMMON_H_ #include <dmlc/json.h> #include <dmlc/logging.h> #include <dmlc/thread_local.h> #include <mxnet/operator.h> #include <mxnet/ndarray.h> #include <mxnet/op_attr_types.h> #include <mxnet/base.h> #include <istream> #include <ostream> #include <string> #include <vector> #include <algorithm> #include "../common/cuda/utils.h" #include "../common/utils.h" namespace mxnet { namespace op { /*! * \brief assign the expression to out according to request * \param out the data to be assigned * \param req the assignment request * \param exp the expression * \tparam OType output type * \tparam Exp expression type */ #define Assign(out, req, exp) \ { \ switch (req) { \ case kNullOp: \ break; \ case kWriteTo: \ case kWriteInplace: \ (out) = (exp); \ break; \ case kAddTo: \ (out) += (exp); \ break; \ default: \ LOG(FATAL) << "not reached"; \ } \ } /*! \brief exception throwed by InferShape error */ struct InferShapeError : public dmlc::Error { /*! \brief analyze message */ std::string msg; /*! \brief corresponding input index */ int index; // constructor InferShapeError(const std::string& msg_, int index) : dmlc::Error(msg_), msg(msg_), index(index) {} }; /*! \brief exception throwed by InferShape error */ struct InferTypeError : public dmlc::Error { /*! \brief analyze message */ std::string msg; /*! \brief corresponding input index */ int index; // constructor InferTypeError(const std::string& msg_, int index) : dmlc::Error(msg_), msg(msg_), index(index) {} }; /*! \brief exception throwed by InferStorageType error */ struct InferStorageTypeError : public dmlc::Error { /*! \brief analyze message */ std::string msg; /*! \brief corresponding input index */ int index; // constructor InferStorageTypeError(const std::string& msg_, int index) : dmlc::Error(msg_), msg(msg_), index(index) {} }; /*! \brief check if shape is empty or contains unknown (0) dim. * DEPRECATED. */ inline bool shape_is_none(const mxnet::TShape& x) { return !mxnet::shape_is_known(x); } /*! \brief check if type is none (-1) */ inline bool type_is_none(const int& x) { return x == -1; } /*! \brief check if type is none (-1) */ inline bool storage_type_is_none(const int& x) { return x == -1; } /*! \brief check if shape is scalar({1}). */ inline bool shape_is_scalar(const mxnet::TShape& x) { return x.ndim() == 0; } /*! \brief get string representation of shape */ inline std::string shape_string(const mxnet::TShape& x) { std::ostringstream os; os << x; return os.str(); } /*! \brief get string representation of data type */ inline std::string type_string(const int& x) { switch (x) { case mshadow::kFloat32: return "float32"; case mshadow::kFloat64: return "float64"; case mshadow::kFloat16: return "float16"; case mshadow::kBfloat16: return "bfloat16"; case mshadow::kInt8: return "int8"; case mshadow::kUint8: return "uint8"; case mshadow::kInt32: return "int32"; case mshadow::kInt64: return "int64"; } return "unknown"; } /*! * \brief Assign x to y. Checks for compatiblity when y is not empty. * Allow missing dim in both x and y (as 0). * \param y target shape. * \param x source shape. * \return whether x and y are compatible. */ inline bool shape_assign(mxnet::TShape* y, const mxnet::TShape& x) { if (!mxnet::ndim_is_known(*y)) { *y = x; return true; } else if (y->ndim() != x.ndim()) { return !mxnet::ndim_is_known(x); } else { for (int i = 0; i < y->ndim(); ++i) { if (!mxnet::dim_size_is_known(*y, i)) { (*y)[i] = x[i]; } else if ((*y)[i] != x[i] && x[i] >= 0) { return false; } } return true; } } /*! * \brief Assign x to y. Checks for compatiblity when y is not -1. * \param y target type. * \param x source type. * \return whether x and y are compatible. */ inline bool type_assign(int* y, const int& x) { if (*y == -1) { *y = x; return true; } else if (*y != x && x != -1) { return false; } return true; } /*! * \brief Assign x to y. Checks for compatiblity when y is not DispatchMode::kUndefined. * \param y target mode. * \param x source mode. * \return whether x and y are compatible. */ inline bool dispatch_mode_assign(DispatchMode* y, const DispatchMode& x) { if (*y == DispatchMode::kUndefined) { *y = x; return true; } else if (*y != x && x != DispatchMode::kUndefined) { return false; } return true; } /*! \brief Register op name as an alias */ #define MXNET_ADD_SPARSE_OP_ALIAS(__name$) .add_alias("_sparse_" #__name$) /*! * \brief macro assign shape to out if out is unknown otherwise check consistency * Use macro so we can see the error file more clearly * \param shape_array the shape array to store the result * \param index the index of in the array * \param shape the inferred shape */ #define SHAPE_ASSIGN_CHECK(shape_array, index, shape) \ { \ if (!::mxnet::op::shape_assign(&(shape_array)[index], mxnet::TShape(shape))) { \ std::ostringstream os; \ os << "Shape inconsistent, Provided = " << (shape_array)[index] << ',' \ << " inferred shape=" << shape; \ throw ::mxnet::op::InferShapeError(os.str(), index); \ } \ } /*! * \brief macro assign type to out if out is unknown (-1) otherwise check consistency * Use macro so we can see the error file more clearly * \param type_array the type array to store the result * \param index the index of in the array * \param type the inferred type */ #define TYPE_ASSIGN_CHECK(type_array, index, type) \ { \ if (!::mxnet::op::type_assign(&(type_array)[index], type)) { \ std::ostringstream os; \ os << "Type inconsistent, Provided = " << ::mxnet::op::type_string((type_array)[index]) \ << ',' << " inferred type = " << ::mxnet::op::type_string(type); \ throw ::mxnet::op::InferTypeError(os.str(), index); \ } \ } /*! * \brief macro assign storage type to out if out is unknown (-1) otherwise check consistency * Use macro so we can see the error file more clearly * \param type_array the type array to store the result * \param index the index of in the array * \param type the inferred storage type */ #define STORAGE_TYPE_ASSIGN_CHECK(type_array, index, type) \ { \ if (!::mxnet::op::type_assign(&(type_array)[index], type)) { \ std::ostringstream os; \ os << "Storage type inconsistent, Provided = " << common::stype_string((type_array)[index]) \ << ',' << " inferred storage type = " << common::stype_string(type); \ throw ::mxnet::op::InferStorageTypeError(os.str(), index); \ } \ } /*! * \brief macro assign type to out if out is unknown (-1) otherwise check consistency * Use macro so we can see the error file more clearly * \param type_array the type array to store the result * \param index the index of in the array * \param type the inferred dispatch type */ #define DISPATCH_MODE_ASSIGN_CHECK(type_array, index, type) \ { \ if (!::mxnet::op::dispatch_mode_assign(&(type_array)[index], type)) { \ std::ostringstream os; \ os << "Dispatch mode inconsistent, Provided = " \ << common::dispatch_mode_string((type_array)[index]) << ',' \ << " inferred mode = " << common::dispatch_mode_string(type); \ throw ::mxnet::op::InferStorageTypeError(os.str(), index); \ } \ } /*! * \brief macro check if type is the same as expected. * \param type the type to be checked * \param expected the expected type */ #define UNIFORM_TYPE_CHECK(type, expected, arg) \ { \ CHECK_EQ(type, expected) << "This layer requires uniform type. " \ << "Expected '" << ::mxnet::op::type_string(expected) \ << "' v.s. given '" << ::mxnet::op::type_string(type) << "' at '" \ << arg << "'"; \ } // helper macro to implement bind dispatch #if MXNET_USE_CUDA #define DO_BIND_DISPATCH(Method, ...) \ if (ctx.dev_mask() == cpu::kDevMask) { \ return Method<cpu>(__VA_ARGS__); \ } else { \ return Method<gpu>(__VA_ARGS__); \ } #else #define DO_BIND_DISPATCH(Method, ...) \ if (ctx.dev_mask() == cpu::kDevMask) { \ return Method<cpu>(__VA_ARGS__); \ } else { \ LOG(FATAL) << "GPU is not enabled"; \ return nullptr; \ } #endif /*! \brief assign stype to target_stype, if successful, * assign dispatch_mode to target_dispatch */ inline bool storage_type_assign(int* stype, const NDArrayStorageType target_stype, DispatchMode* dispatch, const DispatchMode target_dispatch) { if (type_assign(stype, target_stype)) { DISPATCH_MODE_ASSIGN_CHECK(dispatch, 0, target_dispatch); return true; } return false; } /*! \brief assign the stype vector to target_stype, if successful, * assign dispatch_mode to target_dispatch */ inline bool storage_type_assign(StorageTypeVector* stypes, const NDArrayStorageType target_stype, DispatchMode* dispatch, const DispatchMode target_dispatch) { CHECK_GT(stypes->size(), 0); bool success = true; for (int& stype : *stypes) { if (!type_assign(&stype, target_stype)) { success = false; } } if (success) { DISPATCH_MODE_ASSIGN_CHECK(dispatch, 0, target_dispatch); } return success; } /*! \brief update the stype vector to default storage and dispatch_mode to fallback */ inline bool dispatch_fallback(StorageTypeVector* stypes, DispatchMode* dispatch) { for (auto& stype : *stypes) { type_assign(&stype, kDefaultStorage); } DISPATCH_MODE_ASSIGN_CHECK(dispatch, 0, DispatchMode::kFComputeFallback); return true; } inline std::vector<nnvm::NodeEntry> CreateNodeEntries( nnvm::ObjectPtr pNode, const std::vector<nnvm::NodeEntry>* pOgrads = nullptr, const std::vector<nnvm::NodeEntry>* pInputs = nullptr) { if (pOgrads) pNode->inputs.insert(pNode->inputs.end(), pOgrads->begin(), pOgrads->end()); if (pInputs) pNode->inputs.insert(pNode->inputs.end(), pInputs->begin(), pInputs->end()); if (!pNode->is_variable()) { CHECK_EQ(pNode->num_inputs(), pNode->inputs.size()) << "Number of inputs to operator " << pNode->op()->name << " (" << pNode->num_inputs() << ") does not match the actual number of inputs provided to operator " << pNode->attrs.name << " (" << pNode->inputs.size() << ")."; } std::vector<nnvm::NodeEntry> ret; for (uint32_t i = 0; i < pNode->num_outputs(); ++i) ret.emplace_back(nnvm::NodeEntry{pNode, i, 0}); return ret; } // make a new node with operator op_name. Inputs are not filled. inline nnvm::ObjectPtr MakeNode(const char* op_name, const std::string& name, std::vector<nnvm::NodeEntry> const* inputs = nullptr, std::unordered_map<std::string, std::string> const* dict = nullptr, nnvm::ObjectPtr const* fwd_node = nullptr) { auto p = nnvm::Node::Create(); p->attrs.op = nnvm::Op::Get(op_name); p->attrs.name = name; if (dict != nullptr) p->attrs.dict = *dict; if (inputs != nullptr) p->inputs = *inputs; if (fwd_node != nullptr) { p->control_deps.emplace_back(*fwd_node); } if (p->op()->attr_parser != nullptr) { p->op()->attr_parser(&(p->attrs)); } if (inputs != nullptr) { CHECK_EQ(p->num_inputs(), p->inputs.size()) << "Number of inputs to operator " << op_name << " (" << p->num_inputs() << ") does not match the actual number of inputs provided to operator " << name << " (" << p->inputs.size() << ")."; } return p; } inline nnvm::ObjectPtr MakeNode(const char* op_name, const std::string& name, const std::vector<nnvm::NodeEntry>& inputs, std::unordered_map<std::string, std::string> const* dict, nnvm::ObjectPtr const* fwd_node) { return MakeNode(op_name, name, &inputs, dict, fwd_node); } // quick helper to make node inline std::vector<nnvm::NodeEntry> MakeGradNode( const char* op_name, const nnvm::ObjectPtr& n, const std::vector<nnvm::NodeEntry>& inputs, const std::unordered_map<std::string, std::string>& dict) { auto p = MakeNode(op_name, n->attrs.name + "_backward", &inputs, &dict, &n); return CreateNodeEntries(p); } // quick helper to make gradient nodes that simply pass back zero. could be used in output ops. inline std::vector<nnvm::NodeEntry> MakeZeroGradNodes(const nnvm::ObjectPtr& n, const std::vector<nnvm::NodeEntry>& ograds) { std::vector<nnvm::NodeEntry> ret; for (uint32_t i = 0; i < n->num_inputs(); ++i) { std::ostringstream os; if (1 == n->num_inputs()) { os << n->attrs.name << "_backward"; } else { os << n->attrs.name << "_in" << i << "_backward"; } ret.emplace_back(MakeNode("zeros_like", os.str(), {n->inputs[i]}, nullptr, &n)); } return ret; } // check whether all output grads are zero. inline bool CheckGradAllZero(const std::vector<nnvm::NodeEntry>& ograds) { static const auto zero_op = nnvm::Op::Get("_zeros"); static const auto zero_like_op = nnvm::Op::Get("zeros_like"); if (ograds.empty()) return false; for (const auto& grad : ograds) { if (!grad.node) return false; if (grad.node->op() != zero_op && grad.node->op() != zero_like_op) return false; } return true; } // make gradient node that doesn't add to objective. // i.e. igrads are always zero when ograds are zero. inline std::vector<nnvm::NodeEntry> MakeNonlossGradNode( const char* op_name, const nnvm::ObjectPtr& n, const std::vector<nnvm::NodeEntry>& ograds, const std::vector<nnvm::NodeEntry>& inputs, const std::unordered_map<std::string, std::string>& dict) { if (CheckGradAllZero(ograds)) return MakeZeroGradNodes(n, ograds); auto p = MakeNode(op_name, n->attrs.name + "_backward", nullptr, &dict, &n); return CreateNodeEntries(p, &ograds, &inputs); } /*! \brief Parse keyword arguments as PType arguments and save to parsed */ template <typename PType> inline void ParamParser(nnvm::NodeAttrs* attrs) { PType param; try { param.Init(attrs->dict); } catch (const dmlc::ParamError& e) { std::ostringstream os; os << e.what(); os << ", in operator " << attrs->op->name << "(" << "name=\"" << attrs->name << "\""; for (const auto& k : attrs->dict) { os << ", " << k.first << "=\"" << k.second << "\""; } os << ")"; throw dmlc::ParamError(os.str()); } attrs->parsed = std::move(param); } inline void CheckAllRowsPresent(const NDArray& arr, const std::string& func, const std::string& param) { if (arr.storage_type() == kRowSparseStorage) { CHECK(arr.storage_shape()[0] == arr.shape()[0]) << func << " for RowSparse " << param << " is only implemented for " << "RowSparse " << param << " with all rows containing non-zeros. " << "Expects " << param << ".data.shape[0] (" << arr.storage_shape()[0] << ") == " << param << ".shape[0] (" << arr.shape()[0] << ")."; } else { CHECK(arr.storage_type() == kDefaultStorage); } } inline void LogUnimplementedOp(const nnvm::NodeAttrs& attrs, const OpContext& ctx, const std::vector<NDArray>& inputs, const std::vector<OpReqType>& req, const std::vector<NDArray>& outputs) { using common::operator_string; LOG(FATAL) << "Not implemented: " << operator_string(attrs, ctx, inputs, req, outputs); } class OpSignature { std::vector<int64_t> eles; uint64_t hash; public: OpSignature() { hash = 0; } explicit OpSignature(uint64_t hash) { this->hash = hash; } /* * This is to reserve space for the vector. */ void Reserve(size_t num) { eles.reserve(num); } /* * We provide different methods to add signature to an op. * For operations, such as convolutin and fully connected, which determines * the optimal data layout for the op, we only need to use the shape and data * type to sign the op. For other operations, such as activation, which uses * whatever layout in the input array, we have to use the shape, the data type * and the layout to sign the op. */ #if MXNET_USE_ONEDNN == 1 void AddSign(const dnnl::memory::desc& desc) { hash = hash * 2 + desc.data.format_kind; eles.push_back(desc.data.format_kind); hash = hash * 2 + desc.data.data_type; eles.push_back(desc.data.data_type); for (int i = 0; i < desc.data.ndims; i++) { hash = hash * 2 + desc.data.dims[i]; eles.push_back(desc.data.dims[i]); } switch (desc.data.format_kind) { case dnnl_blocked: hash = hash * 2 + desc.data.ndims; eles.push_back(desc.data.ndims); for (int i = 0; i < desc.data.ndims; i++) { hash = hash * 2 + desc.data.format_desc.blocking.strides[i]; eles.push_back(desc.data.format_desc.blocking.strides[i]); } hash = hash * 2 + desc.data.format_desc.blocking.inner_nblks; eles.push_back(desc.data.format_desc.blocking.inner_nblks); for (int i = 0; i < desc.data.format_desc.blocking.inner_nblks; i++) { hash = hash * 2 + desc.data.format_desc.blocking.inner_blks[i]; hash = hash * 2 + desc.data.format_desc.blocking.inner_idxs[i]; eles.push_back(desc.data.format_desc.blocking.inner_blks[i]); eles.push_back(desc.data.format_desc.blocking.inner_idxs[i]); } break; case dnnl_format_kind_wino: hash = hash * 2 + desc.data.format_desc.wino_desc.wino_format; eles.push_back(desc.data.format_desc.wino_desc.wino_format); break; case dnnl_format_kind_rnn_packed: hash = hash * 2 + desc.data.format_desc.rnn_packed_desc.format; eles.push_back(desc.data.format_desc.rnn_packed_desc.format); hash = hash * 2 + desc.data.format_desc.rnn_packed_desc.n_parts; eles.push_back(desc.data.format_desc.rnn_packed_desc.n_parts); for (int i = 0; i < desc.data.format_desc.rnn_packed_desc.n_parts; ++i) { hash = hash * 2 + desc.data.format_desc.rnn_packed_desc.parts[i]; hash = hash * 2 + desc.data.format_desc.rnn_packed_desc.part_pack_size[i]; hash = hash * 2 + desc.data.format_desc.rnn_packed_desc.pack_part[i]; eles.push_back(desc.data.format_desc.rnn_packed_desc.parts[i]); eles.push_back(desc.data.format_desc.rnn_packed_desc.part_pack_size[i]); eles.push_back(desc.data.format_desc.rnn_packed_desc.pack_part[i]); } hash = hash * 2 + desc.data.format_desc.rnn_packed_desc.n; hash = hash * 2 + desc.data.format_desc.rnn_packed_desc.ldb; hash = hash * 2 + desc.data.format_desc.rnn_packed_desc.offset_compensation; hash = hash * 2 + desc.data.format_desc.rnn_packed_desc.size; eles.push_back(desc.data.format_desc.rnn_packed_desc.n); eles.push_back(desc.data.format_desc.rnn_packed_desc.ldb); eles.push_back(desc.data.format_desc.rnn_packed_desc.offset_compensation); eles.push_back(desc.data.format_desc.rnn_packed_desc.size); break; default: // nothing need to add break; } } void AddSign(const dnnl::memory& mem) { auto desc = mem.get_desc(); AddSign(desc); } void AddSign(const std::vector<dnnl::memory::desc>& arrs) { for (auto& arr : arrs) { AddSign(arr); } } #endif void AddSign(const std::string& s) { uint64_t key = static_cast<uint64_t>(std::hash<std::string>{}(s)); eles.push_back(key); } void AddSign(const std::vector<NDArray>& arrs) { for (auto& arr : arrs) { AddSign(arr); } } void AddSign(const NDArray& arr) { #if MXNET_USE_ONEDNN == 1 if (arr.IsDNNLData()) { AddSign(*(arr.GetDNNLData())); } else { #endif hash = hash * 2 + arr.dtype(); eles.push_back(arr.dtype()); AddSign(arr.shape()); #if MXNET_USE_ONEDNN == 1 } #endif } void AddSign(const mxnet::ShapeVector& shapes) { for (auto& shape : shapes) { AddSign(shape); } } void AddSign(const mxnet::TShape& shape) { for (int i = 0; i < shape.ndim(); i++) { hash = hash * 2 + shape[i]; eles.push_back(shape[i]); } } void AddSign(int val) { hash = hash * 2 + val; eles.push_back(val); } void AddSign(const std::vector<float>& vec) { for (auto& val : vec) { AddSign(val); } } void AddSign(float val) { hash = dmlc::HashCombine(hash, val); eles.push_back(val); } bool operator==(const OpSignature& sign) const { if (hash != sign.hash) return false; if (eles.size() != sign.eles.size()) return false; for (size_t i = 0; i < eles.size(); i++) if (eles[i] != sign.eles[i]) return false; return true; } uint64_t GetHash() const { return hash; } }; struct OpHash { size_t operator()(const OpSignature& sign) const { return sign.GetHash(); } }; template <typename ParamType> class ParamOpSign : public OpSignature { const ParamType param; static size_t hash(const ParamType& param) { std::hash<ParamType> fn; return fn(param); } public: explicit ParamOpSign(const ParamType& _param) : OpSignature(hash(_param)), param(_param) {} bool operator==(const ParamOpSign<ParamType>& sign) const { const OpSignature& this_upper = *this; const OpSignature& other_upper = sign; return this_upper == other_upper && param == sign.param; } }; } // namespace op } // namespace mxnet #endif // MXNET_OPERATOR_OPERATOR_COMMON_H_