poc: add tfhe-rs with bazel deps

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

fix sbox timeout

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

remove test files

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

address comments

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

remove unneeded deps

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

Author: asraa

.github/workflows/build_and_test.yml

@@ -74,13 +74,3 @@ jobs:
       - name: "Run `bazel test`"
         run: |
           bazel test --noincompatible_strict_action_env --//:enable_openmp=0 -c fastbuild //...
-
-      # Tests specifically for the tfhe-rs codegen
-      - name: rustup toolchain install
-        uses: dtolnay/rust-toolchain@439cf607258077187679211f12aa6f19af4a0af7 # [email protected]
-        with:
-          toolchain: stable
-
-      - name: Test rust codegen targets
-        run: |
-          bash .github/workflows/run_rust_tests.sh

.github/workflows/run_rust_tests.sh

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

.pre-commit-config.yaml

@@ -110,3 +110,33 @@ 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 = [
+        "boolean",
+        "shortint",
+        "integer",
+        "x86_64-unix",
+    ],
+    package = "tfhe",
+    version = "0.5.3",
+)
+crate.from_specs()
+use_repo(crate, "crates")

MODULE.bazel

@@ -4,5 +4,4 @@ These tests exercise OpenFHE codegen for the
 [OpenFHE](https://github.com/openfheorg/openfhe-development) backend library,
 including compiling the generated C++ source and running the resulting binary.
 
-OpenFHE is added as a project-level dependency (unlike the `tfhe-rs` end-to-end
-tests) and built from source.
+OpenFHE is added as a project-level dependency and built from source.

MODULE.bazel.lock

@@ -1,33 +1,96 @@
-# See README.md for setup required to run these tests
-
-load("//bazel:lit.bzl", "glob_lit_tests")
+load("@heir//tests/Examples/tfhe_rust:test.bzl", "tfhe_rs_end_to_end_test")
 
 package(default_applicable_licenses = ["@heir//:license"])
 
-glob_lit_tests(
-    name = "all_tests",
-    data = [
-        "Cargo.toml",
-        "src/main.rs",
-        "src/main_add_one.rs",
-        "src/main_fully_connected.rs",
-        "src/main_multi_output.rs",
-        "src/main_sbox.rs",
-        "@heir//tests:test_utilities",
-    ],
-    default_tags = [
-        "manual",
-        "notap",
-    ],
-    driver = "@heir//tests:run_lit.sh",
-    size_override = {
-        "test_add.mlir": "large",
-        "test_add_one.mlir": "large",
-        "test_bitand.mlir": "large",
-        "test_simple_lut.mlir": "large",
-        "test_fully_connected.mlir": "large",
-        "test_sbox.mlir": "large",
-        "test_multi_output.mlir": "large",
-    },
-    test_file_exts = ["mlir"],
+tfhe_rs_end_to_end_test(
+    name = "add_one_test",
+    size = "large",
+    heir_opt_flags = [
+        "--forward-store-to-load",
+        "--cggi-to-tfhe-rust",
+        "--canonicalize",
+        "--cse",
+    ],
+    heir_translate_flags = [
+        "--emit-tfhe-rust",
+        "--use-levels",
+    ],
+    mlir_src = "add_one.mlir",
+    test_src = "src/main_add_one.rs",
+)
+
+tfhe_rs_end_to_end_test(
+    name = "add_test",
+    size = "large",
+    heir_translate_flags = ["--emit-tfhe-rust"],
+    mlir_src = "add.mlir",
+    test_src = "src/main_add.rs",
+)
+
+tfhe_rs_end_to_end_test(
+    name = "bitand_test",
+    size = "large",
+    heir_translate_flags = [
+        "--emit-tfhe-rust",
+        "--use-levels",
+    ],
+    mlir_src = "bitand.mlir",
+    test_src = "src/main_bitand.rs",
+)
+
+tfhe_rs_end_to_end_test(
+    name = "simple_lut_test",
+    size = "large",
+    heir_translate_flags = ["--emit-tfhe-rust"],
+    mlir_src = "simple_lut.mlir",
+    test_src = "src/main_simple_lut.rs",
+)
+
+tfhe_rs_end_to_end_test(
+    name = "fully_connected_test",
+    size = "large",
+    heir_opt_flags = [
+        "--mlir-to-cggi",
+        "--scheme-to-tfhe-rs",
+    ],
+    heir_translate_flags = [
+        "--emit-tfhe-rust",
+        "--use-levels",
+    ],
+    mlir_src = "fully_connected.mlir",
+    tags = ["yosys"],
+    test_src = "src/main_fully_connected.rs",
+)
+
+tfhe_rs_end_to_end_test(
+    name = "sbox_test",
+    # Skip in debug mode
+    size = "enormous",
+    heir_opt_flags = [
+        "--mlir-to-cggi",
+        "--scheme-to-tfhe-rs",
+    ],
+    heir_translate_flags = [
+        "--emit-tfhe-rust",
+        "--use-levels",
+    ],
+    mlir_src = "sbox.mlir",
+    tags = ["yosys"],
+    test_src = "src/main_sbox.rs",
+)
+
+tfhe_rs_end_to_end_test(
+    name = "multi_output_test",
+    size = "large",
+    heir_opt_flags = [
+        "--mlir-to-cggi",
+        "--scheme-to-tfhe-rs",
+    ],
+    heir_translate_flags = [
+        "--emit-tfhe-rust",
+        "--use-levels",
+    ],
+    mlir_src = "multi_output.mlir",
+    tags = ["yosys"],
+    test_src = "src/main_multi_output.rs",
 )

tests/Examples/openfhe/README.md

@@ -4,37 +4,7 @@ These tests exercise Rust codegen for the
 [tfhe-rs](https://github.com/zama-ai/tfhe-rs) backend library, including
 compiling the generated Rust source and running the resulting binary.
 
-To avoid introducing these large dependencies into the entire project, these
-tests are manual, and require the system they're running on to have
-[Cargo](https://doc.rust-lang.org/cargo/index.html) installed. During the test,
-cargo will fetch and build the required dependencies, and `Cargo.toml` in this
-directory effectively pins the version of `tfhe` supported.
-
-Use the following command to run the tests in this directory, where the default
-Cargo home `$HOME/.cargo` may need to be replaced by your custom `$CARGO_HOME`,
-if you overrode the default option when installing Cargo.
-
-```bash
-bazel query "filter('.mlir.test$', //tests/Examples/tfhe_rust/...)" \
-  | xargs bazel test --noincompatible_strict_action_env --sandbox_writable_path=$HOME/.cargo "$@"
-```
-
-The `manual` tag is added to the targets in this directory to ensure that they
-are not run when someone runs a glob test like `bazel test //...`.
-
-If you don't do this correctly, you will see an error like this:
-
-```
-# .---command stderr------------
-# |     Updating crates.io index
-# |  Downloading crates ...
-# |   Downloaded memoffset v0.9.0
-# | error: failed to download replaced source registry `crates-io`
-# |
-# | Caused by:
-# |   failed to open `/home/you/.cargo/registry/cache/index.crates.io-6f17d22bba15001f/memoffset-0.9.0.crate`
-# |
-# | Caused by:
-# |   Read-only file system (os error 30)
-# `-----------------------------
-```
+The tests are generated using a custom Bazel macro found in
+[test.bzl](tests/Examples/tfhe_rust/test.bzl). `tfhe-rs` is added a
+project-level dependency, and it's version is pinned in
+[MODULE.bazel](MODULE.bazel).

tests/Examples/tfhe_rust/BUILD

@@ -1,13 +1,9 @@
 // This test ensures the testing harness is working properly with minimal codegen.
 
-// RUN: heir-translate %s --emit-tfhe-rust --use-levels > %S/src/fn_under_test.rs
-// RUN: cargo run --release --manifest-path %S/Cargo.toml --bin main -- 2 3 --message_bits=3 | FileCheck %s
-
 !sks = !tfhe_rust.server_key
 !lut = !tfhe_rust.lookup_table
 !eui3 = !tfhe_rust.eui3
 
-// CHECK: 5
 func.func @fn_under_test(%sks : !sks, %a: !eui3, %b: !eui3) -> !eui3 {
   %res = tfhe_rust.add %sks, %a, %b: (!sks, !eui3, !eui3) -> !eui3
   return %res : !eui3

tests/Examples/tfhe_rust/Cargo.toml

@@ -1,11 +1,7 @@
-// RUN: heir-opt --forward-store-to-load --cggi-to-tfhe-rust --canonicalize --cse %s | heir-translate --emit-tfhe-rust > %S/src/fn_under_test.rs
-// RUN: cargo run --release --manifest-path %S/Cargo.toml --bin main_add_one -- 2 --message_bits=3 | FileCheck %s
-
 #encoding = #lwe.unspecified_bit_field_encoding<cleartext_bitwidth = 3>
 !ct_ty = !lwe.lwe_ciphertext<encoding = #encoding>
 !pt_ty = !lwe.lwe_plaintext<encoding = #encoding>
 
-// CHECK: 00000011
 module {
   func.func @fn_under_test(%arg0: memref<8x!ct_ty>) -> memref<8x!ct_ty> {
     %true = arith.constant true

tests/Examples/tfhe_rust/README.md

@@ -1,12 +1,8 @@
 // This test ensures the testing harness is working properly with minimal codegen.
 
-// RUN: heir-translate %s --emit-tfhe-rust > %S/src/fn_under_test.rs
-// RUN: cargo run --release --manifest-path %S/Cargo.toml --bin main -- 2 3 --message_bits=3 | FileCheck %s
-
 !sks = !tfhe_rust.server_key
 !eui3 = !tfhe_rust.eui3
 
-// CHECK: 0
 func.func @fn_under_test(%sks : !sks, %a: !eui3, %b: !eui3) -> !eui3 {
   %res = tfhe_rust.bitand %sks, %a, %b: (!sks, !eui3, !eui3) -> !eui3
   return %res : !eui3

tests/Examples/tfhe_rust/test_add.mlir renamed to tests/Examples/tfhe_rust/add.mlir

@@ -1,8 +1,4 @@
-// RUN: heir-opt --mlir-to-cggi --scheme-to-tfhe-rs %s | heir-translate --emit-tfhe-rust > %S/src/fn_under_test.rs
-// RUN: cargo run --release --manifest-path %S/Cargo.toml --bin main_fully_connected -- 2 --message_bits=3 | FileCheck %s
-
 // This takes takes the input x and outputs 2 \cdot x + 1.
-// CHECK: 00000101
 #map = affine_map<(d0, d1) -> (0)>
 #map1 = affine_map<(d0, d1) -> (d0, d1)>
 module {

tests/Examples/tfhe_rust/test_add_one.mlir renamed to tests/Examples/tfhe_rust/add_one.mlir

@@ -0,0 +1,9 @@
+module {
+  func.func @multi_output(%arg0: i8 {secret.secret}) -> (i8, i8) {
+    %c2_i8 = arith.constant 2 : i8
+    %c3_i8 = arith.constant 3 : i8
+    %0 = arith.addi %arg0, %c3_i8 : i8
+    %1 = arith.muli %arg0, %c2_i8 : i8
+    return %0, %1 : i8, i8
+  }
+}

tests/Examples/tfhe_rust/test_bitand.mlir renamed to tests/Examples/tfhe_rust/bitand.mlir

@@ -0,0 +1,11 @@
+// Single (mini) sbox table lookup for AES encryption.
+
+#map = affine_map<(d0) -> (d0)>
+module {
+  func.func @sub_bytes(%arg0: i8 {secret.secret}) -> i8 {
+    %cst = arith.constant dense<"0x637C777BF26B6FC5"> : tensor<8xi8>
+    %2 = arith.index_cast %arg0 : i8 to index
+    %extracted = tensor.extract %cst[%2] : tensor<8xi8>
+    return %extracted : i8
+  }
+}

tests/Examples/tfhe_rust/test_fully_connected.mlir renamed to tests/Examples/tfhe_rust/fully_connected.mlir

@@ -1,14 +1,9 @@
 // This test ensures the testing harness is working properly with minimal codegen.
 
-// RUN: heir-translate %s --emit-tfhe-rust > %S/src/fn_under_test.rs
-// RUN: cargo run --release --manifest-path %S/Cargo.toml --bin main -- 1 0 --message_bits=3 | FileCheck %s
-
 !sks = !tfhe_rust.server_key
 !lut = !tfhe_rust.lookup_table
 !eui3 = !tfhe_rust.eui3
 
-// We're computing, effectively (0b00000111 >> (1 << 1)) & 1, i.e., 0b111 >> 2
-// CHECK: 1
 func.func @fn_under_test(%sks : !sks, %a: !eui3, %b: !eui3) -> !eui3 {
   %lut = tfhe_rust.generate_lookup_table %sks {truthTable = 7 : ui8} : (!sks) -> !lut
   %0 = tfhe_rust.scalar_left_shift %sks, %a {shiftAmount = 1 : index} : (!sks, !eui3) -> !eui3