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 的 PTO2TaskId。它不是数值类型——上面没有任何算术运算。Scalar[TASK_ID] 值由 with pl.manual_scope(): 内的 pl.submit(...) 产生（它返回的二元组第二个元素命名 producer task）。Python 字面量 None 是 "暂无 producer" 的哨兵——它用作 TaskId 循环 iter_arg 的种子，也可作为 deps=[None] 条目；当 None 出现在 TaskId 位置时，会下沉为 system.task_invalid builtin → PTO2TaskId::invalid()。TaskId 值通过 pl.submit(...) 的 deps=[tid1, tid2] kwarg 传入。codegen 把 TASK_ID 下沉为 PTO2TaskId。

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
memref = ir.MemRef(ir.ConstInt(0x1000, DataType.INT64, span), 800, 0)
tensor_with_memref = ir.TensorType(shape, DataType.FP32, memref)

TensorType.memory_space 始终是 ir.Mem.DDR。MemRef 只保存地址、大小和 id；内存空间不再存储在 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 上。

带 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(ir.ConstInt(0x2000, DataType.INT64, span), 16384, 1)
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
• valid_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	无法用 flat stride 表达 —— 仅 tile 用，分形布局

同一 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（NZ 仅 tile 使用），并按 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
memref = ir.MemRef(ir.ConstInt(0, DataType.INT64, span), 512, 0)

tile_view = ir.TileView()
tile_view.valid_shape = [ir.ConstInt(16, DataType.INT64, span)] * 2
tile_view.stride = [ir.ConstInt(1, DataType.INT64, span), ir.ConstInt(16, DataType.INT64, span)]
tile_view.start_offset = ir.ConstInt(0, DataType.INT64, span)

tile_with_view = ir.TileType(shape, DataType.FP16, memref, tile_view, ir.Mem.Left)

TileType.memory_space 才是 Tile 放置位置的唯一来源。如果 TileType 携带 MemRef, 请在 TileType 自身上显式提供 tile 内存空间。

对于 Python DSL 类型标注，省略的 TileView 语法会被规范化为一个隐式 TileView：它由 tile shape 以及（如果存在）tile memory space 推导得到。 像 pl.TileView() 这样的冗余显式默认写法，会与省略写法被视为语义等价， 并且在 printer 输出时可能统一成规范形式。

ArrayType

片上定长同构 1-D 数组,存放于标量寄存器堆 / C 栈(memory space ScalarLocal)。 区别于 TensorType(GM/DDR 指针)和 TileType(向量/cube 单元状态)。

arr_type = ir.ArrayType(DataType.INT32, 16)       # 16 个 INT32 元素
# DSL 注解形式:
arr: pl.Array[16, pl.INT32]

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

未知或待推断类型的占位符。

unknown = ir.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 with custom stride
shape = [ir.ConstInt(32, DataType.INT64, span)] * 2
memref = ir.MemRef(ir.ConstInt(0, DataType.INT64, span), 2048, 0)

tile_view = ir.TileView()
tile_view.valid_shape = shape
tile_view.stride = [ir.ConstInt(1, DataType.INT64, span), ir.ConstInt(32, DataType.INT64, span)]
tile_view.start_offset = ir.ConstInt(0, DataType.INT64, span)

tile_type = ir.TileType(shape, DataType.FP16, memref, tile_view, ir.Mem.Left)

类型系统总结

类型	维度	内存信息	使用场景
ScalarType	0	-	单个值
TensorType	N（任意）	可选 MemRef	通用张量
TileType	N（任意）*	可选 MemRef + TileView	硬件优化 Tile
TupleType	-	-	多返回值
PipeType	-	-	硬件同步
UnknownType	-	-	类型推断占位符

常用模式

创建常量：

i32 = ir.ConstInt(42, DataType.INT32, span)
f32 = ir.ConstFloat(3.14, DataType.FP32, span)

创建运算符：

# 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)

语句序列：

seq = ir.SeqStmts([stmt1, stmt2, stmt3], 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

相关文档

• IR 概述 - 核心概念与设计原则
• IR 节点层次结构 - 完整节点类型参考
• 结构比较 - 相等性和哈希工具

总结

PyPTO 的类型系统提供：

• 标量类型 用于原始值
• 张量/Tile 类型 用于带内存布局的多维数据
• 元组类型 用于异构集合
• 流水线类型 用于硬件同步

IR 构建 API 支持：

• 通过共享指针创建不可变节点
• 带编译时检查的类型安全操作
• 通过 MemRef 和 TileView 实现硬件感知的内存管理
• 通过 GlobalVar 实现程序内函数调用
• 通过 IterArg 实现循环携带依赖