mlir.dialects.linalg.opdsl.lang¶

Submodules¶

Attributes¶

`_CONTEXT`
`StructuredOpOuts`
`ContractionOpInterface`
`ConvolutionOpInterface`
`FillOpInterface`
`Canonicalizer`
`D`
`S`
`TV`
`I32`
`I64`
`F32`
`F64`
`T`
`U`
`V`
`ValueList`

Classes¶

`DefinedOpCallable`	Callable that wraps any defined op function.
`TensorExpression`	An expression that can appear on the RHS of a comprehension.
`TensorUse`	A used tensor represented by its (tensor_name, indices).
`TensorFn`	Application of a tensor function.
`TensorReduceFn`	Application of a reduction function.
`const`	Returns the given constant floating point or integer value.
`index`	Returns the iteration index for a given dimension name.
`FunctionKind`	Generic enumeration.
`UnaryFnType`	Unary function.
`UnaryFn`	Unary function namespace.
`BinaryFnType`	Binary function.
`BinaryFn`	Binary function namespace.
`TernaryFnType`	Ternary function.
`TernaryFn`	Ternary function namespace.
`TypeFnType`	Type conversion function.
`TypeFn`	Type conversion function namespace.
`ReduceFnUse`	Reduction function use.
`ReduceFnType`	Reduction function.
`ReduceFn`
`OperandKind`	Generic enumeration.
`OperandDef`	Definition of an operand passed to an operation.
`TensorDef`	Tensor operand definition.
`ScalarDef`	Scalar operand definition.
`IndexAttrDef`	Index attribute definition.
`UnaryFnAttrDef`	Unary function attribute definition.
`BinaryFnAttrDef`	Binary function attribute definition.
`TernaryFnAttrDef`	Ternary function attribute definition.
`TypeFnAttrDef`	Type conversion function attribute definition.
`Comprehension`	Represents a single comprehension.
`OpInterfaceDef`	An interface that an op implements.
`OpDefinitionDef`	A method that an op implements.
`OpMetadataDef`	Metadata about the op (generally not behavior impacting).
`LinalgOpDef`	Definition of a linalg op.
`AffineBuildState`	Internal state for the AffineExprDef._create impls.
`AffineExprDef`	Base class for an affine expression being defined.
`DimDef`	Represents a named dimension.
`SymbolDef`	Represents a named symbol.
`ScalarAssign`	An assignment to a named argument (LHS of a comprehension).
`ScalarFn`	A type of ScalarExpression that applies a function.
`ScalarArg`	A type of ScalarExpression that references a named argument.
`ScalarConst`	A type of ScalarExpression representing a constant.
`ScalarIndex`	A type of ScalarExpression accessing an iteration index.
`ScalarExpression`	An expression on scalar values.
`TypeVar`	A replaceable type variable.
`YAMLObject`	An object that can dump itself to a YAML stream
`LinalgStructuredOpConfig`	Configuration for metadata sufficient to construct a linalg named op.
`LinalgOpConfig`	Container for any supported linalg op type.
`OperandDefConfig`	Wrapper containing an operand definition with additional state.

Functions¶

`_get_op_result_or_value`(→ mlir._mlir_libs._mlir.ir.Value)	Returns the given value or the single result of the given op.
`_get_op_results_or_values`(...)	Returns the given sequence of values or the results of the given op.
`bind_op_def`(op_def)
`current_op_def`(...)
`_prepare_structured_op_outs`(...)
`linalg_structured_op`(→ DefinedOpCallable)
`domain`(*dimensions)
`implements`(*interfaces)
`defines`(*definitions)
`yaml_dump`(data[, sort_keys])
`yaml_dump_all`(data[, sort_keys, explicit_start])
`emit_generic_structured_op`(op_config, ins, outs, *attrs)
`emit_named_structured_op`(op_config, op_name, ...)

Package Contents¶

mlir.dialects.linalg.opdsl.lang._get_op_result_or_value(arg: mlir._mlir_libs._mlir.ir.OpView | mlir._mlir_libs._mlir.ir.Operation | mlir._mlir_libs._mlir.ir.Value | mlir._mlir_libs._mlir.ir.OpResultList) → mlir._mlir_libs._mlir.ir.Value¶

