fix bugs

Author: icemelon

python/tvm/autotvm/graph_tuner/utils/traverse_graph.py

@@ -83,7 +83,7 @@ def _traverse_expr(node):
             return
         node_index = len(node_list)
         node_entry = {"node": node, "inputs": [], "types": [],
-                      "op": "null", "name": None}
+                      "op": None, "name": None}
 
         if isinstance(node, Call):
             op = node.op

python/tvm/autotvm/graph_tuner/utils/utils.py

@@ -46,7 +46,7 @@ def has_multiple_inputs(node_list, node_idx, input_names):
         in_idx = in_idx[0]
         in_node = node_list[in_idx]
         # Exclude parameter nodes
-        if in_node["op"] != "null" or \
+        if in_node["op"] is not None or \
                 ("name" in in_node and in_node["name"] in input_names):
             num_inputs += 1
     return num_inputs > 1
@@ -71,9 +71,10 @@ def is_boundary_node(node_entry, input_names):
         whether node is a boundary node.
     """
     # Operators dependent on original layouts.
-    _LAYOUT_FIXED_OP = ["batch_flatten", "transpose", "reshape",
-                        "multibox_prior", "multibox_transform_loc", "where",
-                        "non_max_suppression", "strided_slice"]
+    _LAYOUT_FIXED_OP = [relay.op.get(name) for name in (
+        "nn.batch_flatten", "transpose", "reshape", "vision.multibox_prior",
+        "vision.multibox_transform_loc", "where", "vision.non_max_suppression",
+        "strided_slice")]
 
     out = node_entry["op"] in _LAYOUT_FIXED_OP or \
           ("name" in node_entry and node_entry["name"] in input_names)
@@ -94,9 +95,7 @@ def is_skipped_node(node_entry):
         whether node is skipped.
     """
     # Operators not counted in graph tuner.
-    _SKIPPED_OP = ["Tuple"]
-
-    return node_entry["op"] in _SKIPPED_OP
+    return isinstance(node_entry["node"], relay.Tuple)
 
 
 def bind_inputs(expr, input_shapes=None, input_dtypes="float32"):

python/tvm/autotvm/record.py

@@ -161,12 +161,12 @@ def clean_json_to_python(x):
         tgt = _target.create(items[0])
         task_tuple = pickle.loads(base64.b64decode(items[1].encode()))
         config = pickle.loads(base64.b64decode(items[2].encode()))
-        result = pickle.loads(base64.b64decode(items[3].encode()))
+        result = MeasureResult(*pickle.loads(base64.b64decode(items[3].encode())))
         config.cost = np.mean(result.costs)
 
         tsk = task.Task(task_tuple[0], task_tuple[1])
         tsk.workload = task_tuple[3]
-        return MeasureInput(tgt, tsk, config), MeasureResult(*result)
+        return MeasureInput(tgt, tsk, config), result
 
     raise RuntimeError("Invalid log protocol: " + protocol)
 

python/tvm/autotvm/task/dispatcher.py

@@ -33,9 +33,6 @@
 import logging
 
 import numpy as np
-from decorator import decorate
-
-from tvm import target as _target
 
 from .space import FallbackConfigEntity
 

python/tvm/autotvm/task/task.py

