Removing the conv2d legalization for x86. Will send a separate PR.

Author: anijain2305

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

@@ -191,13 +191,6 @@ def alter_op_layout_conv2d(attrs, inputs, tinfos):
     return topi.nn.conv2d_alter_layout(attrs, inputs, tinfos, op)
 
 
-@reg.register_legalize("nn.conv2d")
-def rewrite_conv2d(attrs, inputs, tinfos):
-    """Rewrite conv2d"""
-    from ... import op
-    return topi.nn.conv2d_legalize(attrs, inputs, tinfos, op)
-
-
 reg.register_pattern("nn.conv2d", OpPattern.OUT_ELEMWISE_FUSABLE)
 
 

tests/python/relay/test_op_level2.py

@@ -21,7 +21,6 @@
 from tvm import relay
 from tvm.relay import transform
 from tvm.relay.testing import ctx_list
-import tvm.contrib.graph_runtime as runtime
 import topi.testing
 
 def run_infer_type(expr):
@@ -30,15 +29,6 @@ def run_infer_type(expr):
     entry = mod["main"]
     return entry if isinstance(expr, relay.Function) else entry.body
 
-def run_opt_pass(expr, passes):
-    passes = passes if isinstance(passes, list) else [passes]
-    mod = relay.Module.from_expr(expr)
-    seq = transform.Sequential(passes)
-    with transform.PassContext(opt_level=3):
-        mod = seq(mod)
-    entry = mod["main"]
-    return entry if isinstance(expr, relay.Function) else entry.body
-
 def test_conv2d_infer_type():
     # symbolic in batch dimension
     n, c, h, w = tvm.var("n"), 10, 224, 224
@@ -116,7 +106,7 @@ def run_test_conv2d(dtype, out_dtype, scale, dshape, kshape,
                         **attrs):
         if except_targets is None:
           except_targets = []
