From 00142d40234004cdae8714bdbc7df3a43745d72c Mon Sep 17 00:00:00 2001
From: Jeremy Kun <j2kun@users.noreply.github.com>
Date: Fri, 29 Mar 2024 11:33:04 -0700
Subject: [PATCH] port two polynomial examples from HECO, fix lowering to bgv
 test

---
 .../linear_polynomial_64.mlir                 | 20 ++++++++++++++++
 .../quadratic_polynomial.mlir                 | 23 +++++++++++++++++++
 .../secret_to_bgv/hamming_distance_1024.mlir  |  4 +---
 3 files changed, 44 insertions(+), 3 deletions(-)
 create mode 100644 tests/heir_simd_vectorizer/linear_polynomial_64.mlir
 create mode 100644 tests/heir_simd_vectorizer/quadratic_polynomial.mlir

diff --git a/tests/heir_simd_vectorizer/linear_polynomial_64.mlir b/tests/heir_simd_vectorizer/linear_polynomial_64.mlir
new file mode 100644
index 000000000..e88046429
--- /dev/null
+++ b/tests/heir_simd_vectorizer/linear_polynomial_64.mlir
@@ -0,0 +1,20 @@
+// RUN: heir-opt --secretize=entry-function=linear_polynomial --wrap-generic --canonicalize --cse \
+// RUN:   --heir-simd-vectorizer %s | FileCheck %s
+
+// CHECK-LABEL: @linear_polynomial
+// CHECK: secret.generic
+// CHECK-NOT: tensor_ext.rotate
+func.func @linear_polynomial(%a: tensor<64xi16>, %b: tensor<64xi16>, %x: tensor<64xi16>, %y: tensor<64xi16>) -> tensor<64xi16> {
+  %0 = affine.for %i = 0 to 64 iter_args(%iter = %y) -> (tensor<64xi16>) {
+    %ai = tensor.extract %a[%i] : tensor<64xi16>
+    %bi = tensor.extract %b[%i] : tensor<64xi16>
+    %xi = tensor.extract %x[%i] : tensor<64xi16>
+    %yi = tensor.extract %y[%i] : tensor<64xi16>
+    %axi = arith.muli %ai, %xi : i16
+    %t1 = arith.subi %yi, %axi : i16
+    %t2 = arith.subi %t1, %bi : i16
+    %out = tensor.insert %t2 into %iter[%i] : tensor<64xi16>
+    affine.yield %out : tensor<64xi16>
+  }
+  return %0 : tensor<64xi16>
+}
diff --git a/tests/heir_simd_vectorizer/quadratic_polynomial.mlir b/tests/heir_simd_vectorizer/quadratic_polynomial.mlir
new file mode 100644
index 000000000..0efc41094
--- /dev/null
+++ b/tests/heir_simd_vectorizer/quadratic_polynomial.mlir
@@ -0,0 +1,23 @@
+// RUN: heir-opt --secretize=entry-function=quadratic_polynomial --wrap-generic --canonicalize --cse \
+// RUN:   --heir-simd-vectorizer %s | FileCheck %s
+
+// CHECK-LABEL: @quadratic_polynomial
+// CHECK: secret.generic
+// CHECK-NOT: tensor_ext.rotate
+func.func @quadratic_polynomial(%a: tensor<64xi16>, %b: tensor<64xi16>, %c: tensor<64xi16>, %x: tensor<64xi16>, %y: tensor<64xi16>) -> tensor<64xi16> {
+  %0 = affine.for %i = 0 to 64 iter_args(%iter = %y) -> (tensor<64xi16>) {
+    %ai = tensor.extract %a[%i] : tensor<64xi16>
+    %bi = tensor.extract %b[%i] : tensor<64xi16>
+    %ci = tensor.extract %c[%i] : tensor<64xi16>
+    %xi = tensor.extract %x[%i] : tensor<64xi16>
+    %yi = tensor.extract %y[%i] : tensor<64xi16>
+    %axi = arith.muli %ai, %xi : i16
+    %t1 = arith.addi %axi, %bi : i16
+    %t2 = arith.muli %xi, %t1 : i16
+    %t3 = arith.addi %t2, %ci : i16
+    %t4 = arith.subi %yi, %t3 : i16
+    %out = tensor.insert %t4 into %iter[%i] : tensor<64xi16>
+    affine.yield %out : tensor<64xi16>
+  }
+  return %0 : tensor<64xi16>
+}
diff --git a/tests/secret_to_bgv/hamming_distance_1024.mlir b/tests/secret_to_bgv/hamming_distance_1024.mlir
index 37a611b15..392a1053b 100644
--- a/tests/secret_to_bgv/hamming_distance_1024.mlir
+++ b/tests/secret_to_bgv/hamming_distance_1024.mlir
@@ -7,9 +7,7 @@
 // CHECK: bgv.sub
 // CHECK-NEXT: bgv.mul
 // CHECK-NEXT: bgv.relinearize
-
-// TODO(#521): After rotate-and-reduce works, only check for 10 bg.rotate
-// CHECK-COUNT-1023: bgv.rotate
+// CHECK-COUNT-10: bgv.rotate
 // CHECK: bgv.extract
 // CHECK-NEXT: return