| // |
| // Copyright © 2019 Arm Ltd. All rights reserved. |
| // SPDX-License-Identifier: MIT |
| // |
| |
| #include "DepthToSpace.hpp" |
| |
| #include <armnnUtils/DataLayoutIndexed.hpp> |
| #include <armnnUtils/Permute.hpp> |
| |
| #include <armnn/utility/Assert.hpp> |
| |
| using namespace armnnUtils; |
| |
| namespace armnn |
| { |
| |
| void DepthToSpace(const TensorInfo& inputInfo, |
| const DepthToSpaceDescriptor& descriptor, |
| const void* inputData, |
| void* outputData, |
| unsigned int dataTypeSize) |
| { |
| const unsigned int blockSize = descriptor.m_BlockSize; |
| ARMNN_ASSERT(blockSize != 0u); |
| |
| const TensorShape& inputShape = inputInfo.GetShape(); |
| const unsigned int batches = inputShape[0]; |
| |
| armnnUtils::DataLayoutIndexed dataLayoutIndexed(descriptor.m_DataLayout); |
| const unsigned int inDepth = inputShape[dataLayoutIndexed.GetChannelsIndex()]; |
| const unsigned int inHeight = inputShape[dataLayoutIndexed.GetHeightIndex()]; |
| const unsigned int inWidth = inputShape[dataLayoutIndexed.GetWidthIndex()]; |
| |
| const unsigned int outDepth = inDepth / (blockSize * blockSize); |
| |
| // The 4D input data can be interpreted as 6D (implicitly reshaped) as follows: |
| // |
| // [batch, block size, block size, inDepth, inHeight, inWidth] for NCHW and |
| // [batch, inHeight, inWidth, blockSize, blockSize, outDepth] for NHWC. |
| // |
| // DepthToSpace can then be implemented as a permutation in 6D resulting in |
| // the following shapes: |
| // |
| // [batch, outDepth, inHeight, blockSize, inWidth, blockSize] for NCHW and |
| // [batch, inHeight, blockSize, inWidth, blockSize, outDepth] for NHWC. |
| // |
| // NOTE: |
| // Since 6D tensors are not currently supported, in practice we need to handle each |
| // batch separately and execute 5D permutations |
| |
| TensorShape permDestShape; |
| PermutationVector permVector{}; |
| if (descriptor.m_DataLayout == DataLayout::NCHW) |
| { |
| permDestShape = TensorShape({ outDepth, inHeight, blockSize, inWidth, blockSize }); |
| permVector = { 2, 4, 0, 1, 3 }; |
| } |
| else |
| { |
| permDestShape = TensorShape({ inHeight, blockSize, inWidth, blockSize, outDepth }); |
| permVector = { 0, 2, 1, 3, 4 }; |
| } |
| |
| const unsigned int numElementsPerBatch = inputShape.GetNumElements() / batches; |
| |
| for (unsigned int batchIndex = 0u; batchIndex < batches; ++batchIndex) |
| { |
| const uintptr_t batchDataOffset = batchIndex * (numElementsPerBatch * dataTypeSize); |
| |
| armnnUtils::Permute(permDestShape, |
| permVector, |
| static_cast<const void*>(reinterpret_cast<const uint8_t*>(inputData) + batchDataOffset), |
| static_cast<void*>(reinterpret_cast<uint8_t*>(outputData) + batchDataOffset), |
| dataTypeSize); |
| } |
| } |
| |
| } // namespace armnn |