[TileOP/Compiler] Group TMATMUL 应按每 PE 的 local tileC 行数生成 LB0

## 问题背景

在 4 PE multi-thread matmul 中，矩阵 A、B 使用 core-level SharedTile，输出矩阵 C 按 M 轴切成 4 份，每个 PE 只保存自己的 local tileC。

预期的 tile 组织方式为：

- Shared A：`[tM, tK]`，4 个 PE 共同可见。
- Shared B：`[tK, tN]`，4 个 PE 共同可见。
- 每个 PE 的 local C：`[tM / 4, tN]`。
- 4 个 PE 合起来得到完整输出 C：`[tM, tN]`。

当前 TileOP API 可以成功编译该代码，但生成的 Group TMATMUL 反汇编中，`LB0` 仍然使用 Shared A 的完整行数 `tM`，而不是每个 PE 的 local tileC 行数 `tM / 4`。

这导致反汇编中的输出 tile 大小与 `LB0/LB1` 表示的结果矩阵形状不一致，也会导致性能模型无法正确解释该指令。

## 参考用例

用例文件：

[SuperNPUBench matmul_shared.hpp](https://github.com/PTO-ISA/SuperNPUBench/blob/main/benchmark/one-level-arch/kernels/matmul/matmul_shared.hpp)

核心代码如下：

```cpp
template <typename dtype, int gM, int gN, int gK,
          int tM, int tN, int tK>
void matmul_shared(float *c_ptr, dtype *a_ptr, dtype *b_ptr) {
    constexpr int kPeNum = 4;

    using tileALocal = TileLeft<dtype, tM, tK>;
    using tileBLocal = TileRight<dtype, tK, tN>;

    using tileAShared = SharedTile<tileALocal>;
    using tileBShared = SharedTile<tileBLocal>;

    // 每个 PE 只保存完整 C tile 的 1/4 行
    using tileC =
        Tile<Location::Vec, float,
             tM / kPeNum, tN,
             BLayout::RowMajor>;

    tileAShared tAShared;
    tileBShared tBShared;
    tileC tC;

    // SharedTile 由一个 PE 加载，加载完成后 core 内 4 个 PE 可见
    TLOAD<tileALocal, 1>(tAShared, gA);
    TLOAD<tileBLocal, 1>(tBShared, gB);

    // 4 PE 协作执行 Group TMATMUL
    TMATMUL(tC, tAShared, tBShared);

    // 每个 PE 写回自己的 M 轴切片
    auto gC = gIterC(i * kPeNum + tid, j);
    TSTORE(gC, tC);
}
```

测试参数：

```text
gM = 256
gN = 256
gK = 256

tM = 64
tN = 16
tK = 16

PE count = 4
```

对应的 tile 形状和大小：

```text
Shared A = [64, 16], FP32 = 4096 B
Shared B = [16, 16], FP32 = 1024 B
local C  = [16, 16], FP32 = 1024 B / PE
```

## 复现命令

```bash
cd benchmark/one-level-arch/test/kernel/multi_thread/matmul

make \
  TESTCASE=matmul_shared \
  B=1 M=256 N=256 K=256 \
  tM=64 tN=16 tK=16 \
  COMPILER_DIR=<linx-toolchain>/bin \
  diss
```

## 当前反汇编

Shared A 和 Shared B 的加载结果基本符合预期：

```asm
# Shared A: 64 x 16 x FP32 = 4 KB
BSTART.TLSU TLOAD, FP32
B.DIM 16, 0, ->lb0
B.DIM 64, 0, ->lb1
C.B.DIMI 16, ->lb2
B.IOS mask=0001, ->S0<4KB>

# Shared B: 16 x 16 x FP32 = 1 KB
BSTART.TLSU TLOAD, FP32
B.DIM 16, 0, ->lb0
B.DIM 16, 0, ->lb1
C.B.DIMI 16, ->lb2
B.IOS mask=0001, ->S1<1KB>
```

当前 Group TMATMUL 反汇编为：

```asm
BSTART.CUBE TMATMUL, FP32
B.DATR FP32, byte0, Zero
B.FPATR 0, 0, 0, 0, 0, 0, 0

B.DIM 64, 0, ->lb0       # 当前来自 Shared A.Rows
B.DIM 16, 0, ->lb1
B.DIM 16, 0, ->lb2

B.IOS S0, mask=1111
B.IOS S1, mask=1111
B.IOT mask=1111, last, ->u<1KB>
```

这里存在明显的不一致：

```text
LB0 × LB1 × sizeof(float)
= 64 × 16 × 4
= 4096 B
```

但是目的 local tile 实际为：

```text
16 × 16 × 4
= 1024 B
```

因此，指令中的 `LB0=64` 描述的是完整 core-level 输出行数，而 `B.IOT ->u<1KB>` 描述的是每个 PE 的 local 输出 tile。

## 预期反汇编

在 Shared A/B 触发的 Group TMATMUL 中，建议定义：

```text
LB0 = 每个 PE 的 local tileC row
LB1 = local tileC col
LB2 = A/B 的 K
```

因此本用例的预期反汇编为：

```asm
BSTART.CUBE TMATMUL, FP32
B.DATR FP32, byte0, Zero
B.FPATR 0, 0, 0, 0, 0, 0, 0

B.DIM 16, 0, ->lb0       # local tileC.Rows = tM / 4
B.DIM 16, 0, ->lb1       # local tileC.Cols = tN
B.DIM 16, 0, ->lb2       # A/B 的 K = tK

B.IOS S0, mask=1111
B.IOS S1, mask=1111
B.IOT mask=1111, last, ->u<1KB>
```

对应的 Group TMATMUL 语义为：

```text
PE 0 使用 Shared A 的 row [0, 16)
PE 1 使用 Shared A 的 row [16, 32)
PE 2 使用 Shared A 的 row [32, 48)
PE 3 使用 Shared A 的 row [48, 64)

每个 PE 产生 [16,16] 的 local C。
4 个 PE 合起来产生完整的 [64,16] 输出。
```

即：

```text
SharedA.Rows = PECount × LB0
SharedB.Cols = LB1
SharedA.Cols = SharedB.Rows = LB2
localC.Rows  = LB0
localC.Cols  = LB1
```

## 原因分析

当前 TileOP API 的三参数 `TMATMUL` 接口使用 A、B 的形状生成 M/N/K：

```cpp
size_t M = pto_matmul_detail::matrix_valid_row(a);
size_t N = pto_matmul_detail::matrix_valid_col(b);
size_t K = pto_matmul_detail::matrix_valid_col(a);

pto_matmul_detail::matmul<Attr>(c, a, b, M, N, K);
```

因此：

```text
M = SharedA.Rows = 64
```

最终生成：

```asm
B.DIM 64, ->lb0
```

但是该接口没有区分以下两种执行模式：

1. 普通单 PE TMATMUL：

```text
C.Rows == A.Rows
```

2. Shared A/B 的 4 PE Group TMATMUL：

```text
A.Rows == 4 × C_local.Rows
```

在 Group TMATMUL 中继续从 A 取得 M，会导致 `LB0` 与 local destination tile 的实际形状不一致。

## TileOP API 修改建议

建议针对 Shared A + Shared B 的 Group TMATMUL 单独确定 M/N/K：

```cpp
constexpr bool IsGroupMatmul =
    is_shared_tile_v<tile_shape_a> &&
    is_shared_tile_v<tile_shape_b>;

size_t M;
size_t N;
size_t K;

if constexpr (IsGroupMatmul) {
    // Group TMATMUL 中 B.DIM 描述每个 PE 的 local 输出形状
    M = pto_matmul_detail::matrix_valid_row(c);
    N = pto_matmul_detail::matrix_valid_col(c);
    K = pto_matmul_detail::matrix_valid_col(a);
} else {
    M = pto_matmul_detail::matrix_valid_row(a);
    N = pto_matmul_detail::matrix_valid_col(b);
    K = pto_matmul_detail::matrix_valid_col(a);
}
```

同时建议增加编译期形状检查：

```cpp
constexpr int kGroupPeCount = 4;

if constexpr (IsGroupMatmul) {
    static_assert(
        tile_shape_a::Rows ==
            kGroupPeCount * tile_shape_c::Rows,
        "Group TMATMUL requires "
        "SharedA.Rows == PECount * LocalC.Rows");

    static_assert(
        tile_shape_a::Cols == tile_shape_b::Rows,
        "TMATMUL requires A.Cols == B.Rows");

    static_assert(
        tile_shape_b::Cols == tile_shape_c::Cols,
        "Group TMATMUL requires "
        "SharedB.Cols == LocalC.Cols");
} else {
    static_assert(
        tile_shape_a::Rows == tile_shape_c::Rows &&
        tile_shape_b::Cols == tile_shape_c::Cols,
        "TMATMUL output shape must be A.Rows x B.Cols");
}
```

如果 tile 支持运行时 valid shape，也需要增加对应的运行时或 debug 检查：

```text
SharedA.ValidRows == PECount × LocalC.ValidRows
SharedA.ValidCols == SharedB.ValidRows
SharedB.ValidCols == LocalC.ValidCols
```

建议将同样的 shape 和 `LB0` 规则覆盖到：

- `TMATMUL`
- `TMATMUL_ACC`
- 带 `fixp::Options` 的 `TMATMUL`
- 带 `fixp::Options` 的 `TMATMUL_ACC`
- 其他支持 Shared A/B 的 Group TMATMUL 变体

目前带 options 的接口仍包含：

```cpp
static_assert(
    tile_shape_d::Rows == tile_shape_a::Rows &&
    tile_shape_d::Cols == tile_shape_b::Cols,
    "TMATMUL output shape must be M x N");
```

该检查同样无法接受 `D.Rows == A.Rows / 4` 的 Group TMATMUL，需要区分普通模式和 Group 模式。

另外，当前 Shared A/B 路径内联汇编中固定使用：

```asm
B.IOS S0, mask=1111
B.IOS S1, mask=1111
B.IOT mask=1111, last, ->Dst
```

建议明确以下两种方案中的一种：

1. 当前 ISA 固定为 4 PE group，TileOP API 中明确记录并检查 `PECount=4`。
2. 将 PE mask 或 PE count 作为接口属性传入，避免在 API 内硬编码 `1111`。

## LLVM 编译器修改建议

如果 `LB0/LB1/LB2` 的语义在 ISA 层调整为 per-PE local output shape，LLVM 编译器需要同步确认以下模块：

1. 指令选择/降低

确保 Group TMATMUL 的参数映射为：

```text
LB0 = local destination rows
LB1 = local destination cols
LB2 = reduction K
```

不能再次根据 Shared A 的完整行数重建 `LB0`。

2. 汇编器和反汇编器

Group TMATMUL 的编码和译码应保持以下信息一致：

```text
B.DIM LB0/LB1/LB2：每个 PE 的计算形状
B.IOS mask：参与 Group TMATMUL 的 PE
B.IOT TSize：每个 PE 的 local destination tile 大小
```

3. 指令合法性检查

对 Group TMATMUL 检查：

```text
SharedA.Rows == popcount(PEMask) × LB0
SharedB.Cols == LB1
SharedA.Cols == SharedB.Rows == LB2
DstSize >= LB0 × LB1 × sizeof(DstType)
```

不应继续要求：

```text
DstSize >= SharedA.Rows × SharedB.Cols × sizeof(DstType)
```

4. 回归测试

建议增加 FileCheck 测试，固定检查：

```asm
B.DIM 16, 0, ->lb0
B.DIM 16, 0, ->lb1
B.DIM 16, 0, ->lb2
B.IOS S0, mask=1111
B.IOS S1, mask=1111
B.IOT mask=1111, last, ->u<1KB>
```

并增加负例：

```text
SharedA.Rows != 4 × LocalC.Rows
SharedB.Cols != LocalC.Cols
SharedA.Cols != SharedB.Rows
```

## 期望结果

希望 TileOP API 和 LLVM 编译器明确并统一 Group TMATMUL 的语义：

```text
Shared tile：core-level、所有参与 PE 可见
Local tile：PE-private
Group TMATMUL：完整 M 轴按参与 PE 数切分
LB0/LB1：表示每个 PE 的 local C 形状
LB2：表示 A/B 的 K
```

对于本用例，最终反汇编应从：

```asm
B.DIM 64, 0, ->lb0
B.IOT mask=1111, last, ->u<1KB>
```

修改为：

```asm
B.DIM 16, 0, ->lb0
B.IOT mask=1111, last, ->u<1KB>
```

这样 `B.DIM`、目的 tile 大小以及每个 PE 的实际输出数据范围才能保持一致。

## 测试版本

```text
Linx-TileOP-API commit:
8b2ee780ddcc5d7a04c50e337d67eadc3637a17c

llvm-project commit:
86959776bd1fb22dcc8e73b57ec2276c65d44f38
```

Navigation Menu

Uh oh!

[TileOP/Compiler] Group TMATMUL 应按每 PE 的 local tileC 行数生成 LB0 #18

问题背景

参考用例

复现命令

当前反汇编

预期反汇编

原因分析

TileOP API 修改建议

LLVM 编译器修改建议

期望结果

测试版本

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Uh oh!

[TileOP/Compiler] Group TMATMUL 应按每 PE 的 local tileC 行数生成 LB0 #18

Description

问题背景

参考用例

复现命令

当前反汇编

预期反汇编

原因分析

TileOP API 修改建议

LLVM 编译器修改建议

期望结果

测试版本

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions