MLIR

Multi-Level IR Compiler Framework

'shape' Dialect

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).

Type definition 

component type 

shape.element_type represents the element type of the ShapedType. It may be unknown, error or regular element type supported by ShapedType.

element type 

shape.element_type represents the element type of the ShapedType. It may be unknown, error or regular element type supported by ShapedType.

shape 

shape.type 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

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).

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.type. 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.

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.

Operation definition 

shape.add (shape::AddOp) 

Addition of sizes

Adds two valid sizes as follows:

  • lhs + rhs = unknown if either lhs or rhs unknown;
  • lhs + rhs = (int)lhs + (int)rhs if known;

Operands: 

OperandDescription
lhssize
rhssize

Results: 

ResultDescription
resultsize

shape.any (shape::AnyOp) 

Return any combination of the input shapes.

Syntax:

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

This operation takes multiple input shapes 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]

Operands: 

OperandDescription
inputsshape

Results: 

ResultDescription
resultshape

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] // Success
  %w1 = shape.cstr_broadcastable [2,2], [3,2] // Failure
  %w2 = shape.cstr_eq [1,2], [1,2], [1,2] // Success
  %wf = shape.assuming_all %w0, %w1 // Failure
  %wt = shape.assuming_all %w0, %w2 // Success

Operands: 

OperandDescription
inputswitness

Results: 

ResultDescription
resultwitness

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.

Operands: 

OperandDescription
witnesswitness

Results: 

ResultDescription
resultsany type

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 assert_and_exec region. The operation takes variable number of operands and produces no results. The operand number and types must match the return signature of the region that contains the operation.

Operands: 

OperandDescription
operandsany type

shape.broadcast (shape::BroadcastOp) 

Returns the broadcasted output shape of two inputs

Computes the broadcasted output shape following:

  1. If any inputs are unranked, output is unranked;

  2. Else the input array with number of dimensions smaller than the max input dimension, has 1’s prepended to its shapes and the output shape is calculated as follows:

    output[i] = lhs[i] if lhs[i] == rhs[i] or rhs[i] is unknown/undefined
              = rhs[i] if lhs[i] is unknown/undefined
              = lhs[i] if rhs[i] == 1
              = rhs[i] if lhs[i] == 1
              = error  if lhs[i] != rhs[i]
    

Op has an optional string attribute for the error case where there is no broadcastable output shape possible for the given inputs.

Attributes: 

AttributeMLIR TypeDescription
errorStringAttrstring attribute

Operands: 

OperandDescription
lhsshape
rhsshape

Results: 

ResultDescription
resultshape

shape.concat (shape::ConcatOp) 

Concatenates two shapes.

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]

Operands: 

OperandDescription
lhsshape
rhsshape

Results: 

ResultDescription
resultshape

shape.const_shape (shape::ConstShapeOp) 

Creates a constant of !shape.shape type.

Creates a !shape.shape with rank given by the length of shape and with dimension sizes given by the values of shape.

%0 = shape.const_shape []
%1 = shape.const_shape [1, 2, 3]

Attributes: 

AttributeMLIR TypeDescription
shapeDenseIntElementsAttrindex elements attribute

Results: 

ResultDescription
resultshape

shape.const_size (shape::ConstSizeOp) 

Creates a constant of type shape.size

Syntax:

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

Creates a shape.size type representing the constant size given by value.

%x = shape.const_size 10

Attributes: 

AttributeMLIR TypeDescription
valueIntegerAttrindex attribute

Results: 

ResultDescription
resultsize

shape.cstr_broadcastable (shape::CstrBroadcastableOp) 

Determines if 2 shapes can be successfully broadcasted.

Syntax:

operation ::= `shape.cstr_broadcastable` $lhs `,` $rhs attr-dict

Given 2 input shapes, 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] // Success
  %w1 = shape.cstr_broadcastable [2,2], [3,2] // Failure

Operands: 

OperandDescription
lhsshape
rhsshape

Results: 

ResultDescription
resultwitness

shape.cstr_eq (shape::CstrEqOp) 

Determines if all input shapes are equal.

Syntax:

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

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] // Success
  %w1 = shape.cstr_eq [2,2], [1,2] // Failure

Operands: 

OperandDescription
inputsshape

Results: 

ResultDescription
resultwitness

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: 

OperandDescription
inputshape or size

Results: 

ResultDescription
outputshape or size

