[Bug] 输出全是 0，说明 DMA mte_gm_l1_frac (nd2nz 模式) 不支持 M=1同时将M向16对齐pto.cpp 文件中 pto.tgemv 检查左边矩阵（lhs）的valid_shape[0]是否等于1 或是动态值。

[yuqiha]

Component

PTO Dialect / ODS (include/PTO/IR)

Description

tgemv 算子 m=1，在pto.cpp 中 进行了限制：/home/wujiajun/llvm-workspace/PTOAS_matmul0/lib/PTO/IR/PTO.cpp 3799 行，但是这时候会结果都为0，agent定位如下：

如果对齐到 16 ，/home/wujiajun/llvm-workspace/PTOAS_matmul0/lib/PTO/IR/PTO.cpp 3799 行不兼容（pto-isa中没有对齐到16）

对应代码：https://github.com/yuqiha/PTOAS_fork/tree/matmul_1

Reproduction (minimal)

python3 test/tilelang_st/script/run_st.py -r sim -v a5 -t tgemv

Expected behavior

pass

Actual behavior / error logs

M=1时

Model RUN TIME: 109682 ms
[INFO] Total tick: 10022
[INFO] Model stopped successfully.
run command: cp testcase/st_common.py build/testcase/tgemv/st_common.py
run command: cp testcase/tgemv/cases.py build/testcase/tgemv/cases.py
run command: cp testcase/tgemv/gen_data.py build/testcase/tgemv/gen_data.py
run command: cp testcase/tgemv/compare.py build/testcase/tgemv/compare.py
run command: /usr/bin/python3 compare.py
[ERROR] Mismatch: max diff=11.883535385131836 at flat idx=32 (golden=-11.883535385131836, output=0.0)
[ERROR] gemv_f16_1x300x60: compare failed
[ERROR] Mismatch: max diff=18.262971878051758 at flat idx=18 (golden=18.262971878051758, output=0.0)
[ERROR] gemv_bias_f16_1x512x85: compare failed
run command failed with return code 2
compare failed: Command '['/usr/bin/python3', 'compare.py']' returned non-zero exit status 2.
run failed: Command '['/usr/bin/python3', 'compare.py']' returned non-zero exit status 2.

M对齐到 16：
CMake Error at /home/wujiajun/llvm-workspace/PTOAS_matmul0/test/tilelang_st/npu/a5/src/st/testcase/run_ptoas_to_file.cmake:51 (message):
  ptoas failed while generating
  /home/wujiajun/llvm-workspace/PTOAS_matmul0/test/tilelang_st/npu/a5/src/st/build/testcase/tgemv/tgemv_kernel.o:



  loc("/home/wujiajun/llvm-workspace/PTOAS_matmul0/test/tilelang_st/npu/a5/src/st/testcase/tgemv/tgemv.pto":73:5):
  error: 'pto.tgemv' op expects lhs valid_shape[0] to be 1 for tgemv


  loc("/home/wujiajun/llvm-workspace/PTOAS_matmul0/test/tilelang_st/npu/a5/src/st/testcase/tgemv/tgemv.pto":160:5):
  error: 'pto.tgemv.bias' op expects lhs valid_shape[0] to be 1 for tgemv

  Error: Failed to parse MLIR.


make[2]: *** [testcase/tgemv/CMakeFiles/tgemv_kernel.dir/build.make:75: testcase/tgemv/tgemv_kernel.o] Error 1
make[2]: Leaving directory '/home/wujiajun/llvm-workspace/PTOAS_matmul0/test/tilelang_st/npu/a5/src/st/build'
make[1]: *** [CMakeFiles/Makefile2:114: testcase/tgemv/CMakeFiles/tgemv_kernel.dir/all] Error 2
make[1]: Leaving directory '/home/wujiajun/llvm-workspace/PTOAS_matmul0/test/tilelang_st/npu/a5/src/st/build'
make: *** [Makefile:91: all] Error 2

run failed: Command '['make', 'VERBOSE=1', '-j', '64']' returned non-zero exit status 2.
wujiajun@ecs-1030-cba0:~/llvm-workspace/PTOAS_matmul0$ python3 test/tilelang_st/script/run_st.py -r sim -v a5 -t tgemv

