epic(tilelang-dsl): complete cube matrix tileop templates and ST coverage

[Zhendong404]

背景

cube dsl 的基础设施已经基本落地：@pto.ckernel、cube surface、VPTO bridge lowering、pto.tmatmul 的首个 cube template，以及基础 lit/ST 链路都已经打通。

相关基础工作已经在 PR #336 ：

• @pto.ckernel frontend / semantic / lowering
• cube raw ops 与 bridge lowering
• lib/TileOps/tmatmul_template.py
• test/tilelang_st/npu/a5/src/st/testcase/tmatmul/

当前缺口

1. tload/tstore 的问题本质上是 matrix MemorySpace 支持没补齐

tload/tstore 目前还没有把 vec 之外的 MemorySpace 版本实现完整补齐，尤其是 matrix/cube 需要的 mat / acc 等路径。因此在 ST 用例里，很多 matrix/cube 数据流现在还只能直接写裸 VPTO cube 指令，没法通过完整的 matrix tileop authoring 路径表达。

当前现状：

• lib/TileOps/tload_template.py / lib/TileOps/tstore_template.py 仍主要覆盖 vector 路径
• 缺少面向 matrix/cube 使用场景的完整 MemorySpace 实现、匹配与约束
• 结果就是 cube ST 里仍需要直接写 raw VPTO cube op 才能完成关键数据流

2. matrix tileop 的实现与 ST 覆盖都还不完整

目前只有 tmatmul 是这一轮新增的 cube ST，但它更接近一个 demo，而不是 matrix tileop 已经全面可用的证明。

具体来说：

• tmatmul 现有实现和测试还没有把各种 layout / dtype / valid-shape / 约束边界补完整
• docs/PTO_IR_manual.md 里还有一批 matrix op 已经列出，但还没有对应实现与 ST 收口

当前应纳入这轮范围的 matrix op 至少包括：

• TMOV2LEFT
• TMOV2BIAS
• TMOV2VEC
• TMOV2RIGHT
• TMOV2SCALE
• TMOV_FP
• TEXTRACT(CUBE)
• TEXTRACT_FP
• TINSERT(CUBE类)
• TINSERT_FP
• TLOAD.MAT
• TSTORE.MAT
• TSTORE.ACC
• TSTORE_FP
• TIMG2COL

任务拆分

• 任务1: tload/tstore

范围：

• TLOAD.MAT
• TSTORE.MAT
• TSTORE.ACC
• TSTORE_FP

把 matrix 数据搬运路径打通，让 cube/matrix authoring 不再依赖手写裸 VPTO 指令。

工作项：

• 为 tload/tstore 补齐 vec 之外的 MemorySpace 实现，重点覆盖 matrix/cube 需要的 mat / acc 路径
• 梳理 PartitionTensorView <-> tile_buf<mat/...>、tile_buf<acc/...> -> PartitionTensorView 的合法匹配规则
• 补齐 template matcher / constraint / semantic check / diagnostic，使 vector 与 cube 路径能正确分流
• 明确不同 layout、dtype、valid-shape 在 cube/matrix 路径下的支持边界
• 让 ST 可以通过 pto.tload / pto.tstore 表达 cube 数据流，而不是直接写 raw cube op
• 补齐对应 lit：正向 expand、负向诊断、vector/cube matcher 不串线
• 补齐对应 ST：
  • 独立 tload cube testcase
  • 独立 tstore cube testcase
  • 关键 GM <-> MAT、ACC -> GM 路径覆盖

交付标准：

• tload/tstore 能稳定覆盖 matrix/cube 需要的数据搬运主路径

• cube ST 不再需要依赖手写裸 VPTO cube load/store 作为主 authoring 手段

• lit/ST 能把支持边界和非法组合固定下来

• 跑通pto-isa仓上对应的st cases

• 任务2: matmul 类

把 tmatmul 家族从“能跑 demo”补成“约束和测试相对完整的一组 matrix tileop”。

范围：

• TMATMUL
• TMATMUL.ACC
• TMATMUL.BIAS
• TMATMUL.MX

工作项：

• 在现有 tmatmul cube template 基础上补齐家族化实现
• 梳理并固化各变体的输入/输出 loc、layout、dtype、valid-shape、bias/scale operand 约束
• 补齐 expand 结果与底层 cube op 契约之间的映射验证
• 增加负向诊断用例，避免看起来“能匹配”但语义上不合法的组合漏过去
• 补齐 lit/ST，覆盖普通、acc、bias、mx 等关键变体
• 对现有 tmatmul demo/ST 做增强，避免只验证最窄 happy path

交付标准：

• tmatmul 家族主要 public 变体都有可用 template 和稳定约束

• 每个变体至少有基本 lit 覆盖，关键变体有独立 ST

• 已知 layout/dtype/valid-shape 限制能通过诊断或测试明确表达

• 跑通pto-isa仓上对应的st cases

• 任务3: tgemv 类
  补齐第二组 matrix compute tileop，避免整个 matrix 能力只停留在 matmul 家族。

范围：

• TGEMV_MX
• TGEMV
• TGEMV_ACC
• TGEMV_BIAS
• TGEMV
• TGEMV.ACC
• TGEMV.BIAS