-
+          
         x = relay.var("x", shape=dshape, dtype=dtype)
         w = relay.var("w", dtype=dtype)
         y = relay.nn.conv2d(x, w,
@@ -580,12 +570,13 @@ def _compile(input_dtype, weight_dtype, output_dtype, target):
         # Check that intrinisic is present in the assembly.
         assert "pmaddubs" in asm
 
-        # Ensure that code is generated for i8 x i8 conv.
+        # Ensure that code is generated when datatypes are not HW supported.
         asm = _compile(input_dtype="int8",
                        weight_dtype="int8",
                        output_dtype="int32",
                        target=target)
-        assert "pmaddubs" in asm
+        # Check that intrinisic is not present in the assembly.
+        assert "pmaddubs" not in asm
 
         # Ensure that code is generated when datatypes are not HW supported.
         asm = _compile(input_dtype="uint8",
@@ -605,47 +596,6 @@ def _compile(input_dtype, weight_dtype, output_dtype, target):
     # Check that vector int mult and add instructions are generated.
     assert "vpmulld" in asm and "vpadd" in asm
 
-def test_rewrite_conv2d_intel_int8():
-    def verify(orig, rewritten, input_shape, weight_shape):
-        data = np.random.random_integers(-10, 10,
-                                         size=input_shape).astype('int8')
-        weight = np.random.random_integers(-10, 10,
-                                           size=weight_shape).astype('int8')
-        def _get_output(func):
-            params = {"w": weight}
-            with relay.build_config(opt_level=0):
-                graph, lib, params = relay.build(func, 'llvm', params=params)
-
-            ctx = tvm.cpu(0)
-            module = runtime.create(graph, lib, ctx)
-            module.set_input('data', data)
-            module.set_input(**params)
-            module.run()
-            return module.get_output(0).asnumpy()
-        orig_output = _get_output(orig)
-        rewritten_output = _get_output(rewritten)
-        np.testing.assert_equal(orig_output, rewritten_output)
-
-    input_shape = (1, 128, 28, 28)
-    weight_shape = (256, 128, 3, 3)
-    idtype = "int8"
-    wdtype = "int8"
-    odtype = "int32"
-
-    var_input = relay.var("data", shape=input_shape, dtype=idtype)
-    var_weight = relay.var("w", shape=weight_shape, dtype=wdtype)
-
-    f = relay.nn.conv2d(var_input,
-                        var_weight,
-                        kernel_size=(3, 3),
-                        channels=256,
-                        out_dtype=odtype)
-
-    orig = relay.Function([var_input, var_weight], f)
-    with tvm.target.create("llvm"):
-        rewritten = run_opt_pass(orig, transform.InferType())
-        rewritten = run_opt_pass(rewritten, transform.RewriteOp())
-    verify(orig, rewritten, input_shape, weight_shape)
 
 if __name__ == "__main__":
     test_pool2d()
@@ -663,4 +613,3 @@ def _get_output(func):
     test_batch_flatten()
     test_upsampling()
     test_conv2d_int8_intrinsics()
-    test_rewrite_conv2d_intel_int8()

topi/python/topi/nn/conv2d.py

@@ -94,29 +94,6 @@ def conv2d_alter_layout(attrs, inputs, tinfos, F):
     # not to change by default
     return None
 
-@tvm.target.generic_func
-def conv2d_rewrite_op(attrs, inputs, tinfos, F):
-    """Rewrite Conv2D op.
-
-    Parameters
-    ----------
-    attrs : nnvm.top.AttrDict or tvm.attrs.Attrs
-        Attributes of current convolution
-    inputs : nnvm.symbol or tvm.relay.Expr
-        Grouped input symbols
-    tinfos : list
-        Input shape and dtype
-    F: symbol
-        The context, can be either nnvm.sym or relay.op
-
-    Note
-    ----
-    Unlike other TOPI functions, this function operates on both graph level and operator level,
-    so we have to pass 'F' to make it support our two versions of graph IR, NNVM and Relay.
-    """
-    # not to change by default
-    return None
-
 @tvm.target.generic_func
 def conv2d_infer_layout(workload, cfg):
     """Infer input/output shapes and layouts from a workload and cfg.

topi/python/topi/x86/conv2d.py

@@ -27,9 +27,8 @@
 from .. import generic, tag
 from .. import nn
 from ..util import get_const_tuple, get_shape
-from ..nn.conv2d import conv2d, conv2d_NCHWc
-from ..nn.conv2d import conv2d_alter_layout, conv2d_infer_layout, conv2d_rewrite_op
-from ..nn.conv2d import _get_workload as _get_conv2d_workload
+from ..nn.conv2d import conv2d, conv2d_NCHWc, \
+    conv2d_alter_layout, conv2d_infer_layout, _get_workload as _get_conv2d_workload
 from ..nn.depthwise_conv2d import _get_workload as _get_depthwise_conv2d_workload
 from ..nn.depthwise_conv2d import depthwise_conv2d_NCHWc, depthwise_conv2d_nchw
 from ..nn.pad import pad
@@ -38,7 +37,7 @@
 
 logger = logging.getLogger('topi')
 
-def _is_int8_hw_support(data_dtype, kernel_dtype, target, ignore_dtype=False):
+def _is_int8_hw_support(data_dtype, kernel_dtype, target):
     """
     Checks to ensure that we can use Intel DLBoost instructions
     1) The datatypes are correct.
@@ -59,8 +58,6 @@ def _is_int8_hw_support(data_dtype, kernel_dtype, target, ignore_dtype=False):
         if opt == '-mcpu=skylake-avx512':
             is_target_support = True
 
-    if ignore_dtype:
-        return is_llvm_support and is_target_support
     return is_dtype_support and is_llvm_support and is_target_support
 
 def _get_default_config(cfg, data, kernel, strides, padding, out_dtype, is_depthwise=False,
@@ -412,54 +409,6 @@ def _topi_nn_conv2d_NCHWc(*args, **kwargs):
     s = _schedule_conv2d_NCHWc(cfg, [C])
     return s, [new_data, new_kernel, C]
 
-@conv2d_rewrite_op.register("cpu")
-def _conv2d_rewrite_op(attrs, inputs, arg_types, F):
-    if F.__name__ != 'tvm.relay.op':
-        return None
-    data_type, kernel_type = arg_types[0], arg_types[1]
-    target = tvm.target.current_target()
-    data_layout = attrs['data_layout']
-    kernel_layout = attrs['kernel_layout']
-    # Uncomment when this bug is resolved
-    # https://discuss.tvm.ai/t/segfault-in-llvm/3567
-    # if not ((data_layout == 'NCHW' and kernel_layout == 'OIHW')
-    #         or (data_layout == 'NHWC' and kernel_layout == 'HWIO')):
-    #     return None
-    if not (data_layout == 'NCHW' and kernel_layout == 'OIHW'):
-        return None
-
-    if not (data_type.dtype == 'int8' and kernel_type.dtype == 'int8'):
-        return  None
-
-    if not _is_int8_hw_support(data_type.dtype, kernel_type.dtype,
-                               target,
-                               ignore_dtype=True):
-        return None
-
-    # Convert i8 x i8 to u8 x i8
-    # Intel has fast instructions for u8 x i8 conv. For i8 x i8 conv, we can
-    # convert the i8 tensor to u8 by adding 128 and use u8 x i8 conv. Since 128
-    # has been added, the output now has to be adjusted.
-    out_channel = attrs["channels"]
-    data_expr, kernel_expr = inputs
-    data_expr = F.cast(data_expr, "int32")
-    data_expr = F.add(data_expr, F.const(128, "int32"))
-    data_expr = F.clip(data_expr, a_min=0, a_max=255)
-    data_expr = F.cast(data_expr, "uint8")
-    conv = F.nn.conv2d(data_expr, kernel_expr, **attrs)
-    bias_adjust = F.cast(kernel_expr, "int32")
-    if kernel_layout == 'OIHW' and data_layout == 'NCHW':
-        bias_adjust = F.sum(bias_adjust, axis=(1, 2, 3))
-        bias_adjust = F.reshape(bias_adjust,
-                                newshape=(1, out_channel, 1, 1))
-    elif kernel_layout == 'HWIO' and data_layout == 'NHWC':
-        bias_adjust = F.sum(bias_adjust, axis=(0, 1, 2))
-        bias_adjust = F.reshape(bias_adjust,
-                                newshape=(1, 1, 1, out_channel))
-    bias_adjust = F.cast(bias_adjust, 'int32')
-    bias_adjust = F.multiply(bias_adjust, F.const(128, 'int32'))
-    return F.subtract(conv, bias_adjust)
-
 
 @conv2d_alter_layout.register("cpu")
 def _alter_conv2d_layout(attrs, inputs, tinfo, F):