fix(wopbs): Fixing woppbs with tensor operators (close #789)

- Missing offset in woppbs routine
- Better error message for check of tensor result in end to end fixture
- Modify fixture generator for testing purpose

compiler/.gitignore

@@ -12,3 +12,5 @@ concrete-compiler_compilation_artifacts/
 py_test_lib_compile_and_run_custom_perror/
 tests/end_to_end_fixture/end_to_end_linalg_apply_lookup_table.yaml
 tests/end_to_end_fixture/end_to_end_linalg_leveled.yaml
+tests/end_to_end_fixture/end_to_end_linalg_2_apply_lookup_table.yaml
+tests/end_to_end_fixture/bug_report.yaml

compiler/Makefile

@@ -243,7 +243,13 @@ run-end-to-end-tests: build-end-to-end-tests
 build-end-to-end-jit-test: build-initialized
 	cmake --build $(BUILD_DIR) --target end_to_end_jit_test
 
-build-end-to-end-jit-fhe: build-initialized
+generate-end-to-end-tests:
+	$(Python3_EXECUTABLE) ./tests/end_to_end_fixture/end_to_end_linalg_apply_lookup_table_gen.py \
+	--n-lut 2 --n-ct 4 \
+	&> ./tests/end_to_end_fixture/end_to_end_linalg_2_apply_lookup_table.yaml
+	unzip -o tests/end_to_end_fixture/bug_report.zip -d tests/end_to_end_fixture/
+
+build-end-to-end-jit-fhe: build-initialized generate-end-to-end-tests
 	cmake --build $(BUILD_DIR) --target end_to_end_jit_fhe
 
 build-end-to-end-jit-encrypted-tensor: build-initialized

compiler/lib/Runtime/wrappers.cpp

@@ -507,7 +507,7 @@ void memref_wop_pbs_crt_buffer(
     auto nb_bits_to_extract = number_of_bits_per_block[i];
 
     auto delta_log = 64 - nb_bits_to_extract;
-    auto in_block = &in_aligned[lwe_big_size * i];
+    auto in_block = &in_aligned[lwe_big_size * i + in_offset];
 
     // trick ( ct - delta/2 + delta/2^4  )
     uint64_t sub = (uint64_t(1) << (uint64_t(64) - nb_bits_to_extract - 1)) -

compiler/tests/end_to_end_fixture/EndToEndFixture.cpp

@@ -111,14 +111,19 @@ llvm::Error checkResult(const mlir::concretelang::TensorLambdaArgument<
   if (!expectedNumElts)
     return expectedNumElts.takeError();
 
+  auto hasError = false;
+  StreamStringError err("result value differ");
   for (size_t i = 0; i < *expectedNumElts; i++) {
+
     if (resValues[i] != expectedValues[i]) {
-      return StreamStringError("result value differ at pos(")
-             << i << "), got " << resValues[i] << " expected "
-             << expectedValues[i];
+      hasError = true;
+      err << " [pos(" << i << "), got " << resValues[i] << " expected "
+          << expectedValues[i] << "]";
     }
   }
-
+  if (hasError) {
+    return err;
+  }
   return llvm::Error::success();
 }
 

compiler/tests/end_to_end_fixture/end_to_end_linalg_apply_lookup_table_gen.py

@@ -1,17 +1,14 @@
 from platform import mac_ver
 import numpy as np
-
-MIN_PRECISON = 1
-MAX_PRECISION = 16
-N_CT = [1, 64, 128, 1024]
+import argparse
 
 
-def main():
+def generate(args):
     print("# /!\ DO NOT EDIT MANUALLY THIS FILE MANUALLY")
     print("# /!\ THIS FILE HAS BEEN GENERATED THANKS THE end_to_end_levelled_gen.py scripts")
     np.random.seed(0)
-    for n_ct in N_CT:
-        for p in range(MIN_PRECISON, MAX_PRECISION+1):
+    for n_ct in args.n_ct:
+        for p in range(args.min_bitwidth, args.max_bitwidth+1):
             max_value = (2 ** p) - 1
             random_lut = np.random.randint(max_value+1, size=2**p)
             # identity_apply_lookup_table
@@ -19,12 +16,14 @@ def main():
                 "description: apply_lookup_table_{0}bits_{1}ct".format(p, n_ct))
             print("program: |")
             print(
-                "  func.func @main(%arg0: tensor<{1}x!FHE.eint<{0}>>) -> tensor<{1}x!FHE.eint<{0}>> {{".format(p, n_ct))
+                "  func.func @main(%0: tensor<{1}x!FHE.eint<{0}>>) -> tensor<{1}x!FHE.eint<{0}>> {{".format(p, n_ct))
             print("    %tlu = arith.constant dense<[{0}]> : tensor<{1}xi64>".format(
                 ','.join(map(str, random_lut)), 2**p))
