# 'shape' Dialect

Description of operations & types within the Shape dialect as well as their usage.

Types and operations for shape dialect This dialect contains operations for shape inference.

Note: Unless explicitly stated, all functions that return a shape and take shapes as input, return the invalid shape if one of its operands is an invalid shape. This avoids flagging multiple errors for one verification failure. The dialect itself does not specify how errors should be combined (there are multiple different options, from always choosing first operand, concatting etc. on how to combine them).

## Operation definition ¶

`shape.add`

(shape::AddOp) ¶

*Addition of sizes and indices*

Syntax:

```
operation ::= `shape.add` $lhs `,` $rhs attr-dict `:` type($lhs) `,` type($rhs) `->` type($result)
```

Adds two sizes or indices. If either operand is an error it will be
propagated to the result. The operands can be of type `size`

or `index`

. If
at least one of the operands can hold an error, i.e. if it is of type
`size`

, the result must be of type `size`

. If error propagation is not
possible because both operands are of type `index`

then the result may be
of type `size`

or `index`

.

Traits: AlwaysSpeculatableImplTrait, Commutative, InferTypeOpAdaptor

Interfaces: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

#### Operands: ¶

Operand | Description |
---|---|

`lhs` | size or index |

`rhs` | size or index |

#### Results: ¶

Result | Description |
---|---|

`result` | size or index |

`shape.any`

(shape::AnyOp) ¶

*Return any combination of the input shapes*

Syntax:

```
operation ::= `shape.any` $inputs attr-dict `:` type($inputs) `->` type($result)
```

This operation takes multiple input shapes or extent tensors and returns some combination of their dimensions. This can be best seen with examples below.

The result is undefined, but still side-effect free, in cases where the inputs have differing ranks or differ in extents of shared dimensions.

Example:

```
%s0 = shape.any [2,?], [?,3] // [2,3]
%s1 = shape.any [?,?], [1,2] // [1,2]
```

Traits: AlwaysSpeculatableImplTrait, Commutative

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

#### Operands: ¶

Operand | Description |
---|---|

`inputs` | variadic of shape or extent tensor |

#### Results: ¶

Result | Description |
---|---|

`result` | shape or extent tensor |

`shape.assuming`

(shape::AssumingOp) ¶

*Execute the region*

Executes the region assuming all witnesses are true.

“assuming” operations represent an execution order restriction to the compiler, information for dependent code to rely on (by assuming), and nothing else. They should not exist after a program is fully lowered and ready to execute.

Traits: RecursiveMemoryEffects, SingleBlock, SingleBlockImplicitTerminator

Interfaces: RegionBranchOpInterface

#### Operands: ¶

Operand | Description |
---|---|

`witness` |

#### Results: ¶

Result | Description |
---|---|

`results` | variadic of any type |

`shape.assuming_all`

(shape::AssumingAllOp) ¶

*Return a logical AND of all witnesses*

Syntax:

```
operation ::= `shape.assuming_all` $inputs attr-dict
```

Used to simplify constraints as any single failing precondition is enough to prevent execution.

“assuming” operations represent an execution order restriction to the compiler, information for dependent code to rely on (by assuming), and nothing else. They should not exist after a program is fully lowered and ready to execute.

Example:

```
%w0 = shape.cstr_broadcastable [2,2], [3,1,2] // Passing
%w1 = shape.cstr_broadcastable [2,2], [3,2] // Failure
%w2 = shape.cstr_eq [1,2], [1,2], [1,2] // Passing
%wf = shape.assuming_all %w0, %w1 // Failure
%wt = shape.assuming_all %w0, %w2 // Passing
```

Traits: AlwaysSpeculatableImplTrait, Commutative

Interfaces: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

#### Operands: ¶

Operand | Description |
---|---|

`inputs` | variadic of |

#### Results: ¶

Result | Description |
---|---|

`result` |

`shape.assuming_yield`

(shape::AssumingYieldOp) ¶

*Yield operation*

Syntax:

```
operation ::= `shape.assuming_yield` attr-dict ($operands^ `:` type($operands))?
```

This yield operation represents a return operation within the
`shape.assuming`

operation region. The operation takes variable number of
operands and produces no results. The operand number and types must match
the number and types of parent `shape.assuming`

results.

Traits: AlwaysSpeculatableImplTrait, HasParent

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface), RegionBranchTerminatorOpInterface

Effects: MemoryEffects::Effect{}

#### Operands: ¶

Operand | Description |
---|---|

`operands` | variadic of any type |

`shape.broadcast`

(shape::BroadcastOp) ¶

*Returns the broadcasted output shape of two or more inputs*

Syntax:

```
operation ::= `shape.broadcast` $shapes attr-dict `:` type($shapes) `->` type($result)
```

Returns the broadcasted shape for input shapes or extent tensors. The rest
of this description is simplified for the 2 input case but can be extended
to more inputs. Both operands can be of type `shape.shape`

or
`tensor<?xindex>`

. The result is of type `shape.shape`

and, if both
operands are tensors, may be of type `tensor<?xindex>`

.

If the two operand shapes are of different rank the smaller one is padded with 1’s from the left. The resulting broadcasted shape is then defined as

```
result[i] = lhs[i] if lhs[i] == rhs[i]
= lhs[i] if rhs[i] == 1
= rhs[i] if lhs[i] == 1.
```

In case the resulting shape is undefined, i.e. if corresponding extents are different from each other but none is 1, the result is an error shape. Likewise error values are propagated if any of the operands holds an error value. If the result type is an extent tensor (and can therefore not hold the error value) the behavior may be undefined. The optional string attribute can be used to describe the error case.

Traits: AlwaysSpeculatableImplTrait, Commutative

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

#### Attributes: ¶

Attribute | MLIR Type | Description |
---|---|---|

`error` | ::mlir::StringAttr | string attribute |

