Function ldmatrix_impl

python/tvm/s_tir/tensor_intrin/cuda.py:184–220

(warp_handle: T.handle, shared_handle: T.handle)

    @T.prim_func(s_tir=True)
    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.sblock("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_legacy(
                        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