工作项：

• 为 tgemv 家族编写 cube/matrix template
• 梳理 GEMV 特有的 shape、layout、loc、dtype、bias/scale 约束
• 对接 tload/tstore 已打通的数据搬运路径，确保 authoring 链路闭环
• 补齐 lit：正向 expand + 负向诊断
• 补齐 ST：普通、acc、bias、mx 的关键代表场景
• 明确哪些限制与 tmatmul 可共用，哪些是 GEMV 特有约束

交付标准：

• tgemv 家族主要 public 变体具备可用 template

• tgemv 不依赖额外手写 raw cube op 才能完成基本 ST authoring

• 关键 shape/layout 约束由 lit/ST 固定下来

• 跑通pto-isa仓上对应的st cases

• 任务4: tmov bridge 类

把 cube authoring 里 L1/Mat -> L0/Left/Right/Bias/Scale 以及 Acc -> Vec/Mat 中由 tmov 承担的 bridge 语义单独收口。tmov 既是最常见的 cube bridge surface 之一，也和 tmatmul/tgemv 主数据流高度耦合，适合独立派工。

范围：

• TMOV2LEFT
• TMOV2BIAS
• TMOV2VEC
• TMOV2RIGHT
• TMOV2SCALE
• TMOV_FP

工作项：

• 明确 pto.tmov / pto.tmov.fp 在 cube 路径下的 surface 设计与 matcher 分流
• 补齐 Mat -> Left/Right/Bias/Scale、Acc -> Mat/Vec 对应的 DSL template / expand / lowering
• 梳理 fp、preQuantScalar、reluPreMode、accToVecMode 等参数在 cube 路径下的支持边界
• 为 tmov 增加独立 lit/ST，而不是只通过 matmul 间接覆盖
• 对齐 pto-isa 已有 tmov 相关 ST 目录，至少能稳定映射和跑通对应 case

交付标准：

• tmov 的 cube 相关主路径都能通过 TileLang/PTO IR authoring 表达

• L1/Mat -> L0、Acc -> Mat/Vec 的关键 bridge 场景不再依赖手写裸 VPTO 指令拼接

• 对应 lit/ST 能固定 surface 选择、参数约束和非法组合诊断

• 跑通 pto-isa 仓上这一组对应的 ST cases

• 任务5: textract cube 类

范围：

• TEXTRACT(CUBE)
• TEXTRACT(Vec-Vec)
• TEXTRACT_FP

工作项：

• 单独梳理 pto.textract 在 cube 路径下的 surface 语义，不和 tmov/tinsert 混做一个子任务
• 明确 textract 在哪些场景下承担 Mat -> Left/Right、Acc -> Mat/Vec、Vec -> Vec/Mat 的子窗口抽取职责
• 梳理 indexRow/indexCol、fp、preQuantScalar、layout、transpose/compact/mx 等关键约束
• 为 textract 补齐独立 lit/ST，覆盖正向 expand、负向诊断、参数边界和 matcher 选择
• 对齐 pto-isa 已有 textract 相关 ST 目录，补足当前上层 authoring 还没承接的部分

交付标准：

• textract 在 cube 路径下具备独立、清晰且稳定的 authoring surface

• textract 的约束和 bridge 语义由单独的 lit/ST 固定下来，而不是继续寄生在 tmov/tinsert 或 matmul demo 上

• 跑通 pto-isa 仓上这一组对应的 ST cases

• 任务6: tinsert cube 类

范围：

• TINSERT(CUBE类)
• TINSERT(VEC-VEC)
• TINSERT_FP

工作项：

• 单独梳理 pto.tinsert 在 cube 路径下的 surface 语义
• 明确 tinsert 负责的 Vec -> Mat、Acc -> Mat、Acc -> Vec 等回填/插入场景边界，不和 tmov/textract 重叠
• 梳理 indexRow/indexCol、fp、preQuantScalar、reluPreMode、accToVecMode 等关键参数在 cube 路径下的支持边界
• 为 tinsert 补齐独立 lit/ST，覆盖正向 expand、负向诊断、参数边界和 matcher 选择
• 对齐 pto-isa 已有 tinsert 相关 ST 目录，补足当前上层 authoring 还没承接的部分

交付标准：

• tinsert 在 cube 路径下具备独立、清晰且稳定的 authoring surface

• tinsert 的插入/回填语义由单独的 lit/ST 固定下来，而不是继续并入 tmov 或 textract 的大任务

• 跑通 pto-isa 仓上这一组对应的 ST cases

• 任务7: img2col/fmatrix cube 准备类

范围：

• TIMG2COL

工作项：

• 明确这些 op 在 DSL / PTO IR 中是否暴露为独立 surface，还是作为高阶 tileop 的内部 lowering 细节
• 如果需要暴露 authoring surface，补齐 template、约束、诊断和示例
• 如果不直接暴露，至少要补齐从上层 tileop 到这组 pto-isa op 的 lowering 契约测试
• 梳理它们与 tload/tmov/textract、tmatmul 之间的数据流关系，避免后续 conv 场景重复设计
• 补齐最小必要 lit/ST，验证 fmatrix / img2col 相关参数不会在 lowering 过程中丢失或错配

