@ -11,12 +11,18 @@ distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND , either express or implied . WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND , either express or implied .
See the License for the specific language governing permissions and See the License for the specific language governing permissions and
limitations under the License . */ limitations under the License . */
# define GLOG_NO_ABBREVIATED_SEVERITIES
# include "paddle/fluid/memory/detail/system_allocator.h" # include "paddle/fluid/memory/detail/system_allocator.h"
# include <stdlib.h> // for malloc and free # ifdef _WIN32
# include <malloc.h>
# include <windows.h> // VirtualLock/VirtualUnlock
# else
# include <sys/mman.h> // for mlock and munlock # include <sys/mman.h> // for mlock and munlock
# include <algorithm> // for std::max # endif
# include <stdlib.h> // for malloc and free
# include <algorithm> // for std::max
# include "gflags/gflags.h" # include "gflags/gflags.h"
# include "paddle/fluid/platform/assert.h" # include "paddle/fluid/platform/assert.h"
@ -35,31 +41,42 @@ namespace paddle {
namespace memory { namespace memory {
namespace detail { namespace detail {
void * CPUAllocator : : Alloc ( size_t * index , size_t size ) { void * AlignedMalloc ( size_t size ) {
// According to http://www.cplusplus.com/reference/cstdlib/malloc/,
// malloc might not return nullptr if size is zero, but the returned
// pointer shall not be dereferenced -- so we make it nullptr.
if ( size < = 0 ) return nullptr ;
* index = 0 ; // unlock memory
void * p = nullptr ;
void * p = nullptr ;
size_t alignment = 32ul ;
# ifdef PADDLE_WITH_MKLDNN # ifdef PADDLE_WITH_MKLDNN
// refer to https://github.com/01org/mkl-dnn/blob/master/include/mkldnn.hpp
// refer to https://github.com/01org/mkl-dnn/blob/master/include/mkldnn.hpp
// memory alignment
// memory alignment
PADDLE_ENFORCE_EQ ( posix_memalign ( & p , 4096ul , size ) , 0 , " Alloc %ld error! " ,
alignment = 4096ul ;
size ) ;
# endif
# ifdef _WIN32
p = _aligned_malloc ( size , alignment ) ;
# else # else
PADDLE_ENFORCE_EQ ( posix_memalign ( & p , 32ul , size ) , 0 , " Alloc %ld error! " ,
PADDLE_ENFORCE_EQ ( posix_memalign ( & p , alignment , size ) , 0 , " Alloc %ld error! " ,
size ) ;
size ) ;
# endif # endif
PADDLE_ENFORCE ( p , " Fail to allocate CPU memory: size = %d . " , size ) ;
PADDLE_ENFORCE ( p , " Fail to allocate CPU memory: size = %d . " , size ) ;
return p ;
}
void * CPUAllocator : : Alloc ( size_t * index , size_t size ) {
// According to http://www.cplusplus.com/reference/cstdlib/malloc/,
// malloc might not return nullptr if size is zero, but the returned
// pointer shall not be dereferenced -- so we make it nullptr.
if ( size < = 0 ) return nullptr ;
* index = 0 ; // unlock memory
void * p = AlignedMalloc ( size ) ;
if ( p ! = nullptr ) {
if ( p ! = nullptr ) {
if ( FLAGS_use_pinned_memory ) {
if ( FLAGS_use_pinned_memory ) {
* index = 1 ;
* index = 1 ;
# ifdef _WIN32
VirtualLock ( p , size ) ;
# else
mlock ( p , size ) ; // lock memory
mlock ( p , size ) ; // lock memory
# endif
}
}
}
}
@ -68,7 +85,11 @@ void* CPUAllocator::Alloc(size_t* index, size_t size) {
void CPUAllocator : : Free ( void * p , size_t size , size_t index ) { void CPUAllocator : : Free ( void * p , size_t size , size_t index ) {
if ( p ! = nullptr & & index = = 1 ) {
if ( p ! = nullptr & & index = = 1 ) {
# ifdef _WIN32
VirtualUnlock ( p , size ) ;
# else
munlock ( p , size ) ;
munlock ( p , size ) ;
# endif
}
}
free ( p ) ;
free ( p ) ;
} }
tensor-tang
7 years ago