Operations
mhlo.abs (mhlo::AbsOp)
Abs operation
Syntax:
operation ::= `mhlo.abs` $operand attr-dict
`:` custom<SameOperandsAndResultType>(type($operand), type($result))
Performs element-wise abs operation on operand tensor and produces a
result tensor.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#abs
Example:
%result = mhlo.abs %operand : tensor<3xi32>
Traits: AlwaysSpeculatableImplTrait, Elementwise, SameOperandsAndResultShape
Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Operands:
| Operand | Description |
|---|---|
operand |
ranked tensor of 2/4/8/16/32/64-bit signless integer or 4/6/8/16/32/64-bit float or complex type with 32/64-bit float elements or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values |
Results:
| Result | Description |
|---|---|
result |
ranked tensor of 2/4/8/16/32/64-bit signless integer or 4/6/8/16/32/64-bit float or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values |
mhlo.acos (mhlo::AcosOp)
Acos operation
Syntax:
operation ::= `mhlo.acos` $operand attr-dict
`:` custom<SameOperandsAndResultType>(type($operand), type($result))
Performs element-wise acos operation on operand tensor and produces a
result tensor.
Example:
%result = mhlo.acos %operand : tensor<2x2xf32>
Traits: CompatibleOperandsAndResultType, Elementwise, SameOperandsAndResultShape
Interfaces: InferShapedTypeOpInterface, InferTypeOpInterface
Operands:
| Operand | Description |
|---|---|
operand |
tensor of 4/6/8/16/32/64-bit float or complex type with 32/64-bit float elements values |
Results:
| Result | Description |
|---|---|
result |
tensor of 4/6/8/16/32/64-bit float or complex type with 32/64-bit float elements values |
mhlo.acosh (mhlo::AcoshOp)
Acosh operation
Syntax:
operation ::= `mhlo.acosh` $operand attr-dict
`:` custom<SameOperandsAndResultType>(type($operand), type($result))
Performs element-wise acosh operation on operand tensor and produces a
result tensor.
Example:
%result = mhlo.acosh %operand : tensor<2x2xf32>
Traits: CompatibleOperandsAndResultType, Elementwise, SameOperandsAndResultShape
Interfaces: InferShapedTypeOpInterface, InferTypeOpInterface
Operands:
| Operand | Description |
|---|---|
operand |
tensor of 4/6/8/16/32/64-bit float or complex type with 32/64-bit float elements values |
Results:
| Result | Description |
|---|---|
result |
tensor of 4/6/8/16/32/64-bit float or complex type with 32/64-bit float elements values |
mhlo.add (mhlo::AddOp)
Add operation
Syntax:
operation ::= `mhlo.add` $lhs `,` $rhs attr-dict
`:` custom<SameOperandsAndResultType>(type($lhs), type($rhs), type($result))
Performs element-wise addition of two tensors lhs and rhs and produces a
result tensor.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#add
Example:
%result = mhlo.add %lhs, %rhs : tensor<2x2xi32>
Traits: AlwaysSpeculatableImplTrait, Commutative, CompatibleOperandsAndResultType, Elementwise, SameOperandsAndResultShape
Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Operands:
| Operand | Description |
|---|---|
lhs |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
rhs |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
Results:
| Result | Description |
|---|---|
result |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
mhlo.add_dependency (mhlo::AddDependencyOp)
AddDependency operation
Syntax:
operation ::= `mhlo.add_dependency` operands attr-dict `:` functional-type(operands, results)
This operation is private to the XLA compiler, so it is does not yet have a specification.
Informally, this operation two operands: a data operand and a token. The output of the operation is the data operand. When used with AfterAll this operation enables ordering non-side-effecting operations (those that do not produce token values).
Example:
%1 = mhlo.add_dependency %arg0, %0 : (tensor<3x4xf32>, !mhlo.token) -> tensor<3x4xf32>
Traits: AlwaysSpeculatableImplTrait
Interfaces: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Operands:
| Operand | Description |
|---|---|
operand |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized values or ranked tensor of per-axis integer quantized values or token or stablehlo token |
token |
token or stablehlo token |
Results:
| Result | Description |
|---|---|
output |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized values or ranked tensor of per-axis integer quantized values or token or stablehlo token |
mhlo.after_all (mhlo::AfterAllOp)
AfterAll operation
Syntax:
operation ::= `mhlo.after_all` $inputs attr-dict
`:` custom<VariadicSameOperandsAndResultType>(ref($inputs), type($inputs), type($result))
Ensures that the operations producing the inputs are executed before any
operations that depend on result.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#after_all
Example:
%result = mhlo.after_all %input0, %input1 : !mhlo.token
Traits: AlwaysSpeculatableImplTrait
Interfaces: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Operands:
| Operand | Description |
|---|---|
inputs |
variadic of token |
Results:
| Result | Description |
|---|---|
result |
token |
mhlo.all_gather (mhlo::AllGatherOp)
AllGather operation
Within each process group in the process grid, concatenates the values of the
operand tensor from each process along all_gather_dim and produces a
result tensor. The computation is applied separately for each operand in
operands, producing one result per operand.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#all_gather
Example:
%result = "mhlo.all_gather"(%operand) {
all_gather_dim = 1 : i64,
replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>
// channel_id = 0
channel_handle = #mhlo.channel_handle<handle = 0, type = 0>,
// use_global_device_ids = false
} : (tensor<2x2xf32>) -> tensor<2x4xf32>
Traits: SameOperandsAndResultElementType
Attributes:
| Attribute | MLIR Type | Description |
|---|---|---|
all_gather_dim | ::mlir::IntegerAttr | 64-bit signless integer attribute whose value is non-negative |
replica_groups | ::mlir::DenseIntElementsAttr | 64-bit signless integer elements attribute |
channel_handle | ::mlir::mhlo::ChannelHandleAttr | two 64-bit integers 'handle' and 'type' |
use_global_device_ids | ::mlir::UnitAttr | unit attribute |
Operands:
| Operand | Description |
|---|---|
operands |
variadic of ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
Results:
| Result | Description |
|---|---|
| «unnamed» | variadic of ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
mhlo.all_reduce (mhlo::AllReduceOp)
AllReduce operation
Within each process group in the process grid, applies a reduction function
computation to the values of an operand tensor from each process and
produces a result tensor. The computation is applied separately for each
operand in operands, producing one result per operand.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#all_reduce
Example:
%result = "mhlo.all_reduce"(%operand) ({
^bb0(%arg0: tensor<f32>, %arg1: tensor<f32>):
%0 = mhlo.add %arg1, %arg2 : tensor<f32>
mhlo.return %0 : tensor<f32>
}) {
replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>
// channel_id = 0
channel_handle = #mhlo.channel_handle<handle = 0, type = 0>
// use_global_device_ids = false
} : (tensor<4xf32>) -> tensor<4xf32>
Traits: InferTensorType, SingleBlockImplicitTerminator<ReturnOp>, SingleBlock
Interfaces: InferShapedTypeOpInterface, InferTypeOpInterface
Attributes:
| Attribute | MLIR Type | Description |
|---|---|---|
replica_groups | ::mlir::DenseIntElementsAttr | 64-bit signless integer elements attribute |
channel_handle | ::mlir::mhlo::ChannelHandleAttr | two 64-bit integers 'handle' and 'type' |
use_global_device_ids | ::mlir::UnitAttr | unit attribute |
Operands:
| Operand | Description |
|---|---|
operands |
variadic of ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
Results:
| Result | Description |
|---|---|
| «unnamed» | variadic of ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
mhlo.all_to_all (mhlo::AllToAllOp)
AllToAll operation
Within each process group in the process grid, splits the values of the
operand tensor along split_dimension into parts, scatters the split parts
between the processes, concatenates the scattered parts along concat_dimension
and produces a result tensor.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#all_to_all
Example:
%result = "mhlo.all_to_all"(%operand) {
split_dimension = 1 : i64,
concat_dimension = 0 : i64,
split_count = 2 : i64,
replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>
} : (tensor<2x4xf32>) -> tensor<4x2xf32>
Traits: AlwaysSpeculatableImplTrait, InferTensorType, SameOperandsElementType, SameOperandsShape, SameVariadicOperandSize
Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Attributes:
| Attribute | MLIR Type | Description |
|---|---|---|
split_dimension | ::mlir::IntegerAttr | 64-bit signless integer attribute whose value is non-negative |
concat_dimension | ::mlir::IntegerAttr | 64-bit signless integer attribute whose value is non-negative |
split_count | ::mlir::IntegerAttr | 64-bit signless integer attribute whose value is positive |
replica_groups | ::mlir::DenseIntElementsAttr | 64-bit signless integer elements attribute |
channel_handle | ::mlir::mhlo::ChannelHandleAttr | two 64-bit integers 'handle' and 'type' |
Operands:
| Operand | Description |
|---|---|
operand |
variadic of ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
Results:
| Result | Description |
|---|---|
| «unnamed» | variadic of ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
mhlo.and (mhlo::AndOp)
And operation
Syntax:
operation ::= `mhlo.and` $lhs `,` $rhs attr-dict
`:` custom<SameOperandsAndResultType>(type($lhs), type($rhs), type($result))
Performs element-wise AND of two tensors lhs and rhs and produces a
result tensor
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#and
Example:
%result = mhlo.and %lhs, %rhs : tensor<2x2xi32>
Traits: AlwaysSpeculatableImplTrait, Commutative, CompatibleOperandsAndResultType, Elementwise, SameOperandsAndResultShape
Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Operands:
| Operand | Description |
|---|---|
lhs |
ranked tensor of bool or 2/4/8/16/32/64-bit integer values |
rhs |
ranked tensor of bool or 2/4/8/16/32/64-bit integer values |
Results:
| Result | Description |
|---|---|
result |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
mhlo.asin (mhlo::AsinOp)
Asin operation
Syntax:
operation ::= `mhlo.asin` $operand attr-dict
`:` custom<SameOperandsAndResultType>(type($operand), type($result))
Performs element-wise asin operation on operand tensor and produces a
result tensor.
Example:
%result = mhlo.asin %operand : tensor<2x2xf32>
Traits: CompatibleOperandsAndResultType, Elementwise, SameOperandsAndResultShape
Interfaces: InferShapedTypeOpInterface, InferTypeOpInterface
Operands:
| Operand | Description |
|---|---|
operand |
tensor of 4/6/8/16/32/64-bit float or complex type with 32/64-bit float elements values |
Results:
| Result | Description |
|---|---|
result |
tensor of 4/6/8/16/32/64-bit float or complex type with 32/64-bit float elements values |
mhlo.asinh (mhlo::AsinhOp)
Asinh operation
Syntax:
operation ::= `mhlo.asinh` $operand attr-dict
`:` custom<SameOperandsAndResultType>(type($operand), type($result))
Performs element-wise asinh operation on operand tensor and produces a
result tensor.
Example:
%result = mhlo.asinh %operand : tensor<2x2xf32>
Traits: CompatibleOperandsAndResultType, Elementwise, SameOperandsAndResultShape
Interfaces: InferShapedTypeOpInterface, InferTypeOpInterface
Operands:
| Operand | Description |
|---|---|
operand |
tensor of 4/6/8/16/32/64-bit float or complex type with 32/64-bit float elements values |
Results:
| Result | Description |
|---|---|
result |
tensor of 4/6/8/16/32/64-bit float or complex type with 32/64-bit float elements values |
mhlo.async_done (mhlo::AsyncDoneOp)
AsyncDone operation
This operation is private to the XLA compiler, so it is does not yet have a specification.
Informally, this operation blocks until the end of an asynchronous computation. It returns the final result of the asynchronous computation.
See the documentation for AsyncStart for more information.
Interfaces: InferTypeOpInterface
Operands:
| Operand | Description |
|---|---|
bundle |
async_bundle with any combination of ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values or token or stablehlo token values |
Results:
| Result | Description |
|---|---|
| «unnamed» | variadic of ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values or token or stablehlo token or nested tuple with any combination of ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized values or memref of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized values or ranked tensor of per-axis integer quantized values or token values |
mhlo.async_start (mhlo::AsyncStartOp)
AsyncStart operation
This operation is private to the XLA compiler, so it is does not yet have a specification.
Informally, this operation kicks off an asynchronous computation.
This is used when there are functions that contain both asynchronous waits (such as DMAs) and on-thread computation. For example, a function might consist of a computation, a DMA, another computation, a second DMA, and a final computation. This would be represented as an async_start followed by and async_update and an async_done. The async_start would do the first computation on-thread and then start the DMA. The async_update would wait for the DMA to complete if it wasn't yet done, then execute the second computation in the function, and start the second DMA. Finally, the async_done would wait on this last DMA, and then run the last computation that needs to be run on-thread and return the result of that final computation.
operands are passed to the computation directly
called_computation is the function that will be run asynchronously
execution_thread is the name of the thread in which it will be run. The main
thread is called "main". All threads have names.
This returns all the state needed between async ops. After buffer assignment, the return values represents the space needed to hold the input, results, and any scratchpads needed or edited by the async op.
Attributes:
| Attribute | MLIR Type | Description |
|---|---|---|
called_computation | ::mlir::FlatSymbolRefAttr | flat symbol reference attribute |
execution_thread | ::mlir::StringAttr | string attribute |
Operands:
| Operand | Description |
|---|---|
inputs |
variadic of ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values or token or stablehlo token or nested tuple with any combination of ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized values or memref of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized values or ranked tensor of per-axis integer quantized values or token values |
Results:
| Result | Description |
|---|---|
| «unnamed» | async_bundle with any combination of ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values or token or stablehlo token values |
mhlo.async_update (mhlo::AsyncUpdateOp)
AsyncUpdate operation
This operation is private to the XLA compiler, so it is does not yet have a specification.
Informally, this operation blocks on an asynchronous computation until a
sync barrier. This returns bundle after operating on it.
See the documentation for AsyncStart for more information.
Interfaces: InferTypeOpInterface
Operands:
| Operand | Description |
|---|---|
bundle |
async_bundle with any combination of ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values or token or stablehlo token values |
Results:
| Result | Description |
|---|---|
| «unnamed» | async_bundle with any combination of ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values or token or stablehlo token values |
mhlo.atan2 (mhlo::Atan2Op)
Atan2 operation
Syntax:
operation ::= `mhlo.atan2` $lhs `,` $rhs attr-dict
`:` custom<SameOperandsAndResultType>(type($lhs), type($rhs), type($result))
Performs element-wise atan2 operation on lhs and rhs tensor and produces
a result tensor.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#atan2
Example:
%result = mhlo.atan2 %lhs, %rhs : tensor<3xf32>
Traits: AlwaysSpeculatableImplTrait, CompatibleOperandsAndResultType, Elementwise, SameOperandsAndResultShape
Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Operands:
| Operand | Description |
|---|---|
lhs |
ranked tensor of 4/6/8/16/32/64-bit float or complex type with 32/64-bit float elements or per-tensor integer quantized values |
rhs |
ranked tensor of 4/6/8/16/32/64-bit float or complex type with 32/64-bit float elements or per-tensor integer quantized values |
Results:
| Result | Description |
|---|---|
result |
ranked tensor of 4/6/8/16/32/64-bit float or complex type with 32/64-bit float elements or per-tensor integer quantized values |
mhlo.atanh (mhlo::AtanhOp)
Atanh operation
Syntax:
operation ::= `mhlo.atanh` $operand attr-dict
`:` custom<SameOperandsAndResultType>(type($operand), type($result))
Performs element-wise atanh operation on operand tensor and produces a
result tensor.
Example:
%result = mhlo.atanh %operand : tensor<2x2xf32>
Traits: CompatibleOperandsAndResultType, Elementwise, SameOperandsAndResultShape
Interfaces: InferShapedTypeOpInterface, InferTypeOpInterface
Operands:
| Operand | Description |
|---|---|
operand |
tensor of 4/6/8/16/32/64-bit float or complex type with 32/64-bit float elements values |
Results:
| Result | Description |
|---|---|
result |
tensor of 4/6/8/16/32/64-bit float or complex type with 32/64-bit float elements values |
mhlo.batch_norm_grad (mhlo::BatchNormGradOp)
BatchNormGrad operation
Computes gradients of several inputs of BatchNormTrainingOp backpropagating
from grad_output, and produces grad_operand, grad_scale and
grad_offset tensors.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#batch_norm_grad
Example:
%grad_operand, %grad_scale, %grad_offset =
"mhlo.batch_norm_grad"(%operand, %scale, %mean, %variance, %grad_output) {
epsilon = 0.0 : f32,
feature_index = 2 : i64
} : (tensor<2x2x2xf32>, tensor<2xf32>, tensor<2xf32>, tensor<2xf32>,
tensor<2x2x2xf32>) -> (tensor<2x2x2xf32>, tensor<2xf32>, tensor<2xf32>)
Traits: AlwaysSpeculatableImplTrait, InferTensorType
Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Attributes:
| Attribute | MLIR Type | Description |
|---|---|---|
epsilon | ::mlir::FloatAttr | 32-bit float attribute |
feature_index | ::mlir::IntegerAttr | 64-bit signless integer attribute whose value is non-negative |
Operands:
| Operand | Description |
|---|---|
operand |
ranked tensor of 4/6/8/16/32/64-bit float values |
scale |
1D tensor of 4/6/8/16/32/64-bit float values |
mean |
1D tensor of 4/6/8/16/32/64-bit float values |
variance |
1D tensor of 4/6/8/16/32/64-bit float values |
grad_output |
ranked tensor of 4/6/8/16/32/64-bit float values |
Results:
| Result | Description |
|---|---|
grad_operand |
ranked tensor of 4/6/8/16/32/64-bit float values |
grad_scale |
1D tensor of 4/6/8/16/32/64-bit float values |
grad_offset |
1D tensor of 4/6/8/16/32/64-bit float values |
mhlo.batch_norm_inference (mhlo::BatchNormInferenceOp)
BatchNormInference operation
Normalizes the operand tensor across all dimensions except for the
feature_index dimension and produces a result tensor.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#batch_norm_inference
Example:
%result = "mhlo.batch_norm_inference"(%operand, %scale, %offset, %mean, %variance) {
epsilon = 0.0 : f32,
feature_index = 2 : i64
} : (tensor<2x2x2xf32>, tensor<2xf32>, tensor<2xf32>, tensor<2xf32>, tensor<2xf32>) -> tensor<2x2x2xf32>
Traits: AlwaysSpeculatableImplTrait, InferTensorType
Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Attributes:
| Attribute | MLIR Type | Description |
|---|---|---|
epsilon | ::mlir::FloatAttr | 32-bit float attribute |
feature_index | ::mlir::IntegerAttr | 64-bit signless integer attribute whose value is non-negative |
Operands:
| Operand | Description |
|---|---|
operand |
ranked tensor of 4/6/8/16/32/64-bit float values |
scale |
1D tensor of 4/6/8/16/32/64-bit float values |
offset |
1D tensor of 4/6/8/16/32/64-bit float values |
mean |
1D tensor of 4/6/8/16/32/64-bit float values |
variance |
1D tensor of 4/6/8/16/32/64-bit float values |
Results:
| Result | Description |
|---|---|
result |
ranked tensor of 4/6/8/16/32/64-bit float values |
mhlo.batch_norm_training (mhlo::BatchNormTrainingOp)
BatchNormTraining operation
Computes mean and variance across batch and spatial dimensions and
normalizes the operand tensor, for each feature in the feature_index
dimension and produces output, batch_mean and batch_var tensors.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#batch_norm_training
Example:
%output, %batch_mean, %batch_var = "mhlo.batch_norm_training"(%operand, %scale, %offset) {
epsilon = 0.0 : f32,
feature_index = 2 : i64
} : (tensor<2x2x2xf32>, tensor<2xf32>, tensor<2xf32>) -> (tensor<2x2x2xf32>, tensor<2xf32>, tensor<2xf32>)
Traits: AlwaysSpeculatableImplTrait, InferTensorType
Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Attributes:
| Attribute | MLIR Type | Description |
|---|---|---|
epsilon | ::mlir::FloatAttr | 32-bit float attribute |
feature_index | ::mlir::IntegerAttr | 64-bit signless integer attribute whose value is non-negative |
Operands:
| Operand | Description |
|---|---|
operand |
ranked tensor of 4/6/8/16/32/64-bit float values |
scale |
1D tensor of 4/6/8/16/32/64-bit float values |
offset |
1D tensor of 4/6/8/16/32/64-bit float values |
Results:
| Result | Description |
|---|---|
output |
ranked tensor of 4/6/8/16/32/64-bit float values |
batch_mean |
1D tensor of 4/6/8/16/32/64-bit float values |
batch_var |
1D tensor of 4/6/8/16/32/64-bit float values |
mhlo.bitcast (mhlo::BitcastOp)
Bitcast operation
Syntax:
operation ::= `mhlo.bitcast` operands attr-dict `:` functional-type(operands, results)
This operation is private to the XLA compiler, so it is does not yet have a specification.
Informally, this operation changes the shape of the input in the way that the physical arrangement of elements are unchanged.
This operation needs layout information to make sense of "physical arrangement of elements", and layout support in MHLO is currently a work in progress.
Example:
%0 = mhlo.bitcast %arg0 : (tensor<3x4xf32>) -> tensor<3x4x1xf32>
Traits: AlwaysSpeculatableImplTrait
Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Operands:
| Operand | Description |
|---|---|
operand |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
Results:
| Result | Description |
|---|---|
| «unnamed» | ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
mhlo.bitcast_convert (mhlo::BitcastConvertOp)
BitcastConvert operation
Syntax:
operation ::= `mhlo.bitcast_convert` operands attr-dict `:` functional-type(operands, results)
Performs a bitcast operation on operand tensor and produces a result
tensor where the bits of the entire operand tensor are reinterpreted using
the type of the result tensor.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#bitcast_convert
Example:
%result = mhlo.bitcast_convert %operand : (tensor<2xf32>) -> tensor<2x4xi8>
Traits: AlwaysSpeculatableImplTrait
Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Operands:
| Operand | Description |
|---|---|
operand |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
Results:
| Result | Description |
|---|---|
| «unnamed» | ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
mhlo.broadcast (mhlo::BroadcastOp)
Broadcast operation
This operation is on its way out of StableHLO, so it is not included in the specification: https://github.com/openxla/stablehlo/issues/3
Informally, this operation does the same thing as XLA's Broadcast: https://www.tensorflow.org/xla/operation_semantics#broadcast
Example:
%result = mhlo.broadcast %operand, sizes = [1, 2] : (tensor<3xi32>) -> tensor<1x2x3xi32>
Traits: AlwaysSpeculatableImplTrait, InferTensorType, SameOperandsAndResultElementType
Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Attributes:
| Attribute | MLIR Type | Description |
|---|---|---|
broadcast_sizes | ::mlir::DenseIntElementsAttr | 64-bit signless integer elements attribute |
Operands:
| Operand | Description |
|---|---|
operand |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
Results:
| Result | Description |
|---|---|
| «unnamed» | ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
mhlo.broadcast_in_dim (mhlo::BroadcastInDimOp)
BroadcastInDim operation
Expands the dimensions and/or rank of an input tensor by duplicating the
data in the operand tensor and produces a result tensor.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#broadcast_in_dim
Example:
%result = mhlo.broadcast_in_dim %operand, dims = [2, 1] : (tensor<1x3xi32>) -> tensor<2x3x2xi32>
Traits: AlwaysSpeculatableImplTrait, HLO_CompatibleOperandsAndResultElementType
Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Attributes:
| Attribute | MLIR Type | Description |
|---|---|---|
broadcast_dimensions | ::mlir::DenseIntElementsAttr | 64-bit signless integer elements attribute |
Operands:
| Operand | Description |
|---|---|
operand |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
Results:
| Result | Description |
|---|---|
| «unnamed» | statically shaped or single bounded dimension tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
mhlo.case (mhlo::CaseOp)
Case operation
Produces the output from executing exactly one function from branches
depending on the value of index.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#case
Example:
%result0, %result1 = "mhlo.case"(%index) ({
mhlo.return %result_branch0, %result_branch0 : tensor<2xi64>, tensor<2xi64>
}, {
mhlo.return %result_branch1, %result_branch1 : tensor<2xi64>, tensor<2xi64>
}) : (tensor<i32>) -> (tensor<2xi64>, tensor<2xi64>)
Traits: RecursiveMemoryEffects, SingleBlockImplicitTerminator<ReturnOp>, SingleBlock
Interfaces: InferTypeOpInterface
Operands:
| Operand | Description |
|---|---|
index |
tensor of 32-bit signless integer values |
Results:
| Result | Description |
|---|---|
| «unnamed» | variadic of ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized values or ranked tensor of per-axis integer quantized values or token |
mhlo.cbrt (mhlo::CbrtOp)
Cbrt operation
Syntax:
operation ::= `mhlo.cbrt` $operand attr-dict
`:` custom<SameOperandsAndResultType>(type($operand), type($result))
Performs element-wise cubic root operation on operand tensor and produces
a result tensor.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#cbrt
Example:
%result = mhlo.cbrt %operand : tensor<4xf32>
Traits: AlwaysSpeculatableImplTrait, CompatibleOperandsAndResultType, Elementwise, SameOperandsAndResultShape
Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Attributes:
| Attribute | MLIR Type | Description |
|---|---|---|
result_accuracy | ::mlir::mhlo::ResultAccuracyAttr | The requested accuracy for unary ops. |
Operands:
| Operand | Description |
|---|---|
operand |
ranked tensor of 4/6/8/16/32/64-bit float or complex type with 32/64-bit float elements or per-tensor integer quantized values |
Results:
| Result | Description |
|---|---|
result |
ranked tensor of 4/6/8/16/32/64-bit float or complex type with 32/64-bit float elements or per-tensor integer quantized values |
mhlo.ceil (mhlo::CeilOp)
Ceil operation
Syntax:
operation ::= `mhlo.ceil` $operand attr-dict
`:` custom<SameOperandsAndResultType>(type($operand), type($result))
Performs element-wise ceil of operand tensor and produces a result tensor.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#ceil
Example:
%result = mhlo.ceil %operand : tensor<5xf32>
Traits: AlwaysSpeculatableImplTrait, CompatibleOperandsAndResultType, Elementwise, SameOperandsAndResultShape
Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Operands:
| Operand | Description |
|---|---|
operand |
ranked tensor of 4/6/8/16/32/64-bit float or per-tensor integer quantized values |
Results:
| Result | Description |
|---|---|
result |
ranked tensor of 4/6/8/16/32/64-bit float or per-tensor integer quantized values |
mhlo.cholesky (mhlo::CholeskyOp)
Cholesky operation
Computes the Cholesky decomposition of a batch of matrices.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#cholesky
Example:
%result = mhlo.cholesky %a, lower = true : tensor<3x3xf32>
Traits: AlwaysSpeculatableImplTrait, InferTensorType, SameOperandsAndResultElementType
Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Attributes:
| Attribute | MLIR Type | Description |
|---|---|---|
lower | ::mlir::BoolAttr | bool attribute |
Operands:
| Operand | Description |
|---|---|
a |
ranked tensor of 4/6/8/16/32/64-bit float or complex type with 32/64-bit float elements values |
Results:
| Result | Description |
|---|---|
| «unnamed» | ranked tensor of 4/6/8/16/32/64-bit float or complex type with 32/64-bit float elements values |
mhlo.clamp (mhlo::ClampOp)
Clamp operation
Syntax:
operation ::= `mhlo.clamp` $min `,` $operand `,` $max attr-dict
`:` custom<SameOperandsAndResultType>(type($min), type($operand), type($max), type($result))
Clamps every element of the operand tensor between a minimum and maximum
value and produces a result tensor.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#clamp
Example:
%result = mhlo.clamp %min, %operand, %max : tensor<3xi32>
Traits: AlwaysSpeculatableImplTrait, HLO_BroadcastingElementwise, InferTensorType, SameOperandsAndResultElementType
Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Operands:
| Operand | Description |
|---|---|
min |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
operand |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
max |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
Results:
| Result | Description |
|---|---|
result |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
mhlo.collective_broadcast (mhlo::CollectiveBroadcastOp)
CollectiveBroadcast operation
Within each process group in the process grid, send the value of the
operand tensor from the source process to the target processes and produce a
result tensor.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#collective_broadcast
Example:
%result = "mhlo.collective_broadcast"(%operand) {
replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>,
channel_handle = #mhlo.channel_handle<handle = 0, type = 0>
} : (tensor<1x2xi64>) -> tensor<1x2xi64>
Traits: CompatibleOperandsAndResultType
Interfaces: InferShapedTypeOpInterface, InferTypeOpInterface
Attributes:
| Attribute | MLIR Type | Description |
|---|---|---|
replica_groups | ::mlir::DenseIntElementsAttr | 64-bit signless integer elements attribute |
channel_handle | ::mlir::mhlo::ChannelHandleAttr | two 64-bit integers 'handle' and 'type' |
Operands:
| Operand | Description |
|---|---|
operand |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
Results:
| Result | Description |
|---|---|
| «unnamed» | ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
mhlo.collective_permute (mhlo::CollectivePermuteOp)
CollectivePermute operation
Within each process group in the process grid, sends the value of the
operand tensor from the source process to the target process and produces
a result tensor.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#collective_permute
Example:
%result = "mhlo.collective_permute"(%operand) {
source_target_pairs = dense<[[0, 1], [1, 2]]> : tensor<2x2xi64>,
// channel_id = 0
channel_handle = #mhlo.channel_handle<handle = 0, type = 0>
} : (tensor<4x2xf32>) -> tensor<4x2xf32>
Traits: AlwaysSpeculatableImplTrait, CompatibleOperandsAndResultType
Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Attributes:
| Attribute | MLIR Type | Description |
|---|---|---|
source_target_pairs | ::mlir::DenseIntElementsAttr | 64-bit signless integer elements attribute |
channel_handle | ::mlir::mhlo::ChannelHandleAttr | two 64-bit integers 'handle' and 'type' |
Operands:
| Operand | Description |
|---|---|
operand |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
Results:
| Result | Description |
|---|---|
| «unnamed» | ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
mhlo.compare (mhlo::CompareOp)
Compare operation
Syntax:
operation ::= `mhlo.compare` $comparison_direction `,` $lhs `,` $rhs (`,` $compare_type^)?
attr-dict `:` functional-type(operands, results)
Performs element-wise comparison of lhs and rhs tensors according to
comparison_direction and compare_type, and produces a result tensor.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#compare
Example:
%result = mhlo.compare LT, %lhs, %rhs, FLOAT : (tensor<2xf32>, tensor<2xf32>) -> tensor<2xi1>
Traits: AlwaysSpeculatableImplTrait, Elementwise, InferTensorType, SameOperandsAndResultShape, SameOperandsElementType
Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Attributes:
| Attribute | MLIR Type | Description |
|---|---|---|
comparison_direction | ::mlir::mhlo::ComparisonDirectionAttr | Which comparison operation to perform. |
compare_type | ::mlir::mhlo::ComparisonTypeAttr | Which comparison type to use. |
Operands:
| Operand | Description |
|---|---|
lhs |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
rhs |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
Results:
| Result | Description |
|---|---|
| «unnamed» | ranked tensor of bool values |
mhlo.complex (mhlo::ComplexOp)
Complex operation
Syntax:
operation ::= `mhlo.complex` operands attr-dict
`:` custom<ComplexOpType>(type($lhs), type($rhs), type($result))
Performs element-wise conversion to a complex value from a pair of real and
imaginary values, lhs and rhs, and produces a result tensor.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#complex
Example:
%result = mhlo.complex %lhs, %rhs : tensor<2xcomplex<f32>>
Traits: AlwaysSpeculatableImplTrait, Elementwise, SameOperandsAndResultShape, SameOperandsElementType
Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Operands:
| Operand | Description |
|---|---|
lhs |
ranked tensor of 32/64-bit float values |
rhs |
ranked tensor of 32/64-bit float values |
Results:
| Result | Description |
|---|---|
result |
ranked tensor of complex type with 32/64-bit float elements values |
mhlo.composite (mhlo::CompositeOp)
Composite operation
Syntax:
operation ::= `mhlo.composite` $name $inputs attr-dict `:` functional-type(operands, results)
Encapsulates an operation made up (composed) of other StableHLO operations,
taking inputs and composite_attributes and producing results. The
semantics of the op are implemented by the decomposition attribute. The
composite op can be replaced with its decomposition without changing program
semantics. In cases where inlining the decomposition does not provide the same
op semantics, prefer using custom_call.
The version field (defaults to 0) is used to denote when a composite's
semantics change.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#composite
Example:
%results = mhlo.composite "my.op" %arg0, %arg1 {
decomposition = @my_op,
composite_attributes = { my_attribute = "my_value" },
version = 1 : i32
} : (tensor<f32>, tensor<f32>) -> tensor<f32>
Interfaces: SymbolUserOpInterface
Attributes:
| Attribute | MLIR Type | Description |
|---|---|---|
name | ::mlir::StringAttr | string attribute |
composite_attributes | ::mlir::DictionaryAttr | dictionary of named attribute values |
decomposition | ::mlir::FlatSymbolRefAttr | flat symbol reference attribute |
version | ::mlir::IntegerAttr | 32-bit signless integer attribute |
Operands:
| Operand | Description |
|---|---|
inputs |
variadic of ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values or token or nested tuple with any combination of ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized values or memref of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized values or ranked tensor of per-axis integer quantized values or token values |
Results:
| Result | Description |
|---|---|
| «unnamed» | variadic of ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values or token or nested tuple with any combination of ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized values or memref of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized values or ranked tensor of per-axis integer quantized values or token values |
mhlo.concatenate (mhlo::ConcatenateOp)
Concatenate operation
Concatenates a variadic number of tensors in inputs along dimension
dimension in the same order as the given arguments and produces a result
tensor.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#concatenate
Example:
%result = mhlo.concatenate %input0, %input1, dim = 0 : (tensor<3x2xi64>, tensor<1x2xi64>) -> tensor<4x2xi64>
Traits: AlwaysSpeculatableImplTrait, SameOperandsAndResultElementType
Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Attributes:
| Attribute | MLIR Type | Description |
|---|---|---|
dimension | ::mlir::IntegerAttr | 64-bit signless integer attribute whose value is non-negative |
Operands:
| Operand | Description |
|---|---|
val |
variadic of ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
Results:
| Result | Description |
|---|---|
| «unnamed» | ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
mhlo.constant (mhlo::ConstantOp)
Constant operation
Produces an output tensor from a constant value.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#constant
Example:
%output = mhlo.constant dense<[[0.0, 1.0], [2.0, 3.0]]> : tensor<2x2xf32>
Traits: AlwaysSpeculatableImplTrait, ConstantLike
Interfaces: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Attributes:
| Attribute | MLIR Type | Description |
|---|---|---|
value | ::mlir::ElementsAttr | constant vector/tensor attribute |
Results:
| Result | Description |
|---|---|
output |
statically shaped tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
mhlo.convert (mhlo::ConvertOp)
Convert operation
Syntax:
operation ::= `mhlo.convert` $operand attr-dict
`:` custom<SameOperandsAndResultType>(type($operand), type($result))
Performs an element-wise conversion from one element type to another on
operand tensor and produces a result tensor.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#convert
Example:
%result = mhlo.convert %operand : (tensor<3xi32>) -> tensor<3xcomplex<f32>>
Traits: AlwaysSpeculatableImplTrait, Elementwise, SameOperandsAndResultShape
Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Operands:
| Operand | Description |
|---|---|
operand |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
Results:
| Result | Description |
|---|---|
result |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
mhlo.convolution (mhlo::ConvolutionOp)
Convolution operation
Syntax:
operation ::= `mhlo.convolution` `(`operands`)`
`dim_numbers` `=` custom<ConvolutionDimensions>($dimension_numbers) `,`
`window` `=` `{` custom<WindowAttributes>($window_strides, $padding,
$lhs_dilation, $rhs_dilation,
$window_reversal) `}`
attr-dict `:` functional-type(operands, results)
Computes dot products between windows of lhs and slices of rhs and
produces result.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#convolution
Example:
%result = "mhlo.convolution"(%lhs, %rhs) {
window_strides = dense<4> : tensor<2xi64>,
padding = dense<0> : tensor<2x2xi64>,
lhs_dilation = dense<2> : tensor<2xi64>,
rhs_dilation = dense<1> : tensor<2xi64>,
window_reversal = dense<false> : tensor<2xi1>,
dimension_numbers = #mhlo.conv<[b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f]>,
feature_group_count = 1 : i64,
batch_group_count = 1 : i64,
precision_config = [#stablehlo<precision DEFAULT>, #stablehlo<precision DEFAULT>]
} : (tensor<1x4x4x1xi32>, tensor<3x3x1x1xi32>) -> tensor<1x2x2x1xi32>
Traits: AlwaysSpeculatableImplTrait
Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Attributes:
| Attribute | MLIR Type | Description |
|---|---|---|
window_strides | ::mlir::DenseIntElementsAttr | 64-bit signless integer elements attribute |
padding | ::mlir::DenseIntElementsAttr | 64-bit signless integer elements attribute |
lhs_dilation | ::mlir::DenseIntElementsAttr | 64-bit signless integer elements attribute |
rhs_dilation | ::mlir::DenseIntElementsAttr | 64-bit signless integer elements attribute |
window_reversal | ::mlir::DenseElementsAttr | constant boolean vector/tensor attribute |
dimension_numbers | ::mlir::mhlo::ConvDimensionNumbersAttr | Structure of dimension information for conv op |
feature_group_count | ::mlir::IntegerAttr | 64-bit signless integer attribute whose value is positive |
batch_group_count | ::mlir::IntegerAttr | 64-bit signless integer attribute whose value is positive |
precision_config | ::mlir::ArrayAttr | Precision Config attribute |
Operands:
| Operand | Description |
|---|---|
lhs |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
rhs |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
Results:
| Result | Description |
|---|---|
| «unnamed» | ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
mhlo.copy (mhlo::CopyOp)
Copy operation
Syntax:
operation ::= `mhlo.copy` operands attr-dict
`:` custom<SameOperandsAndResultType>(type($operand), type($result))
This operation is private to the XLA compiler, so it is does not yet have a specification.
Informally, this operation a copy of operand. Depending on the metadata
attached to the operation, it can behave quite differently from a no-op.
Example:
%0 = mhlo.copy %arg0 : tensor<f32>
Traits: AlwaysSpeculatableImplTrait, CompatibleOperandsAndResultType, Elementwise
Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Attributes:
| Attribute | MLIR Type | Description |
|---|---|---|
cross_program_prefetch_index | ::mlir::IntegerAttr | 32-bit signless integer attribute |
Operands:
| Operand | Description |
|---|---|
operand |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values or token or nested tuple with any combination of ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized values or memref of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized values or ranked tensor of per-axis integer quantized values or token values |
Results:
| Result | Description |
|---|---|
result |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values or token or nested tuple with any combination of ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized values or memref of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized values or ranked tensor of per-axis integer quantized values or token values |
mhlo.cosh (mhlo::CoshOp)
Cosh operation
Syntax:
operation ::= `mhlo.cosh` $operand attr-dict
`:` custom<SameOperandsAndResultType>(type($operand), type($result))
Performs element-wise cosh operation on operand tensor and produces a
result tensor.
Example:
%result = mhlo.cosh %operand : tensor<2x2xf32>
Traits: CompatibleOperandsAndResultType, Elementwise, SameOperandsAndResultShape
Interfaces: InferShapedTypeOpInterface, InferTypeOpInterface
Operands:
| Operand | Description |
|---|---|
operand |
tensor of 4/6/8/16/32/64-bit float or complex type with 32/64-bit float elements values |
Results:
| Result | Description |
|---|---|
result |
tensor of 4/6/8/16/32/64-bit float or complex type with 32/64-bit float elements values |
mhlo.cosine (mhlo::CosineOp)
Cosine operation
Syntax:
operation ::= `mhlo.cosine` $operand attr-dict
`:` custom<SameOperandsAndResultType>(type($operand), type($result))
Performs element-wise cosine operation on operand tensor and produces a
result tensor.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#cosine
Example:
%result = mhlo.cosine %operand : tensor<2xf32>
Traits: AlwaysSpeculatableImplTrait, CompatibleOperandsAndResultType, Elementwise, SameOperandsAndResultShape
Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Attributes:
| Attribute | MLIR Type | Description |
|---|---|---|
result_accuracy | ::mlir::mhlo::ResultAccuracyAttr | The requested accuracy for unary ops. |
Operands:
| Operand | Description |
|---|---|
operand |
ranked tensor of 4/6/8/16/32/64-bit float or complex type with 32/64-bit float elements or per-tensor integer quantized values |
Results:
| Result | Description |
|---|---|
result |
ranked tensor of 4/6/8/16/32/64-bit float or complex type with 32/64-bit float elements or per-tensor integer quantized values |
mhlo.count_leading_zeros (mhlo::ClzOp)
Clz operation
Syntax:
operation ::= `mhlo.count_leading_zeros` $operand attr-dict
`:` custom<SameOperandsAndResultType>(type($operand), type($result))
Performs element-wise count of the number of leading zero bits in the
operand tensor and produces a result tensor.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#count_leading_zeros
Example:
%result = mhlo.count_leading_zeros %operand : tensor<2x2xi8>
Traits: AlwaysSpeculatableImplTrait, CompatibleOperandsAndResultType, Elementwise, SameOperandsAndResultShape
Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Operands:
| Operand | Description |
|---|---|
operand |
ranked tensor of 2/4/8/16/32/64-bit integer values |
Results:
| Result | Description |
|---|---|
result |
ranked tensor of 2/4/8/16/32/64-bit integer values |
mhlo.create_token (mhlo::CreateTokenOp)
CreateToken operation
Syntax:
operation ::= `mhlo.create_token` attr-dict `:` type(results)
This operation is on its way out of StableHLO, so it is not included in the specification: https://github.com/openxla/stablehlo/issues/3
Informally, this operation does the same thing as AfterAllOp with 0 inputs: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#after_all
Example:
%output = mhlo.create_token : !mhlo.token
Traits: AlwaysSpeculatableImplTrait
Interfaces: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Results:
| Result | Description |
|---|---|
output |
token |
mhlo.cross-replica-sum (mhlo::CrossReplicaSumOp)
CrossReplicaSum operation
This operation is on its way out of StableHLO, so it is not included in the specification: https://github.com/openxla/stablehlo/issues/3
Informally, this operation does the same thing as AllReduceOp with
channel_id = 0, use_global_device_ids = false and computation
implementing addition:
https://github.com/openxla/stablehlo/blob/main/docs/spec.md#all_reduce
Example:
%result = "mhlo.cross-replica-sum"(%operand) {
replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>
} : (tensor<4xf32>) -> tensor<4xf32>
Traits: AlwaysSpeculatableImplTrait, CompatibleOperandsAndResultType
Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Attributes:
| Attribute | MLIR Type | Description |
|---|---|---|
replica_groups | ::mlir::DenseIntElementsAttr | 64-bit signless integer elements attribute |
Operands:
| Operand | Description |
|---|---|
operand |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
Results:
| Result | Description |
|---|---|
| «unnamed» | ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
mhlo.custom_call (mhlo::CustomCallOp)
CustomCall operation
Syntax:
operation ::= `mhlo.custom_call` custom<CustomCallTarget>($call_target_name) `(` $inputs `)`
attr-dict `:` functional-type(operands, results)
Encapsulates an implementation-defined operation call_target_name that
takes inputs and called_computations and produces results.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#custom_call
Example:
%results = "mhlo.custom_call"(%input0) {
call_target_name = "foo",
has_side_effect = false,
backend_config = "bar",
api_version = 1 : i32,
called_computations = [@foo]
} : (tensor<f32>) -> tensor<f32>
A custom call invokes code external to XLA. The `inputs` are passed to the
external code, and the external code is expected to produce a result of the
given type. The exact mechanism is backend-specific. For example, in the CPU
backend, a call instruction is emitted which targets a symbol with the name
`call_target_name`.
If XLA runtime is enabled for a backend, then custom calls use the runtime
custom call calling convention to call into the external functions. This
calling convention defines an ABI for encoding arguments, attributes and
results.
Depending on the API version there are two ways to pass extra bits of static
information to the external function:
1. For `API_VERSION_TYPED_FFI` custom calls `backend_config` must be a
dictionary attribute, that will be encoded according to the custom call
calling convention and passed to the external function as the attributes
argument. External code is expected to use declarative bindings (see
`xla/runtime/custom_call.h`) to decode them at run time. These custom
calls are only supported if XLA uses XLA runtime.
2. For previous API versions it is the user responsibility to encode extra
bits of static information as a string `backend_config` attribute, and
decode it at run time.
Interfaces: MemoryEffectOpInterface
Attributes:
| Attribute | MLIR Type | Description |
|---|---|---|
call_target_name | ::mlir::StringAttr | string attribute |
has_side_effect | ::mlir::BoolAttr | bool attribute |
backend_config | ::mlir::Attribute | string attribute or dictionary of named attribute values |
api_version | ::mlir::mhlo::CustomCallApiVersionAttr | Custom call API version |
called_computations | ::mlir::ArrayAttr | flat symbol ref array attribute |
custom_call_schedule | ::mlir::mhlo::CustomCallScheduleAttr | Specifies the desired schedule for the custom-call. |
operand_layouts | ::mlir::ArrayAttr | Array of layout (1D tensor of index type) attributes |
result_layouts | ::mlir::ArrayAttr | Array of layout (1D tensor of index type) attributes |
output_operand_aliases | ::mlir::ArrayAttr | Aliasing attribute for outputs and operands of CustomCall |
Operands:
| Operand | Description |
|---|---|
inputs |
variadic of tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values or memref of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized values or token or nested tuple with any combination of tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values or memref of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized values or token values |
Results:
| Result | Description |
|---|---|
| «unnamed» | variadic of tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values or memref of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized values or token or nested tuple with any combination of tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values or memref of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized values or token values |
mhlo.divide (mhlo::DivOp)
Div operation
Syntax:
operation ::= `mhlo.divide` $lhs `,` $rhs attr-dict
`:` custom<SameOperandsAndResultType>(type($lhs), type($rhs), type($result))
Performs element-wise division of dividend lhs and divisor rhs tensors
and produces a result tensor.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#divide
Example:
%result = mhlo.divide %lhs, %rhs : tensor<4xf32>
Traits: AlwaysSpeculatableImplTrait, CompatibleOperandsAndResultType, Elementwise, SameOperandsAndResultShape
Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Operands:
| Operand | Description |
|---|---|
lhs |
ranked tensor of 2/4/8/16/32/64-bit integer or 4/6/8/16/32/64-bit float or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
rhs |
ranked tensor of 2/4/8/16/32/64-bit integer or 4/6/8/16/32/64-bit float or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
Results:
| Result | Description |
|---|---|
result |
ranked tensor of 2/4/8/16/32/64-bit integer or 4/6/8/16/32/64-bit float or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
mhlo.domain (mhlo::DomainOp)
Domain operation
This operation is private to the XLA compiler, so it is does not yet have a specification.
Informally, these operations are used to group instructions with the same DomainMetadata property. ShardingMetadata is the main use case today to group instructions on the same device. Domain instructions provide two major benefits:
- Prevent unintentionally optimizing instructions across domains.
- Automatically assign the metadata of the instructions created in the domain. Without domain instructions, each HLO optimization pass would have to check and propagate the metadata, which would be easy to miss and also adds complexity to the compiler. Since domain instructions connect two different domains, each domain instruction is associated with two DomainMetadata -- one on the operand side and one on the user side of the domain.
Traits: AlwaysSpeculatableImplTrait, CompatibleOperandsAndResultType
Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Attributes:
| Attribute | MLIR Type | Description |
|---|---|---|
kind | ::mlir::mhlo::DomainKindAttr | Kind of domain metatdata attached to an HLO domain. |
entry_metadata | ::mlir::StringAttr | string attribute |
exit_metadata | ::mlir::StringAttr | string attribute |
Operands:
| Operand | Description |
|---|---|
operand |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized values or ranked tensor of per-axis integer quantized values or token |
Results:
| Result | Description |
|---|---|
result |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized values or ranked tensor of per-axis integer quantized values or token |
mhlo.dot (mhlo::DotOp)
Dot operation
This operation is on its way out of StableHLO, so it is not included in the specification: https://github.com/openxla/stablehlo/issues/3
Informally, this operation does the same thing as XLA's Dot: https://www.tensorflow.org/xla/operation_semantics#dot
Example:
%0 = mhlo.dot %arg0, %arg1 : (tensor<1x2xi32>, tensor<2x1xi32>) -> tensor<1x1xi32>
Traits: AlwaysSpeculatableImplTrait
Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Attributes:
| Attribute | MLIR Type | Description |
|---|---|---|
precision_config | ::mlir::ArrayAttr | Precision Config attribute |
Operands:
| Operand | Description |
|---|---|
lhs |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
rhs |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
Results:
| Result | Description |
|---|---|
| «unnamed» | ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
mhlo.dot_general (mhlo::DotGeneralOp)
DotGeneral operation
Computes dot products between slices of lhs and slices of rhs and
produces a result tensor.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#dot_general
Example:
%result = "mhlo.dot_general"(%lhs, %rhs) {
dot_dimension_numbers = #mhlo.dot<
lhs_batching_dimensions = [0],
rhs_batching_dimensions = [0],
lhs_contracting_dimensions = [2],
rhs_contracting_dimensions = [1]
>,
precision_config = [#stablehlo<precision DEFAULT>, #stablehlo<precision DEFAULT>]
} : (tensor<2x2x2xi32>, tensor<2x2x2xi32>) -> tensor<2x2x2xi32>
Traits: AlwaysSpeculatableImplTrait
Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Attributes:
| Attribute | MLIR Type | Description |
|---|---|---|
dot_dimension_numbers | ::mlir::mhlo::DotDimensionNumbersAttr | Attribute that models the dimension information for dot. |
precision_config | ::mlir::ArrayAttr | Precision Config attribute |
algorithm | ::mlir::mhlo::DotAlgorithmAttr | Attribute that models the algorithm constraints to use for computing dot. |
Operands:
| Operand | Description |
|---|---|
lhs |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
rhs |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
Results:
| Result | Description |
|---|---|
| «unnamed» | ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
mhlo.dynamic_broadcast_in_dim (mhlo::DynamicBroadcastInDimOp)
DynamicBroadcastInDim operation
This operation is functionally identical to
broadcast_in_dim
op, but the result shape is specified dynamically via output_dimensions.
It also accepts optional attributes to express static knowledge about the expanding behavior of dimensions. If not specified, all dimensions are assumed to be possibly expanding. The sets of dimensions that are known to be expanding and the set of dimensions that are known to be non-expanding must be disjoint and they must be a subset of the operand's dimensions.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#dynamic_broadcast_in_dim
Example:
%operand = mhlo.constant dense<[[1, 2, 3]]> : tensor<1x3xi64>
%output_dimensions = mhlo.constant dense<[2, 3, 2]> : tensor<3xi64>
%result = "mhlo.dynamic_broadcast_in_dim"(%operand, %output_dimensions) {
broadcast_dimensions = array<i64: 2, 1>,
known_expanding_dimensions = array<i64: 0>,
known_nonexpanding_dimensions = array<i64: 1>
} : (tensor<1x3xi64>, tensor<3xi64>) -> tensor<2x3x2xi64>
Traits: AlwaysSpeculatableImplTrait
Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Attributes:
| Attribute | MLIR Type | Description |
|---|---|---|
broadcast_dimensions | ::mlir::DenseIntElementsAttr | 64-bit signless integer elements attribute |
known_expanding_dimensions | ::mlir::DenseIntElementsAttr | 64-bit signless integer elements attribute |
known_nonexpanding_dimensions | ::mlir::DenseIntElementsAttr | 64-bit signless integer elements attribute |
Operands:
| Operand | Description |
|---|---|
operand |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
output_dimensions |
1D tensor of index or 2/4/8/16/32/64-bit integer values |
Results:
| Result | Description |
|---|---|
| «unnamed» | ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
mhlo.dynamic_conv (mhlo::DynamicConvOp)
DynamicConv operation
This operation is a work in progress, so it is not yet included in the specification: https://github.com/openxla/stablehlo/issues/8
Informally, this operation does the same thing as ConvolutionOp except
that padding is specified dynamically via d_padding:
https://github.com/openxla/stablehlo/blob/main/docs/spec.md#convolution
Example:
%result = "mhlo.dynamic_conv"(%lhs, %rhs, %d_padding) {
window_strides = dense<4> : tensor<2xi64>,
lhs_dilation = dense<2> : tensor<2xi64>,
rhs_dilation = dense<1> : tensor<2xi64>,
window_reversal = dense<false> : tensor<2xi1>,
dimension_numbers = #mhlo.conv<[b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f]>,
feature_group_count = 1 : i64,
batch_group_count = 1 : i64,
precision_config = [#stablehlo<precision DEFAULT>, #stablehlo<precision DEFAULT>]
} : (tensor<1x4x4x1xi32>, tensor<3x3x1x1xi32>, tensor<2x2xi64>) -> tensor<1x2x2x1xi32>
Traits: AlwaysSpeculatableImplTrait
Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Attributes:
| Attribute | MLIR Type | Description |
|---|---|---|
window_strides | ::mlir::DenseIntElementsAttr | 64-bit signless integer elements attribute |
padding | ::mlir::DenseIntElementsAttr | 64-bit signless integer elements attribute |
lhs_dilation | ::mlir::DenseIntElementsAttr | 64-bit signless integer elements attribute |
rhs_dilation | ::mlir::DenseIntElementsAttr | 64-bit signless integer elements attribute |
window_reversal | ::mlir::DenseElementsAttr | constant boolean vector/tensor attribute |
dimension_numbers | ::mlir::mhlo::ConvDimensionNumbersAttr | Structure of dimension information for conv op |
feature_group_count | ::mlir::IntegerAttr | 64-bit signless integer attribute whose value is positive |
batch_group_count | ::mlir::IntegerAttr | 64-bit signless integer attribute whose value is positive |
precision_config | ::mlir::ArrayAttr | Precision Config attribute |
Operands:
| Operand | Description |
|---|---|
lhs |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
rhs |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
d_padding |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
Results:
| Result | Description |
|---|---|
| «unnamed» | ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
mhlo.dynamic_gather (mhlo::DynamicGatherOp)
DynamicGather operation
This operation is a work in progress, so it is not yet included in the specification: https://github.com/openxla/stablehlo/issues/8
Informally, this operation does the same thing as GatherOp except
that slice_sizes are specified dynamically:
https://github.com/openxla/stablehlo/blob/main/docs/spec.md#gather
Example:
%result = "mhlo.dynamic_gather"(%operand, %start_indices, %slice_sizes) {
dimension_numbers = #mhlo.gather<
offset_dims = [2, 3],
collapsed_slice_dims = [0],
start_index_map = [0, 2],
index_vector_dim = 2>,
indices_are_sorted = false
} : (tensor<3x4x2xi32>, tensor<2x3x2xi64>, tensor<3xi64>) -> tensor<2x3x2x2xi32>
Traits: AlwaysSpeculatableImplTrait, InferTensorType
Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Attributes:
| Attribute | MLIR Type | Description |
|---|---|---|
dimension_numbers | ::mlir::mhlo::GatherDimensionNumbersAttr | Attribute that models the dimension information for gather |
indices_are_sorted | ::mlir::BoolAttr | bool attribute |
Operands:
| Operand | Description |
|---|---|
operand |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
start_indices |
ranked tensor of 2/4/8/16/32/64-bit integer values |
slice_sizes |
statically shaped 1-dimensional integer tensor of 2/4/8/16/32/64-bit integer values |
Results:
| Result | Description |
|---|---|
| «unnamed» | ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
mhlo.dynamic_iota (mhlo::DynamicIotaOp)
DynamicIota operation
This operation is functionally identical to
iota
op, but the result shape is specified dynamically via output_shape.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#dynamic_iota
Example:
%0 = mhlo.dynamic_iota %arg0, dim = 0 : (tensor<1xindex>) -> tensor<4xi32>
Traits: AlwaysSpeculatableImplTrait
Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Attributes:
| Attribute | MLIR Type | Description |
|---|---|---|
iota_dimension | ::mlir::IntegerAttr | 64-bit signless integer attribute whose value is non-negative |
Operands:
| Operand | Description |
|---|---|
output_shape |
1D tensor of index or 2/4/8/16/32/64-bit integer values |
Results:
| Result | Description |
|---|---|
result |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
mhlo.dynamic_pad (mhlo::DynamicPadOp)
DynamicPad operation
Syntax:
operation ::= `mhlo.dynamic_pad` operands attr-dict `:` functional-type(operands, results)
Dynamically Pads the operand, with amount of padding added at
low-end/high-end/interior is passed through input tensors.
Traits: AlwaysSpeculatableImplTrait
Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Operands:
| Operand | Description |
|---|---|
operand |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
padding_value |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
edge_padding_low |
1D tensor of index or 2/4/8/16/32/64-bit integer values |
edge_padding_high |
1D tensor of index or 2/4/8/16/32/64-bit integer values |
interior_padding |
1D tensor of index or 2/4/8/16/32/64-bit integer values |
Results:
| Result | Description |
|---|---|
result |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
mhlo.dynamic_reshape (mhlo::DynamicReshapeOp)
DynamicReshape operation
Syntax:
operation ::= `mhlo.dynamic_reshape` operands attr-dict `:` functional-type(operands, results)
This operation is functionally identical to
reshape
op, but the result shape is specified dynamically via output_shape.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#dynamic_reshape
Example:
%output_shape = mhlo.constant dense<[3, 2]> : tensor<2xi64>
%result = mhlo.dynamic_reshape %operand, %output_shape : (tensor<2x3xi64>, tensor<2xi64>) -> tensor<3x2xi64>
Traits: AlwaysSpeculatableImplTrait
Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Operands:
| Operand | Description |
|---|---|
operand |
tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
output_shape |
1D tensor of index or 2/4/8/16/32/64-bit integer values |
Results:
| Result | Description |
|---|---|
result |
tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
mhlo.dynamic_slice (mhlo::DynamicSliceOp)
DynamicSlice operation
Extracts a slice from the operand using dynamically-computed starting
indices and produces a result tensor.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#dynamic_slice
Example:
%result = mhlo.dynamic_slice %operand, %start_indices0, %start_indices1, sizes = [2, 2]
: (tensor<4x4xi32>, tensor<i64>, tensor<i64>) -> tensor<2x2xi32>
Traits: AlwaysSpeculatableImplTrait, InferTensorType
Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Attributes:
| Attribute | MLIR Type | Description |
|---|---|---|
slice_sizes | ::mlir::DenseIntElementsAttr | 64-bit signless integer elements attribute |
Operands:
| Operand | Description |
|---|---|
operand |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
start_indices |
variadic of 0D tensor of 2/4/8/16/32/64-bit integer values |
Results:
| Result | Description |
|---|---|
result |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
mhlo.dynamic_update_slice (mhlo::DynamicUpdateSliceOp)
DynamicUpdateSlice operation
Syntax:
operation ::= `mhlo.dynamic_update_slice` operands attr-dict `:` functional-type(operands, results)
Produces a result tensor which is equal to the operand tensor except
that the slice starting at start_indices is updated with the values in
update.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#dynamic_update_slice
Example:
%result = mhlo.dynamic_update_slice %operand, %update, %start_indices0, %start_indices1
: (tensor<4x4xi32>, tensor<2x2xi32>, tensor<i64>, tensor<i64>) -> tensor<4x4xi32>
Traits: AlwaysSpeculatableImplTrait, InferTensorType
Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Operands:
| Operand | Description |
|---|---|
operand |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
update |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
start_indices |
variadic of 0D tensor of 2/4/8/16/32/64-bit integer values |
Results:
| Result | Description |
|---|---|
result |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
mhlo.einsum (mhlo::EinsumOp)
Einsum operation
This operation is on its way out of StableHLO, so it is not included in the specification: https://github.com/openxla/stablehlo/issues/3
Informally, this operation does the same thing as TF's einsum: https://www.tensorflow.org/api_docs/python/tf/einsum
Example:
%result = "mhlo.einsum"(%lhs, %rhs) {
einsum_config = "ab,bc->ac"
} : (tensor<4x16xf32>, tensor<16x4xf32>) -> tensor<4x4xf32>
Traits: AlwaysSpeculatableImplTrait
Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Attributes:
| Attribute | MLIR Type | Description |
|---|---|---|
einsum_config | ::mlir::StringAttr | string attribute |
Operands:
| Operand | Description |
|---|---|
lhs |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
rhs |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
Results:
| Result | Description |
|---|---|
| «unnamed» | ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
mhlo.erf (mhlo::ErfOp)
Erf operation
Syntax:
operation ::= `mhlo.erf` $operand attr-dict
`:` custom<SameOperandsAndResultType>(type($operand), type($result))
Performs element-wise erf operation on operand tensor and produces a
result tensor.
Example:
%result = mhlo.erf %operand : tensor<2x2xf32>
Traits: AlwaysSpeculatableImplTrait, CompatibleOperandsAndResultType, Elementwise, SameOperandsAndResultShape
Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Operands:
| Operand | Description |
|---|---|
operand |
ranked tensor of 4/6/8/16/32/64-bit float values |
Results:
| Result | Description |
|---|---|
result |
ranked tensor of 4/6/8/16/32/64-bit float values |
mhlo.exponential (mhlo::ExpOp)
Exp operation
Syntax:
operation ::= `mhlo.exponential` $operand attr-dict
`:` custom<SameOperandsAndResultType>(type($operand), type($result))
Performs element-wise exponential operation on operand tensor and produces
a result tensor.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#exponential
Example:
%result = mhlo.exponential %operand : tensor<2x2xf64>
Traits: AlwaysSpeculatableImplTrait, CompatibleOperandsAndResultType, Elementwise, SameOperandsAndResultShape
Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Attributes:
| Attribute | MLIR Type | Description |
|---|---|---|
result_accuracy | ::mlir::mhlo::ResultAccuracyAttr | The requested accuracy for unary ops. |
Operands:
| Operand | Description |
|---|---|
operand |
ranked tensor of 4/6/8/16/32/64-bit float or complex type with 32/64-bit float elements or per-tensor integer quantized values |
Results:
| Result | Description |
|---|---|
result |
ranked tensor of 4/6/8/16/32/64-bit float or complex type with 32/64-bit float elements or per-tensor integer quantized values |
mhlo.exponential_minus_one (mhlo::Expm1Op)
Expm1 operation
Syntax:
operation ::= `mhlo.exponential_minus_one` $operand attr-dict
`:` custom<SameOperandsAndResultType>(type($operand), type($result))
Performs element-wise exponential minus one operation on operand tensor
and produces a result tensor.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#exponential_minus_one
Example:
%result = mhlo.exponential_minus_one %operand : tensor<2xf32>
Traits: AlwaysSpeculatableImplTrait, CompatibleOperandsAndResultType, Elementwise, SameOperandsAndResultShape
Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Attributes:
| Attribute | MLIR Type | Description |
|---|---|---|
result_accuracy | ::mlir::mhlo::ResultAccuracyAttr | The requested accuracy for unary ops. |
Operands:
| Operand | Description |
|---|---|
operand |
ranked tensor of 4/6/8/16/32/64-bit float or complex type with 32/64-bit float elements or per-tensor integer quantized values |
Results:
| Result | Description |
|---|---|
result |
ranked tensor of 4/6/8/16/32/64-bit float or complex type with 32/64-bit float elements or per-tensor integer quantized values |
mhlo.fft (mhlo::FftOp)
Fft operation
Performs the forward and inverse Fourier transforms for real and complex inputs/outputs.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#fft
Example:
%result = mhlo.fft %operand, type = FFT, length = [4] : (tensor<4xcomplex<f32>>) -> tensor<4xcomplex<f32>>
Traits: AlwaysSpeculatableImplTrait, InferTensorType
Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Attributes:
| Attribute | MLIR Type | Description |
|---|---|---|
fft_type | ::mlir::mhlo::FftTypeAttr | XLA fast fourier transform type. |
fft_length | ::mlir::DenseIntElementsAttr | 64-bit signless integer elements attribute |
Operands:
| Operand | Description |
|---|---|
operand |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
Results:
| Result | Description |
|---|---|
| «unnamed» | ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
mhlo.floor (mhlo::FloorOp)
Floor operation
Syntax:
operation ::= `mhlo.floor` $operand attr-dict
`:` custom<SameOperandsAndResultType>(type($operand), type($result))
Performs element-wise floor of operand tensor and produces a result
tensor.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#floor
Example:
%result = mhlo.floor %operand : tensor<2xf32>
Traits: AlwaysSpeculatableImplTrait, CompatibleOperandsAndResultType, Elementwise, SameOperandsAndResultShape
Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Operands:
| Operand | Description |
|---|---|
operand |
ranked tensor of 4/6/8/16/32/64-bit float or per-tensor integer quantized values |
Results:
| Result | Description |
|---|---|
result |
ranked tensor of 4/6/8/16/32/64-bit float or per-tensor integer quantized values |
mhlo.fusion (mhlo::FusionOp)
Fusion operation
This operation is private to the XLA compiler, so it is does not yet have a specification.
Informally, this operation consists of a group of basic ops (represented as a region attached to it). It serves as a hint to the backend that it is beneficial to emit the contained ops into a single loop nest or kernel.
Attributes:
| Attribute | MLIR Type | Description |
|---|---|---|
fusion_kind | ::mlir::mhlo::FusionKindAttr | fusion kind |
output_operand_aliases | ::mlir::ArrayAttr | Aliasing attribute for outputs and operands of Fusion |
Operands:
| Operand | Description |
|---|---|
inputs |
variadic of ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized values or ranked tensor of per-axis integer quantized values or token |
Results:
| Result | Description |
|---|---|
results |
variadic of ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values or nested tuple with any combination of ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized values or memref of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized values or ranked tensor of per-axis integer quantized values or token values |
mhlo.gather (mhlo::GatherOp)
Gather operation
Gathers slices from operand tensor from offsets specified in
start_indices and produces a result tensor.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#gather
Example:
%result = "mhlo.gather"(%operand, %start_indices) {
dimension_numbers = #stablehlo.gather<
offset_dims = [3, 4],
collapsed_slice_dims = [1],
operand_batching_dims = [0],
start_indices_batching_dims = [1],
start_index_map = [2, 1],
index_vector_dim = 3>,
slice_sizes = dense<[0, 2, 2]> : tensor<3xi64>,
indices_are_sorted = false
} : (tensor<2x3x4x2xi64>, tensor<2x2x3x2xi64>) -> tensor<2x2x3x2x2xi64>
Traits: AlwaysSpeculatableImplTrait, InferTensorType
Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Attributes:
| Attribute | MLIR Type | Description |
|---|---|---|
dimension_numbers | ::mlir::mhlo::GatherDimensionNumbersAttr | Attribute that models the dimension information for gather |
slice_sizes | ::mlir::DenseIntElementsAttr | 64-bit signless integer elements attribute |
indices_are_sorted | ::mlir::BoolAttr | bool attribute |
Operands:
| Operand | Description |
|---|---|
operand |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
start_indices |
ranked tensor of 2/4/8/16/32/64-bit integer values |
Results:
| Result | Description |
|---|---|
| «unnamed» | ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
mhlo.get_dimension_size (mhlo::GetDimensionSizeOp)
GetDimensionSize operation
Produces the size of the given dimension of the operand.
See https://github.com/openxla/stablehlo/blob/main/docs/spec.md#get_dimension_size
Example:
%result = mhlo.get_dimension_size %operand, dim = 1 : (tensor<2x3xf32>) -> tensor<i32>
Traits: AlwaysSpeculatableImplTrait, InferTensorType
Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Attributes:
| Attribute | MLIR Type | Description |
|---|---|---|
dimension | ::mlir::IntegerAttr | 64-bit signless integer attribute whose value is non-negative |
Operands:
| Operand | Description |
|---|---|
operand |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
Results:
| Result | Description |
|---|---|
| «unnamed» | tensor of 32-bit signless integer values |
mhlo.get_tuple_element (mhlo::GetTupleElementOp)
GetTupleElement operation
Syntax:
operation ::= `mhlo.get_tuple_element` $operand `[` $index `]` attr-dict `:` functional-type(operands, results)
Extracts element at index position of the operand tuple and produces a
result.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#get_tuple_element
Example:
%result = mhlo.get_tuple_element %operand[0] : (tuple<tensor<2xf32>, tuple<tensor<i32>>>) -> tensor<2xf32>
Traits: AlwaysSpeculatableImplTrait
Interfaces: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Attributes:
| Attribute | MLIR Type | Description |
|---|---|---|
index | ::mlir::IntegerAttr | 32-bit signless integer attribute whose value is non-negative |
Operands:
| Operand | Description |
|---|---|
operand |
nested tuple with any combination of ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized values or memref of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized values or ranked tensor of per-axis integer quantized values or token values |
Results:
| Result | Description |
|---|---|
| «unnamed» | ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values or token or nested tuple with any combination of ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized values or memref of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized values or ranked tensor of per-axis integer quantized values or token values or memref of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized values |
mhlo.if (mhlo::IfOp)
If operation
Produces the output from executing exactly one branch from true_branch or
false_branch depending on the value of pred.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#if
Example:
%result = "mhlo.if"(%pred) ({
"mhlo.return"(%result_true_branch) : (tensor
Traits: RecursiveMemoryEffects, SingleBlockImplicitTerminator<ReturnOp>, SingleBlock
Interfaces: InferTypeOpInterface
Operands:
| Operand | Description |
|---|---|
pred |
ranked tensor of bool values |
Results:
| Result | Description |
|---|---|
| «unnamed» | variadic of ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized values or ranked tensor of per-axis integer quantized values or token |
mhlo.imag (mhlo::ImagOp)
Imag operation
Syntax:
operation ::= `mhlo.imag` $operand attr-dict
`:` custom<SameOperandsAndResultType>(type($operand), type($result))
Extracts the imaginary part, element-wise, from the operand and produces a
result tensor.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#imag
Example:
%result = mhlo.imag %operand : (tensor<2xcomplex<f32>>) -> tensor<2xf32>
Traits: AlwaysSpeculatableImplTrait, Elementwise, SameOperandsAndResultShape
Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Operands:
| Operand | Description |
|---|---|
operand |
ranked tensor of 4/6/8/16/32/64-bit float or complex type with 32/64-bit float elements values |
Results:
| Result | Description |
|---|---|
result |
ranked tensor of 4/6/8/16/32/64-bit float values |
mhlo.infeed (mhlo::InfeedOp)
Infeed operation
Reads data from the infeed and produces results.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#infeed
Example:
%results:2 = "mhlo.infeed"(%token) {
infeed_config = ""
} : (!mhlo.token) -> (tensor<3x3x3xi32>, !mhlo.token)
Attributes:
| Attribute | MLIR Type | Description |
|---|---|---|
infeed_config | ::mlir::StringAttr | string attribute |
layout | ::mlir::ArrayAttr | array attribute |
Operands:
| Operand | Description |
|---|---|
token |
token |
Results:
| Result | Description |
|---|---|
| «unnamed» | variadic of statically shaped tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized values or statically shaped tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values or token |
mhlo.iota (mhlo::IotaOp)
Iota operation
Fills an output tensor with values in increasing order starting from zero
along the iota_dimension dimension.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#iota
Example:
%output = mhlo.iota dim = 0 : tensor<4x5xi32>
Traits: AlwaysSpeculatableImplTrait
Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Attributes:
| Attribute | MLIR Type | Description |
|---|---|---|
iota_dimension | ::mlir::IntegerAttr | 64-bit signless integer attribute whose value is non-negative |
Results:
| Result | Description |
|---|---|
output |
statically shaped tensor of 2/4/8/16/32/64-bit integer or 4/6/8/16/32/64-bit float or complex type with 32/64-bit float elements values |
mhlo.is_finite (mhlo::IsFiniteOp)
IsFinite operation
Syntax:
operation ::= `mhlo.is_finite` $x attr-dict `:` functional-type(operands, results)
Performs element-wise check whether the value in x is finite (i.e. is
neither +Inf, -Inf, nor NaN) and produces a y tensor.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#is_finite
Example:
%y = mhlo.is_finite %x : (tensor<7xf32>) -> tensor<7xi1>
Traits: AlwaysSpeculatableImplTrait, Elementwise, SameOperandsAndResultShape
Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Operands:
| Operand | Description |
|---|---|
x |
ranked tensor of 4/6/8/16/32/64-bit float values |
Results:
| Result | Description |
|---|---|
y |
ranked tensor of bool values |
mhlo.log (mhlo::LogOp)
Log operation
Syntax:
operation ::= `mhlo.log` $operand attr-dict
`:` custom<SameOperandsAndResultType>(type($operand), type($result))
Performs element-wise logarithm operation on operand tensor and produces a
result tensor.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#log
Example:
%result = mhlo.log %operand : tensor<2x2xf64>
Traits: AlwaysSpeculatableImplTrait, CompatibleOperandsAndResultType, Elementwise, SameOperandsAndResultShape
Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Attributes:
| Attribute | MLIR Type | Description |
|---|---|---|
result_accuracy | ::mlir::mhlo::ResultAccuracyAttr | The requested accuracy for unary ops. |
Operands:
| Operand | Description |
|---|---|
operand |
ranked tensor of 4/6/8/16/32/64-bit float or complex type with 32/64-bit float elements or per-tensor integer quantized values |
Results:
| Result | Description |
|---|---|
result |
ranked tensor of 4/6/8/16/32/64-bit float or complex type with 32/64-bit float elements or per-tensor integer quantized values |
mhlo.log_plus_one (mhlo::Log1pOp)
Log1p operation
Syntax:
operation ::= `mhlo.log_plus_one` $operand attr-dict
`:` custom<SameOperandsAndResultType>(type($operand), type($result))
Performs element-wise logarithm plus one operation on operand tensor and
produces a result tensor.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#log_plus_one
Example:
%result = mhlo.log_plus_one %operand : tensor<6xf32>
Traits: AlwaysSpeculatableImplTrait, CompatibleOperandsAndResultType, Elementwise, SameOperandsAndResultShape
Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Attributes:
| Attribute | MLIR Type | Description |
|---|---|---|
result_accuracy | ::mlir::mhlo::ResultAccuracyAttr | The requested accuracy for unary ops. |
Operands:
| Operand | Description |
|---|---|
operand |
ranked tensor of 4/6/8/16/32/64-bit float or complex type with 32/64-bit float elements or per-tensor integer quantized values |
Results:
| Result | Description |
|---|---|
result |
ranked tensor of 4/6/8/16/32/64-bit float or complex type with 32/64-bit float elements or per-tensor integer quantized values |
mhlo.logistic (mhlo::LogisticOp)
Logistic operation
Syntax:
operation ::= `mhlo.logistic` $operand attr-dict
`:` custom<SameOperandsAndResultType>(type($operand), type($result))
Performs element-wise logistic operation on operand tensor and produces a
result tensor.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#logistic
Example:
%result = mhlo.logistic %operand : tensor<2x2xf32>
Traits: AlwaysSpeculatableImplTrait, CompatibleOperandsAndResultType, Elementwise, SameOperandsAndResultShape
Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Attributes:
| Attribute | MLIR Type | Description |
|---|---|---|
result_accuracy | ::mlir::mhlo::ResultAccuracyAttr | The requested accuracy for unary ops. |
Operands:
| Operand | Description |
|---|---|
operand |
ranked tensor of 4/6/8/16/32/64-bit float or complex type with 32/64-bit float elements or per-tensor integer quantized values |
Results:
| Result | Description |
|---|---|
result |
ranked tensor of 4/6/8/16/32/64-bit float or complex type with 32/64-bit float elements or per-tensor integer quantized values |
mhlo.map (mhlo::MapOp)
Map operation
Applies a map function computation to inputs along the dimensions and
produces a result tensor.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#map
Example:
%result = "mhlo.map"(%input0, %input1) ({
^bb0(%arg0: tensor<i32>, %arg1: tensor<i32>):
%0 = mhlo.multiply %arg0, %arg1 : tensor<i32>
mhlo.return %0 : tensor<i32>
}) {
dimensions = dense<[0, 1]> : tensor<2xi64>
} : (tensor<2x2xi32>, tensor<2x2xi32>) -> tensor<2x2xi32>
Traits: InferTensorType, RecursiveMemoryEffects, SameOperandsAndResultShape, SingleBlockImplicitTerminator<ReturnOp>, SingleBlock
Interfaces: InferShapedTypeOpInterface, InferTypeOpInterface
Attributes:
| Attribute | MLIR Type | Description |
|---|---|---|
dimensions | ::mlir::DenseIntElementsAttr | 64-bit signless integer elements attribute |
Operands:
| Operand | Description |
|---|---|
inputs |
variadic of ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
Results:
| Result | Description |
|---|---|
| «unnamed» | ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
mhlo.maximum (mhlo::MaxOp)
Max operation
Syntax:
operation ::= `mhlo.maximum` $lhs `,` $rhs attr-dict
`:` custom<SameOperandsAndResultType>(type($lhs), type($rhs), type($result))
Performs element-wise max operation on tensors lhs and rhs and produces
a result tensor.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#maximum
Example:
%result = mhlo.maximum %lhs, %rhs : tensor<4xf32>
Traits: AlwaysSpeculatableImplTrait, Commutative, CompatibleOperandsAndResultType, Elementwise, SameOperandsAndResultShape
Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Operands:
| Operand | Description |
|---|---|
lhs |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
rhs |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
Results:
| Result | Description |
|---|---|
result |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
mhlo.minimum (mhlo::MinOp)
Min operation
Syntax:
operation ::= `mhlo.minimum` $lhs `,` $rhs attr-dict
`:` custom<SameOperandsAndResultType>(type($lhs), type($rhs), type($result))
Performs element-wise min operation on tensors lhs and rhs and produces a
result tensor.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#minimum
Example:
%result = mhlo.minimum %lhs, %rhs : tensor<4xf32>
Traits: AlwaysSpeculatableImplTrait, Commutative, CompatibleOperandsAndResultType, Elementwise, SameOperandsAndResultShape
Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Operands:
| Operand | Description |
|---|---|
lhs |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
rhs |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
Results:
| Result | Description |
|---|---|
result |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
mhlo.minimum_broadcast_shapes (mhlo::MinimumBroadcastShapesOp)
Minimizes the rank of two or more shapes to be broadcasted
Syntax:
operation ::= `mhlo.minimum_broadcast_shapes` $shapes attr-dict `:` type($shapes) `->` type($results)
Given two or more 1D tensors representing shapes, returns one 1D tensor for
each operand, where operand i corresponds to output i.
The returned tensors have the property that they specify a shape which is a reshape of the corresponding input shape, and the broadcasted output shape (using shape::BroadcastOp) of the returned shapes is a reshape of the broadcasted output shape of the input shapes. Among all possibilities with this property, the one is chosen which minimizes the rank of each returned shape.
The general idea of this op is that it can be used for ops which have a broadcasting semantic to operate on shapes with a possibly smaller rank while preserving equivalence of the computed values. After computing the result of the op using reshaped operands, the result can be reshaped to the result that would have been originally computed.
Here is an example with two input shapes:
mhlo.minimum_broadcast_shapes [1, 2, 3, 1, 2, 1],
[1, 1, 1, 2, 3] -> [6, 2, 1], [2, 3]
The broadcasted output shape of the operands is [1, 2, 3, 1, 2, 3], the broadcasted output shape of the outputs is [6, 2, 3]. These two shapes are reshapes of each other, and also each output is a reshape of the corresponding input.
Traits: AlwaysSpeculatableImplTrait
Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Operands:
| Operand | Description |
|---|---|
shapes |
variadic of 1D tensor of index values |
Results:
| Result | Description |
|---|---|
results |
variadic of 1D tensor of index values |
mhlo.multiply (mhlo::MulOp)
Mul operation
Syntax:
operation ::= `mhlo.multiply` $lhs `,` $rhs attr-dict
`:` custom<SameOperandsAndResultType>(type($lhs), type($rhs), type($result))
Performs element-wise product of two tensors lhs and rhs and produces a
result tensor.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#multiply
Example:
%result = mhlo.multiply %lhs, %rhs : tensor<2xi32>
Traits: AlwaysSpeculatableImplTrait, Commutative, CompatibleOperandsAndResultType, Elementwise, SameOperandsAndResultShape
Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Operands:
| Operand | Description |
|---|---|
lhs |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
rhs |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
Results:
| Result | Description |
|---|---|
result |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
mhlo.negate (mhlo::NegOp)
Neg operation
Syntax:
operation ::= `mhlo.negate` $operand attr-dict
`:` custom<SameOperandsAndResultType>(type($operand), type($result))
Performs element-wise negation of operand tensor and produces a result
tensor.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#negate
Example:
%result = mhlo.negate %operand : tensor<2x3xi32>
Traits: AlwaysSpeculatableImplTrait, CompatibleOperandsAndResultType, Elementwise, SameOperandsAndResultShape
Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Operands:
| Operand | Description |
|---|---|
operand |
ranked tensor of 2/4/8/16/32/64-bit integer or 4/6/8/16/32/64-bit float or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
Results:
| Result | Description |
|---|---|
result |
ranked tensor of 2/4/8/16/32/64-bit integer or 4/6/8/16/32/64-bit float or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
mhlo.not (mhlo::NotOp)
Not operation
Syntax:
operation ::= `mhlo.not` $operand attr-dict
`:` custom<SameOperandsAndResultType>(type($operand), type($result))
Performs element-wise NOT of tensor operand of type integer and produces
a result tensor.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#not
Example:
%result = mhlo.not %operand : tensor<5x3x1xi1>
Traits: AlwaysSpeculatableImplTrait, CompatibleOperandsAndResultType, Elementwise, SameOperandsAndResultShape
Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Operands:
| Operand | Description |
|---|---|
operand |
ranked tensor of bool or 2/4/8/16/32/64-bit integer values |
Results:
| Result | Description |
|---|---|
result |
ranked tensor of bool or 2/4/8/16/32/64-bit integer values |
mhlo.optimization_barrier (mhlo::OptimizationBarrierOp)
OptimizationBarrier operation
Syntax:
operation ::= `mhlo.optimization_barrier` attr-dict ($operand^ `:` custom<PairwiseOpType>(type($operand), type($result))):(`(` `)`)?
Ensures that the operations that produce the operand are executed before any
operations that depend on the result and prevents compiler transformations
from moving operations across the barrier. Other than that, the operation is
an identity, i.e. result = operand.
See https://github.com/openxla/stablehlo/blob/main/docs/spec.md#optimization_barrier
Example:
%result0, %result1 = mhlo.optimization_barrier %operand0, %operand1 : tensor<f32>, tensor<f32>
Traits: AlwaysSpeculatableImplTrait, HLO_PairwiseSameOperandAndResultType
Interfaces: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Operands:
| Operand | Description |
|---|---|
operand |
variadic of ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized values or ranked tensor of per-axis integer quantized values or token |
Results:
| Result | Description |
|---|---|
result |
variadic of ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized values or ranked tensor of per-axis integer quantized values or token |
mhlo.or (mhlo::OrOp)
Or operation
Syntax:
operation ::= `mhlo.or` $lhs `,` $rhs attr-dict
`:` custom<SameOperandsAndResultType>(type($lhs), type($rhs), type($result))
Performs element-wise OR of two tensors lhs and rhs and produces a
result tensor.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#or
Example:
%result = mhlo.or %lhs, %rhs : tensor<2xi1>
Traits: AlwaysSpeculatableImplTrait, Commutative, CompatibleOperandsAndResultType, Elementwise, SameOperandsAndResultShape
Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Operands:
| Operand | Description |
|---|---|
lhs |
ranked tensor of bool or 2/4/8/16/32/64-bit integer values |
rhs |
ranked tensor of bool or 2/4/8/16/32/64-bit integer values |
Results:
| Result | Description |
|---|---|
result |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
mhlo.outfeed (mhlo::OutfeedOp)
Outfeed operation
Writes inputs to the outfeed and produces a result token.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#outfeed
Example:
%result = "mhlo.outfeed"(%input0, %token) {
outfeed_config = ""
} : (tensor<3x3x3xi32>, !mhlo.token) -> !mhlo.token
Interfaces: InferTypeOpInterface
Attributes:
| Attribute | MLIR Type | Description |
|---|---|---|
outfeed_config | ::mlir::StringAttr | string attribute |
Operands:
| Operand | Description |
|---|---|
inputs |
variadic of ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
token |
token |
Results:
| Result | Description |
|---|---|
| «unnamed» | token |
mhlo.pad (mhlo::PadOp)
Pad operation
Expands operand by padding around the tensor as well as between the
elements of the tensor with the given padding_value.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#pad
Example:
%0 = mhlo.pad %arg0, %arg1, low = [0, 1], high = [2, 1], interior = [1, 2]
: (tensor<2x3xi32>, tensor<i32>) -> tensor<5x9xi32>
Traits: AlwaysSpeculatableImplTrait, SameOperandsAndResultElementType
Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Attributes:
| Attribute | MLIR Type | Description |
|---|---|---|
edge_padding_low | ::mlir::DenseIntElementsAttr | 64-bit signless integer elements attribute |
edge_padding_high | ::mlir::DenseIntElementsAttr | 64-bit signless integer elements attribute |
interior_padding | ::mlir::DenseIntElementsAttr | 64-bit signless integer elements attribute |
Operands:
| Operand | Description |
|---|---|
operand |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
padding_value |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
Results:
| Result | Description |
|---|---|
| «unnamed» | ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
mhlo.partition_id (mhlo::PartitionIdOp)
PartitionId operation
Syntax:
operation ::= `mhlo.partition_id` attr-dict `:` type(results)
Produces partition_id of the current process.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#partition_id
Example:
%result = mhlo.partition_id : tensor<ui32>
Interfaces: InferTypeOpInterface
Results:
| Result | Description |
|---|---|
| «unnamed» | ranked tensor of 32-bit unsigned integer values |
mhlo.popcnt (mhlo::PopulationCountOp)
PopulationCount operation
Syntax:
operation ::= `mhlo.popcnt` $operand attr-dict
`:` custom<SameOperandsAndResultType>(type($operand), type($result))
Performs element-wise count of the number of bits set in the operand
tensor and produces a result tensor.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#popcnt
Example:
%result = mhlo.popcnt %operand : tensor<4xi8>
Traits: AlwaysSpeculatableImplTrait, CompatibleOperandsAndResultType, Elementwise, SameOperandsAndResultShape
Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Operands:
| Operand | Description |
|---|---|
operand |
ranked tensor of 2/4/8/16/32/64-bit integer values |
Results:
| Result | Description |
|---|---|
result |
ranked tensor of 2/4/8/16/32/64-bit integer values |
mhlo.power (mhlo::PowOp)
Pow operation
Syntax:
operation ::= `mhlo.power` $lhs `,` $rhs attr-dict
`:` custom<SameOperandsAndResultType>(type($lhs), type($rhs), type($result))
Performs element-wise exponentiation of lhs tensor by rhs tensor and
produces a result tensor.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#power
Example:
%result = mhlo.power %lhs, %rhs : tensor<6xf32>
Traits: AlwaysSpeculatableImplTrait, CompatibleOperandsAndResultType, Elementwise, SameOperandsAndResultShape
Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Operands:
| Operand | Description |
|---|---|
lhs |
ranked tensor of 2/4/8/16/32/64-bit integer or 4/6/8/16/32/64-bit float or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
rhs |
ranked tensor of 2/4/8/16/32/64-bit integer or 4/6/8/16/32/64-bit float or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
Results:
| Result | Description |
|---|---|
result |
ranked tensor of 2/4/8/16/32/64-bit integer or 4/6/8/16/32/64-bit float or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
mhlo.ragged_dot (mhlo::RaggedDotOp)
Ragged matrix multiplication over a single ragged dimension
This operation takes three tensor args---lhs, rhs, and group_sizes---and a "ragged_dot_dimension_numbers" attribute. Like dot_general, the lhs and rhs are allowed arbitrary batch and contracting dimensions. Additionally, the lhs is required to have one ragged dimension, and the rhs may have at most one group dimension. The op has three modes, depending on the kind of the lhs ragged dimension.
In mode 1, the shape-signature is [b,m,k], [g,b,k,n], [b,g] -> [b,m,n].
Here the ragged dimension is an lhs non-contracting dimension (m). The
dimensions b and k represent batch and contracting dimensions
respectively. The rhs is required to have a group dimension (g).
In mode 2, the shape-signature is [b,m,k], [b,k,n], [b,g] -> [g,b,m,n].
Here the ragged dimension is an lhs/rhs contracting dimension (k).
In mode 3, the shape-signature is [b,m,k], [b,k,n], [g] -> [b,m,n]. Here
the ragged dimension is an lhs/rhs batch dimension (b).
Traits: AlwaysSpeculatableImplTrait
Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Attributes:
| Attribute | MLIR Type | Description |
|---|---|---|
ragged_dot_dimension_numbers | ::mlir::mhlo::RaggedDotDimensionNumbersAttr | Attribute that models the dimension information for ragged dot. |
precision_config | ::mlir::ArrayAttr | Precision Config attribute |
Operands:
| Operand | Description |
|---|---|
lhs |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
rhs |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
group_sizes |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
Results:
| Result | Description |
|---|---|
result |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized values |
mhlo.real (mhlo::RealOp)
Real operation
Syntax:
operation ::= `mhlo.real` $operand attr-dict
`:` custom<SameOperandsAndResultType>(type($operand), type($result))
Extracts the real part, element-wise, from the operand and produces a
result tensor.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#real
Example:
%result = mhlo.real %operand : (tensor<2xcomplex<f32>>) -> tensor<2xf32>
Traits: AlwaysSpeculatableImplTrait, Elementwise, SameOperandsAndResultShape
Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Operands:
| Operand | Description |
|---|---|
operand |
ranked tensor of 4/6/8/16/32/64-bit float or complex type with 32/64-bit float elements values |
Results:
| Result | Description |
|---|---|
result |
ranked tensor of 4/6/8/16/32/64-bit float values |
mhlo.real_dynamic_slice (mhlo::RealDynamicSliceOp)
RealDynamicSlice operation
Syntax:
operation ::= `mhlo.real_dynamic_slice` operands attr-dict `:` functional-type(operands, results)
This operation is a work in progress, so it is not yet included in the specification: https://github.com/openxla/stablehlo/issues/8
Informally, this operation does the same thing as SliceOp except
that start_indices, limit_indices and strides are specified dynamically:
https://github.com/openxla/stablehlo/blob/main/docs/spec.md#slice
Example:
%result = mhlo.real_dynamic_slice %operand,
%start_indices, %limit_indices, %strides
: (tensor<256x?xf32>, tensor<2xindex>, tensor<2xindex>, tensor<2xindex>) -> tensor<256x?xf32>
Traits: AlwaysSpeculatableImplTrait
Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Operands:
| Operand | Description |
|---|---|
operand |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
start_indices |
1D tensor of index or 2/4/8/16/32/64-bit integer values |
limit_indices |
1D tensor of index or 2/4/8/16/32/64-bit integer values |
strides |
1D tensor of index or 2/4/8/16/32/64-bit integer values |
Results:
| Result | Description |
|---|---|
result |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
mhlo.recv (mhlo::RecvOp)
Recv operation
Receives data from a channel with channel_id and produces results.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#recv
Example:
%results:2 = "mhlo.recv"(%token) {
// channel_id = 5 : i64,
// channel_type = #stablehlo<channel_type DEVICE_TO_DEVICE>,
channel_handle = #mhlo.channel_handle<handle = 5, type = 1>,
is_host_transfer = false,
source_target_pairs = dense<[[0, 1], [1, 2]]> : tensor<2x2xi64>
} : (!mhlo.token) -> (tensor<3x4xi32>, !mhlo.token)
Attributes:
| Attribute | MLIR Type | Description |
|---|---|---|
channel_handle | ::mlir::mhlo::ChannelHandleAttr | two 64-bit integers 'handle' and 'type' |
is_host_transfer | ::mlir::BoolAttr | bool attribute |
source_target_pairs | ::mlir::DenseIntElementsAttr | 64-bit signless integer elements attribute |
Operands:
| Operand | Description |
|---|---|
token |
token |
Results:
| Result | Description |
|---|---|
| «unnamed» | variadic of statically shaped tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized values or statically shaped tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values or token |
mhlo.reduce (mhlo::ReduceOp)
Reduce operation
Applies a reduction function body to inputs and init_values along the
dimensions and produces a result tensor.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#reduce
Example:
%result = "mhlo.reduce"(%input, %init_value) ({
^bb0(%arg0: tensor<i32>, %arg1: tensor<i32>):
%0 = "mhlo.add"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
"mhlo.return"(%0) : (tensor<i32>) -> ()
}) {
dimensions = dense<1> : tensor<1xi64>
} : (tensor<1x6xi32>, tensor<i32>) -> tensor<1xi32>
Traits: InferTensorType, RecursiveMemoryEffects, SameVariadicOperandSize, SingleBlockImplicitTerminator<ReturnOp>, SingleBlock
Interfaces: InferShapedTypeOpInterface, InferTypeOpInterface
Attributes:
| Attribute | MLIR Type | Description |
|---|---|---|
dimensions | ::mlir::DenseIntElementsAttr | 64-bit signless integer elements attribute |
Operands:
| Operand | Description |
|---|---|
inputs |
variadic of ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
init_values |
variadic of ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
Results:
| Result | Description |
|---|---|
| «unnamed» | variadic of ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
mhlo.reduce_precision (mhlo::ReducePrecisionOp)
ReducePrecision operation
Syntax:
operation ::= `mhlo.reduce_precision` $operand `,` `format` `=` custom<ExponentMantissa>($exponent_bits, $mantissa_bits)
attr-dict `:` custom<SameOperandsAndResultType>(type($operand), type($output))
Performs element-wise conversion of operand to another floating-point type
that uses exponent_bits and mantissa_bits and back to the original
floating-point type and produces an output tensor.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#reduce_precision
Example:
%output = mhlo.reduce_precision %operand, format = e5m2 : tensor<6xf32>
Traits: AlwaysSpeculatableImplTrait, CompatibleOperandsAndResultType, Elementwise
Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Attributes:
| Attribute | MLIR Type | Description |
|---|---|---|
exponent_bits | ::mlir::IntegerAttr | 32-bit signless integer attribute whose value is positive |
mantissa_bits | ::mlir::IntegerAttr | 32-bit signless integer attribute whose value is non-negative |
Operands:
| Operand | Description |
|---|---|
operand |
ranked tensor of 4/6/8/16/32/64-bit float values |
Results:
| Result | Description |
|---|---|
output |
ranked tensor of 4/6/8/16/32/64-bit float values |
mhlo.reduce_scatter (mhlo::ReduceScatterOp)
ReduceScatter operation
Within each process group in the process grid, performs reduction, using
computations, over the values of the operand tensor from each process,
splits the reduction result along scatter_dimension into parts, and
scatters the split parts between the processes to produce the result.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#reduce_scatter
Example:
%result = "mhlo.reduce_scatter"(%operand) ({
^bb0(%arg0: tensor<f32>, %arg1: tensor<f32>):
%0 = mhlo.add %arg0, %arg1 : tensor<f32>
mhlo.return %0 : tensor<f32>
}) {
scatter_dimension = 1 : i64,
replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>,
// channel_id = 0
channel_handle = #mhlo.channel_handle<handle = 0, type = 0>
// use_global_device_ids = false
} : (tensor<2x4xf32>) -> tensor<2x2xf32>
Attributes:
| Attribute | MLIR Type | Description |
|---|---|---|
scatter_dimension | ::mlir::IntegerAttr | 64-bit signless integer attribute whose value is non-negative |
replica_groups | ::mlir::DenseIntElementsAttr | 64-bit signless integer elements attribute |
channel_handle | ::mlir::mhlo::ChannelHandleAttr | two 64-bit integers 'handle' and 'type' |
use_global_device_ids | ::mlir::UnitAttr | unit attribute |
Operands:
| Operand | Description |
|---|---|
operand |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
Results:
| Result | Description |
|---|---|
| «unnamed» | ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
mhlo.reduce_window (mhlo::ReduceWindowOp)
ReduceWindow operation
Applies a reduction function body to windows of inputs and init_values
and produces results.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#reduce_window
Example:
%result = "mhlo.reduce_window"(%input, %init_value) ({
^bb0(%arg0: tensor<i32>, %arg1: tensor<i32>):
%0 = mhlo.add %arg0, %arg1 : tensor<i32>
mhlo.return %0 : tensor<i32>
}) {
window_dimensions = dense<[2, 1]> : tensor<2xi64>,
window_strides = dense<[4, 1]> : tensor<2xi64>,
base_dilations = dense<[2, 1]> : tensor<2xi64>,
window_dilations = dense<[3, 1]> : tensor<2xi64>,
padding = dense<[[2, 1], [0, 0]]> : tensor<2x2xi64>
} : (tensor<3x2xi32>, tensor<i32>) -> tensor<2x2xi32>
Traits: InferTensorType, RecursiveMemoryEffects, SameVariadicOperandSize, SingleBlockImplicitTerminator<ReturnOp>, SingleBlock
Interfaces: InferShapedTypeOpInterface, InferTypeOpInterface
Attributes:
| Attribute | MLIR Type | Description |
|---|---|---|
window_dimensions | ::mlir::DenseIntElementsAttr | 64-bit signless integer elements attribute |
window_strides | ::mlir::DenseIntElementsAttr | 64-bit signless integer elements attribute |
base_dilations | ::mlir::DenseIntElementsAttr | 64-bit signless integer elements attribute |
window_dilations | ::mlir::DenseIntElementsAttr | 64-bit signless integer elements attribute |
padding | ::mlir::DenseIntElementsAttr | 64-bit signless integer elements attribute |
Operands:
| Operand | Description |
|---|---|
inputs |
variadic of ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
init_values |
variadic of ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
Results:
| Result | Description |
|---|---|
| «unnamed» | variadic of ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
mhlo.remainder (mhlo::RemOp)
Rem operation
Syntax:
operation ::= `mhlo.remainder` $lhs `,` $rhs attr-dict
`:` custom<SameOperandsAndResultType>(type($lhs), type($rhs), type($result))
Performs element-wise remainder of dividend lhs and divisor rhs tensors
and produces a result tensor.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#remainder
Example:
%result = mhlo.remainder %lhs, %rhs : tensor<4xi64>
Traits: AlwaysSpeculatableImplTrait, CompatibleOperandsAndResultType, Elementwise, SameOperandsAndResultShape
Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Operands:
| Operand | Description |
|---|---|
lhs |
ranked tensor of 2/4/8/16/32/64-bit integer or 4/6/8/16/32/64-bit float or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
rhs |
ranked tensor of 2/4/8/16/32/64-bit integer or 4/6/8/16/32/64-bit float or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
Results:
| Result | Description |
|---|---|
result |
ranked tensor of 2/4/8/16/32/64-bit integer or 4/6/8/16/32/64-bit float or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
mhlo.replica_id (mhlo::ReplicaIdOp)
ReplicaId operation
Syntax:
operation ::= `mhlo.replica_id` attr-dict `:` type(results)
Produces replica_id of the current process.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#replica_id
Example:
%result = mhlo.replica_id : tensor<ui32>
Traits: AlwaysSpeculatableImplTrait
Interfaces: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Results:
| Result | Description |
|---|---|
| «unnamed» | ranked tensor of 32-bit unsigned integer values |
mhlo.reshape (mhlo::ReshapeOp)
Reshape operation
Syntax:
operation ::= `mhlo.reshape` operands attr-dict `:` functional-type(operands, results)
Performs reshape of operand tensor to a result tensor.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#reshape
Example:
%result = mhlo.reshape %operand : (tensor<2xf32>) -> tensor<1x2xf32>
Traits: AlwaysSpeculatableImplTrait, HLO_CompatibleOperandsAndResultElementType
Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Operands:
| Operand | Description |
|---|---|
operand |
tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
Results:
| Result | Description |
|---|---|
| «unnamed» | statically shaped or single bounded dimension tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
mhlo.return (mhlo::ReturnOp)
_This operation is a work in progress, so it is not yet included in the specification: https://github.com/openxla/stablehlo/issues/425
Informally, this operation serves as a terminator for regions defined by
the StableHLO ops. Non-StableHLO ops, e.g. `func.func`, have their own
terminators, e.g. `func.return`.
Example:
```mlir
%result = "mhlo.reduce"(%input, %init_value) ({
^bb0(%arg0: tensor<i32>, %arg1: tensor<i32>):
%0 = "mhlo.add"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
"mhlo.return"(%0) : (tensor<i32>) -> ()
}) {
dimensions = dense<1> : tensor<1xi64>
} : (tensor<1x6xi32>, tensor<i32>) -> tensor<1xi32>
```_
Syntax:
```
operation ::= mhlo.return $results attr-dict (: type($results)^)?
Traits: `AlwaysSpeculatableImplTrait`, `Terminator`
Interfaces: `ConditionallySpeculatable`, `NoMemoryEffect (MemoryEffectOpInterface)`
Effects: `MemoryEffects::Effect{}`
#### Operands:
| Operand | Description |
| :-----: | ----------- |
| `results` | variadic of ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values or token or nested tuple with any combination of ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized values or memref of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized values or ranked tensor of per-axis integer quantized values or token values or memref of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized values |
### `mhlo.reverse` (mhlo::ReverseOp)
_Reverse operation_
Reverses the order of elements in the `operand` along the specified
`dimensions` and produces a `result` tensor.
See:
<a href="https://github.com/openxla/stablehlo/blob/main/docs/spec.md#reverse">https://github.com/openxla/stablehlo/blob/main/docs/spec.md#reverse</a>
Example:
```mlir
%result = mhlo.reverse %operand, dims = [1] : tensor<3x2xi32>
Traits: AlwaysSpeculatableImplTrait, CompatibleOperandsAndResultType
Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Attributes:
| Attribute | MLIR Type | Description |
|---|---|---|
dimensions | ::mlir::DenseIntElementsAttr | 64-bit signless integer elements attribute |
Operands:
| Operand | Description |
|---|---|
operand |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
Results:
| Result | Description |
|---|---|
| «unnamed» | ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
mhlo.rng (mhlo::RngOp)
Rng operation
Generates random numbers using the rng_distribution algorithm and produces
a result tensor of a given shape shape.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#rng
Example:
%result = mhlo.rng %a, %b, %shape, distribution = NORMAL : (tensor<i32>, tensor<i32>, tensor<2xi64>) -> tensor<3x3xi32>
Traits: InferTensorType
Interfaces: InferShapedTypeOpInterface, InferTypeOpInterface
Attributes:
| Attribute | MLIR Type | Description |
|---|---|---|
rng_distribution | ::mlir::mhlo::RngDistributionAttr | XLA PRNG distribution to be used. |
Operands:
| Operand | Description |
|---|---|
a |
0D tensor of bool or 2/4/8/16/32/64-bit integer or 4/6/8/16/32/64-bit float values |
b |
0D tensor of bool or 2/4/8/16/32/64-bit integer or 4/6/8/16/32/64-bit float values |
shape |
1D tensor of index or 2/4/8/16/32/64-bit integer values |
Results:
| Result | Description |
|---|---|
result |
ranked tensor of bool or 2/4/8/16/32/64-bit integer or 4/6/8/16/32/64-bit float values |
mhlo.rng_bit_generator (mhlo::RngBitGeneratorOp)
RngBitGenerator operation
Returns an output filled with uniform random data and an updated output
state output_state given an initial state initial_state using the
pseudorandom number generator algorithm rng_algorithm.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#rng_bit_generator
Example:
%output_state, %output = mhlo.rng_bit_generator %initial_state, algorithm = THREE_FRY : (tensor<2xui64>) -> (tensor<2xui64>, tensor<2x2xui64>)
Traits: AlwaysSpeculatableImplTrait
Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Attributes:
| Attribute | MLIR Type | Description |
|---|---|---|
rng_algorithm | ::mlir::mhlo::RngAlgorithmAttr | XLA PRNG algorithm to be used. |
Operands:
| Operand | Description |
|---|---|
initial_state |
ranked tensor of 2/4/8/16/32/64-bit integer or 4/6/8/16/32/64-bit float values |
Results:
| Result | Description |
|---|---|
output_state |
ranked tensor of 2/4/8/16/32/64-bit integer or 4/6/8/16/32/64-bit float values |
output |
statically shaped tensor of 2/4/8/16/32/64-bit integer or 4/6/8/16/32/64-bit float values |
mhlo.round_nearest_afz (mhlo::RoundOp)
Round operation
Syntax:
operation ::= `mhlo.round_nearest_afz` $operand attr-dict
`:` custom<SameOperandsAndResultType>(type($operand), type($result))
Performs element-wise rounding towards the nearest integer, breaking ties
away from zero, on the operand tensor and produces a result tensor.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#round_nearest_afz
Example:
%result = mhlo.round_nearest_afz %operand : tensor<5xf32>
Traits: AlwaysSpeculatableImplTrait, CompatibleOperandsAndResultType, Elementwise, SameOperandsAndResultShape
Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Operands:
| Operand | Description |
|---|---|
operand |
ranked tensor of 4/6/8/16/32/64-bit float values |
Results:
| Result | Description |
|---|---|
result |
ranked tensor of 4/6/8/16/32/64-bit float values |
mhlo.round_nearest_even (mhlo::RoundNearestEvenOp)
RoundNearestEven operation
Syntax:
operation ::= `mhlo.round_nearest_even` $operand attr-dict
`:` custom<SameOperandsAndResultType>(type($operand), type($result))
Performs element-wise rounding towards the nearest integer, breaking ties
towards the even integer, on the operand tensor and produces a result
tensor.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#round_nearest_even
Example:
%result = mhlo.round_nearest_even %operand : tensor<5xf32>
Traits: AlwaysSpeculatableImplTrait, CompatibleOperandsAndResultType, Elementwise, SameOperandsAndResultShape
Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Operands:
| Operand | Description |
|---|---|
operand |
ranked tensor of 4/6/8/16/32/64-bit float values |
Results:
| Result | Description |
|---|---|
result |
ranked tensor of 4/6/8/16/32/64-bit float values |
mhlo.rsqrt (mhlo::RsqrtOp)
Rsqrt operation
Syntax:
operation ::= `mhlo.rsqrt` $operand attr-dict
`:` custom<SameOperandsAndResultType>(type($operand), type($result))
Performs element-wise reciprocal square root operation on operand tensor
and produces a result tensor, implementing the rSqrt operation from the
IEEE-754 specification.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#rsqrt
Example:
%result = mhlo.rsqrt %operand : tensor<2x2xf32>
Traits: AlwaysSpeculatableImplTrait, CompatibleOperandsAndResultType, Elementwise, SameOperandsAndResultShape
Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Attributes:
| Attribute | MLIR Type | Description |
|---|---|---|
result_accuracy | ::mlir::mhlo::ResultAccuracyAttr | The requested accuracy for unary ops. |
Operands:
| Operand | Description |
|---|---|
operand |
ranked tensor of 4/6/8/16/32/64-bit float or complex type with 32/64-bit float elements or per-tensor integer quantized values |
Results:
| Result | Description |
|---|---|
result |
ranked tensor of 4/6/8/16/32/64-bit float or complex type with 32/64-bit float elements or per-tensor integer quantized values |
mhlo.scatter (mhlo::ScatterOp)
Scatter operation
Produces results tensors which are equal to inputs tensors except that
several slices specified by scatter_indices are updated with the values
updates using update_computation.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#scatter
Example:
%result = "mhlo.scatter"(%input, %scatter_indices, %update) ({
^bb0(%arg0: tensor<i32>, %arg1: tensor<i32>):
%0 = mhlo.add %arg0, %arg1 : tensor<i32>
mhlo.return %0 : tensor<i32>
}) {
scatter_dimension_numbers = #mhlo.scatter<
update_window_dims = [3, 4],
inserted_window_dims = [1],
input_batching_dims = [0],
scatter_indices_batching_dims = [1],
scatter_dims_to_operand_dims = [2, 1],
index_vector_dim = 3>,
indices_are_sorted = false,
unique_indices = false
} : (tensor<2x3x4x2xi64>, tensor<2x2x3x2xi64>, tensor<2x2x3x2x2xi64>) -> tensor<2x3x4x2xi64>
Traits: RecursiveMemoryEffects, SameVariadicOperandSize
Interfaces: InferTypeOpInterface
Attributes:
| Attribute | MLIR Type | Description |
|---|---|---|
scatter_dimension_numbers | ::mlir::mhlo::ScatterDimensionNumbersAttr | Attribute that models the dimension information for scatter |
indices_are_sorted | ::mlir::BoolAttr | bool attribute |
unique_indices | ::mlir::BoolAttr | bool attribute |
Operands:
| Operand | Description |
|---|---|
inputs |
variadic of ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
scatter_indices |
ranked tensor of integer or index values |
updates |
variadic of ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
Results:
| Result | Description |
|---|---|
| «unnamed» | variadic of ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
mhlo.select (mhlo::SelectOp)
Select operation
Syntax:
operation ::= `mhlo.select` operands attr-dict `:`
custom<SelectOpType>(type($pred), type($on_true), type($on_false), type($result))
Produces a result tensor where each element is selected from on_true or
on_false tensor based on the value of the corresponding element of pred.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#select
Example:
%result = mhlo.select %pred, %on_true, %on_false : tensor<2x2xi1>, tensor<2x2xi32>
Traits: AlwaysSpeculatableImplTrait, HLO_BroadcastingElementwise, InferTensorType
Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Operands:
| Operand | Description |
|---|---|
pred |
ranked tensor of bool values |
on_true |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
on_false |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
Results:
| Result | Description |
|---|---|
result |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
mhlo.select_and_scatter (mhlo::SelectAndScatterOp)
SelectAndScatter operation
Scatters the values from the source tensor using scatter based on the
outcome of reduce_window of the input tensor using select and produces
a result tensor.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#select_and_scatter
Example:
%result = "mhlo.select_and_scatter"(%operand, %source, %init_value) ({
^bb0(%arg0: tensor<i32>, %arg1: tensor<i32>):
%0 = "mhlo.compare"(%arg0, %arg1) {
comparison_direction = #stablehlo<comparison_direction GE>
} : (tensor<i32>, tensor<i32>) -> tensor<i1>
"mhlo.return"(%0) : (tensor<i1>) -> ()
}, {
^bb0(%arg0: tensor<i32>, %arg1: tensor<i32>):
%0 = "mhlo.add"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
"mhlo.return"(%0) : (tensor<i32>) -> ()
}) {
window_dimensions = dense<[3, 1]> : tensor<2xi64>,
window_strides = dense<[2, 1]> : tensor<2xi64>,
padding = dense<[[0, 1], [0, 0]]> : tensor<2x2xi64>
} : (tensor<4x2xi32>, tensor<2x2xi32>, tensor<i32>) -> tensor<4x2xi32>
Traits: RecursiveMemoryEffects
Interfaces: InferTypeOpInterface
Attributes:
| Attribute | MLIR Type | Description |
|---|---|---|
window_dimensions | ::mlir::DenseIntElementsAttr | 64-bit signless integer elements attribute |
window_strides | ::mlir::DenseIntElementsAttr | 64-bit signless integer elements attribute |
padding | ::mlir::DenseIntElementsAttr | 64-bit signless integer elements attribute |
Operands:
| Operand | Description |
|---|---|
operand |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
source |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
init_value |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
Results:
| Result | Description |
|---|---|
| «unnamed» | ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
mhlo.send (mhlo::SendOp)
Send operation
Sends inputs to a channel channel_id and produces a result token.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#send
Example:
%result = "mhlo.send"(%operand, %token) {
// channel_id = 5 : i64,
// channel_type = #stablehlo<channel_type DEVICE_TO_DEVICE>,
channel_handle = #mhlo.channel_handle<handle = 5, type = 1>,
is_host_transfer = false,
source_target_pairs = dense<[[0, 1], [1, 2]]> : tensor<2x2xi64>
} : (tensor<3x4xi32>, !mhlo.token) -> !mhlo.token
Interfaces: InferTypeOpInterface
Attributes:
| Attribute | MLIR Type | Description |
|---|---|---|
channel_handle | ::mlir::mhlo::ChannelHandleAttr | two 64-bit integers 'handle' and 'type' |
is_host_transfer | ::mlir::BoolAttr | bool attribute |
source_target_pairs | ::mlir::DenseIntElementsAttr | 64-bit signless integer elements attribute |
Operands:
| Operand | Description |
|---|---|
inputs |
variadic of ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
token |
token |
Results:
| Result | Description |
|---|---|
| «unnamed» | token |
mhlo.set_dimension_size (mhlo::SetDimensionSizeOp)
SetDimensionSize operation
This operation is a work in progress, so it is not yet included in the specification: https://github.com/openxla/stablehlo/issues/8
Informally, this operation does the same thing as XLA's SetDimensionSize: https://www.tensorflow.org/xla/operation_semantics#setdimensionsize
Example:
%0 = mhlo.set_dimension_size %arg0, %arg1, dim = 1 : (tensor<4x2xf32>, tensor<i32>) -> tensor<4x2xf32>
Traits: AlwaysSpeculatableImplTrait, InferTensorType
Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Attributes:
| Attribute | MLIR Type | Description |
|---|---|---|
dimension | ::mlir::IntegerAttr | 64-bit signless integer attribute whose value is non-negative |
Operands:
| Operand | Description |
|---|---|
operand |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
size |
tensor of 32-bit signless integer values |
Results:
| Result | Description |
|---|---|
| «unnamed» | ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
mhlo.shift_left (mhlo::ShiftLeftOp)
ShiftLeft operation
Syntax:
operation ::= `mhlo.shift_left` $lhs `,` $rhs attr-dict
`:` custom<SameOperandsAndResultType>(type($lhs), type($rhs), type($result))
Performs element-wise left-shift operation on the lhs tensor by rhs
number of bits and produces a result tensor.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#shift_left
Example:
%result = mhlo.shift_left %lhs, %rhs : tensor<6xi8>
Traits: AlwaysSpeculatableImplTrait, CompatibleOperandsAndResultType, Elementwise, SameOperandsAndResultShape
Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Operands:
| Operand | Description |
|---|---|
lhs |
ranked tensor of 2/4/8/16/32/64-bit integer values |
rhs |
ranked tensor of 2/4/8/16/32/64-bit integer values |
Results:
| Result | Description |
|---|---|
result |
ranked tensor of 2/4/8/16/32/64-bit integer values |
mhlo.shift_right_arithmetic (mhlo::ShiftRightArithmeticOp)
ShiftRightArithmetic operation
Syntax:
operation ::= `mhlo.shift_right_arithmetic` $lhs `,` $rhs attr-dict
`:` custom<SameOperandsAndResultType>(type($lhs), type($rhs), type($result))
Performs element-wise arithmetic right-shift operation on the lhs tensor
by rhs number of bits and produces a result tensor.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#shift_right_arithmetic
Example:
%result = mhlo.shift_right_arithmetic %lhs, %rhs : tensor<6xi8>
Traits: AlwaysSpeculatableImplTrait, CompatibleOperandsAndResultType, Elementwise, SameOperandsAndResultShape
Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Operands:
| Operand | Description |
|---|---|
lhs |
ranked tensor of 2/4/8/16/32/64-bit integer values |
rhs |
ranked tensor of 2/4/8/16/32/64-bit integer values |
Results:
| Result | Description |
|---|---|
result |
ranked tensor of 2/4/8/16/32/64-bit integer values |
mhlo.shift_right_logical (mhlo::ShiftRightLogicalOp)
ShiftRightLogical operation
Syntax:
operation ::= `mhlo.shift_right_logical` $lhs `,` $rhs attr-dict
`:` custom<SameOperandsAndResultType>(type($lhs), type($rhs), type($result))
Performs element-wise logical right-shift operation on the lhs tensor by
rhs number of bits and produces a result tensor.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#shift_right_logical
Example:
%result = mhlo.shift_right_logical %lhs, %rhs : tensor<6xi8>
Traits: AlwaysSpeculatableImplTrait, CompatibleOperandsAndResultType, Elementwise, SameOperandsAndResultShape
Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Operands:
| Operand | Description |
|---|---|
lhs |
ranked tensor of 2/4/8/16/32/64-bit integer values |
rhs |
ranked tensor of 2/4/8/16/32/64-bit integer values |
Results:
| Result | Description |
|---|---|
result |
ranked tensor of 2/4/8/16/32/64-bit integer values |
mhlo.sign (mhlo::SignOp)
Sign operation
Syntax:
operation ::= `mhlo.sign` $operand attr-dict
`:` custom<SameOperandsAndResultType>(type($operand), type($result))
Returns the sign of the operand element-wise and produces a result
tensor.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#sign
Example:
%result = mhlo.sign %operand : tensor<7xf32>
Traits: AlwaysSpeculatableImplTrait, CompatibleOperandsAndResultType, Elementwise, SameOperandsAndResultShape
Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Operands:
| Operand | Description |
|---|---|
operand |
ranked tensor of 2/4/8/16/32/64-bit signless integer or 4/6/8/16/32/64-bit float or complex type with 32/64-bit float elements or per-tensor integer quantized values |
Results:
| Result | Description |
|---|---|
result |
ranked tensor of 2/4/8/16/32/64-bit signless integer or 4/6/8/16/32/64-bit float or complex type with 32/64-bit float elements or per-tensor integer quantized values |
mhlo.sine (mhlo::SineOp)
Sine operation
Syntax:
operation ::= `mhlo.sine` $operand attr-dict
`:` custom<SameOperandsAndResultType>(type($operand), type($result))
Performs element-wise sine operation on operand tensor and produces a
result tensor.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#sine
Example:
%result = mhlo.sine %operand : tensor<2xf32>
Traits: AlwaysSpeculatableImplTrait, CompatibleOperandsAndResultType, Elementwise, SameOperandsAndResultShape
Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Attributes:
| Attribute | MLIR Type | Description |
|---|---|---|
result_accuracy | ::mlir::mhlo::ResultAccuracyAttr | The requested accuracy for unary ops. |
Operands:
| Operand | Description |
|---|---|
operand |
ranked tensor of 4/6/8/16/32/64-bit float or complex type with 32/64-bit float elements or per-tensor integer quantized values |
Results:
| Result | Description |
|---|---|
result |
ranked tensor of 4/6/8/16/32/64-bit float or complex type with 32/64-bit float elements or per-tensor integer quantized values |
mhlo.sinh (mhlo::SinhOp)
Sinh operation
Syntax:
operation ::= `mhlo.sinh` $operand attr-dict
`:` custom<SameOperandsAndResultType>(type($operand), type($result))
Performs element-wise sinh operation on operand tensor and produces a
result tensor.
Example:
%result = mhlo.sinh %operand : tensor<2x2xf32>
Traits: CompatibleOperandsAndResultType, Elementwise, SameOperandsAndResultShape
Interfaces: InferShapedTypeOpInterface, InferTypeOpInterface
Operands:
| Operand | Description |
|---|---|
operand |
tensor of 4/6/8/16/32/64-bit float or complex type with 32/64-bit float elements values |
Results:
| Result | Description |
|---|---|
result |
tensor of 4/6/8/16/32/64-bit float or complex type with 32/64-bit float elements values |
mhlo.slice (mhlo::SliceOp)
Slice operation
Extracts a slice from the operand using statically-computed starting
indices and produces a result tensor.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#slice
Example:
%result = "mhlo.slice" (%operand) {
start_indices = dense<[1, 2]> : tensor<2xi64>,
limit_indices = dense<[3, 4]> : tensor<2xi64>,
strides = dense<1> : tensor<2xi64>
} : (tensor<3x4xi64>) -> tensor<2x2xi64>
Traits: AlwaysSpeculatableImplTrait, SameOperandsAndResultElementType
Interfaces: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Attributes:
| Attribute | MLIR Type | Description |
|---|---|---|
start_indices | ::mlir::DenseIntElementsAttr | 64-bit signless integer elements attribute |
limit_indices | ::mlir::DenseIntElementsAttr | 64-bit signless integer elements attribute |
strides | ::mlir::DenseIntElementsAttr | 64-bit signless integer elements attribute |
Operands:
| Operand | Description |
|---|---|
operand |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
Results:
| Result | Description |
|---|---|
| «unnamed» | ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
mhlo.sort (mhlo::SortOp)
Sort operation
Sorts a variadic number of tensors in inputs together, according to a
custom comparator, along the given dimension and produces a variadic
number of tensors as results.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#sort
Example:
%result0, %result1 = "mhlo.sort"(%input0, %input1) ({
^bb0(%arg0: tensor<i32>, %arg1: tensor<i32>, %arg2: tensor<i32>, %arg3: tensor<i32>):
%predicate = "mhlo.compare"(%arg0, %arg1) {
comparison_direction = #stablehlo<comparison_direction GT>
} : (tensor<i32>, tensor<i32>) -> tensor<i1>
"mhlo.return"(%predicate) : (tensor<i1>) -> ()
}) {
dimension = 0 : i64,
is_stable = true
} : (tensor<2x3xi32>, tensor<2x3xi32>) -> (tensor<2x3xi32>, tensor<2x3xi32>)
Traits: InferTensorType, RecursiveMemoryEffects, SameOperandsAndResultShape
Interfaces: InferShapedTypeOpInterface, InferTypeOpInterface
Attributes:
| Attribute | MLIR Type | Description |
|---|---|---|
dimension | ::mlir::IntegerAttr | 64-bit signless integer attribute |
is_stable | ::mlir::BoolAttr | bool attribute |
Operands:
| Operand | Description |
|---|---|
inputs |
variadic of ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
Results:
| Result | Description |
|---|---|
| «unnamed» | variadic of ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
mhlo.sqrt (mhlo::SqrtOp)
Sqrt operation
Syntax:
operation ::= `mhlo.sqrt` $operand attr-dict
`:` custom<SameOperandsAndResultType>(type($operand), type($result))
Performs element-wise square root operation on operand tensor and produces
a result tensor.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#sqrt
Example:
%result = mhlo.sqrt %operand : tensor<2x2xf32>
Traits: AlwaysSpeculatableImplTrait, CompatibleOperandsAndResultType, Elementwise, SameOperandsAndResultShape
Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Attributes:
| Attribute | MLIR Type | Description |
|---|---|---|
result_accuracy | ::mlir::mhlo::ResultAccuracyAttr | The requested accuracy for unary ops. |
Operands:
| Operand | Description |
|---|---|
operand |
ranked tensor of 4/6/8/16/32/64-bit float or complex type with 32/64-bit float elements or per-tensor integer quantized values |
Results:
| Result | Description |
|---|---|
result |
ranked tensor of 4/6/8/16/32/64-bit float or complex type with 32/64-bit float elements or per-tensor integer quantized values |
mhlo.stochastic_convert (mhlo::StochasticConvertOp)
StochasticConvert operation
This operation is a work in progress, so it is not yet included in the specification: https://github.com/openxla/stablehlo/issues/295
Informally, this operation performs element-wise conversion of values from a bigger type to a smaller one with stochastic rounding using the random number passed in.
Traits: AlwaysSpeculatableImplTrait, Elementwise
Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Operands:
| Operand | Description |
|---|---|
operand |
ranked tensor of 4/6/8/16/32/64-bit float values |
random |
ranked tensor of 2/4/8/16/32/64-bit unsigned integer values |
Results:
| Result | Description |
|---|---|
result |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
mhlo.subtract (mhlo::SubtractOp)
Subtract operation
Syntax:
operation ::= `mhlo.subtract` $lhs `,` $rhs attr-dict
`:` custom<SameOperandsAndResultType>(type($lhs), type($rhs), type($result))
Performs element-wise subtraction of two tensors lhs and rhs and
produces a result tensor.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#subtract
Example:
%result = mhlo.subtract %lhs, %rhs : tensor<2xi32>
Traits: AlwaysSpeculatableImplTrait, CompatibleOperandsAndResultType, Elementwise, SameOperandsAndResultShape
Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Operands:
| Operand | Description |
|---|---|
lhs |
ranked tensor of 2/4/8/16/32/64-bit integer or 4/6/8/16/32/64-bit float or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
rhs |
ranked tensor of 2/4/8/16/32/64-bit integer or 4/6/8/16/32/64-bit float or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
Results:
| Result | Description |
|---|---|
result |
ranked tensor of 2/4/8/16/32/64-bit integer or 4/6/8/16/32/64-bit float or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
mhlo.tan (mhlo::TanOp)
Tan operation
Syntax:
operation ::= `mhlo.tan` $operand attr-dict
`:` custom<SameOperandsAndResultType>(type($operand), type($result))
This operation is a work in progress, so it is not yet included in the specification: https://github.com/openxla/stablehlo/issues/954
Informally, this operation returns Tan(operand) element-wise.
Example:
%0 = mhlo.tan %arg0 : tensor<2xf32>
Traits: AlwaysSpeculatableImplTrait, CompatibleOperandsAndResultType, Elementwise, SameOperandsAndResultShape
Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Attributes:
| Attribute | MLIR Type | Description |
|---|---|---|
result_accuracy | ::mlir::mhlo::ResultAccuracyAttr | The requested accuracy for unary ops. |
Operands:
| Operand | Description |
|---|---|
operand |
ranked tensor of 4/6/8/16/32/64-bit float or complex type with 32/64-bit float elements or per-tensor integer quantized values |
Results:
| Result | Description |
|---|---|
result |
ranked tensor of 4/6/8/16/32/64-bit float or complex type with 32/64-bit float elements or per-tensor integer quantized values |
mhlo.tanh (mhlo::TanhOp)
Tanh operation
Syntax:
operation ::= `mhlo.tanh` $operand attr-dict
`:` custom<SameOperandsAndResultType>(type($operand), type($result))
Performs element-wise hyperbolic tangent operation on operand tensor and
produces a result tensor.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#tanh
Example:
%result = mhlo.tanh %operand : tensor<2xf32>
Traits: AlwaysSpeculatableImplTrait, CompatibleOperandsAndResultType, Elementwise, SameOperandsAndResultShape
Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Attributes:
| Attribute | MLIR Type | Description |
|---|---|---|
result_accuracy | ::mlir::mhlo::ResultAccuracyAttr | The requested accuracy for unary ops. |
Operands:
| Operand | Description |
|---|---|
operand |
ranked tensor of 4/6/8/16/32/64-bit float or complex type with 32/64-bit float elements or per-tensor integer quantized values |
Results:
| Result | Description |
|---|---|
result |
ranked tensor of 4/6/8/16/32/64-bit float or complex type with 32/64-bit float elements or per-tensor integer quantized values |
mhlo.topk (mhlo::TopKOp)
TopK operation
Syntax:
operation ::= `mhlo.topk` `(`$operand `,` `k` `=` $k (`,` `largest` `=` $largest^)? `)` attr-dict `:`
type($operand) `->` `(`type($values)`,` type($indices)`)`
Returns top k values and their indices, along the last
dimension of the operand if largest=true or the bottom k values if
largest=false.
See: https://www.tensorflow.org/xla/operation_semantics#top-k
Example:
%values, %indices = mhlo.topk(%operand, k=5, largest=true)
: tensor<100xf32> -> (tensor<5xf32>, tensor<5xi32>)
Traits: InferTensorType, RecursiveMemoryEffects
Interfaces: InferShapedTypeOpInterface, InferTypeOpInterface
Attributes:
| Attribute | MLIR Type | Description |
|---|---|---|
k | ::mlir::IntegerAttr | 64-bit signless integer attribute |
largest | ::mlir::BoolAttr | bool attribute |
Operands:
| Operand | Description |
|---|---|
operand |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
Results:
| Result | Description |
|---|---|
values |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
indices |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
mhlo.torch_index_select (mhlo::TorchIndexSelectOp)
TorchIndexSelect operation
This operation is on its way out of StableHLO, so it is not included in the specification: https://github.com/openxla/stablehlo/issues/3
Informally, this operation does the same thing as PyTorch's index_select, augmented with support for batch dimensions: https://pytorch.org/docs/stable/generated/torch.index_select.html
The batch_dims attribute specifies the number of major batch dimensions
(0 or more) that act like a multidimensional loop over both the operand and
the index.
Example:
%result = "mhlo.torch_index_select"(%operand, %index) {
dim = 2 : i64,
batch_dims = 1 : i64
} : (tensor<8x128x3072x64xf32>, tensor<8x16x1024xi32>) -> tensor<8x128x16x1024x64xf32>
Traits: AlwaysSpeculatableImplTrait
Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Attributes:
| Attribute | MLIR Type | Description |
|---|---|---|
dim | ::mlir::IntegerAttr | 64-bit signless integer attribute |
batch_dims | ::mlir::IntegerAttr | 64-bit signless integer attribute |
Operands:
| Operand | Description |
|---|---|
operand |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
index |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
Results:
| Result | Description |
|---|---|
| «unnamed» | ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
mhlo.trace (mhlo::TraceOp)
Trace operation
Syntax:
operation ::= `mhlo.trace` $operand `,` $tag attr-dict `:` type($operand)
This operation is on its way out of StableHLO, so it is not included in the specification: https://github.com/openxla/stablehlo/issues/604
It is not used by JAX, PyTorch or TensorFlow, so it looks like we should've classified it as "Private to XLA" and not included it in StableHLO in the first place. With that in mind, its semantics will not be documented here.
Example:
mhlo.trace %arg0, "In test code." : tensor<5x1x5xi32>
Attributes:
| Attribute | MLIR Type | Description |
|---|---|---|
tag | ::mlir::StringAttr | string attribute |
Operands:
| Operand | Description |
|---|---|
operand |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
mhlo.transpose (mhlo::TransposeOp)
Transpose operation
Permutes the dimensions of operand tensor using permutation and produces
a result tensor.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#transpose
Example:
%0 = mhlo.transpose %arg0, dims = [2, 1, 0] : (tensor<1x2x3xi32>) -> tensor<3x2x1xi32>
Traits: AlwaysSpeculatableImplTrait, HLO_CompatibleOperandsAndResultElementType
Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Attributes:
| Attribute | MLIR Type | Description |
|---|---|---|
permutation | ::mlir::DenseIntElementsAttr | 64-bit signless integer elements attribute |
Operands:
| Operand | Description |
|---|---|
operand |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
Results:
| Result | Description |
|---|---|
| «unnamed» | ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
mhlo.triangular_solve (mhlo::TriangularSolveOp)
TriangularSolve operation
Solves batches of systems of linear equations with lower or upper triangular coefficient matrices.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#triangular_solve
Example:
%result = "mhlo.triangular_solve"(%a, %b) {
left_side = true,
lower = true,
unit_diagonal = false,
transpose_a = #stablehlo<transpose NO_TRANSPOSE>
} : (tensor<3x3xf32>, tensor<3x3xf32>) -> tensor<3x3xf32>
Traits: AlwaysSpeculatableImplTrait, InferTensorType, SameOperandsAndResultElementType
Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Attributes:
| Attribute | MLIR Type | Description |
|---|---|---|
left_side | ::mlir::BoolAttr | bool attribute |
lower | ::mlir::BoolAttr | bool attribute |
unit_diagonal | ::mlir::BoolAttr | bool attribute |
transpose_a | ::mlir::mhlo::TransposeAttr | Transpose options |
Operands:
| Operand | Description |
|---|---|
a |
ranked tensor of 4/6/8/16/32/64-bit float or complex type with 32/64-bit float elements values |
b |
ranked tensor of 4/6/8/16/32/64-bit float or complex type with 32/64-bit float elements values |
Results:
| Result | Description |
|---|---|
| «unnamed» | ranked tensor of 4/6/8/16/32/64-bit float or complex type with 32/64-bit float elements values |
mhlo.tuple (mhlo::TupleOp)
Tuple operation
Syntax:
operation ::= `mhlo.tuple` $val attr-dict `:` custom<TupleOpType>(type($val), type($result))
Produces a result tuple from values val.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#tuple
Example:
%result = mhlo.tuple %val0, %val1 : tuple<tensor<2xf32>, tuple<tensor<i32>>>
Traits: AlwaysSpeculatableImplTrait
Interfaces: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Operands:
| Operand | Description |
|---|---|
val |
variadic of ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values or token or nested tuple with any combination of ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized values or memref of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized values or ranked tensor of per-axis integer quantized values or token values or memref of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized values |
Results:
| Result | Description |
|---|---|
result |
nested tuple with any combination of ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized values or memref of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized values or ranked tensor of per-axis integer quantized values or token values |
mhlo.uniform_dequantize (mhlo::UniformDequantizeOp)
UniformDequantize operation
Syntax:
operation ::= `mhlo.uniform_dequantize` $operand attr-dict
`:` custom<SameOperandsAndResultType>(type($operand), type($result))
Performs element-wise conversion of quantized tensor operand to a
floating-point tensor result according to the quantization parameters
defined by the operand type.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#uniform_dequantize
Example:
%result = mhlo.uniform_dequantize %operand : (tensor<16x16x!quant.uniform<i8:f32, 34.0:16>>) -> tensor<16x16xf32>
Traits: AlwaysSpeculatableImplTrait, Elementwise, InferTensorType, SameOperandsAndResultShape
Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Operands:
| Operand | Description |
|---|---|
operand |
ranked tensor of per-tensor integer quantized or per-axis integer quantized values |
Results:
| Result | Description |
|---|---|
result |
ranked tensor of 4/6/8/16/32/64-bit float values |
mhlo.uniform_quantize (mhlo::UniformQuantizeOp)
UniformQuantize operation
Syntax:
operation ::= `mhlo.uniform_quantize` $operand attr-dict
`:` custom<SameOperandsAndResultType>(type($operand), type($result))
Performs element-wise conversion of floating-point tensor or quantized
tensor operand to a quantized tensor result according to the
quantization parameters defined by the result type.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#uniform_quantize
Example:
%result = mhlo.uniform_quantize %operand : (tensor<16x16xf32>) -> tensor<16x16x!quant.uniform<ui8:f32, 34.0:16>>
Traits: AlwaysSpeculatableImplTrait, Elementwise, SameOperandsAndResultShape
Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Operands:
| Operand | Description |
|---|---|
operand |
ranked tensor of 4/6/8/16/32/64-bit float or 2/4/8/16/32-bit uniform quantized signed integer or 2/4/8/16/32-bit uniform quantized per axis signed integer or 2/4/8/16/32-bit uniform quantized unsigned integer or 2/4/8/16/32-bit uniform quantized per axis unsigned integer values |
Results:
| Result | Description |
|---|---|
result |
ranked tensor of per-tensor integer quantized or per-axis integer quantized values |
mhlo.while (mhlo::WhileOp)
While operation
Produces the output from executing body function 0 or more times while the
cond function outputs true.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#while
Example:
%results0, %results1 = "mhlo.while"(%operand0, %operand1) ({
^bb0(%arg0: tensor<i32>, %arg1: tensor<i32>):
%0 = "mhlo.compare"(%arg0, %arg1) {
comparison_direction = #stablehlo<comparison_direction LT>
} : (tensor<i32>, tensor<i32>) -> tensor<i1>
"mhlo.return"(%0) : (tensor<i1>) -> ()
}, {
^bb0(%arg0: tensor<i32>, %arg1: tensor<i32>):
%0 = "mhlo.add"(%arg0, %constant0) : (tensor<i32>, tensor<i32>) -> tensor<i32>
"mhlo.return"(%0, %arg1) : (tensor<i32>, tensor<i32>) -> ()
}) : (tensor<i32>, tensor<i32>) -> (tensor<i32>, tensor<i32>)
Traits: RecursiveMemoryEffects, SingleBlockImplicitTerminator<ReturnOp>, SingleBlock
Interfaces: InferTypeOpInterface, OpAsmOpInterface
Operands:
| Operand | Description |
|---|---|
operand |
variadic of ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized values or ranked tensor of per-axis integer quantized values or token or memref of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized values |
Results:
| Result | Description |
|---|---|
| «unnamed» | variadic of ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized values or ranked tensor of per-axis integer quantized values or token or memref of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized values |
mhlo.xla.rng_get_and_update_state (mhlo::XlaRngGetAndUpdateStateOp)
XlaRngGetAndUpdateState operation
Syntax:
operation ::= `mhlo.xla.rng_get_and_update_state` attr-dict
This operation is private to the XLA compiler, so it is does not yet have a specification.
Informally, this operation represents the change of the global random number generator state for rng instructions. The global state is incremented by delta and the old state is returned.
The output is currently defined for a single output type. If this changes in the future to support multiple types, lowering to use of a global memref must ensure that a single memref is still used and updated appropriately.
Interfaces: InferTypeOpInterface
Attributes:
| Attribute | MLIR Type | Description |
|---|---|---|
delta | ::mlir::IntegerAttr | 64-bit signless integer attribute |
Results:
| Result | Description |
|---|---|
| «unnamed» | statically shaped tensor of 64-bit unsigned integer values |
mhlo.xor (mhlo::XorOp)
Xor operation
Syntax:
operation ::= `mhlo.xor` $lhs `,` $rhs attr-dict
`:` custom<SameOperandsAndResultType>(type($lhs), type($rhs), type($result))
Performs element-wise XOR of two tensors lhs and rhs and produces a
result tensor.
See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#xor
Example:
%result = mhlo.xor %lhs, %rhs : tensor<2xi32>
Traits: AlwaysSpeculatableImplTrait, Commutative, CompatibleOperandsAndResultType, Elementwise, SameOperandsAndResultShape
Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)
Effects: MemoryEffects::Effect{}
Operands:
| Operand | Description |
|---|---|
lhs |
ranked tensor of bool or 2/4/8/16/32/64-bit integer values |
rhs |
ranked tensor of bool or 2/4/8/16/32/64-bit integer values |
Results:
| Result | Description |
|---|---|
result |
ranked tensor of 4/6/8/16/32/64-bit float or bool or 2/4/8/16/32/64-bit integer or complex type with 32/64-bit float elements or per-tensor integer quantized or per-axis integer quantized values |
Attributes
ArgResultAliasAttr
Attribute that models the alias relationship of entry function argument
This attribute captures the alias relationship of an MHLO main function
argument to one of the results, denoted by resultIndex. The
argTupleIndices and resultTupleIndices are used to index into nested
tuples in operand and result respectively. If isMustAlias is true then the
operand-result pair must alias.
This is meant to be used as an attribute on a function argument in MHLO.
For example, in the following code it expresses that %arg1 may alias 0-th
result.
func @main(%arg0: tensor<2xf32>, %arg1: tensor<3xf32> {mhlo.result_alias =
mhlo.result_alias<result_index = [2], ...>}
) -> tensor<2xf32>, tensor<3xf32> {
// function body ...
}
Parameters:
| Parameter | C++ type | Description |
|---|---|---|
| argTupleIndices | ::llvm::ArrayRef<int64_t> |
Dimension |
| resultIndex | int64_t |
|
| resultTupleIndices | ::llvm::ArrayRef<int64_t> |
Dimension |
| isMustAlias | bool |
ChannelHandleAttr
Two 64-bit integers 'handle' and 'type'
Syntax:
#mhlo.channel_handle<
int64_t, # handle
int64_t # type
>
Parameters:
| Parameter | C++ type | Description |
|---|---|---|
| handle | int64_t |
|
| type | int64_t |
ComparisonDirectionAttr
Which comparison operation to perform.
Syntax:
#mhlo.comparison_direction<
::mlir::mhlo::ComparisonDirection # value
>
Parameters:
| Parameter | C++ type | Description |
|---|---|---|
| value | ::mlir::mhlo::ComparisonDirection |
an enum of type ComparisonDirection |
ComparisonTypeAttr
Which comparison type to use.
Syntax:
#mhlo.comparison_type<
::mlir::mhlo::ComparisonType # value
>
Parameters:
| Parameter | C++ type | Description |
|---|---|---|
| value | ::mlir::mhlo::ComparisonType |
an enum of type ComparisonType |
ConvDimensionNumbersAttr
Structure of dimension information for conv op
Parameters:
| Parameter | C++ type | Description |
|---|---|---|
| inputBatchDimension | int64_t |
|
| inputFeatureDimension | int64_t |
|
| inputSpatialDimensions | ::llvm::ArrayRef<int64_t> |
Dimension |
| kernelInputFeatureDimension | int64_t |
|
| kernelOutputFeatureDimension | int64_t |
|
| kernelSpatialDimensions | ::llvm::ArrayRef<int64_t> |
Dimension |
| outputBatchDimension | int64_t |
|
| outputFeatureDimension | int64_t |
|
| outputSpatialDimensions | ::llvm::ArrayRef<int64_t> |
Dimension |
CrossProgramPrefetchAttr
Argument that is prefetched from another program
Syntax:
#mhlo.cross_program_prefetch<
int64_t, # parameter
::llvm::ArrayRef<int64_t>, # indices
std::optional<int64_t> # offset
>
This attribute captures an argument that is prefetched from another program.
For a given CrossProgramPrefetchAttr, parameter tells us which argument
of the main function of the module is prefetched, and indices is a shape
index telling us what subshape of that argument is prefetched.
A shape has a subshape iff it is a tuple. In that case, the subshape of
the tuple by indices is the shape achieved after indexing by each
element of indices in turn. For example, the [1,0] subshape of
tuple<tuple<token, token>, tuple<tensor<i32>, token>> is tensor<i32>.
An empty value for indices means the whole shape is prefetched.
For example,
module attributes { mhlo.cross_program_prefetch = [ #mhlo.cross_program_prefetch< parameter = 0, indices = [0]> ]} {
func.func @copy(%arg0 : tuple<tensor<2x3xi32>, tensor<i32>>) -> tuple<tensor<2x3xi32>, tensor<i32>> {
%0 = "mhlo.copy"(%arg0) {is_cross_program_prefetch}
return %0 : tuple<tensor<2x3xi32>, tensor<i32>>
}
func.func @main(%arg0 : tuple<tensor<2x3xi32>, tensor<i32>>) -> tuple<tensor<2x3xi32>, tensor<i32>> {
%1 = "mhlo.async_start"(%arg0) {called_computation=@copy}
%2 = "mhlo.async_done"(%1) {called_computation=@copy}
return %2 : tuple<tensor<2x3xi32>, tensor<i32>>
}
}
The parameter = 0 tells us that the async copy of the 0th parameter is
a cross_program_prefetch, while the index of [0] tells us that the
0th element of the tuple is prefetched while the other element of the
tuple is not.
Parameters:
| Parameter | C++ type | Description |
|---|---|---|
| parameter | int64_t |
|
| indices | ::llvm::ArrayRef<int64_t> |
Dimension |
| offset | std::optional<int64_t> |
CustomCallScheduleAttr
Specifies the desired schedule for the custom-call.
Syntax:
#mhlo.custom_call_schedule<
::mlir::mhlo::CustomCallSchedule # value
>
Parameters:
| Parameter | C++ type | Description |
|---|---|---|
| value | ::mlir::mhlo::CustomCallSchedule |
an enum of type CustomCallSchedule |
DequantizeModeAttr
_Dequantization mode. Only MINCOMBINED is supported.
Syntax:
#mhlo.dequantize_mode<
::mlir::mhlo::DequantizeMode # value
>
Parameters:
| Parameter | C++ type | Description |
|---|---|---|
| value | ::mlir::mhlo::DequantizeMode |
an enum of type DequantizeMode |
DomainKindAttr
Kind of domain metatdata attached to an HLO domain.
Syntax:
#mhlo.kind<
::mlir::mhlo::DomainKind # value
>
Parameters:
| Parameter | C++ type | Description |
|---|---|---|
| value | ::mlir::mhlo::DomainKind |
an enum of type DomainKind |
DotAlgorithmAttr
Attribute that models the algorithm constraints to use for computing dot.
Syntax:
#mhlo.dot_algorithm<
Type, # lhsPrecisionType
Type, # rhsPrecisionType
Type, # accumulationType
int64_t, # lhsComponentCount
int64_t, # rhsComponentCount
int64_t, # numPrimitiveOperations
bool # allowImpreciseAccumulation
>
Parameters:
| Parameter | C++ type | Description |
|---|---|---|
| lhsPrecisionType | Type |
|
| rhsPrecisionType | Type |
|
| accumulationType | Type |
|
| lhsComponentCount | int64_t |
|
| rhsComponentCount | int64_t |
|
| numPrimitiveOperations | int64_t |
|
| allowImpreciseAccumulation | bool |
DotDimensionNumbersAttr
Attribute that models the dimension information for dot.
Parameters:
| Parameter | C++ type | Description |
|---|---|---|
| lhsBatchingDimensions | ::llvm::ArrayRef<int64_t> |
Dimension |
| rhsBatchingDimensions | ::llvm::ArrayRef<int64_t> |
Dimension |
| lhsContractingDimensions | ::llvm::ArrayRef<int64_t> |
Dimension |
| rhsContractingDimensions | ::llvm::ArrayRef<int64_t> |
Dimension |
FftTypeAttr
XLA fast fourier transform type.
Syntax:
#mhlo.fft_type<
::mlir::mhlo::FftType # value
>
Parameters:
| Parameter | C++ type | Description |
|---|---|---|
| value | ::mlir::mhlo::FftType |
an enum of type FftType |
FusionKindAttr
Fusion kind
Syntax:
#mhlo.fusion_kind<
::mlir::mhlo::FusionKind # value
>
Parameters:
| Parameter | C++ type | Description |
|---|---|---|
| value | ::mlir::mhlo::FusionKind |
an enum of type FusionKind |
GatherDimensionNumbersAttr
Attribute that models the dimension information for gather
Parameters:
| Parameter | C++ type | Description |
|---|---|---|
| offsetDims | ::llvm::ArrayRef<int64_t> |
Dimension |
| collapsedSliceDims | ::llvm::ArrayRef<int64_t> |
Dimension |
| operandBatchingDims | ::llvm::ArrayRef<int64_t> |
Dimension |
| startIndicesBatchingDims | ::llvm::ArrayRef<int64_t> |
Dimension |
| startIndexMap | ::llvm::ArrayRef<int64_t> |
Dimension |
| indexVectorDim | int64_t |
OutputOperandAliasAttr
Attribute that models the alias relationship of output and operand of a CustomCall op
Syntax:
#mhlo.output_operand_alias<
::llvm::ArrayRef<int64_t>, # outputTupleIndices
int64_t, # operandIndex
::llvm::ArrayRef<int64_t> # operandTupleIndices
>
This attribute captures the alias relationship of the output to one of the
operands for a CustomCall op, denoted by operand_index. The
output_tuple_indices and operand_tuple_indices are used to index into
output and operand types. These indices lists are empty if the corresponding
types are not tuple types, and can be arbitrarily long in case of
arbitrarily nested tuple types.
See https://www.tensorflow.org/xla/aliasing
Example when used as array with in mhlo.custom-call:
%0 = "mhlo.custom_call"(%arg0, %arg1) {
// other attributes
output_operand_alias = [
#mhlo.output_operand_alias<output_tuple_indices = [0],
operand_index = 0,
operand_tuple_indices = [1]>
]
} : (tuple<tensor<1x1xf32>, tensor<2x3xf32>>, tensor<5x5xf32>) -> tuple<tensor<2x3xf32>>
The output and the 0th operand are both tuples. The aliasing shows the
relationship between the 0th element in output tuple with the 1st element in
the 0th operand. And both of them are of the same type: tensor<2x3xf32>.
Parameters:
| Parameter | C++ type | Description |
|---|---|---|
| outputTupleIndices | ::llvm::ArrayRef<int64_t> |
Dimension |
| operandIndex | int64_t |
|
| operandTupleIndices | ::llvm::ArrayRef<int64_t> |
Dimension |
PrecisionAttr
XLA precision for an operand. Has backend specific meaning.
Syntax:
#mhlo.precision<
::mlir::mhlo::Precision # value
>
Parameters:
| Parameter | C++ type | Description |
|---|---|---|
| value | ::mlir::mhlo::Precision |
an enum of type Precision |
RaggedDotDimensionNumbersAttr
Attribute that models the dimension information for ragged dot.
Parameters:
| Parameter | C++ type | Description |
|---|---|---|
| dotDimensionNumbers | ::mlir::mhlo::DotDimensionNumbersAttr |
Attribute that models the dimension information for dot. |
| lhsRaggedDimensions | ::llvm::ArrayRef<int64_t> |
Dimension |
| rhsGroupDimensions | ::llvm::ArrayRef<int64_t> |
Dimension |
ResultAccuracyAttr
The requested accuracy for unary ops.
Parameters:
| Parameter | C++ type | Description |
|---|---|---|
| atol | APFloat |
|
| rtol | APFloat |
|
| ulps | int64_t |
|
| mode | ::mlir::mhlo::ResultAccuracyModeAttr |
XLA result accuracy mode. |
ResultAccuracyModeAttr
XLA result accuracy mode.
Syntax:
#mhlo.result_accuracy_mode<
::mlir::mhlo::ResultAccuracyMode # value
>
Parameters:
| Parameter | C++ type | Description |
|---|---|---|
| value | ::mlir::mhlo::ResultAccuracyMode |
an enum of type ResultAccuracyMode |
RngAlgorithmAttr
XLA PRNG algorithm to be used.
Syntax:
#mhlo.rng_algorithm<
::mlir::mhlo::RngAlgorithm # value
>
Parameters:
| Parameter | C++ type | Description |
|---|---|---|
| value | ::mlir::mhlo::RngAlgorithm |
an enum of type RngAlgorithm |
RngDistributionAttr
XLA PRNG distribution to be used.
Syntax:
#mhlo.rng_distribution<
::mlir::mhlo::RngDistribution # value
>
Parameters:
| Parameter | C++ type | Description |
|---|---|---|
| value | ::mlir::mhlo::RngDistribution |
an enum of type RngDistribution |
ScatterDimensionNumbersAttr
Attribute that models the dimension information for scatter
Parameters:
| Parameter | C++ type | Description |
|---|---|---|
| updateWindowDims | ::llvm::ArrayRef<int64_t> |
Dimension |
| insertedWindowDims | ::llvm::ArrayRef<int64_t> |
Dimension |
| inputBatchingDims | ::llvm::ArrayRef<int64_t> |
Dimension |
| scatterIndicesBatchingDims | ::llvm::ArrayRef<int64_t> |
Dimension |
| scatterDimsToOperandDims | ::llvm::ArrayRef<int64_t> |
Dimension |
| indexVectorDim | int64_t |
TransposeAttr
Transpose options
Syntax:
#mhlo.transpose<
::mlir::mhlo::Transpose # value
>
Parameters:
| Parameter | C++ type | Description |
|---|---|---|
| value | ::mlir::mhlo::Transpose |
an enum of type Transpose |
TypeExtensionsAttr
Attribute that extends tensor type with MHLO type properties.
Syntax:
#mhlo.type_extensions<
::llvm::ArrayRef<int64_t> # bounds
>
This attribute is used to extend MLIR tensor type with MHLO tensor specific
properties. These properties aren't modeled in the MLIR type. This
attribute is set in the encoding field of the tensor type.
See HLO_BoundedAttrInterface for documentation for bounds.
Parameters:
| Parameter | C++ type | Description |
|---|---|---|
| bounds | ::llvm::ArrayRef<int64_t> |
Types
AsyncBundleType
Opaque collection of other types
Syntax:
!mhlo.async_bundle<
::llvm::ArrayRef<Type> # types
>
Parameters:
| Parameter | C++ type | Description |
|---|---|---|
| types | ::llvm::ArrayRef<Type> |
Enums
ComparisonDirection
Which comparison operation to perform.
Cases:
| Symbol | Value | String |
|---|---|---|
| EQ | 0 |
EQ |
| NE | 1 |
NE |
| GE | 2 |
GE |
| GT | 3 |
GT |
| LE | 4 |
LE |
| LT | 5 |
LT |
ComparisonType
Which comparison type to use.
Cases:
| Symbol | Value | String |
|---|---|---|
| NOTYPE | 0 |
NOTYPE |
| FLOAT | 1 |
FLOAT |
| TOTALORDER | 2 |
TOTALORDER |
| SIGNED | 3 |
SIGNED |
| UNSIGNED | 4 |
UNSIGNED |
CustomCallApiVersion
Custom call API version
Cases:
| Symbol | Value | String |
|---|---|---|
| API_VERSION_UNSPECIFIED | 0 |
API_VERSION_UNSPECIFIED |
| API_VERSION_ORIGINAL | 1 |
API_VERSION_ORIGINAL |
| API_VERSION_STATUS_RETURNING | 2 |
API_VERSION_STATUS_RETURNING |
| API_VERSION_STATUS_RETURNING_UNIFIED | 3 |
API_VERSION_STATUS_RETURNING_UNIFIED |
| API_VERSION_TYPED_FFI | 4 |
API_VERSION_TYPED_FFI |
CustomCallSchedule
Specifies the desired schedule for the custom-call.
Cases:
| Symbol | Value | String |
|---|---|---|
| NONE | 0 |
NONE |
| LATEST | 1 |
LATEST |
| EARLIEST | 2 |
EARLIEST |
DequantizeMode
_Dequantization mode. Only MINCOMBINED is supported.
Cases:
| Symbol | Value | String |
|---|---|---|
| MIN_COMBINED | 0 |
MIN_COMBINED |
DomainKind
Kind of domain metatdata attached to an HLO domain.
Cases:
| Symbol | Value | String |
|---|---|---|
| sharding | 0 |
sharding |
FftType
XLA fast fourier transform type.
Cases:
| Symbol | Value | String |
|---|---|---|
| FFT | 0 |
FFT |
| IFFT | 1 |
IFFT |
| RFFT | 2 |
RFFT |
| IRFFT | 3 |
IRFFT |
FusionKind
Fusion kind
Cases:
| Symbol | Value | String |
|---|---|---|
| kLoop | 0 |
kLoop |
| kInput | 1 |
kInput |
| kOutput | 2 |
kOutput |
| kCustom | 3 |
kCustom |
Precision
XLA precision for an operand. Has backend specific meaning.
Cases:
| Symbol | Value | String |
|---|---|---|
| DEFAULT | 0 |
DEFAULT |
| HIGH | 1 |
HIGH |
| HIGHEST | 2 |
HIGHEST |
ResultAccuracyMode
XLA result accuracy mode.
Cases:
| Symbol | Value | String |
|---|---|---|
| DEFAULT | 0 |
DEFAULT |
| HIGHEST | 1 |
HIGHEST |
| TOLERANCE | 2 |
TOLERANCE |
RngAlgorithm
XLA PRNG algorithm to be used.
Cases:
| Symbol | Value | String |
|---|---|---|
| DEFAULT | 0 |
DEFAULT |
| THREE_FRY | 1 |
THREE_FRY |
| PHILOX | 2 |
PHILOX |
RngDistribution
XLA PRNG distribution to be used.
Cases:
| Symbol | Value | String |
|---|---|---|
| UNIFORM | 1 |
UNIFORM |
| NORMAL | 2 |
NORMAL |
Transpose
Transpose options
Cases:
| Symbol | Value | String |
|---|---|---|
| TRANSPOSE_INVALID | 0 |
TRANSPOSE_INVALID |
| NO_TRANSPOSE | 1 |
NO_TRANSPOSE |
| TRANSPOSE | 2 |
TRANSPOSE |
| ADJOINT | 3 |
ADJOINT |