code: def tune_and_evaluate(M, N, L, dtype, layout):
      task = autotvm.task.create(
          "tutorial/auto_tensorcore/test_gemm",
          args=(N, L, M, dtype, layout),
          target="cuda",
      )
      print(task.config_space)
  
      logging.getLogger("autotvm").setLevel(logging.DEBUG)
      logging.getLogger("autotvm").addHandler(
          logging.StreamHandler(sys.stdout)
      )
  
      measure_option = autotvm.measure_option(
          builder="local", runner=autotvm.LocalRunner(number=5)
      )
  
      tuner = autotvm.tuner.XGBTuner(task)
      tuner.tune(
          n_trial=1000,
          measure_option=measure_option,
          callbacks=[autotvm.callback.log_to_file("matmul.log")],
      )
  
      dispatch_context = autotvm.apply_history_best("matmul.log")
      best_config = dispatch_context.query(task.target, task.workload)
      print("\nBest config:")
      print(best_config)
      with autotvm.apply_history_best("matmul.log"):
          with tvm.target.Target("cuda"):
              s, arg_bufs = test_gemm(N, L, M, dtype, layout)
              print(tvm.lower(s, arg_bufs, simple_mode=True))
              func = tvm.build(s, arg_bufs)
      dev_module = func.imported_modules[0]
      print(dev_module.get_source())
  
      # check correctness
      if layout == "NN":
          shape_a = (N, L)
          shape_b = (L, M)
      elif layout == "NT":
          shape_a = (L, N)
          shape_b = (L, M)
      elif layout == "TN":
          shape_a = (N, L)
          shape_b = (M, L)
      elif layout == "TT":
          shape_a = (L, N)
          shape_b = (M, L)
  
      a_np = None
      b_np = None
      c_np = None
      c_np_type = None
      if dtype == "float16":
          c_np_type = np.float32
          a_np = np.random.uniform(size=shape_a).astype(np.float16)
          b_np = np.random.uniform(size=shape_b).astype(np.float16)
          if layout == "NN":
              c_np = np.dot(a_np, b_np)
          elif layout == "NT":
              c_np = np.dot(a_np.T, b_np)
          elif layout == "TN":
              c_np = np.dot(a_np, b_np.T)
          elif layout == "TT":
              c_np = np.dot(a_np.T, b_np.T)
      elif dtype == "int8":
          c_np_type = np.int32
          a_np = np.random.randint(
              low=-128, high=127, size=shape_a
          ).astype(np.int8)
          b_np = np.random.randint(
              low=-128, high=127, size=shape_b
          ).astype(np.int8)
          if layout == "NN":
              c_np = np.dot(
                  a_np.astype(np.int32), b_np.astype(np.int32)
              )
          elif layout == "NT":
              c_np = np.dot(
                  a_np.astype(np.int32).T, b_np.astype(np.int32)
              )
          elif layout == "TN":
              c_np = np.dot(
                  a_np.astype(np.int32), b_np.astype(np.int32).T
              )
          elif layout == "TT":
              c_np = np.dot(
                  a_np.astype(np.int32).T, b_np.astype(np.int32).T
              )
      elif dtype == "int4":
          c_np_type = np.int32
          a_np_int = np.random.randint(
              low=-8, high=7, size=shape_a
          ).astype(np.int32)
          b_np_int = np.random.randint(
              low=-8, high=7, size=shape_b
          ).astype(np.int32)
          # "TN"
          c_np = np.dot(
              a_np_int.astype(np.int32), b_np_int.astype(np.int32).T
          )
          a_np = np.zeros(shape=(N, int(L / 8)), dtype=np.int32)
          b_np = np.zeros(shape=(M, int(L / 8)), dtype=np.int32)
          # a_np --> col_major
          for i in range(N):
              for j in range(int(L / 8)):
                  for k in range(8):
                      a_np[i, j] = a_np[i, j] | (
                          (a_np_int[i, j * 8 + k] & 0xF)
                          << ((7 - k) * 4)
                      )
  
          # b_np --> row_major
          for i in range(M):
              for j in range(int(L / 8)):
                  for k in range(8):
                      b_np[i, j] = b_np[i, j] | (
                          (b_np_int[i, j * 8 + k] & 0xF)
                          << ((7 - k) * 4)
                      )
      elif dtype == "int1":
          c_np_type = np.int32
          a_np_int = np.random.randint(
              low=0, high=1, size=shape_a
          ).astype(np.int32)
          b_np_int = np.random.randint(
              low=0, high=1, size=shape_b
          ).astype(np.int32)
          # "TN"
          c_np = np.dot(
              a_np_int.astype(np.int32), b_np_int.astype(np.int32).T
          )
          a_np = np.zeros(shape=(N, int(L / 32)), dtype=np.int32)
          b_np = np.zeros(shape=(M, int(L / 32)), dtype=np.int32)
          for i in range(N):
              for j in range(int(L / 32)):
                  for k in range(32):
                      a_np[i, j] = a_np[i, j] | (
                          (a_np_int[i, j * 32 + k] & 0xF) << (31 - k)
                      )
  
          for i in range(M):
              for j in range(int(L / 32)):
                  for k in range(32):
                      b_np[i, j] = b_np[i, j] | (
                          (b_np_int[i, j * 32 + k] & 0xF) << (31 - k)
                      )
  
      c_tvm = tvm.nd.array(
          np.zeros(c_np.shape, dtype=c_np_type), device=ctx
      )
      a_tvm = tvm.nd.array(a_np, device=ctx)
      b_tvm = tvm.nd.array(b_np, device=ctx)
      func(a_tvm, b_tvm, c_tvm)
  
      tvm.testing.assert_allclose(c_np, c_tvm.asnumpy(), rtol=1e-3)
  
      evaluator = func.time_evaluator(func.entry_name, ctx, number=100)
      print(
          "Time cost of this operator: %f"
          % evaluator(a_tvm, b_tvm, c_tvm).mean
      )
  
      func.save("gemm.o")
      func.imported_modules[0].save("gemm.ptx")
      cc.create_shared("gemm.so", ["gemm.o"])
  
file_path: lightseq/tests/tvm/matmul_autotune.py
imports: import logging
  import sys
  import numpy
  import tvm
  import tvm.testing
  from tvm import te
  from tvm import autotvm
  from tvm.contrib import nvcc, cc
label_desc: Use image node_id 3 for calling a global function with example usage: tune_and_evaluate(M, N, L, dtype, layout) without return types