2 * Copyright (c) 2020 Samsung Electronics Co., Ltd. All Rights Reserved
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
8 * http://www.apache.org/licenses/LICENSE-2.0
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.
17 #include "compiler/StaticShapeInference.h"
18 #include "util/ShapeInference.h"
19 #include "util/logging.h"
28 bool StaticShapeInferer::infer(const ir::OpSequence &op_seq)
30 bool has_dynamic_tensor = false;
32 for (const auto &operation_idx : op_seq.operations())
34 auto &op = _operations.at(operation_idx);
35 auto opcode = op.opcode();
37 _return_has_dynamic_tensor = false; // this is used as a return value inside operation's visit()
39 // IF: need shape inference for then, else
40 // While: need shape inference for condition, body
41 if (opcode == ir::OpCode::If || opcode == ir::OpCode::While)
47 _return_has_dynamic_tensor = checkDynamicInput(op);
49 if (_return_has_dynamic_tensor)
59 has_dynamic_tensor = has_dynamic_tensor || _return_has_dynamic_tensor;
62 return has_dynamic_tensor;
65 bool StaticShapeInferer::checkDynamicInput(const ir::Operation &op)
67 for (auto input_idx : op.getInputs() | ir::Remove::UNDEFINED | ir::Remove::DUPLICATED)
69 if (_operands.at(input_idx).info().isDynamic())
78 void StaticShapeInferer::setDynamicOutput(const ir::Operation &op)
80 for (auto output_idx : op.getOutputs())
82 _operands.at(output_idx).info().setDynamic();
86 void StaticShapeInferer::handleBinaryArithmeticOp(const ir::Operation &op,
87 const ir::OperandIndex lhs_idx,
88 const ir::OperandIndex rhs_idx)
90 const auto &lhs = _operands.at(lhs_idx);
91 const auto &rhs = _operands.at(rhs_idx);
93 const auto output_idx = op.getOutputs().at(0);
94 ir::Operand &output = _operands.at(output_idx);
96 // re-sizing output shape
97 ir::Shape new_shape = shape_inference::inferEltwiseShape(lhs.info().shape(), rhs.info().shape());
98 output.info().shape(new_shape);
101 void StaticShapeInferer::handleSimpleUnaryOp(const ir::Operation &op,
102 const ir::OperandIndex input_idx)
104 const auto &input = _operands.at(input_idx);
106 // get mutable output operand
107 const auto output_idx = op.getOutputs().at(0);
108 ir::Operand &output = _operands.at(output_idx);
110 // re-sizing output shape
111 ir::Shape new_shape = input.info().shape();
112 output.info().shape(new_shape);
115 void StaticShapeInferer::dump()
117 auto get_shape_str = [](const ir::Shape &shape) {
118 std::stringstream sstream;
119 sstream << "shape : {";
120 for (int i = 0; i < shape.rank(); i++)
123 sstream << shape.dim(i);
125 sstream << " " << shape.dim(i);
128 return sstream.str();
131 for (const auto &pair : _lowered_subgs)
133 const auto index = pair.first;
134 const auto &lowered_subg = pair.second;
135 VERBOSE(StaticShapeInferer) << "SubGraph #" << index.value() << std::endl;
136 lowered_subg->graph().operands().iterate(
137 [&](const ir::OperandIndex &ind, const ir::Operand &operand) {
138 VERBOSE(StaticShapeInferer) << "Operand #" << ind.value() << ", "
139 << (operand.info().isDynamic() ? "Dynamic" : "Static") << ", "
140 << get_shape_str(operand.info().shape()) << std::endl;
145 void StaticShapeInferer::visit(const ir::operation::ArgMax &op)
147 const auto input_idx{op.getInputs().at(ir::operation::ArgMax::Input::INPUT)};
148 const auto &input = _operands.at(input_idx);
150 // get mutable output operand
151 const auto output_idx = op.getOutputs().at(0);
152 ir::Operand &output = _operands.at(output_idx);
153 const auto rank = input.info().shape().rank();
154 const auto axis = ((op.param().axis < 0) ? rank + op.param().axis : op.param().axis);
156 assert(0 <= axis && axis < rank);
158 // re-sizing output shape
159 ir::Shape new_shape = shape_inference::inferArgMaxShape(input.info().shape(), axis, rank);
160 output.info().shape(new_shape);
163 void StaticShapeInferer::visit(const ir::operation::BatchMatMul &op)
165 const auto lhs_index = op.getInputs().at(ir::operation::BatchMatMul::Input::LHS);
166 const auto rhs_index = op.getInputs().at(ir::operation::BatchMatMul::Input::RHS);
167 const auto output_index = op.getOutputs().at(0);
168 const auto lhs = _operands.at(lhs_index);
169 const auto rhs = _operands.at(rhs_index);
170 auto &output = _operands.at(output_index);
171 auto new_shape = shape_inference::inferBatchMatMulShape(lhs.shape(), rhs.shape(), op.param());
172 output.info().shape(new_shape);
175 void StaticShapeInferer::visit(const ir::operation::BinaryArithmetic &op)
177 handleBinaryArithmeticOp(op, op.getInputs().at(ir::operation::BinaryArithmetic::Input::LHS),
178 op.getInputs().at(ir::operation::BinaryArithmetic::Input::RHS));
181 void StaticShapeInferer::visit(const ir::operation::BroadcastTo &op)
183 // get mutable output operand
184 const auto output_idx = op.getOutputs().at(0);
185 ir::Operand &output = _operands.at(output_idx);
187 const auto shape_idx{op.getInputs().at(ir::operation::BroadcastTo::Input::SHAPE)};
188 const auto &shape = _operands.at(shape_idx);
190 if (!shape.isConstant())
192 output.info().setDynamic();
193 _return_has_dynamic_tensor = true;
197 // assert(shape.typeInfo().type() == ir::DataType::INT32);
198 auto shape_buffer = reinterpret_cast<const int32_t *>(shape.data()->base());
200 // re-sizing output shape
201 ir::Shape new_shape = shape_inference::inferBroadcastToShape(shape.info().shape(), shape_buffer);
202 output.info().shape(new_shape);
205 void StaticShapeInferer::visit(const ir::operation::Comparison &op)
207 handleBinaryArithmeticOp(op, op.getInputs().at(ir::operation::Comparison::Input::INPUT0),
208 op.getInputs().at(ir::operation::Comparison::Input::INPUT1));
211 void StaticShapeInferer::visit(const ir::operation::Concat &op)
213 const auto input_count = op.getInputs().size();
215 const auto output_idx = op.getOutputs().at(0);
216 ir::Operand &output = _operands.at(output_idx);
218 shape_inference::Shapes input_shapes;
219 for (uint32_t i = 0; i < input_count; i++)
221 const auto input_idx{op.getInputs().at(i)};
222 const auto &input = _operands.at(input_idx);
223 input_shapes.emplace_back(input.shape());
226 ir::Shape out_shape = shape_inference::inferConcatShape(input_shapes, op.param());
228 // re-sizing output shape
229 output.info().shape(out_shape);
232 void StaticShapeInferer::visit(const ir::operation::Conv2D &op)
234 const auto input_idx{op.getInputs().at(ir::operation::Conv2D::Input::INPUT)};
235 const auto &input = _operands.at(input_idx);
236 const auto ker_idx{op.getInputs().at(ir::operation::Conv2D::Input::KERNEL)};
237 const auto &ker = _operands.at(ker_idx);
238 const auto output_idx = op.getOutputs().at(0);
239 ir::Operand &output = _operands.at(output_idx);
241 // re-sizing output shape
242 ir::Shape new_shape =
243 shape_inference::inferConv2DShape(input.info().shape(), ker.info().shape(), op.param());
244 output.info().shape(new_shape);
247 void StaticShapeInferer::visit(const ir::operation::ElementwiseActivation &op)
249 handleSimpleUnaryOp(op, op.getInputs().at(ir::operation::ElementwiseActivation::Input::INPUT));
252 void StaticShapeInferer::visit(const ir::operation::ElementwiseBinary &op)
254 handleBinaryArithmeticOp(op, op.getInputs().at(ir::operation::ElementwiseBinary::Input::LHS),
255 op.getInputs().at(ir::operation::ElementwiseBinary::Input::RHS));
258 void StaticShapeInferer::visit(const ir::operation::ElementwiseUnary &op)
260 handleSimpleUnaryOp(op, op.getInputs().at(ir::operation::ElementwiseUnary::Input::INPUT));
263 void StaticShapeInferer::visit(const ir::operation::ExpandDims &op)
265 const auto input_idx{op.getInputs().at(ir::operation::ExpandDims::Input::INPUT)};
266 const auto &input = _operands.at(input_idx);
267 const auto axis_idx{op.getInputs().at(ir::operation::ExpandDims::Input::AXIS)};
268 const auto &axis = _operands.at(axis_idx);
269 const auto output_idx = op.getOutputs().at(0);
270 ir::Operand &output = _operands.at(output_idx);
272 if (!axis.isConstant())
274 output.info().setDynamic();
275 _return_has_dynamic_tensor = true;
279 // even when axis is constant, output shape should be recalculated since user might call
280 // nnfw_set_input_tensorinfo(input, some_new_shape)
281 auto axis_buf = reinterpret_cast<const int32_t *>(axis.data()->base());
284 // re-sizing output shape
285 ir::Shape new_shape = shape_inference::inferExpandDimsShape(input.info().shape(), axis_buf[0]);
286 output.info().shape(new_shape);
289 void StaticShapeInferer::visit(const ir::operation::Fill &op)
291 const auto input_idx{op.getInputs().at(ir::operation::Fill::Input::INPUT)};
292 const auto &input = _operands.at(input_idx);
293 const auto output_idx = op.getOutputs().at(0);
294 ir::Operand &output = _operands.at(output_idx);
296 if (!input.isConstant())
298 output.info().setDynamic();
299 _return_has_dynamic_tensor = true;
303 assert(input.typeInfo().type() == ir::DataType::INT32);
305 auto input_buf = reinterpret_cast<const int32_t *>(input.data()->base());
308 // re-sizing output shape
309 ir::Shape new_shape = shape_inference::inferFillShape(input.info().shape(), input_buf);
310 output.info().shape(new_shape);
313 void StaticShapeInferer::visit(const ir::operation::FullyConnected &op)
315 const auto input_idx{op.getInputs().at(ir::operation::FullyConnected::Input::INPUT)};
316 const auto &input = _operands.at(input_idx);
318 const auto ker_idx{op.getInputs().at(ir::operation::FullyConnected::Input::WEIGHT)};
319 const auto &ker = _operands.at(ker_idx);
321 // get mutable output operand
322 const auto output_idx = op.getOutputs().at(0);
323 ir::Operand &output = _operands.at(output_idx);
324 // re-sizing output shape
325 ir::Shape new_shape =
326 shape_inference::inferFullyConnectedShape(input.info().shape(), ker.info().shape());
327 output.info().shape(new_shape);
330 void StaticShapeInferer::visit(const ir::operation::FusedBatchNorm &op)
332 handleSimpleUnaryOp(op, op.getInputs().at(ir::operation::FusedBatchNorm::Input::INPUT));
335 void StaticShapeInferer::visit(const ir::operation::Gather &op)
337 const auto input_idx{op.getInputs().at(ir::operation::Gather::Input::INPUT)};
338 const auto &input = _operands.at(input_idx);
340 // get mutable output operand
341 const auto output_idx = op.getOutputs().at(0);
342 ir::Operand &output = _operands.at(output_idx);
344 const auto indices_idx{op.getInputs().at(ir::operation::Gather::Input::INDICES)};
345 const auto &indices = _operands.at(indices_idx);
346 const auto rank = input.info().shape().rank();
347 const auto axis = ((op.param().axis < 0) ? rank + op.param().axis : op.param().axis);
349 assert(0 <= axis && axis < rank);
351 // re-sizing output shape
352 ir::Shape new_shape =
353 shape_inference::inferGatherShape(input.info().shape(), indices.info().shape(), axis, rank);
354 output.info().shape(new_shape);
357 void StaticShapeInferer::visit(const ir::operation::If &op)
359 auto &then_graph = _lowered_subgs.at(op.param().then_subg_index)->graph();
360 auto &else_graph = _lowered_subgs.at(op.param().else_subg_index)->graph();
361 const std::vector<ir::OperandIndex> inputs{op.getInputs().begin() + 1, op.getInputs().end()};
362 const auto &outputs = op.getOutputs();
364 // re-sizing input shapes of then subgraph
365 const auto &then_inputs = then_graph.getInputs();
366 assert(inputs.size() == then_inputs.size());
367 for (size_t i = 0; i < inputs.size(); ++i)
369 auto &then_input = then_graph.operands().at(then_inputs.at(i));
370 if (_operands.at(inputs.at(i)).info().isDynamic())
372 then_input.info().setDynamic();
376 auto new_shape = _operands.at(inputs.at(i)).info().shape();
377 then_input.info().shape(new_shape);
381 // re-sizing input shapes of else subgraph
382 const auto &else_inputs = else_graph.getInputs();
383 assert(inputs.size() == else_inputs.size());
384 for (size_t i = 0; i < inputs.size(); ++i)
386 auto &else_input = else_graph.operands().at(else_inputs.at(i));
387 if (_operands.at(inputs.at(i)).info().isDynamic())
389 else_input.info().setDynamic();
393 const auto &new_shape = _operands.at(inputs.at(i)).info().shape();
394 else_input.info().shape(new_shape);
398 // re-sizing operands of then subgraph
399 StaticShapeInferer then_inferer(op.param().then_subg_index, _lowered_subgs);
400 _lowered_subgs.at(op.param().then_subg_index)
401 ->iterateTopolOpSeqs([&](const ir::OpSequenceIndex &, ir::OpSequence &op_seq) {
402 bool has_dynamic_tensor = then_inferer.infer(op_seq);
403 op_seq.has_dynamic_tensor(has_dynamic_tensor);
406 // re-sizing operands of else subgraph
407 StaticShapeInferer else_inferer(op.param().else_subg_index, _lowered_subgs);
408 _lowered_subgs.at(op.param().else_subg_index)
409 ->iterateTopolOpSeqs([&](const ir::OpSequenceIndex &, ir::OpSequence &op_seq) {
410 bool has_dynamic_tensor = else_inferer.infer(op_seq);
411 op_seq.has_dynamic_tensor(has_dynamic_tensor);
414 // re-sizing output shapes
415 const auto &then_outputs = _lowered_subgs.at(op.param().then_subg_index)->graph().getOutputs();
416 const auto &else_outputs = _lowered_subgs.at(op.param().else_subg_index)->graph().getOutputs();
417 assert(outputs.size() == then_outputs.size());
418 assert(outputs.size() == else_outputs.size());
419 for (size_t i = 0; i < outputs.size(); ++i)
421 const auto &then_output = then_graph.operands().at(then_outputs.at(i));
422 const auto &else_output = else_graph.operands().at(else_outputs.at(i));
423 auto &output = _operands.at(outputs.at(i));
424 if (!then_output.info().isDynamic() && !else_output.info().isDynamic() &&
425 then_output.shape() == else_output.shape())
427 output.info().shape(then_output.shape());
431 output.info().setDynamic();
432 _return_has_dynamic_tensor = true;
437 void StaticShapeInferer::visit(const ir::operation::L2Normalization &op)
439 handleSimpleUnaryOp(op, op.getInputs().at(ir::operation::L2Normalization::Input::INPUT));
442 void StaticShapeInferer::visit(const ir::operation::MatrixBandPart &op)
444 handleSimpleUnaryOp(op, op.getInputs().at(ir::operation::MatrixBandPart::Input::INPUT));
447 void StaticShapeInferer::visit(const ir::operation::OneHot &op)
449 const auto indice_idx{op.getInputs().at(ir::operation::OneHot::Input::INDICES)};
450 const auto &indice = _operands.at(indice_idx);
451 const auto depth_idx{op.getInputs().at(ir::operation::OneHot::Input::DEPTH)};
452 const auto &depth = _operands.at(depth_idx);
454 const auto axis = op.param().axis;
456 auto output_idx = op.getOutputs().at(0);
457 ir::Operand &output = _operands.at(output_idx);
459 if (!depth.isConstant())
461 output.info().setDynamic();
462 _return_has_dynamic_tensor = true;
466 const auto *depth_buf = reinterpret_cast<const int32_t *>(depth.data()->base());
468 // re-sizing output shape
469 ir::Shape new_shape = shape_inference::inferOnehotShape(indice.info().shape(), *depth_buf, axis);
470 output.info().shape(new_shape);
473 void StaticShapeInferer::visit(const ir::operation::Pack &op)
475 const auto input_idx{op.getInputs().at(0)};
476 const auto &input = _operands.at(input_idx);
478 // get mutable output operand
479 const auto output_idx = op.getOutputs().at(0);
480 ir::Operand &output = _operands.at(output_idx);
482 const auto rank = input.shape().rank() + 1;
483 const auto axis = ((op.param().axis < 0) ? rank + op.param().axis : op.param().axis);
484 const auto num = op.param().num;
486 assert(0 <= axis && axis < rank);
488 // re-sizing output shape
489 ir::Shape new_shape = shape_inference::inferPackShape(input.info().shape(), axis, rank, num);
490 output.info().shape(new_shape);
493 void StaticShapeInferer::visit(const ir::operation::Pad &op)
495 const auto input_idx{op.getInputs().at(ir::operation::Pad::Input::INPUT)};
496 const auto &input = _operands.at(input_idx);
498 const auto pad_idx{op.getInputs().at(ir::operation::Pad::Input::PAD)};
499 const auto &pad = _operands.at(pad_idx);
501 // get mutable output operand
502 const auto output_idx = op.getOutputs().at(0);
503 ir::Operand &output = _operands.at(output_idx);
505 // if pad is not constant, output also becomes dynamic
506 if (!pad.isConstant())
508 output.info().setDynamic();
509 _return_has_dynamic_tensor = true;
513 // re-sizing output shape
514 const auto new_shape = shape_inference::inferPadShape(
515 input.shape(), reinterpret_cast<const int32_t *>(pad.data()->base()),
516 pad.shape().num_elements());
517 output.info().shape(new_shape);
520 void StaticShapeInferer::visit(const ir::operation::Permute &op)
522 const auto input_idx{op.getInputs().at(0)};
523 const auto &input = _operands.at(input_idx);
524 const auto output_idx = op.getOutputs().at(0);
525 ir::Operand &output = _operands.at(output_idx);
527 // re-sizing output shape
528 // Permute is a special operation that layouts of input/output may be different on backend
529 // However, it is not applied here, so input/output have the same layout of frontend. Because
530 // "ExecutorFactory" would convert shape of input/output accoding to the layouts when registering
531 // operand info to "TensorBuilder" after calling "StaticShapeInferer"
532 const auto new_shape = input.info().shape();
533 output.info().shape(new_shape);
536 void StaticShapeInferer::visit(const ir::operation::Pow &op)
538 handleBinaryArithmeticOp(op, op.getInputs().at(ir::operation::Pow::Input::LHS),
539 op.getInputs().at(ir::operation::Pow::Input::RHS));
542 void StaticShapeInferer::visit(const ir::operation::Range &op)
544 const auto start_idx{op.getInputs().at(ir::operation::Range::Input::START)};
545 const auto limit_idx{op.getInputs().at(ir::operation::Range::Input::LIMIT)};
546 const auto delta_idx{op.getInputs().at(ir::operation::Range::Input::DELTA)};
547 const auto &start_op = _operands.at(start_idx);
548 const auto &limit_op = _operands.at(limit_idx);
549 const auto &delta_op = _operands.at(delta_idx);
551 // get mutable output operand
552 const auto output_idx = op.getOutputs().at(0);
553 ir::Operand &output = _operands.at(output_idx);
556 if (start_op.isConstant() && limit_op.isConstant() && delta_op.isConstant())
558 assert(start_op.typeInfo().type() == limit_op.typeInfo().type() &&
559 start_op.typeInfo().type() == delta_op.typeInfo().type());
560 if (output.typeInfo().type() == ir::DataType::FLOAT32)
562 new_shape = shape_inference::inferRangeShape<float>(
563 start_op.asScalar<float>(), limit_op.asScalar<float>(), delta_op.asScalar<float>());
565 else if (output.typeInfo().type() == ir::DataType::INT32)
567 new_shape = shape_inference::inferRangeShape<int32_t>(
568 start_op.asScalar<int32_t>(), limit_op.asScalar<int32_t>(), delta_op.asScalar<int32_t>());
570 assert(output.shape() == new_shape);
574 output.info().setDynamic();
575 _return_has_dynamic_tensor = true;
579 void StaticShapeInferer::visit(const ir::operation::Reduce &op)
581 const auto input_idx{op.getInputs().at(ir::operation::Reduce::Input::INPUT)};
582 const auto &input = _operands.at(input_idx);
584 const auto axes_idx{op.getInputs().at(ir::operation::Reduce::Input::AXES)};
585 const auto &axes = _operands.at(axes_idx);
587 // get mutable output operand
588 const auto output_idx = op.getOutputs().at(0);
589 ir::Operand &output = _operands.at(output_idx);
591 std::vector<int32_t> axes_vec;
592 for (size_t i = 0; i < axes.shape().num_elements(); ++i)
594 switch (axes.typeInfo().type())
596 case ir::DataType::INT32:
598 axes_vec.emplace_back(reinterpret_cast<const int32_t *>(axes.data()->base())[i]);
601 case ir::DataType::INT64:
603 axes_vec.emplace_back(reinterpret_cast<const int64_t *>(axes.data()->base())[i]);
607 throw std::runtime_error("StaticShapeInferer " + op.name() + ": Not supported data type");
611 const auto keep_dims = op.param().keep_dims;
613 // re-sizing output shape
614 ir::Shape new_shape =
615 shape_inference::inferReduceShape(input.info().shape(), axes_vec, keep_dims);
616 output.info().shape(new_shape);
619 void StaticShapeInferer::visit(const ir::operation::Reshape &op)
621 const auto input_idx{op.getInputs().at(ir::operation::Reshape::Input::INPUT)};
622 const auto &input = _operands.at(input_idx);
624 // get mutable output operand
625 const auto output_idx = op.getOutputs().at(0);
626 ir::Operand &output = _operands.at(output_idx);
628 // New shape is given by second input tensor
629 if (op.getInputs().size() == 2)
631 // Let's check the second input
632 const auto shape_idx{op.getInputs().at(ir::operation::Reshape::Input::SHAPE)};
633 const auto &shape = _operands.at(shape_idx);
635 if (shape.isConstant())
637 const auto *shape_buf = reinterpret_cast<const int32_t *>(shape.data()->base());
640 ir::Shape new_shape = shape_inference::inferReshapeShape(
641 shape_buf, shape.shape().num_elements(), input.shape().num_elements());
643 // if shape is from Const, TFLC put the shape of output into tensor
644 if (new_shape != output.shape())
646 // change on output shape
647 output.info().shape(new_shape);
652 // if shape is NOT Const, set output shape to be dynamic_
653 output.info().setDynamic();
654 _return_has_dynamic_tensor = true;
657 // New shape is given by option
658 else if (op.param().new_shape.size() != 0)
660 // Let's check the new_shape option
661 auto shape = op.param().new_shape;
662 ir::Shape new_shape = shape_inference::inferReshapeShape(shape.data(), shape.size(),
663 input.shape().num_elements());
665 if (new_shape != output.shape())
667 // change on output shape
668 output.info().shape(new_shape);
673 throw std::runtime_error("Reshape: new shape is missing");
677 void StaticShapeInferer::visit(const ir::operation::ResizeBilinear &op)
679 const auto input_idx{op.getInputs().at(ir::operation::ResizeBilinear::Input::INPUT)};
680 const auto &input = _operands.at(input_idx);
682 // get mutable output operand
683 const auto output_idx = op.getOutputs().at(0);
684 ir::Operand &output = _operands.at(output_idx);
686 // Shape inferencing logic based on Params
687 ir::Shape new_shape = shape_inference::inferResizeBilinearShape(
688 input.shape(), op.param().height_out, op.param().width_out);
690 // if size_op is from Const, TFLC put the shape of output into tensor
691 if (new_shape != output.shape())
693 // change on output shape
694 output.info().shape(new_shape);
698 void StaticShapeInferer::visit(const ir::operation::Reverse &op)
700 handleSimpleUnaryOp(op, op.getInputs().at(ir::operation::Reverse::Input::INPUT));
703 void StaticShapeInferer::visit(const ir::operation::Select &op)
705 const auto input_cond_idx{op.getInputs().at(ir::operation::Select::Input::CONDITION)};
706 const auto &input_cond = _operands.at(input_cond_idx);
708 const auto input_true_idx{op.getInputs().at(ir::operation::Select::Input::INPUT_TRUE)};
709 const auto &input_true = _operands.at(input_true_idx);
711 const auto input_false_idx{op.getInputs().at(ir::operation::Select::Input::INPUT_FALSE)};
712 const auto &input_false = _operands.at(input_false_idx);
714 auto output_idx = op.getOutputs().at(0);
715 ir::Operand &output = _operands.at(output_idx);
717 // Select output shpae
718 ir::Shape new_shape = shape_inference::inferSelectShape(
719 input_cond.info().shape(), input_true.info().shape(), input_false.info().shape());
720 output.info().shape(new_shape);
723 void StaticShapeInferer::visit(const ir::operation::Shape &op)
725 const auto input_idx{op.getInputs().at(0)};
726 const auto &input = _operands.at(input_idx);
728 // get mutable output operand
729 const auto output_idx = op.getOutputs().at(0);
730 ir::Operand &output = _operands.at(output_idx);
732 // re-sizing output shape
733 ir::Shape output_shape;
734 output_shape.append(input.info().shape().rank());
736 output.info().shape(output_shape);
739 void StaticShapeInferer::visit(const ir::operation::Slice &op)
741 const auto input_index{op.getInputs().at(ir::operation::Slice::Input::INPUT)};
742 const auto &input = _operands.at(input_index);
743 const auto begins_index{op.getInputs().at(ir::operation::Slice::Input::BEGINS)};
744 const auto &begins = _operands.at(begins_index);
745 const auto sizes_index{op.getInputs().at(ir::operation::Slice::Input::SIZES)};
746 const auto &sizes = _operands.at(sizes_index);
747 const auto output_index = op.getOutputs().at(0);
748 ir::Operand &output = _operands.at(output_index);
750 // Whether input is constant or not does not affect whether output is dynamic or not
751 if (!(begins.isConstant() && sizes.isConstant()))
753 output.info().setDynamic();
754 _return_has_dynamic_tensor = true;
758 auto begins_buf = reinterpret_cast<const int32_t *>(begins.data()->base());
759 auto sizes_buf = reinterpret_cast<const int32_t *>(sizes.data()->base());
761 ir::Shape new_shape =
762 shape_inference::inferSliceShape(input.info().shape(), begins_buf, sizes_buf);
763 output.info().shape(new_shape);
766 void StaticShapeInferer::visit(const ir::operation::Softmax &op)
768 handleSimpleUnaryOp(op, op.getInputs().at(ir::operation::Softmax::Input::INPUT));
771 void StaticShapeInferer::visit(const ir::operation::SpaceToBatchND &op)
773 const auto output_index = op.getOutputs().at(0);
774 const auto input_idx{op.getInputs().at(ir::operation::SpaceToBatchND::Input::INPUT)};
775 const auto block_shape_idx{op.getInputs().at(ir::operation::SpaceToBatchND::Input::BLOCK_SIZE)};
776 const auto padding_idx{op.getInputs().at(ir::operation::SpaceToBatchND::Input::PADDINGS)};
778 ir::Operand &output = _operands.at(output_index);
779 const auto &input = _operands.at(input_idx);
780 const auto &block_shape = _operands.at(block_shape_idx);
781 const auto &padding = _operands.at(padding_idx);
783 // Whether input is constant or not does not affect whether output is dynamic or not
784 if (!(block_shape.isConstant() && padding.isConstant()))
786 output.info().setDynamic();
787 _return_has_dynamic_tensor = true;
791 auto input_shape = input.info().shape();
792 auto block_shape_shape = block_shape.info().shape();
793 auto padding_shape = padding.info().shape();
795 auto block_shape_data = reinterpret_cast<const int32_t *>(block_shape.data()->base());
796 auto padding_data = reinterpret_cast<const int32_t *>(padding.data()->base());
798 ir::Shape new_shape = shape_inference::inferSpaceToBatchNDShape(
799 input_shape, block_shape_shape, padding_shape, block_shape_data, padding_data);
801 output.info().shape(new_shape);
804 void StaticShapeInferer::visit(const ir::operation::Split &op)
806 const auto input_idx{op.getInputs().at(0)};
807 const auto &input = _operands.at(input_idx);
809 const auto axis = op.param().axis;
810 const auto num_splits = op.param().num_splits;
812 const auto rank = input.info().shape().rank();
813 auto axis_resolved = axis < 0 ? axis + rank : axis;
815 assert(0 <= axis_resolved && axis_resolved < rank);
817 ir::Shape new_shape =
818 shape_inference::inferSplitShape(input.info().shape(), axis_resolved, num_splits);
819 auto output_tensors = op.getOutputs();
820 for (auto output_idx : output_tensors)
822 ir::Operand &output = _operands.at(output_idx);
823 output.info().shape(new_shape);
827 void StaticShapeInferer::visit(const ir::operation::SquaredDifference &op)
829 handleBinaryArithmeticOp(op, op.getInputs().at(ir::operation::SquaredDifference::Input::LHS),
830 op.getInputs().at(ir::operation::SquaredDifference::Input::RHS));
833 void StaticShapeInferer::visit(const ir::operation::Squeeze &op)
835 const auto input_idx{op.getInputs().at(ir::operation::Squeeze::Input::INPUT)};
836 const auto &input = _operands.at(input_idx);
838 const auto output_idx = op.getOutputs().at(0);
839 ir::Operand &output = _operands.at(output_idx);
841 if (input.info().isDynamic())
843 output.info().setDynamic();
844 _return_has_dynamic_tensor = true;
848 // Squeeze output shpae
849 ir::Shape new_shape = shape_inference::inferSqueezeShape(input.info().shape(), op.param());
850 output.info().shape(new_shape);
853 void StaticShapeInferer::visit(const ir::operation::StridedSlice &op)
855 const auto input_index{op.getInputs().at(ir::operation::StridedSlice::Input::INPUT)};
856 const auto &input = _operands.at(input_index);
857 const auto starts_index{op.getInputs().at(ir::operation::StridedSlice::Input::STARTS)};
858 const auto &starts = _operands.at(starts_index);
859 const auto ends_index{op.getInputs().at(ir::operation::StridedSlice::Input::ENDS)};
860 const auto &ends = _operands.at(ends_index);
861 const auto strides_index{op.getInputs().at(ir::operation::StridedSlice::Input::STRIDES)};
862 const auto &strides = _operands.at(strides_index);
863 const auto output_index = op.getOutputs().at(0);
864 ir::Operand &output = _operands.at(output_index);
866 if (!(starts.isConstant() && ends.isConstant() && strides.isConstant()))
868 output.info().setDynamic();
869 _return_has_dynamic_tensor = true;
873 const auto begin_mask = op.param().begin_mask;
874 const auto end_mask = op.param().end_mask;
875 const auto shrink_axis_mask = op.param().shrink_axis_mask;
876 const auto rank = input.info().shape().rank();
878 auto starts_buf = reinterpret_cast<const uint32_t *>(starts.data()->base());
879 auto ends_buf = reinterpret_cast<const uint32_t *>(ends.data()->base());
880 auto strides_buf = reinterpret_cast<const uint32_t *>(strides.data()->base());
882 auto op_params = shape_inference::buildStridedSliceParams(
883 starts_buf, ends_buf, strides_buf, begin_mask, end_mask, shrink_axis_mask, rank);
885 ir::Shape new_shape =
886 shape_inference::inferStridedSliceShape(input.info().shape(), op_params, rank);
887 output.info().shape(new_shape);
890 void StaticShapeInferer::visit(const ir::operation::Tile &op)
892 const auto input_idx{op.getInputs().at(ir::operation::Tile::Input::INPUT)};
893 const auto &input = _operands.at(input_idx);
895 const auto multiplier_idx{op.getInputs().at(ir::operation::Tile::Input::MULTIPLES)};
896 const auto &multiplier = _operands.at(multiplier_idx);
898 const auto output_idx = op.getOutputs().at(0);
899 ir::Operand &output = _operands.at(output_idx);
901 if (!multiplier.isConstant())
903 output.info().setDynamic();
904 _return_has_dynamic_tensor = true;
908 auto multiplier_buffer = reinterpret_cast<const int32_t *>(multiplier.data()->base());
909 assert(multiplier_buffer);
911 // re-sizing output shape
912 auto new_shape = shape_inference::inferTileShape(input.info().shape(), multiplier_buffer);
913 output.info().shape(new_shape);
916 void StaticShapeInferer::visit(const ir::operation::Transpose &op)
918 const auto input_idx{op.getInputs().at(ir::operation::Transpose::Input::INPUT)};
919 const auto &input = _operands.at(input_idx);
921 // get mutable output operand
922 const auto output_idx = op.getOutputs().at(0);
923 ir::Operand &output = _operands.at(output_idx);
924 const auto perm{op.param().perm};
925 // const auto rank{op.param().rank};
927 // set output shape, based on input and params
928 ir::Shape new_shape = shape_inference::inferTransposeShape(input.info().shape(), perm);
929 output.info().shape(new_shape);
932 void StaticShapeInferer::visit(const ir::operation::Unpack &op)
934 const auto input_idx{op.getInputs().at(0)};
935 const auto &input = _operands.at(input_idx);
936 const auto num = op.param().num;
937 const auto rank = input.shape().rank();
938 const auto axis = ((op.param().axis < 0) ? rank + op.param().axis : op.param().axis);
943 for (int out_tensor_idx = 0; out_tensor_idx < num; out_tensor_idx++)
945 const auto output_idx = op.getOutputs().at(out_tensor_idx);
946 ir::Operand &output = _operands.at(output_idx);
947 output.info().setDynamic();
949 _return_has_dynamic_tensor = true;
953 ir::Shape new_shape = shape_inference::inferUnpackShape(input.info().shape(), axis, rank);
955 // re-sizing output shape
956 for (int out_tensor_idx = 0; out_tensor_idx < num; out_tensor_idx++)
958 const auto output_idx = op.getOutputs().at(out_tensor_idx);
959 ir::Operand &output = _operands.at(output_idx);
960 output.info().shape(new_shape);
964 void StaticShapeInferer::visit(const ir::operation::While &op)
966 auto &cond_graph = _lowered_subgs.at(op.param().cond_subg_index)->graph();
967 auto &body_graph = _lowered_subgs.at(op.param().body_subg_index)->graph();
968 const auto inputs = op.getInputs();
969 const auto &outputs = op.getOutputs();
971 // re-sizing input shapes of then subgraph
972 const auto &cond_inputs = cond_graph.getInputs();
973 assert(inputs.size() == cond_inputs.size());
974 for (size_t i = 0; i < inputs.size(); ++i)
976 const auto &input = _operands.at(inputs.at(i));
977 auto &cond_input = cond_graph.operands().at(cond_inputs.at(i));
978 if (input.info().isDynamic())
980 cond_input.info().setDynamic();
984 auto new_shape = input.info().shape();
985 cond_input.info().shape(new_shape);
989 // re-sizing input shapes of body subgraph
990 const auto &body_inputs = body_graph.getInputs();
991 assert(cond_inputs.size() == body_inputs.size());
992 for (size_t i = 0; i < cond_inputs.size(); ++i)
994 const auto &cond_input = cond_graph.operands().at(cond_inputs.at(i));
995 auto &body_input = body_graph.operands().at(body_inputs.at(i));
996 if (cond_input.info().isDynamic())
998 body_input.info().setDynamic();
1002 const auto &new_shape = cond_input.info().shape();
1003 body_input.info().shape(new_shape);
1007 // re-sizing operands of body subgraph
1008 StaticShapeInferer body_inferer(op.param().body_subg_index, _lowered_subgs);
1009 _lowered_subgs.at(op.param().body_subg_index)
1010 ->iterateTopolOpSeqs([&](const ir::OpSequenceIndex &, ir::OpSequence &op_seq) {
1011 bool has_dynamic_tensor = body_inferer.infer(op_seq);
1012 op_seq.has_dynamic_tensor(has_dynamic_tensor);
1015 // Check whether while operation's shapes are predictable
1016 // If any of shape of body outputs and cond inputs are different, non-constant operands would be
1018 bool check_unpredictable_dynamic = false;
1019 const auto &body_outputs = body_graph.getOutputs();
1020 assert(body_outputs.size() == cond_inputs.size());
1021 for (size_t i = 0; i < body_outputs.size(); ++i)
1023 const auto &body_output = body_graph.operands().at(body_outputs.at(i));
1024 auto &cond_input = cond_graph.operands().at(cond_inputs.at(i));
1025 if ((cond_input.info().isDynamic() != body_output.info().isDynamic()) ||
1026 (cond_input.shape() != body_output.shape()))
1028 check_unpredictable_dynamic = true;
1033 if (check_unpredictable_dynamic)
1035 // Set inputs of body subgraph
1036 for (const auto &input_index : body_inputs)
1038 auto &input = body_graph.operands().at(input_index);
1039 if (!input.isConstant())
1041 input.info().setDynamic();
1045 // Set inputs of cond subgraph
1046 for (const auto &input_index : cond_inputs)
1048 auto &input = cond_graph.operands().at(input_index);
1049 if (!input.isConstant())
1051 input.info().setDynamic();
1055 // Set non-constant operands of body subgraph to dynamic
1056 StaticShapeInferer body_inferer(op.param().body_subg_index, _lowered_subgs);
1057 _lowered_subgs.at(op.param().body_subg_index)
1058 ->iterateTopolOpSeqs([&](const ir::OpSequenceIndex &, ir::OpSequence &op_seq) {
1059 bool has_dynamic_tensor = body_inferer.infer(op_seq);
1060 op_seq.has_dynamic_tensor(has_dynamic_tensor);
1064 // re-sizing operands of cond subgraph
1065 // If check_unpredictable_dynamic is true, non-constant operands of cond subgraph would be set to
1067 StaticShapeInferer cond_inferer(op.param().cond_subg_index, _lowered_subgs);
1068 _lowered_subgs.at(op.param().cond_subg_index)
1069 ->iterateTopolOpSeqs([&](const ir::OpSequenceIndex &, ir::OpSequence &op_seq) {
1070 bool has_dynamic_tensor = cond_inferer.infer(op_seq);
1071 op_seq.has_dynamic_tensor(has_dynamic_tensor);
1074 // re-sizing outputs of while operation
1075 // If check_unpredictable_dynamic is true, outputs of while operation would be set to dynamic
1076 assert(cond_inputs.size() == outputs.size());
1077 for (size_t i = 0; i < cond_inputs.size(); ++i)
1079 const auto &cond_input = cond_graph.operands().at(cond_inputs.at(i));
1080 auto &output = _operands.at(outputs.at(i));
1081 if (cond_input.info().isDynamic())
1083 output.info().setDynamic();
1084 _return_has_dynamic_tensor = true;
1088 const auto new_shape = cond_input.info().shape();
1089 output.info().shape(new_shape);
1094 } // namespace compiler
1096 } // namespace onert