poc: add tfhe-rs with bazel deps

Signed-off-by: Asra Ali <[email protected]>

Author: asraa

.pre-commit-config.yaml

@@ -39,7 +39,7 @@ repos:
     hooks:
       - id: codespell
         args: ["-L", "crate, fpt"]
-
+        exclude: "^(.*[.]lock)"
 
 # Changes tabs to spaces
 -   repo: https://github.com/Lucas-C/pre-commit-hooks

MODULE.bazel

@@ -113,3 +113,31 @@ git_override(
     commit = "c84a140c93352cdabbfb547c531be34515b12228",
     remote = "https://github.com/google/re2",
 )
+
+bazel_dep(name = "rules_rust", version = "0.60.0")
+
+crate = use_extension("@rules_rust//crate_universe:extensions.bzl", "crate")
+crate.spec(
+    features = ["derive"],
+    package = "clap",
+    version = "4.1.8",
+)
+crate.spec(
+    package = "rayon",
+    version = "1.6.1",
+)
+crate.spec(
+    features = ["derive"],
+    package = "serde",
+    version = "1.0.152",
+)
+crate.spec(
+    features = [
+        "shortint",
+        "x86_64-unix",
+    ],
+    package = "tfhe",
+    version = "0.5.3",
+)
+crate.from_specs()
+use_repo(crate, "crates")

MODULE.bazel.lock

@@ -1,5 +1,6 @@
 # See README.md for setup required to run these tests
 
+load("@heir//tools:heir-tfhe-rs.bzl", "tfhe_rs_lib")
 load("//bazel:lit.bzl", "glob_lit_tests")
 
 package(default_applicable_licenses = ["@heir//:license"])
@@ -31,3 +32,16 @@ glob_lit_tests(
     },
     test_file_exts = ["mlir"],
 )
+
+tfhe_rs_lib(
+    name = "test_add_one",
+    heir_opt_flags = [
+        "--forward-store-to-load",
+        "--cggi-to-tfhe-rust",
+        "--canonicalize",
+        "--cse",
+    ],
+    heir_translate_flags = ["--use-levels"],
+    mlir_src = "test_add_one.mlir",
+    rs_lib_target_name = "test_add_one_rs_lib",
+)

tests/Examples/tfhe_rust/BUILD

@@ -284,6 +284,16 @@ bzl_library(
     ],
 )
 
+bzl_library(
+    name = "heir_tfhe_rs",
+    srcs = ["heir-tfhe-rs.bzl"],
+    visibility = ["//visibility:public"],
+    deps = [
+        ":heir_opt_bzl",
+        ":heir_translate_bzl",
+    ],
+)
+
 bzl_library(
     name = "llc_bzl",
     srcs = ["llc.bzl"],

tools/BUILD

@@ -0,0 +1,57 @@
+"""A macro providing an end-to-end library for tfhe-rs codegen."""
+
+load("@rules_rust//rust:defs.bzl", "rust_library")
+load("@heir//tools:heir-opt.bzl", "heir_opt")
+load("@heir//tools:heir-translate.bzl", "heir_translate")
+
+def tfhe_rs_lib(name, mlir_src, rs_lib_target_name, heir_opt_flags = [], heir_translate_flags = [], data = [], tags = [], deps = [], **kwargs):
+    """A rule for running generating tfhe-rs and running a test on it.
+
+    Args:
+      name: The name of the test target and the generated .rs file basename.
+      mlir_src: The source mlir file to run through heir-translate
+      rs_lib_target_name: The name of the generated rust_library target
+      heir_opt_flags: Flags to pass to heir-opt before heir-translate
+      heir_translate_flags: Flags to pass to heir-translate
+      data: Data dependencies to be passed to heir_opt
+      tags: Tags to pass to rust_library
+      deps: Deps to pass to rust_library
+      **kwargs: Keyword arguments to pass to rust_library.
+    """
+    rs_codegen_target = name + ".heir_translate_rs"
+
+    generated_rs_filename = "%s_lib.rs" % name
+    heir_opt_name = "%s_heir_opt" % name
+    generated_heir_opt_name = "%s_heir_opt.mlir" % name
+    heir_translate_flags = heir_translate_flags + ["--emit-tfhe-rust"]
+
+    if heir_opt_flags:
+        heir_opt(
+            name = heir_opt_name,
+            src = mlir_src,
+            pass_flags = heir_opt_flags,
+            generated_filename = generated_heir_opt_name,
+            data = data,
+        )
+    else:
+        generated_heir_opt_name = mlir_src
+
+    heir_translate(
+        name = rs_codegen_target,
+        src = generated_heir_opt_name,
+        pass_flags = heir_translate_flags,
+        generated_filename = generated_rs_filename,
+    )
+    rust_library(
+        name = rs_lib_target_name,
+        srcs = [":" + generated_rs_filename],
+        deps = deps + [
+            "@crates//:serde",
+            "@crates//:tfhe",
+            "@crates//:rayon",
+            "@crates//:clap"
+        ],
+        tags = tags,
+        data = data,
+        **kwargs
+    )