MLIR

Multi-Level IR Compiler Framework

Dialect 'shape' definition

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 chosing 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

dim

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.

Operation definition

shape.add (shape::AddOp)

Addition of sizes

Description:

Adds two valid sizes as follows:

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

Operands:

  1. lhs: shape
  2. rhs: shape

Attributes:

Results:

  1. result: shape

shape.broadcast (shape::BroadcastOp)

Returns the broadcasted output shape of two inputs

Description:

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.

Operands:

  1. lhs: shape
  2. rhs: shape

Attributes:

AttributeMLIR TypeDescription
errorStringAttrstring attribute attribute

Results:

  1. result: shape

shape.constant (shape::ConstantOp)

Creates a shape constant

Description:

An operation that builds a size or shape from integer or array attribute. It allows for creating dynamically valued shapes by using ? for unknown values. A constant shape specified with * will return an unranked shape.

%x = shape.constant 10 : !shape.size

Operands:

Attributes:

AttributeMLIR TypeDescription
valueArrayAttrarray attribute attribute

Results:

  1. result: shape or size

shape.create_shape (shape::CreateShapeOp)

Creates a shape descriptor from a tensor

Description:

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

Operands:

  1. input: tensor of 32-bit integer values

Attributes:

Results:

  1. result: shape

shape.debug_print (shape::DebugPrintOp)

Prints the input shape or size

Description:

Prints the input dim or shape and passes through input.

Note: This is intended for testing and debugging only.

Operands:

  1. input: shape or size

Attributes:

Results:

  1. output: shape or size

shape.join (shape::JoinOp)

Returns the least general shape.size of its operands

Description:

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

Operands:

  1. arg0: shape or size
  2. arg1: shape or size

Attributes:

AttributeMLIR TypeDescription
errorStringAttrstring attribute attribute

Results:

  1. result: shape or size

shape.mul (shape::MulOp)

Multiplication of sizes

Description:

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:

  1. lhs: shape
  2. rhs: shape

Attributes:

Results:

  1. result: shape

shape.reduce (shape::ReduceOp)

Returns an expression reduced over a shape

Description:

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.type) -> !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.return"(%acc) : (!shape.size) -> ()
    }) : (!shape.type, !shape.size) -> (!shape.size)
  return %1 : !shape.size
}

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

Operands:

  1. shape: shape
  2. args: any type

Attributes:

Results:

  1. result: any type

shape.shape_of (shape::ShapeOfOp)

Returns shape of a value or shaped type operand

Description:

Operands:

  1. arg: shaped of any type values or value shape

Attributes:

Results:

  1. result: shape