#### Operands: ¶

Operand | Description |
---|---|

`shapes` | variadic of shape or extent tensor |

#### Results: ¶

Result | Description |
---|---|

`result` | shape or extent tensor |

`shape.concat`

(shape::ConcatOp) ¶

*Concatenates two shapes*

Syntax:

```
operation ::= `shape.concat` $lhs `,` $rhs attr-dict `:` type($lhs) `,` type($rhs) `->` type($result)
```

Creates a shape whose dimensions consist of first the dimensions from `lhs`

followed by the dimensions of `rhs`

.

Example: concat([2,3], [4,5]) -> [2,3,4,5] concat([], []) -> [] concat([], [4,5,6]) -> [4,5,6]

Traits: AlwaysSpeculatableImplTrait

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

#### Operands: ¶

Operand | Description |
---|---|

`lhs` | shape or extent tensor |

`rhs` | shape or extent tensor |

#### Results: ¶

Result | Description |
---|---|

`result` | shape or extent tensor |

`shape.const_shape`

(shape::ConstShapeOp) ¶

*Creates a constant shape or extent tensor*

Creates a constant shape or extent tensor. The individual extents are given
as the `shape`

attribute. The number of these values equals the shape’s
rank.

```
%0 = shape.const_shape [] : !shape.shape
%1 = shape.const_shape [1, 2, 3] : !shape.shape
%2 = shape.const_shape [4, 5, 6] : tensor<3xindex>
```

Traits: AlwaysSpeculatableImplTrait, ConstantLike, InferTypeOpAdaptor

Interfaces: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

#### Attributes: ¶

Attribute | MLIR Type | Description |
---|---|---|

`shape` | ::mlir::DenseIntElementsAttr | index elements attribute |

#### Results: ¶

Result | Description |
---|---|

`result` | shape or extent tensor |

`shape.const_size`

(shape::ConstSizeOp) ¶

*Creates a constant of type shape.size*

Syntax:

```
operation ::= `shape.const_size` $value attr-dict
```

Creates a `shape.size`

type representing the constant size given by `value`

.

```
%x = shape.const_size 10
```

Traits: AlwaysSpeculatableImplTrait, ConstantLike

Interfaces: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface), OpAsmOpInterface

Effects: MemoryEffects::Effect{}

#### Attributes: ¶

Attribute | MLIR Type | Description |
---|---|---|

`value` | ::mlir::IntegerAttr | index attribute |

#### Results: ¶

Result | Description |
---|---|

`result` |

`shape.const_witness`

(shape::ConstWitnessOp) ¶

*An operation that returns a statically known witness value*

Syntax:

```
operation ::= `shape.const_witness` $passing attr-dict
```

This operation represents a statically known witness result. This can be often used to canonicalize/fold constraint and assuming code that will always pass.

```
%0 = shape.const_shape [1,2,3]
%1 = shape.const_shape [1,2,3]
%w0 = shape.cstr_eq(%0, %1) // Can be folded to "const_witness true"
%w1 = shape.const_witness true
%w2 = shape.assuming_all(%w0, %w2) // Can be folded to "const_witness true"
```

Traits: AlwaysSpeculatableImplTrait, ConstantLike

Effects: MemoryEffects::Effect{}

#### Attributes: ¶

Attribute | MLIR Type | Description |
---|---|---|

`passing` | ::mlir::BoolAttr | bool attribute |

#### Results: ¶

Result | Description |
---|---|

`result` |

`shape.cstr_broadcastable`

(shape::CstrBroadcastableOp) ¶

*Determines if 2+ shapes can be successfully broadcasted*

Syntax:

```
operation ::= `shape.cstr_broadcastable` $shapes attr-dict `:` type($shapes)
```

Given input shapes or extent tensors, return a witness specifying if they are broadcastable. This broadcastable follows the same logic as what shape.broadcast documents.

“cstr” operations represent runtime assertions.

Example:

```
%w0 = shape.cstr_broadcastable [2,2], [3,1,2] // Passing
%w1 = shape.cstr_broadcastable [2,2], [3,2] // Failure
```

Traits: Commutative

Interfaces: InferTypeOpInterface

#### Operands: ¶

Operand | Description |
---|---|

`shapes` | variadic of shape or extent tensor |

#### Results: ¶

Result | Description |
---|---|

`result` |

`shape.cstr_eq`

(shape::CstrEqOp) ¶

*Determines if all input shapes are equal*

Syntax:

```
operation ::= `shape.cstr_eq` $shapes attr-dict `:` type($shapes)
```

Given 1 or more input shapes, determine if all shapes are the exact same.

“cstr” operations represent runtime assertions.

Example:

```
%w0 = shape.cstr_eq [1,2], [1,2], [1,2] // Passing
%w1 = shape.cstr_eq [2,2], [1,2] // Failure
```

Traits: Commutative

Interfaces: InferTypeOpInterface

#### Operands: ¶

Operand | Description |
---|---|

`shapes` | variadic of shape or extent tensor |

#### Results: ¶

Result | Description |
---|---|

`result` |

`shape.cstr_require`

(shape::CstrRequireOp) ¶

*Represents a runtime assertion that an i1 is true*

Syntax:

```
operation ::= `shape.cstr_require` $pred `,` $msg attr-dict
```

Represents a runtime assertion that an i1 is true. It returns a !shape.witness to order this assertion.

For simplicity, prefer using other cstr_* ops if they are available for a given constraint.

Example:

```
%bool = ...
%w0 = shape.cstr_require %bool, "msg" // Passing if `%bool` is true.
```

Since this op can be used to express many different possible assertions
(depending on whatever computation calculated `pred`

), the `msg`

should clarify the nature of the assertion for users.

Interfaces: InferTypeOpInterface

#### Attributes: ¶

Attribute | MLIR Type | Description |
---|---|---|

