-
Notifications
You must be signed in to change notification settings - Fork 21.6k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Mobile Backend: NHWC memory layout + XNNPACK integration.
- Loading branch information
Ashkan Aliabadi
committed
Feb 18, 2020
1 parent
4468a7b
commit 4b95293
Showing
22 changed files
with
1,081 additions
and
170 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,56 @@ | ||
| #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 | ||
|
|
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,82 @@ | ||
| #pragma once | ||
|
|
||
| #include <ATen/ATen.h> | ||
|
|
||
| #ifdef USE_XNNPACK | ||
|
|
||
| #include <xnnpack.h> | ||
|
|
||
| namespace at { | ||
| namespace native { | ||
| namespace xnnpack { | ||
| namespace internal { | ||
|
|
||
| struct Layout final { | ||
| // 4D Activation Maps | ||
| struct Activation4D final { | ||
| static constexpr size_t batch = 0u; | ||
| static constexpr size_t channels = 1u; | ||
| static constexpr size_t height = 2u; | ||
| static constexpr size_t width = 3u; | ||
| }; | ||
|
|
||
| // ND Activation Maps | ||
| struct ActivationND final { | ||
| // Some operators may not be limited to 4 dimensional tensors. In that scenario, | ||
| // XNNPACK denotes that operator with an _nc suffix and expects all dimensions, | ||
| // except channels, to be flattened into one argument: batch_size. | ||
| static int64_t batch(const IntArrayRef tensor) { | ||
| if (C10_UNLIKELY(tensor.empty())) { | ||
| return -1; | ||
| } | ||
|
|
||
| // Handle the case where batch size is zero. | ||
| int64_t batch = std::max<int64_t>(1, tensor[0]); | ||
|
|
||
| for (size_t index = 1u; index < (tensor.size() - 1u); ++index) { | ||
| batch *= tensor[index]; | ||
| } | ||
|
|
||
| return batch; | ||
| }; | ||
|
|
||
| static int64_t channel(const IntArrayRef tensor) { | ||
| if (C10_UNLIKELY(tensor.empty())) { | ||
| return -1; | ||
| } | ||
|
|
||
| return tensor.back(); | ||
| }; | ||
| }; | ||
|
|
||
| // Convolution Filters | ||
| struct Filter final { | ||
| static constexpr size_t output = 0u; | ||
| static constexpr size_t input = 1u; | ||
| static constexpr size_t height = 2u; | ||
| static constexpr size_t width = 3u; | ||
| }; | ||
|
|
||
| // Parameters (Pooling Kernels, Dilation, Padding, Stride, etc.) | ||
| struct Parameter final { | ||
| static constexpr size_t height = 0u; | ||
| static constexpr size_t width = 1u; | ||
| }; | ||
| }; | ||
|
|
||
| struct Deleter final { | ||
| void operator()(const xnn_operator_t op) const { | ||
| xnn_delete_operator(op); | ||
| } | ||
| }; | ||
|
|
||
| using Operator = std::unique_ptr<xnn_operator, Deleter>; | ||
|
|
||
| bool available(); | ||
|
|
||
| } // namespace internal | ||
| } // namespace xnnpack | ||
| } // namespace native | ||
| } // namespace at | ||
|
|
||
| #endif /* USE_XNNPACK */ |
Oops, something went wrong.