交付标准：

• img2col/fmatrix 相关能力在上层 IR/DSL 里有明确归属，不再处于“pto-isa 有实现，但上层没人接”的状态
• 至少有一条稳定的 authoring -> lowering -> ST 验证链路覆盖这组 op
• 跑通 pto-isa 仓上这一组对应的 ST cases

-[] 任务8: ttrans类

范围：

• TTRANS

共同约束与收口要求

任务并行时建议遵守以下约束，避免后面再统一返工：

• 新增 tileop 必须同时带最小必要 lit，关键路径要有 ST，而不是只补 template
• 负向诊断要尽量前置，避免非法组合在后端才暴露
• vector 与 cube 路径的 matcher/expand 不能互相串线
• 新增 ST 尽量按独立 testcase 目录组织，避免把 matrix 场景继续散落在其他 op 的间接覆盖里
• 最终版本的ST用例需要使用tload/tstore打通

[KurrinQu]

应该还需要TMOV/TEXTRACT指令实现L1->L0的数据搬运

[Zhendong404]

其他注意事项：

1. 遇到问题先问agent，确定是DSL或PTOAS问题再提issue
2. 提issue时描述清楚复现的环境及命令，建议直接带上PR，避免贴截图

[Zhendong404]

全局参考：关键目录和文件

PTOAS 侧（需要修改的目标项目）

lib/TileOps/                        # 模板文件（每个 op 一个 _template.py）
  tload_template.py                 # 已有，仅覆盖 vector 路径 (ND/ND, DN/DN, NZ/NZ)
  tstore_template.py                # 已有，仅覆盖 vector 路径
  tmatmul_template.py               # 已有，仅覆盖最窄 happy path (f16/bf16/f32)
  tmov_template.py                  # 已有，仅覆盖 UB2UB + ND2ND (vector 路径)
  textract_template.py              # 待创建
  tinsert_template.py               # 待创建
  tgemv_template.py                 # 待创建
  timg2col_template.py              # 待创建
  ttrans_template.py                # 待创建

test/lit/vpto/cube/                 # cube 路径 lit 测试（IR 级验证）
test/tilelang_st/npu/a5/src/st/testcase/  # ST 测试（端到端 NPU 验证）
  tmatmul/                          # 已有 demo
  tmov/                             # 已有，仅 vector 路径
  tload/                            # 已有，仅 vector 路径

pto-isa 侧（参考实现，不要修改）

include/pto/cpu/
  TLoad.hpp        # TLOAD_TILE_IMPL (plain), TLOAD_CONVTILE_IMPL (conv tile)
  TStore.hpp       # TSTORE_IMPL (plain), 多种 AtomicType/ReluPreMode/STPhase 重载
  TMatmul.hpp      # TMATMUL_IMPL, TMATMUL_ACC_IMPL, TMATMUL_BIAS_IMPL, TMATMUL_MX_IMPL, TGEMV_*
  TMov.hpp         # TMOV_IMPL (基础), 带 ReluPreMode/AccToVecMode/FpTileData 的重载
  TExtract.hpp     # TEXTRACT_IMPL (基础), 带 quant/fp/relu/mode 的重载
  TInsert.hpp      # TINSERT_IMPL (基础), 带 relu/fp/preQuantScalar 的重载
  TImg2col.hpp     # TIMG2COL_IMPL (Mat->Left 转换, conv img2col 语义)
  TTrans.hpp       # TTRANS_IMPL (dst, src, tmp), 分块转置
  SetFmatrix.hpp   # SETFMATRIX_IMPL (设置 conv feature map 元信息)

tests/costmodel/st/testcase/   # pto-isa 的 ST 用例（参考 case 选择）
tests/cpu/st/testcase/
tests/npu/a5/src/st/testcase/

---

通用实施模式（每个任务都遵循）

每个 tileop 的实施分为 4 个阶段：

阶段 A：Template 实现

• 文件：lib/TileOps/<op>_template.py
• 使用 @pto.ckernel（cube 路径）或 @pto.vkernel（vector 路径）装饰器注册
• 函数签名遵循 pto-isa 同名 *_IMPL 的参数顺序和语义
• 函数体内使用 pto DSL intrinsic（pto.mad, pto.copy_gm_to_ubuf, pto.vlds, etc.）

阶段 B：Constraints / Matcher

• 在 template 文件的 constraints 参数中实现
• 约束函数返回 bool，检查 tile_buf 的 loc/blayout/slayout/dtype/shape 是否匹配
• 参考 pto-isa 中的 static_assert 和 Check*Valid 函数
• 关键：vector 路径的 matcher 必须返回 false 当 tile loc 是 cube 专有类型（Left/Right/Acc/Mat）

阶段 C：Lit 测试

• 目录：test/lit/vpto/cube/（cube 路径）或 test/lit/vpto/（vector 路径）
• 每个 template 变体至少一个正向 expand 用例
• 每个关键的非法组合至少一个负向诊断用例
• 命名：expand_tile_op_tilelang_<op>.pto（正向）、<op>_verify_invalid.pto（负向）
• CHECK 行验证 expand 后的 IR 包含预期的底层指令

阶段 D：ST 测试

