iree icon indicating copy to clipboard operation
iree copied to clipboard
verified publisher icon iree-org

Refresh Encoding dialect doc

Open hanhanW opened this issue 7 months ago • 2 comments

We had a discussion about what encodings are and the better naming in https://github.com/iree-org/iree/pull/20700#discussion_r2072932076. This is a placeholder issue to put more context to it. As we have better understanding about encodings today, we should revisit the definition and improve the documentation. E.g., we definitely need to refresh the documentation for EncodingBase.td:

https://github.com/iree-org/iree/blob/f6a5d72bb0a1beff3d30741f9050b08141a71318/compiler/src/iree/compiler/Dialect/Encoding/IR/EncodingBase.td#L22-L25

IMO, there are two types of encodings. One is the encodings that will be attached to tensor types, and the other is the encoding resolvers. E.g., iree_encoding.encoding and iree_encoding.layout are tensor encoding attributes, iree_gpu.pad_layout_attr is an encoding resolver. I want to keep a clear line between these two categories.

I'll try to list down what encodings we have today and give them better definition and abstraction.

hanhanW avatar May 06 '25 20:05 hanhanW

Mirror the good conclusion we had from the recent PR. I'll expand the context and refresh the doc after we move LayoutAttrInterface to Encoding dialect.

  • EncodingLayoutResolverAttrInterface is an interface that help resolve verbose encodings into serialized/specialized encodings.
  • SerializableEncodingAttrInterface is an interface that describes encoding properties and should have enough information for later transformation.
  • LayoutAttrInterface is an interface that provides a set of interface method to materialize serialized encoding to physical operations/types/etc.

With the definition, there are three states for an encoding. (a) Verbose encoding -> (b) serialized encoding -> (c) physical operations and types.

EncodingLayoutResolverAttrInterface helps (a) -> (b). SerializableEncodingAttrInterface describes the properties for (b). LayoutAttrInterface helps (b) -> (c).

hanhanW avatar May 09 '25 23:05 hanhanW

Encoding Snapshot

Current Encodings

  • iree_encoding.layout: Describes the potential serialized layouts of an encoding type.
  • iree_encoding.packed_storage: Indicates that it is a back-to-back packed storage in memory.
  • iree_encoding.encoding: Generic encoding attribute for data-tiling.
  • iree_encoding.matmul_k: Describe matmul properties for reduction dimensions.
  • iree_encoding.pad_encoding_layout: Encodes padding values of tensor dimensions.
  • iree_encoding.identity_encoding: A resolver that resolves all the encodings into identity layout.
  • iree_encoding.unsupported_encoding: A resolver that always fails to get layout. (Testing purpose)
  • iree_encoding.testing_encoding: Similar for iree_encoding.layout, but it does not require all the layouts serialized. (Testing purpose)
  • iree_encoding.unknown_encoding: An encoding type attribute that we don't know what it is at all. (Testing purpose)
  • iree_encoding.unspecialized_encoding<seed>: A resolver that converts an encoding into specialize_encoding attribute with the same seed. (Testing purpose) Currently, it is misused as encoding type attribute in some roundtrip tests.
  • iree_encoding.specialized_encoding<seed>: An attribute that indicates the encoding type is specialized. (Testing purpose)
  • iree_cpu.cpu_encoding_layout: A resolver that lowers verboses encodings to serialized encodings. Typically, the format is not trivial, but the resolver itself should be able to interpret the format. It means that this resolver usually implements the SerializableEncodingAttrInterface as well.
  • iree_cpu.vmvx_encoding_layout: Similar to cpu, but it is for VMVX.
  • iree_gpu.gpu_encoding_layout: Similiar to cpu, but it is for GPU backend.
  • iree_gpu.gpu_pad_layout: A resolver that resolves an encoding type attribute to iree_encoding.pad_encoding_layout based on L1 cache information.

Encoding Attribute Category

Encoding Type Attribute

This kind of attribute is expected to attached on encoding types. We may see some tensor encoding terminology in IREE, but it is not accurate IMO. It is fine in IREE because IREE is a tensor-based compiler and encoding is only available for RankedTensorType today. However, people can define their own type system that has encodings. Thus, I think encoding type attribute is more accurate.

Encoding type attributes used in practice:

  • iree_encoding.layout
  • iree_encoding.packed_storage
  • iree_encoding.encoding
  • iree_encoding.matmul_k
  • iree_encoding.pad_encoding_layout

Encoding type attributes used in testing:

  • iree_encoding.testing_encoding
  • iree_encoding.unknown_encoding
  • iree_encoding.specialized_encoding<seed>

Encoding Resolver

