ONE - On-device Neural Engine
Loading...
Searching...
No Matches
CircleBatchMatMul.cpp
Go to the documentation of this file.
1/*
2 * Copyright (c) 2021 Samsung Electronics Co., Ltd. All Rights Reserved
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 * http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
18
19#include "CircleCloneNode.h"
21
22namespace
23{
24
25loco::TensorShape remove_last_two(const loco::TensorShape &original_shape)
26{
27 assert(original_shape.rank() >= 2); // FIX CALLER UNLESS
28
30 ret.rank(original_shape.rank() - 2);
31
32 for (uint i = 0; i < ret.rank(); ++i)
33 {
34 ret.dim(i) = original_shape.dim(i);
35 }
36 return ret;
37}
38
39bool contain_zero(const loco::TensorShape &shape)
40{
41 bool zero_found = false;
42
43 for (uint32_t axis = 0; axis < shape.rank(); ++axis)
44 {
45 const auto &dim = shape.dim(axis);
46 if (dim.known() && dim.value() == 0)
47 {
48 zero_found = true;
49 break;
50 }
51 }
52 return zero_found;
53}
54
55void throw_unless(bool condition_result, const char *exception_msg)
56{
57 if (not condition_result)
58 {
59 INTERNAL_EXN(exception_msg);
60 }
61}
62
63} // namespace
64
65namespace luci
66{
67
69{
70 auto *cloned = _graph->nodes()->create<luci::CircleBatchMatMul>();
71 if (cloned != nullptr)
72 {
73 cloned->adj_x(node->adj_x());
74 cloned->adj_y(node->adj_y());
75 }
76 return cloned;
77}
78
79namespace sinf
80{
81
82// BatchMatMulV2 supports broadcasting in the batch dimensions(BatchMatMul doesn't)
83// TODO Distinguish BatchMatMul and BatchMatMulV2
85{
86 const auto x = loco::must_cast<CircleNode *>(node->x());
87 const auto y = loco::must_cast<CircleNode *>(node->y());
88
89 const auto x_shape = circle_shape(x);
90 const auto y_shape = circle_shape(y);
91
92 uint32_t x_rank = x_shape.rank();
93 uint32_t y_rank = y_shape.rank();
94
95 // throw internal exception if condition not met
96 throw_unless(x_rank >= 2, "x_rank shoud be >= 2");
97 throw_unless(y_rank >= 2, "y_rank shoud be >= 2");
98 throw_unless((not contain_zero(x_shape)), "x_shape should NOT have 0");
99 throw_unless((not contain_zero(y_shape)), "y_shape should NOT have 0");
100
101 // BatchMatMul shape inference rule works with two-part
102 //
103 // 1) Batch dimensions part
104 // - Batch dimensions correspond to the input_shape[:-2]
105 // - General broadcast rules are used to infer the shape of the output batch
106 //
107 // 2) Contracting dimensions part
108 // - Contracting dimensions correspond to the input_shape[-2:]
109 // - General matrix multiplication shape inference applied for this part,
110 // which means '(x_lhs, x_rhs) x (y_lhs, y_rhs) => (x_lhs, y_rhs)'
111
112 uint32_t max_rank = x_rank > y_rank ? x_rank : y_rank;
114 output_shape.rank(max_rank);
115
116 // broadcast in the batch dimensions
117 if (x_rank > 2 || y_rank > 2)
118 {
119 const auto x_batch_dims = remove_last_two(x_shape);
120 const auto y_batch_dims = remove_last_two(y_shape);
121
122 const auto o_batch_dims = sinf::broadcast_shape(x_batch_dims, y_batch_dims);
123
124 const auto o_batch_rank = o_batch_dims.rank();
125 for (uint i = 0u; i < o_batch_rank; ++i)
126 {
127 output_shape.dim(i) = o_batch_dims.dim(i);
128 }
129 }
130
131 // shape inference in contracting dimensions
132 const auto adj_x = node->adj_x();
133 const auto adj_y = node->adj_y();
134
135 loco::Dimension x_lhs = adj_x ? x_shape.dim(x_rank - 1) : x_shape.dim(x_rank - 2);
136 loco::Dimension x_rhs = adj_x ? x_shape.dim(x_rank - 2) : x_shape.dim(x_rank - 1);
137 loco::Dimension y_lhs = adj_y ? y_shape.dim(y_rank - 1) : y_shape.dim(y_rank - 2);
138 loco::Dimension y_rhs = adj_y ? y_shape.dim(y_rank - 2) : y_shape.dim(y_rank - 1);
139
140 if (x_rhs.known() && y_lhs.known() && not(x_rhs == y_lhs))
141 INTERNAL_EXN("x_rhs and y_lhs should be same");
142
143 uint32_t out_rank = output_shape.rank();
144 output_shape.dim(out_rank - 2) = x_lhs;
145 output_shape.dim(out_rank - 1) = y_rhs;
146
147 return output_shape;
148}
149
150} // namespace sinf
151
152} // namespace luci
#define INTERNAL_EXN(msg)
@ brief throw internal exception with message
Definition InternalExn.h:25
The value of one dimension in a tensor shape.
Definition Dimension.h:30
bool known(void) const
Return whether the value is known (or not)
Definition Dimension.h:47
const Dimension & dim(uint32_t axis) const
Definition TensorShape.h:38
uint32_t rank(void) const
Definition TensorShape.h:35
BATCH_MATMUL in Circle.
loco::Node * y(void) const
loco::Node * x(void) const
loco::TensorShape visit(const luci::CircleNode *node) final
Default fallback.
const luci_interpreter::RuntimeShape output_shape
loco::TensorShape broadcast_shape(const loco::TensorShape &x, const loco::TensorShape &y)
loco::TensorShape circle_shape(const luci::CircleNode *node)