• 目录：test/tilelang_st/npu/a5/src/st/testcase/<op>/
• 使用 test/tilelang_st/script/gen_st.py 生成的标准测试框架
• 每个 case 的 cases.py 定义 dtype/shape/valid_shape/eps
• .pto 文件写完整的 cube kernel（GM->L1->L0->compute->L0->GM 全流程）

---

任务1：tload/tstore Matrix 路径补齐

参考 pto-isa 实现

• TLoad.hpp: TLOAD_TILE_IMPL (L384-L149)、TLOAD_CONVTILE_IMPL (L284-L299)
• TStore.hpp: TSTORE_IMPL (L111-L171)

关键分析

当前 tload_template.py 和 tstore_template.py 的三个模板分别覆盖：

• ND2ND（vector, BLayout=RowMajor, SLayout=NoneBox）
• DN2DN（vector, BLayout=ColMajor, SLayout=NoneBox）
• NZ2NZ（vector, BLayout=ColMajor, SLayout=RowMajor）

缺失的 matrix/cube 路径：

• TLOAD.MAT：GM→L1（ConvTile, layout=NC1HWC0/FRACTAL_Z），对应 TLOAD_CONVTILE_IMPL
• TSTORE.MAT：L1→GM（NZ→ND/DN 转换 store），对应 pto-isa 中 TSTORE_IMPL 的非 plain 路径
• TSTORE.ACC：L0C→GM（Acc tile → Global Memory），对应 mte_l0c_gm
• TSTORE_FP：带 FP 量化参数的 store

实施步骤

1a. tload MAT 模板

在 tload_template.py 中新增：

@pto.ckernel(
    target="a5",
    op="pto.tload",
    constraints=[_tload_preconditions_mat],
)
def template_tload_mat(src: pto.PartitionTensorView, dst: pto.Tile):
    # 参考 pto-isa: TLoad5HD / TLoadFractalZ
    # 使用 pto.mte_gm_l1_frac intrinsic
    # dst 类型为 pto.Tile<loc=mat, layout=NC1HWC0 或 FRACTAL_Z>
    ...

约束函数 _tload_preconditions_mat：

• dst.config.loc == pto.TileType.MAT
• dst.layout 为 NC1HWC0 或 FRACTAL_Z
• src.rank == 5（对应 GlobalTensor 的 5 维 shape/stride）

1b. tstore MAT/ACC/FP 模板

在 tstore_template.py 中新增：

@pto.ckernel(target="a5", op="pto.tstore", constraints=[...])
def template_tstore_mat(src: pto.Tile, dst: pto.PartitionTensorView): ...
    # 参考 pto-isa: TSTORE_IMPL 的 NZ→ND 路径

@pto.ckernel(target="a5", op="pto.tstore", constraints=[...])
def template_tstore_acc(src: pto.Tile, dst: pto.PartitionTensorView): ...
    # Acc → GM，参考 pto-isa: mte_l0c_gm

@pto.ckernel(target="a5", op="pto.tstore", constraints=[...])
def template_tstore_fp(src: pto.Tile, dst: pto.PartitionTensorView, fp: pto.Tile): ...
    # 带 FP scale 参数，参考 pto-isa: TSTORE_IMPL 的 FpTileData 重载

1c. 约束完善

• 确保 vector matcher（_tload_preconditions_nd2nd 等）在遇到 loc=mat 时返回 False
• 确保 cube matcher 在遇到 loc=vec 时返回 False
• 这就是 issue 原文要求的 "vector/cube matcher 不串线"

1d. Lit 测试

新建 test/lit/vpto/cube/tload_mat_expand.pto:

// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --enable-tile-op-expand --emit-vpto %s -o - | FileCheck %s
// CHECK: pto.mte_gm_l1_frac

新建负向诊断测试确保 vector/cube 不串线。

1e. ST 测试

新建独立 testcase 目录：

test/tilelang_st/npu/a5/src/st/testcase/tload_mat/
test/tilelang_st/npu/a5/src/st/testcase/tstore_mat/
test/tilelang_st/npu/a5/src/st/testcase/tstore_acc/

每个目录包含: cases.py, CMakeLists.txt, compare.py, gen_data.py, launch.cpp, main.cpp, <op>.pto

---

任务2：tmatmul 家族完善

参考 pto-isa 实现

• TMatmul.hpp:
  • TMATMUL_IMPL (L155-163) → pto.tmatmul
  • TMATMUL_ACC_IMPL (L167-177) → accumulator 输入/输出分离
  • TMATMUL_BIAS_IMPL (L179-197) → 加 bias 参数
  • TMATMUL_MX_IMPL (L221-305) → MX 精度（fp4/fp8 + scale）
  • CheckMadValid (L70-96) → dtype/layout 约束
  • CheckMadMxValid (L110-131) → MX 约束
  • CheckBiasValid (L144-152) → bias 约束
  • CheckDynamicMmad (L133-142) → shape 范围约束

当前状态

tmatmul_template.py 只有一个 template_tmatmul，只处理 (f16, f16, f32), (bf16, bf16, f32), (f32, f32, f32) 三种 dtype 组合，没有 acc/bias/MX 变体。

实施步骤

2a. 扩展 dtype 支持

