PyPTO IR 类型与示例¶
本文档介绍类型 (Type) 系统并提供实用的使用示例。
类型系统¶
ScalarType¶
表示原始标量类型。
from pypto import DataType, ir
int_type = ir.ScalarType(DataType.INT64)
float_type = ir.ScalarType(DataType.FP32)
支持的 DataType: INT8, INT16, INT32, INT64, UINT8, UINT16, UINT32, UINT64, FP16, FP32, FP64, BOOL, INDEX, TASK_ID
注意:
INDEX是用于索引计算(循环变量、维度、偏移量、步长)的独立整数类型。它拥有自己的类型代码和字符串表示("index")。虽然语义上与INT64类似,但INDEX != INT64—— 它们是不同的类型。在代码生成中,INDEX 和 INT64 之间的隐式类型转换会被抑制。注意:
TASK_ID是一个不透明的 64-bit handle(类型代码0x50),表示 runtime 的TaskId。它不是数值类型——上面没有任何算术运算。Scalar[TASK_ID]值由with pl.manual_scope():内的pl.submit(...)产生(它返回的二元组第二个元素命名 producer task)。Python 字面量None是 "暂无 producer" 的哨兵——它用作 TaskId 循环 iter_arg 的种子,也可作为deps=[None]条目;当None出现在 TaskId 位置时,会下沉为system.task_invalidbuiltin →TaskId::invalid()。TaskId 值通过pl.submit(...)的deps=[tid1, tid2]kwarg 传入。codegen 把TASK_ID下沉为TaskId。
寄存器 IR 描述符¶
RegType(shape, dtype, kind) 直接继承 Type。shape 是表示逻辑 lane 数的单个
正整数,与物理向量宽度无关;lanes 直接返回该整数。RegKind.VecReg 禁止
BOOL,RegKind.MaskReg 必须使用 BOOL。非法 shape、未知 dtype 和未知 kind 在构造或
反序列化时立即失败。
vector = ir.RegType(64, DataType.FP32, ir.RegKind.VecReg)
mask = ir.RegType(64, DataType.BOOL, ir.RegKind.MaskReg)
标注形式为 pl.Reg[64, pl.FP32](默认 VecReg)和
pl.Reg[64, pl.BOOL, pl.RegKind.MaskReg]。这些描述符、无专属字段的
ir.RegScopeStmt(name_hint, body, span, leading_comments=[], attrs=None) 和
ir.ForStmt.hardware_loop_policy 属于 IR 基础设施,支持直接构造、visitor/mutator
遍历、结构比较和二进制序列化。P0 寄存器操作和 pl.reg_scope() 支持经过验证的
DSL 构造和文本往返;参见寄存器操作。
VF 代码生成仍由 PR3 实现。
HardwareLoopPolicy.Auto 是新循环及缺少该字段的旧序列化循环的默认值。
Require 是持久化的语义字段,不能被静默重置为 Auto;未知序列化 policy
必须失败。本次基础设施改动不实现 hardware loop 分析或生成。
Simplify 可消除已证明零次执行的 Sequential Require 循环,并返回初始携带值;非空 Require(包括单次循环)必须保留。
内部 Buffer 类型¶
BufferType 描述最终设备 IR 中的可变片上缓冲区,直接继承 Type,
不包含 MemRef、base pointer、地址或运行时表达式字段。
存储身份由定义该缓冲区的 SSA 值表达,所有权由定义算子声明。
buffer_type = ir.BufferType(
[32, 64], DataType.FP32, ir.Mem.Vec, valid_shape=[-1, 64]
)
multi_type = ir.MultiBufferType(buffer_type, slot_count=2)
物理维度目前必须是静态正整数。valid_shape 可以是零到物理维度之间的
静态有效长度,也可以用 -1 标记由算子操作数提供的运行时有效长度;
省略时使用完整物理形状。布局、以字节为单位的 fractal 大小、padding
和 compact mode 都是显式描述符字段。MultiBufferType 描述一次
多缓冲分配中的相同槽位,槽位数必须为正数。控制流类型检查会在分支结果
和循环携带值之间比较完整描述符,包括槽位数及嵌套 tuple 中的元素。
VoidType 表示确定没有 SSA 结果,与 UnknownType 不同。
Void call 应放在 EvalStmt 中,不能绑定变量、用作操作数、放入 tuple,
也不能作为值 yield 或 return。分配大小表达式(包括 WindowBuffer.size)
也必须产生一个值。
Call 和 Submit 的 attrs、kwargs 中的表达式值同样遵循此规则,
在构造时即进行校验,通过 ir.set_call_attrs 附加属性时也会校验。
这些类型是 Buffer IR 的初始基础设施,自动 tile-to-buffer lowering 尚未
启用,公开 Tile DSL 和默认流水线仍使用 TileType。直接 PTO 代码生成支持
显式构造的 buffer 程序:一维或二维、稠密 row-major Vec FP16/BF16/FP32/INT32
描述符,标量参数与控制流,以及下文的 buffer 算子。普通 GM 参数支持
物理形状静态、列数大于一的稠密 ND 二维 FP16/BF16/FP32/INT32 Tensor。
算术配方分别施加更窄的 dtype 契约。规范化后的 Tensor
返回值是这些参数的别名,原生内核仍返回 void。Buffer 参数 ABI、原生函数
结果、分形子视图、分配时的可变视图元数据、slot、helper 和其他物理布局尚未支持,会明确报错。
Buffer 类型 dump 使用原生 pypto.ir.BufferType(...) 构造表达式,
二进制序列化保留完整描述符。目前尚不支持通过 DSL parser 重新解析
完整的 buffer 程序 dump。
Buffer 算子契约¶
算子注册默认为 OpIRStage::Functional。内部 buffer 算子显式选择
OpIRStage::Buffer 并调用 set_internal_only()。结果数量必须声明:
零结果要求 VoidType,单结果使用原生类型,多结果使用数量匹配的
TupleType。Functional 注册仍按现有契约要求至少一个结果。
每个 buffer 操作数通过 set_buffer_arg_effect(i, data, metadata)
分别声明数据与元数据访问;标量操作数使用 set_buffer_non_memory_arg(i)。
两个访问维度都使用 BufferAccess(None、Read、Write、ReadWrite),
缺少声明不会默认视为读取。set_buffer_result_behavior(...) 将结果分类为
分配、别名、借用句柄或原生值,void call 声明 None。别名和借用结果
需要指定来源操作数。Allocate 声明 root 句柄;指定地址的 root 可以相互
重叠,因此它不证明存储独立或已初始化。描述符及内存空间合法性由各算子的
类型推导或显式结果验证检查。
buffer.alloc 使用 f_validate_explicit_type(...),不使用类型推导器:
物理描述符只存在于 Call.type 中。两种模式互斥,避免在 kwargs 中重复
存储描述符。私有 IR 构造器在 span 前接收结果类型:
from pypto.pypto_core import ir as _ir
span = ir.Span.unknown()
valid_rows = ir.Var("valid_rows", ir.ScalarType(DataType.INDEX), span)
descriptor = ir.BufferType([32, 64], DataType.FP32, ir.Mem.Vec, valid_shape=[-1, 64])
allocation = _ir._create_internal_op_call(
"buffer.alloc", [ir.MakeTuple([valid_rows], span)], {}, descriptor, span
)
第一个操作数始终是 MakeTuple,只包含描述符中 -1 维度对应的运行时
valid extent,按维度顺序排列。静态描述符使用空 tuple。可选的第二个操作数
是最终有效字节地址,不再叠加 base 或 offset。省略地址表示请求独立存储;
显式零地址是合法的指定地址分配。负常量地址(包括 -1)会被拒绝。
两个操作数都属于非内存值,运行时值必须是整数或 INDEX 标量。
常量 valid extent 必须介于零和对应物理维度之间;无法静态检查时,运行时
extent 的边界及地址非负性属于构造调用的前置条件。
buffer.set_validshape(buffer, valid_extents) 返回 VoidType,只写入元数据。
其 MakeTuple 操作数包含所有维度。标记为 -1 的维度可在物理边界内
变化;静态 valid 维度必须传入与描述符一致的常量。该操作不改变不可变类型
或 buffer 身份。对于句柄生命周期内会变化的 valid 维度,lowering 必须提前
选择动态描述符。
初始 PTO emitter 还按原生指令要求,限定 set_validshape 使用二维且两个
valid 维度均为动态的描述符(valid_shape=[-1, -1])。固定的初始 extent
仍可作为该动态描述符的常量操作数。分配支持混合静态/动态描述符,但发射
元数据更新时不能隐式将静态维度改为动态。
OpRegistry::ValidateBufferCall 按创建调用时的同一 schema 检查已有 call,
包括其原始结果类型和 kwargs。BufferIR 属性在设备函数中应用此检查,并
拒绝逻辑 tile、隐式 buffer 别名以及携带 buffer 的控制流结果。
它组合 SSA、先定义后使用及赋值类型检查,建立表示层契约;存储生命周期、
重叠和初始化证明属于后续验证。
buffer.copy(src, dst)、buffer.mul(lhs, rhs, dst) 和 buffer.add(lhs, rhs, dst) 写入显式
destination 并返回 VoidType,目前要求所有参数的 Vec buffer 描述符
相同;例外是 buffer.copy 还可以把 Mat 中的 cube 操作数搬入 Left 或 Right:
元素类型与物理/valid extent 保持不变,目标描述符给出自己的分形布局。写效应并不表示所有字节都已初始化。允许输入和 destination 是同一
句柄;构造调用前需要保证运行时 valid extent 一致,并完成部分重叠 view
的合法化。现有 Functional 阶段的
ArgEffect 查询会显式拒绝 buffer 算子,buffer 消费方必须使用
GetBufferArgEffect。直接 codegen 验证 BufferIR 后,按这些调用及其
BufferType 发射分配、destination 写入与 valid 状态更新。地址发射只取决于
分配操作数;旧 emit_tile_addr 标志不能删除或补充 buffer 地址。
动态操作数保留在其词法作用域内,不重建逻辑 TileType 或 MemRef,
也不运行隐式 tile 分配逻辑。设置 enable_buffer_ir=True 后,流水线在发射前通过
LowerTileToBuffer 转换 Tile IR。
静态存储视图 (static storage view) 在同一 Buffer 阶段使用普通 SSA 别名边:
%root = buffer.alloc(()) : Buffer<[128, 32], UINT8, Vec>
%bytes = buffer.subview(%root, (2, 0)) : Buffer<[64, 32], UINT8, Vec>
%tile = buffer.reshape(%bytes) : Buffer<[16, 32], FP32, Vec>
buffer.subview(source, (row, col)[, valid]) 是稠密行主序二维 Vec buffer 上的带步长窗口:
结果保持源的元素类型和行距,其静态形状必须处于源之内。偏移为整数或 INDEX 标量,
可以是常量或运行时值;常量偏移会做边界检查。可选的 valid 元组给出结果每一维的 valid extent
(静态描述符维度使用常量),结果有动态维度或静态 valid extent 小于物理形状时必须提供。
常量 valid 操作数必须为正且不大于对应窗口维度;暂不支持 valid extent 为零的窗口。
存储根上整行的 UINT8[N,32]
字节窗口是其中的静态特例。
buffer.reshape 保持物理字节数完全相同,支持静态、紧密行主序的二维
Vec FP16、BF16、FP32、INT16、INT32 和 UINT8 描述符,每个物理行的字节数必须为 32 的倍数。
在 Mat、Left、Right 和 Acc 中,它只在同一空间内重新标注一个完整分形窗口
(例如 NZ 到 ZN),要求元素类型、字节数相同且 valid extent 为静态值。
两者均声明 Alias(0)、无数据访问和源元数据读效应,不分配存储,
也不初始化数据。每个中间结果的形状、dtype 和 valid extent 均由结果
类型显式表达。buffer.set_validshape 不能修改这些静态视图句柄。
运行时偏移的窗口只知道位于源窗口之内,因此 BufferIR 在证明不重叠时使用整个源范围,
后续静态子视图和 reshape 也保留该保守范围。保守范围相同不能证明精确别名。
BufferIR 缓存视图的根身份、相对偏移、字节范围、继承行距及偏移确定性。
子视图保留父视图的行距,reshape 使用结果布局的行距。逐元素指令及涉及视图的
拷贝会检查源和目标窗口:同根使用相对偏移,不同已分配根使用有效地址。
不同的无地址分配互不重叠;未经证明的不同输入根不能作为不重叠证明。
同根窗口即使使用运行时基地址,仍可证明相对范围不重叠。此检查不建立
初始化或生命周期安全证明。输入非视图 buffer 之间的旧拷贝契约仍保留
原有重叠前提。已放置地址只在同一内存空间内比较;不同空间的窗口不会重叠。
Cube 操作从 Left/Right 操作数写入显式的 Acc 目标:
buffer.matmul(lhs_left, rhs_right, acc) : Void # acc = lhs @ rhs
buffer.matmul_acc(lhs_left, rhs_right, acc) : Void # acc = acc + lhs @ rhs
buffer.extract(src, row, col, dst) : Void # 静态形状窗口拷贝
操作数为二维物理描述符 [M, K] x [K, N] -> [M, N],两个输入元素类型相同,
为 FP16、BF16、FP32 或 INT8,累加器对应 FP32/INT32。静态的 rhs valid K
必须覆盖 lhs valid K。两种形式都恰好写入乘积矩形(lhs valid 行数 × rhs valid 列数),
按 PTOAS 的要求,静态目标的 valid extent 必须与之相等。buffer.matmul 覆盖写入该矩形;
buffer.matmul_acc 原地累加,没有单独的累加器输入。更宽的累加器通过同一存储上
与乘积同形的视图写入。Cube 按完整分形块(Acc 中为 16x16)写入:乘积之外的完整块保持原数据,
但被部分覆盖的块中其余部分不保留(在 a2a3 上实测为 0,与旧路径一致)。buffer.extract 在整数或 INDEX 偏移处
复制与目标形状相同的窗口,方向为 Mat 到 Left/Right 或 Vec 内部;常量偏移会做
边界检查,偏移是非内存操作数。
GM 传输将完整窗口表示为普通操作数:
buffer.load(tensor, offsets_tuple, valid_extents_tuple, dst_buffer) : Void
buffer.store(src_buffer, offsets_tuple, valid_extents_tuple, tensor) : Void
传输契约要求普通二维 Tensor/Vec Buffer 操作数,元素类型匹配且为
FP16、BF16、FP32 或 INT32。buffer.load 还可以填充 Mat 中的 cube 操作数
(FP16、BF16、FP32 或 INT8,传输把 ND 转成分形布局)。buffer.store 还可以
通过 fix-pipe 导出 Acc 累加器;不带 scale 的转换只有保持累加器类型,以及
FP32 转 FP16/BF16。偏移量是非负的
元素索引,两个 tuple 均包含两个整数或 INDEX 标量。传输长度必须等于
buffer 当前的 valid extent;静态描述符维度必须使用完全相同的常量。
常量长度和窗口会对照 buffer 容量、GM 物理形状及 Tensor 的有效区域检查;
有效区域取 TensorView.valid_shape,省略时取完整形状,传输不会扩展 Tensor 的 valid 元数据。
动态值的一致性、
非负性和边界则是 lowering/运行时前提。传输不修改 buffer valid 元数据,
也不初始化未触及的数据。源操作数声明数据读/元数据读效应,目标声明
数据写/元数据读效应;两个 tuple 都是非内存操作数。GM store 写入选定
区域,不读取未触及的其余区域。
直接发射目前要求 Tensor 操作数为函数参数,采用紧密 ND 步幅且无 padding。
load 要求参数方向为 In 或 InOut,store 要求 Out 或 InOut。原生 ABI
总是先排列 GM 指针参数、再排列标量,不受 IR 参数顺序影响。参数方向、
Tensor 返回类型及规范化的输出参数返回值保留在 IR 中供编排使用;区域结果
和 carry 仅支持标量。动态物理形状、非参数 GM view、缓存策略/atomic kwargs
及其他传输布局需要后续发射规则,目前均明确拒绝。
初始 emitter 支持将无符号标量直接作为分配地址、GM 偏移量或 valid extent,扩展位宽时
保留其无符号数值。无符号算术以及除 BOOL 到整数、index 到整数、同位宽整数转换之外的
无符号转换,需要后续补充原生发射规则,目前会在发射前拒绝。For 循环要求
归纳变量为 INDEX,边界为 INDEX 或有符号整数,步长必须可证明为正的常量。
运行时步长、零步长和负步长均在发射前拒绝。这些是 emitter 的限制,
并非 Buffer IR 表示本身的限制。
TensorType¶
带可选内存引用 (MemRef) 的多维张量 (Tensor)。
span = ir.Span.unknown()
# Tensor with shape [10, 20]
shape = [ir.ConstInt(10, DataType.INT64, span), ir.ConstInt(20, DataType.INT64, span)]
tensor_type = ir.TensorType(shape, DataType.FP32)
# Tensor with MemRef: base allocation, byte offset within it, size in bytes
memref = ir.MemRef("mem_ddr_0", 0, 800)
tensor_with_memref = ir.TensorType(shape, DataType.FP32, memref)
TensorType.memory_space 始终是 ir.Mem.DDR。MemRef 标识一块分配
(base_) 以及其中的一段字节区间 (byte_offset_、size_);内存空间不再
存储在 MemRef 本身上。完整字段列表见
内存引用 (MemRef)。
DistributedTensorType¶
DistributedTensorType 是 TensorType 的精确 ObjectKind 子类,作为 chip
orchestrator / InCore 形参的类型注解,用来切片由 CommDomainScopeStmt 划分的 HCCL window buffer。
它的存在让跨 rank op 的 verifier(后续 milestone 引入)可以静态拒绝普通的
Tensor 实参 —— As<TensorType> 不会匹配 DistributedTensorType
(精确 ObjectKind 匹配语义,见
ir-kind-traits.md),跨 rank op 用
As<DistributedTensorType> 派生。
DSL 形式是 pld.DistributedTensor[[shape], dtype]:
import pypto.language.distributed as pld
import pypto.language as pl
@pl.function(type=pl.FunctionType.InCore)
def kernel(self, data: pld.DistributedTensor[[256], pl.FP32]): ...
IR 层:
t = ir.DistributedTensorType([64], DataType.FP32)
assert isinstance(t, ir.TensorType) # C++ 继承关系保留
# As<TensorType>(t) → null;As<DistributedTensorType>(t) → 转型成功
分配侧的元数据(每 rank 大小、host staging 标志)挂在 pld.tensor.alloc_window_buffer
op 所绑定的 ir.WindowBuffer(Var 子类)上。通过
pld.tensor.window(buf, [shape], dtype=...) 物化的切片在
DistributedTensorType.window_buffer 上保留指向源 WindowBuffer 的可选反向
引用,从而让两个 shape/dtype 相同但分配来源不同的切片在结构上保持不同。
用户在签名中写的 pld.DistributedTensor[[shape], dtype] 不填该字段(为
None)。Tile 类型没有 distributed 变体;跨 rank op 始终作用在
DistributedTensor 上。
在 window 上做本地计算。 在 InCore scope 内,一个 window 切片就是本 rank 的
本地 GM,因此普通 tensor op 可以像读写任何 GM tensor 一样读写它。这些 op 用
AsTensorTypeLike(同时匹配两种
kind)而不是精确匹配的 As<TensorType> 来匹配操作数。结果类型取决于该 op 产生的
是 window 的视图还是新数据:
| 接受 window 的 op | 结果 kind |
|---|---|
tensor.slice、tensor.assemble、tensor.view、tensor.write |
DistributedTensorType —— 仍然是同一个 comm-group 分配上的视图 |
逐元素与一元族、各类 reduction、tensor.matmul、tensor.matmul_acc(仅 lhs / rhs) |
普通 TensorType —— 结果是新产生的本地数据 |
tensor.read |
ScalarType —— 单个元素,没有视图 |
两处已记录的拒绝:tensor.reinterpret_view 直接拒绝 window;tensor.matmul_acc
的 acc 操作数必须是普通 TensorType —— 只有矩阵单元会写 L0C,因此不存在从
window 到 Cube 累加器的数据通路。请先在本地累加,再把结果存回 window。
还有不少读写普通 GM 的 tensor op 目前仍然拒绝 window(全部 broadcast、reshape、
transpose、concat、gather / scatter 族等)。
tests/ut/ir/operators/test_window_operand_acceptance.py 保存了逐算子的权威分类,
并负责保证它与实现一致。
带 TensorView 的 TensorType¶
带有布局和步长信息的张量,用于优化内存访问。
# Create tensor with tensor view (stride/valid_shape accept int or Expr)
tensor_view = ir.TensorView(stride=[1, 128], layout=ir.TensorLayout.ND)
tensor_with_view = ir.TensorType([128, 256], DataType.FP32, memref=None, tensor_view=tensor_view)
# With valid_shape
tensor_view = ir.TensorView(stride=[1, 128], layout=ir.TensorLayout.ND, valid_shape=[64, 128])
# With pad mode for out-of-valid-shape accesses (symmetric with TileView)
tensor_view = ir.TensorView(
stride=[1, 128], layout=ir.TensorLayout.ND, valid_shape=[64, 128], pad=ir.PadValue.zero
)
# Different layouts
nd_view = ir.TensorView(stride=[1, 128], layout=ir.TensorLayout.ND) # ND layout
dn_view = ir.TensorView(stride=[1, 128], layout=ir.TensorLayout.DN) # DN layout
nz_view = ir.TensorView(stride=[1, 128], layout=ir.TensorLayout.NZ) # NZ layout
# Expr values also accepted (e.g., symbolic dimensions)
stride = [ir.ConstInt(1, DataType.INT64, span), ir.ConstInt(128, DataType.INT64, span)]
tensor_view = ir.TensorView(stride=stride, layout=ir.TensorLayout.ND)
# Tensor with both MemRef and TensorView
memref = ir.MemRef("mem_ddr_1", 0, 16384)
tensor_with_both = ir.TensorType([128, 256], DataType.FP16, memref=memref, tensor_view=tensor_view)
TensorLayout 值:
ND:ND 布局DN:DN 布局NZ:NZ 布局
TensorView 字段:
stride:每个维度的步长layout:TensorLayout.ND/DN/NZ/MX_A_ZZ/MX_B_NNvalid_shape:可选的有效区域维度(为空表示使用完整 shape)pad:PadValue.null(默认)/zero/max/min,用于访问超出valid_shape部分时的填充模式。与TileView.pad对称;tensor.slice(..., pad_value=PadValue.zero)会写入该字段。
Canonical TensorView 形式(RFC #1300)¶
按 RFC #1300 的设计,(shape, stride, layout) 三元组在各 pass / verifier /
codegen 之间统一为单一可机械读取的形式:
shape是逻辑 shape —— 消费者索引时使用的维度。stride[i]是第 i 个逻辑维递增 1 时的元素步长。layout是(shape, stride)上的派生标签 / 断言,并非独立描述。 ND / DN 各定义 packed canonical(紧致存储)与 strided 家族(sub-view 继承 父 stride)两种合法形态。
Packed canonical 公式(BuildLogicalStridesFromLayout,见
tensor_view_semantics.h):
| Layout | Packed canonical |
|---|---|
ND |
stride[n-1] = 1; stride[k] = stride[k+1] * shape[k+1] |
DN(n ≥ 2) |
stride[n-2] = 1;stride[n-1] = shape[n-2];stride[n-3] = shape[n-2] * shape[n-1];外层按行主序 |
NZ |
对分块后的 rank-5 shape [B, C/c0, R/16, 16, c0] 求行主序 —— 见 BlockNzTensorViews |
同一 canonical TensorView 的两种写法:
- 隐式 ——
view.has_value() && view.stride.empty():layout 已设但 stride 为空,消费者按对应 layout 的 packed canonical 解释。 - 显式 —— 每个维度的 stride 都已写出。
MaterializeTensorStrides Pass
将所有隐式形态展开为显式 packed canonical,让 codegen 看到单一契约。
TensorViewCanonical IRProperty + verifier 强制此不变量:
- 弱模式(registry 默认,
passes.PropertyVerifierRegistry.verify): 接受stride.empty()作为隐式 packed canonical。 - 严格模式(codegen 入口契约,
passes.verify_tensor_view_canonical(program, require_materialized=True)): 必须有非空view.stride且与 layout 家族一致。
两种模式都拒绝 TensorType 上未分块的 NZ shape,并按
relaxed_symbolic 语义接受符号 stride。
TileType¶
专用张量类型,带可选内存和视图信息,用于硬件优化操作。
# Basic 16x16 tile
shape = [ir.ConstInt(16, DataType.INT64, span)] * 2
tile_type = ir.TileType(shape, DataType.FP16)
# 3D tile (supported at IR level)
shape_3d = [ir.ConstInt(4, DataType.INT64, span),
ir.ConstInt(16, DataType.INT64, span),
ir.ConstInt(16, DataType.INT64, span)]
tile_type_3d = ir.TileType(shape_3d, DataType.FP16)
# Tile with MemRef and TileView. TileView is immutable — every field is passed
# to the constructor; valid_shape / stride / start_offset accept int or Expr.
memref = ir.MemRef("mem_left_0", 0, 512)
tile_view = ir.TileView(valid_shape=[8, 16], stride=[1, 16], start_offset=0)
tile_with_view = ir.TileType(shape, DataType.FP16, memref, tile_view, ir.Mem.Left)
TileType.memory_space 才是 Tile 放置位置的唯一来源。如果 TileType
携带 MemRef, 请在 TileType 自身上显式提供 tile 内存空间。
上面的 valid_shape 是一个真正的子区域([16, 16] Tile 中的 [8, 16])。
若 valid_shape 与完整 shape 相同则是冗余的,构造函数会将其清空 ——
参见下文的规范化规则。
对于 Python DSL 类型标注,省略的 TileView 语法会被规范化为一个隐式
TileView:它由 tile shape 以及(如果存在)tile memory space 推导得到。
像 pl.TileView() 这样的冗余显式默认写法,会与省略写法被视为语义等价,
并且在 printer 输出时可能统一成规范形式。TileView.compact 记录部分有效的
boxed tile 是采用 PTO 的有效区域紧凑表示(CompactMode.normal),还是普通的
物理 box 表示(默认的 CompactMode.null)。它只在 fractal 空间——Left / Right /
Acc——有意义,因为它本身描述的就是 N-fractal pitch;AccCompactValid 校验器会拒绝
其它空间上的 compact。编译器会为进入 L0A/L0B 的部分 tile.extract、以及行窄化的
matmul 累加器(其 L0C pitch 由 mad 按 L0A 操作数的有效行数推导)自动设置该字段,
AutoTileMatmulL0 还会通过 tile.create(..., compact=True) 在它合成的累加器种子上
声明该字段。普通用户代码无需手动选择。
隐式 view 依赖 memory space,构造函数只会针对传入的 space 把 view 折叠成
nullopt。凡能确定结果 space 的 f_deduce_type,都必须把该 space 传进来:
若先针对 nullopt 推导、再由 OpRegistry::Create 补盖 space,就会按两套不同的
隐式 layout 规范化两次,结果取决于 view 是否恰好折叠(即 valid_shape 与 pad)。
ArrayType¶
片上定长同构 1-D 数组,存放于标量寄存器堆 / C 栈(memory space ScalarLocal)。
区别于 TensorType(GM/DDR 指针)和 TileType(向量/cube 单元状态)。
v1 约束:
- 元素 dtype 必须是整型(
INT8/16/32/64、UINT8/16/32/64)或BOOL - 仅支持 rank-1;extent 必须是编译期
ConstInt - 不携带
MemRef—— codegen 直接落到 C 栈数组dtype name[N](无 STL 依赖) - 不能跨函数边界(由
ArrayNotEscaped验证器强制)
操作:
| Op | 语义 | Orchestration(C++) | InCore(.pto) |
|---|---|---|---|
array.create(N, dtype) |
分配栈数组 | dtype arr[N] = {0}; |
pto.declare_local_array -> !pto.local_array<NxT> |
array.get_element(arr, i) → Scalar |
读元素 i |
dtype v = arr[i]; |
pto.local_array_get arr[i] : !pto.local_array<NxT> -> T |
array.update_element(arr, i, v) → Array |
函数式更新(SSA-pure) | arr[i] = v;(LHS 别名到入参) |
pto.local_array_set arr[i], v : !pto.local_array<NxT>, T |
array.update_element 是 tensor.assemble 的 SSA-functional 等价物:返回一个新的
ArrayType SSA 值,表示"原数组中第 i 个元素被替换为 v"。两条 codegen 路径都把结果
Var 别名到入参数组的存储,emit 原地写入 —— 不复制。
InCore 路径对齐 PTOAS 的栈数组三件套(pto.declare_local_array /
pto.local_array_get / pto.local_array_set)。下标统一下降为 MLIR index(源类型
非 index 时插入 arith.index_cast),set 的值在与元素 dtype T 不一致时也会被
cast(verifier 允许把 index 类型的值写入整型数组)。
DSL 下标糖:
arr = pl.array.create(8, pl.INT32)
arr[i] = v # desugar 成: arr = pl.array.update_element(arr, i, v)
x = arr[i] # desugar 成: x = pl.array.get_element(arr, i)
arr[i] = v 时 parser 把左边变量重绑定,后续读取看到更新后的数组 —— 与
Tensor/Tile 下标写入糖一致。
TupleType¶
异构类型元组。
# Scalar tuple: (int, float)
scalar_tuple = ir.TupleType([
ir.ScalarType(DataType.INT64),
ir.ScalarType(DataType.FP32)
])
# Nested tuple
nested = ir.TupleType([
ir.TupleType([ir.ScalarType(DataType.INT64)]),
ir.ScalarType(DataType.FP32)
])
PipeType¶
硬件执行流水线或同步屏障。
pipe_s = ir.PipeType(ir.PipeType.S) # Scalar pipe
pipe_v = ir.PipeType(ir.PipeType.V) # Vector pipe
pipe_m = ir.PipeType(ir.PipeType.M) # Matrix pipe
pipe_all = ir.PipeType(ir.PipeType.ALL) # All pipes
UnknownType¶
未知或待推断类型的占位符。
DSL 中的 MemRef 类型注解¶
MemRef 可以在 @pl.program / @pl.function DSL 代码中作为位置参数指定在类型注解中:
import pypto.language as pl
@pl.program
class MyProgram:
@pl.function(type=pl.FunctionType.InCore)
def kernel(self, x: pl.Tensor[[64, 64], pl.FP32]):
# Tile with MemRef and explicit tile memory space
tile_a: pl.Tile[[64, 64], pl.FP32, pl.MemRef(0, 16384, 0), pl.Mem.Vec] = pl.tile.load(
x, offsets=[0, 0], shapes=[64, 64]
)
# Tensor with MemRef (3-arg: shape, dtype, memref)
y: pl.Tensor[[64, 64], pl.FP32, pl.MemRef(0, 16384, 1)] = pl.add(x, 1.0)
# Tensor with layout and MemRef (4-arg: shape, dtype, layout, memref)
z: pl.Tensor[[64, 64], pl.FP32, pl.NZ, pl.MemRef(0, 16384, 2)] = pl.add(x, 1.0)
pl.MemRef(addr, size, id) 参数:
| 参数 | 类型 | 说明 |
|---|---|---|
addr |
int |
基地址偏移 |
size |
int |
内存分配大小(字节) |
id |
int |
内存缓冲区标识符 |
TensorType 注解默认位于 DDR。为了兼容旧代码,解析器仍接受
pl.MemRef(pl.Mem.DDR, addr, size, id),但新代码应优先使用 3 参数形式。
消歧义(3 参数 Tensor): 解析器会自动区分 pl.MemRef(...) 和
pl.NZ/pl.DN/pl.ND 布局枚举。
Tile 规则: 如果在 pl.Tile[...] 注解中使用 pl.MemRef(...),必须再
单独提供一个 pl.Mem.* 参数来声明 tile 的内存空间。
MemorySpace 枚举(别名:Mem)¶
| 值 | 说明 |
|---|---|
DDR |
主存储器(片外) |
Vec |
向量/统一缓冲区(片上) |
Mat |
矩阵/L1 缓冲区 |
Left |
左矩阵操作数缓冲区 |
Right |
右矩阵操作数缓冲区 |
Acc |
累加器缓冲区 |
Bias |
Bias 缓冲区 |
ScalarLocal |
片上标量寄存器堆 / C 栈(用于 ArrayType) |
Python 使用示例¶
示例 1:构建表达式¶
from pypto import DataType, ir
span = ir.Span.unknown()
dtype = DataType.INT64
# Variables and constants
x = ir.Var("x", ir.ScalarType(dtype), span)
y = ir.Var("y", ir.ScalarType(dtype), span)
one = ir.ConstInt(1, dtype, span)
two = ir.ConstInt(2, dtype, span)
# Build: ((x + 1) * (y - 2)) / (x + y)
x_plus_1 = ir.Add(x, one, dtype, span)
y_minus_2 = ir.Sub(y, two, dtype, span)
numerator = ir.Mul(x_plus_1, y_minus_2, dtype, span)
denominator = ir.Add(x, y, dtype, span)
result = ir.FloatDiv(numerator, denominator, dtype, span)
示例 2:控制流(绝对值)¶
# if (x >= 0) then { result = x } else { result = -x }
x = ir.Var("x", ir.ScalarType(dtype), span)
result = ir.Var("result", ir.ScalarType(dtype), span)
zero = ir.ConstInt(0, dtype, span)
condition = ir.Ge(x, zero, dtype, span)
then_assign = ir.AssignStmt(result, x, span)
else_assign = ir.AssignStmt(result, ir.Neg(x, dtype, span), span)
abs_stmt = ir.IfStmt(condition, then_assign, else_assign, [result], span)
示例 3:带累加的循环¶
# for i, (sum,) in pl.range(n, init_values=(0,)):
# sum = pl.yield_(sum + i)
n = ir.Var("n", ir.ScalarType(dtype), span)
i = ir.Var("i", ir.ScalarType(dtype), span)
zero = ir.ConstInt(0, dtype, span)
one = ir.ConstInt(1, dtype, span)
sum_iter = ir.IterArg("sum", ir.ScalarType(dtype), zero, span)
add_expr = ir.Add(sum_iter, i, dtype, span)
yield_stmt = ir.YieldStmt([add_expr], span)
sum_final = ir.Var("sum_final", ir.ScalarType(dtype), span)
loop = ir.ForStmt(i, zero, n, one, [sum_iter], yield_stmt, [sum_final], span)
示例 4:带运算符调用的函数¶
# def matmul(a, b) -> tensor:
# result = tensor.matmul(a, b, out_dtype=FP32)
shape_m = ir.ConstInt(128, DataType.INT64, span)
shape_k = ir.ConstInt(64, DataType.INT64, span)
shape_n = ir.ConstInt(256, DataType.INT64, span)
a = ir.Var("a", ir.TensorType([shape_m, shape_k], DataType.FP16), span)
b = ir.Var("b", ir.TensorType([shape_k, shape_n], DataType.FP16), span)
matmul_call = ir.op.tensor.matmul(a, b, out_dtype=DataType.FP32)
result = ir.Var("result", ir.TensorType([shape_m, shape_n], DataType.FP32), span)
body = ir.AssignStmt(result, matmul_call, span)
return_types = [ir.TensorType([shape_m, shape_n], DataType.FP32)]
func = ir.Function("matmul", [a, b], return_types, body, span)
示例 5:包含多个函数的程序¶
# Helper: square(x) -> int { return x * x }
x = ir.Var("x", ir.ScalarType(dtype), span)
square_result = ir.Var("result", ir.ScalarType(dtype), span)
square_body = ir.AssignStmt(square_result, ir.Mul(x, x, dtype, span), span)
square_func = ir.Function("square", [x], [ir.ScalarType(dtype)], square_body, span)
# Main: sum_squares(a, b) -> int { return square(a) + square(b) }
a = ir.Var("a", ir.ScalarType(dtype), span)
b = ir.Var("b", ir.ScalarType(dtype), span)
program = ir.Program([square_func], "math", span)
square_gvar = program.get_global_var("square")
call_a = ir.Call(square_gvar, [a], span)
call_b = ir.Call(square_gvar, [b], span)
sum_expr = ir.Add(call_a, call_b, dtype, span)
main_result = ir.Var("result", ir.ScalarType(dtype), span)
main_body = ir.AssignStmt(main_result, sum_expr, span)
main_func = ir.Function("sum_squares", [a, b], [ir.ScalarType(dtype)], main_body, span)
program = ir.Program([square_func, main_func], "math", span)
示例 6:使用 TileType 的内存布局¶
# 32x32 tile in Left memory, viewing a 16x32 valid region with custom stride
shape = [ir.ConstInt(32, DataType.INT64, span)] * 2
memref = ir.MemRef("mem_left_0", 0, 2048)
tile_view = ir.TileView(valid_shape=[16, 32], stride=[1, 32], start_offset=0)
tile_type = ir.TileType(shape, DataType.FP16, memref, tile_view, ir.Mem.Left)
类型系统总结¶
| 类型 | 维度 | 内存信息 | 使用场景 |
|---|---|---|---|
| ScalarType | 0 | - | 单个值 |
| TensorType | N(任意) | 可选 MemRef | 通用张量 |
| TileType | N(任意)* | 可选 MemRef + TileView | 硬件优化 Tile |
| BufferType | 静态物理维度 | 定义它的 SSA 句柄 | 显式设备存储 |
| MultiBufferType | 元素 BufferType | 原生 slot 组 | 相同描述符的 buffer slots |
| VoidType | - | - | 确定不存在 SSA 结果 |
| TupleType | - | - | 多返回值 |
| PipeType | - | - | 硬件同步 |
| UnknownType | - | - | 类型推断占位符 |
常用模式¶
创建常量:
创建运算符:
# High-level API (recommended)
call = ir.op.tensor.matmul(a, b, out_dtype=DataType.FP32)
# Generic operator with kwargs
call = ir.create_op_call("tensor.matmul", [a, b], {"out_dtype": DataType.FP32}, span)
语句序列:
类型检查与转换¶
# Check expression types
if isinstance(expr, ir.Var):
print(expr.name_)
# Check type objects
if isinstance(type_obj, ir.TileType):
# Access tile-specific properties
shape = type_obj.shape
相关文档¶
总结¶
PyPTO 的类型系统提供:
- 标量类型 用于原始值
- 张量/Tile 类型 用于带内存布局的多维数据
- 元组类型 用于异构集合
- 流水线类型 用于硬件同步
IR 构建 API 支持:
- 通过共享指针创建不可变节点
- 带编译时检查的类型安全操作
- 通过 MemRef 和 TileView 实现硬件感知的内存管理
- 通过 GlobalVar 实现程序内函数调用
- 通过 IterArg 实现循环携带依赖