text: # Licensed to the Apache Software Foundation (ASF) under one
  # or more contributor license agreements.  See the NOTICE file
  # distributed with this work for additional information
  # regarding copyright ownership.  The ASF licenses this file
  # 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
  # KIND, either express or implied.  See the License for the
  # specific language governing permissions and limitations
  # under the License.
  """ Support level2 operator test cases.
  """
  import numpy as np
  import tvm
  from tvm import relay
  from tvm.relay import transform
  from tvm.relay.testing import ctx_list
  import topi.testing
  
  def run_infer_type(expr):
      mod = relay.Module.from_expr(expr)
      mod = transform.InferType()(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
      x = relay.var("x", relay.ty.TensorType((n, c, h, w), "float32"))
      w = relay.var("w")
      y = relay.nn.conv2d(x, w,
                          kernel_size=(3, 3),
                          padding=(1, 1),
                          channels=2)
      yy = run_infer_type(y)
      assert yy.checked_type ==  relay.TensorType(
          (n, 2, 224, 224), "float32")
      assert yy.args[1].checked_type == relay.TensorType(
          (2, 10, 3, 3), "float32")
  
      # infer by shape of w, mixed precision
      n, c, h, w = tvm.var("n"), 10, 224, 224
      x = relay.var("x", relay.TensorType((n, c, h, w), "int8"))
      w = relay.var("w", relay.TensorType((2, 10, 3, 3), "int8"))
      y = relay.nn.conv2d(x, w, out_dtype="int32")
      assert "out_dtype=\"int32\"" in y.astext()
      yy = run_infer_type(y)
      assert yy.checked_type ==  relay.TensorType(
          (n, 2, 222, 222), "int32")
  
      # infer shape in case of different dtypes for input and weight.
      n, c, h, w = tvm.var("n"), 10, 224, 224
      x = relay.var("x", relay.TensorType((n, c, h, w), "uint8"))
      w = relay.var("w", relay.TensorType((2, 10, 3, 3), "int8"))
      y = relay.nn.conv2d(x, w, out_dtype="int32")
      assert "out_dtype=\"int32\"" in y.astext()
      yy = run_infer_type(y)
      assert yy.checked_type ==  relay.TensorType(
          (n, 2, 222, 222), "int32")
  
      # Infer with a different layout
      n, c, h, w = 4, 32, 224, 224
      x = relay.var("x", relay.TensorType((n//4, c//4, h, w, 4, 4), "int8"))
      wt = relay.var("w")
      y = relay.nn.conv2d(x, wt,
                          kernel_size=(3, 3),
                          padding=(1, 1),
                          channels=16,
                          data_layout="NCHW4n4c",
                          kernel_layout="OIHW4o4i",
                          out_dtype="int32")
      yy = run_infer_type(y)
      assert yy.checked_type ==  relay.TensorType(
          (1, 4, 224, 224, 4, 4), "int32")
      assert yy.args[1].checked_type == relay.TensorType(
          (4, 8, 3, 3, 4, 4), "int8")
  
      # Infer with NHWC
      n, c, h, w = 4, 32, 224, 224
      x = relay.var("x", relay.TensorType((n, h, w, c), "int8"))
      wt = relay.var("w")
      y = relay.nn.conv2d(x, wt,
                          kernel_size=(3, 3),
                          padding=(1, 1),
                          channels=16,
                          data_layout="NHWC",
                          out_dtype="int32")
      yy = run_infer_type(y)
      assert yy.checked_type ==  relay.TensorType(
          (n, h, w, 16), "int32")
  
  
  def test_conv2d_run():
      def run_test_conv2d(dtype, out_dtype, scale, dshape, kshape,
                          padding=(1, 1),
                          fref=None,
                          groups=1,
                          dilation=(1, 1),
                          except_targets=None,
                          **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,
                              padding=padding,
                              dilation=dilation,
                              groups=groups,
                              **attrs)
          func = relay.Function([x, w], y)
          data = np.random.uniform(-scale, scale, size=dshape).astype(dtype)
          kernel = np.random.uniform(-scale, scale, size=kshape).astype(dtype)
          dkernel = topi.testing.dilate_python(kernel, (1, 1) + dilation)
          if fref is None:
              ref_res = topi.testing.conv2d_nchw_python(
                  data.astype(out_dtype), dkernel.astype(out_dtype), 1, padding,
                  groups=groups)
          else:
              ref_res = fref(data.astype(out_dtype), dkernel.astype(out_dtype))
  
  
          for target, ctx in ctx_list():
              if target in except_targets:
                  continue
              intrp1 = relay.create_executor("graph", ctx=ctx, target=target)
              op_res1 = intrp1.evaluate(func)(data, kernel)
              tvm.testing.assert_allclose(op_res1.asnumpy(), ref_res, rtol=1e-5, atol=1e-5)
  
      # depthwise conv2d
      dshape = (1, 32, 18, 18)
      kshape = (32, 1, 3, 3)
      run_test_conv2d("float32", "float32", 1, dshape, kshape,
                      padding=(1, 1), channels=32, groups=32, kernel_size=(3 ,3),
                      fref=lambda x, w: topi.testing.depthwise_conv2d_python_nchw(
                          x, w, (1, 1), "SAME"))
  
      # CUDA is disabled for 'direct' schedule:
      # https://github.com/dmlc/tvm/pull/3070#issuecomment-486597553
      # group conv2d
      dshape = (1, 32, 18, 18)
      kshape = (32, 4, 3, 3)
      run_test_conv2d("float32", "float32", 1, dshape, kshape,
                      padding=(1, 1), channels=32, groups=8, kernel_size=(3 ,3),
                      except_targets=['cuda'])
      # also group conv2d
      dshape = (1, 32, 18, 18)
      kshape = (64, 1, 3, 3)
      run_test_conv2d("float32", "float32", 1, dshape, kshape,
                      padding=(1, 1), channels=64, groups=32, kernel_size=(3 ,3),
                      except_targets=['cuda'])
  
      # normal conv2d
      dshape = (1, 3, 224, 224)
      kshape = (10, 3, 3, 3)
      run_test_conv2d("float32", "float32", 1, dshape, kshape,
                      padding=(1, 1), channels=10, kernel_size=(3 ,3))
      # mixed precision
      run_test_conv2d("int8", "int32", 1, dshape, kshape,
                      padding=(1, 1), channels=10, kernel_size=(3 ,3))
      kshape = (10, 3, 1, 3)
      # mixed precision.
      run_test_conv2d("int8", "int32", 1, dshape, kshape,
                      padding=(0, 1), channels=10, kernel_size=(1 ,3))
      # dilated conv2d
      dshape = (1, 3, 18, 18)
      kshape = (10, 3, 3, 3)
      run_test_conv2d("float32", "float32", 1, dshape, kshape,
                      padding=(1, 1), channels=10, kernel_size=(3 ,3), dilation=(3, 3))
  
  
  def test_conv2d_transpose_infer_type():
      # symbolic in batch dimension
      n, c, h, w = tvm.var("n"), 10, 10, 12
      x = relay.var("x", relay.TensorType((n, c, h, w), "float32"))
      w = relay.var("w", relay.IncompleteType())
      y = relay.nn.conv2d_transpose(x, w,
                                    kernel_size=(3, 3),
                                    padding=(1, 1),
                                    channels=15)
      assert "channels=15" in y.astext()
      yy = run_infer_type(y)
      assert yy.checked_type == relay.TensorType(
          (n, 15, 10, 12), "float32")
      assert yy.args[1].checked_type == relay.TensorType(
          (10, 15, 3, 3), "float32")
  
      # infer by shape of w, mixed precision
      n, c, h, w = tvm.var("n"), 10, 10, 12
      x = relay.var("x", relay.TensorType((n, c, h, w), "float32"))
      w = relay.var("w", relay.TensorType((12, 11, 5, 5), "float32"))
      y = relay.nn.conv2d_transpose(x, w,
                                    output_padding=(1, 1),
                                    channels=11,
                                    data_layout="NHWC")
      yy = run_infer_type(y)
      assert yy.checked_type == relay.TensorType(
          (n, 15, 15, 11), "float32")
  
  
  def test_conv2d_transpose_run():
      dshape = (1, 3, 18, 18)
      kshape = (3, 10, 3, 3)
      oshape = (1, 10, 37, 37)
      x = relay.var("x", shape=dshape)
      w = relay.var("w")
      y = relay.nn.conv2d_transpose(x, w,
                                    channels=10, kernel_size=(3,3), strides=(2,2),
                                    padding=(1,1), output_padding=(2, 2))
      func = relay.Function([x, w], y)
      dtype = "float32"
      data = np.random.uniform(size=dshape).astype(dtype)
      kernel = np.random.uniform(size=kshape).astype(dtype)
      c_np = topi.testing.conv2d_transpose_nchw_python(
          data, kernel, 2, 1)
      d_np = np.zeros(shape=oshape)
      d_np[:,:,0:c_np.shape[2],0:c_np.shape[3]] = c_np
      ref_res = d_np
  
      for target, ctx in ctx_list():
          intrp1 = relay.create_executor("graph", ctx=ctx, target=target)
          op_res1 = intrp1.evaluate(func)(data, kernel)
          tvm.testing.assert_allclose(op_res1.asnumpy(), ref_res, rtol=1e-5, atol=1e-5)
  
  
  
  def test_upsampling_infer_type():
      n, c , h, w = tvm.var("n"), tvm.var("c"), tvm.var("h"), tvm.var("w")
      x = relay.var("x", relay.TensorType((n, c, h, w), "float32"))
      y = relay.nn.upsampling(x, scale=2, layout="NCHW", method="BILINEAR")
      "method=\"BINLINEAR\"" in y.astext()
      yy = run_infer_type(y)
      assert yy.checked_type == relay.TensorType((n, c, h*2, w*2), "float32")
      n, c = tvm.var("n"), tvm.var("c")
      x = relay.var("x", relay.TensorType((n, c, 100, 200), "float32"))
      y = relay.nn.upsampling(x, scale=2, layout="NCHW", method="BILINEAR")
      yy = run_infer_type(y)
      assert yy.checked_type == relay.TensorType((n, c, 200, 400), "float32")
  
  
  def _test_pool2d(opfunc, reffunc):
      n, c, h, w = tvm.var("n"), 10, 224, 224
      x = relay.var("x", relay.TensorType((n, c, h, w), "float32"))
      y = opfunc(x, pool_size=(1, 1))
      assert "pool_size=" in y.astext()
      yy = run_infer_type(y)
      assert yy.checked_type == relay.TensorType((n, 10, 224, 224), "float32")
      # test execution
      dtype = "float32"
      dshape = (1, 3, 28, 28)
      x = relay.var("x", shape=dshape)
      y = opfunc(x, pool_size=(2, 2), strides=(2, 2), padding=(0, 0))
      func = relay.Function([x], y)
      data = np.random.uniform(size=dshape).astype(dtype)
      ref_res = reffunc(data.reshape(1,3,14,2,14,2), axis=(3,5))
      for target, ctx in ctx_list():
          intrp1 = relay.create_executor("graph", ctx=ctx, target=target)
          op_res1 = intrp1.evaluate(func)(data)
          tvm.testing.assert_allclose(op_res1.asnumpy(), ref_res, rtol=1e-5, atol=1e-5)
  
  
  def _test_global_pool2d(opfunc, reffunc):
      n, c, h, w = tvm.var("n"), tvm.var("c"), 224, 224
      x = relay.var("x", relay.TensorType((n, h, w, c), "float32"))
      y = opfunc(x, layout="NHWC")
      yy = run_infer_type(y)
      assert yy.checked_type == relay.TensorType((n, 1, 1, c), "float32")
  
      n, c, h, w = tvm.var("n"), tvm.var("c"), tvm.var("h"), tvm.var("w")
      x = relay.var("x", relay.TensorType((n, c, h, w), "float32"))
      y = opfunc(x)
      yy = run_infer_type(y)
      assert yy.checked_type == relay.TensorType((n, c, 1, 1), "float32")
      # test execution
      dtype = "float32"
      dshape = (1, 1024, 7, 7)
      x = relay.var("x", shape=dshape)
      y = opfunc(x)
      func = relay.Function([x], y)
      data = np.random.uniform(size=dshape).astype(dtype)
      ref_res = reffunc(data, axis=(2,3), keepdims=True)
      for target, ctx in ctx_list():
          intrp1 = relay.create_executor("graph", ctx=ctx, target=target)
          op_res1 = intrp1.evaluate(func)(data)
          tvm.testing.assert_allclose(op_res1.asnumpy(), ref_res, rtol=1e-5, atol=1e-5)
  
  
  def test_pool2d():
      _test_pool2d(relay.nn.max_pool2d, np.max)
      _test_pool2d(relay.nn.avg_pool2d, np.mean)
      _test_global_pool2d(relay.nn.global_max_pool2d, np.max)
      _test_global_pool2d(relay.nn.global_avg_pool2d, np.mean)
  
  
  def test_avg_pool2d_no_count_pad():
      kh, kw = (4, 4)
      sh, sw = (2, 2)
      ph, pw = (2, 2)
      n = 1
      (ic, ih, iw) = (3, 28, 28)
      (oc, oh, ow) = (3, 15, 15)
      dshape = (n, ic, ih, iw)
      x = relay.var("x", shape=dshape)
      y = relay.nn.avg_pool2d(x,
                              pool_size=(kh, kw),
                              strides=(sw, sw),
                              padding=(ph, pw),
                              count_include_pad=False)
      func = relay.Function([x], y)
      dtype = "float32"
      a_np = np.random.uniform(low=0.001, size=(n, ic, ih, iw)).astype(dtype)
      pad_np = np.zeros(shape=(n, ic, ih+2*ph, iw+2*pw)).astype(dtype)
      no_zero = (range(n), range(ic), (range(ph, ih+ph)), (range(pw, iw+pw)))
      pad_np[np.ix_(*no_zero)] = a_np
      b_np = np.zeros(shape=(n, oc, oh, ow)).astype(dtype)
      for i in range(oh):
          for j in range(ow):
              pad_count = np.sum(pad_np[:, :, i*sh:i*sh+kh, j*sw:j*sw+kw] > 0, axis=(2,3))
              b_np[:,:,i,j] = np.sum(pad_np[:, :, i*sh:i*sh+kh, j*sw:j*sw+kw],
                                     axis=(2,3)) / np.maximum(pad_count, 1)
      ref_res = np.maximum(b_np, 0.0)
      data = a_np
  
      for target, ctx in ctx_list():
          intrp1 = relay.create_executor("graph", ctx=ctx, target=target)
          op_res1 = intrp1.evaluate(func)(data)
          tvm.testing.assert_allclose(op_res1.asnumpy(), ref_res, rtol=1e-5, atol=1e-5)
  
  def test_flatten_infer_type():
      d1, d2, d3, d4 = tvm.var("d1"), tvm.var("d2"), tvm.var("d3"), tvm.var("d4")
      x = relay.var("x", relay.TensorType((d1, d2, d3, d4), "float32"))
      y = relay.nn.batch_flatten(x)
      yy = run_infer_type(y)
      assert yy.checked_type == relay.TensorType((d1, ((d2*d3)*d4)), "float32")
  
      x = relay.var("x", relay.TensorType((3, 2, 4, 3), "float32"))
      y = relay.nn.batch_flatten(x)
      yy = run_infer_type(y)
      assert yy.checked_type == relay.TensorType((3, 24), "float32")
  
      x = relay.var("x", relay.TensorType((d1, 2, d3, 3), "float32"))
      y = relay.nn.batch_flatten(x)
      yy = run_infer_type(y)
      assert yy.checked_type == relay.TensorType((d1, ((2*d3)*3)), "float32")
  
      shape = (1, 5, 10, 10)
      o_shape = (1, 500)
      dtype = "float32"
      x = relay.var("x", relay.TensorType(shape, dtype))
      z = relay.nn.batch_flatten(x)
      yy = run_infer_type(z)
      assert yy.checked_type == relay.TensorType(o_shape, dtype)
      func = relay.Function([x], z)
      x_data = np.random.uniform(low=-1, high=1, size=shape).astype(dtype)
      ref_res = x_data.flatten().reshape(o_shape)
  
      for target, ctx in ctx_list():
          intrp1 = relay.create_executor("graph", ctx=ctx, target=target)
          intrp2 = relay.create_executor("debug", ctx=ctx, target=target)
          op_res1 = intrp1.evaluate(func)(x_data)
          tvm.testing.assert_allclose(op_res1.asnumpy(), ref_res, rtol=1e-5)
          op_res2 = intrp2.evaluate(func)(x_data)
          tvm.testing.assert_allclose(op_res2.asnumpy(), ref_res, rtol=1e-5)
  
  def test_pad_infer_type():
      # entirely concrete case
      n, c, h, w = 1, 2, 3, 4
      t = relay.var("t", relay.TensorType((n, c, h, w), "float32"))
      y = relay.nn.pad(t, ((1, 1), (2, 2), (3, 3), (4, 4)))
      "pad_width=" in y.astext()
      yy = run_infer_type(y)
      assert yy.checked_type == relay.TensorType((3, 6, 9, 12), "float32")
  
      # some symbolic values
      n, c, h, w = tvm.var("n"), 2, 3, tvm.var("w")
      t = relay.var("t", relay.TensorType((n, c, h, w), "float32"))
      y = relay.nn.pad(t, ((1, 1), (2, 2), (3, 3), (4, 4)))
      yy = run_infer_type(y)
      assert yy.checked_type == relay.TensorType((n + 2, 6, 9, w + 8), "float32")
  
  def test_pad_run():
      def _test_run(dtype):
          dshape = (4, 10, 7, 7)
          x = relay.var("x", shape=dshape)
          y = relay.nn.pad(x, ((1, 1), (2, 2), (3, 3), (4, 4)))
          func = relay.Function([x], y)
          data = np.random.uniform(size=dshape).astype(dtype)
          ref_res = np.pad(data, ((1, 1), (2, 2), (3, 3), (4, 4)), 'constant')
          for target, ctx in ctx_list():
              intrp1 = relay.create_executor("graph", ctx=ctx, target=target)
              op_res1 = intrp1.evaluate(func)(data)
              tvm.testing.assert_allclose(op_res1.asnumpy(), ref_res, rtol=1e-5, atol=1e-5)
  
      _test_run('float32')
      _test_run('int32')
  
  def test_lrn():
      n, c , h, w = tvm.var("n"), tvm.var("c"), tvm.var("h"), tvm.var("w")
      x = relay.var("x", shape=(n, c , h, w))
      y = relay.nn.lrn(x, size=10, axis=2, bias=0.5, alpha=.00001, beta=0.75)
      "alpha=" in y.astext()
      yy = run_infer_type(y)
      assert yy.checked_type == relay.TensorType((n, c , h, w))
  
      shape = (1, 5, 10, 10)
      dtype = "float32"
      x = relay.var("x", relay.TensorType(shape, dtype))
      size=5
      axis=1
      bias=0.5
      alpha=.00001
      beta=0.75
      z = relay.nn.lrn(x, size=size, axis=axis, bias=bias, alpha=alpha, beta=beta)
      yy = run_infer_type(z)
      assert yy.checked_type == relay.TensorType(shape, dtype)
      func = relay.Function([x], z)
      x_data = np.random.uniform(low=-1, high=1, size=shape).astype(dtype)
      ref_res = topi.testing.lrn_python(x_data, size, axis, bias, alpha, beta)
  
      for target, ctx in ctx_list():
          intrp1 = relay.create_executor("graph", ctx=ctx, target=target)
          intrp2 = relay.create_executor("debug", ctx=ctx, target=target)
          op_res1 = intrp1.evaluate(func)(x_data)
          tvm.testing.assert_allclose(op_res1.asnumpy(), ref_res, rtol=1e-5)
          op_res2 = intrp2.evaluate(func)(x_data)
          tvm.testing.assert_allclose(op_res2.asnumpy(), ref_res, rtol=1e-5)
  
  def test_l2_normalize():
      n, c , h, w = tvm.var("n"), tvm.var("c"), tvm.var("h"), tvm.var("w")
      x = relay.var("x", shape=(n, c , h, w))
      y = relay.nn.l2_normalize(x, eps=0.001, axis=[1])
      "axis=" in y.astext()
      yy = run_infer_type(y)
      assert yy.checked_type == relay.TensorType((n, c , h, w))
  
      shape = (1, 5, 10, 10)
      dtype = "float32"
      x = relay.var("x", relay.TensorType(shape, dtype))
      eps=0.001
      axis=1
      z = relay.nn.l2_normalize(x, eps=0.001, axis=[axis])
      yy = run_infer_type(z)
      assert yy.checked_type == relay.TensorType(shape, dtype)
      func = relay.Function([x], z)
      x_data = np.random.uniform(low=-1, high=1, size=shape).astype(dtype)
      ref_res = topi.testing.l2_normalize_python(x_data, eps, axis)
  
      for target, ctx in ctx_list():
          intrp1 = relay.create_executor("graph", ctx=ctx, target=target)
          intrp2 = relay.create_executor("debug", ctx=ctx, target=target)
          op_res1 = intrp1.evaluate(func)(x_data)
          tvm.testing.assert_allclose(op_res1.asnumpy(), ref_res, rtol=1e-5)
          op_res2 = intrp2.evaluate(func)(x_data)
          tvm.testing.assert_allclose(op_res2.asnumpy(), ref_res, rtol=1e-5)
  
  
  def batch_flatten(data):
      shape = data.shape
      target_dim = 1
      for i in range(len(shape) - 1):
          target_dim = target_dim * shape[i + 1]
      return np.reshape(data, (shape[0], target_dim))
  
  
  def test_batch_flatten():
      t1 = relay.TensorType((5, 10, 5))
      x = relay.Var("x", t1)
      func = relay.Function([x], relay.nn.batch_flatten(x))
  
      data = np.random.rand(5, 10, 5).astype(t1.dtype)
      ref_res = batch_flatten(data)
      for target, ctx in ctx_list():
          intrp = relay.create_executor("graph", ctx=ctx, target=target)
          op_res = intrp.evaluate(func)(data)
          np.testing.assert_allclose(op_res.asnumpy(), ref_res, rtol=0.01)
  
  
  def _test_upsampling(layout, method):
      n, c, h, w = tvm.var("n"), 16, 32, 32
      scale = 2
      dtype = "float32"
      def get_shape():
          if layout == "NCHW":
              return (c, h, w), (c, h*scale, w*scale)
          else:
              return (h, w, c), (h*scale, w*scale, c)
      ishape, oshape = get_shape()
      x = relay.var("x", relay.TensorType((n,) + ishape, dtype))
      y = relay.nn.upsampling(x, scale=scale, layout=layout, method=method)
      yy = run_infer_type(y)
      assert yy.checked_type == relay.TensorType((n,) + oshape, dtype)
      dshape = (1,) + ishape
      x = relay.var("x", shape=dshape)
      y = relay.nn.upsampling(x, scale=scale, layout=layout, method=method)
      func = relay.Function([x], y)
      data = np.random.uniform(size=dshape).astype(dtype)
      if method == "NEAREST_NEIGHBOR":
          ref = topi.testing.upsampling_python(data, (scale, scale), layout)
      else:
          ref = topi.testing.bilinear_resize_python(data, (h*scale, w*scale), layout)
      for target, ctx in ctx_list():
          executor = relay.create_executor("graph", ctx=ctx, target=target)
          out = executor.evaluate(func)(data)
          tvm.testing.assert_allclose(out.asnumpy(), ref, rtol=1e-5, atol=1e-5)
  
  
  def test_upsampling():
      _test_upsampling("NCHW", "NEAREST_NEIGHBOR")
      _test_upsampling("NCHW", "BILINEAR")
      _test_upsampling("NHWC", "NEAREST_NEIGHBOR")
      _test_upsampling("NHWC", "BILINEAR")
  
  
  def test_conv2d_int8_intrinsics():
      def _compile(input_dtype, weight_dtype, output_dtype, target):
          n, ic, h, w, oc, ch, cw = 1, 16, 224, 224, 32, 3, 3
          x = relay.var("x", relay.TensorType((n, ic, h, w), input_dtype))
          w = relay.var("w", relay.TensorType((oc, ic, ch, cw), weight_dtype))
          y = relay.nn.conv2d(x, w,
                              kernel_size=(ch, cw),
                              channels=oc,
                              padding=(1, 1),
                              dilation=(1, 1),
                              out_dtype=output_dtype)
          func = relay.Function([x, w], y)
          wdata = np.random.rand(oc, ic, ch, cw) * 10
          parameters = {"w": tvm.nd.array(wdata.astype(weight_dtype))}
          with relay.build_config(opt_level=3):
              graph, lib, params = relay.build(func, target, params=parameters)
          assembly = lib.get_source("asm")
          return assembly
  
      # compile conv2d for x86 (skylake) and test assembly contains *pmadd* instructions
      target = "llvm -mcpu=skylake-avx512"
      name = "llvm.x86.avx512.pmaddubs.w.512"
      llvm_id = tvm.codegen.llvm_lookup_intrinsic_id(name)
      if llvm_id != 0:
          # Intel Int8 instruction need uint8 data and int8 kernel
          asm = _compile(input_dtype="uint8",
                         weight_dtype="int8",
                         output_dtype="int32",
                         target=target)
          # Check that intrinisic is present in the assembly.
          assert "pmaddubs" in asm
  
          # Ensure that code is generated when datatypes are not HW supported.
          asm = _compile(input_dtype="int8",
                         weight_dtype="int8",
                         output_dtype="int32",
                         target=target)
          # 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",
                         weight_dtype="uint8",
                         output_dtype="int32",
                         target=target)
          # Check that intrinisic is not present in the assembly.
          assert "pmaddubs" not in asm
  
      # Check that a vectorized instruction is generated for older Intel
      # generations, because we default to NCHWc layout.
      target = "llvm -mcpu=core-avx2"
      asm = _compile(input_dtype="int8",
                    weight_dtype="int8",
                    output_dtype="int32",
                    target=target)
      # Check that vector int mult and add instructions are generated.
      assert "vpmulld" in asm and "vpadd" in asm
  
  
  if __name__ == "__main__":
      test_pool2d()
      test_avg_pool2d_no_count_pad()
      test_lrn()
      test_l2_normalize()
      test_conv2d_infer_type()
      test_upsampling_infer_type()
      test_flatten_infer_type()
      test_pad_infer_type()
      test_pad_run()
      test_conv2d_transpose_infer_type()
      test_conv2d_transpose_run()
      test_conv2d_run()
      test_batch_flatten()
      test_upsampling()
      test_conv2d_int8_intrinsics()
  