Function impl

python/tvm/s_tir/tensor_intrin/arm_cpu.py:271–340

()

                    A_t[v_k, v_m] = A[v_m, v_k]

    def impl():
        sub_tile_count = 4

        with IRBuilder() as ib:
            with build_prim_func():
                a = T.arg("a", T.handle())
                a_t = T.arg("a_t", T.handle())

                A = T.match_buffer(
                    a, (SVF2, SVF2), "float32", offset_factor=1, strides=[T.int32(), 1]
                )
                A_t = T.match_buffer(
                    a_t,
                    (SVF2, SVF2),
                    "float32",
                    offset_factor=1,
                    strides=[T.int32(), 1],
                )

                with T.sblock("root"):
                    T.reads(A[0:SVF2, 0:SVF2])
                    T.writes(A_t[0:SVF2, 0:SVF2])

                    # Load rows of the input matrix
                    with T.serial(0, SVF) as slice_idx:
                        for sub_tile_idx in range(0, sub_tile_count):
                            row_offset = SVF if sub_tile_idx >= (sub_tile_count // 2) else 0
                            col_offset = SVF if sub_tile_idx % 2 else 0
                            offset = (slice_idx + row_offset) * A.strides[0] + col_offset

                            input_ptr = A.access_ptr("r", offset=offset)
                            sub_tile = T.int32(sub_tile_idx)
                            predicate = _create_active_lane_mask(
                                A, (row_offset + slice_idx, col_offset), cols
                            )
                            T.evaluate(
                                T.call_llvm_intrin(
                                    "void",
                                    "llvm.aarch64.sme.ld1w.horiz",
                                    predicate,
                                    input_ptr,
                                    sub_tile,
                                    slice_idx,
                                )
                            )

                    # Store columns to the output matrix
                    with T.serial(0, SVF) as slice_idx:
                        for sub_tile_idx in range(0, sub_tile_count):
                            col_offset = SVF if sub_tile_idx >= (sub_tile_count // 2) else 0
                            row_offset = SVF if sub_tile_idx % 2 else 0
                            offset = (slice_idx + row_offset) * A_t.strides[0] + col_offset

                            output_ptr = A_t.access_ptr("w", offset=offset)
                            sub_tile = T.int32(sub_tile_idx)
                            predicate = _create_active_lane_mask(
                                A_t, (row_offset + slice_idx, col_offset), rows
                            )