`msg` | ::mlir::StringAttr | string attribute |

#### Operands: ¶

Operand | Description |
---|---|

`pred` | 1-bit signless integer |

#### Results: ¶

Result | Description |
---|---|

`result` |

`shape.debug_print`

(shape::DebugPrintOp) ¶

*Prints the input shape or size*

Prints the input dim or shape and passes through input.

Note: This is intended for testing and debugging only.

#### Operands: ¶

Operand | Description |
---|---|

`input` | shape or size |

#### Results: ¶

Result | Description |
---|---|

`output` | shape or size |

`shape.dim`

(shape::DimOp) ¶

*Gets the specified extent from the shape of a shaped input*

Syntax:

```
operation ::= `shape.dim` $value `,` $index attr-dict `:` type($value) `,`type($index) `->` type($extent)
```

Gets the extent indexed by `dim`

from the shape of the `value`

operand. If
the index is error or out-of-bound then it returns an invalid size if the
return type carries error information else the behavior is undefined.

This is a convenience op that performs the equivalent of getting the extent
of a shape (e.g., `dim(x, i) == get_extent(shape_of(x), i)`

).

Traits: AlwaysSpeculatableImplTrait, InferTypeOpAdaptor

Effects: MemoryEffects::Effect{}

#### Operands: ¶

Operand | Description |
---|---|

`value` | shaped of any type values |

`index` | size or index |

#### Results: ¶

Result | Description |
---|---|

`extent` | size or index |

`shape.div`

(shape::DivOp) ¶

*Division of sizes and indices*

Syntax:

```
operation ::= `shape.div` $lhs `,` $rhs attr-dict `:` type($lhs) `,` type($rhs) `->` type($result)
```

Divides two sizes or indices. If either operand is an error it will be
propagated to the result. The operands can be of type `size`

or `index`

.
If at least one of the operands can hold an error, i.e. if it is of type
`size`

, the result must be of type `size`

. If error propagation is not
possible because both operands are of type `index`

then the result may be
of type `size`

or `index`

. If both operands and result are of type
`index`

, their runtime values could be negative. The result is rounded
toward negative infinity, i.e. floor(lhs / rhs), such that

```
div(lhs, rhs) * rhs + mod(lhs, rhs) = lhs
```

always holds. If any of the values is of type `size`

, the behavior for
negative value is undefined.

Traits: AlwaysSpeculatableImplTrait, InferTypeOpAdaptor

Effects: MemoryEffects::Effect{}

#### Operands: ¶

Operand | Description |
---|---|

`lhs` | size or index |

`rhs` | size or index |

#### Results: ¶

Result | Description |
---|---|

`result` | size or index |

`shape.from_extent_tensor`

(shape::FromExtentTensorOp) ¶

*Creates a shape from a tensor of extents*

Syntax:

```
operation ::= `shape.from_extent_tensor` $input attr-dict `:` type($input)
```

Creates a shape from a 1D integral tensor of extents. The rank of the resulting shape equals the number of elements in the tensor, and the extents match the values of the elements.

Traits: AlwaysSpeculatableImplTrait

Effects: MemoryEffects::Effect{}

#### Operands: ¶

Operand | Description |
---|---|

`input` | 1D tensor of index values |

#### Results: ¶

Result | Description |
---|---|

`result` |

`shape.from_extents`

(shape::FromExtentsOp) ¶

*Creates a shape from extents*

Syntax:

```
operation ::= `shape.from_extents` $extents attr-dict `:` type($extents)
```

Creates a shape from multiple SSA values representing the extents of the shape.

```
// Rank 2 shape.
%s0 = shape.from_extents %a, %b
// Rank 0 shape.
%s1 = shape.from_extents
```

Traits: AlwaysSpeculatableImplTrait

Effects: MemoryEffects::Effect{}

#### Operands: ¶

Operand | Description |
---|---|

`extents` | variadic of size or index |

#### Results: ¶

Result | Description |
---|---|

`shape` |

`shape.func`

(shape::FuncOp) ¶

*Shape function*

An operation with a name containing a single `SSACFG`

region which
represents a shape transfer function or helper function for shape transfer
function.

Traits: AffineScope, AutomaticAllocationScope, IsolatedFromAbove

Interfaces: CallableOpInterface, FunctionOpInterface, OpAsmOpInterface, Symbol

#### Attributes: ¶

Attribute | MLIR Type | Description |
---|---|---|

`sym_name` | ::mlir::StringAttr | string attribute |

`function_type` | ::mlir::TypeAttr | type attribute of function type |

`arg_attrs` | ::mlir::ArrayAttr | Array of dictionary attributes |

`res_attrs` | ::mlir::ArrayAttr | Array of dictionary attributes |

`sym_visibility` | ::mlir::StringAttr | string attribute |

`shape.function_library`

(shape::FunctionLibraryOp) ¶

*Represents shape functions and corresponding ops*

Represents a list of shape functions and the ops whose shape transfer functions they represent.

Example:

```
shape.function_library {
func @same_result_shape(%arg: !shape.value_shape) -> !shape.shape {
%0 = shape_of %arg : !shape.value_shape -> !shape.shape
return %0 : !shape.shape
}
} mapping {
std.atan = @same_result_shape
}
```

Traits: AffineScope, IsolatedFromAbove, NoRegionArguments, NoTerminator, SingleBlock, SymbolTable

Interfaces: OpAsmOpInterface, Symbol

#### Attributes: ¶

Attribute | MLIR Type | Description |
---|---|---|

`sym_name` | ::mlir::StringAttr | string attribute |

`sym_visibility` | ::mlir::StringAttr | string attribute |

`mapping` | ::mlir::DictionaryAttr | dictionary of named attribute values |

`shape.get_extent`

(shape::GetExtentOp) ¶

*Gets the specified extent from a shape or extent tensor*

