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
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
|
/*
* Copyright (c) 2018 Samsung Electronics Co., Ltd. All Rights Reserved
* Copyright (C) 2017 The Android Open Source Project
* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
#ifndef __ND_ARRAY_H__
#define __ND_ARRAY_H__
#include "Dims.h"
#include "Macro.h"
// From optimized_ops.h in TensorFlow Lite
//
// DO NOT USE THIS STRUCT FOR NEW FUNCTIONALITY BEYOND IMPLEMENTING ELEMENT-WISE
// BROADCASTING.
//
// NdArrayDesc<N> describes the shape and memory layout of an N-dimensional
// rectangular array of numbers.
//
// NdArrayDesc<N> is basically identical to Dims<N> defined in types.h.
// However, as Dims<N> is to be deprecated, this class exists as an adaptor
// to enable simple unoptimized implementations of element-wise broadcasting
// operations.
template <int N> struct NdArrayDesc
{
// The "extent" of each dimension. Indices along dimension d must be in the
// half-open interval [0, extents[d]).
int extents[N];
// The number of *elements* (not bytes) between consecutive indices of each
// dimension.
int strides[N];
};
// From optimized_ops.h in TensorFlow Lite
//
// DO NOT USE THIS FUNCTION FOR NEW FUNCTIONALITY BEYOND IMPLEMENTING
// ELEMENT-WISE BROADCASTING.
//
// Same as Offset(), except takes as NdArrayDesc<N> instead of Dims<N>.
inline int SubscriptToIndex(const NdArrayDesc<4> &desc, int i0, int i1, int i2, int i3)
{
DCHECK(i0 >= 0 && i0 < desc.extents[0]);
DCHECK(i1 >= 0 && i1 < desc.extents[1]);
DCHECK(i2 >= 0 && i2 < desc.extents[2]);
DCHECK(i3 >= 0 && i3 < desc.extents[3]);
return i0 * desc.strides[0] + i1 * desc.strides[1] + i2 * desc.strides[2] + i3 * desc.strides[3];
}
// From optimized_ops.h in TensorFlow Lite
//
// Given the dimensions of the operands for an element-wise binary broadcast,
// adjusts them so that they can be directly iterated over with simple loops.
// Returns the adjusted dims as instances of NdArrayDesc in 'desc0_out' and
// 'desc1_out'. 'desc0_out' and 'desc1_out' cannot be nullptr.
//
// This function assumes that the two input shapes are compatible up to
// broadcasting and the shorter one has already been prepended with 1s to be the
// same length. E.g., if shape0 is (1, 16, 16, 64) and shape1 is (1, 64),
// shape1 must already have been prepended to be (1, 1, 1, 64). Recall that
// Dims<N> refer to shapes in reverse order. In this case, input0_dims will be
// (64, 16, 16, 1) and input1_dims will be (64, 1, 1, 1).
//
// When two shapes are compatible up to broadcasting, for each dimension d,
// the input extents are either equal, or one of them is 1.
//
// This function performs the following for each dimension d:
// - If the extents are equal, then do nothing since the loop that walks over
// both of the input arrays is correct.
// - Otherwise, one (and only one) of the extents must be 1. Say extent0 is 1
// and extent1 is e1. Then set extent0 to e1 and stride0 *to 0*. This allows
// array0 to be referenced *at any index* in dimension d and still access the
// same slice.
template <int N>
inline void
NdArrayDescsForElementwiseBroadcast(const Dims<N> &input0_dims, const Dims<N> &input1_dims,
NdArrayDesc<N> *desc0_out, NdArrayDesc<N> *desc1_out)
{
DCHECK(desc0_out != nullptr);
DCHECK(desc1_out != nullptr);
// Copy dims to desc.
for (int i = 0; i < N; ++i)
{
desc0_out->extents[i] = input0_dims.sizes[i];
desc0_out->strides[i] = input0_dims.strides[i];
desc1_out->extents[i] = input1_dims.sizes[i];
desc1_out->strides[i] = input1_dims.strides[i];
}
// Walk over each dimension. If the extents are equal do nothing.
// Otherwise, set the desc with extent 1 to have extent equal to the other and
// stride 0.
for (int i = 0; i < N; ++i)
{
const int extent0 = ArraySize(input0_dims, i);
const int extent1 = ArraySize(input1_dims, i);
if (extent0 != extent1)
{
if (extent0 == 1)
{
desc0_out->strides[i] = 0;
desc0_out->extents[i] = extent1;
}
else
{
DCHECK_EQ(extent1, 1);
desc1_out->strides[i] = 0;
desc1_out->extents[i] = extent0;
}
}
}
}
inline int NodeOffset(int b, int h, int w, int height, int width)
{
return (b * height + h) * width + w;
}
#endif // __ND_ARRAY_H__
|