Returns the given value or the single result of the given op.

This is useful to implement op constructors so that they can take other ops as arguments instead of requiring the caller to extract results for every op. Raises ValueError if provided with an op that doesn’t have a single result.

mlir.dialects.linalg.opdsl.lang._get_op_results_or_values(arg: mlir._mlir_libs._mlir.ir.OpView | mlir._mlir_libs._mlir.ir.Operation | Sequence[mlir._mlir_libs._mlir.ir.OpView | mlir._mlir_libs._mlir.ir.Operation | mlir._mlir_libs._mlir.ir.Value]) → Sequence[mlir._mlir_libs._mlir.ir.OpView | mlir._mlir_libs._mlir.ir.Operation | mlir._mlir_libs._mlir.ir.Value] | mlir._mlir_libs._mlir.ir.OpResultList¶

Returns the given sequence of values or the results of the given op.

This is useful to implement op constructors so that they can take other ops as lists of arguments instead of requiring the caller to extract results for every op.

mlir.dialects.linalg.opdsl.lang._CONTEXT¶

mlir.dialects.linalg.opdsl.lang.StructuredOpOuts¶

mlir.dialects.linalg.opdsl.lang.bind_op_def(op_def: mlir.dialects.linalg.opdsl.lang.emitter.LinalgOpDef)¶

mlir.dialects.linalg.opdsl.lang.current_op_def() → mlir.dialects.linalg.opdsl.lang.emitter.LinalgOpDef¶

mlir.dialects.linalg.opdsl.lang._prepare_structured_op_outs(outs: StructuredOpOuts) → mlir.dialects.linalg.opdsl.lang.emitter.ValueList¶

class mlir.dialects.linalg.opdsl.lang.DefinedOpCallable(op_name: str, op_def: mlir.dialects.linalg.opdsl.lang.emitter.LinalgOpDef)¶

Callable that wraps any defined op function.

op_name¶

op_def¶

__call__(*ins: mlir.dialects.linalg.opdsl.lang.emitter.Union[mlir.ir.Operation, mlir.ir.OpView, mlir.ir.Value], outs: StructuredOpOuts, **kwargs)¶

Emits the corresponding op definition as IR.

Most arguments are passed through to the underlying emitter. The following keyword argument is interpreted here: emit_generic: Emits a generic form as appropriate (default True). If False, a named form is emitted (which must have been built in to the compiler).

mlir.dialects.linalg.opdsl.lang.linalg_structured_op(dsl_func=None, *, op_name=None, op_class_name=None) → DefinedOpCallable¶

mlir.dialects.linalg.opdsl.lang.domain(*dimensions: mlir.dialects.linalg.opdsl.lang.emitter.DimDef)¶

mlir.dialects.linalg.opdsl.lang.implements(*interfaces: mlir.dialects.linalg.opdsl.lang.emitter.OpInterfaceDef)¶

mlir.dialects.linalg.opdsl.lang.defines(*definitions: mlir.dialects.linalg.opdsl.lang.emitter.OpDefinitionDef)¶

class mlir.dialects.linalg.opdsl.lang.TensorExpression¶

An expression that can appear on the RHS of a comprehension.

abstract to_scalar_expression() → mlir.dialects.linalg.opdsl.lang.scalar_expr.ScalarExpression¶

visit_tensor_exprs(callback: mlir.dialects.linalg.opdsl.lang.scalar_expr.Callable[[TensorExpression], None])¶: Visits all tensor expression reachable by the expression.

collect_dim_uses(uses: mlir.dialects.linalg.opdsl.lang.scalar_expr.Set[mlir.dialects.linalg.opdsl.lang.scalar_expr.DimDef])¶: Collects all DimDefs reachable through this expression.

collect_tensor_uses(uses: mlir.dialects.linalg.opdsl.lang.scalar_expr.Set[TensorUse])¶: Collects all TensorUses reachable through this expression.

collect_indices(indices: mlir.dialects.linalg.opdsl.lang.scalar_expr.Set[index])¶: Collects all index accesses reachable through this expression.

collect_scalar_uses(uses: mlir.dialects.linalg.opdsl.lang.scalar_expr.Set[ScalarDef])¶: Collects all ScalarDefs reachable through this expression.

