'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: ¶

Results: ¶

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: ¶

Results: ¶

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: ¶

Results: ¶

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: ¶

Results: ¶

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: ¶

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: ¶

Operands: ¶

Results: ¶

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: ¶

Results: ¶

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: ¶

Results: ¶

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: ¶

Results: ¶

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: ¶

Results: ¶

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: ¶

Results: ¶

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: ¶

Results: ¶

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: ¶

Results: ¶

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: ¶

Results: ¶

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: ¶

Operands: ¶

Results: ¶

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: ¶

Results: ¶

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: ¶

Operands: ¶

Results: ¶

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: ¶

Results: ¶

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: ¶

Results: ¶

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: ¶

Results: ¶

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: ¶

Results: ¶

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: ¶

Results: ¶

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: ¶

Results: ¶

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: ¶

Results: ¶

shape.yield (shape::YieldOp) ¶

Returns the value to parent op

Syntax:

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

Operands: ¶