Syntax:

```
operation ::= `shape.get_extent` $shape `,` $dim attr-dict `:` type($shape) `,` type($dim) `->` type($extent)
```

Gets the extent indexed by `dim`

from the `shape`

operand. If the shape is
an error then it returns an invalid size.

Traits: AlwaysSpeculatableImplTrait, InferTypeOpAdaptor

Effects: MemoryEffects::Effect{}

#### Operands: ¶

Operand | Description |
---|---|

`shape` | shape or extent tensor |

`dim` | size or index |

#### Results: ¶

Result | Description |
---|---|

`extent` | size or index |

`shape.index_to_size`

(shape::IndexToSizeOp) ¶

*Converts a standard index to a shape size*

Syntax:

```
operation ::= `shape.index_to_size` $arg attr-dict
```

Converts a standard index to a `shape.size`

. This operation and its
inverse, `size_to_index`

, facilitate index conversion between the standard
and the shape dialect.

The behavior is undefined for negative indices.

Traits: AlwaysSpeculatableImplTrait

Effects: MemoryEffects::Effect{}

#### Operands: ¶

Operand | Description |
---|---|

`arg` | index |

#### Results: ¶

Result | Description |
---|---|

`result` |

`shape.is_broadcastable`

(shape::IsBroadcastableOp) ¶

*Determines if 2+ shapes can be successfully broadcasted*

Syntax:

```
operation ::= `shape.is_broadcastable` $shapes attr-dict `:` type($shapes)
```

Given multiple input shapes or extent tensors, return a predicate specifying if they are broadcastable. This broadcastable follows the same logic as what shape.broadcast documents.

Concretely, shape.is_broadcastable returning true implies that shape.broadcast will not give an error, and shape.cstr_broadcastable will not result in an assertion failure. Similarly, false implies an error or assertion failure.

Example:

```
%true = shape.is_broadcastable [2,2], [3,1,2]
%false = shape.is_broadcastable [2,2], [3,2]
```

Traits: Commutative

Interfaces: InferTypeOpInterface

#### Operands: ¶

Operand | Description |
---|---|

`shapes` | variadic of shape or extent tensor |

#### Results: ¶

Result | Description |
---|---|

`result` | 1-bit signless integer |

`shape.max`

(shape::MaxOp) ¶

*Elementwise maximum*

Syntax:

```
operation ::= `shape.max` $lhs `,` $rhs attr-dict `:` type($lhs) `,` type($rhs) `->` type($result)
```

Computes the elementwise maximum of two sizes or shapes with equal ranks. If either operand is an error, then an error will be propagated to the result. If the input types mismatch or the ranks do not match, then the result is an error.

Traits: AlwaysSpeculatableImplTrait, Commutative, InferTypeOpAdaptor

Effects: MemoryEffects::Effect{}

#### Operands: ¶

Operand | Description |
---|---|

`lhs` | shape or size |

`rhs` | shape or size |

#### Results: ¶

Result | Description |
---|---|

`result` | shape or size |

`shape.meet`

(shape::MeetOp) ¶

*Returns the least general shape or size of its operands*

Syntax:

```
operation ::= `shape.meet` $arg0 `,` $arg1 (`,` `error` `=` $error^)? attr-dict `:`
type($arg0) `,` type($arg1) `->` type($result)
```

An operation that computes the least general shape or dim of input operands. This effectively asserts that corresponding static dimensions are equal. The behavior is to match each element of the shape/size and propagate the most restrictive information, returning an invalid shape if there are contradictory requirements. E.g., using pseudo code

```
shape.meet([*], [*]) -> [*]
shape.meet([*], [1, ?]) -> [1, ?]
shape.meet([1, 2], [1, ?]) -> [1, 2]
shape.meet([*], [1, 2]) -> [1, 2]
shape.meet([], []) -> []
shape.meet([], [*]) -> []
shape.meet([], [?, ?]) -> [invalid]
shape.meet([1, ?], [2, ?, ?]) -> [invalid]
```

`shape.meet`

also allows specifying an optional error string, that may be
used to return an error to the user upon mismatch of dimensions.

```
%c = shape.meet %a, %b, error="<reason>" : !shape.shape, !shape.shape -> !shape.shape
```

Traits: Commutative, InferTypeOpAdaptor

Interfaces: InferTypeOpInterface

#### Attributes: ¶

Attribute | MLIR Type | Description |
---|---|---|

`error` | ::mlir::StringAttr | string attribute |

#### Operands: ¶

Operand | Description |
---|---|

`arg0` | any shape or size |

`arg1` | any shape or size |

#### Results: ¶

Result | Description |
---|---|

`result` | any shape or size |

`shape.min`

(shape::MinOp) ¶

*Elementwise minimum*

Syntax:

```
operation ::= `shape.min` $lhs `,` $rhs attr-dict `:` type($lhs) `,` type($rhs) `->` type($result)
```

Computes the elementwise minimum of two sizes or shapes with equal ranks. If either operand is an error, then an error will be propagated to the result. If the input types mismatch or the ranks do not match, then the result is an error.

Traits: AlwaysSpeculatableImplTrait, Commutative, InferTypeOpAdaptor

Effects: MemoryEffects::Effect{}

#### Operands: ¶

Operand | Description |
---|---|

`lhs` | shape or size |

`rhs` | shape or size |

#### Results: ¶

Result | Description |
---|---|

`result` | shape or size |

`shape.mul`

(shape::MulOp) ¶

*Multiplication of sizes and indices*

Syntax:

```
operation ::= `shape.mul` $lhs `,` $rhs attr-dict `:` type($lhs) `,` type($rhs) `->` type($result)
```

Multiplies two sizes or indices. If either operand is an error it will be
propagated to the result. The operands can be of type `size`

or `index`