__add__(rhs: TensorExpression) → TensorExpression¶

__mul__(rhs) → TensorExpression¶

__sub__(rhs) → TensorExpression¶

__truediv__(rhs) → TensorExpression¶

__hash__()¶

class mlir.dialects.linalg.opdsl.lang.TensorUse(operand_def: OperandDef, indices: mlir.dialects.linalg.opdsl.lang.scalar_expr.Sequence[mlir.dialects.linalg.opdsl.lang.scalar_expr.AffineExprDef])¶

Bases: TensorExpression

A used tensor represented by its (tensor_name, indices).

Note that forming a comprehension via direct assignment is performed through setitem on the TensorDef level. However, performing a reduction with compound ops (+=, =, etc) is done by doing a: TensorDef.**getitem* TensorUse.**iadd** TensorDef.**setitem**

operand_def¶

indices¶

to_scalar_expression() → mlir.dialects.linalg.opdsl.lang.scalar_expr.ScalarExpression¶

property tensor_name: str¶

_compute_reduce_dims(rhs: TensorExpression) → mlir.dialects.linalg.opdsl.lang.scalar_expr.Set[mlir.dialects.linalg.opdsl.lang.scalar_expr.DimDef]¶

__iadd__(rhs: TensorExpression) → TensorReduceFn¶

__repr__()¶

class mlir.dialects.linalg.opdsl.lang.TensorFn(kind: FunctionKind, name: mlir.dialects.linalg.opdsl.lang.scalar_expr.Optional[str], operand_def: mlir.dialects.linalg.opdsl.lang.scalar_expr.Optional[OperandDef], type_var: mlir.dialects.linalg.opdsl.lang.scalar_expr.Optional[mlir.dialects.linalg.opdsl.lang.types.TypeVar], args: mlir.dialects.linalg.opdsl.lang.scalar_expr.Sequence[TensorExpression])¶

Bases: TensorExpression

Application of a tensor function.

name¶

kind¶

operand_def¶

type_var¶

args¶

to_scalar_expression() → mlir.dialects.linalg.opdsl.lang.scalar_expr.ScalarExpression¶

visit_tensor_exprs(callback: mlir.dialects.linalg.opdsl.lang.scalar_expr.Callable[[TensorExpression], None])¶: Visits all tensor expression reachable by the expression.

__repr__()¶

class mlir.dialects.linalg.opdsl.lang.TensorReduceFn(reduce_use: ReduceFnUse, args: mlir.dialects.linalg.opdsl.lang.scalar_expr.Sequence[TensorExpression])¶

Bases: TensorExpression

Application of a reduction function.

This captures the lhs (initial value) separately from the rhs.

reduce_use¶

lhs: mlir.dialects.linalg.opdsl.lang.scalar_expr.Optional[TensorUse] = None¶

args¶

to_scalar_expression() → mlir.dialects.linalg.opdsl.lang.scalar_expr.ScalarExpression¶

visit_tensor_exprs(callback: mlir.dialects.linalg.opdsl.lang.scalar_expr.Callable[[TensorExpression], None])¶: Visits all tensor expression reachable by the expression.

__repr__()¶

class mlir.dialects.linalg.opdsl.lang.const(value: mlir.dialects.linalg.opdsl.lang.scalar_expr.Any)¶

Bases: TensorExpression

Returns the given constant floating point or integer value.

to_scalar_expression() → mlir.dialects.linalg.opdsl.lang.scalar_expr.ScalarExpression¶

__repr__()¶

class mlir.dialects.linalg.opdsl.lang.index(dim: mlir.dialects.linalg.opdsl.lang.scalar_expr.DimDef)¶

Bases: TensorExpression

Returns the iteration index for a given dimension name.

Resolves the given dimension name to obtain its position in the iteration domain of the operation.

dim_def¶

dim = -1¶

resolve_dimension_name(affine_state: mlir.dialects.linalg.opdsl.lang.scalar_expr.AffineBuildState)¶

to_scalar_expression() → mlir.dialects.linalg.opdsl.lang.scalar_expr.ScalarExpression¶

__repr__()¶

class mlir.dialects.linalg.opdsl.lang.FunctionKind¶

Bases: mlir.dialects.linalg.opdsl.lang.types.Enum

