NVWSOps

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

Create an asynchronous reference.

Syntax:

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

Create an asynchronous reference.

Takes as inputs a variadic number of operands, and returns an ARef. The inputs are expected to be array-like (i.e., Tensor, MemDesc, etc) and the first axis of the shape should match between all inputs, representing multi-buffering of the values.

Traits: AlwaysSpeculatableImplTrait

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

Operand	Description
operands	variadic of any type

Results:

Result	Description
result	Asynchronous Reference

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

Enter ArefGet region where the buffer can be used to read data

Syntax:

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

Enter a “region” where you can freely read from the buffer) These ArefGet “regions” can span multiple iterations.

Operands:

Operand	Description
aref	Asynchronous Reference
index	32-bit signless integer
phase	32-bit signless integer

Results:

Result	Description
results	variadic of any type

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

Exit ArefGet region, where the buffer should no longer be used

Syntax:

operation ::= `nvws.aref.get.exit` $aref `[` $index `]` attr-dict
              `:` type($aref)

Leave the region where you can freely read from the buffer). These ArefGet “regions” can span multiple iterations.

Operands:

Operand	Description
aref	Asynchronous Reference
index	32-bit signless integer

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

Enter ArefPut region where the buffer can be used to read data

Syntax:

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

Enter a “region” where you can freely write to the buffer) These ArefPut “regions” can span multiple iterations.

Operands:

Operand	Description
aref	Asynchronous Reference
index	32-bit signless integer
phase	32-bit signless integer

Results:

Result	Description
results	variadic of any type

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

Exit ArefPut region, where the buffer should no longer be used

Syntax:

operation ::= `nvws.aref.put.exit` $aref `[` $index `]` attr-dict
              `:` type($aref)

Leave the region where you can freely write to the buffer). These ArefPut “regions” can span multiple iterations.

Operands:

Operand	Description
aref	Asynchronous Reference
index	32-bit signless integer

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

Consumer releases the token

Syntax:

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

The consumer will release the token and signal the producer that the buffers are ready to be filled.

Operands:

Operand	Description
token	tensor of values
idx	32-bit signless integer

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

Consumer awaits buffer readiness

Syntax:

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

The consumer will wait for the buffer to be ready to be consumed. If the buffers are not ready, the consumer will wait to be signalled by the producer which finishes filling the buffers and releases the token.

Operands:

Operand	Description
token	tensor of values
idx	32-bit signless integer
phase	1-bit signless integer

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

Create a token to be used for synchronizations in communication channels

Syntax:

operation ::= `nvws.create_token` attr-dict `:` type($result)

A token will be used by the producer and consumer to synchronize. The producer will acquire and hold the token, until it has filled the buffers, and signal the waiting consumer. The consumer will hold the token until it has consumed the buffers, and will signal the waiting producer trying to acquire the token.

Attributes:

Attribute	MLIR Type	Description
numBuffers	::mlir::IntegerAttr	32-bit signless integer attribute
loadType	::mlir::triton::nvws::TokenLoadTypeAttr	allowed 32-bit signless integer cases: 0, 1, 2, 3, 4

Results:

Result	Description
result	tensor of values

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

Producer acquires a token to fill buffers

Syntax:

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

The producer will try to acquire the token prior to filling the buffers. If the buffers are not ready to be filled, the producer will wait to be signalled by the consumer which finishes consuming the buffers and releases the token.

Operands:

Operand	Description
token	tensor of values
idx	32-bit signless integer
phase	1-bit signless integer

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

Producer commits the buffer changes

Syntax:

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

The producer will release the token and signal the consumer that the buffers are ready to be consumed.

Operands:

Operand	Description
token	tensor of values
idx	32-bit signless integer

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

Container Op for Warp Specialization

Higher level container for Warp Specialization Analysis.

Contains a variadic number warp groups, with the number of warps in each group, plus a region to hold the computation for that warp group.

Regions are not Isolated from Above to aid in analysis, and take inputs purely by reference.

nvws.warp_group should be lowered to ttg.warp_specialize before execution.

Traits: RecursiveMemoryEffects, RecursivelySpeculatableImplTrait

Interfaces: ConditionallySpeculatable

Attributes:

Attribute	MLIR Type	Description
numWarps	::mlir::DenseI32ArrayAttr	i32 dense array attribute

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

Terminator for a warp group region

Syntax:

operation ::= `nvws.warp_group.return` attr-dict

Warp groups are expected to return values via referential modification of their inputs. Thus, the warp_group.return op takes no values to return from the warp group.

Traits: AlwaysSpeculatableImplTrait, HasParent<WarpGroupOp>, Terminator

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}