Keith Davis | a57eccb | 2019-06-14 17:33:22 +0100 | [diff] [blame] | 1 | // |
| 2 | // Copyright © 2017 Arm Ltd. All rights reserved. |
| 3 | // SPDX-License-Identifier: MIT |
| 4 | // |
| 5 | |
| 6 | #include "SpaceToDepth.hpp" |
| 7 | |
Matteo Martincigh | e011d20 | 2019-11-28 11:35:47 +0000 | [diff] [blame] | 8 | #include <armnnUtils/DataLayoutIndexed.hpp> |
Keith Davis | a57eccb | 2019-06-14 17:33:22 +0100 | [diff] [blame] | 9 | |
| 10 | using namespace armnnUtils; |
| 11 | |
| 12 | namespace { |
| 13 | unsigned int GetOffset(const armnn::TensorShape& shape, |
| 14 | unsigned int c, |
| 15 | unsigned int h, |
| 16 | unsigned int w, |
| 17 | unsigned int b, |
| 18 | const DataLayoutIndexed& dataLayout) |
| 19 | { |
| 20 | if (dataLayout.GetDataLayout() == armnn::DataLayout::NHWC) |
| 21 | { |
| 22 | return ((b * shape[dataLayout.GetHeightIndex()] + h) * shape[dataLayout.GetWidthIndex()] + w) * |
| 23 | shape[dataLayout.GetChannelsIndex()] + c; |
| 24 | } |
| 25 | else |
| 26 | { |
| 27 | return ((b * shape[dataLayout.GetChannelsIndex()] + c) * shape[dataLayout.GetHeightIndex()] + h) * |
| 28 | shape[dataLayout.GetWidthIndex()] + w; |
| 29 | } |
| 30 | } |
| 31 | } |
| 32 | |
| 33 | namespace armnn |
| 34 | { |
| 35 | |
| 36 | void SpaceToDepth(const TensorInfo& inputInfo, |
| 37 | const TensorInfo& outputInfo, |
| 38 | const SpaceToDepthDescriptor& params, |
| 39 | Decoder<float>& inputData, |
| 40 | Encoder<float>& outputData) |
| 41 | { |
| 42 | DataLayoutIndexed dataLayout = params.m_DataLayout; |
| 43 | |
| 44 | const TensorShape& inputShape = inputInfo.GetShape(); |
| 45 | const TensorShape& outputShape = outputInfo.GetShape(); |
| 46 | |
| 47 | const unsigned int inputBatchSize = inputShape[0]; |
| 48 | const unsigned int inputChannels = inputShape[dataLayout.GetChannelsIndex()]; |
| 49 | |
| 50 | const unsigned int outputHeight = outputShape[dataLayout.GetHeightIndex()]; |
| 51 | const unsigned int outputWidth = outputShape[dataLayout.GetWidthIndex()]; |
| 52 | const unsigned int outputChannels = outputShape[dataLayout.GetChannelsIndex()]; |
| 53 | |
| 54 | const unsigned int blockSize = params.m_BlockSize; |
| 55 | |
| 56 | if (blockSize == 0) |
| 57 | { |
| 58 | throw InvalidArgumentException( |
| 59 | "Input shape must be divisible by block size in all spatial dimensions: Block size is" |
| 60 | " equal to zero"); |
| 61 | } |
| 62 | |
| 63 | for (unsigned int outChannelIndex = 0; outChannelIndex < outputChannels; outChannelIndex++) |
| 64 | { |
| 65 | unsigned int inChannelIndex = outChannelIndex % inputChannels; |
| 66 | |
| 67 | unsigned int shiftW = (outChannelIndex / inputChannels) % blockSize; |
| 68 | unsigned int shiftH = (outChannelIndex / inputChannels) / blockSize; |
| 69 | |
| 70 | for (unsigned int outH = 0; outH < outputHeight; outH++) |
| 71 | { |
| 72 | for (unsigned int outW = 0; outW < outputWidth; outW++) |
| 73 | { |
| 74 | for (unsigned int inBatchIndex = 0; inBatchIndex < inputBatchSize; inBatchIndex++) |
| 75 | { |
| 76 | unsigned int inOffset = GetOffset(inputShape, |
| 77 | inChannelIndex, |
| 78 | (outH * blockSize + shiftH), |
| 79 | (outW * blockSize + shiftW), |
| 80 | inBatchIndex, |
| 81 | dataLayout); |
| 82 | |
| 83 | unsigned int outOffset = GetOffset(outputShape, |
| 84 | outChannelIndex, |
| 85 | outH, |
| 86 | outW, |
| 87 | inBatchIndex, |
| 88 | dataLayout); |
| 89 | |
| 90 | outputData += outOffset; |
| 91 | inputData += inOffset; |
| 92 | outputData.Set(inputData.Get()); |
| 93 | inputData -= inOffset; |
| 94 | outputData -= outOffset; |
| 95 | } |
| 96 | } |
| 97 | } |
| 98 | } |
| 99 | } |
| 100 | |
| 101 | void SpaceToDepth(const TensorInfo& inputInfo, |
| 102 | const TensorInfo& outputInfo, |
| 103 | const SpaceToDepthDescriptor& params, |
| 104 | Decoder<float>& inputData, |
| 105 | Encoder<float>& outData); |
| 106 | |
| 107 | } //namespace armnn |