Generic enumeration.

Derive from this class to define new enumerations.

UNARY = 0¶

BINARY = 1¶

TERNARY = 2¶

TYPE = 3¶

class mlir.dialects.linalg.opdsl.lang.UnaryFnType(fn_name: str)¶

Unary function.

A unary function takes one tensor expression and returns the function evaluation result.

fn_name¶

__call__(arg: TensorExpression) → TensorFn¶

__repr__()¶

class mlir.dialects.linalg.opdsl.lang.UnaryFn¶

Unary function namespace.

exp¶

log¶

abs¶

ceil¶

floor¶

negf¶

reciprocal¶

round¶

sqrt¶

rsqrt¶

square¶

tanh¶

erf¶

class mlir.dialects.linalg.opdsl.lang.BinaryFnType(fn_name: str)¶

Binary function.

A binary function takes two tensor expressions and returns the function evaluation result.

fn_name¶

__call__(arg0: TensorExpression, arg1: TensorExpression) → TensorFn¶

__repr__()¶

class mlir.dialects.linalg.opdsl.lang.BinaryFn¶

Binary function namespace.

As the integer types are signless, signedness is implement by different functions that treat integers as signed or unsigned values.

Examples:

max -> arith.MaxSIOp
max_unsigned -> arith.MaxUIOp

add¶

sub¶

mul¶

div¶

div_unsigned¶

max_signed¶

min_signed¶

max_unsigned¶

min_unsigned¶

powf¶

class mlir.dialects.linalg.opdsl.lang.TernaryFnType(fn_name: str)¶

Ternary function.

A ternary function takes three tensor expressions and returns the function evaluation result.

fn_name¶

__call__(arg0: TensorExpression, arg1: TensorExpression, arg2: TensorExpression) → TensorFn¶

__repr__()¶

class mlir.dialects.linalg.opdsl.lang.TernaryFn¶

Ternary function namespace.

select¶

class mlir.dialects.linalg.opdsl.lang.TypeFnType(fn_name: str)¶

Type conversion function.

A type conversion function takes a target type and a tensor expression and returns the casted tensor expression.

fn_name¶

__call__(type_var: mlir.dialects.linalg.opdsl.lang.types.TypeVar, arg: TensorExpression) → TensorFn¶

__repr__()¶

class mlir.dialects.linalg.opdsl.lang.TypeFn¶

Type conversion function namespace.

As the integer types are signless, signedness is implement by different cast functions that treat integers as signed (cast_signed) or unsigned (cast_unsigned) values.

Examples:

cast_signed(I32 -> I64) -> arith.ExtSIOp
cast_unsigned(I32 -> I64) -> arith.ExtUIOp

cast_signed¶

cast_unsigned¶

class mlir.dialects.linalg.opdsl.lang.ReduceFnUse(binary_fn: mlir.dialects.linalg.opdsl.lang.scalar_expr.Optional[BinaryFnType], binary_attr: mlir.dialects.linalg.opdsl.lang.scalar_expr.Optional[BinaryFnAttrDef], *reduce_dims: mlir.dialects.linalg.opdsl.lang.scalar_expr.DimDef)¶

Reduction function use.

A reduction use specifies the reduction function and dimensions.

binary_fn¶

binary_attr¶

reduce_dims = ()¶

__call__(*args: TensorExpression) → TensorReduceFn¶

__repr__()¶

class mlir.dialects.linalg.opdsl.lang.ReduceFnType(binary_fn: BinaryFnType)¶

Reduction function.

A binary function that reduces its RHS into its LHS.

binary_fn¶

__getitem__(reduce_dims: mlir.dialects.linalg.opdsl.lang.scalar_expr.Tuple[mlir.dialects.linalg.opdsl.lang.scalar_expr.DimDef]) → ReduceFnUse¶

__repr__()¶

class mlir.dialects.linalg.opdsl.lang.ReduceFn¶

add¶

mul¶

max_signed¶

min_signed¶

max_unsigned¶

min_unsigned¶

class mlir.dialects.linalg.opdsl.lang.OperandKind¶

Bases: mlir.dialects.linalg.opdsl.lang.types.Enum

Generic enumeration.

Derive from this class to define new enumerations.

INPUT_TENSOR = 0¶

