[Relay][AlterOp] NHWC to NCHWc support for Pool, pad, concatenate, sum.

anijain2305 · anijain2305 · commit 87f012e54c91 · 2019-10-09T00:35:15.000Z
diff --git a/python/tvm/relay/frontend/tflite.py b/python/tvm/relay/frontend/tflite.py
@@ -748,10 +748,12 @@ def convert_conv(self, op, conv_type):
         elif padding == Padding.SAME:
             pad_top, pad_bottom = get_pad_value(input_h, dilated_kernel_h, stride_h)
             pad_left, pad_right = get_pad_value(input_w, dilated_kernel_w, stride_w)
-            in_expr = _op.nn.pad(data=in_expr, pad_width=((0, 0),
-                                                          (pad_top, pad_bottom),
-                                                          (pad_left, pad_right),
-                                                          (0, 0)))
+            do_pad = not (pad_top == 0 and pad_bottom == 0 and pad_left == 0 and pad_right == 0)
+            if do_pad:
+                in_expr = _op.nn.pad(data=in_expr, pad_width=((0, 0),
+                                                              (pad_top, pad_bottom),
+                                                              (pad_left, pad_right),
+                                                              (0, 0)))
         else:
             raise tvm.error.OpAttributeUnImplemented(
                 'Padding format {} is not supported for operator Conv.'.format(padding))
diff --git a/src/relay/op/nn/pad.cc b/src/relay/op/nn/pad.cc
@@ -6,9 +6,9 @@
  * to you 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