. If
at least one of the operands can hold an error, i.e. if it is of type
`size`

, the result must be of type `size`

. If error propagation is not
possible because both operands are of type `index`

then the result may be
of type `size`

or `index`

.

Traits: AlwaysSpeculatableImplTrait, Commutative, InferTypeOpAdaptor

Effects: MemoryEffects::Effect{}

#### Operands: ¶

Operand | Description |
---|---|

`lhs` | size or index |

`rhs` | size or index |

#### Results: ¶

Result | Description |
---|---|

`result` | size or index |

`shape.num_elements`

(shape::NumElementsOp) ¶

*Returns the number of elements for a given shape*

Syntax:

```
operation ::= `shape.num_elements` $shape attr-dict `:` type($shape) `->` type($result)
```

Returns the number of elements for a given shape which is the product of
its extents. If the argument is of type `shape`

then the result will be of
type `size`

and potential errors will be propagated. Otherwise, if the
argument is and extent tensor `tensor<?xindex>`

then the result will be of
type `index`

.

Traits: AlwaysSpeculatableImplTrait, InferTypeOpAdaptor

Effects: MemoryEffects::Effect{}

#### Operands: ¶

Operand | Description |
---|---|

`shape` | shape or extent tensor |

#### Results: ¶

Result | Description |
---|---|

`result` | size or index |

`shape.rank`

(shape::RankOp) ¶

*Gets the rank of a shape*

Syntax:

```
operation ::= `shape.rank` $shape attr-dict `:` type($shape) `->` type($rank)
```

Returns the rank of the shape or extent tensor, i.e. the number of extents.

Traits: AlwaysSpeculatableImplTrait, InferTypeOpAdaptor

Effects: MemoryEffects::Effect{}

#### Operands: ¶

Operand | Description |
---|---|

`shape` | shape or extent tensor |

#### Results: ¶

Result | Description |
---|---|

`rank` | size or index |

`shape.reduce`

(shape::ReduceOp) ¶

*Returns an expression reduced over a shape or extent tensor*

An operation that takes as input a shape or extent tensor, and a number of initial values. This operation has a region that is applied repeatedly for every extent of the input. Starting with the initial values, the individual extents are then aggregated as defined by the associated region.

Conceptually this op performs the following reduction:

```
res[] = init;
for (int i = 0, i < shape.rank(); i++) {
res = reduce(i, shape[i], res[0], ..., res[n]);
}
```

Where `reduce`

represents the region attached and the result of the reduce
op is the last computed output of the reduce region. As an example, the
number of elements can be computed as follows:

```
func.func @reduce(%shape : !shape.shape, %init : !shape.size) ->
!shape.size {
%num_elements = shape.reduce(%shape, %init) -> !shape.size {
^bb0(%index: index, %dim: !shape.size, %acc: !shape.size):
%updated_acc = "shape.mul"(%acc, %dim) :
(!shape.size, !shape.size) -> !shape.size
shape.yield %updated_acc : !shape.size
}
return %num_elements : !shape.size
}
```

Traits: SingleBlock, SingleBlockImplicitTerminator

#### Operands: ¶

Operand | Description |
---|---|

`shape` | shape or extent tensor |

`initVals` | variadic of any type |

#### Results: ¶

Result | Description |
---|---|

`result` | variadic of any type |

`shape.return`

(shape::ReturnOp) ¶

*Shape function return operation*

Syntax:

```
operation ::= `shape.return` attr-dict ($operands^ `:` type($operands))?
```

The `shape.return`

operation represents a return operation within a
function. The operation takes variable number of operands and produces no
results.

Traits: AlwaysSpeculatableImplTrait, HasParent

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface), RegionBranchTerminatorOpInterface

Effects: MemoryEffects::Effect{}

#### Operands: ¶

Operand | Description |
---|---|

`operands` | variadic of any type |

`shape.shape_eq`

(shape::ShapeEqOp) ¶

*Returns whether the input shapes or extent tensors are equal*

Syntax:

```
operation ::= `shape.shape_eq` $shapes attr-dict `:` type($shapes)
```

Takes one or more shape or extent tensor operands and determines whether they are equal. When extent tensors are compared to shapes they are regarded as their equivalent non-error shapes. Error shapes can be tested for equality like any other shape value, meaning that the error value is equal to itself.

Traits: AlwaysSpeculatableImplTrait, Commutative

Effects: MemoryEffects::Effect{}

#### Operands: ¶

Operand | Description |
---|---|

`shapes` | variadic of shape or extent tensor |

#### Results: ¶

Result | Description |
---|---|

`result` | 1-bit signless integer |

`shape.shape_of`

(shape::ShapeOfOp) ¶

*Returns shape of a value or shaped type operand*

Syntax:

```
operation ::= `shape.shape_of` $arg attr-dict `:` type($arg) `->` type($result)
```

The operation takes a value or a shaped operand as an argument and it returns a shape or extent tensor.

Traits: AlwaysSpeculatableImplTrait, InferTypeOpAdaptor

Effects: MemoryEffects::Effect{}

#### Operands: ¶

Operand | Description |
---|---|

`arg` | shaped of any type values or |

#### Results: ¶

Result | Description |
---|---|

`result` | shape or extent tensor |

`shape.size_to_index`

(shape::SizeToIndexOp) ¶

*Casts between index types of the shape and standard dialect*

Syntax:

```
operation ::= `shape.size_to_index` $arg attr-dict `:` type($arg)
```

Converts a `shape.size`

to a standard index. This operation and its
inverse, `index_to_size`

, facilitate index conversion between the standard
and the shape dialect. The behavior is undefined for unknown and invalid
arguments.

Traits: AlwaysSpeculatableImplTrait

Interfaces: CastOpInterface, ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

#### Operands: ¶

Operand | Description |
---|---|

`arg` | size or index |

#### Results: ¶

