transforms: lower polynomials to arithmetic operations

tests/filecheck/transforms/polynomial_eval_to_arith.mlir

@@ -0,0 +1,111 @@
+// RUN: xdsl-opt -p polynomial-eval-to-arith %s | filecheck %s
+
+builtin.module {
+  // f64 with domain bounds [-5, 0], degree 2.
+  // Scale = 2/(0 - (-5)) = 0.4, offset = -((-5)+0)/(0-(-5)) = 1.0.
+  func.func @clenshaw_f64_with_domain(%x: f64) -> f64 {
+    %r = polynomial.eval #polynomial.typed_chebyshev_polynomial<[1.000000e+00 : f64, 2.000000e+00 : f64, 3.000000e+00 : f64]> : !polynomial.polynomial<ring = <coefficientType = f64>>, %x {scheme = "clenshaw", domain_lower = -5.000000e+00 : f64, domain_upper = 0.000000e+00 : f64} : f64
+    func.return %r : f64
+  }
+
+  // f64 without domain bounds: skip the affine remap (t = x).
+  func.func @clenshaw_f64_no_domain(%x: f64) -> f64 {
+    %r = polynomial.eval #polynomial.typed_chebyshev_polynomial<[1.000000e+00 : f64, 2.000000e+00 : f64, 3.000000e+00 : f64]> : !polynomial.polynomial<ring = <coefficientType = f64>>, %x {scheme = "clenshaw"} : f64
+    func.return %r : f64
+  }
+
+  // f32 with domain [-1, 1]: scale=1.0, offset=0.0.
+  func.func @clenshaw_f32(%x: f32) -> f32 {
+    %r = polynomial.eval #polynomial.typed_chebyshev_polynomial<[5.000000e-01 : f64, 1.200000e+00 : f64, 3.000000e-01 : f64]> : !polynomial.polynomial<ring = <coefficientType = f64>>, %x {scheme = "clenshaw", domain_lower = -1.000000e+00 : f64, domain_upper = 1.000000e+00 : f64} : f32
+    func.return %r : f32
+  }
+
+  // Vector type, with domain bounds, degree 1.
+  func.func @clenshaw_vec(%x: vector<4xf32>) -> vector<4xf32> {
+    %r = polynomial.eval #polynomial.typed_chebyshev_polynomial<[1.000000e+00 : f64, 2.000000e+00 : f64]> : !polynomial.polynomial<ring = <coefficientType = f64>>, %x {scheme = "clenshaw", domain_lower = -2.000000e+00 : f64, domain_upper = 2.000000e+00 : f64} : vector<4xf32>
+    func.return %r : vector<4xf32>
+  }
+}
+
+// CHECK: builtin.module {
+
+// ===== f64 with domain [-5, 0], coeffs [1, 2, 3] =====
+
+// CHECK:      func.func @clenshaw_f64_with_domain(%[[X:.*]]: f64) -> f64 {
+//             Domain mapping: t = x * 0.4 + 1.0
+// CHECK-NEXT:   %[[SCALE:.*]] = arith.constant 4.000000e-01 : f64
+// CHECK-NEXT:   %[[OFFSET:.*]] = arith.constant 1.000000e+00 : f64
+// CHECK-NEXT:   %[[SCALED:.*]] = arith.mulf %[[X]], %[[SCALE]] : f64
+// CHECK-NEXT:   %[[T:.*]] = arith.addf %[[SCALED]], %[[OFFSET]] : f64
+//             two_t = 2 * t
+// CHECK-NEXT:   %[[TWO:.*]] = arith.constant 2.000000e+00 : f64
+// CHECK-NEXT:   %[[TWO_T:.*]] = arith.mulf %[[TWO]], %[[T]] : f64
+//             b_{n+2} = 0, b_{n+1} = 0
+// CHECK-NEXT:   %{{.*}} = arith.constant 0.000000e+00 : f64
+// CHECK-NEXT:   %{{.*}} = arith.constant 0.000000e+00 : f64
+//             2 Clenshaw iterations
+// CHECK-COUNT-2: arith.addf
+//             Final: result = c_0/2 + t * b_1 - b_2
+// CHECK:        %[[C0_HALF:.*]] = arith.constant 5.000000e-01 : f64
+// CHECK-NEXT:   %[[T_B1:.*]] = arith.mulf %[[T]], %{{.*}} : f64
+// CHECK-NEXT:   %[[ADD:.*]] = arith.addf %[[C0_HALF]], %[[T_B1]] : f64
+// CHECK-NEXT:   %[[RES:.*]] = arith.subf %[[ADD]], %{{.*}} : f64
+// CHECK-NEXT:   func.return %[[RES]] : f64
+// CHECK-NEXT: }
+
+// ===== f64 without domain bounds: t = x, no scale/offset ops =====
+
+// CHECK:      func.func @clenshaw_f64_no_domain(%[[XN:.*]]: f64) -> f64 {
+// CHECK-NEXT:   %[[TWO_N:.*]] = arith.constant 2.000000e+00 : f64
+// CHECK-NEXT:   %[[TWO_T_N:.*]] = arith.mulf %[[TWO_N]], %[[XN]] : f64
+// CHECK-NEXT:   %{{.*}} = arith.constant 0.000000e+00 : f64
+// CHECK-NEXT:   %{{.*}} = arith.constant 0.000000e+00 : f64
+// CHECK-COUNT-2: arith.addf
+// CHECK:        %[[C0H_N:.*]] = arith.constant 5.000000e-01 : f64
+// CHECK-NEXT:   %[[TB1_N:.*]] = arith.mulf %[[XN]], %{{.*}} : f64
+// CHECK-NEXT:   %{{.*}} = arith.addf %[[C0H_N]], %[[TB1_N]] : f64
+// CHECK-NEXT:   %[[RES_N:.*]] = arith.subf %{{.*}}, %{{.*}} : f64
+// CHECK-NEXT:   func.return %[[RES_N]] : f64
+// CHECK-NEXT: }
+
+// ===== f32 with [-1, 1]: coefficients converted to f32 =====
+
+// CHECK:      func.func @clenshaw_f32(%[[X32:.*]]: f32) -> f32 {
+// CHECK-NEXT:   %{{.*}} = arith.constant 1.000000e+00 : f32
+// CHECK-NEXT:   %{{.*}} = arith.constant -0.000000e+00 : f32
+// CHECK-NEXT:   %{{.*}} = arith.mulf %[[X32]], %{{.*}} : f32
+// CHECK-NEXT:   %[[T32:.*]] = arith.addf %{{.*}}, %{{.*}} : f32
+// CHECK-NEXT:   %{{.*}} = arith.constant 2.000000e+00 : f32
+// CHECK-NEXT:   %{{.*}} = arith.mulf %{{.*}}, %[[T32]] : f32
+// CHECK-NEXT:   %{{.*}} = arith.constant 0.000000e+00 : f32
+// CHECK-NEXT:   %{{.*}} = arith.constant 0.000000e+00 : f32
+// CHECK-COUNT-2: arith.addf
+// CHECK:        %{{.*}} = arith.constant 2.500000e-01 : f32
+// CHECK-NEXT:   %{{.*}} = arith.mulf %[[T32]], %{{.*}} : f32
+// CHECK-NEXT:   %{{.*}} = arith.addf %{{.*}}, %{{.*}} : f32
+// CHECK-NEXT:   %[[RES32:.*]] = arith.subf %{{.*}}, %{{.*}} : f32
+// CHECK-NEXT:   func.return %[[RES32]] : f32
+// CHECK-NEXT: }
+
+// ===== vector<4xf32> with [-2, 2], degree 1 =====
+//   scale = 0.5, offset = -0.0 (since (lower+upper) == 0).
+
+// CHECK:      func.func @clenshaw_vec(%[[XV:.*]]: vector<4xf32>) -> vector<4xf32> {
+// CHECK-NEXT:   %{{.*}} = arith.constant dense<5.000000e-01> : vector<4xf32>
+// CHECK-NEXT:   %{{.*}} = arith.constant dense<-0.000000e+00> : vector<4xf32>
+// CHECK-NEXT:   %{{.*}} = arith.mulf %[[XV]], %{{.*}} : vector<4xf32>
+// CHECK-NEXT:   %[[TV:.*]] = arith.addf %{{.*}}, %{{.*}} : vector<4xf32>
+// CHECK-NEXT:   %{{.*}} = arith.constant dense<2.000000e+00> : vector<4xf32>
+// CHECK-NEXT:   %{{.*}} = arith.mulf %{{.*}}, %[[TV]] : vector<4xf32>
+// CHECK-NEXT:   %{{.*}} = arith.constant dense<0.000000e+00> : vector<4xf32>
+// CHECK-NEXT:   %{{.*}} = arith.constant dense<0.000000e+00> : vector<4xf32>
+//             1 Clenshaw iteration (degree=1)
+// CHECK-COUNT-1: arith.addf
+// CHECK:        %{{.*}} = arith.constant dense<5.000000e-01> : vector<4xf32>
+// CHECK-NEXT:   %{{.*}} = arith.mulf %[[TV]], %{{.*}} : vector<4xf32>
+// CHECK-NEXT:   %{{.*}} = arith.addf %{{.*}}, %{{.*}} : vector<4xf32>
+// CHECK-NEXT:   %[[RESV:.*]] = arith.subf %{{.*}}, %{{.*}} : vector<4xf32>
+// CHECK-NEXT:   func.return %[[RESV]] : vector<4xf32>
+// CHECK-NEXT: }
+
+// CHECK: }

