Merge pull request google#223 from j2kun:more-poly-mul-tests

PiperOrigin-RevId: 579300388

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

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)

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<cmod=16, ideal=#ideal_6>
+!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]

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<cmod=16, ideal=#ideal_7>
+!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]

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]]