code: def test_tensor_core_batch_conv():
      # The sizes of inputs and filters
      batch_size = 32
      height = 14
      width = 14
      in_channels = 32
      out_channels = 64
      kernel_h = 3
      kernel_w = 3
      pad_h = 1
      pad_w = 1
      stride_h = 1
      stride_w = 1
      block_size = 16
  
      block_row_warps = 2
      block_col_warps = 4
      warp_row_tiles = 4
      warp_col_tiles = 2
      warp_size = 32
      chunk = 2
  
      # Input feature map: (N, H, W, IC, n, ic)
      data_shape = (
          batch_size // block_size,
          height,
          width,
          in_channels // block_size,
          block_size,
          block_size,
      )
      # Kernel: (H, W, IC, OC, ic, oc)
      kernel_shape = (
          kernel_h,
          kernel_w,
          in_channels // block_size,
          out_channels // block_size,
          block_size,
          block_size,
      )
  
      # Output feature map: (N, H, W, OC, n, oc)
      output_shape = (
          batch_size // block_size,
          height,
          width,
          out_channels // block_size,
          block_size,
          block_size,
      )
  
      assert batch_size % block_size == 0
      assert in_channels % block_size == 0
      assert out_channels % block_size == 0
  
      kh = te.reduce_axis((0, kernel_h), name="kh")
      kw = te.reduce_axis((0, kernel_w), name="kw")
      ic = te.reduce_axis((0, in_channels // block_size), name="ic")
      ii = te.reduce_axis((0, block_size), name="ii")
  
      # Algorithm
      A = te.placeholder(data_shape, name="A", dtype="float16")
      W = te.placeholder(kernel_shape, name="W", dtype="float16")
      Apad = te.compute(
          (
              batch_size // block_size,
              height + 2 * pad_h,
              width + 2 * pad_w,
              in_channels // block_size,
              block_size,
              block_size,
          ),
          lambda n, h, w, i, nn, ii: tvm.tir.if_then_else(
              tvm.tir.all(
                  h >= pad_h,
                  h - pad_h < height,
                  w >= pad_w,
                  w - pad_w < width,
              ),
              A[n, h - pad_h, w - pad_w, i, nn, ii],
              tvm.tir.const(0.0, "float16"),
          ),
          name="Apad",
      )
      Conv = te.compute(
          output_shape,
          lambda n, h, w, o, nn, oo: te.sum(
              Apad[
                  n, h * stride_h + kh, w * stride_w + kw, ic, nn, ii
              ].astype("float32")
              * W[kh, kw, ic, o, ii, oo].astype("float32"),
              axis=[ic, kh, kw, ii],
          ),
          name="Conv",
      )
  
      s = te.create_schedule(Conv.op)
      s[Apad].compute_inline()
  
      AS = s.cache_read(Apad, "shared", [Conv])
      WS = s.cache_read(W, "shared", [Conv])
      AF = s.cache_read(AS, "wmma.matrix_a", [Conv])
      WF = s.cache_read(WS, "wmma.matrix_b", [Conv])
      ConvF = s.cache_write(Conv, "wmma.accumulator")
  
      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")
  
      nc, hc, wc, oc, nnc, ooc = Conv.op.axis
      block_k = s[Conv].fuse(hc, wc)
      s[Conv].bind(block_k, block_z)
      nc, nci = s[Conv].split(nc, factor=warp_row_tiles)
      block_i, nc = s[Conv].split(nc, factor=block_row_warps)
      oc, oci = s[Conv].split(oc, factor=warp_col_tiles)
      block_j, oc = s[Conv].split(oc, factor=block_col_warps)
      s[Conv].reorder(
          block_k, block_i, block_j, nc, oc, nci, oci, nnc, ooc
      )
      s[Conv].bind(block_i, block_x)
      s[Conv].bind(block_j, block_y)
      s[Conv].bind(nc, thread_y)
      s[Conv].bind(oc, thread_z)
  
      s[ConvF].compute_at(s[Conv], oc)
      n, h, w, o, nnf, oof = ConvF.op.axis
      ko, ki = s[ConvF].split(ic, factor=chunk)
      s[ConvF].reorder(ko, kh, ki, kw, n, o, nnf, oof, ii)
  
      s[AF].compute_at(s[ConvF], kw)
      s[WF].compute_at(s[ConvF], kw)
  
      s[WS].compute_at(s[ConvF], kh)
      s[AS].compute_at(s[ConvF], kh)
  
      n, h, w, i, nn, ii = AS.op.axis
      tx, xo = s[AS].split(n, nparts=block_row_warps)
      ty, yo = s[AS].split(xo, nparts=block_col_warps)
      t = s[AS].fuse(nn, ii)
      to, ti = s[AS].split(t, factor=warp_size)
      s[AS].bind(tx, thread_y)
      s[AS].bind(ty, thread_z)
      s[AS].bind(ti, thread_x)
  
      kh, kw, ic, o, ii, oo = WS.op.axis
      tx, xo = s[WS].split(o, nparts=block_row_warps)
      ty, yo = s[WS].split(xo, nparts=block_col_warps)
      t = s[WS].fuse(ii, oo)
      to, ti = s[WS].split(t, nparts=warp_size)
      s[WS].bind(tx, thread_y)
      s[WS].bind(ty, thread_z)
      s[WS].bind(to, thread_x)
      s[WS].vectorize(ti)
  
      s[AF].tensorize(
          AF.op.axis[-2],
          intrin_wmma_load_matrix((16, 16, 16), "wmma.matrix_a"),
      )
      s[WF].tensorize(
          WF.op.axis[-2],
          intrin_wmma_load_matrix((16, 16, 16), "wmma.matrix_b"),
      )
      s[Conv].tensorize(nnc, intrin_wmma_store_matrix((16, 16, 16)))
      s[ConvF].tensorize(nnf, intrin_wmma_gemm((16, 16, 16)))
  
      func = tvm.build(s, [A, W, Conv], "cuda")
  
      dev = tvm.cuda(0)
      a_np = np.random.uniform(size=data_shape).astype(A.dtype)
      w_np = np.random.uniform(size=kernel_shape).astype(W.dtype)
      a = tvm.nd.array(a_np, dev)
      w = tvm.nd.array(w_np, dev)
      c = tvm.nd.array(np.zeros(output_shape, dtype=Conv.dtype), dev)
      evaluator = func.time_evaluator(func.entry_name, dev, number=3)
      print(
          "conv2d with tensor core: %f ms"
          % (evaluator(a, w, c).mean * 1e3)
      )
  
      if VERIFY:
          func(a, w, c)
          a_np = a_np.transpose(0, 4, 1, 2, 3, 5).reshape(
              batch_size, height, width, in_channels
          )
          w_np = w_np.transpose(0, 1, 2, 4, 3, 5).reshape(
              kernel_h, kernel_w, in_channels, out_channels
          )
          c_np = (
              c.numpy()
              .transpose((0, 4, 1, 2, 3, 5))
              .reshape(batch_size, height, width, out_channels)
          )
          c_std = conv2d_nhwc_python(
              a_np.astype(Conv.dtype),
              w_np.astype(Conv.dtype),
              (stride_h, stride_w),
              (pad_h, pad_w),
          ).astype(Conv.dtype)
          np.testing.assert_allclose(c_np, c_std, 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 5 for calling a global function with example usage: test_tensor_core_batch_conv() without return types