code: def test_gemm(N, L, M, dtype, layout):
      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)
      else:
          print("Unsupported layout:", layout)
          sys.exit(1)
      A = te.placeholder(shape_a, name="A", dtype=dtype)
      B = te.placeholder(shape_b, name="B", dtype=dtype)
      C = matmul_nn(A, B, L, dtype, layout)
  
      s = te.create_schedule(C.op)
      y, x = s[C].op.axis
      k = s[C].op.reduce_axis[0]
  
      # storage_align params
      factor = 16
      offset = 8
      if dtype == "int8":
          factor = 32
          offset = 16
      elif dtype == "int4":
          factor = 64
          offset = 32
      elif dtype == "int1":
          factor = 256
          offset = 128
  
      # create cache stages
      AA = s.cache_read(A, "shared", [C])
      if layout == "NN" or layout == "TN":
          s[AA].storage_align(AA.op.axis[0], factor, offset)
      AL = s.cache_read(AA, "local", [C])
      BB = s.cache_read(B, "shared", [C])
      if layout == "TT" or layout == "NT":
          s[BB].storage_align(BB.op.axis[0], factor, offset)
      BL = s.cache_read(BB, "local", [C])
      CL = s.cache_write(C, "local")
  
      # autotvm search space definition
      cfg = autotvm.get_config()
  
      cfg.define_knob("bx", [2, 4, 8])
      cfg.define_knob("by", [8, 16, 32, 64])
      cfg.define_knob("step_k", [1, 2, 4, 8, 16, 32])
      cfg.define_knob("v", [4, 8, 16, 32])
      by = cfg["by"].val
      bx = cfg["bx"].val
      step_k = cfg["step_k"].val
      v = cfg["v"].val
  
      # thread tile
      TX = 8
      TY = 1
      if dtype == "int4" or dtype == "int1":
          TX = 2
      # warp tile
      warp_tile_m = (
          16  # it could also be 8 or 32 on CUDA version >= 10.0
      )
      warp_tile_k = 16  # it must be 16 for fp16/int8 data type
      if dtype == "int4":
          warp_tile_m = 8
          warp_tile_k = 32
      elif dtype == "int1":
          warp_tile_m = 8
          warp_tile_k = 128
      # block tile
      tile_x = bx * TX
      tile_y = by * TY
  
      yo, ty = s[C].split(y, tile_y)
      ty, yi = s[C].split(ty, TY)
  
      # schedule for C stage
      xo, xi = s[C].split(x, tile_x)
      WX = min(warp_tile_m, tile_x)
      tz, xi = s[C].split(xi, WX)
      tx, xi = s[C].split(xi, TX)
      s[C].reorder(yo, xo, tz, ty, tx, yi, xi)
      s[C].bind(yo, te.thread_axis("blockIdx.y"))
      s[C].bind(xo, te.thread_axis("blockIdx.x"))
      s[C].bind(ty, te.thread_axis("threadIdx.y"))
      s[C].bind(tz, te.thread_axis("threadIdx.z"))
      s[C].bind(tx, te.thread_axis("threadIdx.x"))
  
      # schedule for CL stage
      ko, ki = s[CL].split(k, step_k * warp_tile_k)
      kl, ki = s[CL].split(ki, warp_tile_k)
      s[CL].compute_at(s[C], tx)
      yo, xo = CL.op.axis
      s[CL].reorder(ko, kl, ki, yo, xo)
  
      # schedule for AA stage
      s[AA].compute_at(s[CL], ko)
      xo, xi = s[AA].split(s[AA].op.axis[1], factor=bx * v)
      tz, tx = s[AA].split(xi, factor=(WX // TX) * v)
      tx, vec = s[AA].split(tx, factor=v)
      fused = s[AA].fuse(s[AA].op.axis[0], xo)
      _, ty = s[AA].split(fused, factor=by)
      s[AA].bind(ty, te.thread_axis("threadIdx.y"))
      s[AA].bind(tz, te.thread_axis("threadIdx.z"))
      s[AA].bind(tx, te.thread_axis("threadIdx.x"))
      # vectorization is very important for float16/int8 inputs
      s[AA].vectorize(vec)
  
      # schedule for BB stage
      s[BB].compute_at(s[CL], ko)
      xo, xi = s[BB].split(s[BB].op.axis[1], factor=bx * v)
      tz, tx = s[BB].split(xi, factor=(WX // TX) * v)
      tx, vec = s[BB].split(tx, factor=v)
      fused = s[BB].fuse(s[BB].op.axis[0], xo)
      _, ty = s[BB].split(fused, factor=by)
      s[BB].bind(ty, te.thread_axis("threadIdx.y"))
      s[BB].bind(tz, te.thread_axis("threadIdx.z"))
      s[BB].bind(tx, te.thread_axis("threadIdx.x"))
      s[BB].vectorize(vec)
  
      s[AL].compute_at(s[CL], kl)
      s[BL].compute_at(s[CL], kl)
  
      # set the 'tensor_core' pragma for tensorcore codegen
      s[CL].pragma(ko, "tensor_core")
  
      return s, [A, B, C]
  
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 2 for calling a global function with example usage: test_gemm(N, L, M, dtype, layout) and returns: s,  list