code: def matmul(M: int, N: int, K: int, dtype: str):
      @T.prim_func
      def main(a: T.handle, b: T.handle, c: T.handle) -> None:
          A = T.match_buffer(a, [M, K], dtype=dtype)
          B = T.match_buffer(b, [N, K], dtype=dtype)
          C = T.match_buffer(c, [M, N], dtype=dtype)
          for i, j, k in T.grid(M, N, K):
              with T.block():
                  vi, vj, vk = T.axis.remap("SSR", [i, j, k])
                  with T.init():
                      C[vi, vj] = T.float32(0)
                  C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vj, vk]
  
      return main
  
imports: import tempfile
  import numpy
  import torch
  import torch.nn
  import tvm
  from tvm.target.target import Target
  import tvm.testing
  from tvm.contrib.torch import as_torch
  from tvm.script import tir
label_desc: Use image node_id 1 for calling a global function with example usage: matmul(M, N, K, dtype) and returns: main