根据 CheckMadValid 补充 (int8, int8, int32) 组合：

@pto.ckernel(
    target="a5",
    op="pto.tmatmul",
    dtypes=[
        (pto.f16, pto.f16, pto.f32),
        (pto.bf16, pto.bf16, pto.f32),
        (pto.f32, pto.f32, pto.f32),
        (pto.i8, pto.i8, pto.i32),   # 新增
    ],
)

2b. ACC 变体（accumulator in/out 分离）

新增 template_tmatmul_acc:

@pto.ckernel(target="a5", op="pto.tmatmul.acc", ...)
def template_tmatmul_acc(lhs, rhs, acc_in, acc_out):
    # acc_in + lhs @ rhs -> acc_out
    # 对应 TMATMUL_ACC_IMPL
    pto.mad(lhs.as_ptr(), rhs.as_ptr(), acc_out.as_ptr(), m, n, k, acc_in=acc_in.as_ptr())

2c. BIAS 变体

新增 template_tmatmul_bias:

@pto.ckernel(target="a5", op="pto.tmatmul.bias", ...)
def template_tmatmul_bias(lhs, rhs, bias, acc_out):
    # 对应 TMATMUL_BIAS_IMPL: mad 后逐列加 bias
    pto.mad(lhs.as_ptr(), rhs.as_ptr(), acc_out.as_ptr(), m, n, k)
    # bias 广播加法需要在 DSL 层或 lower 层实现

2d. MX 变体

新增 template_tmatmul_mx:

@pto.ckernel(target="a5", op="pto.tmatmul.mx", ...)
def template_tmatmul_mx(lhs, lhs_scale, rhs, rhs_scale, acc_out):
    # 对应 TMATMUL_MX_IMPL: 带 scale 矩阵的 MX 精度 matmul
    pto.mad_mx(lhs.as_ptr(), rhs.as_ptr(), lhs_scale.as_ptr(), rhs_scale.as_ptr(),
               acc_out.as_ptr(), m, n, k)

2e. 约束完善

为每个变体添加 constraint 函数，对应 pto-isa 的 CheckMadValid、CheckBiasValid、CheckMadMxValid、CheckDynamicMmad：

• lhs.loc == Left, rhs.loc == Right, acc.loc == Acc
• lhs.cols == rhs.rows (K 维度)
• lhs.rows == acc.rows, rhs.cols == acc.cols
• blayout/slayout 分形符合 cube 规约
• MX 变体：K 是 64 的倍数；fp4 额外要求 K 是 2 的倍数
• dynamic shape 范围检查：每个维度 [1, 4095]

2f. Lit 测试

为每个新变体创建独立的 lit 测试文件：

test/lit/vpto/cube/expand_tile_op_tilelang_tmatmul_acc.pto
test/lit/vpto/cube/expand_tile_op_tilelang_tmatmul_bias.pto
test/lit/vpto/cube/expand_tile_op_tilelang_tmatmul_mx.pto
test/lit/vpto/cube/mad_bias_verify_invalid.pto
test/lit/vpto/cube/mad_shape_verify_invalid.pto

2g. ST 测试

扩展 test/tilelang_st/npu/a5/src/st/testcase/tmatmul/cases.py 增加更多 case：

• 不同 dtype: s8, bf16, f16, f32
• 不同 shape: 16x16x16, 32x32x32, 128x128x128, 非对齐 shape
• acc 变体独立 testcase 目录
• bias 变体独立 testcase 目录

---

任务3：tgemv 家族

参考 pto-isa 实现

• TMatmul.hpp:
  • TGEMV_IMPL (L199-204): 当前直接 delegate 给 TMATMUL_IMPL
  • TGEMV_ACC_IMPL (L207-211)
  • TGEMV_BIAS_IMPL (L213-218)
  • TGEMV_MX_IMPL (L278-305): 直接 delegate 给 TMATMUL_MX_IMPL

关键分析

GEMV 是 GEMV（N=1 的特例）。在 pto-isa 的实现中，TGEMV 就是 TMATMUL 的 alias。但在 DSL 层，需要：

1. 通过 shape constraint 区分：rhs.valid_shape[1] == 1（即 N=1）
2. 使用 pto.mad(... disable_gemv=False) 显式走 GEMV 路径
3. 如果硬件对 GEMV 有特殊优化路径，lowering 需要反映这一点

实施步骤

3a. Template 实现

新建 tgemv_template.py：

# 基础 GEMV（等价于 TMATMUL 的 N=1 特例）
@pto.ckernel(target="a5", op="pto.tgemv", dtypes=[...])
def template_tgemv(lhs, rhs, acc):
    m, k = lhs.valid_shape
    _, n = rhs.valid_shape
    # n == 1
    pto.mad(lhs.as_ptr(), rhs.as_ptr(), acc.as_ptr(), m, n, k, disable_gemv=False)

# ACC 变体
@pto.ckernel(target="a5", op="pto.tgemv.acc", dtypes=[...])
def template_tgemv_acc(lhs, rhs, acc_in, acc_out): ...

# BIAS 变体 - 与 tmatmul.bias 相同模式
@pto.ckernel(target="a5", op="pto.tgemv.bias", dtypes=[...])
def template_tgemv_bias(lhs, rhs, bias, acc_out): ...

