'iree_encoding' Dialect

A dialect designed for tensor encoding attributes and ops.

The encoding dialect defines IREE tensor encoding attributes and related ops. These encoding hints enable optimization of compute-heavy operations like matmul. E.g., IREE implements data-tiling using the encoding dialect, by introducing encodings on matmul ops and fuse encoding ops with producers.

There are two core operations in this dialect, and they are very simple: - SetEncodingOp: Adds an encoding attribute to a tensor value. - UnsetEncodingOp: Removes the encoding attribute from a tensor value.

Encoding attributes fall into two categories: encoding type attribute and encoding resolver.

An encoding type attribute is expected to be attached on tensor types. They typically go through three states:

+------------------+ +---------------------+ +--------------------+ | Verbose Encoding | -> | Serialized Encoding | -> | Physical Ops/Types | +------------------+ +---------------------+ +--------------------+

Encoding resolvers are responsible for transforming the encodings between states. Each backend can define its own resolver to interpret encodings based on target-specific details.

To support encoding transformations better, IREE defines three encoding interfaces: - LayoutResolverAttr: Converts verbose encodings into serialized ones, based on target-specific information. - SerializableAttr: Represents serialized encoding formats with enough detail for host-side compilation. - LayoutMaterializerAttr: Lowers encodings into physical operations and types.

An encoding resolver must implement LayoutResolverAttr, as it needs to convert verbose encodings into serialized encodings with target specifics.

Sometimes, an encoding resolver also implements SerializableAttr because the layout can be customized format, e.g., packing, etc. Only the resolver itself knows how to interpret the layout.

An encoding resolver must implement LayoutMaterializerAttr, because only the resolver itself knows what the final operations and types are.

Encoding resolvers are conventionally named with a _resolver suffix.

The transition from verbose to serialized encoding happens during the SpecializeEncoding pass within the stream dialect. At stream level, the pass has access to all necessary context (e.g., affinity, etc.) to retrieve encoding resolvers. It bridges the gap between host and device compilation, as it allows IREE to specialize storage size based on the serialized encoding result.

• 'iree_encoding' Dialect
  • Operations
    • iree_encoding.set_encoding (Encoding::SetEncodingOp)
    • iree_encoding.unset_encoding (Encoding::UnsetEncodingOp)

Operations

iree_encoding.set_encoding (Encoding::SetEncodingOp)

Perform pack and pad operation on source.

Syntax:

operation ::= `iree_encoding.set_encoding` attr-dict $source `:` type($source) `->` type($result)

Operation to assign an encoding to a tensor. The operation does not change the rank or extent of a tensor. Instead it adds a LayoutResolverAttr attribute to the tensor type to represent a change in layout.

Traits: AlwaysSpeculatableImplTrait

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface), ReifyRankedShapedTypeOpInterface

Effects: MemoryEffects::Effect{}

Operands:

Operand	Description
source	ranked tensor of any type values

Results:

Result	Description
result	ranked tensor of any type values

iree_encoding.unset_encoding (Encoding::UnsetEncodingOp)

Perform unpack and extract operation on source.

Syntax:

operation ::= `iree_encoding.unset_encoding` attr-dict $source `:` type($source) `->` type($result) (`` `{` $result_dims^ `}`)?

Operation to convert a tensor with LayoutResolverAttr encoding that represents its data layout into a tensor with default layout (i.e. no encoding). For now in IREE the default layout is row-major.

Traits: AlwaysSpeculatableImplTrait

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface), ReifyRankedShapedTypeOpInterface

Effects: MemoryEffects::Effect{}

Operands:

Operand	Description
source	ranked tensor of any type values
result_dims	variadic of index

Results:

Result	Description
result	ranked tensor of any type values