tests/transforms/test_polynomial_eval_to_arith.py

@@ -0,0 +1,39 @@
+import pytest
+
+from xdsl.builder import ImplicitBuilder
+from xdsl.context import Context
+from xdsl.dialects import arith, polynomial, test
+from xdsl.dialects.builtin import Builtin, ModuleOp, f32
+from xdsl.transforms.polynomial_eval_to_arith import PolynomialEvalToArithPass
+from xdsl.utils.exceptions import PassFailedException
+
+
+@pytest.mark.parametrize("scheme", list(polynomial.EvalScheme))
+def test_every_scheme_has_lowering(scheme: polynomial.EvalScheme):
+    """
+    This tests if every member of `EvalScheme` in the polynomial dialect has a defined lowering to arithmetic ops.
+    """
+    ctx = Context()
+    ctx.load_dialect(Builtin)
+    ctx.load_dialect(arith.Arith)
+    ctx.load_dialect(polynomial.Polynomial)
+    ctx.load_dialect(test.Test)
+
+    module = ModuleOp([])
+    with ImplicitBuilder(module.body):
+        x = test.TestOp(result_types=[f32]).results[0]
+        polynomial.EvalOp.get(
+            value=x,
+            coefficients=(1.0, 2.0, 3.0),
+            element_type=f32,
+            scheme=scheme,
+        )
+
+    try:
+        PolynomialEvalToArithPass().apply(ctx, module)
+    except PassFailedException as e:
+        pytest.fail(
+            f"EvalScheme.{scheme.name} has no dispatch branch in "
+            f"polynomial-eval-to-arith. Add a case for it in "
+            f"PolynomialEvalToArith.match_and_rewrite. ({e})"
+        )

