Nattapat Chaimanowong | 3ea76d5 | 2018-11-09 14:10:38 +0000 | [diff] [blame] | 1 | // |
Teresa Charlin | f77cab5 | 2023-06-01 16:15:13 +0100 | [diff] [blame] | 2 | // Copyright © 2017-2019,2023 Arm Ltd and Contributors. All rights reserved. |
Nattapat Chaimanowong | 3ea76d5 | 2018-11-09 14:10:38 +0000 | [diff] [blame] | 3 | // SPDX-License-Identifier: MIT |
| 4 | // |
| 5 | |
| 6 | #include "SpaceToBatchNd.hpp" |
| 7 | |
Matteo Martincigh | e011d20 | 2019-11-28 11:35:47 +0000 | [diff] [blame] | 8 | #include <armnnUtils/DataLayoutIndexed.hpp> |
Matteo Martincigh | 2135015 | 2018-11-28 16:22:22 +0000 | [diff] [blame] | 9 | |
| 10 | using namespace armnnUtils; |
Matthew Bentham | 8800c00 | 2018-11-19 13:19:28 +0000 | [diff] [blame] | 11 | |
Nattapat Chaimanowong | 3ea76d5 | 2018-11-09 14:10:38 +0000 | [diff] [blame] | 12 | namespace armnn |
| 13 | { |
| 14 | |
| 15 | unsigned int GetOffset(const TensorShape& shape, |
| 16 | unsigned int b, |
| 17 | unsigned int h, |
| 18 | unsigned int w, |
| 19 | unsigned int c, |
| 20 | const DataLayoutIndexed& dataLayout) |
| 21 | { |
Teresa Charlin | f77cab5 | 2023-06-01 16:15:13 +0100 | [diff] [blame] | 22 | // 3D Tensors |
| 23 | unsigned int channelDimension3D = dataLayout.GetDataLayout() == DataLayout::NCHW ? 1 : 2; |
| 24 | if (shape.GetNumDimensions() == 3) |
Nattapat Chaimanowong | 3ea76d5 | 2018-11-09 14:10:38 +0000 | [diff] [blame] | 25 | { |
Teresa Charlin | f77cab5 | 2023-06-01 16:15:13 +0100 | [diff] [blame] | 26 | return (b * shape[dataLayout.GetHeightIndex()] + h) * shape[channelDimension3D] + c; |
| 27 | } |
| 28 | // 4D Tensors |
| 29 | else if (shape.GetNumDimensions() == 4) |
| 30 | { |
| 31 | if (dataLayout.GetDataLayout() == DataLayout::NHWC) |
| 32 | { |
| 33 | return ((b * shape[dataLayout.GetHeightIndex()] + h) * shape[dataLayout.GetWidthIndex()] + w) * |
| 34 | shape[dataLayout.GetChannelsIndex()] + c; |
| 35 | } |
| 36 | else |
| 37 | { |
| 38 | return ((b * shape[dataLayout.GetChannelsIndex()] + c) * shape[dataLayout.GetHeightIndex()] + h) * |
| 39 | shape[dataLayout.GetWidthIndex()] + w; |
| 40 | } |
Nattapat Chaimanowong | 3ea76d5 | 2018-11-09 14:10:38 +0000 | [diff] [blame] | 41 | } |
| 42 | else |
| 43 | { |
Teresa Charlin | f77cab5 | 2023-06-01 16:15:13 +0100 | [diff] [blame] | 44 | throw InvalidArgumentException("Tensor rank must be either 3 or 4", CHECK_LOCATION()); |
Nattapat Chaimanowong | 3ea76d5 | 2018-11-09 14:10:38 +0000 | [diff] [blame] | 45 | } |
| 46 | } |
| 47 | |
Nattapat Chaimanowong | 3ea76d5 | 2018-11-09 14:10:38 +0000 | [diff] [blame] | 48 | void SpaceToBatchNd(const TensorInfo& inputInfo, |
| 49 | const TensorInfo& outputInfo, |
| 50 | const SpaceToBatchNdDescriptor& params, |
nikraj01 | 22f0f2b | 2019-05-30 17:29:32 +0100 | [diff] [blame] | 51 | Decoder<float>& inputData, |
| 52 | Encoder<float>& outputData) |
Nattapat Chaimanowong | 3ea76d5 | 2018-11-09 14:10:38 +0000 | [diff] [blame] | 53 | { |
Teresa Charlin | f77cab5 | 2023-06-01 16:15:13 +0100 | [diff] [blame] | 54 | unsigned int rank = inputInfo.GetNumDimensions(); |
| 55 | if (rank != 3 && rank != 4 ) |
| 56 | { |
| 57 | throw InvalidArgumentException("Tensor rank must be either 3 or 4, but it is " + std::to_string(rank), |
| 58 | CHECK_LOCATION()); |
| 59 | } |
| 60 | |
Nattapat Chaimanowong | 3ea76d5 | 2018-11-09 14:10:38 +0000 | [diff] [blame] | 61 | DataLayoutIndexed dataLayout = params.m_DataLayout; |
Teresa Charlin | f77cab5 | 2023-06-01 16:15:13 +0100 | [diff] [blame] | 62 | unsigned int channelDimension3D = params.m_DataLayout == DataLayout::NCHW ? 1 : 2; |
Nattapat Chaimanowong | 3ea76d5 | 2018-11-09 14:10:38 +0000 | [diff] [blame] | 63 | |
| 64 | const TensorShape& inputShape = inputInfo.GetShape(); |
| 65 | const TensorShape& outputShape = outputInfo.GetShape(); |
| 66 | |
Teresa Charlin | f77cab5 | 2023-06-01 16:15:13 +0100 | [diff] [blame] | 67 | const unsigned int inputBatchSize = inputShape[0]; |
Nattapat Chaimanowong | 3ea76d5 | 2018-11-09 14:10:38 +0000 | [diff] [blame] | 68 | const unsigned int outputBatchSize = outputShape[0]; |
Teresa Charlin | f77cab5 | 2023-06-01 16:15:13 +0100 | [diff] [blame] | 69 | |
| 70 | const unsigned int channels = (rank == 3) ? inputShape[channelDimension3D] |
| 71 | : inputShape[dataLayout.GetChannelsIndex()]; |
| 72 | |
| 73 | const unsigned int inputHeight = inputShape[dataLayout.GetHeightIndex()]; |
| 74 | const unsigned int inputWidth = (rank == 3) ? 1 : inputShape[dataLayout.GetWidthIndex()]; |
Nattapat Chaimanowong | 3ea76d5 | 2018-11-09 14:10:38 +0000 | [diff] [blame] | 75 | const unsigned int outputHeight = outputShape[dataLayout.GetHeightIndex()]; |
Teresa Charlin | f77cab5 | 2023-06-01 16:15:13 +0100 | [diff] [blame] | 76 | const unsigned int outputWidth = (rank == 3) ? 1 : outputShape[dataLayout.GetWidthIndex()]; |
Nattapat Chaimanowong | 3ea76d5 | 2018-11-09 14:10:38 +0000 | [diff] [blame] | 77 | |
| 78 | const unsigned int blockHeight = params.m_BlockShape[0]; |
Teresa Charlin | f77cab5 | 2023-06-01 16:15:13 +0100 | [diff] [blame] | 79 | const unsigned int blockWidth = (rank == 3) ? 1 : params.m_BlockShape[1]; |
Nattapat Chaimanowong | 3ea76d5 | 2018-11-09 14:10:38 +0000 | [diff] [blame] | 80 | |
Teresa Charlin | f77cab5 | 2023-06-01 16:15:13 +0100 | [diff] [blame] | 81 | const unsigned int paddingTop = params.m_PadList[0].first; |
| 82 | const unsigned int paddingLeft = (rank == 3) ? 0 : params.m_PadList[1].first; |
Nattapat Chaimanowong | 3ea76d5 | 2018-11-09 14:10:38 +0000 | [diff] [blame] | 83 | |
Teresa Charlin | f77cab5 | 2023-06-01 16:15:13 +0100 | [diff] [blame] | 84 | for (unsigned int outB = 0; outB < outputBatchSize; ++outB) |
Nattapat Chaimanowong | 3ea76d5 | 2018-11-09 14:10:38 +0000 | [diff] [blame] | 85 | { |
| 86 | unsigned int inB = outB % inputBatchSize; |
| 87 | |
| 88 | unsigned int shiftW = (outB / inputBatchSize) % blockWidth; |
| 89 | unsigned int shiftH = (outB / inputBatchSize) / blockWidth; |
| 90 | |
Teresa Charlin | f77cab5 | 2023-06-01 16:15:13 +0100 | [diff] [blame] | 91 | for (unsigned int outH = 0; outH < outputHeight; ++outH) |
Nattapat Chaimanowong | 3ea76d5 | 2018-11-09 14:10:38 +0000 | [diff] [blame] | 92 | { |
Teresa Charlin | f77cab5 | 2023-06-01 16:15:13 +0100 | [diff] [blame] | 93 | for (unsigned int outW = 0; outW < outputWidth; ++outW) |
Nattapat Chaimanowong | 3ea76d5 | 2018-11-09 14:10:38 +0000 | [diff] [blame] | 94 | { |
| 95 | if (outH * blockHeight + shiftH < paddingTop || |
| 96 | outH * blockHeight + shiftH >= paddingTop + inputHeight || |
| 97 | outW * blockWidth + shiftW < paddingLeft || |
| 98 | outW * blockWidth + shiftW >= paddingLeft + inputWidth) |
| 99 | { |
| 100 | for (unsigned int c = 0; c < channels; c++) |
| 101 | { |
| 102 | unsigned int outOffset = GetOffset(outputShape, |
| 103 | outB, |
| 104 | outH, |
| 105 | outW, |
| 106 | c, |
| 107 | dataLayout); |
nikraj01 | 22f0f2b | 2019-05-30 17:29:32 +0100 | [diff] [blame] | 108 | outputData += outOffset; |
| 109 | outputData.Set(0); |
| 110 | outputData -= outOffset; |
Nattapat Chaimanowong | 3ea76d5 | 2018-11-09 14:10:38 +0000 | [diff] [blame] | 111 | } |
| 112 | } |
| 113 | else |
| 114 | { |
| 115 | for (unsigned int c = 0; c < channels; c++) |
| 116 | { |
| 117 | unsigned int inOffset = GetOffset(inputShape, |
| 118 | inB, |
| 119 | (outH * blockHeight + shiftH) - paddingTop, |
| 120 | (outW * blockWidth + shiftW) - paddingLeft, |
| 121 | c, |
| 122 | dataLayout); |
| 123 | |
| 124 | unsigned int outOffset = GetOffset(outputShape, |
| 125 | outB, |
| 126 | outH, |
| 127 | outW, |
| 128 | c, |
| 129 | dataLayout); |
| 130 | |
nikraj01 | 22f0f2b | 2019-05-30 17:29:32 +0100 | [diff] [blame] | 131 | outputData += outOffset; |
| 132 | inputData += inOffset; |
| 133 | outputData.Set(inputData.Get()); |
| 134 | inputData -= inOffset; |
| 135 | outputData -= outOffset; |
Nattapat Chaimanowong | 3ea76d5 | 2018-11-09 14:10:38 +0000 | [diff] [blame] | 136 | } |
| 137 | } |
| 138 | } |
| 139 | } |
| 140 | } |
| 141 | } |
| 142 | |
Nattapat Chaimanowong | 3ea76d5 | 2018-11-09 14:10:38 +0000 | [diff] [blame] | 143 | } //namespace armnn |