code: def get_mma_store_intrin(
      dtype, local_size, scope="global", use_mma_store_intrinic=True
  ):
      # Assume M = N = 16
      index_map = shared_16x16_to_ldmatrix_32x8_layout
      index_map_rev = ldmatrix_32x8_to_shared_16x16_layout
  
      @T.prim_func
      def mma_store_desc(a: T.handle, c: T.handle) -> None:
          C_warp = T.match_buffer(
              a, [WARP_SIZE, local_size], dtype=dtype, scope="warp"
          )
          C = T.match_buffer(
              c, [M_DIM, N_DIM], dtype=dtype, scope=scope
          )
  
          with T.block("root"):
              T.reads(C_warp[0:WARP_SIZE, 0:local_size])
              T.writes(C[0:M_DIM, 0:N_DIM])
              for i0, i1 in T.grid(M_DIM, N_DIM):
                  with T.block("C_warp"):
                      v0, v1 = T.axis.remap("SS", [i0, i1])
                      thread_id, local_id = T.meta_var(
                          index_map(v0, v1)
                      )
                      T.reads(C_warp[thread_id, local_id])
                      T.writes(C[v0, v1])
                      C[v0, v1] = C_warp[thread_id, local_id]
  
      if use_mma_store_intrinic:
  
          @T.prim_func
          def mma_store_impl(a: T.handle, c: T.handle) -> None:
              s0 = T.int32()
              s1 = T.int32()
  
              C_warp = T.match_buffer(
                  a,
                  [WARP_SIZE, local_size],
                  dtype=dtype,
                  scope="warp",
                  offset_factor=1,
              )
              C = T.match_buffer(
                  c,
                  [M_DIM, N_DIM],
                  dtype=dtype,
                  scope=scope,
                  offset_factor=1,
                  strides=[s0, s1],
              )
  
              with T.block("root"):
                  T.reads(C_warp[0:WARP_SIZE, 0:local_size])
                  T.writes(C[0:M_DIM, 0:N_DIM])
  
                  for tx in T.thread_binding(
                      0, WARP_SIZE, "threadIdx.x"
                  ):
                      T.evaluate(
                          T.mma_store(
                              M_DIM,
                              N_DIM,
                              C.access_ptr("w"),
                              C_warp.data,
                              C_warp.elem_offset,
                              s0,
                              dtype=dtype,
                          )
                      )
  
      else:
  
          @T.prim_func
          def mma_store_impl(a: T.handle, c: T.handle) -> None:
              s0 = T.int32()
              s1 = T.int32()
  
              C_warp = T.match_buffer(
                  a,
                  [WARP_SIZE, local_size],
                  dtype=dtype,
                  scope="warp",
                  offset_factor=1,
              )
              C = T.match_buffer(
                  c,
                  [M_DIM, N_DIM],
                  dtype=dtype,
                  scope=scope,
                  offset_factor=1,
                  strides=[s0, s1],
              )
  
              with T.block("root"):
                  T.reads(C_warp[0:WARP_SIZE, 0:local_size])
                  T.writes(C[0:M_DIM, 0:N_DIM])
  
                  for tx in T.thread_binding(
                      0, WARP_SIZE, "threadIdx.x"
                  ):
                      for local_id in T.serial(local_size):
                          row, col = T.meta_var(
                              index_map_rev(tx, local_id)
                          )
                          C[row, col] = C_warp[tx, local_id]
  
      return mma_store_desc, mma_store_impl
  
imports: from typing import Dict, Optional, Tuple, Literal
  from tvm._ffi import register_func
  from tvm.runtime import convert
  from tvm.script import tir
  from tvm.tir.function import PrimFunc
  from tvm.tir import Cast, IntImm, TensorIntrin
label_desc: Use image node_id 9 for calling a global function with example usage: get_mma_store_intrin(dtype, local_size, scope, use_mma_store_intrinic) and returns: mma_store_desc,  mma_store_impl