TMATMUL¶
指令示意图¶
简介¶
矩阵乘法 (GEMM),生成累加器/输出 Tile。
数学语义¶
设:
M = aMatrix.GetValidRow()K = aMatrix.GetValidCol()N = bMatrix.GetValidCol()
对于 0 <= i < M 和 0 <= j < N(有效矩阵乘法域中的输出元素):
\[ \mathrm{C}_{i,j} = \sum_{k=0}^{K-1} \mathrm{A}_{i,k} \cdot \mathrm{B}_{k,j} \]
精确的累加器行为和数据类型提升由目标/实现定义。
汇编语法¶
PTO-AS 形式:参见 PTO-AS 规范。
同步形式:
%acc = tmatmul %a, %b : (!pto.tile<...>, !pto.tile<...>) -> !pto.tile<...>
AS Level 1(SSA)¶
%c = pto.tmatmul %a, %b : (!pto.tile<...>, !pto.tile<...>) -> !pto.tile<...>
AS Level 2(DPS)¶
pto.tmatmul ins(%a, %b : !pto.tile_buf<...>, !pto.tile_buf<...>) outs(%c : !pto.tile_buf<...>)
C++ 内建接口¶
声明于 include/pto/common/pto_instr.hpp:
template <typename TileRes, typename TileLeft, typename TileRight, typename... WaitEvents>
PTO_INST RecordEvent TMATMUL(TileRes &cMatrix, TileLeft &aMatrix, TileRight &bMatrix, WaitEvents &... events);
template <AccPhase Phase, typename TileRes, typename TileLeft, typename TileRight, typename... WaitEvents>
PTO_INST RecordEvent TMATMUL(TileRes &cMatrix, TileLeft &aMatrix, TileRight &bMatrix, WaitEvents &... events);
约束¶
- 实现检查 (A2A3):
- 支持的
(CType, AType, BType)三元组: (int32_t, int8_t, int8_t)(float, half, half)(float, float, float)(float, bfloat16_t, bfloat16_t)- 静态形状约束:
TileLeft::Rows == TileRes::Rows、TileLeft::Cols == TileRight::Rows、TileRight::Cols == TileRes::Cols。 - Tile 位置:
TileLeft::Loc == Left、TileRight::Loc == Right、TileRes::Loc == Acc。 - 运行时:
m/k/n(取自aMatrix.GetValidRow()、aMatrix.GetValidCol()、bMatrix.GetValidCol())必须在[1, 4095]范围内。
- 支持的
- 实现检查 (A5):
- 累加器类型必须是
int32_t或float。 - 如果是
int32_t:AType == int8_t且BType == int8_t。 - 如果是
float:支持half/bfloat16_t/float和选定的 fp8 对(目标定义)。 - 静态形状约束:
TileLeft::Rows == TileRes::Rows、TileLeft::Cols == TileRight::Rows、TileRight::Cols == TileRes::Cols。 - 强制执行分形/布局约束:
- Left:
Loc == Left、!isRowMajor、SFractal == RowMajor - Right:
Loc == Right、isRowMajor、SFractal == ColMajor - Acc:
Loc == Acc、!isRowMajor、SFractal == RowMajor - 运行时:
m/k/n(取自aMatrix.GetValidRow()、aMatrix.GetValidCol()、bMatrix.GetValidCol())必须在[1, 4095]范围内。
- 累加器类型必须是
示例¶
自动(Auto)¶
#include <pto/pto-inst.hpp>
using namespace pto;
void example_auto() {
using A = TileLeft<half, 16, 16>;
using B = TileRight<half, 16, 16>;
using C = TileAcc<float, 16, 16>;
A a;
B b;
C c;
TMATMUL(c, a, b);
}
手动(Manual)¶
#include <pto/pto-inst.hpp>
using namespace pto;
void example_manual() {
using A = TileLeft<half, 16, 16>;
using B = TileRight<half, 16, 16>;
using C = TileAcc<float, 16, 16>;
A a;
B b;
C c;
TASSIGN(a, 0x1000);
TASSIGN(b, 0x2000);
TASSIGN(c, 0x3000);
TMATMUL(c, a, b);
}
汇编示例(ASM)¶
自动模式¶
# 自动模式:由编译器/运行时负责资源放置与调度。
%c = pto.tmatmul %a, %b : (!pto.tile<...>, !pto.tile<...>) -> !pto.tile<...>
手动模式¶
# 手动模式:先显式绑定资源,再发射指令。
# 可选(当该指令包含 tile 操作数时):
# pto.tassign %arg0, @tile(0x1000)
# pto.tassign %arg1, @tile(0x2000)
%c = pto.tmatmul %a, %b : (!pto.tile<...>, !pto.tile<...>) -> !pto.tile<...>
PTO 汇编形式¶
%acc = tmatmul %a, %b : (!pto.tile<...>, !pto.tile<...>) -> !pto.tile<...>
# AS Level 2 (DPS)
pto.tmatmul ins(%a, %b : !pto.tile_buf<...>, !pto.tile_buf<...>) outs(%c : !pto.tile_buf<...>)