Skip to content

[Feature] mte_l1_l0a / mte_l1_l0b 需增加 start-position 参数以支持 TEXTRACT 偏移提取 #465

Description

@bingmeiyou

Summary

背景

C++ Cube API 的 TExtractToACompact(dst, src, indexRow, indexCol, validRow, validCol) 支持从源 tile 的偏移位置开始提取子 tile,indexRowindexCol 传到底层硬件指令的 start-position 配置。

当前 PTO-ISA 的 Cube wrapper op mte_l1_l0a / mte_l1_l0b 没有 start-position 参数,只有 source, destination, m, k/n, transpose。在 VPTOExpandWrapperOps.cppExpandLeftLoadPattern/ExpandRightLoadPattern 中,mte_l1_l0a/l0b expand 为 load_cbuf_to_ca/cb 时,deriveLoadCbufToCaControl 函数硬编码 mStart=0, kStart=0

// VPTOExpandWrapperOps.cpp:772
return LoadCbufToCbControl{zero, zero, mStep, kStep, stride, stride};
//                      mStart=0  kStart=0

然而,硬件层指令 load_cbuf_to_ca 已经有 m_start, k_start 参数(VPTOOps.td:1142-1164),这意味着硬件本身支持 start-position,只是 wrapper op 没有把能力暴露出来。

问题

由于 mte_l1_l0a/l0b 无法传递 start-position,使用 TileLang DSL 编写的 TEXTRACT 模板无法实现与 C++ TExtractToA/TExtractToATransCompact 对等的偏移提取功能。当用户传入 indexRow/indexCol != 0 时,偏移信息被静默丢弃,导致提取结果错误。

建议方案

mte_l1_l0amte_l1_l0b 增加 start_row / start_col keyword 参数(与 transpose keyword 采用相同的设计模式),默认值 0。具体改动:

  1. VPTOOps.tdMteL1L0aOp / MteL1L0bOp 增加 Optional<I64>:$start_rowOptional<I64>:$start_col(默认值 0)
  2. semantic.py_analyze_mte_stage_load 增加 start_row/start_col keyword 解析
  3. frontend_ast.pymte_l1_l0a/l0b 的 allowed keywords 增加 start_row, start_col
  4. VPTOExpandWrapperOps.cppderiveLoadCbufToCaControl / deriveLoadCbufToCbControl 接收 start_row/start_col 参数,传给 LoadCbufToCaOp / LoadCbufToCbOpm_start/k_start,替代当前硬编码的 zero

影响范围

  • 仅影响 mte_l1_l0a / mte_l1_l0b 两个 wrapper op 及其 expand 逻辑
  • start_row/start_col 默认值为 0,不影响现有使用 mte_l1_l0a/l0b 的其他模板(向后兼容)
  • Fix-pipe 路径 mte_l0c_l1 不涉及此问题(C++ FP 实现也不使用 indexRow/indexCol 做偏移,硬件 fix-pipe 是 whole-tile extraction)

Motivation / use case

同上

Proposed API / behavior

No response

Alternatives considered

No response

Additional context

No response

Metadata

Metadata

Assignees

No one assigned

    Labels

    enhancementNew feature or request

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions