-
Notifications
You must be signed in to change notification settings - Fork 21.3k
/
Allocator.h
55 lines (44 loc) · 1.55 KB
/
Allocator.h
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
#pragma once
#include <c10/core/CPUAllocator.h>
namespace at {
namespace native {
// QNNPACK AND XNNPACK may out-of-bound access the input and / or output tensors.
// This behavior will trigger ASAN, and may result in a segfault if the accessed
// memory just so happens to fall on a page the current process has no read access
// to. Here we define a custom allocator that allocates the extra storage required
// to keep this behavior safe.
//
// PreGuardBytes: Number of guard bytes to allocate before the allocation.
// PostGuardBytes: Number of guard bytes to allocate after the allocation.
template <uint32_t PreGuardBytes, uint32_t PostGuardBytes>
class GuardingAllocator final : public at::Allocator {
public:
GuardingAllocator() = default;
virtual ~GuardingAllocator() override = default;
static void deleter(void* pointer) {
const Cast memory{pointer};
c10::free_cpu(memory.as_byte_ptr - kPreGuardBytes);
}
virtual DataPtr allocate(size_t nbytes) const override {
Cast memory{c10::alloc_cpu(kPreGuardBytes + nbytes + kPostGuardBytes)};
memory.as_byte_ptr += kPreGuardBytes;
return {
memory.as_void_ptr,
memory.as_void_ptr,
&deleter,
at::Device(DeviceType::CPU),
};
}
virtual DeleterFnPtr raw_deleter() const override {
return deleter;
}
private:
static constexpr uint32_t kPreGuardBytes = PreGuardBytes;
static constexpr uint32_t kPostGuardBytes = PostGuardBytes;
union Cast final {
void * const as_void_ptr;
uint8_t * as_byte_ptr;
};
};
} // namespace native
} // namespace at