Git commit

8a8960d

Host platform

None

Target Ascend arch (if relevant)

None

PTOAS build level (if relevant)

None

[Zhendong404]

本地继续复现和分析了一轮，这个问题的结论和建议修复方向如下。

结论

当前问题的根因 不只是 verifier 对 lhs.valid_shape[0] == 1 的限制。

pto.cpp 这部分我本地重新确认过：当 tile 的物理行数对齐到 16、但逻辑 valid_row=1 时，tgemv / tgemv.bias 的 verifier 是可以接受的，所以 issue 里看到的“对齐到 16 后编不过”并不是当前主因。

真正导致 gemv_f16_1x300x60 失败的是：

1. M=1 时，L0C -> GM 的 store 需要按物理行对齐后的布局来写，mte_l0c_gm 的 src_stride 不能继续按 1 来写。这里改成按物理 16 行的 stride 后，输出不再是全 0。
2. 即便修正 store 布局，单次 TGEMV 路径在这个 A5 用例上对 K > 256 仍然不可靠。我本地对 gemv_f16_1x300x60 的输出做了离线比对，发现结果几乎和“只计算前 256 列 K”的 matmul 完全一致，而不是完整的 K=300。

也就是说，这个 case 实际上需要走 split-K：先做一段 K=256 的 tgemv，再做剩余 tail 的 tgemv.acc。

证据

1. pto-isa 的 A5 C++ 封装 TMatmul.hpp 里，TGEMV_IMPL / TGEMV_ACC_IMPL / TGEMV_BIAS_IMPL 都只是单次调用 TMatmul / TMatmulBias，内部没有任何按 K 分块的循环，也没有自动 split-K 的逻辑。
2. 仓里的现有样例 test/samples/Gemv/gemvacc.py 已经是按 BASEK = 256 手工切成两段，再显式调用 TGemvOp + TGemvAccOp。
3. 本地把 ST 里的 gemv_f16_1x300x60 改成 256 + 44 的 split-K（第二段物理 K padding 到 64），并且 host 侧显式传第二段 A/B 指针后，python3 test/tilelang_st/script/run_st.py -r sim -v a5 -t tgemv 已经可以通过，两个 case 都 pass。

建议修复方向

建议把修复拆成两层理解：

1. 短期修复 / 回归用例修复

   • 对 K > 256 的 TGEMV 用例，不要再尝试单次 pto.tgemv 直出。
   • 直接按 BASEK = 256 做 split-K：
     • 首段：pto.tgemv 或 pto.tgemv.bias
     • 后续段：pto.tgemv.acc
   • M=1 场景下，ACC store 到 GM 时按物理对齐后的布局设置 stride。

2. 长期修复 / 产品化修复

   • 在 PTOAS 的 TileOp/template/lowering 层，对 pto.tgemv / pto.tgemv.bias 增加自动 split-K 展开逻辑。
   • 触发条件建议至少覆盖 K > 256 的 A5 GEMV 路径。
   • 展开结果建议是：
     • chunk0 -> tgemv / tgemv.bias
     • chunk1...n -> tgemv.acc
   • 这样用户侧仍然只写一个高层 GEMV，而不是手工拆多个 tile/op。

建议补充测试

建议至少保留/新增下面两类回归：

1. M=1, K=300, N=60 这类 K > 256 的 plain GEMV case
2. M=1, K=512, N=85 这类 split-K + bias 的 case

如果需要，我可以继续把这条逻辑往 PTOAS 的通用模板/自动展开层推进，而不是只停留在 ST testcase 修复。

[Zhendong404]

补充一下这次本地验证通过的核心代码修改，方便后面继续抽象成通用修复。

1. case1 从单次 TGEMV 改成 split-K

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

原来 gemv_f16_1x300x60 是：

• 单次加载 A(1x320)
• 单次加载 B(320x64)
• 单次 pto.tgemv

本地通过的版本改成：

