Skip to content

Expand math to polynomials

expand_math_to_polynomials

ExpandExp dataclass

Bases: RewritePattern

Replace math.exp operations with a polynomial expansion.

Only expands when the number of terms is specified, either via an attribute on the operation or via the pass-level default.

Source code in xdsl/transforms/expand_math_to_polynomials.py 
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
@dataclass
class ExpandExp(RewritePattern):
    """
    Replace `math.exp` operations with a polynomial expansion.

    Only expands when the number of terms is specified, either via an
    attribute on the operation or via the pass-level default.
    """

    default_terms: int | None = None
    """Pass-level default for number of terms. None means don't expand
    unless the operation has an explicit terms attribute."""

    @op_type_rewrite_pattern
    def match_and_rewrite(self, op: math.ExpOp, rewriter: PatternRewriter) -> None:
        terms: int | None = None
        if "terms" in op.attributes:
            attr = op.attributes["terms"]
            if isinstance(attr, IntegerAttr):
                terms = attr.value.data
        elif self.default_terms is not None:
            terms = self.default_terms

        if terms is None:
            return

        expanded: Operation = expand_exp(op, rewriter, terms)
        rewriter.replace_op(op, (), (expanded.results[0],))

default_terms: int | None = None class-attribute instance-attribute

Pass-level default for number of terms. None means don't expand unless the operation has an explicit terms attribute.

__init__(default_terms: int | None = None) -> None

match_and_rewrite(op: math.ExpOp, rewriter: PatternRewriter) -> None

Source code in xdsl/transforms/expand_math_to_polynomials.py 
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
@op_type_rewrite_pattern
def match_and_rewrite(self, op: math.ExpOp, rewriter: PatternRewriter) -> None:
    terms: int | None = None
    if "terms" in op.attributes:
        attr = op.attributes["terms"]
        if isinstance(attr, IntegerAttr):
            terms = attr.value.data
    elif self.default_terms is not None:
        terms = self.default_terms

    if terms is None:
        return

    expanded: Operation = expand_exp(op, rewriter, terms)
    rewriter.replace_op(op, (), (expanded.results[0],))

ExpandMathToPolynomialsPass dataclass

Bases: ModulePass

This pass expands math operations to a polynomial expansion using the Taylor series.

Currently only expands math.exp operations.

Operations are only expanded when the number of terms is specified, either via a terms attribute on the operation itself or via the pass-level terms parameter.

Source code in xdsl/transforms/expand_math_to_polynomials.py 
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
@dataclass(frozen=True)
class ExpandMathToPolynomialsPass(ModulePass):
    """
    This pass expands `math` operations to a polynomial expansion using the Taylor series.

    Currently only expands `math.exp` operations.

    Operations are only expanded when the number of terms is specified,
    either via a `terms` attribute on the operation itself or via the
    pass-level `terms` parameter.
    """

    name = "expand-math-to-polynomials"

    terms: int | None = None
    """Number of terms in the resulting polynomial expansion.
    If not set, only operations with an explicit terms attribute are expanded."""

    def apply(self, ctx: Context, op: ModuleOp) -> None:
        PatternRewriteWalker(
            ExpandExp(default_terms=self.terms),
            apply_recursively=False,
        ).rewrite_module(op)

name = 'expand-math-to-polynomials' class-attribute instance-attribute

terms: int | None = None class-attribute instance-attribute

Number of terms in the resulting polynomial expansion. If not set, only operations with an explicit terms attribute are expanded.

__init__(terms: int | None = None) -> None

apply(ctx: Context, op: ModuleOp) -> None

Source code in xdsl/transforms/expand_math_to_polynomials.py 
123
124
125
126
127
def apply(self, ctx: Context, op: ModuleOp) -> None:
    PatternRewriteWalker(
        ExpandExp(default_terms=self.terms),
        apply_recursively=False,
    ).rewrite_module(op)

expand_exp(op: math.ExpOp, rewriter: PatternRewriter, terms: int) -> Operation

Expand exp(x) using a Taylor-series polynomial expansion.

Pseudo-code::

result = 1.0
term = 1.0
for i in range(1, terms): # loop will be unrolled by the rewriter
    term *= x / i
    result += term
return result
Source code in xdsl/transforms/expand_math_to_polynomials.py 
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
def expand_exp(op: math.ExpOp, rewriter: PatternRewriter, terms: int) -> Operation:
    """
    Expand exp(x) using a Taylor-series polynomial expansion.

    Pseudo-code::

        result = 1.0
        term = 1.0
        for i in range(1, terms): # loop will be unrolled by the rewriter
            term *= x / i
            result += term
        return result
    """
    x = op.operands[0]
    tp = x.type
    if not isa(tp, AnyFloat | VectorType[AnyFloat] | TensorType[AnyFloat]):
        raise TypeError(f"Unsupported type for math.exp expansion: {tp}")

    res = _float_constant(1.0, tp, rewriter)
    term = _float_constant(1.0, tp, rewriter)

    for i in range(1, terms):
        i_val = _float_constant(1.0 / float(i), tp, rewriter)
        frac = rewriter.insert(arith.MulfOp(x, i_val.result))
        mul = rewriter.insert(arith.MulfOp(frac.result, term.result))
        add = rewriter.insert(arith.AddfOp(res.result, mul.result))

        term = mul
        res = add

    return res