[Relay][Layout] Add FInferCorrectLayout for L2 norm layout transform. (#12497)

* [Relay][Layout] FInferCorrectLayout for L2 norm layout change.

* [Relay][Layout] Test for L2 norm layout transform.

* [Relay][Layout] Re-edit test to add multi-dimensional axis list.

* Fix cpplint errors

* Use clang-format-10 rules.

* replace uint with size_t.

Author: blackkker

src/relay/op/nn/nn.cc

@@ -653,6 +653,42 @@ Expr MakeL2Normalize(Expr data, double eps, Array<Integer> axis) {
   return Call(op, {data}, Attrs(attrs), {});
 }
 
+InferCorrectLayoutOutput L2NormalizeInferCorrectLayout(
+    const Attrs& attrs, const Array<Layout>& new_in_layouts, const Array<Layout>& old_in_layouts,
+    const Array<tvm::relay::Type>& old_in_types) {
+  const auto* attrs_ptr = attrs.as<L2NormalizeAttrs>();
+  ICHECK(attrs_ptr);
+  ObjectPtr<L2NormalizeAttrs> param = make_object<L2NormalizeAttrs>(*attrs_ptr);
+
+  Array<Array<IndexExpr>> old_in_shapes;
+  for (auto old_in_t : old_in_types) {
+    ICHECK(old_in_t.as<TensorTypeNode>());
+    old_in_shapes.push_back(old_in_t.as<TensorTypeNode>()->shape);
+  }
+  std::vector<size_t> axis_list;
+  for (auto i : param->axis) {
+    int64_t axis = i->value;
+    if (axis < 0) {
+      axis = axis + static_cast<size_t>(old_in_shapes[0].size());
+    }
+    axis_list.emplace_back(axis);
+  }
+
+  Layout ret = Layout::Undef();
+  if (new_in_layouts.defined() && old_in_layouts.defined()) {
+    for (size_t i = 0; i < axis_list.size(); ++i) {
+      const auto& axis_dim = old_in_layouts[0][axis_list[i]];
+      auto axis_index = new_in_layouts[0].IndexOf(axis_dim);
+      param->axis.Set(i, axis_index);
+    }
+    ret = new_in_layouts[0];
+  } else if (old_in_layouts.defined()) {
+    ret = old_in_layouts[0];
+  }
+
+  return InferCorrectLayoutOutput({ret}, {ret}, Attrs(param));
+}
+
 TVM_REGISTER_GLOBAL("relay.op.nn._make.l2_normalize").set_body_typed(MakeL2Normalize);
 
 RELAY_REGISTER_OP("nn.l2_normalize")
@@ -669,7 +705,7 @@ Normalizes along dimension axis using an L2 norm
     .set_num_inputs(1)
     .add_argument("data", "Tensor", "The input tensor.")
     .set_support_level(2)
-    .set_attr<FInferCorrectLayout>("FInferCorrectLayout", ElemwiseArbitraryLayout)
+    .set_attr<FInferCorrectLayout>("FInferCorrectLayout", L2NormalizeInferCorrectLayout)
     .add_type_rel("Identity", IdentityRel);
 
 // Dropout

tests/python/relay/test_pass_convert_op_layout.py

@@ -2697,5 +2697,44 @@ def expected():
     assert tvm.ir.structural_equal(a, b), "Actual = \n" + str(a) + "\n Expect = \n" + str(b)
 
 
+def test_conv_l2n_convert_layout():
+    """Check that layout transforms are propagated through bn."""
+    axis_list = ([3], [-1], [2, 3])
+    expected_axis = ([1], [1], [3, 1])
+    for i, axis in enumerate(axis_list):
+
+        def before():
+            x = relay.var("x", shape=(1, 56, 56, 64))
+            weight = relay.var("weight", shape=(3, 3, 64, 64))
+            y = relay.nn.conv2d(
+                x,
+                weight,
+                channels=64,
+                kernel_size=(3, 3),
+                padding=(1, 1),
+                data_layout="NHWC",
+                kernel_layout="HWIO",
+            )
+            z = relay.nn.l2_normalize(y, eps=0.001, axis=axis)
+            z = relay.Function(analysis.free_vars(z), z)
+            return z
+
+        def expected():
+            x = relay.var("x", shape=(1, 56, 56, 64))
+            w = relay.var("weight", shape=(3, 3, 64, 64))
+            x = relay.layout_transform(x, "NHWC", "NCHW")
+            w = relay.layout_transform(w, "HWIO", "OIHW")
+            y = relay.nn.conv2d(x, w, channels=64, kernel_size=(3, 3), padding=(1, 1))
+            z = relay.nn.l2_normalize(y, eps=0.001, axis=expected_axis[i])
+            z = relay.layout_transform(z, "NCHW", "NHWC")
+            z = relay.Function(analysis.free_vars(z), z)
+            return z
+
+    a = before()
+    a = run_opt_pass(a, transform.ConvertLayout({"nn.conv2d": ["NCHW", "default"]}))
+    b = run_opt_pass(expected(), transform.InferType())
+    assert tvm.ir.structural_equal(a, b), "Actual = \n" + str(a) + "\n\n Expected = \n" + str(b)
+
+
 if __name__ == "__main__":
     pytest.main([__file__])