mlir.dialects.linalg.opdsl.lang.config

Represents configured ops as emitted for code generation.

Classes in this module generally are directly serializable to YAML for use by the code generator.

TODO: These should just be dumb containers or serialization code but they currently encode too many details of how the language is interpreted. Move this to helpers on the comprehension objects themselves.

Classes

OperandDefConfig	Wrapper containing an operand definition with additional state.
LinalgStructuredOpConfig	Configuration for metadata sufficient to construct a linalg named op.
LinalgOpConfig	Container for any supported linalg op type.

Module Contents

class mlir.dialects.linalg.opdsl.lang.config.OperandDefConfig(operand_def: mlir.dialects.linalg.opdsl.lang.comprehension.OperandDef, shape_map: mlir.dialects.linalg.opdsl.lang.comprehension.Optional[mlir.ir.AffineMap] = None, index_attr_map: mlir.dialects.linalg.opdsl.lang.comprehension.Optional[mlir.ir.AffineMap] = None)

Bases: mlir.dialects.linalg.opdsl.lang.yaml_helper.YAMLObject

Wrapper containing an operand definition with additional state.

yaml_tag = '!LinalgOperandDefConfig'

operand_def

shape_map: mlir.dialects.linalg.opdsl.lang.comprehension.Optional[mlir.ir.AffineMap] = None

index_attr_map: mlir.dialects.linalg.opdsl.lang.comprehension.Optional[mlir.ir.AffineMap] = None

indexing_map: mlir.dialects.linalg.opdsl.lang.comprehension.Optional[mlir.ir.AffineMap] = None

property name: str

property kind: mlir.dialects.linalg.opdsl.lang.comprehension.OperandKind

property type_var: mlir.dialects.linalg.opdsl.lang.comprehension.TypeVar

to_yaml_custom_dict()

__repr__()

class mlir.dialects.linalg.opdsl.lang.config.LinalgStructuredOpConfig(comprehension: mlir.dialects.linalg.opdsl.lang.comprehension.Comprehension, domain: mlir.dialects.linalg.opdsl.lang.comprehension.Sequence[mlir.dialects.linalg.opdsl.lang.comprehension.DimDef], registered_operands: mlir.dialects.linalg.opdsl.lang.comprehension.Sequence[mlir.dialects.linalg.opdsl.lang.comprehension.OperandDef], context: mlir.dialects.linalg.opdsl.lang.comprehension.Optional[mlir.ir.Context] = None)

Bases: mlir.dialects.linalg.opdsl.lang.yaml_helper.YAMLObject

Configuration for metadata sufficient to construct a linalg named op.

yaml_tag = '!LinalgStructuredOpConfig'

context = None

affine_state

writes: mlir.dialects.linalg.opdsl.lang.comprehension.List[mlir.dialects.linalg.opdsl.lang.comprehension.Tuple[mlir.dialects.linalg.opdsl.lang.comprehension.TensorUse, mlir.dialects.linalg.opdsl.lang.comprehension.TensorExpression]] = []

operands: mlir.dialects.linalg.opdsl.lang.comprehension.Dict[mlir.dialects.linalg.opdsl.lang.comprehension.OperandDef, OperandDefConfig]

uses: mlir.dialects.linalg.opdsl.lang.comprehension.Dict[mlir.dialects.linalg.opdsl.lang.comprehension.TensorUse, TensorUseConfig]

reduction_dims

assignments

property ordered_operands: mlir.dialects.linalg.opdsl.lang.comprehension.Sequence[OperandDefConfig]

property ordered_dims: mlir.dialects.linalg.opdsl.lang.comprehension.Sequence[mlir.dialects.linalg.opdsl.lang.comprehension.Tuple[str, int]]

Gets the ordered list of dim bindings (symbolic name, position).

TODO: The original parser relies on parse ordering to arrive at the iterator types, but that ordering is not defined on the Python side, so this may be ambiguous.

property indexing_maps: mlir.dialects.linalg.opdsl.lang.comprehension.Sequence[mlir.ir.AffineMap]

property iterator_types: mlir.dialects.linalg.opdsl.lang.comprehension.Sequence[str]

add_operand(operand_def: mlir.dialects.linalg.opdsl.lang.comprehension.OperandDef)

add_indexed_operand(operand_def: mlir.dialects.linalg.opdsl.lang.comprehension.OperandDef)

add_tensor_use(tensor_use: mlir.dialects.linalg.opdsl.lang.comprehension.TensorUse)

_get_scalar_map() → mlir.ir.AffineMap

Create an empty affine map used to index a scalar.

_normalize_affine_map(affine_map: mlir.ir.AffineMap, with_dims: bool = True) → mlir.ir.AffineMap

Normalizes an indexing map to have the max known symbols and dims.

to_yaml_custom_dict()

__repr__()

class mlir.dialects.linalg.opdsl.lang.config.LinalgOpConfig(metadata: mlir.dialects.linalg.opdsl.lang.comprehension.OpMetadataDef, *, structured_op: mlir.dialects.linalg.opdsl.lang.comprehension.Optional[LinalgStructuredOpConfig] = None)

Bases: mlir.dialects.linalg.opdsl.lang.yaml_helper.YAMLObject

Container for any supported linalg op type.

This includes the concrete type by name for ease of parsing by systems that ignore tags.

yaml_tag = '!LinalgOpConfig'

metadata

structured_op = None

to_yaml_custom_dict()

static from_linalg_op_def(op_def: mlir.dialects.linalg.opdsl.lang.comprehension.LinalgOpDef, context: mlir.dialects.linalg.opdsl.lang.comprehension.Optional[mlir.ir.Context] = None) → mlir.dialects.linalg.opdsl.lang.comprehension.Sequence[LinalgOpConfig]

Expands a LinalgOpDef into corresponding Linalg configured ops.

__repr__()