transforms: (exp) lower exp to polynomial dialect

tests/filecheck/transforms/lower_exp_to_polynomial.mlir

@@ -0,0 +1,65 @@
+// RUN: xdsl-opt -p lower-exp-to-polynomial %s | filecheck %s
+
+builtin.module {
+  // No max_bits_lost: pass picks the default (-1 = correctly-rounded target),
+  // domain falls back to [underflow, overflow] of f32.
+  func.func @exp_f32(%x: f32) -> f32 {
+    %r = math.exp %x : f32
+    func.return %r : f32
+  }
+
+  // No max_bits_lost: same default behavior, but f64 domain.
+  func.func @exp_f64(%x: f64) -> f64 {
+    %r = math.exp %x : f64
+    func.return %r : f64
+  }
+
+  // max_bits_lost is preserved on math.exp by step 1 and drives the polynomial
+  // degree picked by _choose_polynomial.
+  func.func @exp_f32_with_max_bits_lost(%x: f32) -> f32 {
+    %r = math.exp %x {max_bits_lost = 2 : i64} : f32
+    func.return %r : f32
+  }
+
+  // user lower_bound = -2 > underflow (~-103.28) -> kept at -2
+  // user upper_bound =  5 < overflow  (~ 88.72)  -> kept at 5
+  func.func @exp_f32_clamped_lower(%x: f32) -> f32 {
+    %r = math.exp %x {max_bits_lost = 2 : i64, lower_bound = -2.000000e+00 : f32, upper_bound = 5.000000e+00 : f32} : f32
+    func.return %r : f32
+  }
+
+  // both user bounds inside representable range -> used as-is
+  func.func @exp_f32_tight_bounds(%x: f32) -> f32 {
+    %r = math.exp %x {max_bits_lost = 2 : i64, lower_bound = -5.000000e-01 : f32, upper_bound = 5.000000e-01 : f32} : f32
+    func.return %r : f32
+  }
+}
+
+// Without max_bits_lost, math.exp is still lowered using the default
+// correctly-rounded target and [underflow, overflow] domain.
+// CHECK:      func.func @exp_f32(%[[X32:.*]]: f32) -> f32 {
+// CHECK-NEXT:   %[[R32:.*]] = polynomial.eval #polynomial.typed_chebyshev_polynomial<[{{.*}}]> : !polynomial.polynomial<ring = <coefficientType = f32>>, %[[X32]] {scheme = "clenshaw", domain_lower = {{.*}} : f32, domain_upper = {{.*}} : f32} : f32
+// CHECK-NEXT:   func.return %[[R32]] : f32
+// CHECK-NEXT: }
+
+// CHECK:      func.func @exp_f64(%[[X64:.*]]: f64) -> f64 {
+// CHECK-NEXT:   %[[R64:.*]] = polynomial.eval #polynomial.typed_chebyshev_polynomial<[{{.*}}]> : !polynomial.polynomial<ring = <coefficientType = f64>>, %[[X64]] {scheme = "clenshaw", domain_lower = {{.*}} : f64, domain_upper = {{.*}} : f64} : f64
+// CHECK-NEXT:   func.return %[[R64]] : f64
+// CHECK-NEXT: }
+
+// Without explicit lower/upper bounds, the polynomial domain defaults to
+// [underflow, overflow] for the precision (f32: ~[-103.28, 88.72]).
+// CHECK:      func.func @exp_f32_with_max_bits_lost(%[[XB:.*]]: f32) -> f32 {
+// CHECK-NEXT:   %[[RB:.*]] = polynomial.eval #polynomial.typed_chebyshev_polynomial<[{{.*}}]> : !polynomial.polynomial<ring = <coefficientType = f32>>, %[[XB]] {scheme = "clenshaw", domain_lower = {{.*}} : f32, domain_upper = {{.*}} : f32} : f32
+// CHECK-NEXT:   func.return %[[RB]] : f32
+// CHECK-NEXT: }
+
+// CHECK:      func.func @exp_f32_clamped_lower(%[[XC:.*]]: f32) -> f32 {
+// CHECK-NEXT:   %[[RC:.*]] = polynomial.eval #polynomial.typed_chebyshev_polynomial<[{{.*}}]> : !polynomial.polynomial<ring = <coefficientType = f32>>, %[[XC]] {scheme = "clenshaw", domain_lower = -2.000000e+00 : f32, domain_upper = 5.000000e+00 : f32} : f32
+// CHECK-NEXT:   func.return %[[RC]] : f32
+// CHECK-NEXT: }
+
+// CHECK:      func.func @exp_f32_tight_bounds(%[[XT:.*]]: f32) -> f32 {
+// CHECK-NEXT:   %[[RT:.*]] = polynomial.eval #polynomial.typed_chebyshev_polynomial<[{{.*}}]> : !polynomial.polynomial<ring = <coefficientType = f32>>, %[[XT]] {scheme = "clenshaw", domain_lower = -5.000000e-01 : f32, domain_upper = 5.000000e-01 : f32} : f32
+// CHECK-NEXT:   func.return %[[RT]] : f32
+// CHECK-NEXT: }

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_lower_exp_to_polynomial.py

