# Distributed Dialect Operations

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

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

Syntax:

```mlir
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:

```mlir
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:

```mlir
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:

```mlir
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:

```mlir
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:

```mlir
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:

```mlir
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:

```mlir
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 |