NVWSOps¶

`nvws.aref.buffer` (triton::nvws::ArefBufferOp)¶

从 aref 获取缓冲区

语法

operation ::= `nvws.aref.buffer` $aref (`[` $stage^ `]`)? `,` $token attr-dict
              `:` type($aref) `,` type($token) `->` type(results)

接口： ArefStageInterface

操作数:¶

操作数	描述
`aref`	异步引用
`token`	异步令牌类型
`stage`	32位无符号整数

结果:¶

结果	描述
`buffers`	Triton IR 类型系统中内存描述符类型（`::mlir::triton::gpu::MemDescType`）的可变参数

`nvws.aref.create` (triton::nvws::ArefCreateOp)¶

创建一个异步引用。

语法

operation ::= `nvws.aref.create` $buffers attr-dict `:` type($result)

创建一个异步引用。

它接受可变数量的缓冲区作为输入，并返回一个 ARef。输入应为类数组（即 Tensor、MemDesc 等），并且所有输入的第一个形状轴应匹配，表示这些值的多缓冲。

特性: AlwaysSpeculatableImplTrait

接口: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

效果: MemoryEffects::Effect{}

操作数：¶

操作数	描述
`buffers`	Triton IR 类型系统中内存描述符类型（`::mlir::triton::gpu::MemDescType`）的可变参数

结果：¶

结果	描述
`result`	异步引用

`nvws.aref.get.enter` (triton::nvws::ArefGetEnterOp)¶

进入 ArefGet 区域，在该区域中缓冲区可用于读取数据

语法

operation ::= `nvws.aref.get.enter` $aref ( `[` $stage^ `,` $phase `]`)? attr-dict
              `:` type($aref) `->` type(results)

进入一个可以自由读取缓冲区的“区域”。这些 ArefGet“区域”可以跨越多次迭代。

特性： AttrSizedOperandSegments

接口： ArefStageInterface

操作数：¶

操作数	描述
`aref`	异步引用
`stage`	32位无符号整数
`phase`	32位无符号整数

结果：¶

结果	描述
`buffers`	Triton IR 类型系统中内存描述符类型（`::mlir::triton::gpu::MemDescType`）的可变参数
`token`	异步令牌类型

`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	异步操作属性数组

操作数:¶

操作数	描述
`aref`	异步引用
`token`	异步令牌类型
`stage`	32位无符号整数

`nvws.aref.put.enter` (triton::nvws::ArefPutEnterOp)¶

进入 ArefPut 区域，在该区域中缓冲区可用于读取数据

语法

operation ::= `nvws.aref.put.enter` $aref ( `[` $stage^ `,` $phase `]`)? attr-dict
              `:` type($aref) `->` type(results)

进入一个可以自由写入缓冲区的“区域”。这些 ArefPut“区域”可以跨越多次迭代。

特性： AttrSizedOperandSegments

接口： ArefStageInterface

操作数 (Operands):¶

操作数	描述
`aref`	异步引用
`stage`	32位无符号整数
`phase`	32位无符号整数

结果：¶

结果	描述
`buffers`	Triton IR 类型系统中内存描述符类型（`::mlir::triton::gpu::MemDescType`）的可变参数
`token`	异步令牌类型

`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	异步操作属性数组

操作数:¶

操作数	描述
`aref`	异步引用
`token`	异步令牌类型
`stage`	32位无符号整数

`nvws.consumer_release` (triton::nvws::ConsumerReleaseOp)¶

消费者释放令牌

语法

operation ::= `nvws.consumer_release` $token `,` $idx attr-dict `:` type(operands)

消费者将释放令牌并通知生产者缓冲区已准备好被填充。

操作数:¶

操作数	描述
`token`	值的张量
`idx`	32位无符号整数

`nvws.consumer_wait` (triton::nvws::ConsumerWaitOp)¶

消费者等待缓冲区就绪

语法

operation ::= `nvws.consumer_wait` $token `,` $idx `,` $phase attr-dict `:` type(operands)

消费者将等待缓冲区准备好被消费。如果缓冲区尚未就绪，消费者将等待生产者的信号，生产者在填充完缓冲区并释放令牌后会发出信号。

操作数：¶

操作数	描述
`token`	值的张量
`idx`	32位无符号整数
`phase`	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

结果：¶

结果	描述
`result`	值的张量

`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位无符号整数属性

操作数：¶

操作数	描述
`desc`	张量描述符类型（`::mlir::triton::TensorDescType`）在 Triton IR 类型系统中
`x_offsets`	32 位无符号整数值的秩张量
`y_offset`	32位无符号整数
`result`	Triton IR 类型系统中的内存描述符类型 (`::mlir::triton::gpu::MemDescType`)

`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

操作数:¶

操作数	描述
`desc`	张量描述符类型（`::mlir::triton::TensorDescType`）在 Triton IR 类型系统中
`indices`	可变参数的 32 位无符号整数
`result`	Triton IR 类型系统中的内存描述符类型 (`::mlir::triton::gpu::MemDescType`)

