NVWSOps
nvws.aref.buffer (triton::nvws::ArefBufferOp)
从 aref 获取缓冲区
语法
operation ::= `nvws.aref.buffer` $aref (`[` $stage^ `]`)? `,` $token attr-dict
`:` type($aref) `,` type($token) `->` type(results)
接口:ArefStageInterface
操作数:
操作数 |
描述 |
|---|---|
|
异步引用 (Asynchronous Reference) |
|
异步令牌类型 |
|
32位无符号整数 |
结果:
结果 |
描述 |
|---|---|
|
Triton IR 类型系统中内存描述符类型 ( |
nvws.aref.create (triton::nvws::ArefCreateOp)
创建一个异步引用。
语法
operation ::= `nvws.aref.create` $buffers attr-dict `:` type($result)
创建一个异步引用。
接收变长数量的缓冲区作为输入,并返回一个 ARef。输入预期为数组类型(如 Tensor、MemDesc 等),且所有输入的形状第一维必须匹配,表示值的多重缓冲 (multi-buffering)。
特性: AlwaysSpeculatableImplTrait
接口: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)
效果: MemoryEffects::Effect{}
操作数:
操作数 |
描述 |
|---|---|
|
Triton IR 类型系统中内存描述符类型 ( |
结果:
结果 |
描述 |
|---|---|
|
异步引用 (Asynchronous Reference) |
nvws.aref.get.enter (triton::nvws::ArefGetEnterOp)
进入 ArefGet 区域,在该区域内缓冲区可用于读取数据
语法
operation ::= `nvws.aref.get.enter` $aref ( `[` $stage^ `,` $phase `]`)? attr-dict
`:` type($aref) `->` type(results)
进入一个可以自由从缓冲区读取数据的“区域”。这些 ArefGet “区域”可以跨越多次迭代。
特性 (Traits):AttrSizedOperandSegments
接口:ArefStageInterface
操作数:
操作数 |
描述 |
|---|---|
|
异步引用 (Asynchronous Reference) |
|
32位无符号整数 |
|
32位无符号整数 |
结果:
结果 |
描述 |
|---|---|
|
Triton IR 类型系统中内存描述符类型 ( |
|
异步令牌类型 |
nvws.aref.get.exit (triton::nvws::ArefGetExitOp)
退出 ArefGet 区域,在该区域之后不应再使用该缓冲区
语法
operation ::= `nvws.aref.get.exit` $aref (`[` $stage^ `]`)? `,` $token $async_ops attr-dict
`:` type($aref) `,` type($token)
离开可以自由从缓冲区读取数据的区域。这些 ArefGet “区域”可以跨越多次迭代。
接口:ArefStageInterface
属性:
| 属性 | MLIR 类型 | 描述 |
|---|---|---|
async_ops | ::mlir::ArrayAttr | 异步操作属性数组 |
操作数:
操作数 |
描述 |
|---|---|
|
异步引用 (Asynchronous Reference) |
|
异步令牌类型 |
|
32位无符号整数 |
nvws.aref.put.enter (triton::nvws::ArefPutEnterOp)
进入 ArefPut 区域,在该区域内缓冲区可用于写入数据
语法
operation ::= `nvws.aref.put.enter` $aref ( `[` $stage^ `,` $phase `]`)? attr-dict
`:` type($aref) `->` type(results)
进入一个可以自由向缓冲区写入数据的“区域”。这些 ArefPut “区域”可以跨越多次迭代。
特性 (Traits):AttrSizedOperandSegments
接口:ArefStageInterface
操作数:
操作数 |
描述 |
|---|---|
|
异步引用 (Asynchronous Reference) |
|
32位无符号整数 |
|
32位无符号整数 |
结果:
结果 |
描述 |
|---|---|
|
Triton IR 类型系统中内存描述符类型 ( |
|
异步令牌类型 |
nvws.aref.put.exit (triton::nvws::ArefPutExitOp)
退出 ArefPut 区域,在该区域之后不应再使用该缓冲区
语法
operation ::= `nvws.aref.put.exit` $aref (`[` $stage^ `]`)? `,` $token $async_ops attr-dict
`:` type($aref) `,` type($token)
离开可以自由向缓冲区写入数据的区域。这些 ArefPut “区域”可以跨越多次迭代。
接口:ArefStageInterface
属性:
| 属性 | MLIR 类型 | 描述 |
|---|---|---|
async_ops | ::mlir::ArrayAttr | 异步操作属性数组 |
操作数:
操作数 |
描述 |
|---|---|
|
异步引用 (Asynchronous Reference) |
|
异步令牌类型 |
|
32位无符号整数 |
nvws.consumer_release (triton::nvws::ConsumerReleaseOp)
消费者释放令牌
语法
operation ::= `nvws.consumer_release` $token `,` $idx attr-dict `:` type(operands)
消费者将释放令牌并向生产者发送信号,告知缓冲区已准备好可以被填充。
操作数:
操作数 |
描述 |
|---|---|
|
值张量 |
|
32位无符号整数 |
nvws.consumer_wait (triton::nvws::ConsumerWaitOp)
消费者等待缓冲区就绪
语法
operation ::= `nvws.consumer_wait` $token `,` $idx `,` $phase attr-dict `:` type(operands)
消费者将等待缓冲区准备好被消费。如果缓冲区尚未就绪,消费者将等待生产者的信号,生产者在完成缓冲区填充并释放令牌后会发送该信号。
操作数:
操作数 |
描述 |
|---|---|
|
值张量 |
|
32位无符号整数 |
|
1位无符号整数 |
nvws.create_token (triton::nvws::CreateTokenOp)
创建一个用于通信通道同步的令牌
语法
operation ::= `nvws.create_token` attr-dict `:` type($result)
令牌将由生产者和消费者用于同步。生产者将获取并持有令牌,直到完成缓冲区填充,然后向等待的消费者发送信号。消费者将持有令牌直到消费完缓冲区,然后向试图获取令牌的等待生产者发送信号。
属性:
| 属性 | MLIR 类型 | 描述 |
|---|---|---|
numBuffers | ::mlir::IntegerAttr | 32位无符号整数属性 |
loadType | ::mlir::triton::nvws::TokenLoadTypeAttr | 允许的 32 位无符号整数情况:0, 1, 2, 3, 4 |
结果:
结果 |
描述 |
|---|---|
|
值张量 |
nvws.descriptor_gather (triton::nvws::DescriptorGatherOp)
从描述符中收集多行数据到共享内存
语法
operation ::= `nvws.descriptor_gather` $desc `[` $x_offsets `,` $y_offset `]` $txCount $result
attr-dict `:` type(operands)
此操作的行为与 Triton 方言中同名操作完全一致,但加载结果存储在共享内存中。执行仍然是同步的。
接口:NVWS_DescriptorLoadOpInterface, TT_DescriptorOpInterface
属性:
| 属性 | MLIR 类型 | 描述 |
|---|---|---|
txCount | ::mlir::IntegerAttr | 32位无符号整数属性 |
操作数:
操作数 |
描述 |
|---|---|
|
张量描述符类型( |
|
32 位无符号整数值的秩张量 |
|
32位无符号整数 |
|
Triton IR 类型系统中的内存描述符类型 ( |
nvws.descriptor_load (triton::nvws::DescriptorLoadOp)
从描述符加载并存储到共享内存
语法
operation ::= `nvws.descriptor_load` $desc `[` $indices `]` $txCount $result
oilist(
`cacheModifier` `=` $cache |
`evictionPolicy` `=` $evict
)
attr-dict `:` type(operands)
此操作的行为与 Triton 方言中同名操作完全一致,但加载结果存储在共享内存中。执行仍然是同步的。
接口:NVWS_DescriptorLoadOpInterface, TT_DescriptorOpInterface
属性:
| 属性 | MLIR 类型 | 描述 |
|---|---|---|
txCount | ::mlir::IntegerAttr | 32位无符号整数属性 |
cache | ::mlir::triton::CacheModifierAttr | 允许的32位无符号整数情况:1, 2, 3, 4, 5, 6, 7 |
evict | ::mlir::triton::EvictionPolicyAttr | 允许的32位无符号整数情况:1, 2, 3 |
操作数:
操作数 |
描述 |
|---|---|
|
张量描述符类型( |
|
可变参数的 32 位无符号整数 |
|
Triton IR 类型系统中的内存描述符类型 ( |
nvws.producer_acquire (triton::nvws::ProducerAcquireOp)
生产者获取令牌以填充缓冲区
语法
operation ::= `nvws.producer_acquire` $token `,` $idx `,` $phase attr-dict `:` type(operands)
生产者将在填充缓冲区前尝试获取令牌。如果缓冲区未准备好被填充,生产者将等待消费者的信号,消费者在消费完缓冲区并释放令牌后会发送该信号。
操作数:
操作数 |
描述 |
|---|---|
|
值张量 |
|
32位无符号整数 |
|
1位无符号整数 |
nvws.producer_commit (triton::nvws::ProducerCommitOp)
生产者提交缓冲区更改
语法
operation ::= `nvws.producer_commit` $token `,` $idx attr-dict `:` type(operands)
生产者将释放令牌并向消费者发送信号,告知缓冲区已准备好被消费。
操作数:
操作数 |
描述 |
|---|---|
|
值张量 |
|
32位无符号整数 |
nvws.warp_group (triton::nvws::WarpGroupOp)
用于 Warp 专门化的容器操作
用于 Warp 专门化分析的高级容器。
包含变长数量的 Warp 组(带有每个组中的 Warp 数量),以及一个用于容纳该 Warp 组计算的区域。
此操作的结果(如果有)即为第一个区域的结果,由 nvws.warp_group.yield 操作返回。
nvws.warp_group 在执行前应被降低 (lowered) 为 ttg.warp_specialize。
特性:RecursiveMemoryEffects, RecursivelySpeculatableImplTrait
接口:ConditionallySpeculatable
属性:
| 属性 | MLIR 类型 | 描述 |
|---|---|---|
numWarps | ::mlir::DenseI32ArrayAttr | i32 密集数组属性 |
结果:
结果 |
描述 |
|---|---|
|
任意类型的可变参数 |
nvws.warp_group.return (triton::nvws::WarpGroupReturnOp)
Warp 组区域的终结符
语法
operation ::= `nvws.warp_group.return` attr-dict
Warp 组预期通过对其输入进行引用修改来返回数值。因此,warp_group.return 操作不接受来自 Warp 组的返回值。
特性:AlwaysSpeculatableImplTrait, HasParent<WarpGroupOp>, Terminator
接口: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)
效果: MemoryEffects::Effect{}
nvws.warp_group.yield (triton::nvws::WarpGroupYieldOp)
从 nvws.warp_group 的第一个区域中 Yield(产出)
语法
operation ::= `nvws.warp_group.yield` ($values^)? attr-dict (`:` type($values)^)?
此操作等同于 ttg.warp_specialize 操作的 ttg.warp_yield 操作。
待办事项:决定是否应将 nvws.warp_group 移至 TritonGPU,或继续让 TritonGPU 依赖于 NVWS。如果是前者,则可以移除此操作。后者涉及 TritonGPU 和 NVWS 之间的循环依赖。
特性:AlwaysSpeculatableImplTrait, HasParent<WarpGroupOp>, ReturnLike, Terminator
接口 (Interfaces): ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface), RegionBranchTerminatorOpInterface
效果: MemoryEffects::Effect{}
操作数:
操作数 |
描述 |
|---|---|
|
任意类型的可变参数 |