code: def matmul_nn(A, B, L, dtype="float16", layout="NN"):
      k = te.reduce_axis((0, L), name="k")
      if dtype == "float16":
          out_type = "float"
      elif dtype == "int8":
          out_type = "int"
      elif dtype == "int4" or dtype == "int1":
          out_type = "int"
      if layout == "NN":
          return te.compute(
              (N, M),
              lambda i, j: te.sum(
                  A[i, k].astype(out_type) * B[k, j].astype(out_type),
                  axis=k,
              ),
          )
      if layout == "NT":
          return te.compute(
              (N, M),
              lambda i, j: te.sum(
                  A[k, i].astype(out_type) * B[k, j].astype(out_type),
                  axis=k,
              ),
          )
      if layout == "TN":
          return te.compute(
              (N, M),
              lambda i, j: te.sum(
                  A[i, k].astype(out_type) * B[j, k].astype(out_type),
                  axis=k,
              ),
          )
      if layout == "TT":
          return te.compute(
              (N, M),
              lambda i, j: te.sum(
                  A[k, i].astype(out_type) * B[j, k].astype(out_type),
                  axis=k,
              ),
          )
  
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 1 for calling a global function with example usage: matmul_nn(A, B, L, dtype, layout) and returns: te, te, te, te