`nvws.producer_acquire` (triton::nvws::ProducerAcquireOp)¶

生产者获取令牌以填充缓冲区

语法

operation ::= `nvws.producer_acquire` $token `,` $idx `,` $phase attr-dict `:` type(operands)

生产者在填充缓冲区之前会尝试获取令牌。如果缓冲区尚未准备好被填充，生产者将等待消费者的信号，消费者在消费完缓冲区并释放令牌后会发出信号。

操作数：¶

操作数	描述
`token`	值的张量
`idx`	32位无符号整数
`phase`	1位无符号整数

`nvws.producer_commit` (triton::nvws::ProducerCommitOp)¶

生产者提交缓冲区更改

语法

operation ::= `nvws.producer_commit` $token `,` $idx attr-dict `:` type(operands)

生产者将释放令牌并通知消费者缓冲区已准备好被消费。

操作数 (Operands):¶

操作数	描述
`token`	值的张量
`idx`	32位无符号整数

`nvws.warp_group` (triton::nvws::WarpGroupOp)¶

用于 Warp 专门化的容器操作

用于 Warp 专门化分析的更高级别的容器。

包含可变数量的 warp 组，每个组中有指定数量的 warp，以及一个用于保存该 warp 组计算的区域。

此操作的结果（如果有）是第一个区域的结果，由 nvws.warp_group.yield 操作返回。

nvws.warp_group 在执行前应被降级为 ttg.warp_specialize。

特性： RecursiveMemoryEffects, RecursivelySpeculatableImplTrait

接口： ConditionallySpeculatable

属性：¶

属性	MLIR 类型	描述
`numWarps`	::mlir::DenseI32ArrayAttr	i32 密集数组属性

结果：¶

结果	描述
`results`	任意类型的可变参数

`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{}

操作数：¶

操作数	描述
`values`	任意类型的可变参数

NVWSOps¶

nvws.aref.buffer (triton::nvws::ArefBufferOp)¶

操作数:¶

结果:¶

nvws.aref.create (triton::nvws::ArefCreateOp)¶

操作数：¶

结果：¶

nvws.aref.get.enter (triton::nvws::ArefGetEnterOp)¶

操作数：¶

结果：¶

nvws.aref.get.exit (triton::nvws::ArefGetExitOp)¶

属性：¶

操作数:¶

nvws.aref.put.enter (triton::nvws::ArefPutEnterOp)¶

操作数 (Operands):¶

结果：¶

nvws.aref.put.exit (triton::nvws::ArefPutExitOp)¶

属性：¶

操作数:¶

nvws.consumer_release (triton::nvws::ConsumerReleaseOp)¶

操作数:¶

nvws.consumer_wait (triton::nvws::ConsumerWaitOp)¶

操作数：¶

nvws.create_token (triton::nvws::CreateTokenOp)¶

属性：¶

结果：¶

nvws.descriptor_gather (triton::nvws::DescriptorGatherOp)¶

属性：¶

操作数：¶

nvws.descriptor_load (triton::nvws::DescriptorLoadOp)¶

属性：¶

操作数:¶

nvws.producer_acquire (triton::nvws::ProducerAcquireOp)¶

操作数：¶

nvws.producer_commit (triton::nvws::ProducerCommitOp)¶

操作数 (Operands):¶

nvws.warp_group (triton::nvws::WarpGroupOp)¶

属性：¶

结果：¶

nvws.warp_group.return (triton::nvws::WarpGroupReturnOp)¶

nvws.warp_group.yield (triton::nvws::WarpGroupYieldOp)¶

操作数：¶

`nvws.aref.buffer` (triton::nvws::ArefBufferOp)¶

`nvws.aref.create` (triton::nvws::ArefCreateOp)¶

`nvws.aref.get.enter` (triton::nvws::ArefGetEnterOp)¶

`nvws.aref.get.exit` (triton::nvws::ArefGetExitOp)¶

`nvws.aref.put.enter` (triton::nvws::ArefPutEnterOp)¶

`nvws.aref.put.exit` (triton::nvws::ArefPutExitOp)¶

`nvws.consumer_release` (triton::nvws::ConsumerReleaseOp)¶

`nvws.consumer_wait` (triton::nvws::ConsumerWaitOp)¶

`nvws.create_token` (triton::nvws::CreateTokenOp)¶

`nvws.descriptor_gather` (triton::nvws::DescriptorGatherOp)¶

`nvws.descriptor_load` (triton::nvws::DescriptorLoadOp)¶

`nvws.producer_acquire` (triton::nvws::ProducerAcquireOp)¶

`nvws.producer_commit` (triton::nvws::ProducerCommitOp)¶

`nvws.warp_group` (triton::nvws::WarpGroupOp)¶

`nvws.warp_group.return` (triton::nvws::WarpGroupReturnOp)¶

`nvws.warp_group.yield` (triton::nvws::WarpGroupYieldOp)¶