# MX 变体
@pto.ckernel(target="a5", op="pto.tgemv.mx", dtypes=[...])
def template_tgemv_mx(lhs, lhs_scale, rhs, rhs_scale, acc_out): ...

3b. 约束

• rhs.valid_shape[1] == 1（GEMV 的核心约束）
• 其余约束与 tmatmul 相同（loc/layout/dtype/shape 范围）

3c. Lit 测试

test/lit/vpto/cube/expand_tile_op_tilelang_tgemv.pto
test/lit/vpto/cube/expand_tile_op_tilelang_tgemv_mx.pto

3d. ST 测试

test/tilelang_st/npu/a5/src/st/testcase/tgemv/

与 tmatmul 的共享策略

GEMV 的 acc/bias/mx 变体与 TMATMUL 模式高度一致。建议：

• 先完成 tmatmul 的 acc/bias/mx 模板，然后 tgemv 直接复用相同模式
• tgemv 的约束层只需额外检查 N==1
• pto-isa 参考 ST: tests/costmodel/st/testcase/tmatmul/main.cpp 中包含 GEMV 场景的 case

---

任务4：tmov bridge 类

参考 pto-isa 实现

• TMov.hpp:
  • TMOV_IMPL (L21-51): 基础 tile-to-tile copy
  • TMOV_IMPL<Dst, Src, ReluPreMode> (L53-68): 带 ReLU
  • TMOV_IMPL<Dst, Src, FpTileData, ...> (L78-107): 带 FP 量化参数
  • TMOV_IMPL<Dst, Src, preQuantScalar, ...> (L95-106): 带预量化标量

当前状态

tmov_template.py 只覆盖 UB2UB + ND2ND（即 vector Vec→Vec 路径）。缺失的 cube bridge 路径：

• TMOV2LEFT: Mat/L1 → Left/L0A
• TMOV2BIAS: Vec/Acc → Bias/L0B（bias 广播设置）
• TMOV2VEC: Acc → Vec
• TMOV2RIGHT: Mat/L1 → Right/L0B
• TMOV2SCALE: Mat → ScaleLeft/ScaleRight
• TMOV_FP: 带 FP tile 的量化搬运

实施步骤

4a. Cube bridge 模板

在 tmov_template.py 中新增多个 @pto.ckernel 模板：

# L1(Mat) → L0A(Left)
@pto.ckernel(target="a5", op="pto.tmov", constraints=[_tmov_mat2left])
def template_tmov_mat2left(src: pto.Tile, dst: pto.Tile):
    # 对应 pto-isa: mte_l1_l0a
    rows, cols = dst.valid_shape
    pto.mte_l1_l0a(src.as_ptr(), dst.as_ptr(), rows, cols)

# L1(Mat) → L0B(Right)，需要 transpose
@pto.ckernel(target="a5", op="pto.tmov", constraints=[_tmov_mat2right])
def template_tmov_mat2right(src: pto.Tile, dst: pto.Tile):
    rows, cols = dst.valid_shape
    pto.mte_l1_l0b(src.as_ptr(), dst.as_ptr(), rows, cols, transpose=True)

# Acc → Vec
@pto.ckernel(target="a5", op="pto.tmov.acc2vec", constraints=[...])
def template_tmov_acc2vec(src: pto.Tile, dst: pto.Tile): ...

# Acc → Mat
@pto.ckernel(target="a5", op="pto.tmov.acc2mat", constraints=[...])
def template_tmov_acc2mat(src: pto.Tile, dst: pto.Tile): ...

# Vec → Bias（bias 广播）
@pto.ckernel(target="a5", op="pto.tmov.vec2bias", constraints=[...])
def template_tmov_vec2bias(src: pto.Tile, dst: pto.Tile): ...

# 带 FP tile 参数的变体
@pto.ckernel(target="a5", op="pto.tmov.fp", constraints=[...])
def template_tmov_fp(src: pto.Tile, dst: pto.Tile, fp: pto.Tile): ...

# 带 preQuantScalar 的变体
@pto.ckernel(target="a5", op="pto.tmov", constraints=[...])
def template_tmov_pre_quant(src: pto.Tile, dst: pto.Tile, pre_quant_scalar: int): ...

4b. 约束设计

每个 bridge 方向有独特的 loc 匹配规则：

def _tmov_mat2left(src, dst) -> bool:
    return (src.config.loc == pto.TileType.MAT and
            dst.config.loc == pto.TileType.LEFT and
            src.memory_space == "l1" and
            dst.memory_space == "l0a")

def _tmov_mat2right(src, dst) -> bool:
    return (src.config.loc == pto.TileType.MAT and
            dst.config.loc == pto.TileType.RIGHT)

def _tmov_acc2vec(src, dst) -> bool:
    return (src.config.loc == pto.TileType.ACC and
            dst.config.loc == pto.TileType.VEC and
            src.config.s_layout == pto.SLayout.ROW_MAJOR and
            dst.config.s_layout == pto.SLayout.NONE_BOX)

4c. 现有 vector matcher 保护

确保 _tmov_ub2ub_nd2nd_constraint 仅在 src/dst 都是 Vec + NoneBox 时返回 True，遇到 cube loc 返回 False。

