@@ -11,12 +11,18 @@ distributed under the License is distributed on an "AS IS" BASIS,
11
11
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
12
See the License for the specific language governing permissions and
13
13
limitations under the License. */
14
+ #define GLOG_NO_ABBREVIATED_SEVERITIES
14
15
15
16
#include " paddle/fluid/memory/detail/system_allocator.h"
16
17
17
- #include < stdlib.h> // for malloc and free
18
+ #ifdef _WIN32
19
+ #include < malloc.h>
20
+ #include < windows.h> // VirtualLock/VirtualUnlock
21
+ #else
18
22
#include < sys/mman.h> // for mlock and munlock
19
- #include < algorithm> // for std::max
23
+ #endif
24
+ #include < stdlib.h> // for malloc and free
25
+ #include < algorithm> // for std::max
20
26
21
27
#include " gflags/gflags.h"
22
28
#include " paddle/fluid/platform/assert.h"
@@ -35,31 +41,42 @@ namespace paddle {
35
41
namespace memory {
36
42
namespace detail {
37
43
38
- void * CPUAllocator::Alloc (size_t * index, size_t size) {
39
- // According to http://www.cplusplus.com/reference/cstdlib/malloc/,
40
- // malloc might not return nullptr if size is zero, but the returned
41
- // pointer shall not be dereferenced -- so we make it nullptr.
42
- if (size <= 0 ) return nullptr ;
43
-
44
- *index = 0 ; // unlock memory
45
-
44
+ void * AlignedMalloc (size_t size) {
46
45
void * p = nullptr ;
47
-
46
+ size_t alignment = 32ul ;
48
47
#ifdef PADDLE_WITH_MKLDNN
49
48
// refer to https://github.com/01org/mkl-dnn/blob/master/include/mkldnn.hpp
50
49
// memory alignment
51
- PADDLE_ENFORCE_EQ (posix_memalign (&p, 4096ul , size), 0 , " Alloc %ld error!" ,
52
- size);
50
+ alignment = 4096ul ;
51
+ #endif
52
+ #ifdef _WIN32
53
+ p = _aligned_malloc (size, alignment);
53
54
#else
54
- PADDLE_ENFORCE_EQ (posix_memalign (&p, 32ul , size), 0 , " Alloc %ld error!" ,
55
+ PADDLE_ENFORCE_EQ (posix_memalign (&p, alignment , size), 0 , " Alloc %ld error!" ,
55
56
size);
56
57
#endif
57
58
PADDLE_ENFORCE (p, " Fail to allocate CPU memory: size = %d ." , size);
59
+ return p;
60
+ }
61
+
62
+ void * CPUAllocator::Alloc (size_t * index, size_t size) {
63
+ // According to http://www.cplusplus.com/reference/cstdlib/malloc/,
64
+ // malloc might not return nullptr if size is zero, but the returned
65
+ // pointer shall not be dereferenced -- so we make it nullptr.
66
+ if (size <= 0 ) return nullptr ;
67
+
68
+ *index = 0 ; // unlock memory
69
+
70
+ void * p = AlignedMalloc (size);
58
71
59
72
if (p != nullptr ) {
60
73
if (FLAGS_use_pinned_memory) {
61
74
*index = 1 ;
75
+ #ifdef _WIN32
76
+ VirtualLock (p, size);
77
+ #else
62
78
mlock (p, size); // lock memory
79
+ #endif
63
80
}
64
81
}
65
82
@@ -68,7 +85,11 @@ void* CPUAllocator::Alloc(size_t* index, size_t size) {
68
85
69
86
void CPUAllocator::Free (void * p, size_t size, size_t index) {
70
87
if (p != nullptr && index == 1 ) {
88
+ #ifdef _WIN32
89
+ VirtualUnlock (p, size);
90
+ #else
71
91
munlock (p, size);
92
+ #endif
72
93
}
73
94
free (p);
74
95
}
0 commit comments