Result | Description |
---|---|

`result` | index |

`shape.split_at`

(shape::SplitAtOp) ¶

*Splits a shape at a given index*

Splits a shape at a given dimension `index`

, returning two shapes. If
`index`

is negative, it is treated as indexing from the back of the shape.
This negative-handling behavior is important when handling unranked shapes,
where the positive index is not necessarily knowable due to a dynamic
number of leading dimensions. If the result is in extent tensor form out of
bounds indices result in undefined behavior.

Examples:

- split_at([4,5,6], index=0) -> [], [4,5,6]
- split_at([4,5,6], index=1) -> [4], [5,6]
- split_at([4,5,6], index=2) -> [4,5], [6]
- split_at([4,5,6], index=3) -> [4,5,6], []
- split_at([4,5,6], index=4) -> error
- split_at([4,5,6], index=-1) -> [4,5], [6]
- split_at([4,5,6], index=-2) -> [4], [5,6]
- split_at([4,5,6], index=-3) -> [], [4,5,6]
- split_at([4,5,6], index=-4) -> error

Requires:

`index`

is in the range [-rank(operand),rank(operand)]

Traits: AlwaysSpeculatableImplTrait

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

#### Operands: ¶

Operand | Description |
---|---|

`operand` | shape or extent tensor |

`index` | size or index |

#### Results: ¶

Result | Description |
---|---|

`head` | shape or extent tensor |

`tail` | shape or extent tensor |

`shape.to_extent_tensor`

(shape::ToExtentTensorOp) ¶

*Creates a dimension tensor from a shape*

Syntax:

```
operation ::= `shape.to_extent_tensor` $input attr-dict `:` type($input) `->` type($result)
```

Converts a shape to a 1D integral tensor of extents. The number of elements in the tensor equals the rank of the shape, and the elements equal the extents of the shape.

If the shape represents an error, this op’s behavior is undefined.

Traits: AlwaysSpeculatableImplTrait

Interfaces: CastOpInterface, ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

#### Operands: ¶

Operand | Description |
---|---|

`input` | shape or extent tensor |

#### Results: ¶

Result | Description |
---|---|

`result` | tensor of index values |

`shape.value_as_shape`

(shape::ValueAsShapeOp) ¶

*Returns value as a shape*

Syntax:

```
operation ::= `shape.value_as_shape` $arg attr-dict `:` type($arg) `->` type($result)
```

The operations takes a ValueShape and returns a Shape corresponding to the value. If the input value cannot be shape (e.g., not a 1D tensor of integral value representing sizes) then this propagages the error shape. E.g.,

```
// The following
%0 = arith.constant dense<[1,2]> : tensor<2xi32>
%shape = shape.value_as_shape %0 : tensor<2xi32> -> !shape.shape
// is equivalent to
%shape' = shape.const_shape [1, 2] : !shape.shape
```

This operation is the complement of `shape_of`

wrt ValueShape values.

Traits: AlwaysSpeculatableImplTrait

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

#### Operands: ¶

Operand | Description |
---|---|

`arg` | 1D tensor of integer or index values or |

#### Results: ¶

Result | Description |
---|---|

`result` | shape or extent tensor |

`shape.value_of`

(shape::ValueOfOp) ¶

*Returns value of a !shape.value_shape operand*

Syntax:

```
operation ::= `shape.value_of` $arg attr-dict `:` type($result)
```

The operation takes !shape.value_shape, a.k.a. (value, shape) tuple as an argument, and returns its value. The behavior is undefined for unknown and invalid arguments.

Traits: AlwaysSpeculatableImplTrait

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

#### Operands: ¶

Operand | Description |
---|---|

`arg` |

#### Results: ¶

Result | Description |
---|---|

`result` | shaped of any type values |

`shape.with_shape`

(shape::WithOp) ¶

*Returns ValueShape with given shape*

Syntax:

```
operation ::= `shape.with_shape` operands attr-dict `:` type($operand) `,` type($shape)
```

Returns ValueShape with the shape updated to match the shape operand. That
is a new ValueShape tuple is created with value equal to `operand`

’s
value and shape equal to `shape`

. If the ValueShape and given `shape`

are
non-conformant, then the returned ValueShape will represent an error of
this mismatch. Similarly if either inputs are in an error state, then an
error is propagated.

Usage: %0 = shape.with_shape %1, %2 : tensor<…>, !shape.shape

This is used, for example, where one combines shape function calculations and/or call one shape function from another. E.g.,

```
func.func @shape_foobah(%a: !shape.value_shape,
%b: !shape.value_shape,
%c: !shape.value_shape) -> !shape.shape {
%0 = call @shape_foo(%a, %b) :
(!shape.value_shape, !shape.value_shape) -> !shape.shape
%1 = shape.with_shape %b, %0 : !shape.value_shape, !shape.shape
%2 = call @shape_bah(%c, %1) :
(!shape.value_shape, !shape.value_shape) -> !shape.shape
return %2 : !shape.shape
}
```

This op need not be a refinement of the shape. In non-error cases the input
ValueShape’s value and shape are conformant and so too for the output, but
the result may be less specified than `operand`

’s shape as `shape`

is
merely used to construct the new ValueShape. If join behavior is desired
then a join op should be used.

Traits: AlwaysSpeculatableImplTrait

Effects: MemoryEffects::Effect{}

#### Operands: ¶

Operand | Description |
---|---|

`operand` | shaped of any type values or |

`shape` | shape or extent tensor |

#### Results: ¶

Result | Description |
---|---|

`result` |

`shape.yield`

(shape::YieldOp) ¶

*Returns the value to parent op*

Syntax:

```
operation ::= `shape.yield` attr-dict ($operands^ `:` type($operands))?
```

Traits: AlwaysSpeculatableImplTrait, HasParent<ReduceOp, FunctionLibraryOp>, ReturnLike, Terminator

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface), RegionBranchTerminatorOpInterface

