code: def get_ldmatrix_intrin(
      k_dim: int,
      dtype: str,
      matrix_name: Literal["A", "B"],
      transposed: bool,
      shared_scope: str = "shared",
  ):
      local_size = (M_DIM * k_dim) // WARP_SIZE
      smem_offset = None
      index_map = None
  
      if matrix_name == "A":
          transpose_in_ldmatrix = transposed
          # transpose_layout_for_ldmatrix_input: Every thread loads 8 bytes data. This determines
          # which 8 bytes every thread loads.
          # If transpose_layout_for_ldmatrix_input is False, the load pattern is
          # T0  T0  T0  T0  T16 T16 T16 T16
          # T1  T1  T1  T1  T17 T17 T17 T17
          # ...
          # T8  T8  T8  T8  T24 T24 T24 T24
          # T9  T9  T9  T9  T25 T25 T25 T25
          # ...
          # T15 T15 T15 T15 T31 T31 T31 T31
          # Otherwise, the load pattern is
          # T0  T0  T0  T0  T8  T8  T8  T8
          # T1  T1  T1  T1  T9  T9  T9  T9
          # ...
          # T7  T7  T7  T7  T15 T15 T15 T15
          # T16 T16 T16 T16 T24 T24 T24 T24
          # T17 T17 T17 T17 T25 T25 T25 T25
          # ...
          # T23 T23 T23 T23 T31 T31 T31 T31
          transpose_layout_for_ldmatrix_input = transposed
          smem_tile_row, smem_tile_col = (
              (M_DIM, k_dim) if not transposed else (k_dim, M_DIM)
          )
      else:
          assert matrix_name == "B"
          transpose_in_ldmatrix = not transposed
          transpose_layout_for_ldmatrix_input = transposed
          smem_tile_row, smem_tile_col = (
              (k_dim, N_DIM) if not transposed else (N_DIM, k_dim)
          )
  
      if k_dim == 16:
          assert dtype == "float16"
  
          index_map = shared_16x16_to_ldmatrix_32x8_layout
  
          if transpose_layout_for_ldmatrix_input:
              smem_offset = (
                  lambda tx, stride: stride * 8 * (tx // HALF_WARP_expr)
                  + stride * (tx % 8)
                  + 8 * ((tx % HALF_WARP_expr) // 8)
              )
          else:
              smem_offset = lambda tx, stride: stride * (
                  tx % HALF_WARP_expr
              ) + 8 * (tx // HALF_WARP_expr)
      else:
          # TODO(yixin): Support TN and TT matmul for int8
          assert (
              matrix_name == "B" or not transposed
          ), "Now only B matrix can be transposed for int8 matmul"
          assert (
              k_dim == 32 and dtype == "int8"
          ), "Only k_dim == 16 (float16) or k_dim == 32 (int8) supported for now"
  
          if matrix_name == "B" and not transposed:
              index_map = shared_32x16_to_ldmatrix_32x16_layout
              # A dummy offset, ldmatrix cannot be used for int8 + trans case.
              # We still use the ldmatrix intrinsic, but lower it to a manual loop in the codegen.
              # Only the stride information is required.
              smem_offset = lambda _, stride: stride
          elif matrix_name == "B" and transposed:
              index_map = shared_16x32_to_ldmatrix_32x16_layout
              smem_offset = (
                  lambda tx, stride: stride * 8 * (tx // HALF_WARP_expr)
                  + (tx % 8) * stride
                  + 16 * ((tx % HALF_WARP_expr) // 8)
              )
          else:  # A, not transposed
              index_map = shared_16x32_to_ldmatrix_32x16_layout
              smem_offset = lambda tx, stride: stride * (
                  tx % 16
              ) + 16 * (tx // 16)
  
      offset_factor = smem_tile_col
  
      @T.prim_func
      def ldmatrix_desc(
          warp_handle: T.handle, shared_handle: T.handle
      ) -> None:
          shared = T.match_buffer(
              shared_handle,
              (smem_tile_row, smem_tile_col),
              dtype,
              align=64,
              offset_factor=offset_factor,
              scope=shared_scope,
          )
          warp = T.match_buffer(
              warp_handle,
              (WARP_SIZE, local_size),
              dtype,
              align=64,
              offset_factor=offset_factor,
              scope="warp",
          )
  
          with T.block("root"):
              T.reads(shared[0:smem_tile_row, 0:smem_tile_col])
              T.writes(warp[0:WARP_SIZE, 0:local_size])
  
              for ax0, ax1 in T.grid(smem_tile_row, smem_tile_col):
                  with T.block("shared_warp"):
                      v0, v1 = T.axis.remap("SS", [ax0, ax1])
                      T.reads(shared[v0, v1])
  
                      thread_id, local_id = T.meta_var(
                          index_map(v0, v1)
                      )
                      T.writes(warp[thread_id, local_id])
                      warp[thread_id, local_id] = shared[v0, v1]
  
      @T.prim_func
      def ldmatrix_impl(
          warp_handle: T.handle, shared_handle: T.handle
      ) -> None:
          s0 = T.int32()
          s1 = T.int32()
          shared = T.match_buffer(
              shared_handle,
              (smem_tile_row, smem_tile_col),
              dtype,
              align=64,
              offset_factor=offset_factor,
              scope=shared_scope,
              strides=[s0, s1],
          )
          warp = T.match_buffer(
              warp_handle,
              (WARP_SIZE, local_size),
              dtype,
              align=64,
              offset_factor=offset_factor,
              scope="warp",
          )
  
          with T.block("root"):
              T.reads(shared[0:smem_tile_row, 0:smem_tile_col])
              T.writes(warp[0:WARP_SIZE, 0:local_size])
              for tx in T.thread_binding(0, WARP_SIZE, "threadIdx.x"):
                  T.evaluate(
                      T.ptx_ldmatrix(
                          transpose_in_ldmatrix,
                          4,  # Always load 4 matrices
                          ".b16",
                          warp.data,
                          warp.elem_offset + lift(local_size) * tx,
                          shared.access_ptr("r"),
                          smem_offset(tx, s0),
                          dtype=dtype,
                      )
                  )
  
      return ldmatrix_desc, ldmatrix_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 6 for calling a global function with example usage: get_ldmatrix_intrin(k_dim, dtype, matrix_name, transposed, shared_scope) and returns: ldmatrix_desc,  ldmatrix_impl