shape.from_extent_tensor (shape::FromExtentTensorOp) 

Creates a shape from a tensor of extents

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.

Operands: 

OperandDescription
inputtensor of index values

Results: 

ResultDescription
resultshape

shape.from_extents (shape::FromExtentsOp) 

Creates a shape from extents

Syntax:

operation ::= `shape.from_extents` attr-dict $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

Operands: 

OperandDescription
extentsindex

Results: 

ResultDescription
shapeshape

shape.get_extent (shape::GetExtentOp) 

Gets the specified extent from a shape

Syntax:

operation ::= `shape.get_extent` $shape `,` $dim attr-dict

Gets the extent indexed by dim from shape.

If the shape is an error, it returns an error size.

Attributes: 

AttributeMLIR TypeDescription
dimIntegerAttr64-bit signless integer attribute whose value is non-negative

Operands: 

OperandDescription
shapeshape

Results: 

ResultDescription
extentsize

shape.index_to_size (shape::IndexToSizeOp) 

Converts a standard index to a shape size

Syntax:

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

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.

Operands: 

OperandDescription
argindex

Results: 

ResultDescription
resultsize

shape.join (shape::JoinOp) 

Returns the least general shape.size of its operands

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

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

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

%c = shape.join %a, %b, error="<reason>" : !shape.shape

Attributes: 

AttributeMLIR TypeDescription
errorStringAttrstring attribute

Operands: 

OperandDescription
arg0shape or size
arg1shape or size

Results: 

ResultDescription
resultshape or size

shape.mul (shape::MulOp) 

Multiplication of sizes

Multiplies two valid sizes as follows:

  • lhs * rhs = unknown if either lhs or rhs unknown;
  • lhs * rhs = (int)lhs * (int)rhs if both known;

Operands: 

OperandDescription
lhssize
rhssize

Results: 

ResultDescription
resultsize

shape.num_elements (shape::NumElementsOp) 

Returns the number of elements for a given shape

Syntax:

operation ::= `shape.num_elements` attr-dict $shape

Returns the number of elements for a given shape which is the product of its dimensions.

Operands: 

OperandDescription
shapeshape

Results: 

ResultDescription
resultsize

shape.reduce (shape::ReduceOp) 

Returns an expression reduced over a shape

An operation that takes as input a shape, number of initial values and has a region/function that is applied repeatedly for every dimension of the shape.

Conceptually this op performs the following reduction:

res[] = init;
for (int i = 0, e = shape.rank(); i != e; ++i) {
  res = fn(i, shape[i], res[0], ..., res[n]);
}

Where fn is provided by the user and the result of the reduce op is the last computed output of the reduce function. As an example, computing the number of elements

func @shape_num_elements(%shape : !shape.shape) -> !shape.size {
  %0 = "shape.constant_dim"() {value = 1 : i32} : () -> !shape.size
  %1 = "shape.reduce"(%shape, %0) ( {
    ^bb0(%index: i32, %dim: !shape.size, %lci: !shape.size):
      %acc = "shape.mul"(%lci, %dim) :
        (!shape.size, !shape.size) -> !shape.size
      shape.yield %acc : !shape.size
    }) : (!shape.shape, !shape.size) -> (!shape.size)
  return %1 : !shape.size
}

If the shape is unranked, then the results of the op is also unranked.

Operands: 

OperandDescription
shapeshape
argsany type

Results: 

ResultDescription
resultany type

shape.shape_of (shape::ShapeOfOp) 

Returns shape of a value or shaped type operand

Syntax:

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

Operands: 

OperandDescription
argshaped of any type values or value shape

Results: 

ResultDescription
resultshape

shape.size_to_index (shape::SizeToIndexOp) 

Casts between index types of the shape and standard dialect

Syntax:

operation ::= `shape.size_to_index` attr-dict $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.

Operands: 

OperandDescription
argsize

Results: 

ResultDescription
resultindex

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.

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)]

Operands: 

OperandDescription
operandshape
index32-bit signless integer

Results: 

ResultDescription
headshape
tailshape

shape.to_extent_tensor (shape::ToExtentTensorOp) 

Creates a dimension tensor from a shape

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, then this op currently aborts the program.

Operands: 

OperandDescription
inputshape

Results: 

ResultDescription
resulttensor of index values

shape.yield (shape::YieldOp) 

Returns the value to parent op

Syntax:

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

Operands: 

OperandDescription
operandsany type