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 "luci/Pass/ConvertNCHWToNHWCPass.h"
18 #include "CircleOptimizerUtils.h"
20 #include <luci/IR/CircleNodes.h>
21 #include <luci/IR/CircleNodeVisitor.h>
22 #include <luci/Profile/CircleNodeOrigin.h>
23 #include <luci/Service/Nodes/CircleConst.h>
31 // Return true if from can be broadcasted to to
32 // to's shape is [N, C, H, W]
33 bool broadcastable(const luci::CircleConst *from, const luci::CircleNode *to)
35 assert(to->rank() == 4); // FIX_CALLER_UNLESS
37 const auto from_rank = from->rank();
41 // Scalar is always broadcastable
45 for (uint32_t i = 1; i <= from_rank; i++)
47 auto to_index = 4 - i;
48 auto from_index = from_rank - i;
50 if (from->dim(from_index).value() != to->dim(to_index).value() and
51 from->dim(from_index).value() != 1)
58 // Expand node to rank 4
59 // node should have rank less than or equal to 4
60 void expand_to_rank_4(luci::CircleConst *node)
62 auto original_rank = node->rank();
64 assert(original_rank <= 4); // FIX_CALLER_UNLESS
66 if (original_rank == 4)
69 std::vector<uint32_t> original_shape;
70 for (uint32_t i = 0; i < original_rank; i++)
72 original_shape.emplace_back(node->dim(i).value());
76 for (uint32_t i = 0; i < (4 - original_rank); i++)
79 for (uint32_t i = 0; i < original_rank; i++)
80 node->dim(i + (4 - original_rank)) = original_shape.at(i);
83 bool is_output(const loco::Node *node)
85 auto cnode = loco::must_cast<const luci::CircleNode *>(node);
86 auto opcode = cnode->opcode();
87 if (opcode == luci::CircleOpcode::CIRCLEOUTPUT ||
88 opcode == luci::CircleOpcode::CIRCLEOUTPUTEXCLUDE)
94 bool is_same_shape(const luci::CircleNode *node, const std::vector<loco::Dimension> &shape)
99 if (shape.size() != node->rank())
102 for (uint32_t i = 0; i < shape.size(); i++)
104 if (not(node->dim(i) == shape[i]))
110 enum class DataFormat
117 * @brief Set annotation for DataFormat (NCHW, NHWC)
119 * @note DataFormatAnnotation will live longer than this Pass (until the
120 * annotated loco::Node is erased). So, do not use large data in the
121 * annotation to avoid excessive memory usage.
123 class DataFormatAnnotation final : public loco::NodeAnnotation
126 DataFormatAnnotation(const DataFormat &format) : _format{format}
132 const DataFormat &format(void) const { return _format; }
138 void set_data_format(loco::Node *node, const DataFormat &format)
140 node->annot(std::make_unique<DataFormatAnnotation>(format));
143 DataFormat get_data_format(loco::Node *node)
145 assert(node->annot<DataFormatAnnotation>() != nullptr);
146 return node->annot<DataFormatAnnotation>()->format();
149 bool has_data_format(loco::Node *node) { return node->annot<DataFormatAnnotation>() != nullptr; }
151 bool check_4d_transpose(loco::Node *node, const std::vector<int32_t> indices)
153 assert(indices.size() == 4);
155 auto trans = dynamic_cast<luci::CircleTranspose *>(node);
159 if (not trans->perm())
162 auto perm = dynamic_cast<luci::CircleConst *>(trans->perm());
163 // Only const perm is supported
167 if (perm->dtype() != loco::DataType::S32)
170 if (perm->size<loco::DataType::S32>() != 4)
173 for (uint32_t i = 0; i < 4; i++)
175 if (perm->at<loco::DataType::S32>(i) != indices[i])
182 luci::CircleTranspose *create_4d_transpose(luci::CircleNode *node,
183 const std::vector<int32_t> indices)
185 assert(indices.size() == 4);
187 auto name = node->name();
188 assert(name.length() > 0);
190 auto perm = node->graph()->nodes()->create<luci::CircleConst>();
191 perm->dtype(loco::DataType::S32);
192 perm->size<loco::DataType::S32>(4);
195 for (uint32_t i = 0; i < 4; i++)
196 perm->at<loco::DataType::S32>(i) = indices[i];
197 perm->shape_status(luci::ShapeStatus::VALID);
199 auto make_string = [](const std::vector<int32_t> &nums) {
201 for (auto num : nums)
203 if (str.length() > 0)
205 str += std::to_string(num);
210 auto str_indices = make_string(indices);
212 perm->name(name + "/Transpose_" + str_indices + "/perm");
214 auto trans = node->graph()->nodes()->create<luci::CircleTranspose>();
216 trans->name(name + "/Transpose_" + str_indices);
217 luci::add_origin(trans, luci::get_origin(node));
222 luci::CircleTranspose *create_Nd_transpose(luci::CircleNode *node,
223 const std::vector<int32_t> indices)
225 auto name = node->name();
226 assert(name.length() > 0);
228 auto perm = node->graph()->nodes()->create<luci::CircleConst>();
229 perm->dtype(loco::DataType::S32);
230 perm->size<loco::DataType::S32>(indices.size());
232 perm->dim(0) = indices.size();
233 for (uint32_t i = 0; i < indices.size(); i++)
234 perm->at<loco::DataType::S32>(i) = indices[i];
235 perm->shape_status(luci::ShapeStatus::VALID);
237 auto make_string = [](const std::vector<int32_t> &nums) {
239 for (auto num : nums)
241 if (str.length() > 0)
243 str += std::to_string(num);
248 auto str_indices = make_string(indices);
250 perm->name(name + "/Transpose_" + str_indices + "/perm");
252 auto trans = node->graph()->nodes()->create<luci::CircleTranspose>();
254 trans->name(name + "/Transpose_" + str_indices);
255 luci::add_origin(trans, luci::get_origin(node));
260 int32_t nchw_axis_to_nhwc(int32_t axis)
262 uint32_t pos_axis = axis >= 0 ? static_cast<uint32_t>(axis) : static_cast<uint32_t>(axis + 4);
263 static const uint32_t to_nhwc[4] = {0, 3, 1, 2};
265 throw std::runtime_error("Concat axis must be in range [-4, 4)");
266 return to_nhwc[pos_axis];
269 luci::CircleTranspose *create_post_transpose(luci::CircleNode *node)
271 return create_4d_transpose(node, {0, 3, 1, 2});
274 luci::CircleTranspose *create_pre_transpose(luci::CircleNode *node)
276 return create_4d_transpose(node, {0, 2, 3, 1});
279 bool check_4d_reshape(loco::Node *node, const std::vector<int32_t> indices)
281 assert(indices.size() == 4); // FIX_CALLER_UNLESS
283 auto reshape = dynamic_cast<luci::CircleReshape *>(node);
287 if (reshape->rank() != 4)
290 auto input = loco::must_cast<luci::CircleNode *>(reshape->tensor());
291 if (input->shape_status() != luci::ShapeStatus::VALID)
294 if (input->rank() != 4)
297 if (reshape->shape_status() != luci::ShapeStatus::VALID)
300 if (!(input->dim(0) == reshape->dim(indices[0])) ||
301 !(input->dim(1) == reshape->dim(indices[1])) ||
302 !(input->dim(2) == reshape->dim(indices[2])) || !(input->dim(3) == reshape->dim(indices[3])))
308 // Check if Reshape that converts NCHW -> NHWC
309 bool is_pre_reshape(loco::Node *node) { return check_4d_reshape(node, {0, 3, 1, 2}); }
311 // Check if Reshape that converts NHWC -> NCHW
312 bool is_post_reshape(loco::Node *node) { return check_4d_reshape(node, {0, 2, 3, 1}); }
314 bool is_post_transpose(loco::Node *node) { return check_4d_transpose(node, {0, 3, 1, 2}); }
316 bool is_pre_transpose(loco::Node *node) { return check_4d_transpose(node, {0, 2, 3, 1}); }
318 uint32_t cal_offset(const loco::TensorShape &dimension, const uint32_t *indices)
320 return indices[0] * dimension.dim(1).value() * dimension.dim(2).value() *
321 dimension.dim(3).value() +
322 indices[1] * dimension.dim(2).value() * dimension.dim(3).value() +
323 indices[2] * dimension.dim(3).value() + indices[3];
326 luci::CircleConst *create_NHWC_paddings(luci::CircleConst *paddings)
328 // paddings shape is (4,2) (it was checked by is_NCHW)
329 assert(paddings != nullptr);
330 assert(paddings->rank() == 2);
331 assert(paddings->dim(0).value() == 4);
332 assert(paddings->dim(1).value() == 2);
334 // paddings for idx 0~3 are 0 (checked by is_NCHW)
335 assert(paddings->at<loco::DataType::S32>(0) == 0);
336 assert(paddings->at<loco::DataType::S32>(1) == 0);
337 assert(paddings->at<loco::DataType::S32>(2) == 0);
338 assert(paddings->at<loco::DataType::S32>(3) == 0);
340 auto name = paddings->name();
341 assert(name.length() > 0);
343 auto nhwc_paddings = paddings->graph()->nodes()->create<luci::CircleConst>();
344 nhwc_paddings->dtype(loco::DataType::S32);
345 nhwc_paddings->shape({4, 2});
346 nhwc_paddings->shape_status(luci::ShapeStatus::VALID);
347 nhwc_paddings->size<loco::DataType::S32>(4 * 2);
348 nhwc_paddings->name(name + "_NHWC");
350 for (uint32_t dim = 0; dim < 4; dim++)
352 for (uint32_t i = 0; i < 2; i++)
358 // get third dimension (H in NCHW)
359 data = paddings->at<loco::DataType::S32>(2 * 2 + i);
363 // get fourth dimension (W in NCHW)
364 data = paddings->at<loco::DataType::S32>(3 * 2 + i);
367 nhwc_paddings->at<loco::DataType::S32>(dim * 2 + i) = data;
370 return nhwc_paddings;
373 luci::CircleConst *create_NHWC_rindices(luci::CircleConst *rindices)
375 assert(rindices != nullptr); // FIX_CALLER_UNLESS
377 if (rindices->dtype() != loco::DataType::S32)
380 auto nhwc_rindices = luci::clone(rindices);
381 auto name = rindices->name();
382 assert(name.length() > 0); // FIX_CALLER_UNLESS
383 nhwc_rindices->name(name + "_NHWC");
385 auto size = nhwc_rindices->size<loco::DataType::S32>();
386 for (uint32_t i = 0; i < size; i++)
388 nhwc_rindices->at<loco::DataType::S32>(i) =
389 nchw_axis_to_nhwc(rindices->at<loco::DataType::S32>(i));
392 return nhwc_rindices;
395 luci::CircleConst *create_NHWC_from_NCHW(luci::CircleConst *constant)
398 assert(constant->rank() == 4);
400 // TODO: Support non-float types
401 if (constant->dtype() != loco::DataType::FLOAT32)
403 INFO(l) << "Non-float type constant: " << constant->name() << std::endl;
407 loco::TensorShape nchw_dimension{constant->dim(0), constant->dim(1), constant->dim(2),
409 loco::TensorShape nhwc_dimension{constant->dim(0), constant->dim(2), constant->dim(3),
412 auto name = constant->name();
413 assert(name.length() > 0);
415 auto nhwc_const = constant->graph()->nodes()->create<luci::CircleConst>();
416 nhwc_const->dtype(constant->dtype());
418 nhwc_const->dim(0).set(constant->dim(0).value());
419 nhwc_const->dim(1).set(constant->dim(2).value());
420 nhwc_const->dim(2).set(constant->dim(3).value());
421 nhwc_const->dim(3).set(constant->dim(1).value());
422 nhwc_const->shape_status(luci::ShapeStatus::VALID);
423 nhwc_const->size<loco::DataType::FLOAT32>(constant->size<loco::DataType::FLOAT32>());
424 nhwc_const->name(name + "_NHWC");
426 for (uint32_t n = 0; n < nchw_dimension.dim(0).value(); n++)
428 for (uint32_t c = 0; c < nchw_dimension.dim(1).value(); c++)
430 for (uint32_t h = 0; h < nchw_dimension.dim(2).value(); h++)
432 for (uint32_t w = 0; w < nchw_dimension.dim(3).value(); w++)
434 uint32_t nchw_indices[4] = {n, c, h, w};
435 uint32_t nhwc_indices[4] = {n, h, w, c};
437 constant->at<loco::DataType::FLOAT32>(cal_offset(nchw_dimension, nchw_indices));
438 nhwc_const->at<loco::DataType::FLOAT32>(cal_offset(nhwc_dimension, nhwc_indices)) = data;
446 // NOTE Following conditions can be extended later
448 // Find PAD with an NCHW pattern described below
449 // - Paddings shape : [4, 2]
450 // - Paddings value : [[0, 0], [0, 0], [h_t, h_b], [w_t, w_b]]]
451 bool is_NCHW(const luci::CirclePad *node)
453 const auto paddings = dynamic_cast<luci::CircleConst *>(node->paddings());
454 // Non-const paddings is not supported
455 if (paddings == nullptr)
458 if (paddings->rank() != 2)
461 if (paddings->dim(0).value() != 4 || paddings->dim(1).value() != 2)
464 // Only check the first two dimensions
465 for (uint32_t dim = 0; dim < 2; dim++)
467 for (uint32_t i = 0; i < 2; i++)
469 auto data = paddings->at<loco::DataType::S32>(dim * 2 + i);
478 // NOTE Copied from is_NCHW(CirclePad)
479 bool is_NCHW(const luci::CirclePadV2 *node)
481 const auto paddings = dynamic_cast<luci::CircleConst *>(node->paddings());
482 // Non-const paddings is not supported
483 if (paddings == nullptr)
486 if (paddings->rank() != 2)
489 if (paddings->dim(0).value() != 4 || paddings->dim(1).value() != 2)
492 // Only check the first two dimensions
493 for (uint32_t dim = 0; dim < 2; dim++)
495 for (uint32_t i = 0; i < 2; i++)
497 auto data = paddings->at<loco::DataType::S32>(dim * 2 + i);
506 bool is_const(const loco::Node *node)
508 if (not dynamic_cast<const luci::CircleConst *>(node))
514 bool is_scalar_const(const loco::Node *node)
516 auto const_node = dynamic_cast<const luci::CircleConst *>(node);
520 const auto const_rank = const_node->rank();
523 // 2. rank = 1, dimension = 1
527 if (const_rank == 1 && const_node->dim(0).value() == 1)
533 // NOTE Following conditions can be extended later
535 // Find MUL with an NCHW pattern described below
536 // - Input (non-constant) shape : [N, C, H, W]
537 // - Input (constant) shape : broadcastable to [N, C, H, W]
538 // - Output shape : [N, C, H, W]
539 bool is_NCHW_with_const(const luci::CircleMul *node, luci::CircleNode *&pred_node,
540 luci::CircleConst *&multiplier)
542 auto x = dynamic_cast<luci::CircleConst *>(node->x());
543 auto y = dynamic_cast<luci::CircleConst *>(node->y());
545 if (x != nullptr && y == nullptr)
547 pred_node = loco::must_cast<luci::CircleNode *>(node->y());
550 else if (x == nullptr && y != nullptr)
552 pred_node = loco::must_cast<luci::CircleNode *>(node->x());
557 // Ignore if MUL does not have a multiplier input.
561 if (pred_node->rank() != 4)
564 if (not broadcastable(multiplier, node))
567 expand_to_rank_4(multiplier);
572 // We assume ADD with const input is NCHW if,
573 // Input shape: (N, C, H, W)
574 // Output shape: (N, C, H, W)
575 // 1. Const shape is (1, C, 1, 1), (N, C, H, W) or a scalar (1)
576 // 2. Input, Output, Const have the same C.
577 bool is_NCHW_with_const(const luci::CircleAdd *node, luci::CircleNode *&pred_node,
578 luci::CircleConst *&beta)
580 auto x = dynamic_cast<luci::CircleConst *>(node->x());
581 auto y = dynamic_cast<luci::CircleConst *>(node->y());
583 if (x != nullptr && y == nullptr)
585 pred_node = loco::must_cast<luci::CircleNode *>(node->y());
588 else if (x == nullptr && y != nullptr)
590 pred_node = loco::must_cast<luci::CircleNode *>(node->x());
595 // Ignore if ADD does not have a constant input.
599 if (pred_node->rank() != 4)
602 if (not broadcastable(beta, node))
605 expand_to_rank_4(beta);
610 // We assume SUB with const input is NCHW if,
611 // Input shape: (N, C, H, W)
612 // Output shape: (N, C, H, W)
613 // 1. Const shape is (1, C, 1, 1), (N, C, H, W) or a scalar (1)
614 // 2. Input, Output, Const have the same C.
615 bool is_NCHW_with_const(const luci::CircleSub *node, const luci::CircleNode *pred_node,
616 const luci::CircleConst *subtract)
618 assert(pred_node != nullptr);
619 assert(subtract != nullptr);
621 if (pred_node->rank() != 4)
624 const auto const_rank = subtract->rank();
625 // Support Rank 4 or scalar (rank 0 or 1)
626 if (const_rank != 4 && const_rank != 0 && const_rank != 1)
629 const auto input_cdim = pred_node->dim(1);
630 const auto output_cdim = node->dim(1);
634 bool supported_shape = false;
636 // Check subtract is (1, C, 1, 1)
637 if (is_same_shape(subtract, {1, node->dim(1), 1, 1}))
638 supported_shape = true;
640 // Check subtract is (N, C, H, W)
641 if (is_same_shape(subtract, {node->dim(0), node->dim(1), node->dim(2), node->dim(3)}))
642 supported_shape = true;
644 return supported_shape;
646 if (input_cdim == output_cdim)
652 template <class T> bool convert_unary_features(T *node)
654 const auto pred_node = loco::must_cast<luci::CircleNode *>(node->features());
655 auto pre_trans = create_pre_transpose(node);
656 pre_trans->a(pred_node);
657 node->features(pre_trans);
659 // Do shape inference for this node again.
660 node->shape_status(luci::ShapeStatus::UNDEFINED);
662 auto post_trans = create_post_transpose(node);
663 loco::replace(node).with(post_trans);
670 template <class T> bool convert_unary_x(T *node)
672 const auto pred_node = loco::must_cast<luci::CircleNode *>(node->x());
673 auto pre_trans = create_pre_transpose(node);
674 pre_trans->a(pred_node);
677 // Do shape inference for this node again.
678 node->shape_status(luci::ShapeStatus::UNDEFINED);
680 auto post_trans = create_post_transpose(node);
681 loco::replace(node).with(post_trans);
688 template <class T> bool convert_unary_logits(T *node)
690 const auto pred_node = loco::must_cast<luci::CircleNode *>(node->logits());
691 auto pre_trans = create_pre_transpose(node);
692 pre_trans->a(pred_node);
693 node->logits(pre_trans);
695 // Do shape inference for this node again.
696 node->shape_status(luci::ShapeStatus::UNDEFINED);
698 auto post_trans = create_post_transpose(node);
699 loco::replace(node).with(post_trans);
706 class ConvertNCHWToNHWC final : public luci::CircleNodeMutableVisitor<bool>
709 bool visit(luci::CircleNode *node)
711 throw std::runtime_error(node->name() + " is an unsupported operator.");
714 bool visit(luci::CircleInput *node)
716 const auto n = node->dim(0);
717 const auto c = node->dim(1);
718 const auto h = node->dim(2);
719 const auto w = node->dim(3);
725 // Do shape inference for this node again.
726 node->shape_status(luci::ShapeStatus::UNDEFINED);
728 // Insert post-tranpose
729 auto post_trans = create_post_transpose(node);
730 loco::replace(node).with(post_trans);
734 // Update graph input
735 auto graph_inputs = node->graph()->inputs();
736 auto graph_input = graph_inputs->at(node->index());
737 graph_input->shape({n, h, w, c});
742 bool visit(luci::CircleOutput *node)
744 // Insert pre-transpose
745 auto pre_trans = create_pre_transpose(node);
746 pre_trans->a(node->from());
748 node->from(pre_trans);
750 // Do shape inference for this node again.
751 node->shape_status(luci::ShapeStatus::UNDEFINED);
753 // Update graph output
754 const auto n = node->dim(0).value();
755 const auto c = node->dim(1).value();
756 const auto h = node->dim(2).value();
757 const auto w = node->dim(3).value();
759 auto graph_outputs = node->graph()->outputs();
760 auto graph_output = graph_outputs->at(node->index());
761 graph_output->shape({n, h, w, c});
766 bool visit(luci::CircleAdd *node)
768 luci::CircleNode *pred_node = nullptr;
769 luci::CircleConst *beta = nullptr;
771 if (is_NCHW_with_const(node, pred_node, beta))
773 assert(beta->rank() == 4); // FIX is_NCHW_with_const unless
774 auto nhwc_const = create_NHWC_from_NCHW(beta);
775 if (nhwc_const == nullptr)
779 auto pre_trans = create_pre_transpose(node);
780 pre_trans->a(pred_node);
783 else if (beta == nullptr)
785 // Both inputs are not constant.
786 // In this case, we cannot distinguish NCHW from NHWC,
787 // so just insert Transpose Ops.
788 auto pre_trans_x = create_pre_transpose(node);
789 pre_trans_x->a(node->x());
790 node->x(pre_trans_x);
792 auto pre_trans_y = create_pre_transpose(node);
793 pre_trans_y->a(node->y());
794 node->y(pre_trans_y);
801 // Do shape inference for this node again.
802 node->shape_status(luci::ShapeStatus::UNDEFINED);
804 auto post_trans = create_post_transpose(node);
805 loco::replace(node).with(post_trans);
811 bool visit(luci::CircleConcatenation *node)
813 const auto num_values = node->numValues();
814 for (uint32_t i = 0; i < num_values; i++)
816 auto pred_node = loco::must_cast<luci::CircleNode *>(node->values(i));
817 auto pre_trans = create_pre_transpose(node);
818 pre_trans->a(pred_node);
819 node->values(i, pre_trans);
822 // Do shape inference for this node again.
823 node->shape_status(luci::ShapeStatus::UNDEFINED);
825 node->axis(nchw_axis_to_nhwc(node->axis()));
827 auto post_trans = create_post_transpose(node);
828 loco::replace(node).with(post_trans);
835 bool visit(luci::CircleElu *node) { return convert_unary_features<luci::CircleElu>(node); }
837 bool visit(luci::CircleLeakyRelu *node)
839 return convert_unary_features<luci::CircleLeakyRelu>(node);
842 bool visit(luci::CircleLogistic *node) { return convert_unary_x<luci::CircleLogistic>(node); }
844 bool visit(luci::CircleMaximum *node)
846 if ((not is_const(node->x())) and is_scalar_const(node->y()))
848 auto pre_trans = create_pre_transpose(node);
849 pre_trans->a(node->x());
852 else if (is_scalar_const(node->x()) and (not is_const(node->y())))
854 auto pre_trans = create_pre_transpose(node);
855 pre_trans->a(node->y());
858 else if ((not is_const(node->x())) and (not is_const(node->y())))
860 auto pre_trans_x = create_pre_transpose(node);
861 pre_trans_x->a(node->x());
862 node->x(pre_trans_x);
864 auto pre_trans_y = create_pre_transpose(node);
865 pre_trans_y->a(node->y());
866 node->y(pre_trans_y);
870 // TODO support other cases
874 // Do shape inference for this node again.
875 node->shape_status(luci::ShapeStatus::UNDEFINED);
877 auto post_trans = create_post_transpose(node);
878 loco::replace(node).with(post_trans);
884 bool visit(luci::CircleMean *node)
886 auto input = loco::must_cast<luci::CircleNode *>(node->input());
887 if (input->rank() != 4)
890 auto rindices = dynamic_cast<luci::CircleConst *>(node->reduction_indices());
894 auto nhwc_rindices = create_NHWC_rindices(rindices);
895 if (not nhwc_rindices)
898 auto pre_trans = create_pre_transpose(node);
900 node->input(pre_trans);
902 // Do shape inference for this node again.
903 node->shape_status(luci::ShapeStatus::UNDEFINED);
905 node->reduction_indices(nhwc_rindices);
907 if (node->keep_dims())
909 auto post_trans = create_post_transpose(node);
910 loco::replace(node).with(post_trans);
917 // node->keep_dims() == false
918 // 1D output never needs a transpose
919 if (node->rank() <= 1)
922 std::vector<bool> reduced_dims_nhwc(4, false);
923 uint32_t num_reduced_indices = nhwc_rindices->size<loco::DataType::S32>();
925 for (uint32_t ri = 0; ri < num_reduced_indices; ++ri)
927 reduced_dims_nhwc[nhwc_rindices->at<loco::DataType::S32>(ri)] = true;
930 // if channel dimension has been reduced, we don't need a transpose
931 if (reduced_dims_nhwc[3])
934 // likewise, if both space dimensions are reduced, no transpose is needed
935 if (reduced_dims_nhwc[1] && reduced_dims_nhwc[2])
938 std::vector<int32_t> post_trans_ind;
939 // case 1: only N is reduced
940 if (num_reduced_indices == 1 && reduced_dims_nhwc[0])
941 post_trans_ind = {2, 0, 1};
943 // case 2: only H or W is reduced
944 if (num_reduced_indices == 1 && (reduced_dims_nhwc[1] || reduced_dims_nhwc[2]))
945 post_trans_ind = {0, 2, 1};
947 // case 3: N and either H or W are reduced
948 if (num_reduced_indices == 2)
949 post_trans_ind = {1, 0};
951 auto post_trans = create_Nd_transpose(node, post_trans_ind);
952 loco::replace(node).with(post_trans);
959 bool visit(luci::CircleMinimum *node)
961 if ((not is_const(node->x())) and is_scalar_const(node->y()))
963 auto pre_trans = create_pre_transpose(node);
964 pre_trans->a(node->x());
967 else if (is_scalar_const(node->x()) and (not is_const(node->y())))
969 auto pre_trans = create_pre_transpose(node);
970 pre_trans->a(node->y());
975 // TODO support other cases
979 // Do shape inference for this node again.
980 node->shape_status(luci::ShapeStatus::UNDEFINED);
982 auto post_trans = create_post_transpose(node);
983 loco::replace(node).with(post_trans);
989 bool visit(luci::CircleMul *node)
993 luci::CircleNode *pred_node = nullptr;
994 luci::CircleConst *multiplier = nullptr;
996 if (is_NCHW_with_const(node, pred_node, multiplier))
998 assert(multiplier->rank() == 4); // FIX is_NCHW_with_const unless
999 auto nhwc_const = create_NHWC_from_NCHW(multiplier);
1000 if (nhwc_const == nullptr)
1002 node->y(nhwc_const);
1004 auto pre_trans = create_pre_transpose(node);
1005 pre_trans->a(pred_node);
1008 else if (multiplier == nullptr)
1010 // Only support for input rank 4
1011 auto input_x = loco::must_cast<luci::CircleNode *>(node->x());
1012 if (input_x->rank() != 4)
1014 auto input_y = loco::must_cast<luci::CircleNode *>(node->y());
1015 if (input_y->rank() != 4)
1018 auto pre_trans_x = create_pre_transpose(node);
1019 pre_trans_x->a(input_x);
1020 node->x(pre_trans_x);
1022 auto pre_trans_y = create_pre_transpose(node);
1023 pre_trans_y->a(input_y);
1024 node->y(pre_trans_y);
1031 // Do shape inference for this node again.
1032 node->shape_status(luci::ShapeStatus::UNDEFINED);
1034 auto post_trans = create_post_transpose(node);
1035 loco::replace(node).with(post_trans);
1037 post_trans->a(node);
1041 bool visit(luci::CircleNeg *node) { return convert_unary_x<luci::CircleNeg>(node); }
1043 bool visit(luci::CirclePad *node)
1048 const auto pred_node = loco::must_cast<luci::CircleNode *>(node->input());
1049 auto pre_trans = create_pre_transpose(node);
1050 pre_trans->a(pred_node);
1051 node->input(pre_trans);
1053 auto nchw_paddings = loco::must_cast<luci::CircleConst *>(node->paddings());
1054 const auto nhwc_paddings = create_NHWC_paddings(nchw_paddings);
1055 node->paddings(nhwc_paddings);
1057 // Do shape inference for this node again.
1058 node->shape_status(luci::ShapeStatus::UNDEFINED);
1060 auto post_trans = create_post_transpose(node);
1061 loco::replace(node).with(post_trans);
1063 post_trans->a(node);
1068 bool visit(luci::CirclePadV2 *node)
1073 const auto pred_node = loco::must_cast<luci::CircleNode *>(node->input());
1074 auto pre_trans = create_pre_transpose(node);
1075 pre_trans->a(pred_node);
1076 node->input(pre_trans);
1078 auto nchw_paddings = loco::must_cast<luci::CircleConst *>(node->paddings());
1079 const auto nhwc_paddings = create_NHWC_paddings(nchw_paddings);
1080 node->paddings(nhwc_paddings);
1082 // Do shape inference for this node again.
1083 node->shape_status(luci::ShapeStatus::UNDEFINED);
1085 auto post_trans = create_post_transpose(node);
1086 loco::replace(node).with(post_trans);
1088 post_trans->a(node);
1093 // TODO Reduce duplicate code with CircleMean
1094 bool visit(luci::CircleReduceMax *node)
1096 auto input = loco::must_cast<luci::CircleNode *>(node->input());
1097 if (input->rank() != 4)
1100 auto rindices = dynamic_cast<luci::CircleConst *>(node->reduction_indices());
1104 auto nhwc_rindices = create_NHWC_rindices(rindices);
1105 if (not nhwc_rindices)
1108 auto pre_trans = create_pre_transpose(node);
1109 pre_trans->a(input);
1110 node->input(pre_trans);
1112 // Do shape inference for this node again.
1113 node->shape_status(luci::ShapeStatus::UNDEFINED);
1115 node->reduction_indices(nhwc_rindices);
1117 if (node->keep_dims())
1119 auto post_trans = create_post_transpose(node);
1120 loco::replace(node).with(post_trans);
1122 post_trans->a(node);
1127 // The below codes handle the cases where node->keep_dims() == false
1128 // 1D output never needs a transpose
1129 if (node->rank() <= 1)
1132 std::vector<bool> reduced_dims_nhwc(4, false);
1133 uint32_t num_reduced_indices = nhwc_rindices->size<loco::DataType::S32>();
1135 for (uint32_t ri = 0; ri < num_reduced_indices; ++ri)
1137 reduced_dims_nhwc[nhwc_rindices->at<loco::DataType::S32>(ri)] = true;
1140 // if channel dimension has been reduced, we don't need a transpose
1141 if (reduced_dims_nhwc[3])
1144 // likewise, if both space dimensions are reduced, no transpose is needed
1145 if (reduced_dims_nhwc[1] && reduced_dims_nhwc[2])
1148 std::vector<int32_t> post_trans_ind;
1149 // case 1: only N is reduced
1150 if (num_reduced_indices == 1 && reduced_dims_nhwc[0])
1151 post_trans_ind = {2, 0, 1};
1153 // case 2: only H or W is reduced
1154 if (num_reduced_indices == 1 && (reduced_dims_nhwc[1] || reduced_dims_nhwc[2]))
1155 post_trans_ind = {0, 2, 1};
1157 // case 3: N and either H or W are reduced
1158 if (num_reduced_indices == 2)
1159 post_trans_ind = {1, 0};
1161 auto post_trans = create_Nd_transpose(node, post_trans_ind);
1162 loco::replace(node).with(post_trans);
1164 post_trans->a(node);
1169 // TODO Reduce duplicate codes with CircleReduceMax
1170 bool visit(luci::CircleReduceMin *node)
1172 auto input = loco::must_cast<luci::CircleNode *>(node->input());
1173 if (input->rank() != 4)
1176 auto rindices = dynamic_cast<luci::CircleConst *>(node->reduction_indices());
1180 auto nhwc_rindices = create_NHWC_rindices(rindices);
1181 if (not nhwc_rindices)
1184 auto pre_trans = create_pre_transpose(node);
1185 pre_trans->a(input);
1186 node->input(pre_trans);
1188 // Do shape inference for this node again.
1189 node->shape_status(luci::ShapeStatus::UNDEFINED);
1191 node->reduction_indices(nhwc_rindices);
1193 if (node->keep_dims())
1195 auto post_trans = create_post_transpose(node);
1196 loco::replace(node).with(post_trans);
1198 post_trans->a(node);
1203 // The below codes handle the cases where node->keep_dims() == false
1204 // 1D output never needs a transpose
1205 if (node->rank() <= 1)
1208 std::vector<bool> reduced_dims_nhwc(4, false);
1209 uint32_t num_reduced_indices = nhwc_rindices->size<loco::DataType::S32>();
1211 for (uint32_t ri = 0; ri < num_reduced_indices; ++ri)
1213 reduced_dims_nhwc[nhwc_rindices->at<loco::DataType::S32>(ri)] = true;
1216 // if channel dimension has been reduced, we don't need a transpose
1217 if (reduced_dims_nhwc[3])
1220 // likewise, if both space dimensions are reduced, no transpose is needed
1221 if (reduced_dims_nhwc[1] && reduced_dims_nhwc[2])
1224 std::vector<int32_t> post_trans_ind;
1225 // case 1: only N is reduced
1226 if (num_reduced_indices == 1 && reduced_dims_nhwc[0])
1227 post_trans_ind = {2, 0, 1};
1229 // case 2: only H or W is reduced
1230 if (num_reduced_indices == 1 && (reduced_dims_nhwc[1] || reduced_dims_nhwc[2]))
1231 post_trans_ind = {0, 2, 1};
1233 // case 3: N and either H or W are reduced
1234 if (num_reduced_indices == 2)
1235 post_trans_ind = {1, 0};
1237 auto post_trans = create_Nd_transpose(node, post_trans_ind);
1238 loco::replace(node).with(post_trans);
1240 post_trans->a(node);
1245 bool visit(luci::CircleRelu *node) { return convert_unary_features<luci::CircleRelu>(node); }
1247 bool visit(luci::CircleRelu6 *node) { return convert_unary_features<luci::CircleRelu6>(node); }
1249 bool visit(luci::CircleRsqrt *node) { return convert_unary_x<luci::CircleRsqrt>(node); }
1251 bool visit(luci::CircleSplitV *node)
1253 // Change split dimension
1254 auto axis = dynamic_cast<luci::CircleConst *>(node->split_dim());
1258 if (axis->dtype() != loco::DataType::S32)
1261 if (axis->size<loco::DataType::S32>() != 1)
1264 axis->at<loco::DataType::S32>(0) = nchw_axis_to_nhwc(axis->at<loco::DataType::S32>(0));
1266 // Insert pre-transpose
1267 const auto pred_node = loco::must_cast<luci::CircleNode *>(node->input());
1268 auto pre_trans = create_pre_transpose(node);
1269 pre_trans->a(pred_node);
1270 node->input(pre_trans);
1272 // Do shape inference for this node again.
1273 node->shape_status(luci::ShapeStatus::UNDEFINED);
1275 // Insert post-transposes
1276 for (auto succ : loco::succs(node))
1278 auto svo = loco::must_cast<luci::CircleSplitVOut *>(succ);
1280 auto post_trans = create_post_transpose(svo);
1281 loco::replace(svo).with(post_trans);
1288 bool visit(luci::CircleSquaredDifference *node)
1290 // TODO support CircleConst input
1291 if (dynamic_cast<luci::CircleConst *>(node->x()) != nullptr)
1293 if (dynamic_cast<luci::CircleConst *>(node->y()) != nullptr)
1296 auto input_x = loco::must_cast<luci::CircleNode *>(node->x());
1297 if (input_x->rank() != 4)
1299 auto input_y = loco::must_cast<luci::CircleNode *>(node->y());
1300 if (input_y->rank() != 4)
1303 auto pre_trans_x = create_pre_transpose(node);
1304 pre_trans_x->a(input_x);
1305 node->x(pre_trans_x);
1307 auto pre_trans_y = create_pre_transpose(node);
1308 pre_trans_y->a(input_y);
1309 node->y(pre_trans_y);
1311 // Do shape inference for this node again.
1312 node->shape_status(luci::ShapeStatus::UNDEFINED);
1314 auto post_trans = create_post_transpose(node);
1315 loco::replace(node).with(post_trans);
1317 post_trans->a(node);
1321 bool visit(luci::CircleSub *node)
1323 luci::CircleNode *pred_node = nullptr;
1324 luci::CircleConst *subtract = nullptr;
1326 auto const_x = dynamic_cast<luci::CircleConst *>(node->x());
1327 auto const_y = dynamic_cast<luci::CircleConst *>(node->y());
1329 if (const_x != nullptr && const_y == nullptr)
1331 // case of subtract - pred_node
1332 pred_node = loco::must_cast<luci::CircleNode *>(node->y());
1335 if (!is_NCHW_with_const(node, pred_node, subtract))
1338 auto pre_trans = create_pre_transpose(node);
1339 pre_trans->a(pred_node);
1341 if (subtract->rank() == 4)
1343 auto nhwc_const = create_NHWC_from_NCHW(subtract);
1344 if (nhwc_const == nullptr)
1346 node->x(nhwc_const);
1350 else if (const_x == nullptr && const_y != nullptr)
1352 // case of pred_node - subtract
1353 pred_node = loco::must_cast<luci::CircleNode *>(node->x());
1356 if (!is_NCHW_with_const(node, pred_node, subtract))
1359 auto pre_trans = create_pre_transpose(node);
1360 pre_trans->a(pred_node);
1362 if (subtract->rank() == 4)
1364 auto nhwc_const = create_NHWC_from_NCHW(subtract);
1365 if (nhwc_const == nullptr)
1367 node->y(nhwc_const);
1372 else if (const_x == nullptr && const_y == nullptr)
1374 // Both inputs are not constant.
1375 // In this case, we cannot distinguish NCHW from NHWC,
1376 // so just insert Transpose Ops.
1377 // Only support for input rank 4.
1378 auto input_x = loco::must_cast<luci::CircleNode *>(node->x());
1379 if (input_x->rank() != 4)
1381 auto input_y = loco::must_cast<luci::CircleNode *>(node->y());
1382 if (input_y->rank() != 4)
1385 auto pre_trans_x = create_pre_transpose(node);
1386 pre_trans_x->a(input_x);
1387 node->x(pre_trans_x);
1389 auto pre_trans_y = create_pre_transpose(node);
1390 pre_trans_y->a(input_y);
1391 node->y(pre_trans_y);
1394 // Do shape inference for this node again.
1395 node->shape_status(luci::ShapeStatus::UNDEFINED);
1397 auto post_trans = create_post_transpose(node);
1398 loco::replace(node).with(post_trans);
1400 post_trans->a(node);
1410 bool ConvertNCHWToNHWCPass::run(loco::Graph *g)
1413 INFO(l) << "ConvertNCHWToNHWCPass Start" << std::endl;
1415 // Annotate NHWC operators
1416 // NHWC operators are detected by pattern matching
1419 // pre-Transose (or pre-Reshape) + [intermediate Ops] + post-Transpose (or post-Reshape)
1421 // [intermediate Ops] are annotated as NHWC
1423 // NOTE A single pre-Transpose/Reshape can have multiple post-Transpose/Reshape.
1425 // pre-Transpose --- [intermediate Ops] --- post-Transpose
1427 // +--[intermediate Ops] --- post-Transpose
1429 // NOTE Intermediate Ops SHOULD NOT contain pre-Transpose/Reshape
1430 for (auto node : loco::postorder_traversal(loco::output_nodes(g)))
1432 if (has_data_format(node))
1435 if (is_pre_transpose(node) || is_pre_reshape(node))
1437 std::set<loco::Node *> intermediate;
1439 // Variable to check intermediate Ops contain pre-Transpose/Reshape
1440 bool has_pre = false;
1442 // Variable to check the pattern is closed with post-Transpose/Reshape
1443 bool is_closed = true;
1445 // For recursive call of lambda
1446 std::function<void(loco::Node *)> collect_intermediate;
1447 collect_intermediate = [&](loco::Node *n) {
1448 for (auto succ : loco::succs(n))
1450 // Skip unnecessary traversal
1451 if (intermediate.find(succ) != intermediate.end())
1455 if (is_post_transpose(succ) || is_post_reshape(succ))
1458 if (is_pre_transpose(succ) || is_pre_reshape(succ))
1464 if (is_output(succ))
1470 intermediate.emplace(succ);
1472 collect_intermediate(succ);
1476 collect_intermediate(node);
1478 if (has_pre or not is_closed)
1481 for (auto inter : intermediate)
1483 if (not has_data_format(inter))
1484 set_data_format(inter, DataFormat::NHWC);
1489 // Annotate NCHW operators
1490 for (auto node : loco::active_nodes(loco::output_nodes(g)))
1492 auto circle_node = loco::must_cast<luci::CircleNode *>(node);
1493 switch (circle_node->opcode())
1495 // List of supported Ops
1496 case luci::CircleOpcode::CIRCLEINPUT:
1497 if (!_preserve_input && !has_data_format(node))
1499 set_data_format(node, DataFormat::NCHW);
1502 case luci::CircleOpcode::CIRCLEOUTPUT:
1503 if (!_preserve_output && !has_data_format(node))
1505 set_data_format(node, DataFormat::NCHW);
1508 // SOFTMAX, LOG_SOFTMAX are not converted, because
1509 // tflite/circle assumes the last channel is always axis
1510 case luci::CircleOpcode::ADD:
1511 case luci::CircleOpcode::CONCATENATION:
1512 case luci::CircleOpcode::ELU:
1513 case luci::CircleOpcode::LEAKY_RELU:
1514 case luci::CircleOpcode::LOGISTIC:
1515 case luci::CircleOpcode::MAXIMUM:
1516 case luci::CircleOpcode::MEAN:
1517 case luci::CircleOpcode::MINIMUM:
1518 case luci::CircleOpcode::MUL:
1519 case luci::CircleOpcode::NEG:
1520 case luci::CircleOpcode::PAD:
1521 case luci::CircleOpcode::PADV2:
1522 case luci::CircleOpcode::REDUCE_MAX:
1523 case luci::CircleOpcode::REDUCE_MIN:
1524 case luci::CircleOpcode::RELU:
1525 case luci::CircleOpcode::RELU6:
1526 case luci::CircleOpcode::RSQRT:
1527 case luci::CircleOpcode::SPLIT_V:
1528 case luci::CircleOpcode::SQUARED_DIFFERENCE:
1529 case luci::CircleOpcode::SUB:
1530 if (!has_data_format(node))
1532 set_data_format(node, DataFormat::NCHW);
1540 bool changed = false;
1541 for (auto node : loco::active_nodes(loco::output_nodes(g)))
1543 if (!has_data_format(node))
1548 else if (get_data_format(node) == DataFormat::NHWC)
1550 // Already converted to NHWC
1553 else if (has_dynamic_shape(node))
1555 // This pass only works for static-shaped node
1556 INFO(l) << "Skip the node with a dynamic shape." << std::endl;
1561 ConvertNCHWToNHWC converter;
1562 auto circle_node = loco::must_cast<luci::CircleNode *>(node);
1563 if (circle_node->rank() != 4)
1565 // TODO replace the check above with the input rank check, and remove the condition below
1566 if (not dynamic_cast<luci::CircleMean *>(node) and
1567 not dynamic_cast<luci::CircleReduceMax *>(node) and
1568 not dynamic_cast<luci::CircleReduceMin *>(node))
1572 if (circle_node->accept(&converter))
1574 set_data_format(node, DataFormat::NHWC);
1584 INFO(l) << "ConvertNCHWToNHWCPass End" << std::endl;