diff --git a/tests/Dialect/PISA/BUILD b/tests/Dialect/PISA/IR/BUILD similarity index 100% rename from tests/Dialect/PISA/BUILD rename to tests/Dialect/PISA/IR/BUILD diff --git a/tests/Dialect/PISA/ops.mlir b/tests/Dialect/PISA/IR/ops.mlir similarity index 100% rename from tests/Dialect/PISA/ops.mlir rename to tests/Dialect/PISA/IR/ops.mlir diff --git a/tests/Dialect/PISA/end_to_end/BUILD b/tests/Dialect/PISA/end_to_end/BUILD deleted file mode 100644 index b0d0fbc9d..000000000 --- a/tests/Dialect/PISA/end_to_end/BUILD +++ /dev/null @@ -1,14 +0,0 @@ -# See README.md for setup required to run these tests - -load("@heir//tests/Dialect/PISA/end_to_end:test.bzl", "pisa_end_to_end_test") - -package(default_applicable_licenses = ["@heir//:license"]) - -pisa_end_to_end_test( - name = "basic", - generated_lib_header = "basic_inc.h", - heir_opt_flags = ["--mlir-to-openfhe-bgv=ciphertext-degree=8192"], - mlir_src = "basic.mlir", - tags = ["notap"], - test_src = "basic.cpp", -) diff --git a/tests/Dialect/PISA/end_to_end/README.md b/tests/Dialect/PISA/end_to_end/README.md deleted file mode 100644 index ba3dcedfc..000000000 --- a/tests/Dialect/PISA/end_to_end/README.md +++ /dev/null @@ -1,8 +0,0 @@ -# End to End PISA Tests - -These tests use HEIR's ability to generate -[OpenFHE](https://github.com/openfheorg/openfhe-development) code for both or -client-side components such as key generation and encryption/decryption, and to -generate a ground-truth for the sever-side computation. They also produce the -polynomial/pISA code for the server-side computation to allow for testing and -benchmarking. diff --git a/tests/Dialect/PISA/end_to_end/basic.cpp b/tests/Dialect/PISA/end_to_end/basic.cpp deleted file mode 100644 index df4634b00..000000000 --- a/tests/Dialect/PISA/end_to_end/basic.cpp +++ /dev/null @@ -1,44 +0,0 @@ -#include -#include - -#include "gmock/gmock.h" // from @googletest -#include "gtest/gtest.h" // from @googletest -#include "src/pke/include/openfhe.h" // from @openfhe - -// Generated headers (block clang-format from messing up order) -#include "tests/Dialect/PISA/end_to_end/basic_inc.h" - -using namespace lbcrypto; -using ::testing::ElementsAre; - -namespace mlir { -namespace heir { -namespace openfhe { - -TEST(BasicTest, TestInput1) { - // Test Inputs - std::vector x = {1, 2, 3, 4, 5, 6, 7}; - std::vector y = {3, 2, 1, 4, 5, 6, 7}; - - // HEIR-generated func that sets parameters and generates OpenFHE context - auto context = basic_test__generate_crypto_context(); - - auto [publicKey, secretKey] = context->KeyGen(); - - // HEIR-generated funcs for encoding & encrypting of cleartext values - auto xEnc = basic_test__encrypt__arg0(context, x, publicKey); - auto yEnc = basic_test__encrypt__arg1(context, y, publicKey); - - // HEIR-generated func for actual computation (here: simple addition) - auto ciphertextActual = basic_test(context, xEnc, yEnc); - - // HEIR-generated func for - auto r = basic_test__decrypt__result0(context, ciphertextActual, secretKey); - - r.resize(x.size()); - EXPECT_THAT(r, ElementsAre(4, 4, 4, 8, 10, 12, 14)); -} - -} // namespace openfhe -} // namespace heir -} // namespace mlir diff --git a/tests/Dialect/PISA/end_to_end/basic.csv b/tests/Dialect/PISA/end_to_end/basic.csv deleted file mode 100644 index 54191d48b..000000000 --- a/tests/Dialect/PISA/end_to_end/basic.csv +++ /dev/null @@ -1,2 +0,0 @@ -13, add, v4, v0, v2, 0 -13, add, v5, v1, v3, 0 diff --git a/tests/Dialect/PISA/end_to_end/basic.mlir b/tests/Dialect/PISA/end_to_end/basic.mlir deleted file mode 100644 index c92744649..000000000 --- a/tests/Dialect/PISA/end_to_end/basic.mlir +++ /dev/null @@ -1,24 +0,0 @@ -// This is effectively the pipeline run for this test, but it is run by bazel -// and not lit, so the official definition of what command is run can be found -// in the BUILD file for this directory, and the openfhe_end_to_end_test macro -// in test.bzl -// -// heir-opt --mlir-to-openfhe-bgv='entry-function=basic_test ciphertext-degree=8192' %s | bazel-bin/tools/heir-translate --emit-openfhe-pke - -// Below is the pipeline used for pISA: -// heir-opt --mlir-to-bgv='entry-function=basic_test ciphertext-degree=8192' --bgv-to-lwe --lwe-to-polynomial --convert-elementwise-to-affine --full-loop-unroll --convert-tensor-to-scalars polynomial-to-pisa %s | bazel-bin/tools/heir-translate --emit-pisa - - -!t = tensor<8192xi16> - -func.func @basic_test(%x: !t {secret.secret}, %y: !t {secret.secret}) -> !t { - %r0 = tensor.empty() : !t - %r = affine.for %i = 0 to 8192 iter_args(%r = %r0) -> !t { - %xi = tensor.extract %x[%i] : !t - %yi = tensor.extract %y[%i] : !t - %si = arith.addi %xi, %yi : i16 - %ri = tensor.insert %si into %r[%i] : !t - affine.yield %ri : !t - } - return %r : !t -} diff --git a/tests/Dialect/PISA/end_to_end/basic_bgv.mlir b/tests/Dialect/PISA/end_to_end/basic_bgv.mlir deleted file mode 100644 index 3acc3c62b..000000000 --- a/tests/Dialect/PISA/end_to_end/basic_bgv.mlir +++ /dev/null @@ -1,6 +0,0 @@ -module { - func.func @basic_test(%arg0: !lwe.rlwe_ciphertext, rlwe_params = >>, underlying_type = tensor<8192xi16>>, %arg1: !lwe.rlwe_ciphertext, rlwe_params = >>, underlying_type = tensor<8192xi16>>) -> !lwe.rlwe_ciphertext, rlwe_params = >>, underlying_type = tensor<8192xi16>> { - %0 = bgv.add %arg0, %arg1 : !lwe.rlwe_ciphertext, rlwe_params = >>, underlying_type = tensor<8192xi16>> - return %0 : !lwe.rlwe_ciphertext, rlwe_params = >>, underlying_type = tensor<8192xi16>> - } -} diff --git a/tests/Dialect/PISA/end_to_end/basic_lwe.mlir b/tests/Dialect/PISA/end_to_end/basic_lwe.mlir deleted file mode 100644 index 2998158c8..000000000 --- a/tests/Dialect/PISA/end_to_end/basic_lwe.mlir +++ /dev/null @@ -1,6 +0,0 @@ -module { - func.func @basic_test(%arg0: !lwe.rlwe_ciphertext, rlwe_params = >>, underlying_type = tensor<8192xi16>>, %arg1: !lwe.rlwe_ciphertext, rlwe_params = >>, underlying_type = tensor<8192xi16>>) -> !lwe.rlwe_ciphertext, rlwe_params = >>, underlying_type = tensor<8192xi16>> { - %0 = lwe.radd %arg0, %arg1 : !lwe.rlwe_ciphertext, rlwe_params = >>, underlying_type = tensor<8192xi16>> - return %0 : !lwe.rlwe_ciphertext, rlwe_params = >>, underlying_type = tensor<8192xi16>> - } -} diff --git a/tests/Dialect/PISA/end_to_end/basic_openfhe.mlir b/tests/Dialect/PISA/end_to_end/basic_openfhe.mlir deleted file mode 100644 index 14bf42f9a..000000000 --- a/tests/Dialect/PISA/end_to_end/basic_openfhe.mlir +++ /dev/null @@ -1,29 +0,0 @@ -module { - func.func @basic_test(%arg0: !openfhe.crypto_context, %arg1: !lwe.rlwe_ciphertext, rlwe_params = >>, underlying_type = tensor<4096xi16>>, %arg2: !lwe.rlwe_ciphertext, rlwe_params = >>, underlying_type = tensor<4096xi16>>) -> !lwe.rlwe_ciphertext, rlwe_params = >>, underlying_type = tensor<4096xi16>> { - %0 = openfhe.add %arg0, %arg1, %arg2 : (!openfhe.crypto_context, !lwe.rlwe_ciphertext, rlwe_params = >>, underlying_type = tensor<4096xi16>>, !lwe.rlwe_ciphertext, rlwe_params = >>, underlying_type = tensor<4096xi16>>) -> !lwe.rlwe_ciphertext, rlwe_params = >>, underlying_type = tensor<4096xi16>> - return %0 : !lwe.rlwe_ciphertext, rlwe_params = >>, underlying_type = tensor<4096xi16>> - } - func.func @basic_test__encrypt__arg0(%arg0: !openfhe.crypto_context, %arg1: tensor<4096xi16>, %arg2: !openfhe.public_key) -> !lwe.rlwe_ciphertext, rlwe_params = >>, underlying_type = tensor<4096xi16>> { - %0 = lwe.rlwe_encode %arg1 {encoding = #lwe.polynomial_evaluation_encoding, ring = #polynomial.ring>} : tensor<4096xi16> -> !lwe.rlwe_plaintext, ring = >, underlying_type = tensor<4096xi16>> - %1 = openfhe.encrypt %arg0, %0, %arg2 : (!openfhe.crypto_context, !lwe.rlwe_plaintext, ring = >, underlying_type = tensor<4096xi16>>, !openfhe.public_key) -> !lwe.rlwe_ciphertext, rlwe_params = >>, underlying_type = tensor<4096xi16>> - return %1 : !lwe.rlwe_ciphertext, rlwe_params = >>, underlying_type = tensor<4096xi16>> - } - func.func @basic_test__encrypt__arg1(%arg0: !openfhe.crypto_context, %arg1: tensor<4096xi16>, %arg2: !openfhe.public_key) -> !lwe.rlwe_ciphertext, rlwe_params = >>, underlying_type = tensor<4096xi16>> { - %0 = lwe.rlwe_encode %arg1 {encoding = #lwe.polynomial_evaluation_encoding, ring = #polynomial.ring>} : tensor<4096xi16> -> !lwe.rlwe_plaintext, ring = >, underlying_type = tensor<4096xi16>> - %1 = openfhe.encrypt %arg0, %0, %arg2 : (!openfhe.crypto_context, !lwe.rlwe_plaintext, ring = >, underlying_type = tensor<4096xi16>>, !openfhe.public_key) -> !lwe.rlwe_ciphertext, rlwe_params = >>, underlying_type = tensor<4096xi16>> - return %1 : !lwe.rlwe_ciphertext, rlwe_params = >>, underlying_type = tensor<4096xi16>> - } - func.func @basic_test__decrypt__result0(%arg0: !openfhe.crypto_context, %arg1: !lwe.rlwe_ciphertext, rlwe_params = >>, underlying_type = tensor<4096xi16>>, %arg2: !openfhe.private_key) -> tensor<4096xi16> { - %0 = openfhe.decrypt %arg0, %arg1, %arg2 : (!openfhe.crypto_context, !lwe.rlwe_ciphertext, rlwe_params = >>, underlying_type = tensor<4096xi16>>, !openfhe.private_key) -> !lwe.rlwe_plaintext, ring = >, underlying_type = tensor<4096xi16>> - %1 = lwe.rlwe_decode %0 {encoding = #lwe.polynomial_evaluation_encoding, ring = #polynomial.ring>} : !lwe.rlwe_plaintext, ring = >, underlying_type = tensor<4096xi16>> -> tensor<4096xi16> - return %1 : tensor<4096xi16> - } - func.func @basic_test__generate_crypto_context() -> !openfhe.crypto_context { - %0 = openfhe.gen_params {mulDepth = 0 : i64, plainMod = 4295294977 : i64} : () -> !openfhe.cc_params - %1 = openfhe.gen_context %0 : (!openfhe.cc_params) -> !openfhe.crypto_context - return %1 : !openfhe.crypto_context - } - func.func @basic_test__configure_crypto_context(%arg0: !openfhe.crypto_context, %arg1: !openfhe.private_key) -> !openfhe.crypto_context { - return %arg0 : !openfhe.crypto_context - } -} diff --git a/tests/Dialect/PISA/end_to_end/basic_pisa.mlir b/tests/Dialect/PISA/end_to_end/basic_pisa.mlir deleted file mode 100644 index f6392083a..000000000 --- a/tests/Dialect/PISA/end_to_end/basic_pisa.mlir +++ /dev/null @@ -1,7 +0,0 @@ -module { - func.func @basic_test(%arg0: tensor<8192xi32>, %arg1: tensor<8192xi32>, %arg2: tensor<8192xi32>, %arg3: tensor<8192xi32>) -> (tensor<8192xi32>, tensor<8192xi32>) { - %0 = pisa.add %arg0, %arg2 {i = 0 : i32, q = 463187969 : i32} : tensor<8192xi32> - %1 = pisa.add %arg1, %arg3 {i = 0 : i32, q = 463187969 : i32} : tensor<8192xi32> - return %0, %1 : tensor<8192xi32>, tensor<8192xi32> - } -} diff --git a/tests/Dialect/PISA/end_to_end/basic_poly.mlir b/tests/Dialect/PISA/end_to_end/basic_poly.mlir deleted file mode 100644 index 0f7fcacd0..000000000 --- a/tests/Dialect/PISA/end_to_end/basic_poly.mlir +++ /dev/null @@ -1,6 +0,0 @@ -module { - func.func @basic_test(%arg0: tensor<2x!polynomial.polynomial>>>, %arg1: tensor<2x!polynomial.polynomial>>>) -> tensor<2x!polynomial.polynomial>>> { - %0 = polynomial.add %arg0, %arg1 : tensor<2x!polynomial.polynomial>>> - return %0 : tensor<2x!polynomial.polynomial>>> - } -} diff --git a/tests/Dialect/PISA/end_to_end/basic_spoly.mlir b/tests/Dialect/PISA/end_to_end/basic_spoly.mlir deleted file mode 100644 index 852be18d9..000000000 --- a/tests/Dialect/PISA/end_to_end/basic_spoly.mlir +++ /dev/null @@ -1,7 +0,0 @@ -module { - func.func @basic_test(%arg0: !polynomial.polynomial>>, %arg1: !polynomial.polynomial>>, %arg2: !polynomial.polynomial>>, %arg3: !polynomial.polynomial>>) -> (!polynomial.polynomial>>, !polynomial.polynomial>>) { - %0 = polynomial.add %arg0, %arg2 : !polynomial.polynomial>> - %1 = polynomial.add %arg1, %arg3 : !polynomial.polynomial>> - return %0, %1 : !polynomial.polynomial>>, !polynomial.polynomial>> - } -} diff --git a/tests/Dialect/PISA/end_to_end/test.bzl b/tests/Dialect/PISA/end_to_end/test.bzl deleted file mode 100644 index 4f8fb9112..000000000 --- a/tests/Dialect/PISA/end_to_end/test.bzl +++ /dev/null @@ -1,70 +0,0 @@ -"""A macro providing an end-to-end test for OpenFHE codegen.""" - -load("@heir//tools:heir-opt.bzl", "heir_opt") -load("@heir//tools:heir-translate.bzl", "heir_translate") - -def pisa_end_to_end_test(name, mlir_src, test_src, generated_lib_header, heir_opt_flags = [], data = [], tags = [], deps = [], **kwargs): - """A rule for end-to-end tests with OpenFHE and PISA. - - Args: - name: The name of the cc_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. - generated_lib_header: The name of the generated .h file (explicit - because it needs to be manually #include'd in the test_src file) - heir_opt_flags: Flags to pass to heir-opt before heir-translate - data: Data dependencies to be passed to cc_test - tags: Tags to pass to cc_test - deps: Deps to pass to cc_test and cc_library - **kwargs: Keyword arguments to pass to cc_library and cc_test. - """ - cc_codegen_target = name + ".heir_translate_cc" - h_codegen_target = name + ".heir_translate_h" - cc_lib_target_name = "%s_cc_lib" % name - generated_cc_filename = "%s_lib.inc.cc" % name - heir_opt_name = "%s_heir_opt" % name - generated_heir_opt_name = "%s_heir_opt.mlir" % name - - if heir_opt_flags: - heir_opt( - name = heir_opt_name, - src = mlir_src, - pass_flags = heir_opt_flags, - generated_filename = generated_heir_opt_name, - ) - else: - generated_heir_opt_name = mlir_src - - heir_translate( - name = cc_codegen_target, - src = generated_heir_opt_name, - pass_flags = ["--emit-openfhe-pke"], - generated_filename = generated_cc_filename, - ) - heir_translate( - name = h_codegen_target, - src = generated_heir_opt_name, - pass_flags = ["--emit-openfhe-pke-header"], - generated_filename = generated_lib_header, - ) - native.cc_library( - name = cc_lib_target_name, - srcs = [":" + generated_cc_filename], - hdrs = [":" + generated_lib_header], - deps = deps + ["@openfhe//:pke"], - tags = tags, - **kwargs - ) - native.cc_test( - name = name, - srcs = [test_src], - deps = deps + [ - ":" + cc_lib_target_name, - "@openfhe//:pke", - "@openfhe//:core", - "@googletest//:gtest_main", - ], - tags = tags, - data = data, - **kwargs - ) diff --git a/tests/Dialect/Polynomial/Conversions/polynomial_to_pisa/end_to_end.mlir b/tests/Dialect/Polynomial/Conversions/polynomial_to_pisa/end_to_end.mlir new file mode 100644 index 000000000..341e597ba --- /dev/null +++ b/tests/Dialect/Polynomial/Conversions/polynomial_to_pisa/end_to_end.mlir @@ -0,0 +1,16 @@ +// RUN: heir-opt --mlir-to-bgv='ciphertext-degree=8192' --bgv-to-lwe --lwe-to-polynomial --convert-elementwise-to-affine --full-loop-unroll --convert-tensor-to-scalars --polynomial-to-pisa %s + +// FIXME: ADD FILECHECK +!t = tensor<8192xi16> + +func.func @basic_test(%x: !t {secret.secret}, %y: !t {secret.secret}) -> !t { + %r0 = tensor.empty() : !t + %r = affine.for %i = 0 to 8192 iter_args(%r = %r0) -> !t { + %xi = tensor.extract %x[%i] : !t + %yi = tensor.extract %y[%i] : !t + %si = arith.addi %xi, %yi : i16 + %ri = tensor.insert %si into %r[%i] : !t + affine.yield %ri : !t + } + return %r : !t +} diff --git a/tests/Dialect/Polynomial/Conversions/polynomial_to_pisa/non_rns.mlir b/tests/Dialect/Polynomial/Conversions/polynomial_to_pisa/non_rns.mlir index 2cfce8edf..945535da5 100644 --- a/tests/Dialect/Polynomial/Conversions/polynomial_to_pisa/non_rns.mlir +++ b/tests/Dialect/Polynomial/Conversions/polynomial_to_pisa/non_rns.mlir @@ -1,7 +1,7 @@ //RUN: heir-opt --polynomial-to-pisa %s | FileCheck %s -#ring = #polynomial.ring> -!p = !polynomial.polynomial +!coeff_ty = !mod_arith.int<33538049:i32> +!p = !polynomial.polynomial>> //CHECK-LABEL: @test_add //CHECK: [[X:%.+]]: tensor<8192xi32>, [[Y:%.+]]: tensor<8192xi32> diff --git a/tests/Dialect/Polynomial/Conversions/polynomial_to_pisa/polynomial_to_pisa.mlir b/tests/Dialect/Polynomial/Conversions/polynomial_to_pisa/polynomial_to_pisa.mlir new file mode 100644 index 000000000..46236c680 --- /dev/null +++ b/tests/Dialect/Polynomial/Conversions/polynomial_to_pisa/polynomial_to_pisa.mlir @@ -0,0 +1,18 @@ +// RUN: heir-opt --convert-elementwise-to-affine --full-loop-unroll --convert-tensor-to-scalars --polynomial-to-pisa %s + +// FIXME: ADD FILECHECK + +!coeff_ty = !mod_arith.int<33538049:i32> +!p = !polynomial.polynomial>> + +func.func @tensor(%arg0: tensor<2x!poly>, %arg1: tensor<2x!poly>) -> tensor<2x!poly> { + %0 = polynomial.add %arg0, %arg1 : tensor<2x!poly> + return %0 : tensor<2x!poly> +} + + +func.func @scalar(%arg0: !poly, %arg1: !poly, %arg2: !poly, %arg3: !poly) -> (!poly, !poly) { + %0 = polynomial.add %arg0, %arg2 : !poly + %1 = polynomial.add %arg1, %arg3 : !poly + return %0, %1 : !poly, !poly +} diff --git a/tests/Emitter/PISA/BUILD b/tests/Emitter/PISA/BUILD new file mode 100644 index 000000000..c571e6fc6 --- /dev/null +++ b/tests/Emitter/PISA/BUILD @@ -0,0 +1,10 @@ +load("//bazel:lit.bzl", "glob_lit_tests") + +package(default_applicable_licenses = ["@heir//:license"]) + +glob_lit_tests( + name = "all_tests", + data = ["@heir//tests:test_utilities"], + driver = "@heir//tests:run_lit.sh", + test_file_exts = ["mlir"], +) diff --git a/tests/Dialect/PISA/emit.mlir b/tests/Emitter/PISA/emit.mlir similarity index 100% rename from tests/Dialect/PISA/emit.mlir rename to tests/Emitter/PISA/emit.mlir