code: def test_mma_tune():
      arch = tvm.contrib.nvcc.get_target_compute_version()
      major, _ = tvm.contrib.nvcc.parse_compute_version(arch)
      if major < 8:
          # At least sm80 is required
          return
  
      # pylint: disable=import-outside-toplevel
      from tvm.contrib import cublas
  
      def tune(out_dtype):
          M, N, K = 1024, 1024, 1024
          target = Target("nvidia/geforce-rtx-3080")
          func = te.create_prim_func(
              matmul_fp16(N=N, M=M, K=K, out_dtype=out_dtype)
          ).with_attr({"global_symbol": "main"})
          mod = tvm.IRModule({"main": func})
  
          with tempfile.TemporaryDirectory() as work_dir:
              db = ms.tir_integration.tune_tir(
                  mod=mod,
                  target=target,
                  work_dir=work_dir,
                  max_trials_global=8,
                  builder=LocalBuilder(
                      f_build="meta_schedule.builder.async_build",
                      initializer=initializer,
                  ),
                  space=ms.space_generator.PostOrderApply(
                      sch_rules=[
                          multi_level_tiling_mma(out_dtype=out_dtype)
                      ],
                  ),
              )
              sch = db.query_schedule(
                  mod, target=target, workload_name="main"
              )
              with tvm.transform.PassContext(
                  config={"tir.use_async_copy": 1}
              ):
                  rt_mod = tvm.build(sch.mod, target=target)
  
          a_np = np.random.uniform(0, 1, size=(M, K)).astype("float16")
          b_np = np.random.uniform(0, 1, size=(K, N)).astype("float16")
          A_cublas = te.placeholder((M, K), name="A", dtype="float16")
          B_cublas = te.placeholder((K, N), name="B", dtype="float16")
          C_cublas = cublas.matmul(A_cublas, B_cublas, dtype=out_dtype)
          s = te.create_schedule(C_cublas.op)
          dev = tvm.cuda(0)
          f_cublas = tvm.build(
              s, [A_cublas, B_cublas, C_cublas], target
          )
          a_cublas = tvm.nd.array(a_np.astype("float16"), dev)
          b_cublas = tvm.nd.array(b_np.astype("float16"), dev)
          c_cublas = tvm.nd.array(
              np.zeros((M, N), dtype=C_cublas.dtype), dev
          )
          f_cublas(a_cublas, b_cublas, c_cublas)
          a_tvm = tvm.nd.array(a_np, device=tvm.cuda(0))
          b_tvm = tvm.nd.array(b_np, device=tvm.cuda(0))
          c_tvm = tvm.nd.array(
              np.empty((M, N)).astype(out_dtype), device=tvm.cuda(0)
          )
          rt_mod(a_tvm, b_tvm, c_tvm)
          assert np.allclose(c_tvm.numpy(), c_cublas.numpy(), rtol=1e-2)
  
      tune("float16")
      tune("float32")
  
imports: import tempfile
  import numpy
  import tvm
  from tvm import te
  from tvm import meta_schedule
  from tvm._ffi import register_func
  from tvm.meta_schedule.testing.space_generation import check_sketches, generate_design_space
  from tvm.meta_schedule.builder import LocalBuilder
  from tvm.script import ir
  from tvm.script import tir
  from tvm.target import Target
  from tvm.tir import Schedule
  from tvm.tir.schedule import Trace
  from tvm.tir.tensor_intrin import cuda
  import tvm.testing
label_desc: Use image node_id 7 for calling a global function with example usage: test_mma_tune() without return types