Added member attr_name, attr_scope, and data_struct to the Storage cl…

…ass.

include/mxnet/storage.h

@@ -26,10 +26,34 @@
 #define MXNET_STORAGE_H_
 
 #include <memory>
+#include <string>
 #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;
   }