This kind of attribute is used to resolve and transform encoding type attributes. There are three states of encoding type attribute:

  1. Verbose encoding.
  2. Serialized encoding
  3. Physical operations and types.

An encoding resolver implements some Encoding interfaces that transforms the encoding type attribute between the states:

  • EncodingLayoutResolverAttrInterface: an interface that help resolve verbose encodings into serialized/specialized encodings.
  • SerializableEncodingAttrInterface: an interface that describes encoding properties and should have enough information for later transformation.
  • LayoutAttrInterface: an interface that provides a set of interface method to materialize serialized encoding to physical operations/types/etc.

To be more specfic, EncodingLayoutResolverAttrInterface helps (a) -> (b). SerializableEncodingAttrInterface describes the properties for (b). LayoutAttrInterface helps (b) -> (c).

An encoding resolver must implement the EncodingLayoutResolverAttrInterface interface. An verbose encoding will be transformed by an encoding resolver into a serialized encoding that implements SerializableEncodingAttrInterface. It can be either

  • An encoding type attribute that is serialized, or
  • A custom serialized format that is wrapped within an encoding resolver; the encoding resolver can interpret the serialized format.

IREE performs encoding specialization at Stream phase. It is the phase to perform such transformation because Stream models explicitly scheduled asynchronous programs by partitioning the dispatchable work, specifying target affinities, encoding tensors into target-specific forms, and scheduling the work to run concurrently.

Rename

The attributes and interfaces are already defined within encoding namespace/dialect. IMO, there is a hidden convension that we prefer fully specify type name in IREE core dialects. Because it is easier to do large scale replacements, avoids conflicts with MLIR upstream types, and makes it easier to port passes between dialects. In this context, Encoding is redundant in attribute names and interface names.

I also want to propose that all the resolvers should have _resolver suffix in their name, which makes a clear line between encoding type attributes and resolvers.

The last point is that Interface suffix can be dropped, but it introduces a conflict between iree_encoding.layout encoding type attribute and the LayoutAttrInterface interface. I think we can rename the interface to LayoutMaterializerAttr.

Below is the names before and after the rename. (Starting with * means unchanged.)

Encoding Type Attribute:

  • *iree_encoding.layout
  • *iree_encoding.packed_storage
  • ?iree_encoding.encoding -> iree_encoding.encoding or iree_encoding.generic
  • *iree_encoding.matmul_k
  • iree_encoding.pad_encoding_layout -> iree_encoding.pad_physical_only or just iree_encoding.padding
  • iree_encoding.testing_encoding -> iree_encoding.testing
  • iree_encoding.unknown_encoding -> iree_encoding.unknown
  • iree_encoding.specialized_encoding<seed> -> iree_encoding.specialized<seed>

Encoding Resolvers:

  • iree_encoding.identity_encoding -> iree_encoding.identity_resolver
  • iree_encoding.unsupported_encoding -> iree_encoding.unsupported_resolver
  • iree_encoding.unspecialized_encoding<seed> -> iree_encoding.specialization_resolver<seed>
  • iree_cpu.cpu_encoding_layout -> iree_cpu.cpu_encoding_resolver
  • iree_cpu.vmvx_encoding_layout -> iree_cpu.vmvx_encoding_resolver
  • iree_gpu.gpu_encoding_layout -> iree_gpu.gpu_encoding_resolver
  • iree_gpu.gpu_pad_layout -> iree_gpu.gpu_padding_resolver

Interfaces:

  • EncodingLayoutResolverAttrInterface -> LayoutResolverAttr
  • SerializableEncodingAttrInterface -> SerializableAttr
  • LayoutAttrInterface -> LayoutMaterializerAttr
  • (Codegen) PackedLayoutAttrInterface -> PackedLayoutAttr or DataTilingLayoutAttr. This one is optional because Codegen seems to always have Interface suffix.

Completeness

IMO, we are missing some classic encoding type attributes in current codebase:

  • iree_encoding.identity: indicates that the encoded type is as the same as the type without encodings. Sometimes you can't just use the non-encoded types because types could mismatch. Today, we abuse iree_encoding.pad_encoding_layout with zero values to represent an identity encoding.
  • iree_encoding.unpacked_storage: we may use a larger container, e.g., the next power of two, if the element bit-width is odd. This is more like an inverse behavior of iree_encoding.packed_storage, and the frontend can decide what to use. It also isolates whether IREE should use packed or unpacked storage by default when the types do not have encodings. Because encoding type attributes do not really care what the IREE default behavior is.

cc @jtuyls @Max191 @benvanik @MaheshRavishankar

hanhanW avatar May 27 '25 14:05 hanhanW