Effects: MemoryEffects::Effect{}

#### Operands: ¶

Operand | Description |
---|---|

`operands` | variadic of any type |

## Type definition ¶

### ShapeType ¶

Syntax: `!shape.shape`

`shape.shape`

represents either an unranked shape, a ranked shape with
possibly unknown dimensions or an invalid shape. The rank is of type
`shape.size`

and, if rank is known, the extent is a 1D tensor of type
`shape.size`

.

Shape is printed:

`[*]`

if it is an unranked shape`[?, 2]`

if a rank 2 tensor with one unknown dimension`[3, 4]`

is a rank 2 static tensor`[]`

is a scalar`[1]`

is a rank 1 tensor with 1 element`[invalid]`

for an invalid shape

### SizeType ¶

Syntax: `!shape.size`

`shape.size`

represents a non-negative integer with support for being
unknown and invalid.

Operations on `shape.size`

types are specialized to handle unknown/dynamic
value. So, for example, `<unknown> + x == <unknown>`

for all non-error `x : !shape.size`

(e.g., an unknown value does not become known due to addition).

### ValueShapeType ¶

Syntax: `!shape.value_shape`

`shape.value_shape`

represents the value produced by an operation (this
corresponds to `Value`

in the compiler) and a shape. Conceptually this is a
tuple of a value (potentially unknown) and `shape.shape`

. The value and
shape can either or both be unknown. If both the `value`

and `shape`

are
known, then the shape of `value`

is conformant with `shape`

. That is, the
shape of the value conforms to the shape of the ValueShape, so that if we
have ValueShape `(value, shape)`

then `join(shape_of(value), shape)`

would
be error free and in particular it means that if both are statically known,
then they are equal.

### WitnessType ¶

Syntax: `!shape.witness`

A witness is a structural device in the compiler to maintain ordering of code relying on information obtained from passing assertions. Witnesses do not represent any physical data.

“cstr_” operations will return witnesses and be lowered into assertion logic when not resolvable at compile time.

“assuming_” operations will take witnesses as input and represent only information to the compiler, so they do not exist in executing code. Code that is dependent on “assuming_” operations can assume all cstr operations transitively before are honored as true.

These abstractions are intended to allow the compiler more freedom with assertions by merely showing the assertion through dataflow at this time rather than a side effecting operation that acts as a barrier. This can be viewed similarly to a compiler representation of promises from asynchronous, possibly crashing assertions. Reliant code will not be reordered to before the code and non-reliant code can be reordered freely, and there are no guarantees on the final ordering of the assertions or their related code.

## Different stages of lowering Shape dialect ¶

In this section we shall give a brief overview of the different uses of the shape dialect and the lowering between these uses. Currently we have 3 worlds / stages of lowering of shape functions:

*Error monadic/error carrying/user specification*: This “input” form carries both the shape and whether in error state as value. Hence at this level all operations are pure operations producing and consuming values where the values could represent an error.*Constrained*: This form uses a variant of explicit evidence passing to allow leveraging existing compiler infrastructure to preserve safety information during optimization.*Side-effecting/asserting*: This final lowered form is imperative form with side-effecting ops (e.g., assert) for final codegen.

We are going to do a quick step through of the lowering using the example of a matmul.

Starting from the shape function of matmul in the error monadic form
below^{1}:

```
shape.function_library @shplib {
func.func @matmul(%lhs: !shape.value_shape, %rhs: !shape.value_shape) -> !shape.shape {
%c1 = shape.const_size 1
%c2 = shape.const_size 2
// We could also allow rank etc operations directly on value_shape too, that
// would make it nicer as "input" language, but keeping it explicit inside the
// IR instead and then we could have helper methods in front-end language.
%lhs_shape = shape.shape_of %lhs : !shape.value_shape -> !shape.shape
%rhs_shape = shape.shape_of %rhs : !shape.value_shape -> !shape.shape
%lhs_rank = shape.rank %lhs_shape : !shape.shape -> !shape.size
%rhs_rank = shape.rank %rhs_shape : !shape.shape -> !shape.size
// This is not minimal as one could ensure the ranks are the same below, also a
// variadic meet would make it more concise too.
%r = "shape.meet"(%lhs_rank, %rhs_rank) : (!shape.size, !shape.size) -> !shape.size
%rank = shape.meet %c2, %r, error="requires rank 2 operands" :
!shape.size, !shape.size -> !shape.size
%l0, %l1 = "shape.split_at"(%lhs_shape, %c1) :
(!shape.shape, !shape.size) -> (!shape.shape, !shape.shape)
%r0, %r1 = "shape.split_at"(%rhs_shape, %c1) :
(!shape.shape, !shape.size) -> (!shape.shape, !shape.shape)
%c = shape.meet %l1, %r0, error="inner dimensions required to match" :
!shape.shape, !shape.shape -> !shape.shape
%res = shape.concat %l0, %r1
// Should have `shape.return %res requires %c, %rank` to enable
return %res : !shape.shape
}
} mapping {
foo.matmul = @matmul
}
```

We are using the default builtin func and return here. Preferably we’d use ‘shape_func’ as a special function op that allows passing multiple results back that affect correct execution (e.g., serves as an error join)

- This would also means one can’t reify it inside a regular function without handling the shape.return - that is a feature here as these are more of a template.
- Currently we also have not marked
`meet`

as having no side-effects to avoid DCE until we have`shape.return`

, at which point computing the meet could be treated as purely computational returning error.

Meet represents a constraint that should hold, so should not be used to see

*if*something is equal. E.g., this means`meet`

can’t be used to represent`either(meet(x, y), meet(y,z))`

This could have been written more concisely as something like

`concat(lhs[0], rhs[1]) if rank(lhs) == 2 && rank(rhs) == 2 && lhs[1] == rhs[0]`