@@ -36,6 +36,80 @@ namespace relay {
 // relay.nn.pad
 TVM_REGISTER_NODE_TYPE(PadAttrs);
 
+Array<Array<Layout> > PadInferCorrectLayout(
+    const Attrs& attrs,
+    const Array<Layout>& new_in_layouts,
+    const Array<Layout>& old_in_layouts,
+    const Array<Array<IndexExpr>> &old_in_shapes) {
+  // NOTE: Discard "const" qualifier here.
+  PadAttrs *params = const_cast<PadAttrs*>(attrs.as<PadAttrs>());
+
+  Layout ret;
+  // If new_in_layouts are defined, this code tries to modify the layout.
+  bool is_layout_modified = new_in_layouts.defined();
+  if (new_in_layouts.defined()) {
+    // Create a map of axis to param_width. For the new layout, a new param_width is generated using
+    // the map. The new layout is rejected, if the padding is happening along the axis which was
+    // split.
+
+    // 1) Create a map from axis to param_width using old layout.
+    std::map<std::string, tvm::Array<tvm::Expr>> axis_pad_width;
+    int index_counter = 0;
+    CHECK_EQ(new_in_layouts.size(), 1);
+    for (auto iter_var : old_in_layouts[0]->axes) {
+      const auto& old_layout_axis = LayoutAxis::Get(iter_var);
+      axis_pad_width.emplace(old_layout_axis.name(), params->pad_width[index_counter]);
+      index_counter++;
+    }
+
+    // 2) Create new pad width by walking over the new layout and using the map.
+    tvm::Array<tvm::Array<tvm::Expr>> new_pad_width;
+    for (auto iter_var : new_in_layouts[0]->axes) {
+      const auto& new_layout_axis = LayoutAxis::Get(iter_var);
+      auto axis_name = new_layout_axis.name();
+      if (axis_pad_width.count(axis_name) != 0 && new_layout_axis.IsPrimal()) {
+        // This is primal axis. So, directly use the original pad_width.
+        new_pad_width.push_back(axis_pad_width.at(axis_name));
+      } else {
+        // This is the axis that got split. So, check that pad_width was [0, 0] originally.
+        const auto& dual_axis = new_layout_axis.ToPrimal();
+        auto dual_axis_name = dual_axis.name();
+        CHECK(axis_pad_width.count(dual_axis_name));
+        new_pad_width.push_back(axis_pad_width.at(dual_axis_name));
+
+        // If all pad_width elements are not zero, do not change the layout.
+        for (auto width : axis_pad_width.at(dual_axis_name)) {
+          if (auto* width_imm = width.as<IntImm>()) {
+            if (width_imm->value != 0) {
+              is_layout_modified = false;
+            }
+          } else {
+            is_layout_modified = false;
+          }
+        }
+      }
+    }
+
+    // If the above conditions satisfied, we can set the newly created pad_width and use the new
+    // layout.
+    if (is_layout_modified) {
+      ret = new_in_layouts[0];
+      params->pad_width = new_pad_width;
+    }
+  }
+
+  if (!is_layout_modified) {
+    if (old_in_layouts.defined()) {
+      CHECK_EQ(old_in_layouts.size(), 1);
+      ret = old_in_layouts[0];
+    } else {
+      ret = Layout::Undef();
+    }
+  }
+
+  return Array<Array<Layout> >{{ret}, {ret}};
+}
+
 bool PadRel(const Array<Type>& types,
             int num_inputs,
             const Attrs& attrs,
@@ -133,6 +207,7 @@ RELAY_REGISTER_OP("nn.pad")
 .add_argument("data", "Tensor", "The input tensor.")
 .set_support_level(2)
 .add_type_rel("Pad", PadRel)
+.set_attr<FInferCorrectLayout>("FInferCorrectLayout", PadInferCorrectLayout)
 .set_attr<TOpPattern>("TOpPattern", kInjective)
 .set_attr<FTVMCompute>("FTVMCompute", PadCompute);
 
diff --git a/src/relay/op/nn/pooling.cc b/src/relay/op/nn/pooling.cc
@@ -47,15 +47,9 @@ Array<Array<Layout> > Pool2DInferCorrectLayout(
   T *params = const_cast<T*>(attrs.as<T>());
 
   if (new_in_layouts.defined()) {
+    // Set the pool with the new layout.
     CHECK_EQ(new_in_layouts.size(), 1);
-
-    Layout raw_layout(params->layout);
-    Layout input = new_in_layouts[0];
-    if (input.IndexOf(LayoutAxis::Get('W')) == raw_layout.IndexOf(LayoutAxis::Get('W')) &&
-    input.IndexOf(LayoutAxis::Get('H')) == raw_layout.IndexOf(LayoutAxis::Get('H')) &&
-        !input.Contains(LayoutAxis::Get('w')) && !input.Contains(LayoutAxis::Get('h'))) {
-      params->layout = input.name();  // modify self to follow the input layout
-    }
+    params->layout = new_in_layouts[0].name();
   }
 
   Layout inferred_layout(params->layout);
diff --git a/src/relay/op/tensor/reduce.cc b/src/relay/op/tensor/reduce.cc
@@ -6,9 +6,9 @@
  * to you 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
@@ -119,6 +119,59 @@ Array<Integer> GetExcludeAxes(size_t indim,
   return r_axes;
 }
 
+// Return the modified layout for AlterOpLayout pass.
+Array<Array<Layout>> ReduceInferCorrectLayout(const Attrs& attrs,
+                                              const Array<Layout>& new_in_layouts,
+                                              const Array<Layout>& old_in_layouts,
+                                              const Array<Array<IndexExpr>>& old_in_shapes) {
+  // NOTE: Discard "const" qualifier here.
+  ReduceAttrs* params = const_cast<ReduceAttrs*>(attrs.as<ReduceAttrs>());
+
+  // Get the reduce axes.
+  uint32_t indim = old_in_shapes[0].size();
+  auto r_axes = GetReduceAxes(indim, params->axis, params->exclude);
+
+  Layout ret = Layout::Undef();
+  if (new_in_layouts.defined() && r_axes.size()) {
+    // Adapt to new layout. The axis has to change. Record original reduce axes. Convert to the
+    // modified layout axes.
+    CHECK_EQ(new_in_layouts.size(), 1);
+    CHECK_EQ(old_in_layouts.size(), 1);
+
+    // 1) Collect the original axes
+    std::unordered_set<std::string> old_r_dims;
+    for (auto r_axis : r_axes) {
+      old_r_dims.emplace(old_in_layouts[0][r_axis].name());
+    }
+
+    // 2) Collect the new axes by walking new_layout.
+    tvm::Array<tvm::Integer> new_r_axes;
+    std::string new_layout_string = "";
+    int axis_index = 0;
+    for (auto iter_var : new_in_layouts[0]->axes) {
+      const auto& layout_axis = LayoutAxis::Get(iter_var);
+      const std::string& layout_dim = layout_axis.name();
+      if (old_r_dims.count(layout_dim)) {
+        new_r_axes.push_back(tvm::Integer(axis_index));
+      }
+      // Collect only the primal axis.
+      if (layout_axis.IsPrimal()) {
+        new_layout_string += layout_dim;
+        axis_index++;
+      }
+    }
+
+    // 3) Set the new axis and layout.
+    ret = Layout(new_layout_string);
+    params->axis = new_r_axes;
+  } else if (old_in_layouts.defined()) {
+    // If the new layout is undefined, set the old layout as the inferred layout.
+    CHECK_EQ(old_in_layouts.size(), 1);
+    ret = old_in_layouts[0];
+  }
+
+  return Array<Array<Layout>>{{ret}, {ret}};
+}
 
 template<typename F>
 Array<Tensor> ReduceCompute(const Attrs& attrs,
@@ -325,6 +378,7 @@ Example::
 .set_attrs_type_key("relay.attrs.ReduceAttrs")
 .set_support_level(4)
 .add_type_rel("Reduce", ReduceRel)
+.set_attr<FInferCorrectLayout>("FInferCorrectLayout", ReduceInferCorrectLayout)
 .set_attr<FTVMCompute>("FTVMCompute", SumCompute)
 .set_attr<TOpPattern>("TOpPattern", kCommReduce);
 
diff --git a/src/relay/op/tensor/transform.cc b/src/relay/op/tensor/transform.cc
@@ -283,22 +283,34 @@ Array<Array<Layout>> ConcatenateLayout(
     const Array<Layout>& new_in_layouts,
     const Array<Layout>& old_in_layouts,
     const Array<Array<IndexExpr>> &old_in_shapes) {
-  const ConcatenateAttrs* param = attrs.as<ConcatenateAttrs>();
+  ConcatenateAttrs* param = const_cast<ConcatenateAttrs*>(attrs.as<ConcatenateAttrs>());
 
   size_t axis = param->axis < 0 ? param->axis + old_in_shapes[0].size() :
                 static_cast<size_t>(param->axis);
 
   Layout ret;
+  bool is_new_layout_selected = false;
   if (new_in_layouts.defined()) {  // this function is called after some operators are alternated.
+    // If all the new input layouts are same, the new in layout gets selected.  For axis, the new
+    // axis in the new layout is identified. The param->axis is then modified on the fly to conform
+    // to the new input layout.
     const auto& concate_dim = old_in_layouts[0][axis];
-    for (size_t i = 0; i < new_in_layouts.size(); ++i) {
-      if (new_in_layouts[i].ndim() > axis &&
-          new_in_layouts[i][axis] == concate_dim) {
-        ret = new_in_layouts[i];
-        break;
+    bool all_input_layouts_same = true;
+    for (auto new_layout : new_in_layouts) {
+      if (!new_layout.Equals(new_in_layouts[0])) {
+        all_input_layouts_same = false;
       }
     }
-  } else {  // this function is called on the original correct relay ir
+    if (all_input_layouts_same) {
+      auto new_index = new_in_layouts[0].IndexOf(concate_dim);
+      ret = new_in_layouts[0];
+      param->axis = new_index;
+      is_new_layout_selected = true;
+    }
+  }
+
+  if (!is_new_layout_selected) {
+    // this function is called on the original correct relay ir
     for (size_t i = 0; i < old_in_layouts.size(); ++i) {
       if (old_in_layouts[i].defined()) {
         ret = old_in_layouts[i];
diff --git a/tests/python/relay/test_pass_alter_op_layout.py b/tests/python/relay/test_pass_alter_op_layout.py