SCALAR = 1¶

OUTPUT_TENSOR = 2¶

INDEX_ATTR = 3¶

UNARY_FN_ATTR = 4¶

BINARY_FN_ATTR = 5¶

TERNARY_FN_ATTR = 6¶

TYPE_FN_ATTR = 7¶

class mlir.dialects.linalg.opdsl.lang.OperandDef(kind: OperandKind, type_var: mlir.dialects.linalg.opdsl.lang.scalar_expr.Optional[mlir.dialects.linalg.opdsl.lang.types.TypeVar] = None, size_exprs: mlir.dialects.linalg.opdsl.lang.scalar_expr.Optional[mlir.dialects.linalg.opdsl.lang.scalar_expr.Sequence[mlir.dialects.linalg.opdsl.lang.scalar_expr.AffineExprDef]] = None, index_dims: mlir.dialects.linalg.opdsl.lang.scalar_expr.Optional[mlir.dialects.linalg.opdsl.lang.scalar_expr.Sequence[mlir.dialects.linalg.opdsl.lang.scalar_expr.DimDef]] = None, default_indices: mlir.dialects.linalg.opdsl.lang.scalar_expr.Optional[mlir.dialects.linalg.opdsl.lang.scalar_expr.Sequence[int]] = None, default_fn: mlir.dialects.linalg.opdsl.lang.scalar_expr.Optional[str] = None)¶

Definition of an operand passed to an operation.

Keep the meta information of Tensor, Scalar, and Attribute operands and provide the shared registration functionality.

owner: mlir.dialects.linalg.opdsl.lang.scalar_expr.Optional[LinalgOpDef] = None¶

type_var = None¶

size_exprs = None¶

index_dims = None¶

default_indices = None¶

default_fn = None¶

kind¶

name: mlir.dialects.linalg.opdsl.lang.scalar_expr.Optional[str] = None¶

registered_index: int = -1¶

attach(index: int, name: str, owner: LinalgOpDef)¶

is_input() → bool¶

is_tensor() → bool¶

is_attribute() → bool¶

__hash__()¶

__repr__()¶

class mlir.dialects.linalg.opdsl.lang.TensorDef(type_var: mlir.dialects.linalg.opdsl.lang.types.TypeVar, *shape: mlir.dialects.linalg.opdsl.lang.scalar_expr.AffineExprDef, index_dims: mlir.dialects.linalg.opdsl.lang.scalar_expr.Optional[mlir.dialects.linalg.opdsl.lang.scalar_expr.Sequence[mlir.dialects.linalg.opdsl.lang.scalar_expr.DimDef]] = None, output: bool = False)¶

Tensor operand definition.

Tensor operands are indexed using the associated indexing_map when forwarded to the body of the structured op. A unique name identifies the tensor operands and an index determines their position in the operation’s parameter list. A tensor definition takes type, a shape, and an optional flag to mark output tensors. Additionally, a tuple of index dimensions may be used to map the tensor to the loop dimensions of the operation. This mapping is needed to compute the indexing map of shape-only tensors that have no uses.

operand_def¶

__getitem__(dims: mlir.dialects.linalg.opdsl.lang.scalar_expr.Sequence[mlir.dialects.linalg.opdsl.lang.scalar_expr.AffineExprDef]) → TensorUse¶

__setitem__(dims: mlir.dialects.linalg.opdsl.lang.scalar_expr.Sequence[mlir.dialects.linalg.opdsl.lang.scalar_expr.AffineExprDef], value: TensorExpression)¶

Creates a new 1:1 comprehension by binding this tensor to an expression.

Note that due to the way assignment works in Python, we have to capture direct assignment as a setitem on the TensorDef.

class mlir.dialects.linalg.opdsl.lang.ScalarDef(type_var: mlir.dialects.linalg.opdsl.lang.types.TypeVar)¶

Bases: TensorExpression

Scalar operand definition.

Scalar operands are forwarded to the body of the structured op as they are. A unique name identifies the scalars and an index determines their position in the operation’s parameter list.

operand_def¶

property scalar_name: str¶

to_scalar_expression() → mlir.dialects.linalg.opdsl.lang.scalar_expr.ScalarExpression¶