@@ -418,7 +418,8 @@ def get_config():
     cfg: ConfigSpace or ConfigEntity
         The current config
     """
-    return DispatchContext.current.query(None, None)
+    tgt = _target.current_target(allow_none=True)
+    return DispatchContext.current.query(tgt, None)
 
 class FlopCalculationError(RuntimeError):
     """Error happens when estimating FLOP for a compute op"""

python/tvm/relay/backend/compile_engine.py

@@ -256,7 +256,7 @@ def create_tensors(typ, tensors):
         self.func_name = "fused"
         outputs = self.visit(prim_func.body)
         if len(self.func_name) > ScheduleGetter.MAX_FUNC_NAME_LENGTH:
-            hash_digest = int(hashlib.sha1(self.func_name).hexdigest(), 16)
+            hash_digest = int(hashlib.sha1(self.func_name.encode("utf-8")).hexdigest(), 16)
             self.func_name = "%s_%s" % (
                 self.func_name[:ScheduleGetter.MAX_FUNC_NAME_LENGTH], hash_digest)
 
@@ -270,7 +270,8 @@ def create_tensors(typ, tensors):
             # print('master op:', self.master_op.name)
             sch = self.master_implement.schedule(self.master_attrs, tensor_outs, self.target)
             for scalar in self.scalars:
-                sch[scalar].compute_inline()
+                if scalar in sch.stage_map:
+                    sch[scalar].compute_inline()
         return CachedFunc(self.target, self.func_name, inputs, outputs, sch)
 
     def visit_var(self, var):
@@ -381,10 +382,10 @@ def visit_tuple(self, tup):
         return fields
 
     def visit_tuple_getitem(self, t):
-        tup = self.visit(t.tuple)
-        assert len(tup) == len(t.tuple.checked_type.fields)
+        tup = self.visit(t.tuple_value)
+        assert len(tup) == len(t.tuple_value.checked_type.fields)
         assert t.index >= 0
-        assert t.index < tup.size()
+        assert t.index < len(tup)
         return [tup[t.index]]
 
 

python/tvm/relay/frontend/tensorflow.py

@@ -310,6 +310,7 @@ def _impl(inputs, attr, params):
             flip_layout = True
 
         if attr['data_format'] == 'NHWC':
+            in_channels = input_shape[3]
             kernel_h, kernel_w, _, depth_mult = weights_shape
             attr['kernel_shape'] = (weights_shape[0], weights_shape[1])
             if opname == 'conv':
@@ -323,6 +324,7 @@ def _impl(inputs, attr, params):
                 attr['dilations'] = (attr['dilations'][1], attr['dilations'][2])
             attr['strides'] = (attr['strides'][1], attr['strides'][2])
         elif attr['data_format'] == 'NCHW':
+            in_channels = input_shape[1]
             _, depth_mult, kernel_h, kernel_w = weights_shape
             attr['kernel_shape'] = (weights_shape[2], weights_shape[3])
             if opname == 'conv':
@@ -343,7 +345,7 @@ def _impl(inputs, attr, params):
             raise tvm.error.OpAttributeInvalid(msg.format(attr['data_format']))
 
         if opname == 'depthwise':
-            attr['groups'] = attr['channels']
+            attr['groups'] = in_channels
 
         # Fix padding
         attr['padding'] = attr['padding'].decode("utf-8")

python/tvm/relay/op/nn/_nn.py

@@ -307,7 +307,7 @@ def compute_upsampling3d(attrs, inputs, out_dtype):
 
 # mirror_pad
 @reg.register_compute("nn.mirror_pad")
-def compute_mirror_pad(attrs, inputs, out_dtype, target):
+def compute_mirror_pad(attrs, inputs, out_dtype):
     pad_before, pad_after = list(zip(*attrs.pad_width))
     mode = attrs.mode
     out = topi.nn.mirror_pad(inputs[0], pad_before=pad_before, pad_after=pad_after, mode=mode)

python/tvm/relay/op/strategy/x86.py

@@ -116,7 +116,8 @@ def conv2d_strategy_cpu(attrs, inputs, out_type, target):
                     wrap_topi_schedule(topi.generic.schedule_depthwise_conv2d_nchw))
         elif layout == "NHWC":
             assert kernel_layout == "HWOI"
-            logger.warning("For x86 target, NCHW layout is recommended for depthwise_conv2d.")
+            logger.warning("For x86 target, depthwise_conv2d with NCHW layout is "
+                           "not optimized.")
             strategy.add_implement(
                 wrap_compute_conv2d(topi.nn.depthwise_conv2d_nhwc),
                 wrap_topi_schedule(topi.generic.schedule_depthwise_conv2d_nhwc))

src/relay/op/nn/convolution.h

@@ -153,6 +153,16 @@ bool Conv2DRel(const Array<Type>& types, int num_inputs, const Attrs& attrs,
       << " But got " << out_layout;
 
   Array<IndexExpr> dshape_nchw = trans_in_layout.ForwardShape(data->shape);
+  bool is_depthwise = false;
+  if (param->groups > 1) {
+    CHECK(weight->shape.defined()) << "Weight shape must be specified " <<
+        "when groups is greater than 1.";
+    Array<IndexExpr> wshape_oihw = trans_kernel_layout.ForwardShape(weight->shape);
+    if (tvm::tir::Equal(param->groups, dshape_nchw[1]) &&
+        tvm::tir::Equal(param->groups, wshape_oihw[0])) {
+      is_depthwise = true;
+    }
+  }
 
   IndexExpr channels, dilated_ksize_y, dilated_ksize_x;
   // infer weight if the kernel_size and channels are defined
@@ -161,9 +171,9 @@ bool Conv2DRel(const Array<Type>& types, int num_inputs, const Attrs& attrs,
     CHECK_EQ(param->dilation.size(), 2);
     Array<IndexExpr> wshape;
 
-    if (tvm::tir::Equal(param->channels, param->groups) && !tvm::tir::Equal(param->channels, 1)) {
+    if (is_depthwise) {
       // infer weight's shape for depthwise convolution
-      wshape = {{dshape_nchw[1], indexdiv(param->groups, dshape_nchw[1]), param->kernel_size[0],
+      wshape = {{dshape_nchw[1], indexdiv(param->channels, dshape_nchw[1]), param->kernel_size[0],
                  param->kernel_size[1]}};
     } else {
       wshape = {{param->channels, indexdiv(dshape_nchw[1], param->groups), param->kernel_size[0],