diff --git a/tests/polynomial/lower_mul.mlir b/tests/polynomial/lower_mul.mlir index 3d2b2d428..5da24e196 100644 --- a/tests/polynomial/lower_mul.mlir +++ b/tests/polynomial/lower_mul.mlir @@ -91,6 +91,3 @@ func.func @lower_poly_mul(%poly0: !poly_ty, %poly1: !poly_ty) -> !poly_ty { %poly2 = polynomial.mul(%poly0, %poly1) {ring = #ring} : !poly_ty return %poly2 : !poly_ty } - -// TODO(https://github.com/google/heir/issues/199): add proper tests for a -// non-1 divisor leading coefficient diff --git a/tests/polynomial/runner/generate_test_cases.py b/tests/polynomial/runner/generate_test_cases.py index 62cf06d2d..2f22bdf3e 100644 --- a/tests/polynomial/runner/generate_test_cases.py +++ b/tests/polynomial/runner/generate_test_cases.py @@ -1,4 +1,5 @@ """Generate mlir-cpu-runner tests for lowering polynomial.mul ops from a config.""" + import argparse import sys @@ -126,6 +127,7 @@ def parse_polynomial(poly_str: str) -> list[tuple[int, int]]: terms = [x.strip().split('x**') for x in poly_str.split('+')] term_dict = dict() for term in terms: + term = [x.strip() for x in term] if term[0]: term_coeff = int(term[0]) else: @@ -179,6 +181,7 @@ def main(args: argparse.Namespace) -> None: parsed_ideal = parse_to_sympy(ideal, x, cmod) parsed_p0 = parse_to_sympy(p0, x, cmod) parsed_p1 = parse_to_sympy(p1, x, cmod) + domain = parsed_p0.domain expected_remainder = sympy.rem(parsed_p0 * parsed_p1, parsed_ideal, x) print( @@ -187,6 +190,17 @@ def main(args: argparse.Namespace) -> None: ) coeff_list_len = parsed_ideal.degree() expected_coeffs = list(reversed(expected_remainder.all_coeffs())) + + # For whatever reason, sympy won't preserve the domain of the coefficients + # after `rem`, so I have to manually convert any fractional coefficients to + # their modular inverse equivalents. + for j, exp_coeff in enumerate(expected_coeffs): + if exp_coeff.is_rational and not exp_coeff.is_integer: + domain_p = domain.convert(exp_coeff.p) + domain_q = domain.convert(sympy.mod_inverse(exp_coeff.q, cmod)) + result = (domain_p * domain_q) % cmod + expected_coeffs[j] = result + if len(expected_coeffs) < coeff_list_len: expected_coeffs = expected_coeffs + [0] * ( coeff_list_len - len(expected_coeffs) diff --git a/tests/polynomial/runner/lower_mul_6.mlir b/tests/polynomial/runner/lower_mul_6.mlir new file mode 100644 index 000000000..d0eaae3ec --- /dev/null +++ b/tests/polynomial/runner/lower_mul_6.mlir @@ -0,0 +1,40 @@ +// WARNING: this file is autogenerated. Do not edit manually, instead see +// tests/polynomial/runner/generate_test_cases.py + +//------------------------------------------------------- +// entry and check_prefix are re-set per test execution +// DEFINE: %{entry} = +// DEFINE: %{check_prefix} = + +// DEFINE: %{compile} = heir-opt %s --heir-polynomial-to-llvm +// DEFINE: %{run} = mlir-cpu-runner -e %{entry} -entry-point-result=void --shared-libs="%mlir_lib_dir/libmlir_c_runner_utils%shlibext,%mlir_runner_utils" +// DEFINE: %{check} = FileCheck %s --check-prefix=%{check_prefix} +//------------------------------------------------------- + +func.func private @printMemrefI32(memref<*xi32>) attributes { llvm.emit_c_interface } + +// REDEFINE: %{entry} = test_6 +// REDEFINE: %{check_prefix} = CHECK_TEST_6 +// RUN: %{compile} | %{run} | %{check} + +#ideal_6 = #polynomial.polynomial<3 + x**12> +#ring_6 = #polynomial.ring +!poly_ty_6 = !polynomial.polynomial<#ring_6> + +func.func @test_6() { + %const0 = arith.constant 0 : index + %0 = polynomial.constant <1 + x**10> : !poly_ty_6 + %1 = polynomial.constant <1 + x**11> : !poly_ty_6 + %2 = polynomial.mul(%0, %1) : !poly_ty_6 + + + %3 = polynomial.to_tensor %2 : !poly_ty_6 -> tensor<12xi4> + %tensor = arith.extsi %3 : tensor<12xi4> to tensor<12xi32> + + %ref = bufferization.to_memref %tensor : memref<12xi32> + %U = memref.cast %ref : memref<12xi32> to memref<*xi32> + func.call @printMemrefI32(%U) : (memref<*xi32>) -> () + return +} +// expected_result: Poly(x**11 + x**10 - 3*x**9 + 1, x, domain='ZZ[16]') +// CHECK_TEST_6: [1, 0, 0, 0, 0, 0, 0, 0, 0, -3, 1, 1] diff --git a/tests/polynomial/runner/lower_mul_7.mlir b/tests/polynomial/runner/lower_mul_7.mlir new file mode 100644 index 000000000..46a9ecd84 --- /dev/null +++ b/tests/polynomial/runner/lower_mul_7.mlir @@ -0,0 +1,40 @@ +// WARNING: this file is autogenerated. Do not edit manually, instead see +// tests/polynomial/runner/generate_test_cases.py + +//------------------------------------------------------- +// entry and check_prefix are re-set per test execution +// DEFINE: %{entry} = +// DEFINE: %{check_prefix} = + +// DEFINE: %{compile} = heir-opt %s --heir-polynomial-to-llvm +// DEFINE: %{run} = mlir-cpu-runner -e %{entry} -entry-point-result=void --shared-libs="%mlir_lib_dir/libmlir_c_runner_utils%shlibext,%mlir_runner_utils" +// DEFINE: %{check} = FileCheck %s --check-prefix=%{check_prefix} +//------------------------------------------------------- + +func.func private @printMemrefI32(memref<*xi32>) attributes { llvm.emit_c_interface } + +// REDEFINE: %{entry} = test_7 +// REDEFINE: %{check_prefix} = CHECK_TEST_7 +// RUN: %{compile} | %{run} | %{check} + +#ideal_7 = #polynomial.polynomial<3 + 5 x**12> +#ring_7 = #polynomial.ring +!poly_ty_7 = !polynomial.polynomial<#ring_7> + +func.func @test_7() { + %const0 = arith.constant 0 : index + %0 = polynomial.constant <1 + x**10> : !poly_ty_7 + %1 = polynomial.constant <1 + x**11> : !poly_ty_7 + %2 = polynomial.mul(%0, %1) : !poly_ty_7 + + + %3 = polynomial.to_tensor %2 : !poly_ty_7 -> tensor<12xi4> + %tensor = arith.extsi %3 : tensor<12xi4> to tensor<12xi32> + + %ref = bufferization.to_memref %tensor : memref<12xi32> + %U = memref.cast %ref : memref<12xi32> to memref<*xi32> + func.call @printMemrefI32(%U) : (memref<*xi32>) -> () + return +} +// expected_result: Poly(x**11 + x**10 - 3/5*x**9 + 1, x, domain='ZZ(16)') +// CHECK_TEST_7: [1, 0, 0, 0, 0, 0, 0, 0, 0, -7, 1, 1] diff --git a/tests/polynomial/runner/lower_mul_tests.toml b/tests/polynomial/runner/lower_mul_tests.toml index 0ce6dd676..a60ac4a3f 100644 --- a/tests/polynomial/runner/lower_mul_tests.toml +++ b/tests/polynomial/runner/lower_mul_tests.toml @@ -40,6 +40,20 @@ p0 = "1 + x**2" p1 = "1 + x**3" container_type = "i4" +[[test]] +ideal = "3 + x**12" +cmod = 16 +p0 = "1 + x**10" +p1 = "1 + x**11" +container_type = "i4" + +[[test]] +ideal = "3 + 5 x**12" +cmod = 16 +p0 = "1 + x**10" +p1 = "1 + x**11" +container_type = "i4" + # TODO(https://github.com/google/heir/issues/220): restore once # we can use emulate-wide-int in the pipeline # [[test]]