mlir.dialects.linalg.opdsl.lang.affine

DSL for constructing affine expressions and maps.

These python wrappers allow construction of affine expressions in a more pythonic fashion that is later instantiated as an IR AffineExpr. Separating the AST from construction of the map allows for manipulations of symbols and dims beyond the scope of one expression.

Affine expression construction:

with _ir.Context(): … s = AffineBuildState() … (S.K + S.M).build(s) … (S.K * S.M).build(s) … (S.K // S.M).build(s) … (S.K / S.M).build(s) … (S.K % 4).build(s) … (D.i + D.j * 4).build(s) … s AffineExpr(s0 + s1) AffineExpr(s0 * s1) AffineExpr(s0 floordiv s1) AffineExpr(s0 ceildiv s1) AffineExpr(s0 mod 4) AffineExpr(d0 + d1 * 4) AffineBuildState< symbols={‘K’: 0, ‘M’: 1} dims={‘i’: 0, ‘j’: 1}>

In the DSL, dimensions and symbols are name-uniqued instances of DimDef and SymbolDef. There are shortcut “expando” instances that will create a corresponding DimDef/SymbolDef upon accessing an attribute:

Referencing a named dimension:

D.i Dim(i) D.a is D.b False D.a is D.a True

Referencing a named symbol:

S.foobar Symbol(foobar) S.a is S.b False S.a is S.a True

Attributes

D
S

Classes

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.

Module Contents

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

class mlir.dialects.linalg.opdsl.lang.affine.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.

class mlir.dialects.linalg.opdsl.lang.affine.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.

mlir.dialects.linalg.opdsl.lang.affine.D

mlir.dialects.linalg.opdsl.lang.affine.S