-            print(
-                "    %1 = \"FHELinalg.apply_lookup_table\"(%arg0, %tlu): (tensor<{2}x!FHE.eint<{0}>>, tensor<{1}xi64>) -> (tensor<{2}x!FHE.eint<{0}>>)".format(p, 2**p, n_ct))
-            print("    return %1: tensor<{1}x!FHE.eint<{0}>>".format(p, n_ct))
+            for i in range(0, args.n_lut):
+                print(
+                    "    %{4} = \"FHELinalg.apply_lookup_table\"(%{3}, %tlu): (tensor<{2}x!FHE.eint<{0}>>, tensor<{1}xi64>) -> (tensor<{2}x!FHE.eint<{0}>>)".format(p, 2**p, n_ct, i, i+1))
+            print("    return %{2}: tensor<{1}x!FHE.eint<{0}>>".format(
+                p, n_ct, args.n_lut))
             print("  }")
             random_input = np.random.randint(max_value+1, size=n_ct)
             print("tests:")
@@ -32,11 +31,36 @@ def main():
             print(
                 "    - tensor: [{0}]".format(','.join(map(str, random_input))))
             print("      shape: [{0}]".format(n_ct))
-            outputs = np.vectorize(lambda i: random_lut[i])(random_input)
+            outputs = random_input
+            for i in range(0, args.n_lut):
+                outputs = [random_lut[v] for v in outputs]
+
             print("    outputs:")
             print("    - tensor: [{0}]".format(','.join(map(str, outputs))))
             print("      shape: [{0}]".format(n_ct))
             print("---")
 
 
-main()
+CLI = argparse.ArgumentParser()
+CLI.add_argument(
+    "--min-bitwidth",
+    type=int,
+    default=1,
+)
+CLI.add_argument(
+    "--max-bitwidth",
+    type=int,
+    default=16,
+)
+CLI.add_argument(
+    "--n-ct",
+    nargs="+",
+    type=int,
+    default=[1, 64, 128, 1024],
+)
+CLI.add_argument(
+    "--n-lut",
+    type=int,
+    default=1,
+)
+generate(CLI.parse_args())

compiler/tests/end_to_end_tests/end_to_end_jit_fhe.cc

@@ -220,7 +220,13 @@ testParam(std::vector<EndToEndDesc> descs,
       "tests/end_to_end_fixture/end_to_end_leveled.yaml")                      \
   INSTANTIATE_END_TO_END_TEST_SUITE_FROM_FILE(                                 \
       FHEApplyLookupTable, suite, options, lambdasupport,                      \
-      "tests/end_to_end_fixture/end_to_end_apply_lookup_table.yaml")
+      "tests/end_to_end_fixture/end_to_end_apply_lookup_table.yaml")           \
+  INSTANTIATE_END_TO_END_TEST_SUITE_FROM_FILE(                                 \
+      FHELinalgLookupTable, suite, options, lambdasupport,                     \
+      "tests/end_to_end_fixture/end_to_end_linalg_2_apply_lookup_table.yaml")  \
+  INSTANTIATE_END_TO_END_TEST_SUITE_FROM_FILE(                                 \
+      BugReport000, suite, options, lambdasupport,                             \
+      "tests/end_to_end_fixture/bug_report.yaml")
 
 mlir::concretelang::CompilationOptions defaultOptions() {
   mlir::concretelang::CompilationOptions o("main");