# Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
# See https://llvm.org/LICENSE.txt for license information.
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
"""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
"""
from typing import Callable, Dict, Optional, Tuple, Union
from ..... import ir as _ir
__all__ = [
"AffineBuildState",
"AffineExprDef",
"D",
"DimDef",
"S",
"SymbolDef",
]
class AffineBuildState:
"""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.
"""
def __init__(
self,
*,
global_state: "AffineBuildState" = None,
allow_new_symbols: bool = True,
allow_new_dims: bool = True,
):
if not global_state:
self.all_symbols = dict() # type: Dict[str, int]
self.all_dims = dict() # type: Dict[str, int]
else:
# Alias the global dict.
self.all_symbols = global_state.all_symbols
self.all_dims = global_state.all_dims
# Map of symbols and dims in the current build.
self.local_symbols = dict() # type: Dict[str, int]
self.local_dims = dict() # type: Dict[str, int]
self.allow_new_symbols = allow_new_symbols
self.allow_new_dims = allow_new_dims
def get_dim(self, dimname: str) -> int:
"""Gets the dim position given a name."""
pos = self.all_dims.get(dimname)
if pos is None:
if not self.allow_new_dims:
raise ValueError(
f"New dimensions not allowed in the current affine expression: "
f"Requested '{dimname}', Availble: {self.all_dims}"
)
pos = len(self.all_dims)
self.all_dims[dimname] = pos
self.local_dims[dimname] = pos
return pos
def get_symbol(self, symname: str) -> int:
"""Geta a symbol position given a name."""
pos = self.all_symbols.get(symname)
if pos is None:
if not self.allow_new_symbols:
raise ValueError(
f"New symbols not allowed in the current affine expression: "
f"Requested '{symname}', Availble: {self.all_symbols}"
)
pos = len(self.all_symbols)
self.all_symbols[symname] = pos
self.local_symbols[symname] = pos
return pos
@property
def local_dim_count(self) -> int:
return len(self.local_dims)
@property
def local_symbol_count(self) -> int:
return len(self.local_symbols)
@property
def dim_count(self) -> int:
return len(self.all_dims)
@property
def symbol_count(self) -> int:
return len(self.all_symbols)
def __repr__(self):
lines = [f"AffineBuildState<"]
lines.append(f" symbols={self.local_symbols}")
lines.append(f" dims={self.local_dims}>")
return "\n".join(lines)
class AffineExprDef:
"""Base class for an affine expression being defined."""
def build(self, state: Optional[AffineBuildState] = None) -> _ir.AffineExpr:
"""Builds the corresponding _ir.AffineExpr from the definitions."""
state = AffineBuildState() if state is None else state
expr = self._create(state)
return expr
def _create(self, state: AffineBuildState) -> _ir.AffineExpr:
raise NotImplementedError()
@staticmethod
def coerce_from(py_value):
if isinstance(py_value, int):
return AffineConstantExpr(py_value)
assert isinstance(py_value, AffineExprDef)
return py_value
def visit_affine_exprs(self, callback):
"""Visits all AffineExprDefs including self."""
callback(self)
def __add__(lhs, rhs):
rhs = AffineExprDef.coerce_from(rhs)
return AffineBinaryExprDef(_ir.AffineAddExpr, lhs, rhs)
def __mul__(lhs, rhs):
rhs = AffineExprDef.coerce_from(rhs)
return AffineBinaryExprDef(_ir.AffineMulExpr, lhs, rhs)
def __mod__(lhs, rhs):
rhs = AffineExprDef.coerce_from(rhs)
return AffineBinaryExprDef(_ir.AffineModExpr, lhs, rhs)
def __floordiv__(lhs, rhs):
rhs = AffineExprDef.coerce_from(rhs)
return AffineBinaryExprDef(_ir.AffineFloorDivExpr, lhs, rhs)
def __truediv__(lhs, rhs):
# TODO: Not really a ceil div - taking liberties for the DSL.
rhs = AffineExprDef.coerce_from(rhs)
return AffineBinaryExprDef(_ir.AffineCeilDivExpr, lhs, rhs)
class AffineConstantExpr(AffineExprDef):
"""An affine constant being defined."""
def __init__(self, value: int):
assert isinstance(value, int)
self.value = value
def _create(self, state: AffineBuildState) -> _ir.AffineExpr:
return _ir.AffineConstantExpr.get(self.value)
def __repr__(self):
return f"Const({self.value})"
class AffineBinaryExprDef(AffineExprDef):
"""An affine binary expression being defined."""
def __init__(self, ir_ctor, lhs: AffineExprDef, rhs: AffineExprDef):
self.ir_ctor = ir_ctor
self.lhs = lhs
self.rhs = rhs
def _create(self, state: AffineBuildState) -> _ir.AffineExpr:
return self.ir_ctor.get(self.lhs._create(state), self.rhs._create(state))
def visit_affine_exprs(self, callback):
"""Visits all AffineExprDefs including self."""
super().visit_affine_exprs(callback)
self.lhs.visit_affine_exprs(callback)
self.rhs.visit_affine_exprs(callback)
def __repr__(self):
return f"{self.ir_ctor.__name__}({repr(self.lhs)}, {repr(self.rhs)})"
class DimDef(AffineExprDef):
"""Represents a named dimension."""
ALL_DIMS = dict() # type: Dict[str, "DimDef"]
def __new__(cls, dimname: str):
existing = cls.ALL_DIMS.get(dimname)
if existing is not None:
return existing
new = super().__new__(cls)
new.dimname = dimname
cls.ALL_DIMS[dimname] = new
return new
def __repr__(self):
return f"Dim({self.dimname})"
def _create(self, state: AffineBuildState) -> _ir.AffineExpr:
pos = state.get_dim(self.dimname)
return _ir.AffineDimExpr.get(position=pos)
@classmethod
def create_expando(cls):
"""Create an expando class that creates unique symbols based on attr access."""
class ExpandoDims:
def __getattr__(self, n):
return cls(n)
return ExpandoDims()
class SymbolDef(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() # type: Dict[str, "SymbolDef"]
def __new__(cls, symname: str):
existing = cls.ALL_SYMBOLS.get(symname)
if existing is not None:
return existing
new = super().__new__(cls)
new.symname = symname
cls.ALL_SYMBOLS[symname] = new
return new
def __repr__(self):
return f"Symbol({self.symname})"
def _create(self, state: AffineBuildState) -> _ir.AffineExpr:
pos = state.get_symbol(self.symname)
return _ir.AffineSymbolExpr.get(position=pos)
@classmethod
def create_expando(cls):
"""Create an expando class that creates unique symbols based on attr access."""
class ExpandoSymbols:
def __getattr__(self, n):
return cls(n)
return ExpandoSymbols()
# Global accessor for on-demand dims and symbols.
D = DimDef.create_expando()
S = SymbolDef.create_expando()
Codebase Browser
llvm