class mlir.dialects.linalg.opdsl.lang.IndexAttrDef(*sizes: mlir.dialects.linalg.opdsl.lang.scalar_expr.SymbolDef, default: mlir.dialects.linalg.opdsl.lang.scalar_expr.Sequence[int])¶

Index attribute definition.

Index attributes provide a way to define and set symbols that can be used in indexing expressions. Every attribute specifies a tuple of symbols that at compile-time are replaced by integer values as well as their default values.

operand_def¶

class mlir.dialects.linalg.opdsl.lang.UnaryFnAttrDef(default: UnaryFnType)¶

Unary function attribute definition.

Unary function attributes provide a way to make the arithmetic computation parametrizable. Every attribute specifies a default unary function that may be overwritten at operation instantiation time.

operand_def¶

__call__(arg: TensorExpression) → TensorFn¶

class mlir.dialects.linalg.opdsl.lang.BinaryFnAttrDef(default: BinaryFnType)¶

Binary function attribute definition.

Binary function attributes provide a way to make the arithmetic computation parametrizable. Every attribute specifies a default binary function that may be overwritten at operation instantiation time.

operand_def¶

__call__(arg0: TensorExpression, arg1: TensorExpression) → TensorFn¶

__getitem__(reduce_dims: mlir.dialects.linalg.opdsl.lang.scalar_expr.Tuple[mlir.dialects.linalg.opdsl.lang.scalar_expr.DimDef]) → ReduceFnUse¶

class mlir.dialects.linalg.opdsl.lang.TernaryFnAttrDef(default: TernaryFnType)¶

Ternary function attribute definition.

Ternary function attributes provide a way to make the arithmetic computation parametrizable. Every attribute specifies a default Ternary function that may be overwritten at operation instantiation time.

operand_def¶

__call__(arg0: TensorExpression, arg1: TensorExpression) → TensorFn¶

__getitem__(reduce_dims: mlir.dialects.linalg.opdsl.lang.scalar_expr.Tuple[mlir.dialects.linalg.opdsl.lang.scalar_expr.DimDef]) → ReduceFnUse¶

class mlir.dialects.linalg.opdsl.lang.TypeFnAttrDef(default: TypeFnType)¶

Type conversion function attribute definition.

Type conversion function attributes provide a way to make type conversions parameterizable. Every attribute specifies a default type conversion function that may be overwritten at operation instantiation time.

operand_def¶

__call__(type_var: mlir.dialects.linalg.opdsl.lang.types.TypeVar, arg: TensorExpression) → TensorFn¶

class mlir.dialects.linalg.opdsl.lang.Comprehension(*bindings: mlir.dialects.linalg.opdsl.lang.scalar_expr.Tuple[TensorUse, TensorExpression])¶

Represents a single comprehension.

definitions = []¶

values = []¶

property all_reduction_dims: mlir.dialects.linalg.opdsl.lang.scalar_expr.Set[mlir.dialects.linalg.opdsl.lang.scalar_expr.Tuple[mlir.dialects.linalg.opdsl.lang.scalar_expr.DimDef, Ellipsis]]¶: Gets the reduction dims for the comprehension or None.

__repr__()¶

class mlir.dialects.linalg.opdsl.lang.OpInterfaceDef(cpp_name: str)¶

An interface that an op implements.

cpp_name¶

mlir.dialects.linalg.opdsl.lang.ContractionOpInterface¶

mlir.dialects.linalg.opdsl.lang.ConvolutionOpInterface¶

mlir.dialects.linalg.opdsl.lang.FillOpInterface¶

class mlir.dialects.linalg.opdsl.lang.OpDefinitionDef(def_name: str)¶

A method that an op implements.

def_name¶

mlir.dialects.linalg.opdsl.lang.Canonicalizer¶

class mlir.dialects.linalg.opdsl.lang.OpMetadataDef(name: str, cpp_class_name: mlir.dialects.linalg.opdsl.lang.scalar_expr.Optional[str], doc: mlir.dialects.linalg.opdsl.lang.scalar_expr.Optional[str])¶

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

Metadata about the op (generally not behavior impacting).

yaml_tag = '!LinalgOpMetadata'¶

name¶

cpp_class_name¶

doc¶

implements: mlir.dialects.linalg.opdsl.lang.scalar_expr.List[OpInterfaceDef] = []¶