4d. Lit 测试

test/lit/vpto/cube/expand_tile_op_tilelang_tmov_mat2left.pto
test/lit/vpto/cube/expand_tile_op_tilelang_tmov_mat2right.pto
test/lit/vpto/cube/expand_tile_op_tilelang_tmov_acc2vec.pto
test/lit/vpto/cube/tmov_cube_vec_verify_invalid.pto  # vector/cube matcher 不串线

4e. ST 测试

对齐 pto-isa 的 ST 目录：

• 参考 tests/costmodel/st/testcase/tmov/main.cpp（当前仅 Vec→Vec costmodel case）
• 新建 test/tilelang_st/npu/a5/src/st/testcase/tmov_cube/ 覆盖 cube bridge 场景
• ST 中 tmov bridge 的场景应嵌入 matmul 全流程中验证（GM→L1→L0A/L0B→MAD→L0C→GM）

---

任务5：textract cube 类

参考 pto-isa 实现

• TExtract.hpp:
  • TEXTRACT_IMPL (L47-115): 多种重载
  • 核心语义：从 src tile 中按 (idxRow, idxCol) 偏移提取子窗口到 dst tile
  • 支持 ReluPreMode、AccToVecMode、preQuantScalar、FpTileData
  • QuantMode 枚举控制量化行为

实施步骤

5a. Template 实现

新建 textract_template.py：

# 基础 extract: Acc → Vec (cube 中最常见场景)
@pto.ckernel(target="a5", op="pto.textract", constraints=[...])
def template_textract_acc2vec(src: pto.Tile, dst: pto.Tile,
                               idx_row: int, idx_col: int):
    # 从 Acc tile 提取子窗口到 Vec tile
    # 参考 TEXTRACT_IMPL 的双重循环：逐 col、逐 row 搬运
    ...

# Acc → Mat/Vec (通过 accToVecMode 控制)
@pto.ckernel(target="a5", op="pto.textract", constraints=[...])
def template_textract_acc2mat(src: pto.Tile, dst: pto.Tile,
                               idx_row: int, idx_col: int, mode): ...

# Vec → Vec (vector 路径，子窗口提取)
@pto.vkernel(target="a5", op="pto.textract", constraints=[...])
def template_textract_vec2vec(src: pto.Tile, dst: pto.Tile,
                               idx_row: int, idx_col: int): ...

# FP 变体（带 FP tile 量化参数）
@pto.ckernel(target="a5", op="pto.textract.fp", constraints=[...])
def template_textract_fp(src: pto.Tile, dst: pto.Tile, fp: pto.Tile,
                          idx_row: int, idx_col: int): ...

5b. 约束

• dst.valid_shape[0] + idxRow <= src.valid_shape[0]
• dst.valid_shape[1] + idxCol <= src.valid_shape[1]
• Acc→Vec: 需要处理 Acc 的 RowMajor SFractal → Vec 的 NoneBox 布局转换
• FP 变体：fp tile 的 shape 覆盖 dst 的列数

5c. Lit 测试

test/lit/vpto/cube/expand_tile_op_tilelang_textract_acc2vec.pto
test/lit/vpto/cube/textract_bounds_verify_invalid.pto

5d. ST 测试

对齐 pto-isa:

• 参考 tests/costmodel/st/testcase/textract/main.cpp
• 参考 tests/cpu/st/testcase/textract/main.cpp
• 参考 tests/cpu/st/testcase/textract_quant/main.cpp（量化变体）
• 新建 test/tilelang_st/npu/a5/src/st/testcase/textract/

---

任务6：tinsert cube 类

参考 pto-isa 实现

• TInsert.hpp:
  • TINSERT_IMPL (L61-95): 多种重载
  • 核心语义：将 src tile 的内容插入/回填到 dst tile 的 (idxRow, idxCol) 偏移处
  • 与 textract 互为逆操作

实施步骤

6a. Template 实现

新建 tinsert_template.py：

# 基础 insert: Vec → Acc (回填)
@pto.ckernel(target="a5", op="pto.tinsert", constraints=[...])
def template_tinsert_vec2acc(src: pto.Tile, dst: pto.Tile,
                              idx_row: int, idx_col: int):
    ...

# Vec → Mat
@pto.ckernel(target="a5", op="pto.tinsert", constraints=[...])
def template_tinsert_vec2mat(src: pto.Tile, dst: pto.Tile,
                              idx_row: int, idx_col: int): ...

# FP 变体
@pto.ckernel(target="a5", op="pto.tinsert.fp", constraints=[...])
def template_tinsert_fp(src: pto.Tile, dst: pto.Tile, fp: pto.Tile,
                         idx_row: int, idx_col: int): ...

6b. 约束

• src.valid_shape[0] + idxRow <= dst.valid_shape[0]
• src.valid_shape[1] + idxCol <= dst.valid_shape[1]
• 注意区分与 tmov 的语义边界：tinsert 是有偏移的部分插入，tmov 是全量 copy

6c. Lit 测试

test/lit/vpto/cube/expand_tile_op_tilelang_tinsert_vec2acc.pto
test/lit/vpto/cube/tinsert_bounds_verify_invalid.pto
test/lit/vpto/cube/tinsert_vs_tmov_matcher_separation.pto

