code: def get_template_op(**kargs):
      batch = op_attributes["B"]
      M = op_attributes["N"]
      K = op_attributes["K"]
      N = op_attributes["M"]
      pose = op_attributes["P"]
  
      if pose == "NN":
          A = tvm.placeholder((batch, M, K), name="A", dtype="float32")
          B = tvm.placeholder((batch, K, N), name="B", dtype="float32")
          k = tvm.reduce_axis((0, K), name="k")
          C = tvm.compute(
              (batch, M, N),
              lambda b, i, j: tvm.sum(A[b, i, k] * B[b, k, j], axis=k),
              name="C",
          )
      elif pose == "NT":
          A = tvm.placeholder((batch, M, K), name="A", dtype="float32")
          B = tvm.placeholder((batch, N, K), name="B", dtype="float32")
          k = tvm.reduce_axis((0, K), name="k")
          C = tvm.compute(
              (batch, M, N),
              lambda b, i, j: tvm.sum(A[b, i, k] * B[b, j, k], axis=k),
              name="C",
          )
      elif pose == "TN":
          A = tvm.placeholder((batch, K, M), name="A", dtype="float32")
          B = tvm.placeholder((batch, K, N), name="B", dtype="float32")
          k = tvm.reduce_axis((0, K), name="k")
          C = tvm.compute(
              (batch, M, N),
              lambda b, i, j: tvm.sum(A[b, k, i] * B[b, k, j], axis=k),
              name="C",
          )
      elif pose == "TT":
          A = tvm.placeholder((batch, K, M), name="A", dtype="float32")
          B = tvm.placeholder((batch, N, K), name="B", dtype="float32")
          k = tvm.reduce_axis((0, K), name="k")
          C = tvm.compute(
              (batch, M, N),
              lambda b, i, j: tvm.sum(A[b, k, i] * B[b, j, k], axis=k),
              name="C",
          )
      else:
          raise
  
      cfg = autotvm.get_config()
      s = tvm.create_schedule(C.op)
      AA = s.cache_read(A, "shared", [C])
      AL = s.cache_read(AA, "local", [C])
      BB = s.cache_read(B, "shared", [C])
      BL = s.cache_read(BB, "local", [C])
      CC = s.cache_write(C, "local")
  
      b, y, x = C.op.axis
      k = CC.op.reduce_axis[0]
  
      cfg.define_split("B", cfg.axis(b), num_outputs=2)
      bo, bi = cfg["B"].apply(s, C, b)
  
      cfg.define_split("K", cfg.axis(k), num_outputs=3)
      ko, kt, ki = cfg["K"].apply(s, CC, k)
  
      block_x = tvm.thread_axis("blockIdx.x")
      block_y = tvm.thread_axis("blockIdx.y")
      block_z = tvm.thread_axis("blockIdx.z")
      thread_x = tvm.thread_axis("threadIdx.x")
      thread_y = tvm.thread_axis("threadIdx.y")
      thread_z = tvm.thread_axis("threadIdx.z")
  
      cfg.define_split("X", cfg.axis(y), num_outputs=4)
      cfg.define_split("Y", cfg.axis(x), num_outputs=4)
  
      by, tyz, ty, yi = cfg["X"].apply(s, C, y)
      bx, txz, tx, xi = cfg["Y"].apply(s, C, x)
  
      s[C].bind(bo, block_z)
      s[C].bind(by, block_y)
      s[C].bind(bx, block_x)
      s[C].bind(tyz, tvm.thread_axis("vthread"))
      s[C].bind(txz, tvm.thread_axis("vthread"))
      s[C].bind(bi, thread_z)
      s[C].bind(ty, thread_y)
      s[C].bind(tx, thread_x)
      s[C].reorder(by, bx, tyz, txz, ty, tx, yi, xi)
  
      s[CC].compute_at(s[C], tx)
  
      bo, yo, xo = CC.op.axis
      s[CC].reorder(ko, kt, yo, xo, ki)
      s[CC].unroll(kt)
  
      for stage in [AL, BL]:
          s[stage].compute_at(s[CC], kt)
          s[stage].double_buffer()
  
      for stage in [AA, BB]:
          s[stage].compute_at(s[CC], ko)
  
          fused = s[stage].fuse(*s[stage].op.axis)
          ty, tx = s[stage].split(fused, nparts=cfg["X"].size[2])
          tx, xi = s[stage].split(tx, nparts=cfg["Y"].size[2])
          _, xi = s[stage].split(xi, factor=4)
  
          s[stage].bind(ty, thread_y)
          s[stage].bind(tx, thread_x)
          s[stage].vectorize(xi)
          s[stage].double_buffer()
  
      cfg.add_flop(batch * M * K * N * 2.0)
      return s, [A, B, C]
  
file_path: nni/examples/trials/systems_auto_tuning/opevo/src/templates/batch_matmul.py
imports: import numpy
  import tvm
  import logging
  import sys, time, subprocess
  from tvm import autotvm
  import topi
  import json
  from topi.util import get_const_tuple
  import os
label_desc: Use image node_id 1 for calling a global function with example usage: get_template_op() and returns: s,  list