code: def test_tensor_core_batch_matmal():
      batch_size = 4
      n = 512
      m, l = n, n
      assert n % 32 == 0
      assert m % 8 == 0
      assert l % 16 == 0
      nn, mm, ll = n // 32, m // 8, l // 16
      A = te.placeholder(
          (batch_size, nn, ll, 32, 16), name="A", dtype="float16"
      )
      B = te.placeholder(
          (batch_size, ll, mm, 16, 8), name="B", dtype="float16"
      )
      k1 = te.reduce_axis((0, ll), name="k1")
      k2 = te.reduce_axis((0, 16), name="k2")
      C = te.compute(
          (batch_size, nn, mm, 32, 8),
          lambda b, i, j, ii, jj: te.sum(
              A[b, i, k1, ii, k2].astype("float")
              * B[b, k1, j, k2, jj].astype("float"),
              axis=[k1, k2],
          ),
          name="Fragment_C",
      )
      s = te.create_schedule(C.op)
  
      warp_size = 32
      kernel_size = 16
      block_row_warps = 2
      block_col_warps = 4
      warp_row_tiles = 4
      warp_col_tiles = 2
      chunk = 4
  
      block_x = te.thread_axis("blockIdx.x")
      block_y = te.thread_axis("blockIdx.y")
      block_z = te.thread_axis("blockIdx.z")
      thread_x = te.thread_axis("threadIdx.x")
      thread_y = te.thread_axis("threadIdx.y")
      thread_z = te.thread_axis("threadIdx.z")
  
      AS = s.cache_read(A, "shared", [C])
      BS = s.cache_read(B, "shared", [C])
      AF = s.cache_read(AS, "wmma.matrix_a", [C])
      BF = s.cache_read(BS, "wmma.matrix_b", [C])
      CF = s.cache_write(C, "wmma.accumulator")
  
      b, i, j, kernel_i, kernel_j = s[C].op.axis
      i, ii = s[C].split(i, factor=warp_row_tiles)
      block_i, i = s[C].split(i, factor=block_row_warps)
      j, jj = s[C].split(j, factor=warp_col_tiles)
      block_j, j = s[C].split(j, factor=block_col_warps)
      s[C].reorder(block_i, block_j, i, j, ii, jj, kernel_i, kernel_j)
      s[C].bind(b, block_z)
      s[C].bind(block_i, block_x)
      s[C].bind(block_j, block_y)
      s[C].bind(i, thread_y)
      s[C].bind(j, thread_z)
  
      s[CF].compute_at(s[C], j)
      b, warp_i, warp_j, _i, _j = s[CF].op.axis
      k, _k = CF.op.reduce_axis
      ko, ki = s[CF].split(k, factor=chunk)
      s[CF].reorder(ko, ki, warp_i, warp_j, _i, _j, _k)
  
      s[AF].compute_at(s[CF], ki)
      s[BF].compute_at(s[CF], ki)
  
      s[AS].compute_at(s[CF], ko)
      b, xo, yo, xi, yi = AS.op.axis
      tx, xo = s[AS].split(xo, nparts=block_row_warps)
      ty, yo = s[AS].split(yo, nparts=block_col_warps)
      t = s[AS].fuse(xi, yi)
      to, ti = s[AS].split(t, nparts=warp_size)
      s[AS].bind(tx, thread_y)
      s[AS].bind(ty, thread_z)
      s[AS].bind(to, thread_x)
  
      s[BS].compute_at(s[CF], ko)
      b, xo, yo, xi, yi = BS.op.axis
      tx, xo = s[BS].split(xo, nparts=block_row_warps)
      ty, yo = s[BS].split(yo, nparts=block_col_warps)
      t = s[BS].fuse(xi, yi)
      to, ti = s[BS].split(t, nparts=warp_size)
      s[BS].bind(tx, thread_y)
      s[BS].bind(ty, thread_z)
      s[BS].bind(to, thread_x)
  
      s[AF].tensorize(
          AF.op.axis[-2],
          intrin_wmma_load_matrix((32, 8, 16), "wmma.matrix_a"),
      )
      s[BF].tensorize(
          BF.op.axis[-2],
          intrin_wmma_load_matrix((32, 8, 16), "wmma.matrix_b"),
      )
      s[C].tensorize(kernel_i, intrin_wmma_store_matrix((32, 8, 16)))
      s[CF].tensorize(_i, intrin_wmma_gemm((32, 8, 16)))
  
      func = tvm.build(s, [A, B, C], "cuda")
  
      dev = tvm.cuda(0)
      a_np = np.random.uniform(
          size=(batch_size, nn, ll, 32, 16)
      ).astype(A.dtype)
      b_np = np.random.uniform(size=(batch_size, ll, mm, 16, 8)).astype(
          B.dtype
      )
      a = tvm.nd.array(a_np, dev)
      b = tvm.nd.array(b_np, dev)
      c = tvm.nd.array(
          np.zeros((batch_size, nn, mm, 32, 8), dtype=C.dtype), dev
      )
      func(a, b, c)
      evaluator = func.time_evaluator(func.entry_name, dev, number=3)
      print(
          "gemm with tensor core: %f ms"
          % (evaluator(a, b, c).mean * 1e3)
      )
  
      if VERIFY:
          func(a, b, c)
          a_np = a_np.transpose((0, 1, 3, 2, 4)).reshape(
              batch_size, n, n
          )
          b_np = b_np.transpose((0, 1, 3, 2, 4)).reshape(
              batch_size, n, n
          )
          c_np = (
              c.numpy()
              .transpose((0, 1, 3, 2, 4))
              .reshape(batch_size, n, n)
          )
          np.testing.assert_allclose(
              c_np,
              np.matmul(a_np.astype(C.dtype), b_np.astype(C.dtype)),
              rtol=1e-4,
              atol=1e-4,
          )
  
file_path: tvm/tests/python/te/test_te_schedule_tensor_core.py
imports: import tvm
  from tvm import te
  import numpy
  from tvm.topi.testing import conv2d_nhwc_python
  import tvm.testing
label_desc: Use image node_id 4 for calling a global function with example usage: test_tensor_core_batch_matmal() without return types