but not focusing on front-end proper here.

We are going to lower to “most” nested form directly (see test for an example reification along with legalization). In the above this was in a separate shape function library, while here we would normally reify it as part of lowering, but for simplicity will show as a standalone shape function.

```
func.func @matmul_shape1(%lhs: tensor<*xf32>, %rhs: tensor<*xindex>) -> tensor<?xindex> {
%c1 = shape.const_size 1
%c2 = shape.const_size 2
// We allow `shape.shape_of` to return either a `!shape.shape` or
// `tensor<?xindex>` type, in the case where the input is a tensor the most
// refined type is a tensor of `index` but not required.
%lhs_shape = shape.shape_of %lhs : tensor<*xf32> -> !shape.shape
%rhs_shape = shape.shape_of %rhs : tensor<*xf32> -> !shape.shape
%lhs_rank = shape.rank %lhs_shape : !shape.shape -> !shape.size
%rhs_rank = shape.rank %rhs_shape : !shape.shape -> !shape.size
%w1 = shape.cstr_eq %lhs_rank, %rhs_rank : !shape.witness
%res = shape.assuming %w1 -> tensor<?xindex> {
%r1 = shape.any %lhs_rank, %rhs_rank : (!shape.size, !shape.size) -> !shape.size
// Error message needs an addition, currently only on cstr_require.
%w2 = shape.cstr_eq %c2, %r1, error="requires rank 2 operands"
%res_1 = shape.assuming %w2 -> tensor<?xindex> {
// Here the lowered
// %rank = shape.any %c2, %r1 (!shape.size, !shape.size) -> !shape.size
// is dead and so elided further. But if `%rank` was actually consumed,
// then it could have been folded in `shape.any`.
%l0, %r0 = "shape.split_at"(%lhs_shape, %c1) :
(!shape.shape, !shape.size) -> !shape.shape
%l1, %r1 = "shape.split_at"(%lhs_shape, %c1) :
(!shape.shape, !shape.size) -> !shape.shape
%c = shape.meet %l1, %r0, error="inner dimensions required to match" :
!shape.size, !shape.size -> !shape.size
%res = concat(%l0, %r1)
shape.assuming_yield %res
}
shape.assuming_yield %res_1
}
return %res : tensor<?xindex>
}
```

We can now hoist computations of constraint were possible (which in the case below is not too many as we need to verify the rank before we can split)

```
func.func @matmul_shape2(%lhs: tensor<*xf32>, %lhs: tensor<*xf32>) -> tensor<?xindex> {
%c1 = shape.const_size 1
%c2 = shape.const_size 2
%lhs_shape = shape.shape_of %lhs : tensor<*xf32> -> tensor<?xindex>
%rhs_shape = shape.shape_of %rhs : tensor<*xf32> -> tensor<?xindex>
%lhs_rank = shape.rank %lhs_shape : tensor<?xindex> -> tensor<index>
%rhs_rank = shape.rank %rhs_shape : tensor<?xindex> -> tensor<index>
%w1 = shape.cstr_eq %c2, %lhs_rank, error="requires rank 2 operands"
%w2 = shape.cstr_eq %c2, %rhs_rank, error="requires rank 2 operands"
%w = shape.assuming_all %w1, %w2
%res = shape.assuming %w -> tensor<?xindex> {
%l0, %r0 = "shape.split_at"(%lhs_shape, %c1) :
(tensor<?xindex>, !shape.size) -> tensor<?xindex>
%l1, %r1 = "shape.split_at"(%lhs_shape, %c1) :
(tensor<?xindex>, !shape.size) -> tensor<?xindex>
%w3 = shape.cstr_eq %l1, %r0, error="inner dimensions required to match"
%res_2 = shape.assuming %w3 {
%res = concat(%l0, %r1)
shape.assuming_yield %res
}
shape.assuming_yield %res_1
}
return %res
}
```

The above form can now be lowered to the fully imperative form (see test for example).

```
func.func @matmul_shape3(%lhs: tensor<*xf32>, %lhs: tensor<*xf32>) -> tensor<?xindex> {
%c1 = arith.constant 1 : index
%c2 = arith.constant 2 : index
%lhs_shape = shape.shape_of %lhs : tensor<*xf32> -> tensor<?xindex>
%rhs_shape = shape.shape_of %rhs : tensor<*xf32> -> tensor<?xindex>
%lhs_rank = shape.rank %lhs_shape : tensor<?xindex> -> tensor<index>
%rhs_rank = shape.rank %rhs_shape : tensor<?xindex> -> tensor<index>
%w1 = shape.shape_eq %lhs_rank, %rhs_rank
%w2 = shape.shape_eq %c2, %lhs_rank
%w3 = and %w1, %w2
assert %w3, "requires rank 2 operands"
%l0, %l1 = shape.split_at(%lhs_shape, %c1) : tensor<?xindex>
%r0, %r1 = shape.split_at(%rhs_shape, %c1) : tensor<?xindex>
%w4 = shape.eq %l1, %r0
assert %w4, "inner dimensions required to match"
%res = concat(%l0, %r1)
return %res
}
```

- In this case form 3 is as easy and closer to form 1 (but only as no reordering was required). So it is a good question if the frontend authoring language could be more similar to the imperative form (under discussion).
- The above form presented here is an intermittent form during a lowering
pass. If used as input we would need to restrict the optimizations on it as
the
`shape`

dialect operations are no longer connected by producer-consumer to enforce guard checking.

The above could be further lowered by using `tensor.dim`

, `tensor.from_elements`

etc (or one could even lower these by way of, say, MHLO or TOSA dialect).

This form is least use inside the current workflows and needs more work. In particular in the example we use

`shape_func`

where in the code we instead use standard func as first form 1 isn’t used explicitly. ↩︎