Merge pull request google#210 from j2kun:poly-mul-cpu-runner-test-fra…

…mework

PiperOrigin-RevId: 576859638

.bazelrc

@@ -2,3 +2,4 @@ common --action_env=BAZEL_CXXOPTS=-std=c++17
 common --cxxopt='-std=c++17'
 common --copt=-fdiagnostics-color=always
 common --test_output=errors
+common -c dbg

requirements.txt

@@ -1 +1,6 @@
 lit==16.0.6
+
+# for tests/poly/runner/generate_test_cases.py
+tomli==2.0.1
+sympy==1.12
+mpmath==1.3.0

tests/poly/runner/BUILD

@@ -0,0 +1,27 @@
+load("@rules_python//python:py_binary.bzl", "py_binary")
+load("//bazel:lit.bzl", "glob_lit_tests")
+
+package(
+    default_applicable_licenses = ["@heir//:license"],
+    default_visibility = ["//visibility:public"],
+)
+
+py_binary(
+    name = "generate_test_cases",
+    srcs = [
+        "generate_test_cases.py",
+    ],
+    python_version = "PY3",
+    srcs_version = "PY3",
+    deps = [
+        "@heir_pip_deps_sympy//:pkg",
+        "@heir_pip_deps_tomli//:pkg",
+    ],
+)
+
+glob_lit_tests(
+    name = "all_tests",
+    data = ["@heir//tests:test_utilities"],
+    driver = "@heir//tests:run_lit.sh",
+    test_file_exts = ["mlir"],
+)

tests/poly/runner/README.md

@@ -0,0 +1,7 @@
+To re-generate the tests in this directory:
+
+```
+cd tests/poly/runner
+rm lower_mul_*.mlir
+bazel run generate_test_cases -- --tests_toml_path $PWD/lower_mul_tests.toml --output_test_stem=$PWD/lower_mul_
+```

tests/poly/runner/generate_test_cases.py