xdsl/dialects/polynomial.py

@@ -217,9 +217,12 @@ class EvalOp(IRDLOperation):
     """
     Evaluate a polynomial at a given point.
 
-    This op is *unevaluated* but carries all information needed for
+    This op is unevaluated but carries all information needed for
     later lowering to arithmetic ops, dispatched on `scheme`.
 
+    Domain convention: `domain_lower` / `domain_upper` describe the interval
+    the Chebyshev coefficients were computed for.  Default is [-1, 1].
+
     Syntax: polynomial.eval $polynomial `,` $value attr-dict `:` type($value)
     Example:
         %result = polynomial.eval

xdsl/transforms/__init__.py

@@ -91,6 +91,11 @@ def get_expand_math_to_polynomials():
 
         return expand_math_to_polynomials.ExpandMathToPolynomialsPass
 
+    def get_polynomial_eval_to_arith():
+        from xdsl.transforms import polynomial_eval_to_arith
+
+        return polynomial_eval_to_arith.PolynomialEvalToArithPass
+
     def get_convert_linalg_to_loops():
         from xdsl.transforms import convert_linalg_to_loops
 
@@ -693,6 +698,7 @@ def get_verify_register_allocation():
         "convert-func-to-riscv-func": get_convert_func_to_riscv_func,
         "convert-func-to-x86-func": get_convert_func_to_x86_func,
         "expand-math-to-polynomials": get_expand_math_to_polynomials,
+        "polynomial-eval-to-arith": get_polynomial_eval_to_arith,
         "convert-linalg-to-loops": get_convert_linalg_to_loops,
         "convert-linalg-to-memref-stream": get_convert_linalg_to_memref_stream,
         "convert-memref-stream-to-loops": get_convert_memref_stream_to_loops,

xdsl/transforms/polynomial_eval_to_arith.py

@@ -0,0 +1,150 @@
+"""
+Expand `polynomial.eval` ops to arithmetic operations.
+
+This pass dispatches on the `scheme` attribute and emits the corresponding arith ops.
+
+Currently supported schemes:
+  - "clenshaw": Chebyshev series evaluated via Clenshaw's recurrence.
+
+"""
+
+from dataclasses import dataclass
+
+from xdsl.context import Context
+from xdsl.dialects import arith, polynomial
+from xdsl.dialects.builtin import (
+    AnyFloat,
+    DenseIntOrFPElementsAttr,
+    FloatAttr,
+    ModuleOp,
+    TensorType,
+    VectorType,
+)
+from xdsl.ir import Operation, SSAValue
+from xdsl.irdl import isa
+from xdsl.passes import ModulePass
+from xdsl.pattern_rewriter import (
+    PatternRewriter,
+    PatternRewriteWalker,
+    RewritePattern,
+    op_type_rewrite_pattern,
+)
+from xdsl.utils.exceptions import PassFailedException
+
+
+def _float_constant(
+    value: float,
+    tp: AnyFloat | VectorType[AnyFloat] | TensorType[AnyFloat],
+    rewriter: PatternRewriter,
+) -> arith.ConstantOp:
+    """Create and insert a float constant, handling scalar/vector/tensor types."""
+    if isa(tp, VectorType[AnyFloat]):
+        attr = DenseIntOrFPElementsAttr.from_list(tp, [value])
+    elif isa(tp, TensorType[AnyFloat]):
+        attr = DenseIntOrFPElementsAttr.from_list(tp, [value])
+    elif isa(tp, AnyFloat):
+        attr = FloatAttr(value, tp)
+    else:
+        raise TypeError(f"Unsupported type for float constant: {tp}")
+    return rewriter.insert(arith.ConstantOp(attr))
+
+
+def expand_clenshaw(
+    op: polynomial.EvalOp,
+    rewriter: PatternRewriter,
+    coeffs: tuple[float, ...],
+    lower: float | None,
+    upper: float | None,
+) -> Operation:
+    """
+    Expand a Chebyshev series via Clenshaw's recurrence into arith ops.
+
+    If no domain bounds are provided, the input is assumed to already be
+    in the canonical Chebyshev domain [-1, 1].
+
+    Returns the final operation whose result is the evaluation result.
+    """
+    x = op.value
+    tp = x.type
+    if not isa(tp, AnyFloat | VectorType[AnyFloat] | TensorType[AnyFloat]):
+        raise TypeError(f"Unsupported type for polynomial.eval expansion: {tp}")
+
+    n = len(coeffs) - 1
+
+    # --- domain mapping: t = x * scale + offset --------------------------
+    t: SSAValue
+    if lower is not None and upper is not None:
+        scale = 2.0 / (upper - lower)
+        offset = -(upper + lower) / (upper - lower)
+        scale_op = _float_constant(scale, tp, rewriter)
+        offset_op = _float_constant(offset, tp, rewriter)
+        scaled = rewriter.insert(arith.MulfOp(x, scale_op.result))
+        t = rewriter.insert(arith.AddfOp(scaled.result, offset_op.result)).result
+    else:
+        t = x  # canonical Chebyshev domain [-1, 1]
+
+    # --- Clenshaw recurrence --------------------------------------------
+    two = _float_constant(2.0, tp, rewriter)
+    two_t = rewriter.insert(arith.MulfOp(two.result, t))
+
+    b_prev2 = _float_constant(0.0, tp, rewriter)  # b_{n+2}
+    b_prev1 = _float_constant(0.0, tp, rewriter)  # b_{n+1}
+
+    for k in range(n, 0, -1):
+        c_k = _float_constant(coeffs[k], tp, rewriter)
+        two_t_b = rewriter.insert(arith.MulfOp(two_t.result, b_prev1.result))
+        sub = rewriter.insert(arith.SubfOp(two_t_b.result, b_prev2.result))
+        b_k = rewriter.insert(arith.AddfOp(sub.result, c_k.result))
+        b_prev2 = b_prev1
+        b_prev1 = b_k
+
+    # --- final: result = c_0/2 + t * b_1 - b_2 ---------------------------
+    c0_half = _float_constant(coeffs[0] / 2.0, tp, rewriter)
+    t_b1 = rewriter.insert(arith.MulfOp(t, b_prev1.result))
+    add = rewriter.insert(arith.AddfOp(c0_half.result, t_b1.result))
+    return rewriter.insert(arith.SubfOp(add.result, b_prev2.result))
+
+
+class PolynomialEvalToArith(RewritePattern):
+    """Replace each `polynomial.eval` op with the arith ops for its scheme."""
+
+    @op_type_rewrite_pattern
+    def match_and_rewrite(
+        self, op: polynomial.EvalOp, rewriter: PatternRewriter
+    ) -> None:
+        coeffs = op.polynomial.coeff_values
+        lower = op.domain_lower.value.data if op.domain_lower is not None else None
+        upper = op.domain_upper.value.data if op.domain_upper is not None else None
+
+        scheme = op.eval_scheme
+
+        match scheme:
+            case polynomial.EvalScheme.CLENSHAW:
+                expanded = expand_clenshaw(op, rewriter, coeffs, lower, upper)
+            case _:
+                # Verifier already restricts `scheme` to known EvalScheme members,
+                # so this is only reachable if a new scheme is added without a
+                # corresponding lowering branch.
+                raise PassFailedException(
+                    f"polynomial.eval scheme {scheme.value!r} has no lowering"
+                )
+        rewriter.replace_op(op, (), (expanded.results[0],))
+
+
+@dataclass(frozen=True)
+class PolynomialEvalToArithPass(ModulePass):
+    """
+    Expand `polynomial.eval` ops to arithmetic operations.
+
+    All information needed for lowering (coefficients, scheme, domain
+    bounds) is read directly from each op, so this pass takes no
+    parameters.
+    """
+
+    name = "polynomial-eval-to-arith"
+
+    def apply(self, ctx: Context, op: ModuleOp) -> None:
+        PatternRewriteWalker(
+            PolynomialEvalToArith(),
+            apply_recursively=False,
+        ).rewrite_module(op)

xdsl/transforms/test_lower_linalg_to_snitch.py

@@ -29,6 +29,7 @@
     memref_stream_tile_outer_loops,
     memref_stream_unnest_out_parameters,
     memref_streamify,
+    polynomial_eval_to_arith,
     reconcile_unrealized_casts,
     riscv_allocate_registers,
     riscv_lower_parallel_mov,
@@ -58,6 +59,7 @@
     lower_affine.LowerAffinePass(),
     convert_scf_to_riscv_scf.ConvertScfToRiscvPass(),
     expand_math_to_polynomials.ExpandMathToPolynomialsPass(),
+    polynomial_eval_to_arith.PolynomialEvalToArithPass(),
     convert_arith_to_riscv_snitch.ConvertArithToRiscvSnitchPass(),
     convert_arith_to_riscv.ConvertArithToRiscvPass(),
     convert_func_to_riscv_func.ConvertFuncToRiscvFuncPass(),