diff --git a/include/mxnet/storage.h b/include/mxnet/storage.h index 4d1fc3d2c6da..11e64de36a86 100644 --- a/include/mxnet/storage.h +++ b/include/mxnet/storage.h @@ -26,10 +26,34 @@ #define MXNET_STORAGE_H_ #include +#include #include "./base.h" namespace mxnet { +namespace { +/// \brief Given a path, extract the filename. +inline std::string __extract_fname(const std::string& path) { + std::size_t last_dir_pos = path.find_last_of("/\\"); + if (last_dir_pos == std::string::npos) { + return path; + } + return path.substr(last_dir_pos + 1); +} +} // namespace anonymous + +#if __GNUG__ // if compiled with GCC +#define MXNET_STORAGE_HANDLE_DEFAULT_ATTR_NAME(tag) \ + + ":" + __extract_fname(__FILE__) \ + + " +" + std::to_string(__LINE__) \ + + ":" + __extract_fname(__builtin_FILE()) \ + + " +" + std::to_string(__builtin_LINE()) +#else // !__GNUG__ +#define MXNET_STORAGE_HANDLE_DEFAULT_ATTR_NAME(tag) \ + + ":" + __extract_fname(__FILE__) \ + + " +" + std::to_string(__LINE__) +#endif // __GNUG__ + /*! * \brief Storage manager across multiple devices. */ @@ -55,18 +79,47 @@ class Storage { * \brief Id for IPC shared memory */ int shared_pid{-1}; - int shared_id{-1}; + int shared_id {-1}; + /*! + * \brief Attribute Name & Scope for tracking storage allocations. + */ + std::string attr_name {"unknown"}; + std::string attr_scope{"unknown"}; + /*! + * \brief Data Structure categorizes storage allocations + * based on their functionality. + * It is also used for tracking storage allocations. + */ + enum class DataStruct { + kDataEntry, ///<- Data Entries (!Important) + kTempSpace, ///<- Temporary Workspace + kParameters, ///<- Weight Parameters + kParameterGrads, ///<- Weight Parameter Gradients + kOptimizerStates, ///<- Optimizer States (e.g., Adam Mean & Var) + kAuxStates, ///<- Auxiliary States + kUnknown} data_struct; }; /*! * \brief Allocate a new contiguous memory for a given size. * \param size Total size of memory in bytes. * \param ctx Context information about the device and ID. + * \param attr_name Attribute name. + * \param attr_scope Attribute scope. + * \param data_struct Data structure. * \return Handle struct. */ - Handle Alloc(size_t size, Context ctx) { + Handle Alloc(size_t size, Context ctx, + const std::string& attr_name= + MXNET_STORAGE_HANDLE_DEFAULT_ATTR_NAME("unknown"), + const std::string& attr_scope="unknown", + const Handle::DataStruct& data_struct= + Handle::DataStruct::kUnknown) { Handle hd; hd.size = size; hd.ctx = ctx; + hd.attr_name = attr_name; + hd.attr_scope = attr_scope; + hd.data_struct = data_struct; this->Alloc(&hd); return hd; }