Distributed Dialect Operations¶

distributed.consume_token (triton::distributed::ConsumeTokenOp)¶

Consumes a token and builds data dependency between wait ops and input.

Syntax:

operation ::= `distributed.consume_token` $input `,` $token `,`
              attr-dict `:` type($input) `,` type($token) `->` type($result)

The operator creates a consume node that represents a data dependency between a wait node that produces a token and another node that uses the input data (e.g., load).

Example:

distributed.consume_token() {input, token}

Traits: Elementwise

Interfaces: InferTypeOpInterface, MemoryEffectOpInterface

Operands (distributed.consume_token)¶

Operand	Description
`input`	floating-point or ranked tensor of floating-point values or integer or ranked tensor of integer values or ptr or ranked tensor of ptr values or ptr or Tensor descriptor type (`::mlir::triton::TensorDescType`) in Triton IR type system
`token`	integer or ranked tensor of integer values

Results (distributed.consume_token)¶

Result	Description
`result`	floating-point or ranked tensor of floating-point values or integer or ranked tensor of integer values or ptr or ranked tensor of ptr values or ptr or Tensor descriptor type (`::mlir::triton::TensorDescType`) in Triton IR type system

distributed.extern_call (triton::distributed::ExternCallOp)¶

Syntax:

operation ::= `distributed.extern_call` operands attr-dict `:` functional-type(operands, $result)

call an external function $symbol implemented in $libpath/$libname with $args return $libpath/$libname:$symbol($args…)

Interfaces: ConditionallySpeculatable, MemoryEffectOpInterface

Attributes (distributed.extern_call)¶

Attribute	MLIR Type	Description
`libname`	`::mlir::StringAttr`	string attribute
`libpath`	`::mlir::StringAttr`	string attribute
`symbol`	`::mlir::StringAttr`	string attribute
`pure`	`::mlir::BoolAttr`	bool attribute

Operands (distributed.extern_call)¶

Operand	Description
`srcs`	variadic of floating-point or ranked tensor of floating-point values or integer or ranked tensor of integer values or ptr or ranked tensor of ptr values or ptr

Results (distributed.extern_call)¶

Result	Description
`result`	variadic of floating-point or ranked tensor of floating-point values or integer or ranked tensor of integer values or ptr or ranked tensor of ptr values or ptr

distributed.get_num_ranks (triton::distributed::GetNumRanksOp)¶

Get the number of PEs running in the program.

Syntax:

operation ::= `distributed.get_num_ranks` $axis attr-dict `:` type($result)

Traits: AlwaysSpeculatableImplTrait

Interfaces: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands (distributed.get_num_ranks)¶

Operand	Description
`axis`	32-bit signless integer

Results (distributed.get_num_ranks)¶

Result	Description
`result`	32-bit signless integer

distributed.get_rank (triton::distributed::GetRankOp)¶

Rank of the calling PE.

Syntax:

operation ::= `distributed.get_rank` $axis attr-dict `:` type($result)

Traits: AlwaysSpeculatableImplTrait

Interfaces: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands (distributed.get_rank)¶

Operand	Description
`axis`	32-bit signless integer

Results (distributed.get_rank)¶

Result	Description
`result`	32-bit signless integer

distributed.notify (triton::distributed::NotifyOp)¶

Update sigAddr with signalVal using operation sigOp on the specified rank.

Interfaces: MemoryEffectOpInterface

Attributes (distributed.notify)¶

Attribute	MLIR Type	Description
`sigOp`	`::mlir::triton::distributed::SignalOpAttr`	allowed 32-bit signless integer cases: 1, 2. Enum cases: set (`SET`), add (`ADD`)
`commScope`	`::mlir::triton::distributed::CommScopeAttr`	allowed 32-bit signless integer cases: 1, 2, 3. Enum cases: gpu (`GPU`), intra_node (`INTRA_NODE`), inter_node (`INTER_NODE`)

Operands (distributed.notify)¶

Operand	Description
`sigAddr`	ptr
`signalVal`	64-bit signless integer
`rank`	32-bit signless integer

distributed.symm_at (triton::distributed::SymmAtOp)¶

Obtain the symmetric address on the specified Rank.

Syntax:

operation ::= `distributed.symm_at` $symmAddr `,` $rank `,` attr-dict `:` qualified(type($symmAddr))

Interfaces: InferTypeOpInterface, MemoryEffectOpInterface

Operands (distributed.symm_at)¶

Operand	Description
`symmAddr`	ptr
`rank`	32-bit signless integer

Results (distributed.symm_at)¶

Result	Description
`remoteAddr`	ptr

distributed.wait (triton::distributed::WaitOp)¶

Wait for a signal to be ready (becomes 1) and produce a token.

Syntax:

operation ::= `distributed.wait` $barrierPtr `,` $numBarriers `,` $waitValue `,`
              `scope` `=` $scope `semantic` `=` $semantic
              attr-dict `:` type($barrierPtr) `,` type($numBarriers) `,` type($waitValue) `->` type($token)

The operator creates a wait node that represents a spin lock waiting for a list of signals pointed by the barrierPtr to be set ready (set to 1).

Example:

token = distributed.wait() {barrierPtr, num_barriers, scope = 1 : i32, semantic = 0 : i32}

Interfaces: MemoryEffectOpInterface

Attributes (distributed.wait)¶

Attribute	MLIR Type	Description
`scope`	`::mlir::triton::MemSyncScopeAttr`	allowed 32-bit signless integer cases: 1, 2, 3. Enum cases: gpu (`GPU`), cta (`CTA`), sys (`SYSTEM`)
`semantic`	`::mlir::triton::MemSemanticAttr`	allowed 32-bit signless integer cases: 1, 2, 3, 4. Enum cases: relaxed (`RELAXED`), acquire (`ACQUIRE`), release (`RELEASE`), acq_rel (`ACQUIRE_RELEASE`)

Operands (distributed.wait)¶

Operand	Description
`barrierPtr`	ptr or ranked tensor of ptr values
`numBarriers`	integer or ranked tensor of integer values
`waitValue`	integer

Results (distributed.wait)¶

Result	Description
`token`	integer or ranked tensor of integer values