defines: mlir.dialects.linalg.opdsl.lang.scalar_expr.List[OpDefinitionsDef] = []¶

to_yaml_custom_dict()¶

class mlir.dialects.linalg.opdsl.lang.LinalgOpDef(name: str, cpp_class_name: mlir.dialects.linalg.opdsl.lang.scalar_expr.Optional[str] = None, doc: mlir.dialects.linalg.opdsl.lang.scalar_expr.Optional[str] = None)¶

Definition of a linalg op.

metadata¶

registered_operands: mlir.dialects.linalg.opdsl.lang.types.Dict[str, OperandDef]¶

domain: mlir.dialects.linalg.opdsl.lang.scalar_expr.List[mlir.dialects.linalg.opdsl.lang.scalar_expr.DimDef] = []¶

comprehensions: mlir.dialects.linalg.opdsl.lang.scalar_expr.List[Comprehension] = []¶

_affine_state¶

add_operand(name: str, operand: OperandDef)¶: Registers an operand.

__repr__()¶

class mlir.dialects.linalg.opdsl.lang.AffineBuildState(*, global_state: AffineBuildState = None, allow_new_symbols: bool = True, allow_new_dims: bool = True)¶

Internal state for the AffineExprDef._create impls.

Note that a “local” AffineBuildState can be created relative to a “global” AffineBuildState. In that case, any affine expressions built will inherit symbol and dim bindings from the global state and will update both as new ones are discovered. This allows for building expressions across contexts which share a common symbol and dim space.

local_symbols: Dict[str, int]¶

local_dims: Dict[str, int]¶

allow_new_symbols = True¶

allow_new_dims = True¶

get_dim(dimname: str) → int¶: Gets the dim position given a name.

get_symbol(symname: str) → int¶: Geta a symbol position given a name.

property local_dim_count: int¶

property local_symbol_count: int¶

property dim_count: int¶

property symbol_count: int¶

__repr__()¶

class mlir.dialects.linalg.opdsl.lang.AffineExprDef¶

Base class for an affine expression being defined.

build(state: AffineBuildState | None = None) → mlir.ir.AffineExpr¶: Builds the corresponding _ir.AffineExpr from the definitions.

abstract _create(state: AffineBuildState) → mlir.ir.AffineExpr¶

static coerce_from(py_value)¶

visit_affine_exprs(callback)¶: Visits all AffineExprDefs including self.

__add__(rhs)¶

__mul__(rhs)¶

__mod__(rhs)¶

__floordiv__(rhs)¶

__truediv__(rhs)¶

mlir.dialects.linalg.opdsl.lang.D¶

class mlir.dialects.linalg.opdsl.lang.DimDef¶

Bases: AffineExprDef

Represents a named dimension.

ALL_DIMS: Dict[str, DimDef]¶

__repr__()¶

_create(state: AffineBuildState) → mlir.ir.AffineExpr¶

classmethod create_expando()¶: Create an expando class that creates unique symbols based on attr access.

mlir.dialects.linalg.opdsl.lang.S¶

class mlir.dialects.linalg.opdsl.lang.SymbolDef¶

Bases: AffineExprDef

Represents a named symbol.

s1 = SymbolDef(“s1”) s1 Symbol(s1) s2 = SymbolDef(“s2”) s1 is s2 False s1 is SymbolDef(“s1”) True

ALL_SYMBOLS: Dict[str, SymbolDef]¶

__repr__()¶

_create(state: AffineBuildState) → mlir.ir.AffineExpr¶

classmethod create_expando()¶: Create an expando class that creates unique symbols based on attr access.

class mlir.dialects.linalg.opdsl.lang.ScalarAssign(arg: str, value: ScalarExpression)¶

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

An assignment to a named argument (LHS of a comprehension).

yaml_tag = '!ScalarAssign'¶

arg¶

value¶

to_yaml_custom_dict()¶

__repr__()¶