• pass0: K=256，用 pto.tgemv
• pass1: tail K=44，物理 padding 到 64，用 pto.tgemv.acc

核心片段如下：

// Pass 0: A[:,0:256], B[0:256,:]
pto.mte_gm_l1 %a0_gm, %l1_a0, %c512_i64 nburst(%c1_i64, %c0_i64, %c0_i64)
pto.mte_gm_l1_frac %b0_gm, %l1_b0, nd2nz,
  shape(%c256_i64, %c64_i64), src_layout(%c128_i64),
  dst_group(%c1_i64, %c1_i64, %c256_i64, %c0_i64), ctrl(%c0_i64, %false)
...
pto.tgemv ins(%l0a0_tile, %l0b0_tile ...)
         outs(%l0c_tile ...)

// Pass 1: A[:,256:320], B[256:320,:], logical K=44 / physical K=64
pto.mte_gm_l1 %a1_gm, %l1_a1, %c128_i64 nburst(%c1_i64, %c0_i64, %c0_i64)
pto.mte_gm_l1_frac %b1_gm, %l1_b1, nd2nz,
  shape(%c64_i64, %c64_i64), src_layout(%c128_i64),
  dst_group(%c1_i64, %c1_i64, %c64_i64, %c0_i64), ctrl(%c0_i64, %false)
...
pto.tgemv.acc ins(%l0c_tile, %l0a1_tile, %l0b1_tile ...)
            outs(%l0c_tile ...)

2. host 侧显式拆第二段 A/B 指针

因为当前 lowering 不允许把 gm addptr 直接喂给这里的 MTE 路径，所以 case1 需要像 bias case 一样，在 launch wrapper 里把第二段 A/B 指针提前拆好传进内核。

test/tilelang_st/npu/a5/src/st/testcase/tgemv/launch.cpp

extern "C" __global__ AICORE void TGEMV_f16_1x300x60(
    __gm__ uint16_t *a0,
    __gm__ uint16_t *b0,
    __gm__ uint16_t *a1,
    __gm__ uint16_t *b1,
    __gm__ float *c
);

void LaunchTGEMV_f16_1x300x60(void *a, void *b, void *c, void *stream) {
    uint16_t *a_ = (uint16_t *)a;
    uint16_t *b_ = (uint16_t *)b;
    TGEMV_f16_1x300x60<<<1, nullptr, stream>>>(
        (__gm__ uint16_t *)(a_),
        (__gm__ uint16_t *)(b_),
        (__gm__ uint16_t *)(a_ + 256),
        (__gm__ uint16_t *)(b_ + 256 * 64),
        (__gm__ float *)c
    );
}

这里 b_ + 256 * 64 对应的是 padded B 的第二个 K chunk 起点。

3. M=1 场景下 L0C -> GM store 的 stride 要按物理对齐布局写

这个也是现象从“全 0”变成“非 0”的关键条件之一。对 acc tile 物理 rows=16、逻辑 valid_row=1 的场景，mte_l0c_gm 不能继续按 src_stride=1 写。

本地通过版本里 case1/case2 都用了：

pto.mte_l0c_gm %l0c, %c_gm, %c1_i64, %c64_i64, %c16_i64, %c64_i64, %c0_i64, %c0_i64, nz2nd

这里关键是 src_stride = 16，对应物理对齐后的 ACC 行布局。

4. 本地验证结果

用下面命令验证：

python3 test/tilelang_st/script/run_st.py -r sim -v a5 -t tgemv

结果：

[INFO] gemv_f16_1x300x60: compare passed
[INFO] gemv_bias_f16_1x512x85: compare passed
[INFO] all cases passed

5. 这批修改的意义

这批代码修改更像是“确认真实可用执行路径”的证据：

• 单次 K=300 的 TGEMV 路径在这个 A5 case 上不可靠
• 但 256 + tail 的 tgemv + tgemv.acc 路径是可用的
• 所以后续真正产品化的修复，建议仍然放在 PTOAS 的 TileOp/template/lowering 层，自动把 K > 256 的 TGEMV 展开成 split-K，而不是要求用户手工写多个 kernel 参数和多个 tile chunk