6d. ST 测试

对齐 pto-isa:

• 参考 tests/cpu/st/testcase/tinsert/
• 参考 tests/npu/a5/src/st/testcase/tinsert/
• 新建 test/tilelang_st/npu/a5/src/st/testcase/tinsert/

---

任务7：img2col cube 准备类

参考 pto-isa 实现

• TImg2col.hpp:
  • TIMG2COL_IMPL (L107-139): ConvTile(Mat) → TileLeft 的 img2col 转换
  • ExtractImg2ColParams: 提取 conv 参数（stride/dilation/filter/pad）
  • CalculateValue: 计算每个元素值（含 padding 处理）

关键分析

• img2col 是 Conv 场景的前处理步骤，将 image 数据重排为 matmul 可用的矩阵格式
• 当前在 tilelang DSL 层没有暴露 authoring surface
• 需要决策：是暴露为独立 op 还是作为高阶 tileop 的内部 lowering 细节

实施步骤

7a. 设计决策

建议：先暴露为独立 op，后续 conv 场景可以直接复用。独立 op 的测试更容易隔离问题。

7b. Template 实现

新建 timg2col_template.py：

@pto.ckernel(target="a5", op="pto.timg2col", constraints=[...])
def template_timg2col(src: pto.Tile, dst: pto.Tile,
                       pos_m: int, pos_k: int,
                       fmatrix_mode):
    # 对应 TIMG2COL_IMPL 的核心逻辑
    # src: ConvTile (loc=mat, layout=NC1HWC0/NDC1HWC0)
    # dst: TileLeft (loc=left)
    ...

7c. 约束

• src.loc == Mat, dst.loc == Left
• src.layout in [NC1HWC0, NDC1HWC0]
• dst.s_layout == RowMajor, dst.isRowMajor == False
• dtype 一致

7d. Lit 测试

test/lit/vpto/cube/expand_tile_op_tilelang_timg2col.pto

7e. ST 测试

对齐 pto-isa:

• 参考 tests/cpu/st/testcase/timg2col/
• 参考 tests/npu/a5/src/st/testcase/timg2col/
• 新建 test/tilelang_st/npu/a5/src/st/testcase/timg2col/

---

任务8：ttrans 类

参考 pto-isa 实现

• TTrans.hpp:
  • TTRANS_IMPL (L48-55): dst, src, tmp 三个参数
  • 核心语义：矩阵转置，dst.rows == src.cols, dst.cols == src.rows
  • 使用 tmp tile 作为中间缓冲

实施步骤

8a. Template 实现

新建 ttrans_template.py：

@pto.ckernel(target="a5", op="pto.ttrans", constraints=[...])
def template_ttrans(src: pto.Tile, dst: pto.Tile, tmp: pto.Tile):
    # 对应 TTRANS_IMPL
    # 分块转置需要考虑 SFractal 格式的变换
    ...

8b. 约束

• src.valid_shape == (dst.valid_cols, dst.valid_rows)（行列交换）
• tmp tile 的 size 至少与 src 一致
• 布局约束参考 TTrans_Impl 中的分块转置逻辑

8c. Lit 测试

test/lit/vpto/cube/expand_tile_op_tilelang_ttrans.pto

8d. ST 测试

对齐 pto-isa:

• 参考 tests/costmodel/st/testcase/ttrans/
• 参考 tests/cpu/st/testcase/ttrans/
• 参考 tests/npu/a5/src/st/testcase/ttrans/
• 新建 test/tilelang_st/npu/a5/src/st/testcase/ttrans/

---

任务依赖关系和推荐执行顺序

任务1 (tload/tstore MAT)
 ├── 是其余所有 cube 任务的前置依赖
 │   （没有 tload/tstore MAT，cube ST 只能用裸 VPTO 指令搬运数据）
 │
 ├──> 任务2 (tmatmul 家族)  ← 最先开始的 compute op
 │    └──> 任务3 (tgemv 家族)  ← 可以复用 tmatmul 模式
 │
 ├──> 任务4 (tmov bridge)  ← L1→L0 的桥接语义，cube 全流程必备
 │    ├──> 任务5 (textract)  ← 与 tmov 有语义边界，需和任务4协调
 │    └──> 任务6 (tinsert)   ← textract 的逆操作，可参考任务5模式
 │
 └──> 任务7 (img2col)  ← conv 场景专用，相对独立
      └──> 任务8 (ttrans)      ← 也相对独立

---

共同约束检查清单（每个任务提交前自查）

• 新增 template 有对应的 @pto.ckernel / @pto.vkernel 装饰器，constraints 正确
• vector 路径 matcher 遇到 cube loc 返回 False，反之亦然（不串线）
• 至少一个正向 lit 测试（expand 通过 FileCheck）
• 关键非法组合有负向 lit 测试（verify/diagnostic 报错）
• ST 用例使用独立的 testcase 目录
• ST 用例通过 tload/tstore 打通完整数据流，不手写裸 VPTO cube load/store
• 代码文件包含 OAT.3 license header
• 参考的 pto-isa ST case 已在对应的 pto-isa tests 目录下确认存在