Fix write-after-free of the device memory accountant at exit Backports: https://github.com/ROCm/rocm-libraries/commit/d2ef3e645eb93bd28f01b022209e3de9d80af46f --- a/library/CMakeLists.txt +++ b/library/CMakeLists.txt @@ -52,6 +52,15 @@ source_group( "Header Files\\Public" FILES ${hipfft_headers_public} ) # Include sources include( src/CMakeLists.txt ) +# Add a definition for internal library targets so code can be +# conditional on being in-library or not +function( rocfft_add_internal_build_flag target ) + get_target_property( target_is_imported ${target} IMPORTED ) + if( NOT target_is_imported ) + target_compile_definitions( ${target} PRIVATE ROCFFT_BUILD_INTERNAL ) + endif() +endfunction() + # Create hipFFT library (and hipfftw on ROCm) add_library( hipfft ${hipfft_source} ${hipfft_headers_public} ) add_library( hip::hipfft ALIAS hipfft ) @@ -89,6 +98,10 @@ else() list( APPEND HIPFFT_EXPORT_TARGETS hip::hipfftw ) # hipfftw.h header uses the same hipfft-export.h as hipfft.h target_compile_definitions( hipfftw PRIVATE hipfft_EXPORTS ) + + rocfft_add_internal_build_flag( hipfft ) + rocfft_add_internal_build_flag( hipfftw ) + endif() if( WIN32 ) --- a/library/src/amd_detail/hipfft.cpp +++ b/library/src/amd_detail/hipfft.cpp @@ -1098,10 +1098,6 @@ catch(hipfftResult e) { return e; } -catch(const DEVICEBUF_MEM_USAGE& e) -{ - return HIPFFT_ALLOC_FAILED; -} catch(const std::exception& e) { return HIPFFT_INTERNAL_ERROR; --- a/shared/gpubuf.h +++ b/shared/gpubuf.h @@ -31,6 +31,8 @@ #include #include +#ifndef ROCFFT_BUILD_INTERNAL + struct DEVICEBUF_MEM_USAGE : public std::runtime_error { using std::runtime_error::runtime_error; @@ -190,146 +192,167 @@ private: } }; -// Simple RAII class for GPU buffers. T is the type of pointer that -// data() returns -template -class gpubuf_t -{ -public: - gpubuf_t() {} - // buffers are movable but not copyable - gpubuf_t(gpubuf_t&& other) - { - std::swap(buf, other.buf); - std::swap(owned, other.owned); - std::swap(bsize, other.bsize); - std::swap(device, other.device); - std::swap(is_managed_memory, other.is_managed_memory); - } - gpubuf_t& operator=(gpubuf_t&& other) - { - std::swap(buf, other.buf); - std::swap(owned, other.owned); - std::swap(bsize, other.bsize); - std::swap(device, other.device); - std::swap(is_managed_memory, other.is_managed_memory); - return *this; - } - gpubuf_t(const gpubuf_t&) = delete; - gpubuf_t& operator=(const gpubuf_t&) = delete; - - static gpubuf_t make_nonowned(T* p, size_t size_bytes = 0) - { - gpubuf_t ret; - ret.owned = false; - ret.buf = p; - ret.bsize = size_bytes; - ret.is_managed_memory = false; // irrelevant if not owned - return ret; - } +#endif - ~gpubuf_t() - { - free(); - } +#ifdef ROCFFT_BUILD_INTERNAL +// Put the internal version of gpubuf_t in a different namespace - +// the intention is for only one to be used in any module, but defend +// against ODR violations if both are used simultaneously. +inline namespace rocfft_internal +{ +#endif - static bool use_alloc_managed() + // Simple RAII class for GPU buffers. T is the type of pointer that + // data() returns + template + class gpubuf_t { - return std::getenv("ROCFFT_MALLOC_MANAGED"); - } + public: + gpubuf_t() {} + // buffers are movable but not copyable + gpubuf_t(gpubuf_t&& other) + { + std::swap(buf, other.buf); + std::swap(owned, other.owned); + std::swap(bsize, other.bsize); + std::swap(device, other.device); + std::swap(is_managed_memory, other.is_managed_memory); + } + gpubuf_t& operator=(gpubuf_t&& other) + { + std::swap(buf, other.buf); + std::swap(owned, other.owned); + std::swap(bsize, other.bsize); + std::swap(device, other.device); + std::swap(is_managed_memory, other.is_managed_memory); + return *this; + } + gpubuf_t(const gpubuf_t&) = delete; + gpubuf_t& operator=(const gpubuf_t&) = delete; - hipError_t alloc(const size_t size, bool make_it_shared = false) - { - free(); - // remember the device that was current as of alloc, so we can - // free on the correct device - auto ret = hipGetDevice(&device); - if(ret != hipSuccess) + static gpubuf_t make_nonowned(T* p, size_t size_bytes = 0) + { + gpubuf_t ret; + ret.owned = false; + ret.buf = p; + ret.bsize = size_bytes; + ret.is_managed_memory = false; // irrelevant if not owned return ret; + } - if(size > device_memory_accountant::singleton().get_usable_bytes(device)) + ~gpubuf_t() { - std::stringstream msg; - msg << "Unauthorized device allocation (device ID: " << device << ").\n" - << "\tRequested size is " << byte_size_to_str(size) << "\n" - << device_memory_accountant::singleton().get_details(device); - throw DEVICEBUF_MEM_USAGE{msg.str()}; + free(); } - bsize = size; - is_managed_memory = use_alloc_managed() || make_it_shared; - ret = is_managed_memory ? hipMallocManaged(&buf, bsize) : hipMalloc(&buf, bsize); - if(ret != hipSuccess) + static bool use_alloc_managed() { - buf = nullptr; - bsize = 0; + return std::getenv("ROCFFT_MALLOC_MANAGED"); } - device_memory_accountant::singleton().record_used_bytes(bsize, device); + hipError_t alloc(const size_t size, bool make_it_shared = false) + { + free(); + // remember the device that was current as of alloc, so we can + // free on the correct device + auto ret = hipGetDevice(&device); + if(ret != hipSuccess) + return ret; + +#ifndef ROCFFT_BUILD_INTERNAL + if(size > device_memory_accountant::singleton().get_usable_bytes(device)) + { + std::stringstream msg; + msg << "Unauthorized device allocation (device ID: " << device << ").\n" + << "\tRequested size is " << byte_size_to_str(size) << "\n" + << device_memory_accountant::singleton().get_details(device); + throw DEVICEBUF_MEM_USAGE{msg.str()}; + } +#endif - return ret; - } + bsize = size; + is_managed_memory = use_alloc_managed() || make_it_shared; + ret = is_managed_memory ? hipMallocManaged(&buf, bsize) : hipMalloc(&buf, bsize); + if(ret != hipSuccess) + { + buf = nullptr; + bsize = 0; + } - size_t size() const - { - return bsize; - } +#ifndef ROCFFT_BUILD_INTERNAL + device_memory_accountant::singleton().record_used_bytes(bsize, device); +#endif - void free() - { - if(buf != nullptr) + return ret; + } + + size_t size() const { - if(owned) + return bsize; + } + + void free() + { + if(buf != nullptr) { - // free on the device we allocated on - rocfft_scoped_device dev(device); - (void)hipFree(buf); - device_memory_accountant::singleton().release_used_bytes(bsize, device); + if(owned) + { + // free on the device we allocated on + rocfft_scoped_device dev(device); + (void)hipFree(buf); +#ifndef ROCFFT_BUILD_INTERNAL + device_memory_accountant::singleton().release_used_bytes(bsize, device); +#endif + } + buf = nullptr; + bsize = 0; } - buf = nullptr; - bsize = 0; + owned = true; } - owned = true; - } - // return a pointer to the allocated memory, offset by the - // specified number of bytes - T* data_offset(size_t offset_bytes = 0) const - { - void* ptr = static_cast(buf) + offset_bytes; - return static_cast(ptr); - } + // return a pointer to the allocated memory, offset by the + // specified number of bytes + T* data_offset(size_t offset_bytes = 0) const + { + void* ptr = static_cast(buf) + offset_bytes; + return static_cast(ptr); + } - T* data() const - { - return static_cast(buf); - } + T* data() const + { + return static_cast(buf); + } - // equality/bool tests - bool operator==(std::nullptr_t n) const - { - return buf == n; - } - bool operator!=(std::nullptr_t n) const - { - return buf != n; - } - operator bool() const - { - return buf; - } + // equality/bool tests + bool operator==(std::nullptr_t n) const + { + return buf == n; + } + bool operator!=(std::nullptr_t n) const + { + return buf != n; + } + operator bool() const + { + return buf; + } -private: - // The GPU buffer - void* buf = nullptr; - // whether this object owns the 'buf' pointer (and hence needs to - // free it) - bool owned = true; - bool is_managed_memory = false; - size_t bsize = 0; - int device = 0; -}; + private: + // The GPU buffer + void* buf = nullptr; + // whether this object owns the 'buf' pointer (and hence needs to + // free it) + bool owned = true; + bool is_managed_memory = false; + size_t bsize = 0; + int device = 0; + }; + + // default gpubuf that gives out void* pointers + typedef gpubuf_t<> gpubuf; + +#ifdef ROCFFT_BUILD_INTERNAL +} +#endif -// default gpubuf that gives out void* pointers -typedef gpubuf_t<> gpubuf; #endif --- a/shared/rocfft_params.h +++ b/shared/rocfft_params.h @@ -405,7 +405,9 @@ public: std::ostringstream oss; oss << "work buffer allocation failed (" << byte_size_to_str(workbuffersizes[device]) << " requested)"; +#ifndef ROCFFT_BUILD_INTERNAL oss << "\n" << device_memory_accountant::singleton().get_details(device); +#endif throw work_buffer_alloc_failure(oss.str(), workbuffersizes[device], hip_status); }