'tensor' Dialect

The tensor dialect is intended to hold core tensor creation and manipulation ops, which are not strongly associated with any particular other dialect or domain abstraction. The primary smoke test of this is ops that make sense for any tensor element type.

We leave it to other dialects to hold the vast swath of possible computations one might want to do on a tensor.

The tensor type is (for better or for worse) used to represent all kinds of things, and supports an open-ended set of element types. Examples:

• representing large, dense aggregations of primitive types, suitable for high-performance numerical computing.
• representing shapes in the shape dialect, which consist of small 1D tensors of index data type.
• representing aggregations of strings or “variant” types.
• representing large, sparse aggregations of primitive types, suitable for high-performance numerical computing.

Thus, for the tensor dialect, we prefer for now to constrain the scope as much as possible. The expectation is that at some point in the future, the tensor dialect’s scope may be broadened through a careful discussion of the tradeoffs.

The tensor type is actually a builtin type (it lives in the builtin dialect), and does not live in this dialect.

Operation definition ¶

tensor.cast (::mlir::tensor::CastOp) ¶

tensor cast operation

Syntax:

operation ::= `tensor.cast` $source attr-dict `:` type($source) `to` type($dest)

Convert a tensor from one type to an equivalent type without changing any data elements. The source and destination types must both be tensor types with the same element type. If both are ranked, then the rank should be the same and static dimensions should match. The operation is invalid if converting to a mismatching constant dimension.

Example:

// Convert from unknown rank to rank 2 with unknown dimension sizes.
%2 = tensor.cast %1 : tensor<*xf32> to tensor<?x?xf32>

// Convert to a type with more known dimensions.
%3 = tensor.cast %2 : tensor<?x?xf32> to tensor<4x?xf32>

// Discard static dimension and rank information.
%4 = tensor.cast %3 : tensor<4x?xf32> to tensor<?x?xf32>
%5 = tensor.cast %4 : tensor<?x?xf32> to tensor<*xf32>

Operands: ¶

Results: ¶

tensor.extract (::mlir::tensor::ExtractOp) ¶

element extraction operation

Syntax:

operation ::= `tensor.extract` $tensor `[` $indices `]` attr-dict `:` type($tensor)

The tensor.extract op reads a tensor and returns one element from it specified by an index list. The output of the op is a new value with the same type as the elements of the tensor. The arity of indices must match the rank of the accessed value (i.e., if a tensor is of rank 3, then 3 indices are required for the extract. The indices should all be of index type.

Example:

%4 = tensor.extract %t[%1, %2] : tensor<4x4xi32>
%5 = tensor.extract %rt[%1, %2] : tensor<?x?xi32>
%6 = tensor.extract %ut[%1, %2] : tensor<*xi32>

Operands: ¶

Results: ¶

tensor.from_elements (::mlir::tensor::FromElementsOp) ¶

tensor from elements operation.

Syntax:

operation ::= `tensor.from_elements` $elements attr-dict `:` type($result)

Create a 1D tensor from a range of same-type arguments.

Example:

tensor.from_elements(i_1, ..., i_N) :  tensor<Nxindex>

Operands: ¶

Results: ¶

tensor.generate (::mlir::tensor::GenerateOp) ¶

Creates a dynamically sized tensor from elements

Syntax:

operation ::= `tensor.generate` $dynamicExtents $body attr-dict `:` type($result)

This operation creates a dynamically sized tensor with elements of any type. It expects one index operand per dynamic extent of the result tensor.

The body region defines the tensor’s elements. It takes index operands as its region arguments that span the index space. The element at the given position is yielded with the yield operation (see YieldOp). There is no defined ordering to the invocations of the body. It is conceptually a “parallel map” operation.

Example:

  %tnsr = tensor.generate %m, %n {
  ^bb0(%i : index, %j : index, %k : index):
    ...
    yield %elem : f32
  } : tensor<?x3x?f32>

Operands: ¶

Results: ¶

tensor.yield (::mlir::tensor::YieldOp) ¶

Yield a value from a region

Syntax:

operation ::= `tensor.yield` $value attr-dict `:` type($value)

This operation is used to yield a single value from a within a region. It is used to create dynamically sized tensors (see tensor.generate op).

Operands: ¶