Éanna Ó Catháin | 4e1e136 | 2018-11-12 11:36:34 +0000 | [diff] [blame] | 1 | // |
Teresa Charlin | f77cab5 | 2023-06-01 16:15:13 +0100 | [diff] [blame] | 2 | // Copyright © 2017-2020,2023 Arm Ltd and Contributors. All rights reserved. |
Éanna Ó Catháin | 4e1e136 | 2018-11-12 11:36:34 +0000 | [diff] [blame] | 3 | // SPDX-License-Identifier: MIT |
| 4 | // |
| 5 | |
| 6 | #include "BatchToSpaceNd.hpp" |
| 7 | |
Teresa Charlin | f77cab5 | 2023-06-01 16:15:13 +0100 | [diff] [blame] | 8 | #include <armnnUtils/DataLayoutIndexed.hpp> |
Éanna Ó Catháin | 4e1e136 | 2018-11-12 11:36:34 +0000 | [diff] [blame] | 9 | |
Matteo Martincigh | 2135015 | 2018-11-28 16:22:22 +0000 | [diff] [blame] | 10 | using namespace armnnUtils; |
| 11 | |
Éanna Ó Catháin | 4e1e136 | 2018-11-12 11:36:34 +0000 | [diff] [blame] | 12 | namespace armnn |
| 13 | { |
| 14 | |
Teresa Charlin | f77cab5 | 2023-06-01 16:15:13 +0100 | [diff] [blame] | 15 | unsigned int Offset(const TensorShape& shape, |
| 16 | unsigned int batch, |
| 17 | unsigned int height, |
| 18 | unsigned int width, |
| 19 | unsigned int channels, |
| 20 | const DataLayoutIndexed& dataLayout) |
Éanna Ó Catháin | 4e1e136 | 2018-11-12 11:36:34 +0000 | [diff] [blame] | 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) |
Éanna Ó Catháin | 4e1e136 | 2018-11-12 11:36:34 +0000 | [diff] [blame] | 25 | { |
Teresa Charlin | f77cab5 | 2023-06-01 16:15:13 +0100 | [diff] [blame] | 26 | return (batch * shape[dataLayout.GetHeightIndex()] + height) * shape[channelDimension3D] + channels; |
| 27 | } |
| 28 | // 4D Tensors |
| 29 | else if (shape.GetNumDimensions() == 4) |
| 30 | { |
| 31 | if (dataLayout.GetDataLayout() == DataLayout::NHWC) |
| 32 | { |
| 33 | return ((batch * shape[dataLayout.GetHeightIndex()] + height) * |
| 34 | shape[dataLayout.GetWidthIndex()] + width) * |
| 35 | shape[dataLayout.GetChannelsIndex()] + channels; |
| 36 | } |
| 37 | else |
| 38 | { |
| 39 | return ((batch * shape[dataLayout.GetChannelsIndex()] + channels) * |
| 40 | shape[dataLayout.GetHeightIndex()] + height) * |
| 41 | shape[dataLayout.GetWidthIndex()] + width; |
| 42 | } |
Éanna Ó Catháin | 4e1e136 | 2018-11-12 11:36:34 +0000 | [diff] [blame] | 43 | } |
| 44 | else |
| 45 | { |
Teresa Charlin | f77cab5 | 2023-06-01 16:15:13 +0100 | [diff] [blame] | 46 | throw InvalidArgumentException("Tensor rank must be either 3 or 4", CHECK_LOCATION()); |
Éanna Ó Catháin | 4e1e136 | 2018-11-12 11:36:34 +0000 | [diff] [blame] | 47 | } |
| 48 | } |
| 49 | |
Teresa Charlin | f77cab5 | 2023-06-01 16:15:13 +0100 | [diff] [blame] | 50 | void BatchToSpaceNd(const TensorInfo& inputInfo, |
| 51 | const TensorInfo& outputInfo, |
| 52 | const BatchToSpaceNdDescriptor& params, |
| 53 | Decoder<float>& inputData, |
| 54 | Encoder<float>& outputData) |
Éanna Ó Catháin | 4e1e136 | 2018-11-12 11:36:34 +0000 | [diff] [blame] | 55 | { |
Teresa Charlin | f77cab5 | 2023-06-01 16:15:13 +0100 | [diff] [blame] | 56 | unsigned int rank = inputInfo.GetNumDimensions(); |
| 57 | if (rank != 3 && rank != 4 ) |
| 58 | { |
| 59 | throw InvalidArgumentException("Tensor rank must be either 3 or 4, but it is " + std::to_string(rank), |
| 60 | CHECK_LOCATION()); |
| 61 | } |
Éanna Ó Catháin | 95807ce | 2018-11-12 17:14:43 +0000 | [diff] [blame] | 62 | |
Teresa Charlin | f77cab5 | 2023-06-01 16:15:13 +0100 | [diff] [blame] | 63 | DataLayoutIndexed dataLayout = params.m_DataLayout; |
| 64 | unsigned int channelDimension3D = params.m_DataLayout == DataLayout::NCHW ? 1 : 2; |
Éanna Ó Catháin | 4e1e136 | 2018-11-12 11:36:34 +0000 | [diff] [blame] | 65 | |
Teresa Charlin | f77cab5 | 2023-06-01 16:15:13 +0100 | [diff] [blame] | 66 | TensorShape inputShape = inputInfo.GetShape(); |
| 67 | TensorShape outputShape = outputInfo.GetShape(); |
Éanna Ó Catháin | 95807ce | 2018-11-12 17:14:43 +0000 | [diff] [blame] | 68 | |
Teresa Charlin | f77cab5 | 2023-06-01 16:15:13 +0100 | [diff] [blame] | 69 | const unsigned int inputBatchSize = inputShape[0]; |
Éanna Ó Catháin | 4e1e136 | 2018-11-12 11:36:34 +0000 | [diff] [blame] | 70 | const unsigned int outputBatchSize = outputShape[0]; |
Teresa Charlin | f77cab5 | 2023-06-01 16:15:13 +0100 | [diff] [blame] | 71 | |
| 72 | const unsigned int channels = (rank == 3) ? inputShape[channelDimension3D] |
| 73 | : inputShape[dataLayout.GetChannelsIndex()]; |
| 74 | |
| 75 | const unsigned int inputHeight = inputShape[dataLayout.GetHeightIndex()]; |
| 76 | const unsigned int inputWidth = (rank == 3) ? 1 : inputShape[dataLayout.GetWidthIndex()]; |
Éanna Ó Catháin | 4e1e136 | 2018-11-12 11:36:34 +0000 | [diff] [blame] | 77 | const unsigned int outputHeight = outputShape[dataLayout.GetHeightIndex()]; |
Teresa Charlin | f77cab5 | 2023-06-01 16:15:13 +0100 | [diff] [blame] | 78 | const unsigned int outputWidth = (rank == 3) ? 1 : outputShape[dataLayout.GetWidthIndex()]; |
Éanna Ó Catháin | 4e1e136 | 2018-11-12 11:36:34 +0000 | [diff] [blame] | 79 | |
Teresa Charlin | f77cab5 | 2023-06-01 16:15:13 +0100 | [diff] [blame] | 80 | const unsigned int blockHeight = params.m_BlockShape[0]; |
| 81 | const unsigned int blockWidth = (rank == 3) ? 1 : params.m_BlockShape[1]; |
Éanna Ó Catháin | 95807ce | 2018-11-12 17:14:43 +0000 | [diff] [blame] | 82 | |
Teresa Charlin | f77cab5 | 2023-06-01 16:15:13 +0100 | [diff] [blame] | 83 | const unsigned int cropsTop = params.m_Crops[0].first; |
| 84 | const unsigned int cropsLeft = (rank == 3) ? 0 : params.m_Crops[1].first; |
Éanna Ó Catháin | 4e1e136 | 2018-11-12 11:36:34 +0000 | [diff] [blame] | 85 | |
| 86 | for (unsigned int inBatch = 0; inBatch < inputBatchSize; ++inBatch) |
| 87 | { |
| 88 | const unsigned int outBatch = inBatch % outputBatchSize; |
| 89 | const unsigned int spatialOffset = inBatch / outputBatchSize; |
| 90 | |
Teresa Charlin | f77cab5 | 2023-06-01 16:15:13 +0100 | [diff] [blame] | 91 | for (unsigned int inH = 0; inH < inputHeight; ++inH) |
| 92 | { |
| 93 | const unsigned int outH = inH * blockHeight + spatialOffset / blockWidth - cropsTop; |
Éanna Ó Catháin | 4e1e136 | 2018-11-12 11:36:34 +0000 | [diff] [blame] | 94 | |
| 95 | if (outH >= outputHeight) |
| 96 | { |
| 97 | continue; |
| 98 | } |
| 99 | |
Teresa Charlin | f77cab5 | 2023-06-01 16:15:13 +0100 | [diff] [blame] | 100 | for (unsigned int inW = 0; inW < inputWidth; ++inW) |
| 101 | { |
| 102 | const unsigned int outW = inW * blockWidth + spatialOffset % blockWidth - cropsLeft; |
Éanna Ó Catháin | 4e1e136 | 2018-11-12 11:36:34 +0000 | [diff] [blame] | 103 | |
| 104 | if (outW >= outputWidth) |
| 105 | { |
| 106 | continue; |
| 107 | } |
| 108 | |
| 109 | for (unsigned int c = 0; c < channels; c++) |
| 110 | { |
| 111 | unsigned int outOffset = Offset(outputShape, outBatch, outH, outW, c, dataLayout); |
| 112 | unsigned int inOffset = Offset(inputShape, inBatch, inH, inW, c, dataLayout); |
Francis Murtagh | 47ea3c0 | 2019-06-20 12:07:19 +0100 | [diff] [blame] | 113 | |
Teresa Charlin | f77cab5 | 2023-06-01 16:15:13 +0100 | [diff] [blame] | 114 | outputData[outOffset]; |
| 115 | inputData[inOffset]; |
| 116 | outputData.Set(inputData.Get()); |
Éanna Ó Catháin | 4e1e136 | 2018-11-12 11:36:34 +0000 | [diff] [blame] | 117 | } |
| 118 | } |
| 119 | } |
| 120 | } |
| 121 | } |
| 122 | |
| 123 | } //namespace armnn |