Add a bazel file for HEIR-Jaxite path

PiperOrigin-RevId: 683692013

lib/Transforms/YosysOptimizer/yosys/BUILD

@@ -11,3 +11,7 @@ filegroup(
         "*.v",
     ]),
 )
+
+exports_files([
+    "*.v",
+])

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",
+])

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

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
+    )