[Bug] PTODSL's `pto.mte_gm_ub` does not handle tail block correctly when `len_burst` is not 32-byte aligned

[liggest]

Component

PTODSL

Description

PTODSL 中，pto.mte_gm_ub 的 len_burst 不是 32 字节整数倍且不传 pad 时，对于尾块未与 32 字节对齐的部分没有精确处理，实际仍按照 32 字节对齐的模式写入 UB

例如：

• len_burst=72，实际写入 96 字节
• len_burst=76，实际写入 96 字节
• len_burst=96 时行为正常

Reproduction (minimal)

from ptodsl import pto


VL = 64


@pto.jit(
    name="mte_gm_ub_unaligned",
    kernel_kind="vector",
    target="a5",
    backend="vpto",
    mode="explicit",
)
def kernel(
    input_ptr: pto.ptr(pto.i32, "gm"),
    output_ptr: pto.ptr(pto.i32, "gm"),
    valid_elems: pto.i32,
):
    ub = pto.castptr(pto.i64(0), pto.ptr(pto.i32, "ub"))

    # Clear UB so writes beyond len_burst are directly observable.
    with pto.vecscope():
        full = pto.pset_b32("PAT_ALL")
        zero = pto.vdup(pto.const(0, dtype=pto.i32), full)
        pto.vsts(zero, ub, pto.const(0), full, dist="NORM_B32")

    pto.set_flag("V", "MTE2", event_id=0)
    pto.wait_flag("V", "MTE2", event_id=0)

    valid_bytes = valid_elems * 4
    pto.mte_gm_ub(
        input_ptr,
        ub,
        0,
        valid_bytes,
        nburst=(1, valid_bytes, valid_bytes),
    )

    pto.set_flag("MTE2", "MTE3", event_id=0)
    pto.wait_flag("MTE2", "MTE3", event_id=0)
    pto.mte_ub_gm(
        ub,
        output_ptr,
        VL * 4,
        nburst=(1, 0, 0),
    )


def main():
    import torch
    # import torch_npu

    # torch.npu.config.allow_internal_format = False
    # torch_npu.npu.set_compile_mode(jit_compile=False)
    # torch.npu.set_device("npu:0")

    compiled = kernel.compile()
    stream = torch.npu.current_stream()._as_parameter_

    # GM is all ones; UB is cleared to zero inside the kernel.
    source = torch.ones(VL, dtype=torch.int32, device="npu")

    for valid_elems in (18, 19, 24):
        output = torch.empty(VL, dtype=torch.int32, device="npu")

        compiled[1, stream](
            source.data_ptr(),
            output.data_ptr(),
            valid_elems,
        )
        torch.npu.synchronize()

        result = output.cpu()
        changed_lanes = int(result.sum())
        tail = result[valid_elems : valid_elems + 8].tolist()

        print(
            f"valid={valid_elems}, "
            f"requested_bytes={valid_elems * 4}, "
            f"changed_lanes={changed_lanes}, "
            f"tail={tail}, "
            f"over_copied={changed_lanes - valid_elems}"
        )


if __name__ == "__main__":
    main()

Expected behavior

pto.mte_gm_ub 只搬运 len_burst 指定数量的字节：

valid=18, requested_bytes=72, changed_lanes=18
valid=19, requested_bytes=76, changed_lanes=19
valid=24, requested_bytes=96, changed_lanes=24

len_burst 之外的 UB 内容应保持为搬运前的值

Actual behavior / error logs

valid=18, requested_bytes=72, changed_lanes=24, tail=[1, 1, 1, 1, 1, 1, 0, 0], over_copied=6
valid=19, requested_bytes=76, changed_lanes=24, tail=[1, 1, 1, 1, 1, 0, 0, 0], over_copied=5
valid=24, requested_bytes=96, changed_lanes=24, tail=[0, 0, 0, 0, 0, 0, 0, 0], over_copied=0

不管 valid_elems 是 18、19 还是 24，pto.mte_ub_gm 都从 GM 搬运了 24 * 4 = 96 字节到 UB

Git commit

bdcb319

Host platform

Linux (x86_64)

Target Ascend arch (if relevant)

a5

PTOAS build level (if relevant)

None