From 3a760c04f3fb14cd1eb041b321f71626b3619807 Mon Sep 17 00:00:00 2001
From: Jeremy Kun <j2kun@users.noreply.github.com>
Date: Fri, 8 Mar 2024 07:23:32 -0800
Subject: [PATCH] add a more precise test for box_blur 64

---
 tests/simd/box_blur_64x64.mlir | 31 +++++++++++++++++++++++++++----
 1 file changed, 27 insertions(+), 4 deletions(-)

diff --git a/tests/simd/box_blur_64x64.mlir b/tests/simd/box_blur_64x64.mlir
index 96e63b130..21fb28d95 100644
--- a/tests/simd/box_blur_64x64.mlir
+++ b/tests/simd/box_blur_64x64.mlir
@@ -1,12 +1,35 @@
-// RUN: heir-opt --secretize=entry-function=box_blur --wrap-generic --canonicalize --cse \
-// RUN:   --full-loop-unroll \
+// RUN: heir-opt --secretize=entry-function=box_blur --wrap-generic --canonicalize --cse --full-loop-unroll \
 // RUN:   --insert-rotate --cse --canonicalize --collapse-insertion-chains --canonicalize --cse \
 // RUN:   %s | FileCheck %s
 
 module  {
   // CHECK-LABEL: @box_blur
-  // CHECK-NOT: tensor.extract
-  // CHECK-COUNT-7: tensor_ext.rotate
+  // CHECK-SAME: %[[arg0:.*]]: !secret.secret<tensor<4096xi16>>) -> !secret.secret<tensor<4096xi16>> {
+  // CHECK-NEXT:   %[[c65:.*]] = arith.constant 65 : i32
+  // CHECK-NEXT:   %[[c127:.*]] = arith.constant 127 : index
+  // CHECK-NEXT:   %[[c4032:.*]] = arith.constant 4032 : index
+  // CHECK-NEXT:   %[[c3968:.*]] = arith.constant 3968 : index
+  // CHECK-NEXT:   %[[c63:.*]] = arith.constant 63 : index
+  // CHECK-NEXT:   %[[v0:.*]] = secret.generic ins(%[[arg0]] : !secret.secret<tensor<4096xi16>>) {
+  // CHECK-NEXT:   ^bb0(%[[arg1:.*]]: tensor<4096xi16>):
+  // CHECK-NEXT:     %[[v1:.*]] = tensor_ext.rotate %[[arg1]], %[[c3968]]
+  // CHECK-NEXT:     %[[v2:.*]] = tensor_ext.rotate %[[arg1]], %[[c4032]]
+  // CHECK-NEXT:     %[[v3:.*]] = arith.addi %[[v1]], %[[v2]]
+  // CHECK-NEXT:     %[[v4:.*]] = arith.addi %[[v3]], %arg1
+  // CHECK-NEXT:     %[[v5:.*]] = tensor_ext.rotate %[[v4]], %[[c63]]
+  // CHECK-NEXT:     %[[v6:.*]] = arith.addi %[[v5]], %[[v2]]
+  // CHECK-NEXT:     %[[v7:.*]] = arith.addi %[[v6]], %arg1
+  // CHECK-NEXT:     %[[v8:.*]] = tensor_ext.rotate %[[v7]], %[[c63]]
+  // CHECK-NEXT:     %[[v9:.*]] = tensor_ext.rotate %arg1, %[[c127]]
+  // CHECK-NEXT:     %[[v10:.*]] = arith.addi %[[v8]], %[[v9]]
+  // CHECK-NEXT:     %[[v11:.*]] = arith.addi %[[v10]], %arg1
+  // CHECK-NEXT:     %[[v12:.*]] = tensor_ext.rotate %[[v11]], %[[c3968]]
+  // CHECK-NEXT:     %[[v13:.*]] = arith.addi %[[v12]], %[[v2]]
+  // CHECK-NEXT:     %[[v14:.*]] = arith.addi %[[v13]], %arg1
+  // CHECK-NEXT:     %[[v15:.*]] = tensor_ext.rotate %[[v14]], %[[c65]]
+  // CHECK-NEXT:     secret.yield %[[v15]]
+  // CHECK-NEXT:   } -> !secret.secret<tensor<4096xi16>>
+  // CHECK-NEXT:   return %[[v0]]
   func.func @box_blur(%arg0: tensor<4096xi16> {secret.secret}) -> tensor<4096xi16> {
     %c4096 = arith.constant 4096 : index
     %c64 = arith.constant 64 : index