diff --git a/include/mxnet/resource.h b/include/mxnet/resource.h index f8ee6364807c..4ca123b2c850 100644 --- a/include/mxnet/resource.h +++ b/include/mxnet/resource.h @@ -120,16 +120,19 @@ struct Resource { * when running on device, so the launched kernels that depend on the temp space * can finish correctly. * - * \param shape the Shape of returning tensor. - * \param stream the stream of retruning tensor. + * \param shape the Shape of returning tensor. + * \param stream the Stream of returning tensor. + * \param attr_name the Name of the operator requesting the resource. * \return the mshadow tensor requested. * \tparam xpu the device type of random number generator. * \tparam ndim the number of dimension of the tensor requested. */ template inline mshadow::Tensor get_space( - mshadow::Shape shape, mshadow::Stream *stream) const { - return get_space_typed(shape, stream); + mshadow::Shape shape, mshadow::Stream *stream, + std::string &&attr_name=MXNET_STORAGE_HANDLE_DEFAULT_ATTR_NAME( + "temp_space")) const { + return get_space_typed(shape, stream, attr_name); } /*! * \brief Get cpu space requested as mshadow Tensor. @@ -148,33 +151,39 @@ struct Resource { * \brief Get space requested as mshadow Tensor in specified type. * The caller can request arbitrary size. * - * \param shape the Shape of returning tensor. - * \param stream the stream of retruning tensor. + * \param shape the Shape of returning tensor. + * \param stream the Stream of returning tensor. + * \parma attr_name the Name of the operator requesting the resource. * \return the mshadow tensor requested. * \tparam xpu the device type of random number generator. * \tparam ndim the number of dimension of the tensor requested. */ template inline mshadow::Tensor get_space_typed( - mshadow::Shape shape, mshadow::Stream *stream) const { + mshadow::Shape shape, mshadow::Stream *stream, + std::string &&attr_name=MXNET_STORAGE_HANDLE_DEFAULT_ATTR_NAMER( + "temp_space")) const { CHECK_EQ(req.type, ResourceRequest::kTempSpace); return mshadow::Tensor( - reinterpret_cast(get_space_internal(shape.Size() * sizeof(DType))), + reinterpret_cast(get_space_internal( + shape.Size() * sizeof(DType), attr_name)), shape, shape[ndim - 1], stream); } #if MXNET_USE_CUDNN == 1 /*! - * \brief Get cudnn dropout descriptor from shared state space. + * \brief Get cuDNN dropout descriptor from shared state space. * - * \param dropout_desc reference to previously created cudnn dropout descriptor. - * \param stream the stream of retruning tensor. + * \param dropout_desc Reference to previously created cuDNN dropout descriptor. + * \param stream the Stream of returning tensor. + * \param attr_name the Name of the operator requesting the resource. * \return the mshadow tensor requested. */ void get_cudnn_dropout_desc( - cudnnDropoutDescriptor_t* dropout_desc, + cudnnDropoutDescriptor_t *dropout_desc, mshadow::Stream *stream, - const float dropout, - uint64_t seed) const; + const float dropout, uint64_t seed, + std::string &&attr_name=MXNET_STORAGE_HANDLE_DEFAULT_ATTR_NAME( + "cudnn_dropout_state")) const; #endif // MXNET_USE_CUDNN == 1 /*! @@ -195,10 +204,11 @@ struct Resource { } /*! * \brief internal function to get space from resources. - * \param size The size of the space. + * \param size the Size of the space. + * \param attr_name the Name of the operator requesting the resource. * \return The allocated space. */ - void* get_space_internal(size_t size) const; + void* get_space_internal(size_t size, std::string &&attr_name) const; /*! * \brief internal function to get cpu space from resources. * \param size The size of space. diff --git a/include/mxnet/storage.h b/include/mxnet/storage.h index 11e64de36a86..16b7c7be4234 100644 --- a/include/mxnet/storage.h +++ b/include/mxnet/storage.h @@ -44,14 +44,14 @@ inline std::string __extract_fname(const std::string& path) { #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()) + + "_" + __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__) + + "_" + __extract_fname(__FILE__) \ + + "+" + std::to_string(__LINE__) #endif // __GNUG__ /*! @@ -91,12 +91,12 @@ class Storage { * 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 + 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; }; /*! @@ -109,9 +109,9 @@ class Storage { * \return Handle struct. */ Handle Alloc(size_t size, Context ctx, - const std::string& attr_name= + std::string&& attr_name= MXNET_STORAGE_HANDLE_DEFAULT_ATTR_NAME("unknown"), - const std::string& attr_scope="unknown", + std::string&& attr_scope="unknown", const Handle::DataStruct& data_struct= Handle::DataStruct::kUnknown) { Handle hd; diff --git a/src/resource.cc b/src/resource.cc index 3f461243e499..46006d67ebe3 100644 --- a/src/resource.cc +++ b/src/resource.cc @@ -65,11 +65,13 @@ struct SpaceAllocator { host_handle.size = 0; } - inline void* GetSpace(size_t size) { + inline void* GetSpace(size_t size, std::string &&attr_name) { if (handle.size >= size) return handle.dptr; Storage::Get()->DirectFree(handle); - handle = Storage::Get()->Alloc(size, ctx); + handle = Storage::Get()->Alloc(size, ctx, + attr_name, "unknown", + Storage::Handle::DataStruct::kTempSpace); return handle.dptr; } @@ -410,8 +412,9 @@ class ResourceManagerImpl : public ResourceManager { }; } // namespace resource -void* Resource::get_space_internal(size_t size) const { - return static_cast(ptr_)->GetSpace(size); +void* Resource::get_space_internal(size_t size, + std::string &&attr_name) const { + return static_cast(ptr_)->GetSpace(size, attr_name); } void* Resource::get_host_space_internal(size_t size) const { @@ -420,27 +423,29 @@ void* Resource::get_host_space_internal(size_t size) const { #if MXNET_USE_CUDNN == 1 void Resource::get_cudnn_dropout_desc( - cudnnDropoutDescriptor_t* dropout_desc, + cudnnDropoutDescriptor_t *dropout_desc, mshadow::Stream *stream, - const float dropout, - uint64_t seed) const { + const float dropout, uint64_t seed, + std::string &&attr_name) const { CHECK_EQ(req.type, ResourceRequest::kCuDNNDropoutDesc); auto state_space = static_cast(ptr_); CHECK_EQ(state_space->ctx.dev_id, stream->dev_id) - << "The device id of cudnn dropout state space doesn't match that from stream."; + << "The device id of cuDNN dropout state space doesn't match that from stream."; if (!state_space->handle.size) { // not initialized yet. size_t dropout_state_size; CUDNN_CALL(cudnnDropoutGetStatesSize(stream->dnn_handle_, &dropout_state_size)); // reserve GPU space Storage::Get()->DirectFree( - Storage::Get()->Alloc(dropout_state_size, state_space->ctx)); + Storage::Get()->Alloc( + dropout_state_size, + state_space->ctx, "skip me")); CUDNN_CALL(cudnnSetDropoutDescriptor(*dropout_desc, stream->dnn_handle_, dropout, - state_space->GetSpace(dropout_state_size), - dropout_state_size, - seed)); + state_space->GetSpace(dropout_state_size, + attr_name), + dropout_state_size, seed)); } else { // cudnnRestoreDropoutDescriptor() introduced with cuDNN v7 STATIC_ASSERT_CUDNN_VERSION_GE(7000);