@@ -0,0 +1,191 @@
+import pytest
+
+from xdsl.context import Context
+from xdsl.dialects import math as math_dialect
+from xdsl.dialects import polynomial
+from xdsl.dialects.builtin import (
+    Float16Type,
+    Float32Type,
+    Float64Type,
+    FloatAttr,
+    IntegerAttr,
+    ModuleOp,
+    i64,
+)
+from xdsl.transforms.lower_exp_to_polynomial import (
+    OVERFLOW_UPPER_BOUND,
+    UNDERFLOW_LOWER_BOUND,
+    LowerExpToPolynomialPass,
+)
+from xdsl.utils.test_value import create_ssa_value
+
+_FloatTy = Float16Type | Float32Type | Float64Type
+
+PRECISIONS = pytest.mark.parametrize(
+    "elem_ty", [Float16Type(), Float32Type(), Float64Type()]
+)
+
+
+def _underflow_at(elem_ty: _FloatTy) -> float:
+    """UNDERFLOW_LOWER_BOUND rounded to elem_ty's precision (matches what
+    the pass stores in `domain_lower` after `FloatAttr(_, elem_ty)`)."""
+    return FloatAttr(UNDERFLOW_LOWER_BOUND[type(elem_ty)], elem_ty).value.data
+
+
+def _overflow_at(elem_ty: _FloatTy) -> float:
+    """OVERFLOW_UPPER_BOUND rounded to elem_ty's precision."""
+    return FloatAttr(OVERFLOW_UPPER_BOUND[type(elem_ty)], elem_ty).value.data
+
+
+def _run_pass(elem_ty: _FloatTy, attrs: dict[str, int | float]) -> ModuleOp:
+    operand = create_ssa_value(elem_ty)
+    exp_op = math_dialect.ExpOp(operand)
+    for name, value in attrs.items():
+        if name == "max_bits_lost":
+            exp_op.attributes[name] = IntegerAttr(int(value), i64)
+        else:
+            exp_op.attributes[name] = FloatAttr(float(value), elem_ty)
+    module = ModuleOp([exp_op])
+    ctx = Context()
+    ctx.load_op(ModuleOp)
+    LowerExpToPolynomialPass().apply(ctx, module)
+    return module
+
+
+@PRECISIONS
+def test_default_no_bounds(elem_ty: _FloatTy):
+    """No max_bits_lost, no [lower, upper] -> lowered with default target
+    (correctly-rounded) and default [underflow, overflow] domain."""
+    module = _run_pass(elem_ty, {})
+    eval_op = next(iter(module.body.block.ops))
+    assert isinstance(eval_op, polynomial.EvalOp)
+    assert eval_op.domain_lower is not None
+    assert eval_op.domain_upper is not None
+    assert eval_op.domain_lower.value.data == _underflow_at(elem_ty)
+    assert eval_op.domain_upper.value.data == _overflow_at(elem_ty)
+
+
+@PRECISIONS
+def test_default_with_bounds(elem_ty: _FloatTy):
+    """No max_bits_lost, [lower, upper] given -> lowered with default target
+    and the user-supplied domain."""
+    module = _run_pass(elem_ty, {"lower_bound": -0.5, "upper_bound": 0.5})
+    eval_op = next(iter(module.body.block.ops))
+    assert isinstance(eval_op, polynomial.EvalOp)
+    assert eval_op.domain_lower is not None
+    assert eval_op.domain_upper is not None
+    assert eval_op.domain_lower.value.data == -0.5
+    assert eval_op.domain_upper.value.data == 0.5
+
+
+@PRECISIONS
+def test_max_bits_lost_only(elem_ty: _FloatTy):
+    """max_bits_lost only -> default interval [underflow, overflow]."""
+    module = _run_pass(elem_ty, {"max_bits_lost": 2})
+    eval_op = next(iter(module.body.block.ops))
+    assert isinstance(eval_op, polynomial.EvalOp)
+    assert eval_op.domain_lower is not None
+    assert eval_op.domain_upper is not None
+    assert eval_op.domain_lower.value.data == _underflow_at(elem_ty)
+    assert eval_op.domain_upper.value.data == _overflow_at(elem_ty)
+
+
+@PRECISIONS
+def test_max_bits_lost_bounds_in_range(elem_ty: _FloatTy):
+    """max_bits_lost + bounds entirely in range -> bounds used as-is."""
+    module = _run_pass(
+        elem_ty,
+        {"max_bits_lost": 2, "lower_bound": -0.5, "upper_bound": 0.5},
+    )
+    eval_op = next(iter(module.body.block.ops))
+    assert isinstance(eval_op, polynomial.EvalOp)
+    assert eval_op.domain_lower is not None
+    assert eval_op.domain_upper is not None
+    assert eval_op.domain_lower.value.data == -0.5
+    assert eval_op.domain_upper.value.data == 0.5
+
+
+@PRECISIONS
+def test_max_bits_lost_lower_out_of_range(elem_ty: _FloatTy):
+    """max_bits_lost + lower < underflow -> lower clamped to underflow.
+
+    -1000 is below every supported precision's underflow (f16 ~= -16.64,
+    bf16 ~= -92.18, f32 ~= -103.28, f64 ~= -744.44).
+    """
+    module = _run_pass(
+        elem_ty,
+        {"max_bits_lost": 2, "lower_bound": -1000.0, "upper_bound": 0.5},
+    )
+    eval_op = next(iter(module.body.block.ops))
+    assert isinstance(eval_op, polynomial.EvalOp)
+    assert eval_op.domain_lower is not None
+    assert eval_op.domain_upper is not None
+    assert eval_op.domain_lower.value.data == _underflow_at(elem_ty)
+    assert eval_op.domain_upper.value.data == 0.5
+
+
+@PRECISIONS
+def test_max_bits_lost_upper_out_of_range(elem_ty: _FloatTy):
+    """max_bits_lost + upper > overflow -> upper clamped to overflow.
+
+    1000 exceeds every supported precision's overflow (f16 ~= 11.09,
+    f32/bf16 ~= 88.72, f64 ~= 709.78).
+    """
+    module = _run_pass(
+        elem_ty,
+        {"max_bits_lost": 2, "lower_bound": -0.5, "upper_bound": 1000.0},
+    )
+    eval_op = next(iter(module.body.block.ops))
+    assert isinstance(eval_op, polynomial.EvalOp)
+    assert eval_op.domain_lower is not None
+    assert eval_op.domain_upper is not None
+    assert eval_op.domain_lower.value.data == -0.5
+    assert eval_op.domain_upper.value.data == _overflow_at(elem_ty)
+
+
+@PRECISIONS
+def test_max_bits_lost_lower_only_in_range(elem_ty: _FloatTy):
+    """max_bits_lost + lower only (in range) -> [lower, overflow]."""
+    module = _run_pass(elem_ty, {"max_bits_lost": 2, "lower_bound": -0.5})
+    eval_op = next(iter(module.body.block.ops))
+    assert isinstance(eval_op, polynomial.EvalOp)
+    assert eval_op.domain_lower is not None
+    assert eval_op.domain_upper is not None
+    assert eval_op.domain_lower.value.data == -0.5
+    assert eval_op.domain_upper.value.data == _overflow_at(elem_ty)
+
+
+@PRECISIONS
+def test_max_bits_lost_lower_only_out_of_range(elem_ty: _FloatTy):
+    """max_bits_lost + lower only (out of range) -> [underflow, overflow]."""
+    module = _run_pass(elem_ty, {"max_bits_lost": 2, "lower_bound": -1000.0})
+    eval_op = next(iter(module.body.block.ops))
+    assert isinstance(eval_op, polynomial.EvalOp)
+    assert eval_op.domain_lower is not None
+    assert eval_op.domain_upper is not None
+    assert eval_op.domain_lower.value.data == _underflow_at(elem_ty)
+    assert eval_op.domain_upper.value.data == _overflow_at(elem_ty)
+
+
+@PRECISIONS
+def test_max_bits_lost_upper_only_in_range(elem_ty: _FloatTy):
+    """max_bits_lost + upper only (in range) -> [underflow, upper]."""
+    module = _run_pass(elem_ty, {"max_bits_lost": 2, "upper_bound": 0.5})
+    eval_op = next(iter(module.body.block.ops))
+    assert isinstance(eval_op, polynomial.EvalOp)
+    assert eval_op.domain_lower is not None
+    assert eval_op.domain_upper is not None
+    assert eval_op.domain_lower.value.data == _underflow_at(elem_ty)
+    assert eval_op.domain_upper.value.data == 0.5
+
+
+@PRECISIONS
+def test_max_bits_lost_upper_only_out_of_range(elem_ty: _FloatTy):
+    """max_bits_lost + upper only (out of range) -> [underflow, overflow]."""
+    module = _run_pass(elem_ty, {"max_bits_lost": 2, "upper_bound": 1000.0})
+    eval_op = next(iter(module.body.block.ops))
+    assert isinstance(eval_op, polynomial.EvalOp)
+    assert eval_op.domain_lower is not None
+    assert eval_op.domain_upper is not None
+    assert eval_op.domain_lower.value.data == _underflow_at(elem_ty)
+    assert eval_op.domain_upper.value.data == _overflow_at(elem_ty)