add a more precise test for box_blur 64

google · Mar 8, 2024 · dabcbfd · dabcbfd
1 parent 0adf17d
commit dabcbfd
Showing 1 changed file with 27 additions and 4 deletions.
diff --git 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]]_i32
+  // 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