Skip to content

Expand math to polynomials

expand_math_to_polynomials

ExpandExp dataclass

Bases: RewritePattern

Replace math.exp operations with a polynomial expansion.

Source code in xdsl/transforms/expand_math_to_polynomials.py 
24
25
26
27
28
29
30
31
32
33
34
35
36
@dataclass
class ExpandExp(RewritePattern):
    """
    Replace `math.exp` operations with a polynomial expansion.
    """

    terms: int
    """Number of terms to use when expanding `math.exp`."""

    @op_type_rewrite_pattern
    def match_and_rewrite(self, op: math.ExpOp, rewriter: PatternRewriter) -> None:
        expanded: Operation = expand_exp(op, rewriter, self.terms)
        rewriter.replace_op(op, (), (expanded.results[0],))

terms: int instance-attribute

Number of terms to use when expanding math.exp.

__init__(terms: int) -> None

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

Source code in xdsl/transforms/expand_math_to_polynomials.py 
33
34
35
36
@op_type_rewrite_pattern
def match_and_rewrite(self, op: math.ExpOp, rewriter: PatternRewriter) -> None:
    expanded: Operation = expand_exp(op, rewriter, self.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.

Source code in xdsl/transforms/expand_math_to_polynomials.py 
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
@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.
    """

    name = "expand-math-to-polynomials"

    terms: int = 4
    """Number of terms in the resulting polynomial expansion."""

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

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

terms: int = 4 class-attribute instance-attribute

Number of terms in the resulting polynomial expansion.

__init__(terms: int = 4) -> None

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

Source code in xdsl/transforms/expand_math_to_polynomials.py 
101
102
103
104
105
def apply(self, ctx: Context, op: ModuleOp) -> None:
    PatternRewriteWalker(
        ExpandExp(self.terms),
        apply_recursively=False,
    ).rewrite_module(op)

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

Expand exp(x) using the Taylor-series loop from the pseudo-code:

terms = 75
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 
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
def expand_exp(op: math.ExpOp, rewriter: PatternRewriter, terms: int) -> Operation:
    """
    Expand exp(x) using the Taylor-series loop from the pseudo-code:

        terms = 75
        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