1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
|
//
// Copyright © 2019 Arm Ltd. All rights reserved.
// SPDX-License-Identifier: MIT
//
#include "DepthToSpace.hpp"
#include <DataLayoutIndexed.hpp>
#include <Permute.hpp>
#include <boost/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;
BOOST_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
|
