Skip to content

[Bug] trowexpand op do not support col major tile #826

Description

@Likai-19

Component

PTO Dialect / ODS (include/PTO/IR)

Description

trowexpand 系列 op(pto.trowexpandmulpto.trowexpandsubpto.trowexpanddiv)不支持 src1 为 col major 布局的场景。当 src1 的 shape 为 [M, 1](per-row scalar column)且布局为 col_major 时,模板中对 src1[row, :]vlds 调用会访问非 512B 对齐的 UB 地址,在 A5 设备上触发 error 340。

此前模板的布局约束要求所有操作数均为 row_major,但 Qwen3 row-expand 类 kernel 中常见的 scalar 输入布局是 col_major, cols=1,导致模板匹配失败或运行时访存异常。

此外,tilelang-dsl 的 lowering 层在生成 vldas/vldus 指令时,原有的 tile → memref → subview → castptr 路径存在结构性问题:vldas 只接受 !pto.ptr 类型,而 memref 中间表示引入了不必要的 subview 和 cast 操作,对 col major [M, 1] tile 的地址计算也不正确。

Reproduction (minimal)

使用 Qwen3 row-expand kernel,其中 `src1` 为 col major 且 shape 为 `[M, 1]`# src0: row_major [M, N]
# src1: col_major [M, 1]   ← per-row scalar
# dst:  row_major [M, N]
pto.trowexpandmul(src0, src1, dst)


在 A5 设备上运行会触发 error 340(非对齐访问异常)。

Expected behavior

  • trowexpandmul/trowexpandsub/trowexpanddiv 模板应同时支持 src1row_majorcol_major 两种布局。
  • col_major 路径,应使用非对齐访问微指令 vldas + vldus 流水线加载 scalar 向量,避免 512B 对齐限制。
  • row_major 路径保持原有的 vlds 对齐访问方式不变。
  • lowering 层应能将 tile 直接转换为 !pto.ptr(而非 memref),并通过 arith.muli + arith.addi + pto.addptr 正确计算元素偏移。

Actual behavior / error logs

- 模板布局约束 `_constraint_trowexpanddiv_row_major` 要求所有操作数为 `row_major`,`col_major` 的 `src1` 无法匹配。
- 即使绕过约束,运行时 `vlds` 在 col major `[M, 1]` tile 切片上访问非 512B 对齐地址,触发 error 340。
- lowering 层原有的 memref → subview → castptr 路径对 `vldas`/`vldus` 不适用,生成的 IR 结构不正确。

Git commit

25d8508

Host platform

None

Target Ascend arch (if relevant)

None

PTOAS build level (if relevant)

None

Metadata

Metadata

Assignees

No one assigned

    Labels

    bugSomething isn't working

    Type

    No type
    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions