Skip to content

【Matmul、FlashAttention 算子】TileOp 使用情况与 PTO ISA 映射总结(17 类 TileOp,3 核心需补 / 3 可选需补 / 11 暂无明显需补) #16

Description

@lvhao7896

展示 test/kernel 中 Matmul 与 FlashAttention 相关 kernel 使用到的 TileOp,并整理这些 TileOp 到 PTO ISA 的接口映射、语义差异、映射难点和 blockisa 当前支持情况。

算子tileop汇总

Category Test Type / MODE Actual Kernel TileOps Used Source
Matmul MASK_FP32 / MASK_FP16 / MASK_FP8 matmul_mask TCOPYIN, MATMUL, MATMACC, TMOV(NZ2ND), TCOPYOUT kernels/matmul/matmul.hpp
Matmul *_REUSEA matmul_mask_reuseA TCOPYIN, MATMUL, MATMACC, TMOV(NZ2ND), TCOPYOUT kernels/matmul/matmul.hpp
Matmul *_REUSEB matmul_mask_reuseB TCOPYIN, MATMUL, MATMACC, TMOV(NZ2ND), TCOPYOUT kernels/matmul/matmul.hpp
Matmul HIF4_HIF4 VER=MX_NOGATHER matmul_mxfp_notcvt TCOPYIN, MATMULMX, MATMACCMX, TMOV(NZ2ND), TCOPYOUT kernels/matmul/matmul_mx.hpp
Matmul HIF4_HIF4 VER=MX_NOGATHER_REUSEA matmul_mxfp_notcvt_reuseA TCOPYIN, MATMULMX, MATMACCMX, TMOV(NZ2ND), TCOPYOUT kernels/matmul/matmul_mx.hpp
Matmul A16W4 matmul_mp TEXPANDSCALAR, TCOPYIN, MATMUL, TMOV(NZ2ND), TCOPYOUT kernels/matmul/matmul_mx.hpp
Flash Attention fa_2d_unroll flash_attention_2d_unroll TCOPYIN, TEXPANDSCALAR, MATMUL, TCVT, TMOV_DN2NZ, MATMACC, TMOV_NORM, TCOPYOUT kernels/fa/fa_2d_unroll.hpp
Flash Attention fa_HIF4_HIF4 MODE=BF16_NOGATHER flash_attention_2d_unroll_hif4_nogather TCOPYIN, TEXPANDSCALAR, MATMULMX, TCVT, TMOV_DN2NZ, MATMACCMX, TMOV_NORM, TCOPYOUT kernels/fa/fa_hif4.hpp
flash_attention_vector_kernels rowmax,tsub,texp,rowsum,tdiv,fcvt(2xbf16->hif4x2),tmadd,tmul,tadd,tgemv kernels/fa/fa_hif4.hpp

LinxISA TileOp 到 PTO ISA 映射总结

  • op范围:TCOPYIN, TEXPANDSCALAR, MATMULMX, TCVT, TMOV_DN2NZ, MATMACCMX, TMOV_NORM, TCOPYOUT, rowmax, tsub, texp, rowsum, tdiv, fcvt(2xbf16->hif4x2), tmadd, tmul, tadd。
tileop PTO接口 / 组合 LInxISA接口 / 原算子 共同语义 语义差异 映射难点 当前编译器需补充内容(对齐 PTO ISA)
MATMACCMX %c_out = tmatmul.mx.acc %c_in, %a, %a_scale, %b, %b_scale : (!pto.tile<...>, !pto.tile<...>, !pto.tile<...>, !pto.tile<...>, !pto.tile<...>) -> !pto.tile<...> TMATMULMX.ACC <LB0: M, LB1: N, LB2: K, DataTypeA, DataTypeB> SrcTile0<.reuse>, SrcScale0, SrcTile1<.reuse>, SrcScale1, ACC, ->ACC<Size> MX/block-scale 矩阵乘累加:ACC_out = ACC_in + A*B。 pto.tmatmul.mx.acc,显式把 c_in/c_out 建成数据流 value表达能力比LinxISA 使用acc寄存器强一些。 基本一致 缺少与 PTO tmatmul.mx.acc 对齐的显式 c_in -> c_out 对外实现;需补独立 acc 输入/输出的数据流表达。
fcvt(2xbf16->hif4x2) 优先尝试:%dst = tcvt %src {rmode = ...} : !pto.tile<bf16x2/...> -> !pto.tile<hif4x2/...><br>无支持时组合:tcvt/量化(scale) + tpack/tmov 或自定义转换 kernel FCVT 2xbf16 -> hif4x2;当前 repo 中相关路径用 tohif4_bf16x2 生成 hif4 tile 和 scale tile,注释写有 FCVT TODO。 把两个 BF16 元素量化/打包成一个 HIF4x2/FP4x2 元素,并通常伴随 scale 元数据。 PTO 有通用 pto.tcvt,但没有在已确认文档中看到 2xbf16->hif4x2 这一特定单条转换 缺乏2xbf16->hif4x2的数据类型转换 缺少 2xbf16 到 hif4x2 的对应实现;需补 PTO tcvt/fcvt 转换实现,包括量化/打包路径。
tgemv %acc = tgemv %a, %b : (!pto.tile<...>, !pto.tile<...>) -> !pto.tile<...><br>另有:tgemv.acc / tgemv.bias TGEMV <LB0: M, LB1: N, LB2: K, DataType> SrcTile0<.reuse>, SrcTile1<.reuse>, ->ACC<Size> 矩阵-向量乘:C[0,j] = sum_k A[0,k] * B[k,j]。 - 基本一致。 缺少对应 TGEMV 实现;需补普通 tgemv,若 PTO 使用 tgemv.acc/tgemv.bias,再补 acc/bias 变体即可。
TCVT %dst = tcvt %src {rmode = #pto.round_mode<CAST_RINT>} : !pto.tile<...> -> !pto.tile<...> TCVT <LB0: ValidCol, LB1: ValidRow, DataTypeDst, RoundMode/SaturationMode>, SrcTile, ->DstTile<Size> 逐元素类型转换,可涉及 rounding/saturation,也常用于 ACC -> Vec。 - 基本一致 需补充/确认 PTO tcvt 的 rounding mode 表达;若目标 PTO 流程依赖 saturation,也需补对应模式。
TCOPYOUT tstore %t1, %sv_out[%c0, %c0]<br>fix-pipe: tstore.fp %t1, %fp, %sv_out[%c0, %c0] TSTORE.Layout <LB0: ValidCol, LB1: ValidRow, LB2: Col, DataType>, SrcTile0<.reuse>..., [RegSrc0, RegSrc1], DepSrc..., ->DepDst 从 tile 写回全局内存。 - 基本一致 可对齐 PTO 基本 tstore
TMOV_NORM %dst = tmov.v2v %src : !pto.tile<...> -> !pto.tile<...><br>或 %dst = tmov.s2d %src : !pto.tile<...> -> !pto.tile<...> TMOV NORM, <LB0: ValidCol, LB1: ValidRow, LB2: Col, DataType, PadValue>, SrcTile<.reuse>, ->DstTile<Size> Tile 到 Tile 的普通复制,不改变 layout。 - 基本一致 可对齐 PTO 基本 tmov
TCOPYIN %t0 = tload %sv[%c0, %c0] : (!pto.memref<...>, index, index) -> !pto.tile<...> TLOAD.Layout <LB0: ValidCol, LB1: ValidRow, LB2: Col, DataType, PadValue>, [RegSrc0, RegSrc1], DepSrc..., ->DstTile<Size>, DepDst 从全局内存搬入 tile。 LinxISA TLOAD 额外支持带 ValidCol/ValidRow/Col/PadValue。 处理尾块/matmul_mask功能/padding时需要考虑把 LinxISA 的 valid col/row、 Col/Row、PadValue 转成 PTO t额外的 padding 语义;load-time padding 不能只按参数替换。 基本可对齐 PTO tload;当前无需补 LinxISA 额外 PadValue/Dep 等特性。若 PTO 侧使用更丰富的 padding 语义,再补对应实现。
TEXPANDSCALAR %dst = texpands %scalar : <scalar type> -> !pto.tile<...> TEXPANDSCALAR <LB0: ValidCol, LB1: ValidRow, LB2: Col, DataType>, scalar, ->DstTile<Size> 把一个标量广播/填充成 tile。 - 基本一致。 可对齐 PTO texpands,暂无需要补充的 PTO 功能。
MATMULMX %c = tmatmul.mx %a, %a_scale, %b, %b_scale : (!pto.tile<...>, !pto.tile<...>, !pto.tile<...>, !pto.tile<...>) -> !pto.tile<...> TMATMULMX <LB0: M, LB1: N, LB2: K, DataTypeA, DataTypeB> SrcTile0<.reuse>, SrcScale0, SrcTile1<.reuse>, SrcScale1, ->ACC<Size> MX/block-scale 矩阵乘。 - PTO似乎暂未支持hif4。 可对齐 PTO tmatmul.mx 的 A/B scale 输入形式;暂无需要补充的 PTO 功能。
TMOV_DN2NZ %dst = tmov.* %src : !pto.tile<...> -> !pto.tile<...><br>或组合:tmov/treshape/ttrans/tpack/tstore+tload TMOV DN2NZ, <LB0: ValidCol, LB1: ValidRow, LB2: Col, DataType, PadValue>, SrcTile<.reuse>, ->DstTile<Size> Tile 内 layout 转换,从 DN/ND 类普通布局转成 NZ 分形布局。 - 基本一致 PTO 侧没有必须一一对应的 DN2NZ 专用语义;可通过现有 layout conversion/load 相关实现对齐,暂无需补 LinxISA 专有 layout 字段。
rowmax %dst = trowmax %src : !pto.tile<...> -> !pto.tile<...> TROWMAX <LB0: ValidCol, LB1: ValidRow, LB2: Col, DataType>, SrcTile, ->DstTile<Size> 对每一行做最大值归约。 - 基本一致 可对齐 PTO trowmax,暂无需要补充的 PTO 功能。
tsub %dst = tsub %lhs, %rhs : (!pto.tile<...>, !pto.tile<...>) -> !pto.tile<...> TSUB <LB0: ValidCol, LB1: ValidRow, LB2: Col, DataType>, Src0, Src1, ->DstTile<Size> 逐元素相减。 - 基本一致 可对齐 PTO tsub,暂无需要补充的 PTO 功能。
texp %dst = texp %src : !pto.tile<...> -> !pto.tile<...> TEXP <LB0: ValidCol, LB1: ValidRow, LB2: Col, DataType>, SrcTile, ->DstTile<Size> 逐元素指数函数。 - 基本一致 可对齐 PTO texp,暂无需要补充的 PTO 功能。
rowsum %dst = trowsum %src : !pto.tile<...> -> !pto.tile<...> TROWSUM <LB0: ValidCol, LB1: ValidRow, LB2: Col, DataType>, SrcTile, ->DstTile<Size> 对每一行做求和归约。 - 基本一致 可对齐 PTO trowsum,暂无需要补充的 PTO 功能。
tdiv %dst = tdiv %lhs, %rhs : (!pto.tile<...>, !pto.tile<...>) -> !pto.tile<...> TDIV <LB0: ValidCol, LB1: ValidRow, LB2: Col, DataType>, Src0, Src1, ->DstTile<Size> 逐元素除法。 - 基本一致 可对齐 PTO tdiv,暂无需要补充的 PTO 功能。
tmul %dst = tmul %lhs, %rhs : (!pto.tile<...>, !pto.tile<...>) -> !pto.tile<...> TMUL <LB0: ValidCol, LB1: ValidRow, LB2: Col, DataType>, Src0, Src1, ->DstTile<Size> 逐元素乘法。 基本一致 基本一致 可对齐 PTO tmul,暂无需要补充的 PTO 功能。
tadd %dst = tadd %lhs, %rhs : (!pto.tile<...>, !pto.tile<...>) -> !pto.tile<...> TADD <LB0: ValidCol, LB1: ValidRow, LB2: Col, DataType>, Src0, Src1, ->DstTile<Size> 逐元素加法。 基本一致 基本一致 可对齐 PTO tadd,暂无需要补充的 PTO 功能。

Metadata

Metadata

Assignees

No one assigned

    Labels

    enhancementNew feature or request

    Type

    No type

    Fields

    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