add checkes for query_layout

Author: Ivy

python/tvm/relay/op/contrib/dnnl.py

@@ -211,20 +211,23 @@ def pattern_table():
     return dnnl_patterns
 
 
-def get_optimal_layout_for_conv(weight_shape, out_shape, paddings, strides, dilates, groups):
+def get_optimal_layout_for_conv(data_layout, kernel_layout, weight_shape,
+                                out_shape, paddings, strides, dilates, groups):
     """Get the optimal layout of dnnl, given shape of conv2d.
 
     Parameters
     ----------
-    weight_shape, out_shape, paddings, strides, dilates, groups : Int, String
-                                                                  Input argument.
+    data_layout, kernel_layout,weight_shape, out_shape, paddings, strides, dilates, groups :String
+                                                                                            Input argument.
 
     Returns
     -------
     layouts : string
               The result.
     """
     return _ffi_api.get_optimal_layout_for_conv(
+        data_layout,
+        kernel_layout,
         weight_shape,
         out_shape,
         paddings,
@@ -235,13 +238,13 @@ def get_optimal_layout_for_conv(weight_shape, out_shape, paddings, strides, dila
 
 
 def get_optimal_layout_for_conv_transpose(
-    weight_shape, out_shape, paddings, output_paddings, strides, dilates, groups
+    data_layout, kernel_layout, weight_shape, out_shape, paddings, output_paddings, strides, dilates, groups
 ):
     """Get the optimal layout of dnnl, given shape of tranposed conv2d.
 
     Parameters
     ----------
-    weight_shape, out_shape, paddings, output_paddings, strides, dilates, groups
+    data_layout, kernel_layout, weight_shape, out_shape, paddings, output_paddings, strides, dilates, groups
         : Int, String
           Input argument.
 
@@ -251,6 +254,8 @@ def get_optimal_layout_for_conv_transpose(
               The result.
     """
     return _ffi_api.get_optimal_layout_for_conv_transpose(
+        data_layout,
+        kernel_layout, 
         weight_shape,
         out_shape,
         paddings,
@@ -341,7 +346,8 @@ def alter_conv(attrs, inputs, tinfos, out_type):
     conv_type = type(attrs).__name__.split("Attrs")[0]
 
     res = get_optimal_layout_for_conv(
-        weight_shape, out_shape, paddings, strides, dilates, groups
+        attrs["data_layout"], attrs["kernel_layout"], weight_shape, out_shape, paddings,
+        strides, dilates, groups,
     )
     src_df, weight_df, dst_df = res.split(",")
     new_attrs["data_layout"] = tag2layout(src_df, is_weight=False, conv_type=conv_type)
@@ -373,6 +379,8 @@ def alter_conv_transpose(attrs, inputs, tinfos, out_type):
     conv_type = type(attrs).__name__.split("Attrs")[0]
 
     res = get_optimal_layout_for_conv_transpose(
+        attrs["data_layout"],
+        attrs["kernel_layout"], 
         weight_shape,
         out_shape,
         paddings,

src/relay/backend/contrib/dnnl/query_layout.cc

@@ -172,9 +172,32 @@ dnnl::memory::dims str2dims(const std::string& str_shape,
   return out_dims;
 }
 
-std::string get_optimal_layout_for_conv(std::string weight_shape,
-                                        std::string out_shape, std::string paddings,
-                                        std::string strides, std::string dilates, std::string G) {
+void check_shapes(const std::vector<std::string> shapes) {
+  std::regex valid_pat("(\\d*)(,(\\d*))*");
+  bool checked = std::regex_match(shapes[0], valid_pat);
+  for (size_t i = 1; i < shapes.size()-1; i++) {
+    checked &= std::regex_match(shapes[i], valid_pat);
+  }
+  checked &= std::regex_match(shapes[shapes.size()-1], std::regex("\\d*"));
+  if (!checked) {
+    LOG(FATAL) << "Invalid input args for query dnnl optimal layout.";
+  }
+}
+
+void check_layout(bool var, bool ref) {
+  if (var != ref) {
+    LOG(FATAL) << "Invalid input layout for query dnnl optimal layout.";
+  }
+}
+
+std::string get_optimal_layout_for_conv(std::string data_layout, std::string kernel_layout,
+                                        std::string weight_shape, std::string out_shape,
+                                        std::string paddings, std::string strides,
+                                        std::string dilates, std::string G) {   
+  check_layout(std::regex_match(data_layout, std::regex("NC(D?)(H?)W")), true);
+  check_layout(std::regex_match(kernel_layout, std::regex("(G?)OI(D?)(H?)W")), true);
+  check_shapes({weight_shape, out_shape, paddings, strides, dilates, G});
+
   dnnl::engine eng(dnnl::engine::kind::cpu, 0);
   dnnl::stream s(eng);
   using tag = dnnl::memory::format_tag;
@@ -249,10 +272,15 @@ std::string get_optimal_layout_for_conv(std::string weight_shape,
   return res;
 }
 
-std::string get_optimal_layout_for_conv_transpose(std::string weight_shape,
-                                                  std::string out_shape, std::string paddings,
-                                                  std::string output_paddings, std::string strides,
-                                                  std::string dilates, std::string G) {
+std::string get_optimal_layout_for_conv_transpose(std::string data_layout, std::string kernel_layout,
+                                                  std::string weight_shape, std::string out_shape,
+                                                  std::string paddings, std::string output_paddings,
+                                                  std::string strides, std::string dilates,
+                                                  std::string G) {
+  check_layout(std::regex_match(data_layout, std::regex("NC(D?)(H?)W")), true);
+  check_layout(std::regex_match(kernel_layout, std::regex("(G?)((IO)|(OI))(D?)(H?)W")), true);
+  check_shapes({weight_shape, out_shape, paddings, output_paddings, strides, dilates, G});
+
   dnnl::engine eng(dnnl::engine::kind::cpu, 0);
   dnnl::stream s(eng);
   using tag = dnnl::memory::format_tag;
@@ -335,13 +363,14 @@ std::string get_optimal_layout_for_conv_transpose(std::string weight_shape,
 
 TVM_REGISTER_GLOBAL("relay.ir.get_optimal_layout_for_conv")
     .set_body([](TVMArgs args, TVMRetValue* rv) {
-      *rv = get_optimal_layout_for_conv(args[0], args[1], args[2], args[3], args[4], args[5]);
+      *rv = get_optimal_layout_for_conv(args[0], args[1], args[2], args[3], args[4], args[5],
+                                        args[6], args[7]);
     });
 
 TVM_REGISTER_GLOBAL("relay.ir.get_optimal_layout_for_conv_transpose")
     .set_body([](TVMArgs args, TVMRetValue* rv) {
       *rv = get_optimal_layout_for_conv_transpose(args[0], args[1], args[2], args[3], args[4],
-                                                  args[5], args[6]);
+                                                  args[5], args[6], args[7], args[8]);
     });
 
 }  // namespace contrib