From d8aa9ff76bebe9f0109e97366e9eac413a653d62 Mon Sep 17 00:00:00 2001
From: youben11 <ayoub.benaissa@zama.ai>
Date: Tue, 22 Mar 2022 11:20:55 +0100
Subject: [PATCH] tests: test execution of transpose

---
 .../tests/unittest/end_to_end_fhelinalg.yaml  | 29 +++++++++++++++++++
 compiler/tests/unittest/end_to_end_jit_fhe.cc |  7 +++--
 2 files changed, 34 insertions(+), 2 deletions(-)
 create mode 100644 compiler/tests/unittest/end_to_end_fhelinalg.yaml

diff --git a/compiler/tests/unittest/end_to_end_fhelinalg.yaml b/compiler/tests/unittest/end_to_end_fhelinalg.yaml
new file mode 100644
index 000000000..0ddacce9b
--- /dev/null
+++ b/compiler/tests/unittest/end_to_end_fhelinalg.yaml
@@ -0,0 +1,29 @@
+description: transpose1d
+program: |
+  func @main(%input: tensor<3x!FHE.eint<6>>) -> tensor<3x!FHE.eint<6>> {
+    %1 = "FHELinalg.transpose"(%input): (tensor<3x!FHE.eint<6>>) ->
+    tensor<3x!FHE.eint<6>> return %1 : tensor<3x!FHE.eint<6>>
+  }
+tests:
+  - inputs:
+    - tensor: [1, 2, 3]
+      shape: [3]
+      width: 8
+    outputs:
+    - tensor: [1, 2, 3]
+      shape: [3]
+---
+description: transpose2d
+program: |
+  func @main(%input: tensor<3x2x!FHE.eint<6>>) -> tensor<2x3x!FHE.eint<6>> {
+    %1 = "FHELinalg.transpose"(%input): (tensor<3x2x!FHE.eint<6>>) -> tensor<2x3x!FHE.eint<6>>
+    return %1 : tensor<2x3x!FHE.eint<6>>
+  }
+tests:
+  - inputs:
+    - tensor: [1, 2, 3, 4, 5, 6]
+      shape: [3, 2]
+      width: 8
+    outputs:
+    - tensor: [1, 3, 5, 2, 4, 6]
+      shape: [2, 3]
diff --git a/compiler/tests/unittest/end_to_end_jit_fhe.cc b/compiler/tests/unittest/end_to_end_jit_fhe.cc
index 17bbe2eb9..e6bff4499 100644
--- a/compiler/tests/unittest/end_to_end_jit_fhe.cc
+++ b/compiler/tests/unittest/end_to_end_jit_fhe.cc
@@ -84,7 +84,10 @@ void compile_and_run(EndToEndDesc desc, LambdaSupport support) {
       FHE, suite, lambdasupport, "tests/unittest/end_to_end_fhe.yaml")         \
   INSTANTIATE_END_TO_END_TEST_SUITE_FROM_FILE(                                 \
       EncryptedTensor, suite, lambdasupport,                                   \
-      "tests/unittest/end_to_end_encrypted_tensor.yaml")
+      "tests/unittest/end_to_end_encrypted_tensor.yaml")                       \
+  INSTANTIATE_END_TO_END_TEST_SUITE_FROM_FILE(                                 \
+      FHELinalg, suite, lambdasupport,                                         \
+      "tests/unittest/end_to_end_fhelinalg.yaml")
 
 /// Instantiate the test suite for Jit
 INSTANTIATE_END_TO_END_TEST_SUITE_FROM_ALL_TEST_FILES(
@@ -94,4 +97,4 @@ INSTANTIATE_END_TO_END_TEST_SUITE_FROM_ALL_TEST_FILES(
 INSTANTIATE_END_TO_END_TEST_SUITE_FROM_ALL_TEST_FILES(
     LibraryTest,
     mlir::concretelang::LibraryLambdaSupport("/tmp/end_to_end_test_" +
-                                             desc.description))
\ No newline at end of file
+                                             desc.description))