@@ -0,0 +1,228 @@
+"""Generate mlir-cpu-runner tests for lowering poly.mul ops from a config."""
+import argparse
+import sys
+
+import sympy
+import tomli
+
+HEADER = """// WARNING: this file is autogenerated. Do not edit manually, instead see
+// tests/poly/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 }
+
+"""
+
+
+def test_setup(i: int):
+  entry_def = '// REDEFINE: %{entry} = test_' + str(i)
+  check_def = '// REDEFINE: %{check_prefix} = CHECK_TEST_' + str(i)
+  run = '// RUN: %{compile} | %{run} | %{check}'
+  return '\n'.join([entry_def, check_def, run])
+
+
+# Params:
+#  - ideal
+#  - cmod
+#  - p0
+#  - p1
+#  - test_number
+#  - gen_tensor_op
+#  - degree
+#  - expected
+#  - coefficient_list
+TEST_TEMPLATE = """
+#ideal_{{test_number}} = #poly.polynomial<{{ideal}}>
+#ring_{{test_number}} = #poly.ring<cmod={{cmod}}, ideal=#ideal_{{test_number}}>
+!poly_ty_{{test_number}} = !poly.poly<#ring_{{test_number}}>
+
+func.func @test_{{test_number}}() {{{{
+  %const0 = arith.constant 0 : index
+  %0 = poly.constant <{{p0}}> : !poly_ty_{{test_number}}
+  %1 = poly.constant <{{p1}}> : !poly_ty_{{test_number}}
+  %2 = poly.mul(%0, %1) : !poly_ty_{{test_number}}
+
+{gen_tensor_op}
+  %ref = bufferization.to_memref %tensor : memref<{{degree}}xi32>
+  %U = memref.cast %ref : memref<{{degree}}xi32> to memref<*xi32>
+  func.call @printMemrefI32(%U) : (memref<*xi32>) -> ()
+  return
+}}}}
+// expected_result: {{expected}}
+// CHECK_TEST_{{test_number}}: {{coefficient_list}}
+"""
+
+# We need to compare the final output in i32s because that's the only type
+# supported by the mlir-cpu-runner. So we need to truncate or extend the bits.
+# I have not looked into the possibility of adding more functions to the
+# upstream mlir-cpu-runner.
+GEN_TENSOR_OP_TEMPLATE_WITH_TRUNC = """
+  %3 = poly.to_tensor %2 : !poly_ty_{{test_number}} -> tensor<{{degree}}x{{container_type}}>
+  %tensor = arith.{trunc_ext_op} %3 : tensor<{{degree}}x{{container_type}}> to tensor<{{degree}}xi32>
+"""
+GEN_TENSOR_OP_TEMPLATE = """
+  %tensor = poly.to_tensor %2 : !poly_ty_{test_number} -> tensor<{degree}x{container_type}>
+"""
+
+
+def make_test_template(container_bit_width: int):
+  if container_bit_width != 32:
+    trunc_ext_op = 'trunci' if container_bit_width > 32 else 'extsi'
+    return TEST_TEMPLATE.format(
+        gen_tensor_op=GEN_TENSOR_OP_TEMPLATE_WITH_TRUNC.format(
+            trunc_ext_op=trunc_ext_op
+        )
+    )
+  else:
+    return TEST_TEMPLATE.format(gen_tensor_op=GEN_TENSOR_OP_TEMPLATE)
+
+
+parser = argparse.ArgumentParser(
+    description=(
+        'Generate a list of mlir-cpu-runner integration tests for lowering'
+        ' poly.mul ops'
+    )
+)
+parser.add_argument(
+    '--tests_toml_path',
+    type=str,
+    help='A filepath to a toml file containing a list of tests to generate.',
+)
+parser.add_argument(
+    '--output_test_file',
+    type=str,
+    help='A filepath to the output file containing all the tests.',
+)
+parser.add_argument(
+    '--output_test_stem',
+    type=str,
+    help=(
+        'If specified, generate each test as its own file with this argument'
+        ' as the path stem.'
+    ),
+)
+
+
+def get_key_or_err(the_dict, key: str, error: str = None):
+  error = error or f'[[test]] Missing key {key}, parsed dict={the_dict}'
+  try:
+    return the_dict[key]
+  except KeyError:
+    print(error)
+    sys.exit(1)
+
+
+def parse_poly(poly_str: str) -> list[tuple[int, int]]:
+  """Parse a polynomial string into a list of coeff-degree pairs."""
+  terms = [x.strip().split('x**') for x in poly_str.split('+')]
+  term_dict = dict()
+  for term in terms:
+    if term[0]:
+      term_coeff = int(term[0])
+    else:
+      term_coeff = 1
+    if len(term) == 1:
+      term_dict[0] = term_coeff
+    else:
+      degree = int(term[1])
+      term_dict[degree] = term_coeff
+
+  return list((coeff, degree) for (degree, coeff) in term_dict.items())
+
+
+def parse_to_sympy(poly_str: str, var: sympy.Symbol, cmod: int):
+  terms = parse_poly(poly_str)
+  poly = 0
+  for coeff, degree in terms:
+    poly += coeff * var**degree
+  return poly.as_poly(domain=f'ZZ[{cmod}]')
+
+
+def main(args: argparse.Namespace) -> None:
+  if not args.tests_toml_path:
+    print('No test config was passed via --tests_toml_path')
+    sys.exit(1)
+
+  if not args.output_test_file and not args.output_test_stem:
+    print('Must pass one of --output_test_file or --output_test_stem')
+    sys.exit(1)
+
+  with open(args.tests_toml_path, 'r') as infile:
+    config = tomli.load(infile)
+
+  tests = []
+  try:
+    tests = config['test']
+  except KeyError:
+    print('TOML file must contain one or more sections like [[test]]')
+    sys.exit(1)
+  test_count = len(tests)
+
+  print(f'Generating {test_count} tests...')
+  output_tests = []
+  for i, test in enumerate(tests):
+    (ideal, cmod, p0, p1, container_type) = (
+        get_key_or_err(test, s)
+        for s in ['ideal', 'cmod', 'p0', 'p1', 'container_type']
+    )
+
+    x = sympy.Symbol('x')
+    parsed_ideal = parse_to_sympy(ideal, x, cmod)
+    parsed_p0 = parse_to_sympy(p0, x, cmod)
+    parsed_p1 = parse_to_sympy(p1, x, cmod)
+
+    expected_remainder = sympy.rem(parsed_p0 * parsed_p1, parsed_ideal, x)
+    print(
+        f'{expected_remainder.domain} : ({p0}) * ({p1}) ='
+        f' {expected_remainder.as_expr()} mod ({ideal})'
+    )
+    coeff_list_len = parsed_ideal.degree()
+    expected_coeffs = list(reversed(expected_remainder.all_coeffs()))
+    if len(expected_coeffs) < coeff_list_len:
+      expected_coeffs = expected_coeffs + [0] * (
+          coeff_list_len - len(expected_coeffs)
+      )
+
+    container_width = int(container_type[1:])
+
+    output_tests.append(
+        '\n'.join([
+            test_setup(i),
+            make_test_template(container_width).format(
+                ideal=ideal,
+                cmod=cmod,
+                p0=p0,
+                p1=p1,
+                test_number=i,
+                container_type=container_type,
+                degree=parsed_ideal.degree(),
+                expected=expected_remainder,
+                coefficient_list=expected_coeffs,
+            ),
+        ])
+    )
+
+  if args.output_test_stem:
+    for i, test in enumerate(output_tests):
+      with open(f'{args.output_test_stem}{i}.mlir', 'w') as outfile:
+        outfile.write(HEADER)
+        outfile.write(test)
+  else:
+    with open(args.output_test_file, 'w') as outfile:
+      outfile.write(HEADER)
+      outfile.write('\n'.join(output_tests))
+
+  print('Done')
+
+
+if __name__ == '__main__':
+  main(parser.parse_args())