class mlir.dialects.linalg.opdsl.lang.ScalarFn(kind: mlir.dialects.linalg.opdsl.lang.comprehension.FunctionKind, fn_name: mlir.dialects.linalg.opdsl.lang.comprehension.Optional[str], attr_name: mlir.dialects.linalg.opdsl.lang.comprehension.Optional[str], type_var: mlir.dialects.linalg.opdsl.lang.comprehension.Optional[mlir.dialects.linalg.opdsl.lang.types.TypeVar], operands: mlir.dialects.linalg.opdsl.lang.comprehension.Sequence[ScalarExpression])¶

A type of ScalarExpression that applies a function.

kind¶

fn_name¶

attr_name¶

type_var¶

operands¶

expr() → ScalarExpression¶

__repr__()¶

class mlir.dialects.linalg.opdsl.lang.ScalarArg(arg: str)¶

A type of ScalarExpression that references a named argument.

arg¶

expr() → ScalarExpression¶

__repr__()¶

class mlir.dialects.linalg.opdsl.lang.ScalarConst(value: str)¶

A type of ScalarExpression representing a constant.

value¶

expr() → ScalarExpression¶

__repr__()¶

class mlir.dialects.linalg.opdsl.lang.ScalarIndex(dim: int)¶

A type of ScalarExpression accessing an iteration index.

dim¶

expr() → ScalarExpression¶

__repr__()¶

class mlir.dialects.linalg.opdsl.lang.ScalarExpression(scalar_fn: mlir.dialects.linalg.opdsl.lang.comprehension.Optional[ScalarFn] = None, scalar_arg: mlir.dialects.linalg.opdsl.lang.comprehension.Optional[ScalarArg] = None, scalar_const: mlir.dialects.linalg.opdsl.lang.comprehension.Optional[ScalarConst] = None, scalar_index: mlir.dialects.linalg.opdsl.lang.comprehension.Optional[ScalarIndex] = None)¶

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

An expression on scalar values.

Can be one of:

ScalarFn
ScalarArg
ScalarConst
ScalarIndex

yaml_tag = '!ScalarExpression'¶

scalar_fn = None¶

scalar_arg = None¶

scalar_const = None¶

scalar_index = None¶

to_yaml_custom_dict()¶

class mlir.dialects.linalg.opdsl.lang.TypeVar¶

A replaceable type variable.

Type variables are uniqued by name.

ALL_TYPEVARS: Dict[str, TypeVar]¶

__repr__()¶

classmethod create_expando()¶: Create an expando class that creates unique type vars on attr access.

mlir.dialects.linalg.opdsl.lang.TV¶

mlir.dialects.linalg.opdsl.lang.I32¶

mlir.dialects.linalg.opdsl.lang.I64¶

mlir.dialects.linalg.opdsl.lang.F32¶

mlir.dialects.linalg.opdsl.lang.F64¶

mlir.dialects.linalg.opdsl.lang.T¶

mlir.dialects.linalg.opdsl.lang.U¶

mlir.dialects.linalg.opdsl.lang.V¶

mlir.dialects.linalg.opdsl.lang.yaml_dump(data, sort_keys=False, **kwargs)¶

mlir.dialects.linalg.opdsl.lang.yaml_dump_all(data, sort_keys=False, explicit_start=True, **kwargs)¶

class mlir.dialects.linalg.opdsl.lang.YAMLObject¶

Bases: yaml.YAMLObject

An object that can dump itself to a YAML stream and load itself from a YAML stream.

classmethod to_yaml(dumper, self)¶: Default to a custom dictionary mapping.

abstract to_yaml_custom_dict()¶

as_linalg_yaml()¶

class mlir.dialects.linalg.opdsl.lang.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.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__()¶

class mlir.dialects.linalg.opdsl.lang.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__()¶

mlir.dialects.linalg.opdsl.lang.emit_generic_structured_op(op_config: mlir.dialects.linalg.opdsl.lang.config.LinalgStructuredOpConfig, *ins: Value, outs: ValueList, **attrs: mlir.dialects.linalg.opdsl.lang.comprehension.Sequence[int])¶

mlir.dialects.linalg.opdsl.lang.emit_named_structured_op(op_config: mlir.dialects.linalg.opdsl.lang.config.LinalgStructuredOpConfig, op_name: str, op_class_name: str, *ins: Value, outs: ValueList, **attrs: mlir.dialects.linalg.opdsl.lang.comprehension.Sequence[int])¶

mlir.dialects.linalg.opdsl.lang.ValueList¶