code: # 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.
  # pylint: disable=import-self, invalid-name, unused-argument, too-many-lines, len-as-condition
  
  import tvm
  from tvm import te
  import numpy as np
  from topi.x86.tensor_intrin import dot_16x1x16_uint8_int8_int32_cascadelake
  from topi.x86.tensor_intrin import dot_16x1x16_uint8_int8_int32
  import pytest
  
  
  @pytest.mark.skip("skip because feature not enabled")
  def test_fc_int8_acc32():
      m = 1024
      n = 1024
      k = 1024
  
      X = te.placeholder((m, k), name='X', dtype="uint8")
      W = te.placeholder((n, k), name='W', dtype="int8")
  
      peak = 280
      print("Peak {} Gops/s".format(peak))
      memory_ops = m * k + n * k + 2 * m * n
      gops_per_mm = 2 * m * n * k
  
      # For LLVM < 8.0, it shows "'cascadelake' is not a recognized processor for this target
      # (ignoring processor)" error with the following setting. After LLVM 8.0 is enabled in the
      # test, we should use cascadelake setting.
      def verify(target="llvm -mcpu=cascadelake"):
          if not tvm.runtime.enabled(target):
              print("skip because %s is not enabled..." % target)
              return
  
          ctx = tvm.context(target, 0)
          pc = dot_16x1x16_uint8_int8_int32_cascadelake()
          ak = te.reduce_axis((0, k), name='k')
          packedW = te.placeholder(
              (n // 16, 16 * (k // 4), 4), name='packedW', dtype="int8")
  
          t_fc = te.compute((m, n), lambda i, j: te.sum(X[i, ak].astype(
              "int32") * packedW[j / 16, (ak / 4) * 16 + j % 16, ak % 4].astype("int32"), axis=ak), name="F")
          t_sch = te.create_schedule(t_fc.op)
          a_x, a_y = t_fc.op.axis
          a_k, = t_fc.op.reduce_axis
  
          a_yo, a_yi = t_sch[t_fc].split(a_y, factor=16)
          a_xo, a_xi = t_sch[t_fc].split(a_x, factor=32)
          a_ko, a_ki = t_sch[t_fc].split(a_k, factor=4)
          a_koo, a_koi = t_sch[t_fc].split(a_ko, factor=4)
          t_sch[t_fc].reorder(a_yo, a_xo, a_xi, a_koo, a_koi, a_yi, a_ki)
  
          t_sch[t_fc].unroll(a_koi)
          t_sch[t_fc].tensorize(a_yi, pc)
  
          t_func = tvm.build(t_sch, [X, packedW, t_fc], target, name="intrinsic")
          t_evaluator = t_func.time_evaluator(t_func.entry_name, ctx, number=10)
  
          # generate the plain data
          a_ = np.random.uniform(1, 10, size=(m, k)).astype("uint8")
          b_ = np.random.uniform(1, 10, size=(n, k)).astype("int8")
  
          packW = np.random.uniform(1, 10, size=(
              n // 16, 16 * (k // 4), 4)).astype("int8")
          # This occurs in pre_compute stage
          for r_idx in range(n // 16):
              for s_idx in range(16 * (k // 4)):
                  for t_idx in range(4):
                      packW[r_idx][s_idx][t_idx] = b_[r_idx * 16 + s_idx %
                                                      16][(s_idx // 16) * 4 + t_idx]
  
          x = tvm.nd.array(a_, ctx)
          w = tvm.nd.array(packW, ctx)
          y = tvm.nd.array(np.zeros((m, n), dtype="int32"), ctx)
          result = t_evaluator(x, w, y)
  
          gops_per_sec = gops_per_mm / result.mean / 1e9
          # verify the correctness
          tvm.testing.assert_allclose(y.asnumpy(), np.dot(a_, b_.T), rtol=0)
          print('Tensorization: running time: {:.3f} ms, {:.2f} Gops/s, effiency: {:.2f}'.format(
              result.mean * 1000, gops_per_sec, gops_per_sec / peak))
          t_func.export_library("tensorize_acc32.o")
  
      verify()
  
  
  if __name__ == "__main__":
      # The test requires Cascade Lake and newer Intel machines to generate the
      # correct AVX512 VNNI instruction. So, disabling the test.
  
      # test_fc_int8_acc32()
      pass
   # 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.
  import numpy as np
  import pytest
  
  import tvm
  from tvm import te
  from tvm import relay
  from tvm.relay import testing
  from tvm.relay.expr import Call
  
  
  def quantize_and_build(out):
      f = relay.Function(relay.analysis.free_vars(out), out)
      mod, params = testing.create_workload(f)
  
      with relay.quantize.qconfig(skip_conv_layers=[]):
          qmod = relay.quantize.quantize(mod, params)
  
      relay.build(qmod, "llvm", params=params)
  
      return qmod
  
  def test_mul_rewrite():
      """a test case where rhs of mul is not constant"""
      data = relay.var("data", shape=(1, 16, 64, 64))
      multiplier = relay.sigmoid(relay.var("data", shape=(1, 16, 1, 1)))
      conv = relay.nn.conv2d(data, relay.var("weight"),
                             kernel_size=(3, 3),
                             padding=(1, 1),
                             channels=16)
      act = relay.nn.relu(data=conv)
  
      quantize_and_build(act * multiplier)
  
      pool = relay.nn.global_avg_pool2d(data=act)
  
      quantize_and_build(act * pool)
  
  def test_batch_flatten_rewrite():
  
      data = relay.var("data", shape=(1, 16, 64, 64), dtype="float32")
  
      out = relay.nn.conv2d(data, relay.var("weight"),
                            kernel_size=(3, 3),
                            padding=(1, 1),
                            channels=16)
  
      out = relay.nn.batch_flatten(out)
  
      qmod = quantize_and_build(out)
  
      def _check_batch_flatten(node):
          if isinstance(node, Call):
              if(node.op.name == "nn.batch_flatten"):
                 assert node.checked_type.dtype == "int8"
  
      # check if batch_flatten is quantized
      relay.analysis.post_order_visit(qmod["main"], _check_batch_flatten)
  
  def get_calibration_dataset(input_name):
      dataset = []
      for i in range(5):
          data = np.random.uniform(size=(1, 3, 224, 224))
          dataset.append({input_name: data})
      return dataset
  
  
  @pytest.mark.parametrize("create_target", [True, False])
  def test_calibrate_target(create_target):
      mod, params = testing.resnet.get_workload(num_layers=18)
      dataset = get_calibration_dataset("data")
      with relay.quantize.qconfig(calibrate_mode="kl_divergence"):
          if create_target:
              with tvm.target.create("llvm"):
                  relay.quantize.quantize(mod, params, dataset)
          else:
              # current_target = None
              relay.quantize.quantize(mod, params, dataset)
  
  
  def test_calibrate_memory_bound():
      mod, params = testing.resnet.get_workload(num_layers=18)
      dataset = get_calibration_dataset("data")
      import multiprocessing
      num_cpu = multiprocessing.cpu_count()
      with relay.quantize.qconfig(calibrate_mode="kl_divergence",
                                  calibrate_chunk_by=num_cpu):
          relay.quantize.quantize(mod, params, dataset)
  
  
  if __name__ == "__main__":
      test_mul_rewrite()
      test_batch_flatten_rewrite()
      test_calibrate_target(False)
      test_calibrate_target(True)
      test_calibrate_memory_bound()
   # 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.
  # pylint: disable=invalid-name, too-many-locals, too-many-statements, unused-argument
  """Test code for dense tensorcore operator"""
  import numpy as np
  import tvm
  import topi
  import topi.testing
  from topi.util import get_const_tuple
  from tvm import te
  from tvm.contrib.pickle_memoize import memoize
  from tvm.contrib import nvcc
  
  
  _dense_implement = {
      "gpu": [(topi.cuda.dense_tensorcore, topi.cuda.schedule_dense_tensorcore)]
  }
  
  def verify_dense(batch, in_dim, out_dim, use_bias=True):
      """Dense tensorcore verify function"""
      A = te.placeholder((batch, in_dim), name='A')
      B = te.placeholder((out_dim, in_dim), name='B')
      C = te.placeholder((out_dim,), name='C')
      dtype = A.dtype
  
      # use memoize to pickle the test data for next time use
      @memoize("topi.tests.test_topi_dense_tensorcore")
      def get_ref_data():
          a_np = np.random.uniform(size=(batch, in_dim)).astype(dtype)
          b_np = np.random.uniform(size=(out_dim, in_dim)).astype(dtype)
          c_np = np.random.uniform(size=(out_dim,)).astype(dtype)
          if use_bias:
              d_np = np.maximum(np.dot(a_np, b_np.T) + c_np, 0.0)
          else:
              d_np = np.maximum(np.dot(a_np, b_np.T), 0.0)
          return (a_np, b_np, c_np, d_np)
      # get the test data
      a_np, b_np, c_np, d_np = get_ref_data()
  
      def check_device(device):
          ctx = tvm.context(device, 0)
          if not ctx.exist:
              print("Skip because %s is not enabled" % device)
              return
          if not nvcc.have_tensorcore(ctx.compute_version):
              print("skip because gpu does not support Tensor Cores")
              return
          print("Running on target: %s" % device)
          for fcompute, fschedule in topi.testing.dispatch(device, _dense_implement):
              with tvm.target.create(device):
                  D = fcompute(A, B, C if use_bias else None)
                  D = topi.nn.relu(D)
                  s = fschedule([D])
              a = tvm.nd.array(a_np, ctx)
              b = tvm.nd.array(b_np, ctx)
              c = tvm.nd.array(c_np, ctx)
              d = tvm.nd.array(np.zeros(get_const_tuple(D.shape), dtype=dtype), ctx)
              f = tvm.build(s, [A, B, C, D], device, name="dense")
              f(a, b, c, d)
              tvm.testing.assert_allclose(d.asnumpy(), d_np, rtol=1e-3)
  
  
      for device in ['cuda']:
          check_device(device)
  
  
  def test_dense_tensorcore():
      """Test cases"""
      verify_dense(8, 16, 32, use_bias=True)
      verify_dense(16, 32, 16, use_bias=True)
      verify_dense(256, 1024, 1024, use_bias=True)
      verify_dense(1000, 1024, 1024, use_bias=False)
      verify_dense(256, 2048, 1000, use_bias=False)
  
  
  if __name__ == "__main__":
      test_dense_tensorcore()
  
apis: numpy.random.uniform numpy.dot numpy.zeros
apis: numpy.random.uniform
apis: numpy.random.uniform numpy.dot