tests/poly/runner/lower_mul_0.mlir

@@ -0,0 +1,39 @@
+// WARNING: this file is autogenerated. Do not edit manually, instead see
+// tests/poly/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_0
+// REDEFINE: %{check_prefix} = CHECK_TEST_0
+// RUN: %{compile} | %{run} | %{check}
+
+#ideal_0 = #poly.polynomial<1 + x**12>
+#ring_0 = #poly.ring<cmod=4294967296, ideal=#ideal_0>
+!poly_ty_0 = !poly.poly<#ring_0>
+
+func.func @test_0() {
+  %const0 = arith.constant 0 : index
+  %0 = poly.constant <1 + x**10> : !poly_ty_0
+  %1 = poly.constant <1 + x**11> : !poly_ty_0
+  %2 = poly.mul(%0, %1) : !poly_ty_0
+
+
+  %tensor = poly.to_tensor %2 : !poly_ty_0 -> 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 - x**9 + 1, x, domain='ZZ[4294967296]')
+// CHECK_TEST_0: [1, 0, 0, 0, 0, 0, 0, 0, 0, -1, 1, 1]

tests/poly/runner/lower_mul_tests.toml

@@ -0,0 +1,6 @@
+[[test]]
+ideal = "1 + x**12"
+cmod = 4294967296  # 2**32
+p0 = "1 + x**10"
+p1 = "1 + x**11"
+container_type = "i32"

tools/BUILD

@@ -23,21 +23,34 @@ cc_binary(
         "@heir//lib/Dialect/Secret/IR:Dialect",
         "@heir//lib/Dialect/Secret/Transforms",
         "@llvm-project//mlir:AffineDialect",
+        "@llvm-project//mlir:AffineToStandard",
         "@llvm-project//mlir:AffineTransforms",
         "@llvm-project//mlir:AllPassesAndDialects",
         "@llvm-project//mlir:ArithDialect",
+        "@llvm-project//mlir:ArithToLLVM",
         "@llvm-project//mlir:ArithTransforms",
         "@llvm-project//mlir:BufferizationTransforms",
+        "@llvm-project//mlir:ControlFlowToLLVM",
         "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:FuncToLLVM",
+        "@llvm-project//mlir:FuncTransforms",
+        "@llvm-project//mlir:IndexToLLVM",
         "@llvm-project//mlir:LLVMDialect",
         "@llvm-project//mlir:LinalgTransforms",
         "@llvm-project//mlir:MemRefDialect",
+        "@llvm-project//mlir:MemRefToLLVM",
         "@llvm-project//mlir:MemRefTransforms",
         "@llvm-project//mlir:MlirOptLib",
         "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:ReconcileUnrealizedCasts",
         "@llvm-project//mlir:SCFDialect",
+        "@llvm-project//mlir:SCFToControlFlow",
+        "@llvm-project//mlir:TensorToLinalg",
         "@llvm-project//mlir:TensorTransforms",
+        "@llvm-project//mlir:TosaDialect",
+        "@llvm-project//mlir:TosaToArith",
         "@llvm-project//mlir:TosaToLinalg",
+        "@llvm-project//mlir:TosaToTensor",
         "@llvm-project//mlir:Transforms",
     ],
 )