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. (#32509)
Summary: In order to improve CPU performance on floating-point models on mobile, this PR introduces a new CPU backend for mobile that implements the most common mobile operators with NHWC memory layout support through integration with XNNPACK. XNNPACK itself, and this codepath, are currently only included in the build, but the actual integration is gated with USE_XNNPACK preprocessor guards. This preprocessor symbol is intentionally not passed on to the compiler, so as to enable this rollout in multiple stages in follow up PRs. This changeset will build XNNPACK as part of the build if the identically named USE_XNNPACK CMAKE variable, defaulted to ON, is enabled, but will not actually expose or enable this code path in any other way. Furthermore, it is worth pointing out that in order to efficiently map models to these operators, some front-end method of exposing this backend to the user is needed. The less efficient implementation would be to hook these operators into their corresponding **native** implementations, granted that a series of XNNPACK-specific conditions are met, much like how NNPACK is integrated with PyTorch today for instance. Having said that, while the above implementation is still expected to outperform NNPACK based on the benchmarks I ran, the above integration would be leave a considerable gap between the performance achieved and the maximum performance potential XNNPACK enables, as it does not provide a way to compute and factor out one-time operations out of the inner most forward() loop. The more optimal solution, and one we will decide on soon, would involve either providing a JIT pass that maps nn operators onto these newly introduced operators, while allowing one-time calculations to be factored out, much like quantized mobile models. Alternatively, new eager-mode modules can also be introduced that would directly call into these implementations either through c10 or some other mechanism, also allowing for decoupling of op creation from op execution. This PR does not include any of the front end changes mentioned above. Neither does it include the mobile threadpool unification present in the original #30644. Furthermore, this codepath seems to be faster than NNPACK in a good number of use cases, which can potentially allow us to remove NNPACK from aten to make the codebase a little simpler, granted that there is widespread support for such a move. Regardless, these changes will be introduced gradually and in a more controlled way in subsequent PRs. Pull Request resolved: #32509 Reviewed By: dreiss Differential Revision: D19521853 Pulled By: AshkanAliabadi fbshipit-source-id: 99a1fab31d0ece64961df074003bb852c36acaaa
- Loading branch information
1 parent
7aa605e
commit 941b424
Showing
22 changed files
with
1,080 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,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 |
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.