From 5323843564756a116311cb59b57152a539fe9746 Mon Sep 17 00:00:00 2001 From: szerdick Date: Mon, 4 May 2026 18:09:18 +0100 Subject: [PATCH 01/28] add polynomial ops and heir compat tests --- tests/dialects/test_polynomial.py | 84 ++++++++++ tests/filecheck/dialects/polynomial/ops.mlir | 36 ++++ .../dialects/polynomial/ops_invalid.mlir | 31 ++++ .../with-heir/dialects/polynomial/ops.mlir | 1 + xdsl/dialects/polynomial.py | 156 +++++++++++++++++- 5 files changed, 307 insertions(+), 1 deletion(-) create mode 100644 tests/filecheck/dialects/polynomial/ops.mlir create mode 100644 tests/filecheck/dialects/polynomial/ops_invalid.mlir create mode 100644 tests/filecheck/with-heir/dialects/polynomial/ops.mlir diff --git a/tests/dialects/test_polynomial.py b/tests/dialects/test_polynomial.py index 5d47a10885..09056be296 100644 --- a/tests/dialects/test_polynomial.py +++ b/tests/dialects/test_polynomial.py @@ -3,14 +3,23 @@ from xdsl.dialects.builtin import ( ArrayAttr, FloatAttr, + StringAttr, + TensorType, + VectorType, + f16, + f32, f64, ) from xdsl.dialects.polynomial import ( ChebyshevPolynomialAttr, + EvalOp, + EvalScheme, PolynomialType, RingAttr, TypedChebyshevPolynomialAttr, ) +from xdsl.ir import Attribute +from xdsl.utils.test_value import create_ssa_value # --- ChebyshevPolynomialAttr construction --- @@ -64,3 +73,78 @@ def test_typed_chebyshev_polynomial_from_attr(): inner = ChebyshevPolynomialAttr((1.0, 2.0)) typed = TypedChebyshevPolynomialAttr(poly_ty, inner) assert typed.value is inner + + +# --- EvalOp construction --- + + +def test_eval_basic_construction(): + x = create_ssa_value(f32) + op = EvalOp(x, (0.5, 1.2, 0.3), EvalScheme.CLENSHAW) + assert op.value.type == f32 + assert op.result.type == f32 + assert op.degree == 2 + assert op.polynomial.coeff_values == (0.5, 1.2, 0.3) + assert op.scheme.data == EvalScheme.CLENSHAW.value + assert op.eval_scheme == EvalScheme.CLENSHAW + + +def test_eval_pre_built_typed_attr(): + poly_ty = PolynomialType(RingAttr(f64)) + typed = TypedChebyshevPolynomialAttr(poly_ty, (1.0, 2.0, 3.0)) + x = create_ssa_value(f64) + op = EvalOp(x, typed, EvalScheme.CLENSHAW) + assert op.polynomial is typed + + +def test_eval_pre_built_untyped_attr_is_wrapped(): + """An untyped ChebyshevPolynomialAttr is auto-wrapped with the default type.""" + untyped = ChebyshevPolynomialAttr((1.0, 2.0, 3.0)) + x = create_ssa_value(f64) + op = EvalOp(x, untyped, EvalScheme.CLENSHAW) + assert isinstance(op.polynomial, TypedChebyshevPolynomialAttr) + assert op.polynomial.value is untyped + + +@pytest.mark.parametrize( + "tp", + [f16, f32, f64, VectorType(f32, [4]), TensorType(f64, [8])], +) +def test_eval_type_polymorphism(tp: Attribute): + x = create_ssa_value(tp) + op = EvalOp(x, (1.0, 2.0), EvalScheme.CLENSHAW) + assert op.value.type == tp + assert op.result.type == tp + + +def test_eval_domain_bounds_propagated(): + x = create_ssa_value(f32) + op = EvalOp( + x, (1.0, 2.0), EvalScheme.CLENSHAW, domain_lower=-10.0, domain_upper=0.0 + ) + assert op.domain_lower is not None + assert op.domain_upper is not None + assert op.domain_lower.value.data == -10.0 + assert op.domain_upper.value.data == 0.0 + + +def test_eval_scheme_propagated(): + x = create_ssa_value(f32) + op = EvalOp(x, (1.0, 2.0), EvalScheme.CLENSHAW) + assert op.scheme.data == EvalScheme.CLENSHAW.value + assert op.eval_scheme == EvalScheme.CLENSHAW + + +def test_eval_scheme_accepts_string_attr(): + x = create_ssa_value(f32) + attr = StringAttr("clenshaw") + op = EvalOp(x, (1.0, 2.0), attr) + assert op.scheme is attr + assert op.eval_scheme == EvalScheme.CLENSHAW + + +def test_eval_scheme_accepts_str(): + x = create_ssa_value(f32) + op = EvalOp(x, (1.0, 2.0), "clenshaw") + assert op.scheme.data == "clenshaw" + assert op.eval_scheme == EvalScheme.CLENSHAW diff --git a/tests/filecheck/dialects/polynomial/ops.mlir b/tests/filecheck/dialects/polynomial/ops.mlir new file mode 100644 index 0000000000..f20c597aed --- /dev/null +++ b/tests/filecheck/dialects/polynomial/ops.mlir @@ -0,0 +1,36 @@ +// RUN: XDSL_ROUNDTRIP + +// CHECK: builtin.module { + +// CHECK-NEXT: %[[X_F32:.*]] = "test.op"() : () -> f32 +// CHECK-NEXT: %[[X_F16:.*]] = "test.op"() : () -> f16 +// CHECK-NEXT: %[[X_F64:.*]] = "test.op"() : () -> f64 +// CHECK-NEXT: %[[X_VEC:.*]] = "test.op"() : () -> vector<4xf32> +// CHECK-NEXT: %[[X_TENSOR:.*]] = "test.op"() : () -> tensor<8xf64> +%x_f32 = "test.op"() : () -> f32 +%x_f16 = "test.op"() : () -> f16 +%x_f64 = "test.op"() : () -> f64 +%x_vec = "test.op"() : () -> vector<4xf32> +%x_tensor = "test.op"() : () -> tensor<8xf64> + +// Scalar f32, degree 2, with domain bounds +// CHECK-NEXT: %{{.*}} = polynomial.eval #polynomial.typed_chebyshev_polynomial<[5.000000e-01 : f64, 1.200000e+00 : f64, 3.000000e-01 : f64]> : !polynomial.polynomial>, %[[X_F32]] {domain_lower = -1.000000e+00 : f64, domain_upper = 1.000000e+00 : f64, scheme = "clenshaw"} : f32 +%r0 = polynomial.eval #polynomial.typed_chebyshev_polynomial<[5.000000e-01 : f64, 1.200000e+00 : f64, 3.000000e-01 : f64]> : !polynomial.polynomial>, %x_f32 {domain_lower = -1.000000e+00 : f64, domain_upper = 1.000000e+00 : f64, scheme = "clenshaw"} : f32 + +// Scalar f16, no domain bounds +// CHECK-NEXT: %{{.*}} = polynomial.eval #polynomial.typed_chebyshev_polynomial<[1.000000e+00 : f64, 2.000000e+00 : f64]> : !polynomial.polynomial>, %[[X_F16]] {scheme = "clenshaw"} : f16 +%r1 = polynomial.eval #polynomial.typed_chebyshev_polynomial<[1.000000e+00 : f64, 2.000000e+00 : f64]> : !polynomial.polynomial>, %x_f16 {scheme = "clenshaw"} : f16 + +// Scalar f64, custom domain [-10, 0] +// CHECK-NEXT: %{{.*}} = polynomial.eval #polynomial.typed_chebyshev_polynomial<[1.000000e+00 : f64, 2.000000e+00 : f64, 3.000000e+00 : f64]> : !polynomial.polynomial>, %[[X_F64]] {domain_lower = -1.000000e+01 : f64, domain_upper = 0.000000e+00 : f64, scheme = "clenshaw"} : f64 +%r2 = polynomial.eval #polynomial.typed_chebyshev_polynomial<[1.000000e+00 : f64, 2.000000e+00 : f64, 3.000000e+00 : f64]> : !polynomial.polynomial>, %x_f64 {domain_lower = -1.000000e+01 : f64, domain_upper = 0.000000e+00 : f64, scheme = "clenshaw"} : f64 + +// Vector type, no domain bounds +// CHECK-NEXT: %{{.*}} = polynomial.eval #polynomial.typed_chebyshev_polynomial<[1.000000e+00 : f64, 2.000000e+00 : f64]> : !polynomial.polynomial>, %[[X_VEC]] {scheme = "clenshaw"} : vector<4xf32> +%r3 = polynomial.eval #polynomial.typed_chebyshev_polynomial<[1.000000e+00 : f64, 2.000000e+00 : f64]> : !polynomial.polynomial>, %x_vec {scheme = "clenshaw"} : vector<4xf32> + +// Tensor type, with domain bounds +// CHECK-NEXT: %{{.*}} = polynomial.eval #polynomial.typed_chebyshev_polynomial<[5.000000e-01 : f64, 1.500000e+00 : f64]> : !polynomial.polynomial>, %[[X_TENSOR]] {domain_lower = -2.000000e+00 : f64, domain_upper = 2.000000e+00 : f64, scheme = "clenshaw"} : tensor<8xf64> +%r4 = polynomial.eval #polynomial.typed_chebyshev_polynomial<[5.000000e-01 : f64, 1.500000e+00 : f64]> : !polynomial.polynomial>, %x_tensor {domain_lower = -2.000000e+00 : f64, domain_upper = 2.000000e+00 : f64, scheme = "clenshaw"} : tensor<8xf64> + +// CHECK-NEXT: } diff --git a/tests/filecheck/dialects/polynomial/ops_invalid.mlir b/tests/filecheck/dialects/polynomial/ops_invalid.mlir new file mode 100644 index 0000000000..8a395a2834 --- /dev/null +++ b/tests/filecheck/dialects/polynomial/ops_invalid.mlir @@ -0,0 +1,31 @@ +// RUN: xdsl-opt %s --verify-diagnostics --split-input-file | filecheck %s + +// domain_lower > domain_upper +%x0 = "test.op"() : () -> f32 +%r0 = polynomial.eval #polynomial.typed_chebyshev_polynomial<[1.000000e+00 : f64, 2.000000e+00 : f64]> : !polynomial.polynomial>, %x0 {scheme = "clenshaw", domain_lower = 1.000000e+00 : f64, domain_upper = -1.000000e+00 : f64} : f32 + +// CHECK: domain_lower (1.0) must be strictly less than domain_upper (-1.0) + +// ----- + +// domain_lower == domain_upper +%x1 = "test.op"() : () -> f32 +%r1 = polynomial.eval #polynomial.typed_chebyshev_polynomial<[1.000000e+00 : f64, 2.000000e+00 : f64]> : !polynomial.polynomial>, %x1 {scheme = "clenshaw", domain_lower = 0.000000e+00 : f64, domain_upper = 0.000000e+00 : f64} : f32 + +// CHECK: domain_lower (0.0) must be strictly less than domain_upper (0.0) + +// ----- + +// degree 0 (single coefficient) +%x2 = "test.op"() : () -> f32 +%r2 = polynomial.eval #polynomial.typed_chebyshev_polynomial<[1.000000e+00 : f64]> : !polynomial.polynomial>, %x2 {scheme = "clenshaw"} : f32 + +// CHECK: Chebyshev polynomial must have at least degree 1 (got 1 coefficients) + +// ----- + +// unknown evaluation scheme +%x3 = "test.op"() : () -> f32 +%r3 = polynomial.eval #polynomial.typed_chebyshev_polynomial<[1.000000e+00 : f64, 2.000000e+00 : f64]> : !polynomial.polynomial>, %x3 {scheme = "horner"} : f32 + +// CHECK: unknown evaluation scheme 'horner' diff --git a/tests/filecheck/with-heir/dialects/polynomial/ops.mlir b/tests/filecheck/with-heir/dialects/polynomial/ops.mlir new file mode 100644 index 0000000000..0f9b8bb09b --- /dev/null +++ b/tests/filecheck/with-heir/dialects/polynomial/ops.mlir @@ -0,0 +1 @@ +// RUN: xdsl-opt %S/../../../dialects/polynomial/ops.mlir | heir-opt --allow-unregistered-dialect | xdsl-opt | filecheck %S/../../../dialects/polynomial/ops.mlir diff --git a/xdsl/dialects/polynomial.py b/xdsl/dialects/polynomial.py index 11b6e458d7..216219589b 100644 --- a/xdsl/dialects/polynomial.py +++ b/xdsl/dialects/polynomial.py @@ -13,23 +13,53 @@ from __future__ import annotations from collections.abc import Sequence +from typing import ClassVar from xdsl.dialects.builtin import ( + AnyFloatConstr, ArrayAttr, + ContainerOf, FloatAttr, + StringAttr, f64, ) from xdsl.ir import ( Attribute, Dialect, + Operation, ParametrizedAttribute, + SSAValue, TypeAttribute, ) from xdsl.irdl import ( + IRDLOperation, + VarConstraint, + attr_def, irdl_attr_definition, + irdl_op_definition, + operand_def, + opt_attr_def, + prop_def, + result_def, + traits_def, ) from xdsl.parser import AttrParser from xdsl.printer import Printer +from xdsl.traits import Pure, SameOperandsAndResultType +from xdsl.utils.exceptions import VerifyException +from xdsl.utils.str_enum import StrEnum + + +class EvalScheme(StrEnum): + """ + Evaluation scheme used to lower a `polynomial.eval` op to arithmetic ops. + + Stored on the op as a builtin string attribute so that compilers like HEIR + that don't know about this xDSL extension can preserve it on round-trip + (the same way they ignore `domain_lower`/`domain_upper`). + """ + + CLENSHAW = "clenshaw" @irdl_attr_definition @@ -181,9 +211,133 @@ def coeff_values(self) -> tuple[float, ...]: return self.value.coeff_values +def _default_polynomial_type() -> PolynomialType: + """Default polynomial type for f64 Chebyshev coefficients.""" + return PolynomialType(RingAttr(f64)) + + +@irdl_op_definition +class EvalOp(IRDLOperation): + """ + Evaluate a polynomial at a given point. + + This op is *unevaluated* but carries all information needed for + later lowering to arithmetic ops, dispatched on `scheme`. + + Syntax: polynomial.eval $polynomial `,` $value attr-dict `:` type($value) + Example: + %result = polynomial.eval + #polynomial.typed_chebyshev_polynomial<[0.5, 1.2]> : + !polynomial.polynomial>, + %x {scheme = "clenshaw", + domain_lower = -1.0 : f64, domain_upper = 1.0 : f64} + : f32 + """ + + name = "polynomial.eval" + + T: ClassVar = VarConstraint("T", ContainerOf(AnyFloatConstr)) + + value = operand_def(T) + result = result_def(T) + + polynomial = prop_def(TypedChebyshevPolynomialAttr) + + scheme = attr_def(StringAttr) + domain_lower = opt_attr_def(FloatAttr) + domain_upper = opt_attr_def(FloatAttr) + + traits = traits_def(Pure(), SameOperandsAndResultType()) + + # `qualified(...)` forces the full `#polynomial.typed_chebyshev_polynomial<...>` + # form when printing/parsing, instead of the elided `<...>` form. This is + # required for HEIR round-trip: HEIR's parser only accepts the qualified form. + assembly_format = "qualified($polynomial) `,` $value attr-dict `:` type($value)" + + def __init__( + self, + value: Operation | SSAValue, + polynomial: ( + TypedChebyshevPolynomialAttr | ChebyshevPolynomialAttr | tuple[float, ...] + ), + scheme: StringAttr | EvalScheme | str, + domain_lower: float | FloatAttr | None = None, + domain_upper: float | FloatAttr | None = None, + ): + if isinstance(polynomial, tuple): + polynomial = TypedChebyshevPolynomialAttr( + _default_polynomial_type(), polynomial + ) + elif isinstance(polynomial, ChebyshevPolynomialAttr): + polynomial = TypedChebyshevPolynomialAttr( + _default_polynomial_type(), polynomial + ) + + if isinstance(scheme, EvalScheme): + scheme = StringAttr(scheme.value) + elif isinstance(scheme, str): + scheme = StringAttr(scheme) + + value = SSAValue.get(value) + + if isinstance(domain_lower, (int, float)): + domain_lower = FloatAttr(float(domain_lower), f64) + if isinstance(domain_upper, (int, float)): + domain_upper = FloatAttr(float(domain_upper), f64) + + attrs: dict[str, StringAttr | FloatAttr] = {"scheme": scheme} + if domain_lower is not None: + attrs["domain_lower"] = domain_lower + if domain_upper is not None: + attrs["domain_upper"] = domain_upper + + super().__init__( + operands=[value], + result_types=[value.type], + properties={ + "polynomial": polynomial, + }, + attributes=attrs, + ) + + def verify_(self) -> None: + if self.domain_lower is not None and self.domain_upper is not None: + lower = self.domain_lower.value.data + upper = self.domain_upper.value.data + if lower >= upper: + raise VerifyException( + f"domain_lower ({lower}) must be strictly " + f"less than domain_upper ({upper})" + ) + if self.polynomial.degree < 1: + raise VerifyException( + "Chebyshev polynomial must have at least degree 1 " + f"(got {self.polynomial.degree + 1} coefficients)" + ) + try: + EvalScheme(self.scheme.data) + except ValueError: + valid = ", ".join(repr(s.value) for s in EvalScheme) + raise VerifyException( + f"unknown evaluation scheme {self.scheme.data!r}; " + f"expected one of: {valid}" + ) + + @property + def degree(self) -> int: + return self.polynomial.degree + + @property + def eval_scheme(self) -> EvalScheme: + """Return the scheme as an EvalScheme enum.""" + return EvalScheme(self.scheme.data) + + Polynomial = Dialect( "polynomial", - [], + [ + EvalOp, + ], [ ChebyshevPolynomialAttr, TypedChebyshevPolynomialAttr, From 6eee7182471538abe41aa5944ee51cf12d92fd88 Mon Sep 17 00:00:00 2001 From: szerdick Date: Wed, 13 May 2026 11:29:05 +0200 Subject: [PATCH 02/28] make default polynomialtype macro --- xdsl/dialects/polynomial.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/xdsl/dialects/polynomial.py b/xdsl/dialects/polynomial.py index 216219589b..4317f3b108 100644 --- a/xdsl/dialects/polynomial.py +++ b/xdsl/dialects/polynomial.py @@ -211,9 +211,8 @@ def coeff_values(self) -> tuple[float, ...]: return self.value.coeff_values -def _default_polynomial_type() -> PolynomialType: - """Default polynomial type for f64 Chebyshev coefficients.""" - return PolynomialType(RingAttr(f64)) +_DEFAULT_POLYNOMIAL_TYPE = PolynomialType(RingAttr(f64)) +"""Default polynomial type for f64 Chebyshev coefficients.""" @irdl_op_definition @@ -266,11 +265,11 @@ def __init__( ): if isinstance(polynomial, tuple): polynomial = TypedChebyshevPolynomialAttr( - _default_polynomial_type(), polynomial + _DEFAULT_POLYNOMIAL_TYPE, polynomial ) elif isinstance(polynomial, ChebyshevPolynomialAttr): polynomial = TypedChebyshevPolynomialAttr( - _default_polynomial_type(), polynomial + _DEFAULT_POLYNOMIAL_TYPE, polynomial ) if isinstance(scheme, EvalScheme): From 594db9d825e4469277ce40c9b10b04cde592f7b6 Mon Sep 17 00:00:00 2001 From: szerdick Date: Wed, 13 May 2026 12:14:53 +0200 Subject: [PATCH 03/28] get method --- tests/dialects/test_polynomial.py | 86 ++++++++++++++----------------- xdsl/dialects/polynomial.py | 82 ++++++++++++++++------------- 2 files changed, 85 insertions(+), 83 deletions(-) diff --git a/tests/dialects/test_polynomial.py b/tests/dialects/test_polynomial.py index 09056be296..f482dd2be9 100644 --- a/tests/dialects/test_polynomial.py +++ b/tests/dialects/test_polynomial.py @@ -78,73 +78,67 @@ def test_typed_chebyshev_polynomial_from_attr(): # --- EvalOp construction --- -def test_eval_basic_construction(): +def test_eval_init_stores_attributes_verbatim(): + """`__init__` takes pre-built attributes and stores them as-is.""" + poly_ty = PolynomialType(RingAttr(f64)) + typed = TypedChebyshevPolynomialAttr(poly_ty, (1.0, 2.0, 3.0)) + scheme_attr = StringAttr("clenshaw") + x = create_ssa_value(f64) + op = EvalOp(x, typed, scheme_attr) + assert op.polynomial is typed + assert op.scheme is scheme_attr + + +def test_eval_get_basic_construction(): x = create_ssa_value(f32) - op = EvalOp(x, (0.5, 1.2, 0.3), EvalScheme.CLENSHAW) + op = EvalOp.get(x, (0.5, 1.2, 0.3), f32, EvalScheme.CLENSHAW) assert op.value.type == f32 assert op.result.type == f32 assert op.degree == 2 assert op.polynomial.coeff_values == (0.5, 1.2, 0.3) + assert op.polynomial.type == PolynomialType(RingAttr(f32)) assert op.scheme.data == EvalScheme.CLENSHAW.value assert op.eval_scheme == EvalScheme.CLENSHAW -def test_eval_pre_built_typed_attr(): - poly_ty = PolynomialType(RingAttr(f64)) - typed = TypedChebyshevPolynomialAttr(poly_ty, (1.0, 2.0, 3.0)) - x = create_ssa_value(f64) - op = EvalOp(x, typed, EvalScheme.CLENSHAW) - assert op.polynomial is typed - - -def test_eval_pre_built_untyped_attr_is_wrapped(): - """An untyped ChebyshevPolynomialAttr is auto-wrapped with the default type.""" - untyped = ChebyshevPolynomialAttr((1.0, 2.0, 3.0)) - x = create_ssa_value(f64) - op = EvalOp(x, untyped, EvalScheme.CLENSHAW) - assert isinstance(op.polynomial, TypedChebyshevPolynomialAttr) - assert op.polynomial.value is untyped - - @pytest.mark.parametrize( - "tp", - [f16, f32, f64, VectorType(f32, [4]), TensorType(f64, [8])], + "value_type, element_type", + [ + (f16, f16), + (f32, f32), + (f64, f64), + (VectorType(f32, [4]), f32), + (TensorType(f64, [8]), f64), + ], ) -def test_eval_type_polymorphism(tp: Attribute): - x = create_ssa_value(tp) - op = EvalOp(x, (1.0, 2.0), EvalScheme.CLENSHAW) - assert op.value.type == tp - assert op.result.type == tp +def test_eval_get_type_polymorphism(value_type: Attribute, element_type: Attribute): + x = create_ssa_value(value_type) + op = EvalOp.get(x, (1.0, 2.0), element_type, EvalScheme.CLENSHAW) + assert op.value.type == value_type + assert op.result.type == value_type + assert op.polynomial.type == PolynomialType(RingAttr(element_type)) -def test_eval_domain_bounds_propagated(): +def test_eval_get_domain_bounds_use_element_type(): x = create_ssa_value(f32) - op = EvalOp( - x, (1.0, 2.0), EvalScheme.CLENSHAW, domain_lower=-10.0, domain_upper=0.0 + op = EvalOp.get( + x, + (1.0, 2.0), + f32, + EvalScheme.CLENSHAW, + domain_lower=-10.0, + domain_upper=0.0, ) assert op.domain_lower is not None assert op.domain_upper is not None assert op.domain_lower.value.data == -10.0 assert op.domain_upper.value.data == 0.0 + assert op.domain_lower.type == f32 + assert op.domain_upper.type == f32 -def test_eval_scheme_propagated(): - x = create_ssa_value(f32) - op = EvalOp(x, (1.0, 2.0), EvalScheme.CLENSHAW) - assert op.scheme.data == EvalScheme.CLENSHAW.value - assert op.eval_scheme == EvalScheme.CLENSHAW - - -def test_eval_scheme_accepts_string_attr(): - x = create_ssa_value(f32) - attr = StringAttr("clenshaw") - op = EvalOp(x, (1.0, 2.0), attr) - assert op.scheme is attr - assert op.eval_scheme == EvalScheme.CLENSHAW - - -def test_eval_scheme_accepts_str(): +def test_eval_get_accepts_str_scheme(): x = create_ssa_value(f32) - op = EvalOp(x, (1.0, 2.0), "clenshaw") + op = EvalOp.get(x, (1.0, 2.0), f32, "clenshaw") assert op.scheme.data == "clenshaw" assert op.eval_scheme == EvalScheme.CLENSHAW diff --git a/xdsl/dialects/polynomial.py b/xdsl/dialects/polynomial.py index 4317f3b108..8809281b95 100644 --- a/xdsl/dialects/polynomial.py +++ b/xdsl/dialects/polynomial.py @@ -211,10 +211,6 @@ def coeff_values(self) -> tuple[float, ...]: return self.value.coeff_values -_DEFAULT_POLYNOMIAL_TYPE = PolynomialType(RingAttr(f64)) -"""Default polynomial type for f64 Chebyshev coefficients.""" - - @irdl_op_definition class EvalOp(IRDLOperation): """ @@ -255,40 +251,17 @@ class EvalOp(IRDLOperation): def __init__( self, - value: Operation | SSAValue, - polynomial: ( - TypedChebyshevPolynomialAttr | ChebyshevPolynomialAttr | tuple[float, ...] - ), - scheme: StringAttr | EvalScheme | str, - domain_lower: float | FloatAttr | None = None, - domain_upper: float | FloatAttr | None = None, + value: SSAValue, + polynomial: TypedChebyshevPolynomialAttr, + scheme: StringAttr, + domain_lower: FloatAttr | None = None, + domain_upper: FloatAttr | None = None, ): - if isinstance(polynomial, tuple): - polynomial = TypedChebyshevPolynomialAttr( - _DEFAULT_POLYNOMIAL_TYPE, polynomial - ) - elif isinstance(polynomial, ChebyshevPolynomialAttr): - polynomial = TypedChebyshevPolynomialAttr( - _DEFAULT_POLYNOMIAL_TYPE, polynomial - ) - - if isinstance(scheme, EvalScheme): - scheme = StringAttr(scheme.value) - elif isinstance(scheme, str): - scheme = StringAttr(scheme) - - value = SSAValue.get(value) - - if isinstance(domain_lower, (int, float)): - domain_lower = FloatAttr(float(domain_lower), f64) - if isinstance(domain_upper, (int, float)): - domain_upper = FloatAttr(float(domain_upper), f64) - - attrs: dict[str, StringAttr | FloatAttr] = {"scheme": scheme} - if domain_lower is not None: - attrs["domain_lower"] = domain_lower - if domain_upper is not None: - attrs["domain_upper"] = domain_upper + attrs = { + "scheme": scheme, + "domain_lower": domain_lower, + "domain_upper": domain_upper, + } super().__init__( operands=[value], @@ -299,6 +272,41 @@ def __init__( attributes=attrs, ) + @classmethod + def get( + cls, + value: Operation | SSAValue, + coefficients: tuple[float, ...], + element_type: Attribute, + scheme: EvalScheme | str, + domain_lower: float | None = None, + domain_upper: float | None = None, + ) -> EvalOp: + """ + Build an `EvalOp` from raw Python floats. + + `element_type` is the floating-point type used for the polynomial + coefficients, the ring's coefficient type, and the domain bounds. + """ + value = SSAValue.get(value) + + coeff_attrs = ArrayAttr([FloatAttr(c, element_type) for c in coefficients]) + polynomial = TypedChebyshevPolynomialAttr( + PolynomialType(RingAttr(element_type)), + ChebyshevPolynomialAttr(coeff_attrs), + ) + + scheme_attr = StringAttr( + scheme.value if isinstance(scheme, EvalScheme) else scheme + ) + lower_attr = ( + FloatAttr(domain_lower, element_type) if domain_lower is not None else None + ) + upper_attr = ( + FloatAttr(domain_upper, element_type) if domain_upper is not None else None + ) + return cls(value, polynomial, scheme_attr, lower_attr, upper_attr) + def verify_(self) -> None: if self.domain_lower is not None and self.domain_upper is not None: lower = self.domain_lower.value.data From e31e996a3673772ffa24e233ce13a9847e761279 Mon Sep 17 00:00:00 2001 From: szerdick Date: Wed, 13 May 2026 13:25:05 +0200 Subject: [PATCH 04/28] edit test to match numbers --- tests/dialects/test_polynomial.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/dialects/test_polynomial.py b/tests/dialects/test_polynomial.py index f482dd2be9..2f95e79cb9 100644 --- a/tests/dialects/test_polynomial.py +++ b/tests/dialects/test_polynomial.py @@ -91,11 +91,11 @@ def test_eval_init_stores_attributes_verbatim(): def test_eval_get_basic_construction(): x = create_ssa_value(f32) - op = EvalOp.get(x, (0.5, 1.2, 0.3), f32, EvalScheme.CLENSHAW) + op = EvalOp.get(x, (0.5, 1.25, 0.375), f32, EvalScheme.CLENSHAW) assert op.value.type == f32 assert op.result.type == f32 assert op.degree == 2 - assert op.polynomial.coeff_values == (0.5, 1.2, 0.3) + assert op.polynomial.coeff_values == (0.5, 1.25, 0.375) assert op.polynomial.type == PolynomialType(RingAttr(f32)) assert op.scheme.data == EvalScheme.CLENSHAW.value assert op.eval_scheme == EvalScheme.CLENSHAW From 0c2d643183b79f69aac9aedab1e0d03022fb8a1e Mon Sep 17 00:00:00 2001 From: szerdick Date: Wed, 13 May 2026 13:31:37 +0200 Subject: [PATCH 05/28] add generic format and heir-opt first tests --- tests/filecheck/with-heir/dialects/polynomial/ops.mlir | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/filecheck/with-heir/dialects/polynomial/ops.mlir b/tests/filecheck/with-heir/dialects/polynomial/ops.mlir index 0f9b8bb09b..8db67d1306 100644 --- a/tests/filecheck/with-heir/dialects/polynomial/ops.mlir +++ b/tests/filecheck/with-heir/dialects/polynomial/ops.mlir @@ -1 +1,5 @@ // RUN: xdsl-opt %S/../../../dialects/polynomial/ops.mlir | heir-opt --allow-unregistered-dialect | xdsl-opt | filecheck %S/../../../dialects/polynomial/ops.mlir +// RUN: xdsl-opt %S/../../../dialects/polynomial/ops.mlir --print-op-generic | heir-opt --allow-unregistered-dialect | xdsl-opt | filecheck %S/../../../dialects/polynomial/ops.mlir + +// RUN: heir-opt --allow-unregistered-dialect %S/../../../dialects/polynomial/ops.mlir | xdsl-opt | filecheck %S/../../../dialects/polynomial/ops.mlir +// RUN: heir-opt --allow-unregistered-dialect --mlir-print-op-generic %S/../../../dialects/polynomial/ops.mlir | xdsl-opt | filecheck %S/../../../dialects/polynomial/ops.mlir From a7c6d8d7f21e3b8f4d4a0f27668e7a2445abecbc Mon Sep 17 00:00:00 2001 From: szerdick Date: Wed, 13 May 2026 13:33:38 +0200 Subject: [PATCH 06/28] resolve pyright issues --- xdsl/dialects/polynomial.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/xdsl/dialects/polynomial.py b/xdsl/dialects/polynomial.py index 8809281b95..e7f5599539 100644 --- a/xdsl/dialects/polynomial.py +++ b/xdsl/dialects/polynomial.py @@ -16,6 +16,7 @@ from typing import ClassVar from xdsl.dialects.builtin import ( + AnyFloat, AnyFloatConstr, ArrayAttr, ContainerOf, @@ -277,7 +278,7 @@ def get( cls, value: Operation | SSAValue, coefficients: tuple[float, ...], - element_type: Attribute, + element_type: AnyFloat, scheme: EvalScheme | str, domain_lower: float | None = None, domain_upper: float | None = None, From 73054d324699739449df21e89119780ca8787ed8 Mon Sep 17 00:00:00 2001 From: szerdick Date: Wed, 13 May 2026 13:40:11 +0200 Subject: [PATCH 07/28] pyrithgt issue resolve --- tests/dialects/test_polynomial.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/dialects/test_polynomial.py b/tests/dialects/test_polynomial.py index 2f95e79cb9..bd86a826cb 100644 --- a/tests/dialects/test_polynomial.py +++ b/tests/dialects/test_polynomial.py @@ -1,6 +1,7 @@ import pytest from xdsl.dialects.builtin import ( + AnyFloat, ArrayAttr, FloatAttr, StringAttr, @@ -111,7 +112,7 @@ def test_eval_get_basic_construction(): (TensorType(f64, [8]), f64), ], ) -def test_eval_get_type_polymorphism(value_type: Attribute, element_type: Attribute): +def test_eval_get_type_polymorphism(value_type: Attribute, element_type: AnyFloat): x = create_ssa_value(value_type) op = EvalOp.get(x, (1.0, 2.0), element_type, EvalScheme.CLENSHAW) assert op.value.type == value_type From fd27cf2df510159d0faae6c4a406ac073ea1dd06 Mon Sep 17 00:00:00 2001 From: szerdick Date: Mon, 27 Apr 2026 17:12:37 +0100 Subject: [PATCH 08/28] add lowering pass from polynomial dialect to arithmetic operations --- .../transforms/expand_polynomial_eval.mlir | 111 +++++++++++++ xdsl/dialects/polynomial.py | 10 ++ xdsl/transforms/__init__.py | 6 + xdsl/transforms/expand_polynomial_eval.py | 148 ++++++++++++++++++ 4 files changed, 275 insertions(+) create mode 100644 tests/filecheck/transforms/expand_polynomial_eval.mlir create mode 100644 xdsl/transforms/expand_polynomial_eval.py diff --git a/tests/filecheck/transforms/expand_polynomial_eval.mlir b/tests/filecheck/transforms/expand_polynomial_eval.mlir new file mode 100644 index 0000000000..38994fce98 --- /dev/null +++ b/tests/filecheck/transforms/expand_polynomial_eval.mlir @@ -0,0 +1,111 @@ +// RUN: xdsl-opt -p expand-polynomial-eval %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>, %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>, %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>, %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>, %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: } diff --git a/xdsl/dialects/polynomial.py b/xdsl/dialects/polynomial.py index e7f5599539..e3b2e94530 100644 --- a/xdsl/dialects/polynomial.py +++ b/xdsl/dialects/polynomial.py @@ -220,6 +220,16 @@ class EvalOp(IRDLOperation): 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. When both are present, + consumers (lowerings, cost models, interpreters) must affinely remap the + input from `[domain_lower, domain_upper]` to the canonical Chebyshev + domain `[-1, 1]` before evaluating the series. When both are absent, the + convention is that the coefficients were computed on `[-1, 1]` and the + input is already in that range -- no remap is performed. Producers that + computed coefficients on any non-canonical interval MUST set both bounds; + omitting them silently yields wrong numerical results. + Syntax: polynomial.eval $polynomial `,` $value attr-dict `:` type($value) Example: %result = polynomial.eval diff --git a/xdsl/transforms/__init__.py b/xdsl/transforms/__init__.py index 78cea41103..65b8865011 100644 --- a/xdsl/transforms/__init__.py +++ b/xdsl/transforms/__init__.py @@ -91,6 +91,11 @@ def get_expand_math_to_polynomials(): return expand_math_to_polynomials.ExpandMathToPolynomialsPass + def get_expand_polynomial_eval(): + from xdsl.transforms import expand_polynomial_eval + + return expand_polynomial_eval.ExpandPolynomialEvalPass + def get_convert_linalg_to_loops(): from xdsl.transforms import convert_linalg_to_loops @@ -688,6 +693,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, + "expand-polynomial-eval": get_expand_polynomial_eval, "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, diff --git a/xdsl/transforms/expand_polynomial_eval.py b/xdsl/transforms/expand_polynomial_eval.py new file mode 100644 index 0000000000..eb7341f3f0 --- /dev/null +++ b/xdsl/transforms/expand_polynomial_eval.py @@ -0,0 +1,148 @@ +""" +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, +) + + +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: list[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 ExpandPolynomialEval(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 + if scheme is polynomial.EvalScheme.CLENSHAW: + expanded = expand_clenshaw(op, rewriter, coeffs, lower, upper) + else: + # 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 NotImplementedError( + f"polynomial.eval scheme {scheme.value!r} has no lowering" + ) + + rewriter.replace_op(op, (), (expanded.results[0],)) + + +@dataclass(frozen=True) +class ExpandPolynomialEvalPass(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 = "expand-polynomial-eval" + + def apply(self, ctx: Context, op: ModuleOp) -> None: + PatternRewriteWalker( + ExpandPolynomialEval(), + apply_recursively=False, + ).rewrite_module(op) From e6f7a534b2354d99a4a0d3a3215e63fa24707db2 Mon Sep 17 00:00:00 2001 From: szerdick Date: Wed, 6 May 2026 15:15:12 +0100 Subject: [PATCH 09/28] add to pipeline --- xdsl/transforms/test_lower_linalg_to_snitch.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/xdsl/transforms/test_lower_linalg_to_snitch.py b/xdsl/transforms/test_lower_linalg_to_snitch.py index 443d13cc7c..3208c5da06 100644 --- a/xdsl/transforms/test_lower_linalg_to_snitch.py +++ b/xdsl/transforms/test_lower_linalg_to_snitch.py @@ -19,6 +19,7 @@ convert_memref_stream_to_snitch_stream, convert_riscv_scf_for_to_frep, expand_math_to_polynomials, + expand_polynomial_eval, lower_affine, lower_snitch, memref_stream_fold_fill, @@ -58,6 +59,7 @@ lower_affine.LowerAffinePass(), convert_scf_to_riscv_scf.ConvertScfToRiscvPass(), expand_math_to_polynomials.ExpandMathToPolynomialsPass(), + expand_polynomial_eval.ExpandPolynomialEvalPass(), convert_arith_to_riscv_snitch.ConvertArithToRiscvSnitchPass(), convert_arith_to_riscv.ConvertArithToRiscvPass(), convert_func_to_riscv_func.ConvertFuncToRiscvFuncPass(), From 2add8e6c2bf0111a104267e5314e0b93c5b84572 Mon Sep 17 00:00:00 2001 From: szerdick Date: Thu, 14 May 2026 13:55:24 +0200 Subject: [PATCH 10/28] weird comments --- xdsl/dialects/polynomial.py | 11 ++--------- 1 file changed, 2 insertions(+), 9 deletions(-) diff --git a/xdsl/dialects/polynomial.py b/xdsl/dialects/polynomial.py index e3b2e94530..2de402967f 100644 --- a/xdsl/dialects/polynomial.py +++ b/xdsl/dialects/polynomial.py @@ -217,18 +217,11 @@ 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. When both are present, - consumers (lowerings, cost models, interpreters) must affinely remap the - input from `[domain_lower, domain_upper]` to the canonical Chebyshev - domain `[-1, 1]` before evaluating the series. When both are absent, the - convention is that the coefficients were computed on `[-1, 1]` and the - input is already in that range -- no remap is performed. Producers that - computed coefficients on any non-canonical interval MUST set both bounds; - omitting them silently yields wrong numerical results. + the Chebyshev coefficients were computed for. Default is [-1, 1]. Syntax: polynomial.eval $polynomial `,` $value attr-dict `:` type($value) Example: From ff69cee3428e9a8f52656017924c358dbdb06d1d Mon Sep 17 00:00:00 2001 From: szerdick Date: Thu, 14 May 2026 14:06:20 +0200 Subject: [PATCH 11/28] pyritht fix --- xdsl/transforms/expand_polynomial_eval.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xdsl/transforms/expand_polynomial_eval.py b/xdsl/transforms/expand_polynomial_eval.py index eb7341f3f0..1bdcd03bfb 100644 --- a/xdsl/transforms/expand_polynomial_eval.py +++ b/xdsl/transforms/expand_polynomial_eval.py @@ -51,7 +51,7 @@ def _float_constant( def expand_clenshaw( op: polynomial.EvalOp, rewriter: PatternRewriter, - coeffs: list[float], + coeffs: tuple[float], lower: float | None, upper: float | None, ) -> Operation: From 80923bab6c212b66c00e241b426a571647483a17 Mon Sep 17 00:00:00 2001 From: szerdick Date: Thu, 14 May 2026 14:09:51 +0200 Subject: [PATCH 12/28] pyritht fix2 --- xdsl/transforms/expand_polynomial_eval.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xdsl/transforms/expand_polynomial_eval.py b/xdsl/transforms/expand_polynomial_eval.py index 1bdcd03bfb..b49565844b 100644 --- a/xdsl/transforms/expand_polynomial_eval.py +++ b/xdsl/transforms/expand_polynomial_eval.py @@ -51,7 +51,7 @@ def _float_constant( def expand_clenshaw( op: polynomial.EvalOp, rewriter: PatternRewriter, - coeffs: tuple[float], + coeffs: tuple[float, ...], lower: float | None, upper: float | None, ) -> Operation: From 06c44d7483c32c4d956124d877ad11550a4f2262 Mon Sep 17 00:00:00 2001 From: szerdick Date: Thu, 30 Apr 2026 14:50:44 +0100 Subject: [PATCH 13/28] draft on lowering exp to polynomial dialect --- testlowering.mlir | 6 + testlowering.sh | 7 ++ .../transforms/lower_exp_to_polynomial.mlir | 35 ++++++ xdsl/transforms/__init__.py | 6 + xdsl/transforms/lower_exp_to_polynomial.py | 115 ++++++++++++++++++ .../transforms/test_lower_linalg_to_snitch.py | 2 + 6 files changed, 171 insertions(+) create mode 100644 testlowering.mlir create mode 100644 testlowering.sh create mode 100644 tests/filecheck/transforms/lower_exp_to_polynomial.mlir create mode 100644 xdsl/transforms/lower_exp_to_polynomial.py diff --git a/testlowering.mlir b/testlowering.mlir new file mode 100644 index 0000000000..a7ef28eef0 --- /dev/null +++ b/testlowering.mlir @@ -0,0 +1,6 @@ +builtin.module { + func.func @softmax(%input: tensor<2x3xf32>, %output: tensor<2x3xf32>) -> tensor<2x3xf32> { + %r = linalg.softmax dimension(1) ins(%input : tensor<2x3xf32>) outs(%output : tensor<2x3xf32>) {acc_bound = 1.000000e-03 : f32} -> tensor<2x3xf32> + func.return %r : tensor<2x3xf32> + } +} diff --git a/testlowering.sh b/testlowering.sh new file mode 100644 index 0000000000..cda0cb95a6 --- /dev/null +++ b/testlowering.sh @@ -0,0 +1,7 @@ +uv run xdsl-opt -p decompose-softmax testlowering.mlir + + +uv run xdsl-opt -p decompose-softmax,lower-exp-to-polynomial testlowering.mlir + + +uv run xdsl-opt -p decompose-softmax,lower-exp-to-polynomial,expand-polynomial-eval testlowering.mlir diff --git a/tests/filecheck/transforms/lower_exp_to_polynomial.mlir b/tests/filecheck/transforms/lower_exp_to_polynomial.mlir new file mode 100644 index 0000000000..39f1629a74 --- /dev/null +++ b/tests/filecheck/transforms/lower_exp_to_polynomial.mlir @@ -0,0 +1,35 @@ +// RUN: xdsl-opt -p lower-exp-to-polynomial %s | filecheck %s + +builtin.module { + func.func @exp_f32(%x: f32) -> f32 { + %r = math.exp %x : f32 + func.return %r : f32 + } + + func.func @exp_f64(%x: f64) -> f64 { + %r = math.exp %x : f64 + func.return %r : f64 + } + + // acc_bound is preserved on math.exp by step 1; this pass currently ignores + // it (the polynomial choice is a fixed dummy) but should still rewrite the op. + func.func @exp_f32_with_acc_bound(%x: f32) -> f32 { + %r = math.exp %x {acc_bound = 1.000000e-03 : f32} : f32 + func.return %r : f32 + } +} + +// CHECK: func.func @exp_f32(%[[X32:.*]]: f32) -> f32 { +// CHECK-NEXT: %[[R32:.*]] = polynomial.eval #polynomial.typed_chebyshev_polynomial<[{{.*}}]> : !polynomial.polynomial>, %[[X32]] {scheme = "clenshaw", domain_lower = -103.27892990343184 : f64, domain_upper = 0.000000e+00 : f64} : 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>, %[[X64]] {scheme = "clenshaw", domain_lower = -744.44007192138122 : f64, domain_upper = 0.000000e+00 : f64} : f64 +// CHECK-NEXT: func.return %[[R64]] : f64 +// CHECK-NEXT: } + +// CHECK: func.func @exp_f32_with_acc_bound(%[[XB:.*]]: f32) -> f32 { +// CHECK-NEXT: %[[RB:.*]] = polynomial.eval #polynomial.typed_chebyshev_polynomial<[{{.*}}]> : !polynomial.polynomial>, %[[XB]] {scheme = "clenshaw", domain_lower = -103.27892990343184 : f64, domain_upper = 0.000000e+00 : f64} : f32 +// CHECK-NEXT: func.return %[[RB]] : f32 +// CHECK-NEXT: } diff --git a/xdsl/transforms/__init__.py b/xdsl/transforms/__init__.py index 3367e87fd3..d6a3d55cf0 100644 --- a/xdsl/transforms/__init__.py +++ b/xdsl/transforms/__init__.py @@ -96,6 +96,11 @@ def get_expand_polynomial_eval(): return expand_polynomial_eval.ExpandPolynomialEvalPass + def get_lower_exp_to_polynomial(): + from xdsl.transforms import lower_exp_to_polynomial + + return lower_exp_to_polynomial.LowerExpToPolynomialPass + def get_convert_linalg_to_loops(): from xdsl.transforms import convert_linalg_to_loops @@ -699,6 +704,7 @@ def get_verify_register_allocation(): "convert-func-to-x86-func": get_convert_func_to_x86_func, "expand-math-to-polynomials": get_expand_math_to_polynomials, "expand-polynomial-eval": get_expand_polynomial_eval, + "lower-exp-to-polynomial": get_lower_exp_to_polynomial, "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, diff --git a/xdsl/transforms/lower_exp_to_polynomial.py b/xdsl/transforms/lower_exp_to_polynomial.py new file mode 100644 index 0000000000..09886132b6 --- /dev/null +++ b/xdsl/transforms/lower_exp_to_polynomial.py @@ -0,0 +1,115 @@ +import math as pymath +from collections.abc import Callable +from dataclasses import dataclass + +from xdsl.context import Context +from xdsl.dialects import math, polynomial +from xdsl.dialects.builtin import ( + BFloat16Type, + Float16Type, + Float32Type, + Float64Type, + FloatAttr, + ModuleOp, + TensorType, + VectorType, +) +from xdsl.ir import Attribute +from xdsl.passes import ModulePass +from xdsl.pattern_rewriter import ( + PatternRewriter, + PatternRewriteWalker, + RewritePattern, + op_type_rewrite_pattern, +) + +_UNDERFLOW_LOWER_BOUND: dict[type, float] = { + Float16Type: pymath.log(2.0**-24), # ≈ -16.64 + BFloat16Type: pymath.log(2.0**-133), # ≈ -92.18 + Float32Type: pymath.log(2.0**-149), # ≈ -103.28 + Float64Type: pymath.log(2.0**-1074), # ≈ -744.44 +} + +_DEFAULT_DEGREE = 10 + + +def _chebyshev_coefficients( + f: Callable[[float], float], + degree: int, + lower: float, + upper: float, +) -> list[float]: + """Chebyshev coefficients c_0..c_n for `f` on [lower, upper] via DCT-I.""" + n = degree + nodes = [pymath.cos(pymath.pi * j / n) for j in range(n + 1)] + mid = (upper + lower) / 2.0 + half = (upper - lower) / 2.0 + values = [f(half * t + mid) for t in nodes] + coeffs: list[float] = [] + for k in range(n + 1): + s = 0.0 + for j in range(n + 1): + w = 0.5 if (j == 0 or j == n) else 1.0 + s += w * values[j] * pymath.cos(pymath.pi * k * j / n) + coeffs.append(2.0 * s / n) + return coeffs + + +def _element_float_type(tp: Attribute) -> Attribute: + if isinstance(tp, (VectorType, TensorType)): + return tp.get_element_type() + return tp + + +def _choose_polynomial( + acc_bound: float | None, + lower: float, + upper: float, +) -> list[float]: + """Dummy: ignore acc_bound, always return Chebyshev coefficients of a + fixed degree on [lower, upper]. A real chooser would derive the degree + (and possibly family) from acc_bound and the domain width.""" + return _chebyshev_coefficients(pymath.exp, _DEFAULT_DEGREE, lower, upper) + + +class LowerExpToPolynomial(RewritePattern): + @op_type_rewrite_pattern + def match_and_rewrite(self, op: math.ExpOp, rewriter: PatternRewriter) -> None: + # derive lower and upper bounds from precision format + # this is specific to exp + elem_ty = _element_float_type(op.operand.type) + lower = _UNDERFLOW_LOWER_BOUND.get(type(elem_ty)) + if lower is None: + return + upper = 0.0 + + # accuracy bound carried over from softmax (step 1); drives the + # polynomial degree the chooser picks. Absent => no constraint. + acc_bound_attr = op.attributes.get("acc_bound") + acc_bound = ( + acc_bound_attr.value.data if isinstance(acc_bound_attr, FloatAttr) else None + ) + + # choose which polynomial family and degree to use, and compute coefficients + coeffs = _choose_polynomial(acc_bound, lower, upper) + + # insert polynomial into the IR replacing exp + rewriter.replace_matched_op( + polynomial.EvalOp( + value=op.operand, + polynomial=tuple(coeffs), + scheme=polynomial.EvalScheme.CLENSHAW, + domain_lower=lower, + domain_upper=upper, + ) + ) + + +@dataclass(frozen=True) +class LowerExpToPolynomialPass(ModulePass): + """Lower `math.exp` to `polynomial.eval`.""" + + name = "lower-exp-to-polynomial" + + def apply(self, ctx: Context, op: ModuleOp) -> None: + PatternRewriteWalker(LowerExpToPolynomial()).rewrite_module(op) diff --git a/xdsl/transforms/test_lower_linalg_to_snitch.py b/xdsl/transforms/test_lower_linalg_to_snitch.py index 3208c5da06..234770df34 100644 --- a/xdsl/transforms/test_lower_linalg_to_snitch.py +++ b/xdsl/transforms/test_lower_linalg_to_snitch.py @@ -21,6 +21,7 @@ expand_math_to_polynomials, expand_polynomial_eval, lower_affine, + lower_exp_to_polynomial, lower_snitch, memref_stream_fold_fill, memref_stream_generalize_fill, @@ -59,6 +60,7 @@ lower_affine.LowerAffinePass(), convert_scf_to_riscv_scf.ConvertScfToRiscvPass(), expand_math_to_polynomials.ExpandMathToPolynomialsPass(), + lower_exp_to_polynomial.LowerExpToPolynomialPass(), expand_polynomial_eval.ExpandPolynomialEvalPass(), convert_arith_to_riscv_snitch.ConvertArithToRiscvSnitchPass(), convert_arith_to_riscv.ConvertArithToRiscvPass(), From 683c2d59ee4038ed2ae8bbd3c41085cb6309bfbf Mon Sep 17 00:00:00 2001 From: szerdick Date: Fri, 1 May 2026 10:13:42 +0100 Subject: [PATCH 14/28] lower default Chebyshev degree to 5 --- xdsl/transforms/lower_exp_to_polynomial.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xdsl/transforms/lower_exp_to_polynomial.py b/xdsl/transforms/lower_exp_to_polynomial.py index 09886132b6..62edea4da1 100644 --- a/xdsl/transforms/lower_exp_to_polynomial.py +++ b/xdsl/transforms/lower_exp_to_polynomial.py @@ -30,7 +30,7 @@ Float64Type: pymath.log(2.0**-1074), # ≈ -744.44 } -_DEFAULT_DEGREE = 10 +_DEFAULT_DEGREE = 5 def _chebyshev_coefficients( From 6f6d5d7b9e1736ba45e7b2ef80007f5522aa09c6 Mon Sep 17 00:00:00 2001 From: szerdick Date: Fri, 1 May 2026 13:08:34 +0100 Subject: [PATCH 15/28] change polynomial degree --- xdsl/transforms/lower_exp_to_polynomial.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xdsl/transforms/lower_exp_to_polynomial.py b/xdsl/transforms/lower_exp_to_polynomial.py index 62edea4da1..258c8073bf 100644 --- a/xdsl/transforms/lower_exp_to_polynomial.py +++ b/xdsl/transforms/lower_exp_to_polynomial.py @@ -30,7 +30,7 @@ Float64Type: pymath.log(2.0**-1074), # ≈ -744.44 } -_DEFAULT_DEGREE = 5 +_DEFAULT_DEGREE = 12 def _chebyshev_coefficients( From 709714633c2fbb4413d070fc112bf4dc2c1375b0 Mon Sep 17 00:00:00 2001 From: szerdick Date: Fri, 1 May 2026 14:24:55 +0100 Subject: [PATCH 16/28] change interval bounds to be smaller --- .../transforms/lower_exp_to_polynomial.mlir | 10 +++---- xdsl/transforms/lower_exp_to_polynomial.py | 30 +++++++++++++------ 2 files changed, 26 insertions(+), 14 deletions(-) diff --git a/tests/filecheck/transforms/lower_exp_to_polynomial.mlir b/tests/filecheck/transforms/lower_exp_to_polynomial.mlir index 39f1629a74..fd897ee9f1 100644 --- a/tests/filecheck/transforms/lower_exp_to_polynomial.mlir +++ b/tests/filecheck/transforms/lower_exp_to_polynomial.mlir @@ -11,8 +11,8 @@ builtin.module { func.return %r : f64 } - // acc_bound is preserved on math.exp by step 1; this pass currently ignores - // it (the polynomial choice is a fixed dummy) but should still rewrite the op. + // acc_bound is preserved on math.exp by step 1 and drives the polynomial + // degree picked by _choose_polynomial. func.func @exp_f32_with_acc_bound(%x: f32) -> f32 { %r = math.exp %x {acc_bound = 1.000000e-03 : f32} : f32 func.return %r : f32 @@ -20,16 +20,16 @@ builtin.module { } // CHECK: func.func @exp_f32(%[[X32:.*]]: f32) -> f32 { -// CHECK-NEXT: %[[R32:.*]] = polynomial.eval #polynomial.typed_chebyshev_polynomial<[{{.*}}]> : !polynomial.polynomial>, %[[X32]] {scheme = "clenshaw", domain_lower = -103.27892990343184 : f64, domain_upper = 0.000000e+00 : f64} : f32 +// CHECK-NEXT: %[[R32:.*]] = polynomial.eval #polynomial.typed_chebyshev_polynomial<[{{.*}}]> : !polynomial.polynomial>, %[[X32]] {scheme = "clenshaw", domain_lower = -1.000000e+00 : f64, domain_upper = 0.000000e+00 : f64} : 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>, %[[X64]] {scheme = "clenshaw", domain_lower = -744.44007192138122 : f64, domain_upper = 0.000000e+00 : f64} : f64 +// CHECK-NEXT: %[[R64:.*]] = polynomial.eval #polynomial.typed_chebyshev_polynomial<[{{.*}}]> : !polynomial.polynomial>, %[[X64]] {scheme = "clenshaw", domain_lower = -1.000000e+00 : f64, domain_upper = 0.000000e+00 : f64} : f64 // CHECK-NEXT: func.return %[[R64]] : f64 // CHECK-NEXT: } // CHECK: func.func @exp_f32_with_acc_bound(%[[XB:.*]]: f32) -> f32 { -// CHECK-NEXT: %[[RB:.*]] = polynomial.eval #polynomial.typed_chebyshev_polynomial<[{{.*}}]> : !polynomial.polynomial>, %[[XB]] {scheme = "clenshaw", domain_lower = -103.27892990343184 : f64, domain_upper = 0.000000e+00 : f64} : f32 +// CHECK-NEXT: %[[RB:.*]] = polynomial.eval #polynomial.typed_chebyshev_polynomial<[{{.*}}]> : !polynomial.polynomial>, %[[XB]] {scheme = "clenshaw", domain_lower = -1.000000e+00 : f64, domain_upper = 0.000000e+00 : f64} : f32 // CHECK-NEXT: func.return %[[RB]] : f32 // CHECK-NEXT: } diff --git a/xdsl/transforms/lower_exp_to_polynomial.py b/xdsl/transforms/lower_exp_to_polynomial.py index 258c8073bf..be2dbc79a5 100644 --- a/xdsl/transforms/lower_exp_to_polynomial.py +++ b/xdsl/transforms/lower_exp_to_polynomial.py @@ -24,13 +24,17 @@ ) _UNDERFLOW_LOWER_BOUND: dict[type, float] = { - Float16Type: pymath.log(2.0**-24), # ≈ -16.64 - BFloat16Type: pymath.log(2.0**-133), # ≈ -92.18 - Float32Type: pymath.log(2.0**-149), # ≈ -103.28 - Float64Type: pymath.log(2.0**-1074), # ≈ -744.44 + # Float16Type: pymath.log(2.0**-24), # ≈ -16.64 + # BFloat16Type: pymath.log(2.0**-133), # ≈ -92.18 + # Float32Type: pymath.log(2.0**-149), # ≈ -103.28 + # Float64Type: pymath.log(2.0**-1074), # ≈ -744.44 + Float16Type: -1, + BFloat16Type: -1, + Float32Type: -1, + Float64Type: -1, } -_DEFAULT_DEGREE = 12 +_DEFAULT_DEGREE = 10 def _chebyshev_coefficients( @@ -66,10 +70,18 @@ def _choose_polynomial( lower: float, upper: float, ) -> list[float]: - """Dummy: ignore acc_bound, always return Chebyshev coefficients of a - fixed degree on [lower, upper]. A real chooser would derive the degree - (and possibly family) from acc_bound and the domain width.""" - return _chebyshev_coefficients(pymath.exp, _DEFAULT_DEGREE, lower, upper) + """Smallest degree such that the Lobatto Chebyshev bound for exp on [a,b] is <= acc_bound.""" + if acc_bound is None: + return _chebyshev_coefficients(pymath.exp, _DEFAULT_DEGREE, lower, upper) + + width = upper - lower + degree = 0 + bound = pymath.exp(upper) * width + while bound > acc_bound: + degree += 1 + bound *= width / (4 * (degree + 1)) + + return _chebyshev_coefficients(pymath.exp, degree, lower, upper) class LowerExpToPolynomial(RewritePattern): From 38b1b5b8e67a66e34e4a3850a82a2f253dfaa23c Mon Sep 17 00:00:00 2001 From: szerdick Date: Wed, 6 May 2026 15:26:15 +0100 Subject: [PATCH 17/28] remove scratch files that tested softmax lowering --- testlowering.mlir | 6 ------ testlowering.sh | 7 ------- 2 files changed, 13 deletions(-) delete mode 100644 testlowering.mlir delete mode 100644 testlowering.sh diff --git a/testlowering.mlir b/testlowering.mlir deleted file mode 100644 index a7ef28eef0..0000000000 --- a/testlowering.mlir +++ /dev/null @@ -1,6 +0,0 @@ -builtin.module { - func.func @softmax(%input: tensor<2x3xf32>, %output: tensor<2x3xf32>) -> tensor<2x3xf32> { - %r = linalg.softmax dimension(1) ins(%input : tensor<2x3xf32>) outs(%output : tensor<2x3xf32>) {acc_bound = 1.000000e-03 : f32} -> tensor<2x3xf32> - func.return %r : tensor<2x3xf32> - } -} diff --git a/testlowering.sh b/testlowering.sh deleted file mode 100644 index cda0cb95a6..0000000000 --- a/testlowering.sh +++ /dev/null @@ -1,7 +0,0 @@ -uv run xdsl-opt -p decompose-softmax testlowering.mlir - - -uv run xdsl-opt -p decompose-softmax,lower-exp-to-polynomial testlowering.mlir - - -uv run xdsl-opt -p decompose-softmax,lower-exp-to-polynomial,expand-polynomial-eval testlowering.mlir From 834e2e5fd2525fd7c7a60171d7610571d6d76135 Mon Sep 17 00:00:00 2001 From: szerdick Date: Wed, 6 May 2026 16:11:05 +0100 Subject: [PATCH 18/28] edit attributes for exp->polynomial pass and add more extensive testing --- .../transforms/lower_exp_to_polynomial.mlir | 34 +++- .../test_lower_exp_to_polynomial.py | 171 ++++++++++++++++++ xdsl/transforms/lower_exp_to_polynomial.py | 52 ++++-- 3 files changed, 237 insertions(+), 20 deletions(-) create mode 100644 tests/transforms/test_lower_exp_to_polynomial.py diff --git a/tests/filecheck/transforms/lower_exp_to_polynomial.mlir b/tests/filecheck/transforms/lower_exp_to_polynomial.mlir index fd897ee9f1..930c98d11d 100644 --- a/tests/filecheck/transforms/lower_exp_to_polynomial.mlir +++ b/tests/filecheck/transforms/lower_exp_to_polynomial.mlir @@ -14,22 +14,48 @@ builtin.module { // acc_bound is preserved on math.exp by step 1 and drives the polynomial // degree picked by _choose_polynomial. func.func @exp_f32_with_acc_bound(%x: f32) -> f32 { - %r = math.exp %x {acc_bound = 1.000000e-03 : f32} : f32 + %r = math.exp %x {acc_bound = 1.000000e-03 : f64} : f32 + func.return %r : f32 + } + + // user lower_bound = -2 < underflow (-1) -> clamped to -1 + // user upper_bound = 5 < overflow (~88.72) -> kept at 5 + func.func @exp_f32_clamped_lower(%x: f32) -> f32 { + %r = math.exp %x {acc_bound = 1.000000e-03 : f64, 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 {acc_bound = 1.000000e-03 : f64, lower_bound = -5.000000e-01 : f32, upper_bound = 5.000000e-01 : f32} : f32 func.return %r : f32 } } +// Without acc_bound, math.exp is left untouched. // CHECK: func.func @exp_f32(%[[X32:.*]]: f32) -> f32 { -// CHECK-NEXT: %[[R32:.*]] = polynomial.eval #polynomial.typed_chebyshev_polynomial<[{{.*}}]> : !polynomial.polynomial>, %[[X32]] {scheme = "clenshaw", domain_lower = -1.000000e+00 : f64, domain_upper = 0.000000e+00 : f64} : f32 +// CHECK-NEXT: %[[R32:.*]] = math.exp %[[X32]] : 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>, %[[X64]] {scheme = "clenshaw", domain_lower = -1.000000e+00 : f64, domain_upper = 0.000000e+00 : f64} : f64 +// CHECK-NEXT: %[[R64:.*]] = math.exp %[[X64]] : 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 overflow ~= 88.72). // CHECK: func.func @exp_f32_with_acc_bound(%[[XB:.*]]: f32) -> f32 { -// CHECK-NEXT: %[[RB:.*]] = polynomial.eval #polynomial.typed_chebyshev_polynomial<[{{.*}}]> : !polynomial.polynomial>, %[[XB]] {scheme = "clenshaw", domain_lower = -1.000000e+00 : f64, domain_upper = 0.000000e+00 : f64} : f32 +// CHECK-NEXT: %[[RB:.*]] = polynomial.eval #polynomial.typed_chebyshev_polynomial<[{{.*}}]> : !polynomial.polynomial>, %[[XB]] {scheme = "clenshaw", domain_lower = -1.000000e+00 : f64, domain_upper = {{.*}} : f64} : 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>, %[[XC]] {scheme = "clenshaw", domain_lower = -1.000000e+00 : f64, domain_upper = 5.000000e+00 : f64} : 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>, %[[XT]] {scheme = "clenshaw", domain_lower = -5.000000e-01 : f64, domain_upper = 5.000000e-01 : f64} : f32 +// CHECK-NEXT: func.return %[[RT]] : f32 +// CHECK-NEXT: } diff --git a/tests/transforms/test_lower_exp_to_polynomial.py b/tests/transforms/test_lower_exp_to_polynomial.py new file mode 100644 index 0000000000..352e3af8a5 --- /dev/null +++ b/tests/transforms/test_lower_exp_to_polynomial.py @@ -0,0 +1,171 @@ +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, + ModuleOp, +) +from xdsl.transforms.lower_exp_to_polynomial import ( + _OVERFLOW_UPPER_BOUND, + _UNDERFLOW_LOWER_BOUND, + LowerExpToPolynomialPass, +) +from xdsl.utils.test_value import create_ssa_value + +PRECISIONS = pytest.mark.parametrize( + "elem_ty", [Float16Type(), Float32Type(), Float64Type()] +) + +# acc_bound is always provided in f64 +_F64 = Float64Type() + + +def _run_pass(elem_ty, attrs: dict[str, float]) -> ModuleOp: + operand = create_ssa_value(elem_ty) + exp_op = math_dialect.ExpOp(operand) + for name, value in attrs.items(): + attr_ty = _F64 if name == "acc_bound" else elem_ty + exp_op.attributes[name] = FloatAttr(value, attr_ty) + module = ModuleOp([exp_op]) + ctx = Context() + ctx.load_op(ModuleOp) + LowerExpToPolynomialPass().apply(ctx, module) + return module + + +@PRECISIONS +def test_no_acc_no_bounds(elem_ty): + """No acc_bound, no [lower, upper] -> math.exp left alone.""" + module = _run_pass(elem_ty, {}) + ops = list(module.body.block.ops) + assert len(ops) == 1 + assert isinstance(ops[0], math_dialect.ExpOp) + + +@PRECISIONS +def test_no_acc_with_bounds(elem_ty): + """No acc_bound, [lower, upper] given -> math.exp left alone.""" + module = _run_pass(elem_ty, {"lower_bound": -0.5, "upper_bound": 0.5}) + ops = list(module.body.block.ops) + assert len(ops) == 1 + assert isinstance(ops[0], math_dialect.ExpOp) + + +@PRECISIONS +def test_acc_only(elem_ty): + """acc_bound only -> default interval [underflow, overflow].""" + module = _run_pass(elem_ty, {"acc_bound": 1e-3}) + 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 == pytest.approx( + _UNDERFLOW_LOWER_BOUND[type(elem_ty)] + ) + assert eval_op.domain_upper.value.data == pytest.approx( + _OVERFLOW_UPPER_BOUND[type(elem_ty)] + ) + + +@PRECISIONS +def test_acc_bounds_in_range(elem_ty): + """acc_bound + bounds entirely in range -> bounds used as-is.""" + module = _run_pass( + elem_ty, + {"acc_bound": 1e-3, "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.value.data == pytest.approx(-0.5) + assert eval_op.domain_upper.value.data == pytest.approx(0.5) + + +@PRECISIONS +def test_acc_lower_out_of_range(elem_ty): + """acc_bound + lower < underflow -> lower clamped to underflow.""" + module = _run_pass( + elem_ty, + {"acc_bound": 1e-3, "lower_bound": -100.0, "upper_bound": 0.5}, + ) + eval_op = next(iter(module.body.block.ops)) + assert isinstance(eval_op, polynomial.EvalOp) + assert eval_op.domain_lower.value.data == pytest.approx( + _UNDERFLOW_LOWER_BOUND[type(elem_ty)] + ) + assert eval_op.domain_upper.value.data == pytest.approx(0.5) + + +@PRECISIONS +def test_acc_upper_out_of_range(elem_ty): + """acc_bound + 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, + {"acc_bound": 1e-3, "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.value.data == pytest.approx(-0.5) + assert eval_op.domain_upper.value.data == pytest.approx( + _OVERFLOW_UPPER_BOUND[type(elem_ty)] + ) + + +@PRECISIONS +def test_acc_lower_only_in_range(elem_ty): + """acc_bound + lower only (in range) -> [lower, overflow].""" + module = _run_pass(elem_ty, {"acc_bound": 1e-3, "lower_bound": -0.5}) + eval_op = next(iter(module.body.block.ops)) + assert isinstance(eval_op, polynomial.EvalOp) + assert eval_op.domain_lower.value.data == pytest.approx(-0.5) + assert eval_op.domain_upper.value.data == pytest.approx( + _OVERFLOW_UPPER_BOUND[type(elem_ty)] + ) + + +@PRECISIONS +def test_acc_lower_only_out_of_range(elem_ty): + """acc_bound + lower only (out of range) -> [underflow, overflow].""" + module = _run_pass(elem_ty, {"acc_bound": 1e-3, "lower_bound": -100.0}) + eval_op = next(iter(module.body.block.ops)) + assert isinstance(eval_op, polynomial.EvalOp) + assert eval_op.domain_lower.value.data == pytest.approx( + _UNDERFLOW_LOWER_BOUND[type(elem_ty)] + ) + assert eval_op.domain_upper.value.data == pytest.approx( + _OVERFLOW_UPPER_BOUND[type(elem_ty)] + ) + + +@PRECISIONS +def test_acc_upper_only_in_range(elem_ty): + """acc_bound + upper only (in range) -> [underflow, upper].""" + module = _run_pass(elem_ty, {"acc_bound": 1e-3, "upper_bound": 0.5}) + eval_op = next(iter(module.body.block.ops)) + assert isinstance(eval_op, polynomial.EvalOp) + assert eval_op.domain_lower.value.data == pytest.approx( + _UNDERFLOW_LOWER_BOUND[type(elem_ty)] + ) + assert eval_op.domain_upper.value.data == pytest.approx(0.5) + + +@PRECISIONS +def test_acc_upper_only_out_of_range(elem_ty): + """acc_bound + upper only (out of range) -> [underflow, overflow].""" + module = _run_pass(elem_ty, {"acc_bound": 1e-3, "upper_bound": 1000.0}) + eval_op = next(iter(module.body.block.ops)) + assert isinstance(eval_op, polynomial.EvalOp) + assert eval_op.domain_lower.value.data == pytest.approx( + _UNDERFLOW_LOWER_BOUND[type(elem_ty)] + ) + assert eval_op.domain_upper.value.data == pytest.approx( + _OVERFLOW_UPPER_BOUND[type(elem_ty)] + ) diff --git a/xdsl/transforms/lower_exp_to_polynomial.py b/xdsl/transforms/lower_exp_to_polynomial.py index be2dbc79a5..c81a7686f1 100644 --- a/xdsl/transforms/lower_exp_to_polynomial.py +++ b/xdsl/transforms/lower_exp_to_polynomial.py @@ -34,7 +34,16 @@ Float64Type: -1, } -_DEFAULT_DEGREE = 10 +# Largest x such that exp(x) is still representable in the given precision. +_OVERFLOW_UPPER_BOUND: dict[type, float] = { + Float16Type: 16.0 * pymath.log(2.0), # ≈ 11.09 + BFloat16Type: 128.0 * pymath.log(2.0), # ≈ 88.72 + Float32Type: 128.0 * pymath.log(2.0), # ≈ 88.72 + Float64Type: 1024.0 * pymath.log(2.0), # ≈ 709.78 +} + +# Hard cap on the polynomial degree the chooser will grow to. +_MAX_DEGREE = 30 def _chebyshev_coefficients( @@ -66,18 +75,16 @@ def _element_float_type(tp: Attribute) -> Attribute: def _choose_polynomial( - acc_bound: float | None, + acc_bound: float, lower: float, upper: float, ) -> list[float]: - """Smallest degree such that the Lobatto Chebyshev bound for exp on [a,b] is <= acc_bound.""" - if acc_bound is None: - return _chebyshev_coefficients(pymath.exp, _DEFAULT_DEGREE, lower, upper) - + """Smallest degree such that the Lobatto Chebyshev bound for exp on [a,b] is <= acc_bound. + Capped at _MAX_DEGREE.""" width = upper - lower degree = 0 bound = pymath.exp(upper) * width - while bound > acc_bound: + while bound > acc_bound and degree < _MAX_DEGREE: degree += 1 bound *= width / (4 * (degree + 1)) @@ -87,19 +94,32 @@ def _choose_polynomial( class LowerExpToPolynomial(RewritePattern): @op_type_rewrite_pattern def match_and_rewrite(self, op: math.ExpOp, rewriter: PatternRewriter) -> None: - # derive lower and upper bounds from precision format - # this is specific to exp + # representable range for the precision format; specific to exp elem_ty = _element_float_type(op.operand.type) - lower = _UNDERFLOW_LOWER_BOUND.get(type(elem_ty)) - if lower is None: + underflow = _UNDERFLOW_LOWER_BOUND.get(type(elem_ty)) + overflow = _OVERFLOW_UPPER_BOUND.get(type(elem_ty)) + if underflow is None or overflow is None: return - upper = 0.0 - # accuracy bound carried over from softmax (step 1); drives the - # polynomial degree the chooser picks. Absent => no constraint. + # if no acc_bound attr, do not optimize exp acc_bound_attr = op.attributes.get("acc_bound") - acc_bound = ( - acc_bound_attr.value.data if isinstance(acc_bound_attr, FloatAttr) else None + if not isinstance(acc_bound_attr, FloatAttr): + return + acc_bound = acc_bound_attr.value.data + + # Polynomial domain: each side independently clamped to its representable + # extreme. Missing lower_bound -> underflow; missing upper_bound -> overflow. + lower_attr = op.attributes.get("lower_bound") + upper_attr = op.attributes.get("upper_bound") + lower = ( + max(underflow, lower_attr.value.data) + if isinstance(lower_attr, FloatAttr) + else underflow + ) + upper = ( + min(overflow, upper_attr.value.data) + if isinstance(upper_attr, FloatAttr) + else overflow ) # choose which polynomial family and degree to use, and compute coefficients From d7c8248e305903e7f989810ccaf3cdd605dc2d48 Mon Sep 17 00:00:00 2001 From: szerdick Date: Thu, 7 May 2026 10:10:46 +0100 Subject: [PATCH 19/28] fix some pyright errors --- .../test_lower_exp_to_polynomial.py | 94 +++++++++---------- xdsl/transforms/lower_exp_to_polynomial.py | 11 ++- 2 files changed, 51 insertions(+), 54 deletions(-) diff --git a/tests/transforms/test_lower_exp_to_polynomial.py b/tests/transforms/test_lower_exp_to_polynomial.py index 352e3af8a5..7edba4ad58 100644 --- a/tests/transforms/test_lower_exp_to_polynomial.py +++ b/tests/transforms/test_lower_exp_to_polynomial.py @@ -11,12 +11,14 @@ ModuleOp, ) from xdsl.transforms.lower_exp_to_polynomial import ( - _OVERFLOW_UPPER_BOUND, - _UNDERFLOW_LOWER_BOUND, + 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()] ) @@ -25,7 +27,7 @@ _F64 = Float64Type() -def _run_pass(elem_ty, attrs: dict[str, float]) -> ModuleOp: +def _run_pass(elem_ty: _FloatTy, attrs: dict[str, float]) -> ModuleOp: operand = create_ssa_value(elem_ty) exp_op = math_dialect.ExpOp(operand) for name, value in attrs.items(): @@ -39,7 +41,7 @@ def _run_pass(elem_ty, attrs: dict[str, float]) -> ModuleOp: @PRECISIONS -def test_no_acc_no_bounds(elem_ty): +def test_no_acc_no_bounds(elem_ty: _FloatTy): """No acc_bound, no [lower, upper] -> math.exp left alone.""" module = _run_pass(elem_ty, {}) ops = list(module.body.block.ops) @@ -48,7 +50,7 @@ def test_no_acc_no_bounds(elem_ty): @PRECISIONS -def test_no_acc_with_bounds(elem_ty): +def test_no_acc_with_bounds(elem_ty: _FloatTy): """No acc_bound, [lower, upper] given -> math.exp left alone.""" module = _run_pass(elem_ty, {"lower_bound": -0.5, "upper_bound": 0.5}) ops = list(module.body.block.ops) @@ -57,23 +59,19 @@ def test_no_acc_with_bounds(elem_ty): @PRECISIONS -def test_acc_only(elem_ty): +def test_acc_only(elem_ty: _FloatTy): """acc_bound only -> default interval [underflow, overflow].""" module = _run_pass(elem_ty, {"acc_bound": 1e-3}) 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 == pytest.approx( - _UNDERFLOW_LOWER_BOUND[type(elem_ty)] - ) - assert eval_op.domain_upper.value.data == pytest.approx( - _OVERFLOW_UPPER_BOUND[type(elem_ty)] - ) + assert eval_op.domain_lower.value.data == UNDERFLOW_LOWER_BOUND[type(elem_ty)] + assert eval_op.domain_upper.value.data == OVERFLOW_UPPER_BOUND[type(elem_ty)] @PRECISIONS -def test_acc_bounds_in_range(elem_ty): +def test_acc_bounds_in_range(elem_ty: _FloatTy): """acc_bound + bounds entirely in range -> bounds used as-is.""" module = _run_pass( elem_ty, @@ -81,12 +79,14 @@ def test_acc_bounds_in_range(elem_ty): ) eval_op = next(iter(module.body.block.ops)) assert isinstance(eval_op, polynomial.EvalOp) - assert eval_op.domain_lower.value.data == pytest.approx(-0.5) - assert eval_op.domain_upper.value.data == pytest.approx(0.5) + 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_acc_lower_out_of_range(elem_ty): +def test_acc_lower_out_of_range(elem_ty: _FloatTy): """acc_bound + lower < underflow -> lower clamped to underflow.""" module = _run_pass( elem_ty, @@ -94,14 +94,14 @@ def test_acc_lower_out_of_range(elem_ty): ) eval_op = next(iter(module.body.block.ops)) assert isinstance(eval_op, polynomial.EvalOp) - assert eval_op.domain_lower.value.data == pytest.approx( - _UNDERFLOW_LOWER_BOUND[type(elem_ty)] - ) - assert eval_op.domain_upper.value.data == pytest.approx(0.5) + assert eval_op.domain_lower is not None + assert eval_op.domain_upper is not None + assert eval_op.domain_lower.value.data == UNDERFLOW_LOWER_BOUND[type(elem_ty)] + assert eval_op.domain_upper.value.data == 0.5 @PRECISIONS -def test_acc_upper_out_of_range(elem_ty): +def test_acc_upper_out_of_range(elem_ty: _FloatTy): """acc_bound + upper > overflow -> upper clamped to overflow. 1000 exceeds every supported precision's overflow (f16 ~= 11.09, @@ -113,59 +113,55 @@ def test_acc_upper_out_of_range(elem_ty): ) eval_op = next(iter(module.body.block.ops)) assert isinstance(eval_op, polynomial.EvalOp) - assert eval_op.domain_lower.value.data == pytest.approx(-0.5) - assert eval_op.domain_upper.value.data == pytest.approx( - _OVERFLOW_UPPER_BOUND[type(elem_ty)] - ) + 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_UPPER_BOUND[type(elem_ty)] @PRECISIONS -def test_acc_lower_only_in_range(elem_ty): +def test_acc_lower_only_in_range(elem_ty: _FloatTy): """acc_bound + lower only (in range) -> [lower, overflow].""" module = _run_pass(elem_ty, {"acc_bound": 1e-3, "lower_bound": -0.5}) eval_op = next(iter(module.body.block.ops)) assert isinstance(eval_op, polynomial.EvalOp) - assert eval_op.domain_lower.value.data == pytest.approx(-0.5) - assert eval_op.domain_upper.value.data == pytest.approx( - _OVERFLOW_UPPER_BOUND[type(elem_ty)] - ) + 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_UPPER_BOUND[type(elem_ty)] @PRECISIONS -def test_acc_lower_only_out_of_range(elem_ty): +def test_acc_lower_only_out_of_range(elem_ty: _FloatTy): """acc_bound + lower only (out of range) -> [underflow, overflow].""" module = _run_pass(elem_ty, {"acc_bound": 1e-3, "lower_bound": -100.0}) eval_op = next(iter(module.body.block.ops)) assert isinstance(eval_op, polynomial.EvalOp) - assert eval_op.domain_lower.value.data == pytest.approx( - _UNDERFLOW_LOWER_BOUND[type(elem_ty)] - ) - assert eval_op.domain_upper.value.data == pytest.approx( - _OVERFLOW_UPPER_BOUND[type(elem_ty)] - ) + assert eval_op.domain_lower is not None + assert eval_op.domain_upper is not None + assert eval_op.domain_lower.value.data == UNDERFLOW_LOWER_BOUND[type(elem_ty)] + assert eval_op.domain_upper.value.data == OVERFLOW_UPPER_BOUND[type(elem_ty)] @PRECISIONS -def test_acc_upper_only_in_range(elem_ty): +def test_acc_upper_only_in_range(elem_ty: _FloatTy): """acc_bound + upper only (in range) -> [underflow, upper].""" module = _run_pass(elem_ty, {"acc_bound": 1e-3, "upper_bound": 0.5}) eval_op = next(iter(module.body.block.ops)) assert isinstance(eval_op, polynomial.EvalOp) - assert eval_op.domain_lower.value.data == pytest.approx( - _UNDERFLOW_LOWER_BOUND[type(elem_ty)] - ) - assert eval_op.domain_upper.value.data == pytest.approx(0.5) + assert eval_op.domain_lower is not None + assert eval_op.domain_upper is not None + assert eval_op.domain_lower.value.data == UNDERFLOW_LOWER_BOUND[type(elem_ty)] + assert eval_op.domain_upper.value.data == 0.5 @PRECISIONS -def test_acc_upper_only_out_of_range(elem_ty): +def test_acc_upper_only_out_of_range(elem_ty: _FloatTy): """acc_bound + upper only (out of range) -> [underflow, overflow].""" module = _run_pass(elem_ty, {"acc_bound": 1e-3, "upper_bound": 1000.0}) eval_op = next(iter(module.body.block.ops)) assert isinstance(eval_op, polynomial.EvalOp) - assert eval_op.domain_lower.value.data == pytest.approx( - _UNDERFLOW_LOWER_BOUND[type(elem_ty)] - ) - assert eval_op.domain_upper.value.data == pytest.approx( - _OVERFLOW_UPPER_BOUND[type(elem_ty)] - ) + assert eval_op.domain_lower is not None + assert eval_op.domain_upper is not None + assert eval_op.domain_lower.value.data == UNDERFLOW_LOWER_BOUND[type(elem_ty)] + assert eval_op.domain_upper.value.data == OVERFLOW_UPPER_BOUND[type(elem_ty)] diff --git a/xdsl/transforms/lower_exp_to_polynomial.py b/xdsl/transforms/lower_exp_to_polynomial.py index c81a7686f1..b7590ae285 100644 --- a/xdsl/transforms/lower_exp_to_polynomial.py +++ b/xdsl/transforms/lower_exp_to_polynomial.py @@ -1,6 +1,7 @@ import math as pymath from collections.abc import Callable from dataclasses import dataclass +from typing import cast from xdsl.context import Context from xdsl.dialects import math, polynomial @@ -23,7 +24,7 @@ op_type_rewrite_pattern, ) -_UNDERFLOW_LOWER_BOUND: dict[type, float] = { +UNDERFLOW_LOWER_BOUND: dict[type, float] = { # Float16Type: pymath.log(2.0**-24), # ≈ -16.64 # BFloat16Type: pymath.log(2.0**-133), # ≈ -92.18 # Float32Type: pymath.log(2.0**-149), # ≈ -103.28 @@ -35,7 +36,7 @@ } # Largest x such that exp(x) is still representable in the given precision. -_OVERFLOW_UPPER_BOUND: dict[type, float] = { +OVERFLOW_UPPER_BOUND: dict[type, float] = { Float16Type: 16.0 * pymath.log(2.0), # ≈ 11.09 BFloat16Type: 128.0 * pymath.log(2.0), # ≈ 88.72 Float32Type: 128.0 * pymath.log(2.0), # ≈ 88.72 @@ -70,7 +71,7 @@ def _chebyshev_coefficients( def _element_float_type(tp: Attribute) -> Attribute: if isinstance(tp, (VectorType, TensorType)): - return tp.get_element_type() + return cast(Attribute, tp.get_element_type()) return tp @@ -96,8 +97,8 @@ class LowerExpToPolynomial(RewritePattern): def match_and_rewrite(self, op: math.ExpOp, rewriter: PatternRewriter) -> None: # representable range for the precision format; specific to exp elem_ty = _element_float_type(op.operand.type) - underflow = _UNDERFLOW_LOWER_BOUND.get(type(elem_ty)) - overflow = _OVERFLOW_UPPER_BOUND.get(type(elem_ty)) + underflow = UNDERFLOW_LOWER_BOUND.get(type(elem_ty)) + overflow = OVERFLOW_UPPER_BOUND.get(type(elem_ty)) if underflow is None or overflow is None: return From 4ebd8f609e5aac75305499cce44988b8babee446 Mon Sep 17 00:00:00 2001 From: szerdick Date: Thu, 7 May 2026 10:21:58 +0100 Subject: [PATCH 20/28] newer version of replace --- xdsl/transforms/lower_exp_to_polynomial.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/xdsl/transforms/lower_exp_to_polynomial.py b/xdsl/transforms/lower_exp_to_polynomial.py index b7590ae285..e704ffb923 100644 --- a/xdsl/transforms/lower_exp_to_polynomial.py +++ b/xdsl/transforms/lower_exp_to_polynomial.py @@ -127,14 +127,15 @@ def match_and_rewrite(self, op: math.ExpOp, rewriter: PatternRewriter) -> None: coeffs = _choose_polynomial(acc_bound, lower, upper) # insert polynomial into the IR replacing exp - rewriter.replace_matched_op( + rewriter.replace_op( + op, polynomial.EvalOp( value=op.operand, polynomial=tuple(coeffs), scheme=polynomial.EvalScheme.CLENSHAW, domain_lower=lower, domain_upper=upper, - ) + ), ) From 7116f8c65cc968a949cbfbb09b76aaebe910fadf Mon Sep 17 00:00:00 2001 From: szerdick Date: Sun, 10 May 2026 23:09:00 +0200 Subject: [PATCH 21/28] forgot to uncomment underflow default bound --- xdsl/transforms/lower_exp_to_polynomial.py | 12 ++++-------- 1 file changed, 4 insertions(+), 8 deletions(-) diff --git a/xdsl/transforms/lower_exp_to_polynomial.py b/xdsl/transforms/lower_exp_to_polynomial.py index e704ffb923..ab24a0328d 100644 --- a/xdsl/transforms/lower_exp_to_polynomial.py +++ b/xdsl/transforms/lower_exp_to_polynomial.py @@ -25,14 +25,10 @@ ) UNDERFLOW_LOWER_BOUND: dict[type, float] = { - # Float16Type: pymath.log(2.0**-24), # ≈ -16.64 - # BFloat16Type: pymath.log(2.0**-133), # ≈ -92.18 - # Float32Type: pymath.log(2.0**-149), # ≈ -103.28 - # Float64Type: pymath.log(2.0**-1074), # ≈ -744.44 - Float16Type: -1, - BFloat16Type: -1, - Float32Type: -1, - Float64Type: -1, + Float16Type: pymath.log(2.0**-24), # ≈ -16.64 + BFloat16Type: pymath.log(2.0**-133), # ≈ -92.18 + Float32Type: pymath.log(2.0**-149), # ≈ -103.28 + Float64Type: pymath.log(2.0**-1074), # ≈ -744.44 } # Largest x such that exp(x) is still representable in the given precision. From d7985b0e565718c5a3ed940dfce8faedd1df5883 Mon Sep 17 00:00:00 2001 From: szerdick Date: Sun, 10 May 2026 23:20:15 +0200 Subject: [PATCH 22/28] update tests to match new lower underflow bound --- .../filecheck/transforms/lower_exp_to_polynomial.mlir | 10 +++++----- tests/transforms/test_lower_exp_to_polynomial.py | 10 +++++++--- 2 files changed, 12 insertions(+), 8 deletions(-) diff --git a/tests/filecheck/transforms/lower_exp_to_polynomial.mlir b/tests/filecheck/transforms/lower_exp_to_polynomial.mlir index 930c98d11d..8ecb95c8cf 100644 --- a/tests/filecheck/transforms/lower_exp_to_polynomial.mlir +++ b/tests/filecheck/transforms/lower_exp_to_polynomial.mlir @@ -18,8 +18,8 @@ builtin.module { func.return %r : f32 } - // user lower_bound = -2 < underflow (-1) -> clamped to -1 - // user upper_bound = 5 < overflow (~88.72) -> kept at 5 + // 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 {acc_bound = 1.000000e-03 : f64, lower_bound = -2.000000e+00 : f32, upper_bound = 5.000000e+00 : f32} : f32 func.return %r : f32 @@ -44,14 +44,14 @@ builtin.module { // CHECK-NEXT: } // Without explicit lower/upper bounds, the polynomial domain defaults to -// [underflow, overflow] for the precision (f32 overflow ~= 88.72). +// [underflow, overflow] for the precision (f32: ~[-103.28, 88.72]). // CHECK: func.func @exp_f32_with_acc_bound(%[[XB:.*]]: f32) -> f32 { -// CHECK-NEXT: %[[RB:.*]] = polynomial.eval #polynomial.typed_chebyshev_polynomial<[{{.*}}]> : !polynomial.polynomial>, %[[XB]] {scheme = "clenshaw", domain_lower = -1.000000e+00 : f64, domain_upper = {{.*}} : f64} : f32 +// CHECK-NEXT: %[[RB:.*]] = polynomial.eval #polynomial.typed_chebyshev_polynomial<[{{.*}}]> : !polynomial.polynomial>, %[[XB]] {scheme = "clenshaw", domain_lower = {{.*}} : f64, domain_upper = {{.*}} : f64} : 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>, %[[XC]] {scheme = "clenshaw", domain_lower = -1.000000e+00 : f64, domain_upper = 5.000000e+00 : f64} : f32 +// CHECK-NEXT: %[[RC:.*]] = polynomial.eval #polynomial.typed_chebyshev_polynomial<[{{.*}}]> : !polynomial.polynomial>, %[[XC]] {scheme = "clenshaw", domain_lower = -2.000000e+00 : f64, domain_upper = 5.000000e+00 : f64} : f32 // CHECK-NEXT: func.return %[[RC]] : f32 // CHECK-NEXT: } diff --git a/tests/transforms/test_lower_exp_to_polynomial.py b/tests/transforms/test_lower_exp_to_polynomial.py index 7edba4ad58..c3fb3856c9 100644 --- a/tests/transforms/test_lower_exp_to_polynomial.py +++ b/tests/transforms/test_lower_exp_to_polynomial.py @@ -87,10 +87,14 @@ def test_acc_bounds_in_range(elem_ty: _FloatTy): @PRECISIONS def test_acc_lower_out_of_range(elem_ty: _FloatTy): - """acc_bound + lower < underflow -> lower clamped to underflow.""" + """acc_bound + 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, - {"acc_bound": 1e-3, "lower_bound": -100.0, "upper_bound": 0.5}, + {"acc_bound": 1e-3, "lower_bound": -1000.0, "upper_bound": 0.5}, ) eval_op = next(iter(module.body.block.ops)) assert isinstance(eval_op, polynomial.EvalOp) @@ -134,7 +138,7 @@ def test_acc_lower_only_in_range(elem_ty: _FloatTy): @PRECISIONS def test_acc_lower_only_out_of_range(elem_ty: _FloatTy): """acc_bound + lower only (out of range) -> [underflow, overflow].""" - module = _run_pass(elem_ty, {"acc_bound": 1e-3, "lower_bound": -100.0}) + module = _run_pass(elem_ty, {"acc_bound": 1e-3, "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 From 1f9777b8a59747c90da8fc26f7bdcbf46ce70fcb Mon Sep 17 00:00:00 2001 From: szerdick Date: Wed, 13 May 2026 11:08:45 +0200 Subject: [PATCH 23/28] changed acc_bound representation to be provided by the user --- xdsl/transforms/lower_exp_to_polynomial.py | 59 +++++++++++++++++----- 1 file changed, 45 insertions(+), 14 deletions(-) diff --git a/xdsl/transforms/lower_exp_to_polynomial.py b/xdsl/transforms/lower_exp_to_polynomial.py index ab24a0328d..88e344b724 100644 --- a/xdsl/transforms/lower_exp_to_polynomial.py +++ b/xdsl/transforms/lower_exp_to_polynomial.py @@ -39,6 +39,14 @@ Float64Type: 1024.0 * pymath.log(2.0), # ≈ 709.78 } +# Mantissa width including the hidden bit, per IEEE-754 / brain-float layout. +_PRECISION_BITS: dict[type, int] = { + Float16Type: 11, # 10 stored + 1 hidden + BFloat16Type: 8, # 7 stored + 1 hidden + Float32Type: 24, # 23 stored + 1 hidden + Float64Type: 53, # 52 stored + 1 hidden +} + # Hard cap on the polynomial degree the chooser will grow to. _MAX_DEGREE = 30 @@ -72,18 +80,39 @@ def _element_float_type(tp: Attribute) -> Attribute: def _choose_polynomial( - acc_bound: float, + max_bits_lost: float, lower: float, upper: float, + precision_bits: int, ) -> list[float]: - """Smallest degree such that the Lobatto Chebyshev bound for exp on [a,b] is <= acc_bound. - Capped at _MAX_DEGREE.""" + """Smallest degree such that the Chebyshev-Lobatto *bits-lost* upper bound + for exp on [lower, upper] in the target precision is <= max_bits_lost. + + bits_lost = log2(ULP error) — the number of low-order mantissa bits that + are unreliable. + max_bits_lost = 0 targets 1-ULP-quality (libm-grade); each + additional unit doubles the allowed ULP error. + max_bits_lost = -1 targets orrectly-rounded (<= 0.5 ULP). + + Derivation: the conservative Chebyshev approximation bound gives + |err|_inf <= exp(upper) * width^(d+1) / (2 * 4^d * (d+1)!). + Dividing by min|exp| = exp(lower) yields the relative error bound: + rel_err <= exp(width) * width^(d+1) / (2 * 4^d * (d+1)!). + Multiplying by 2^(precision_bits - 1) converts to ULP error, and taking + log2 gives bits_lost. Inherits the original bound's ~2x conservatism; + expect the picked degree to be ~1 higher than empirically optimal. + """ width = upper - lower degree = 0 - bound = pymath.exp(upper) * width - while bound > acc_bound and degree < _MAX_DEGREE: + # bits_lost upper bound at degree 0: log2(exp(width) * width) + (p - 1) + bits_lost = ( + width * pymath.log2(pymath.e) + pymath.log2(width) + (precision_bits - 1) + ) + while bits_lost > max_bits_lost and degree < _MAX_DEGREE: degree += 1 - bound *= width / (4 * (degree + 1)) + # multiplier bound(d+1) / bound(d) = width / (4 * (degree + 1)) + # in log2 space: = log2(width) - 2 - log2(degree + 1) + bits_lost += pymath.log2(width) - 2 - pymath.log2(degree + 1) return _chebyshev_coefficients(pymath.exp, degree, lower, upper) @@ -91,18 +120,20 @@ def _choose_polynomial( class LowerExpToPolynomial(RewritePattern): @op_type_rewrite_pattern def match_and_rewrite(self, op: math.ExpOp, rewriter: PatternRewriter) -> None: - # representable range for the precision format; specific to exp elem_ty = _element_float_type(op.operand.type) underflow = UNDERFLOW_LOWER_BOUND.get(type(elem_ty)) overflow = OVERFLOW_UPPER_BOUND.get(type(elem_ty)) - if underflow is None or overflow is None: + precision_bits = _PRECISION_BITS.get(type(elem_ty)) + if underflow is None or overflow is None or precision_bits is None: return - # if no acc_bound attr, do not optimize exp - acc_bound_attr = op.attributes.get("acc_bound") - if not isinstance(acc_bound_attr, FloatAttr): - return - acc_bound = acc_bound_attr.value.data + # If no acc bound, compute correctly-rounded result + max_bits_lost_attr = op.attributes.get("max_bits_lost") + max_bits_lost = ( + max_bits_lost_attr.value.data + if isinstance(max_bits_lost_attr, FloatAttr) + else -1.0 + ) # Polynomial domain: each side independently clamped to its representable # extreme. Missing lower_bound -> underflow; missing upper_bound -> overflow. @@ -120,7 +151,7 @@ def match_and_rewrite(self, op: math.ExpOp, rewriter: PatternRewriter) -> None: ) # choose which polynomial family and degree to use, and compute coefficients - coeffs = _choose_polynomial(acc_bound, lower, upper) + coeffs = _choose_polynomial(max_bits_lost, lower, upper, precision_bits) # insert polynomial into the IR replacing exp rewriter.replace_op( From 2411bd1f476582e821a6e56bda6c9d5f6d03353c Mon Sep 17 00:00:00 2001 From: szerdick Date: Wed, 13 May 2026 11:25:11 +0200 Subject: [PATCH 24/28] updated tests to use ulp error --- .../transforms/lower_exp_to_polynomial.mlir | 22 +++-- .../test_lower_exp_to_polynomial.py | 89 ++++++++++--------- xdsl/transforms/lower_exp_to_polynomial.py | 11 ++- 3 files changed, 69 insertions(+), 53 deletions(-) diff --git a/tests/filecheck/transforms/lower_exp_to_polynomial.mlir b/tests/filecheck/transforms/lower_exp_to_polynomial.mlir index 8ecb95c8cf..2f43a9ef84 100644 --- a/tests/filecheck/transforms/lower_exp_to_polynomial.mlir +++ b/tests/filecheck/transforms/lower_exp_to_polynomial.mlir @@ -1,51 +1,55 @@ // 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 } - // acc_bound is preserved on math.exp by step 1 and drives the polynomial + // 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_acc_bound(%x: f32) -> f32 { - %r = math.exp %x {acc_bound = 1.000000e-03 : f64} : f32 + 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 {acc_bound = 1.000000e-03 : f64, lower_bound = -2.000000e+00 : f32, upper_bound = 5.000000e+00 : 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 {acc_bound = 1.000000e-03 : f64, lower_bound = -5.000000e-01 : f32, upper_bound = 5.000000e-01 : 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 acc_bound, math.exp is left untouched. +// 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:.*]] = math.exp %[[X32]] : f32 +// CHECK-NEXT: %[[R32:.*]] = polynomial.eval #polynomial.typed_chebyshev_polynomial<[{{.*}}]> : !polynomial.polynomial>, %[[X32]] {scheme = "clenshaw", domain_lower = {{.*}} : f64, domain_upper = {{.*}} : f64} : f32 // CHECK-NEXT: func.return %[[R32]] : f32 // CHECK-NEXT: } // CHECK: func.func @exp_f64(%[[X64:.*]]: f64) -> f64 { -// CHECK-NEXT: %[[R64:.*]] = math.exp %[[X64]] : f64 +// CHECK-NEXT: %[[R64:.*]] = polynomial.eval #polynomial.typed_chebyshev_polynomial<[{{.*}}]> : !polynomial.polynomial>, %[[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_acc_bound(%[[XB:.*]]: f32) -> f32 { +// CHECK: func.func @exp_f32_with_max_bits_lost(%[[XB:.*]]: f32) -> f32 { // CHECK-NEXT: %[[RB:.*]] = polynomial.eval #polynomial.typed_chebyshev_polynomial<[{{.*}}]> : !polynomial.polynomial>, %[[XB]] {scheme = "clenshaw", domain_lower = {{.*}} : f64, domain_upper = {{.*}} : f64} : f32 // CHECK-NEXT: func.return %[[RB]] : f32 // CHECK-NEXT: } diff --git a/tests/transforms/test_lower_exp_to_polynomial.py b/tests/transforms/test_lower_exp_to_polynomial.py index c3fb3856c9..5a8c2fb389 100644 --- a/tests/transforms/test_lower_exp_to_polynomial.py +++ b/tests/transforms/test_lower_exp_to_polynomial.py @@ -8,7 +8,9 @@ Float32Type, Float64Type, FloatAttr, + IntegerAttr, ModuleOp, + i64, ) from xdsl.transforms.lower_exp_to_polynomial import ( OVERFLOW_UPPER_BOUND, @@ -23,16 +25,15 @@ "elem_ty", [Float16Type(), Float32Type(), Float64Type()] ) -# acc_bound is always provided in f64 -_F64 = Float64Type() - -def _run_pass(elem_ty: _FloatTy, attrs: dict[str, float]) -> ModuleOp: +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(): - attr_ty = _F64 if name == "acc_bound" else elem_ty - exp_op.attributes[name] = FloatAttr(value, attr_ty) + 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) @@ -41,27 +42,35 @@ def _run_pass(elem_ty: _FloatTy, attrs: dict[str, float]) -> ModuleOp: @PRECISIONS -def test_no_acc_no_bounds(elem_ty: _FloatTy): - """No acc_bound, no [lower, upper] -> math.exp left alone.""" +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, {}) - ops = list(module.body.block.ops) - assert len(ops) == 1 - assert isinstance(ops[0], math_dialect.ExpOp) + 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_LOWER_BOUND[type(elem_ty)] + assert eval_op.domain_upper.value.data == OVERFLOW_UPPER_BOUND[type(elem_ty)] @PRECISIONS -def test_no_acc_with_bounds(elem_ty: _FloatTy): - """No acc_bound, [lower, upper] given -> math.exp left alone.""" +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}) - ops = list(module.body.block.ops) - assert len(ops) == 1 - assert isinstance(ops[0], math_dialect.ExpOp) + 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_acc_only(elem_ty: _FloatTy): - """acc_bound only -> default interval [underflow, overflow].""" - module = _run_pass(elem_ty, {"acc_bound": 1e-3}) +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 @@ -71,11 +80,11 @@ def test_acc_only(elem_ty: _FloatTy): @PRECISIONS -def test_acc_bounds_in_range(elem_ty: _FloatTy): - """acc_bound + bounds entirely in range -> bounds used as-is.""" +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, - {"acc_bound": 1e-3, "lower_bound": -0.5, "upper_bound": 0.5}, + {"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) @@ -86,15 +95,15 @@ def test_acc_bounds_in_range(elem_ty: _FloatTy): @PRECISIONS -def test_acc_lower_out_of_range(elem_ty: _FloatTy): - """acc_bound + lower < underflow -> lower clamped to underflow. +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, - {"acc_bound": 1e-3, "lower_bound": -1000.0, "upper_bound": 0.5}, + {"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) @@ -105,15 +114,15 @@ def test_acc_lower_out_of_range(elem_ty: _FloatTy): @PRECISIONS -def test_acc_upper_out_of_range(elem_ty: _FloatTy): - """acc_bound + upper > overflow -> upper clamped to overflow. +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, - {"acc_bound": 1e-3, "lower_bound": -0.5, "upper_bound": 1000.0}, + {"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) @@ -124,9 +133,9 @@ def test_acc_upper_out_of_range(elem_ty: _FloatTy): @PRECISIONS -def test_acc_lower_only_in_range(elem_ty: _FloatTy): - """acc_bound + lower only (in range) -> [lower, overflow].""" - module = _run_pass(elem_ty, {"acc_bound": 1e-3, "lower_bound": -0.5}) +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 @@ -136,9 +145,9 @@ def test_acc_lower_only_in_range(elem_ty: _FloatTy): @PRECISIONS -def test_acc_lower_only_out_of_range(elem_ty: _FloatTy): - """acc_bound + lower only (out of range) -> [underflow, overflow].""" - module = _run_pass(elem_ty, {"acc_bound": 1e-3, "lower_bound": -1000.0}) +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 @@ -148,9 +157,9 @@ def test_acc_lower_only_out_of_range(elem_ty: _FloatTy): @PRECISIONS -def test_acc_upper_only_in_range(elem_ty: _FloatTy): - """acc_bound + upper only (in range) -> [underflow, upper].""" - module = _run_pass(elem_ty, {"acc_bound": 1e-3, "upper_bound": 0.5}) +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 @@ -160,9 +169,9 @@ def test_acc_upper_only_in_range(elem_ty: _FloatTy): @PRECISIONS -def test_acc_upper_only_out_of_range(elem_ty: _FloatTy): - """acc_bound + upper only (out of range) -> [underflow, overflow].""" - module = _run_pass(elem_ty, {"acc_bound": 1e-3, "upper_bound": 1000.0}) +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 diff --git a/xdsl/transforms/lower_exp_to_polynomial.py b/xdsl/transforms/lower_exp_to_polynomial.py index 88e344b724..4ad316abb1 100644 --- a/xdsl/transforms/lower_exp_to_polynomial.py +++ b/xdsl/transforms/lower_exp_to_polynomial.py @@ -11,6 +11,7 @@ Float32Type, Float64Type, FloatAttr, + IntegerAttr, ModuleOp, TensorType, VectorType, @@ -92,7 +93,7 @@ def _choose_polynomial( are unreliable. max_bits_lost = 0 targets 1-ULP-quality (libm-grade); each additional unit doubles the allowed ULP error. - max_bits_lost = -1 targets orrectly-rounded (<= 0.5 ULP). + max_bits_lost = -1 targets correctly-rounded (<= 0.5 ULP). Derivation: the conservative Chebyshev approximation bound gives |err|_inf <= exp(upper) * width^(d+1) / (2 * 4^d * (d+1)!). @@ -127,12 +128,14 @@ def match_and_rewrite(self, op: math.ExpOp, rewriter: PatternRewriter) -> None: if underflow is None or overflow is None or precision_bits is None: return - # If no acc bound, compute correctly-rounded result + # Optional max_bits_lost integer attribute; if absent, default to -1 + # (target correctly-rounded — all mantissa bits guaranteed correct, + # capped at _MAX_DEGREE). max_bits_lost_attr = op.attributes.get("max_bits_lost") max_bits_lost = ( max_bits_lost_attr.value.data - if isinstance(max_bits_lost_attr, FloatAttr) - else -1.0 + if isinstance(max_bits_lost_attr, IntegerAttr) + else -1 ) # Polynomial domain: each side independently clamped to its representable From 8e959b221692e9e45861403a4c9bbb80adb6ab5d Mon Sep 17 00:00:00 2001 From: szerdick Date: Wed, 13 May 2026 14:42:46 +0200 Subject: [PATCH 25/28] update API to use Evalop.get() --- xdsl/transforms/lower_exp_to_polynomial.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/xdsl/transforms/lower_exp_to_polynomial.py b/xdsl/transforms/lower_exp_to_polynomial.py index 4ad316abb1..bb4bb08888 100644 --- a/xdsl/transforms/lower_exp_to_polynomial.py +++ b/xdsl/transforms/lower_exp_to_polynomial.py @@ -6,6 +6,7 @@ from xdsl.context import Context from xdsl.dialects import math, polynomial from xdsl.dialects.builtin import ( + AnyFloat, BFloat16Type, Float16Type, Float32Type, @@ -159,9 +160,10 @@ def match_and_rewrite(self, op: math.ExpOp, rewriter: PatternRewriter) -> None: # insert polynomial into the IR replacing exp rewriter.replace_op( op, - polynomial.EvalOp( + polynomial.EvalOp.get( value=op.operand, - polynomial=tuple(coeffs), + coefficients=tuple(coeffs), + element_type=cast(AnyFloat, elem_ty), scheme=polynomial.EvalScheme.CLENSHAW, domain_lower=lower, domain_upper=upper, From b7fd02b6c010d1e601d0f66d8daa0cd85994ea37 Mon Sep 17 00:00:00 2001 From: szerdick Date: Wed, 13 May 2026 15:09:55 +0200 Subject: [PATCH 26/28] fit filecheck tests --- tests/filecheck/transforms/lower_exp_to_polynomial.mlir | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/filecheck/transforms/lower_exp_to_polynomial.mlir b/tests/filecheck/transforms/lower_exp_to_polynomial.mlir index 2f43a9ef84..c3914c06f8 100644 --- a/tests/filecheck/transforms/lower_exp_to_polynomial.mlir +++ b/tests/filecheck/transforms/lower_exp_to_polynomial.mlir @@ -38,7 +38,7 @@ builtin.module { // 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>, %[[X32]] {scheme = "clenshaw", domain_lower = {{.*}} : f64, domain_upper = {{.*}} : f64} : f32 +// CHECK-NEXT: %[[R32:.*]] = polynomial.eval #polynomial.typed_chebyshev_polynomial<[{{.*}}]> : !polynomial.polynomial>, %[[X32]] {scheme = "clenshaw", domain_lower = {{.*}} : f32, domain_upper = {{.*}} : f32} : f32 // CHECK-NEXT: func.return %[[R32]] : f32 // CHECK-NEXT: } @@ -50,16 +50,16 @@ builtin.module { // 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>, %[[XB]] {scheme = "clenshaw", domain_lower = {{.*}} : f64, domain_upper = {{.*}} : f64} : f32 +// CHECK-NEXT: %[[RB:.*]] = polynomial.eval #polynomial.typed_chebyshev_polynomial<[{{.*}}]> : !polynomial.polynomial>, %[[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>, %[[XC]] {scheme = "clenshaw", domain_lower = -2.000000e+00 : f64, domain_upper = 5.000000e+00 : f64} : f32 +// CHECK-NEXT: %[[RC:.*]] = polynomial.eval #polynomial.typed_chebyshev_polynomial<[{{.*}}]> : !polynomial.polynomial>, %[[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>, %[[XT]] {scheme = "clenshaw", domain_lower = -5.000000e-01 : f64, domain_upper = 5.000000e-01 : f64} : f32 +// CHECK-NEXT: %[[RT:.*]] = polynomial.eval #polynomial.typed_chebyshev_polynomial<[{{.*}}]> : !polynomial.polynomial>, %[[XT]] {scheme = "clenshaw", domain_lower = -5.000000e-01 : f32, domain_upper = 5.000000e-01 : f32} : f32 // CHECK-NEXT: func.return %[[RT]] : f32 // CHECK-NEXT: } From 0a124ffdb5871c432f390de053be16abe6e5f91f Mon Sep 17 00:00:00 2001 From: szerdick Date: Thu, 14 May 2026 16:29:06 +0200 Subject: [PATCH 27/28] rename pass and add test --- ...val.mlir => polynomial_eval_to_arith.mlir} | 2 +- .../test_polynomial_eval_to_arith.py | 39 +++++++++++++++++++ xdsl/transforms/__init__.py | 8 ++-- ...al_eval.py => polynomial_eval_to_arith.py} | 28 ++++++------- .../transforms/test_lower_linalg_to_snitch.py | 4 +- 5 files changed, 61 insertions(+), 20 deletions(-) rename tests/filecheck/transforms/{expand_polynomial_eval.mlir => polynomial_eval_to_arith.mlir} (98%) create mode 100644 tests/transforms/test_polynomial_eval_to_arith.py rename xdsl/transforms/{expand_polynomial_eval.py => polynomial_eval_to_arith.py} (85%) diff --git a/tests/filecheck/transforms/expand_polynomial_eval.mlir b/tests/filecheck/transforms/polynomial_eval_to_arith.mlir similarity index 98% rename from tests/filecheck/transforms/expand_polynomial_eval.mlir rename to tests/filecheck/transforms/polynomial_eval_to_arith.mlir index 38994fce98..76b26e10fe 100644 --- a/tests/filecheck/transforms/expand_polynomial_eval.mlir +++ b/tests/filecheck/transforms/polynomial_eval_to_arith.mlir @@ -1,4 +1,4 @@ -// RUN: xdsl-opt -p expand-polynomial-eval %s | filecheck %s +// RUN: xdsl-opt -p polynomial-eval-to-arith %s | filecheck %s builtin.module { // f64 with domain bounds [-5, 0], degree 2. diff --git a/tests/transforms/test_polynomial_eval_to_arith.py b/tests/transforms/test_polynomial_eval_to_arith.py new file mode 100644 index 0000000000..df28de6bc5 --- /dev/null +++ b/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})" + ) diff --git a/xdsl/transforms/__init__.py b/xdsl/transforms/__init__.py index d6a3d55cf0..5ba73179a7 100644 --- a/xdsl/transforms/__init__.py +++ b/xdsl/transforms/__init__.py @@ -91,10 +91,10 @@ def get_expand_math_to_polynomials(): return expand_math_to_polynomials.ExpandMathToPolynomialsPass - def get_expand_polynomial_eval(): - from xdsl.transforms import expand_polynomial_eval + def get_polynomial_eval_to_arith(): + from xdsl.transforms import polynomial_eval_to_arith - return expand_polynomial_eval.ExpandPolynomialEvalPass + return polynomial_eval_to_arith.PolynomialEvalToArithPass def get_lower_exp_to_polynomial(): from xdsl.transforms import lower_exp_to_polynomial @@ -703,7 +703,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, - "expand-polynomial-eval": get_expand_polynomial_eval, + "polynomial-eval-to-arith": get_polynomial_eval_to_arith, "lower-exp-to-polynomial": get_lower_exp_to_polynomial, "convert-linalg-to-loops": get_convert_linalg_to_loops, "convert-linalg-to-memref-stream": get_convert_linalg_to_memref_stream, diff --git a/xdsl/transforms/expand_polynomial_eval.py b/xdsl/transforms/polynomial_eval_to_arith.py similarity index 85% rename from xdsl/transforms/expand_polynomial_eval.py rename to xdsl/transforms/polynomial_eval_to_arith.py index b49565844b..ee2fa50265 100644 --- a/xdsl/transforms/expand_polynomial_eval.py +++ b/xdsl/transforms/polynomial_eval_to_arith.py @@ -29,6 +29,7 @@ RewritePattern, op_type_rewrite_pattern, ) +from xdsl.utils.exceptions import PassFailedException def _float_constant( @@ -104,7 +105,7 @@ def expand_clenshaw( return rewriter.insert(arith.SubfOp(add.result, b_prev2.result)) -class ExpandPolynomialEval(RewritePattern): +class PolynomialEvalToArith(RewritePattern): """Replace each `polynomial.eval` op with the arith ops for its scheme.""" @op_type_rewrite_pattern @@ -116,21 +117,22 @@ def match_and_rewrite( upper = op.domain_upper.value.data if op.domain_upper is not None else None scheme = op.eval_scheme - if scheme is polynomial.EvalScheme.CLENSHAW: - expanded = expand_clenshaw(op, rewriter, coeffs, lower, upper) - else: - # 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 NotImplementedError( - f"polynomial.eval scheme {scheme.value!r} has no lowering" - ) + 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 ExpandPolynomialEvalPass(ModulePass): +class PolynomialEvalToArithPass(ModulePass): """ Expand `polynomial.eval` ops to arithmetic operations. @@ -139,10 +141,10 @@ class ExpandPolynomialEvalPass(ModulePass): parameters. """ - name = "expand-polynomial-eval" + name = "polynomial-eval-to-arith" def apply(self, ctx: Context, op: ModuleOp) -> None: PatternRewriteWalker( - ExpandPolynomialEval(), + PolynomialEvalToArith(), apply_recursively=False, ).rewrite_module(op) diff --git a/xdsl/transforms/test_lower_linalg_to_snitch.py b/xdsl/transforms/test_lower_linalg_to_snitch.py index 234770df34..7b0c45ff0c 100644 --- a/xdsl/transforms/test_lower_linalg_to_snitch.py +++ b/xdsl/transforms/test_lower_linalg_to_snitch.py @@ -19,7 +19,6 @@ convert_memref_stream_to_snitch_stream, convert_riscv_scf_for_to_frep, expand_math_to_polynomials, - expand_polynomial_eval, lower_affine, lower_exp_to_polynomial, lower_snitch, @@ -31,6 +30,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, @@ -61,7 +61,7 @@ convert_scf_to_riscv_scf.ConvertScfToRiscvPass(), expand_math_to_polynomials.ExpandMathToPolynomialsPass(), lower_exp_to_polynomial.LowerExpToPolynomialPass(), - expand_polynomial_eval.ExpandPolynomialEvalPass(), + polynomial_eval_to_arith.PolynomialEvalToArithPass(), convert_arith_to_riscv_snitch.ConvertArithToRiscvSnitchPass(), convert_arith_to_riscv.ConvertArithToRiscvPass(), convert_func_to_riscv_func.ConvertFuncToRiscvFuncPass(), From 5e7d8c42b9675c5c90c805bf68689a24cffbecff Mon Sep 17 00:00:00 2001 From: szerdick Date: Thu, 14 May 2026 17:55:53 +0200 Subject: [PATCH 28/28] one test failing bc of acc bound in f64 --- .../test_lower_exp_to_polynomial.py | 35 ++++++++++++------- 1 file changed, 23 insertions(+), 12 deletions(-) diff --git a/tests/transforms/test_lower_exp_to_polynomial.py b/tests/transforms/test_lower_exp_to_polynomial.py index 5a8c2fb389..f2526e4e63 100644 --- a/tests/transforms/test_lower_exp_to_polynomial.py +++ b/tests/transforms/test_lower_exp_to_polynomial.py @@ -26,6 +26,17 @@ ) +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) @@ -50,8 +61,8 @@ def test_default_no_bounds(elem_ty: _FloatTy): 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_LOWER_BOUND[type(elem_ty)] - assert eval_op.domain_upper.value.data == OVERFLOW_UPPER_BOUND[type(elem_ty)] + assert eval_op.domain_lower.value.data == _underflow_at(elem_ty) + assert eval_op.domain_upper.value.data == _overflow_at(elem_ty) @PRECISIONS @@ -75,8 +86,8 @@ def test_max_bits_lost_only(elem_ty: _FloatTy): 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_LOWER_BOUND[type(elem_ty)] - assert eval_op.domain_upper.value.data == OVERFLOW_UPPER_BOUND[type(elem_ty)] + assert eval_op.domain_lower.value.data == _underflow_at(elem_ty) + assert eval_op.domain_upper.value.data == _overflow_at(elem_ty) @PRECISIONS @@ -109,7 +120,7 @@ def test_max_bits_lost_lower_out_of_range(elem_ty: _FloatTy): 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_LOWER_BOUND[type(elem_ty)] + assert eval_op.domain_lower.value.data == _underflow_at(elem_ty) assert eval_op.domain_upper.value.data == 0.5 @@ -129,7 +140,7 @@ def test_max_bits_lost_upper_out_of_range(elem_ty: _FloatTy): 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_UPPER_BOUND[type(elem_ty)] + assert eval_op.domain_upper.value.data == _overflow_at(elem_ty) @PRECISIONS @@ -141,7 +152,7 @@ def test_max_bits_lost_lower_only_in_range(elem_ty: _FloatTy): 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_UPPER_BOUND[type(elem_ty)] + assert eval_op.domain_upper.value.data == _overflow_at(elem_ty) @PRECISIONS @@ -152,8 +163,8 @@ def test_max_bits_lost_lower_only_out_of_range(elem_ty: _FloatTy): 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_LOWER_BOUND[type(elem_ty)] - assert eval_op.domain_upper.value.data == OVERFLOW_UPPER_BOUND[type(elem_ty)] + assert eval_op.domain_lower.value.data == _underflow_at(elem_ty) + assert eval_op.domain_upper.value.data == _overflow_at(elem_ty) @PRECISIONS @@ -164,7 +175,7 @@ def test_max_bits_lost_upper_only_in_range(elem_ty: _FloatTy): 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_LOWER_BOUND[type(elem_ty)] + assert eval_op.domain_lower.value.data == _underflow_at(elem_ty) assert eval_op.domain_upper.value.data == 0.5 @@ -176,5 +187,5 @@ def test_max_bits_lost_upper_only_out_of_range(elem_ty: _FloatTy): 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_LOWER_BOUND[type(elem_ty)] - assert eval_op.domain_upper.value.data == OVERFLOW_UPPER_BOUND[type(elem_ty)] + assert eval_op.domain_lower.value.data == _underflow_at(elem_ty) + assert eval_op.domain_upper.value.data == _overflow_at(elem_ty)