diff --git a/lib/Transforms/YosysOptimizer/yosys/BUILD b/lib/Transforms/YosysOptimizer/yosys/BUILD index e89c5930c..648361243 100644 --- a/lib/Transforms/YosysOptimizer/yosys/BUILD +++ b/lib/Transforms/YosysOptimizer/yosys/BUILD @@ -11,3 +11,7 @@ filegroup( "*.v", ]), ) + +exports_files([ + "*.v", +]) diff --git a/tests/jaxite/end_to_end/BUILD b/tests/jaxite/end_to_end/BUILD index f749c8a43..87d5be89a 100644 --- a/tests/jaxite/end_to_end/BUILD +++ b/tests/jaxite/end_to_end/BUILD @@ -22,8 +22,12 @@ py_library( jaxite_end_to_end_test( name = "add_one_lut3", - heir_opt_flags = "--tosa-to-boolean-jaxite=entry-function=test_add_one_lut3", + entry_function_flag = "entry-function=test_add_one_lut3", mlir_src = "add_one_lut3.mlir", test_src = "add_one_lut3_test.py", deps = [":test_utils"], ) + +exports_files([ + "add_one_lut3.mlir", +]) diff --git a/tests/jaxite/end_to_end/test.bzl b/tests/jaxite/end_to_end/test.bzl index bd1191931..eda234940 100644 --- a/tests/jaxite/end_to_end/test.bzl +++ b/tests/jaxite/end_to_end/test.bzl @@ -1,53 +1,11 @@ """A macro providing an end-to-end test for jaxite codegen.""" -load("@heir//tools:heir-opt.bzl", "heir_opt") -load("@heir//tools:heir-translate.bzl", "heir_translate") -load("@rules_python//python:py_library.bzl", "py_library") +load("@heir//tools:heir-jaxite.bzl", "fhe_jaxite_lib") load("@rules_python//python:py_test.bzl", "py_test") -def jaxite_end_to_end_test(name, mlir_src, test_src, heir_opt_flags = "", tags = [], deps = [], **kwargs): - """A rule for running generating OpenFHE and running a test on it. - - Args: - name: The name of the py_test target and the generated .cc file basename. - mlir_src: The source mlir file to run through heir-translate. - test_src: The C++ test harness source file. - heir_opt_flags: Flags to pass to heir-opt. - tags: Tags to pass to py_test. - deps: Deps to pass to py_test and py_library. - **kwargs: Keyword arguments to pass to py_library and py_test. - """ - heir_opt_name = name + ".heir_opt" - generated_heir_opt_name = "%s.heir_opt.mlir" % name - py_codegen_target = name + ".heir_translate_py" +def jaxite_end_to_end_test(name, mlir_src, test_src, entry_function_flag = "", tags = [], deps = [], **kwargs): py_lib_target_name = "%s_py_lib" % name - generated_py_filename = "%s_lib.py" % name - - if heir_opt_flags: - heir_opt( - name = heir_opt_name, - src = mlir_src, - pass_flag = heir_opt_flags, - generated_filename = generated_heir_opt_name, - HEIR_YOSYS = True, - data = ["@heir//lib/Transforms/YosysOptimizer/yosys:share_files"], - ) - else: - generated_heir_opt_name = mlir_src - - heir_translate( - name = py_codegen_target, - src = generated_heir_opt_name, - pass_flags = ["--emit-jaxite"], - generated_filename = generated_py_filename, - ) - py_library( - name = py_lib_target_name, - srcs = [":" + generated_py_filename], - deps = deps + ["@heir_pip_deps_jaxite//:pkg"], - tags = tags, - **kwargs - ) + fhe_jaxite_lib(name, mlir_src, entry_function_flag, py_lib_target_name, tags, deps, **kwargs) py_test( name = name, srcs = [test_src], diff --git a/tools/heir-jaxite.bzl b/tools/heir-jaxite.bzl new file mode 100644 index 000000000..0aaeff68f --- /dev/null +++ b/tools/heir-jaxite.bzl @@ -0,0 +1,51 @@ +"""A macro providing an end-to-end library for jaxite codegen.""" + +load("@heir//tools:heir-opt.bzl", "heir_opt") +load("@heir//tools:heir-translate.bzl", "heir_translate") +load("@rules_python//python:py_library.bzl", "py_library") + +def fhe_jaxite_lib(name, mlir_src, entry_function_flag = "", py_lib_target_name = "", tags = [], deps = [], **kwargs): + """A rule for generating Jaxite code. + + Args: + name: The name of the py_test target and the generated .cc file basename. + mlir_src: The source mlir file to run through heir-translate. + entry_function_flag: Flags for entry function. + py_lib_target_name: target_name for the py_library. + tags: Tags to pass to py_test. + deps: Deps to pass to py_test and py_library. + **kwargs: Keyword arguments to pass to py_library and py_test. + """ + heir_opt_name = name + ".heir_opt" + generated_heir_opt_name = "%s.heir_opt.mlir" % name + py_codegen_target = name + ".heir_translate_py" + generated_py_filename = "%s_lib.py" % name + if not py_lib_target_name: + py_lib_target_name = "%s_py_lib" % name + + if entry_function_flag: + heir_opt_pass_flag = "--tosa-to-boolean-jaxite=%s" % entry_function_flag + heir_opt( + name = heir_opt_name, + src = mlir_src, + pass_flag = heir_opt_pass_flag, + generated_filename = generated_heir_opt_name, + HEIR_YOSYS = True, + data = ["@heir//lib/Transforms/YosysOptimizer/yosys:share_files"], + ) + else: + generated_heir_opt_name = mlir_src + + heir_translate( + name = py_codegen_target, + src = generated_heir_opt_name, + pass_flags = ["--emit-jaxite"], + generated_filename = generated_py_filename, + ) + py_library( + name = py_lib_target_name, + srcs = [":" + generated_py_filename], + deps = deps + ["@heir_pip_deps_jaxite//:pkg"], + tags = tags, + **kwargs + )