Fix stablehlo lit test

Signed-off-by: philass <[email protected]>
onnx · Nov 8, 2023 · f152ff5 · f152ff5
1 parent 99ff3f8
commit f152ff5
Showing 1 changed file with 20 additions and 17 deletions.
diff --git a/test/mlir/conversion/onnx_to_stablehlo/Tensor/Reshape.mlir b/test/mlir/conversion/onnx_to_stablehlo/Tensor/Reshape.mlir
@@ -4,6 +4,8 @@
 /// Test the reshape op inference when constants are present.
 //===----------------------------------------------------------------------===//
 
+// -----
+
 func.func @test_reshape_dynamic(%arg0 : tensor<5x5x1x32xf32>, %arg1 : tensor<4xi64>) -> tensor<*xf32> {
   %0 = "onnx.Reshape"(%arg0, %arg1) : (tensor<5x5x1x32xf32>, tensor<4xi64>) -> tensor<*xf32>
   "func.return"(%0) : (tensor<*xf32>) -> ()
@@ -35,23 +37,24 @@ func.func @test_reshape_dynamic(%arg0 : tensor<5x5x1x32xf32>, %arg1 : tensor<4xi
 // CHECK:           [[VAR_13_:%.+]] = arith.cmpi eq, [[VAR_12_]], [[CST_0_]] : index
 // CHECK:           [[VAR_14_:%.+]] = arith.select [[VAR_13_]], [[CST_1_]], [[VAR_12_]] : index
 // CHECK:           [[VAR_15_:%.+]] = arith.cmpi eq, [[VAR_14_]], [[CST_minus_1_]] : index
-// CHECK:           [[VAR_16_:%.+]] = arith.select [[VAR_15_]], [[CST_1_]], [[VAR_14_]] : index
-// CHECK-DAG:       [[VAR_17_:%.+]] = arith.muli [[VAR_11_]], [[VAR_16_]] : index
-// CHECK-DAG:       [[VAR_18_:%.+]] = shape.get_extent [[VAR_0_]], [[CST_3_]] : tensor<4xindex>, index -> index
-// CHECK:           [[VAR_19_:%.+]] = arith.cmpi eq, [[VAR_18_]], [[CST_0_]] : index
-// CHECK:           [[VAR_20_:%.+]] = arith.select [[VAR_19_]], [[CST_32_]], [[VAR_18_]] : index
-// CHECK:           [[VAR_21_:%.+]] = arith.cmpi eq, [[VAR_20_]], [[CST_minus_1_]] : index
-// CHECK:           [[VAR_22_:%.+]] = arith.select [[VAR_21_]], [[CST_1_]], [[VAR_20_]] : index
-// CHECK:           [[VAR_23_:%.+]] = arith.muli [[VAR_17_]], [[VAR_22_]] : index
-// CHECK:           [[VAR_24_:%.+]] = arith.floordivsi [[CST_800_]], [[VAR_23_]] : index
-// CHECK-DAG:       [[VAR_25_:%.+]] = arith.select [[VAR_4_]], [[VAR_24_]], [[VAR_3_]] : index
-// CHECK-DAG:       [[VAR_26_:%.+]] = arith.select [[VAR_9_]], [[VAR_24_]], [[VAR_8_]] : index
-// CHECK-DAG:       [[VAR_27_:%.+]] = arith.select [[VAR_15_]], [[VAR_24_]], [[VAR_14_]] : index
-// CHECK-DAG:       [[VAR_28_:%.+]] = arith.select [[VAR_21_]], [[VAR_24_]], [[VAR_20_]] : index
-// CHECK:           [[VAR_29_:%.+]] = shape.from_extents [[VAR_25_]], [[VAR_26_]], [[VAR_27_]], [[VAR_28_]] : index, index, index, index
-// CHECK:           [[VAR_30_:%.+]] = shape.to_extent_tensor [[VAR_29_]] : !shape.shape -> tensor<4xindex>
-// CHECK:           [[VAR_31_:%.+]] = stablehlo.dynamic_reshape [[PARAM_0_]], [[VAR_30_]] : (tensor<5x5x1x32xf32>, tensor<4xindex>) -> tensor<?x?x?x?xf32>
-// CHECK:           return [[VAR_31_]] : tensor<?x?x?x?xf32>
+// CHECK:           [[VAR_16_:%.+]] = arith.ori [[VAR_15_]], [[VAR_13_]] : i1
+// CHECK:           [[VAR_17_:%.+]] = arith.select [[VAR_16_]], [[CST_1_]], [[VAR_12_]] : index
+// CHECK-DAG:       [[VAR_18_:%.+]] = arith.muli [[VAR_11_]], [[VAR_17_]] : index
+// CHECK-DAG:       [[VAR_19_:%.+]] = shape.get_extent [[VAR_0_]], [[CST_3_]] : tensor<4xindex>, index -> index
+// CHECK:           [[VAR_20_:%.+]] = arith.cmpi eq, [[VAR_19_]], [[CST_0_]] : index
+// CHECK:           [[VAR_21_:%.+]] = arith.select [[VAR_20_]], [[CST_32_]], [[VAR_19_]] : index
+// CHECK:           [[VAR_22_:%.+]] = arith.cmpi eq, [[VAR_21_]], [[CST_minus_1_]] : index
+// CHECK:           [[VAR_23_:%.+]] = arith.select [[VAR_22_]], [[CST_1_]], [[VAR_21_]] : index
+// CHECK:           [[VAR_24_:%.+]] = arith.muli [[VAR_18_]], [[VAR_23_]] : index
+// CHECK:           [[VAR_25_:%.+]] = arith.floordivsi [[CST_800_]], [[VAR_24_]] : index
+// CHECK-DAG:       [[VAR_26_:%.+]] = arith.select [[VAR_4_]], [[VAR_25_]], [[VAR_3_]] : index
+// CHECK-DAG:       [[VAR_27_:%.+]] = arith.select [[VAR_9_]], [[VAR_25_]], [[VAR_8_]] : index
+// CHECK-DAG:       [[VAR_28_:%.+]] = arith.select [[VAR_15_]], [[VAR_25_]], [[VAR_14_]] : index
+// CHECK-DAG:       [[VAR_29_:%.+]] = arith.select [[VAR_22_]], [[VAR_25_]], [[VAR_21_]] : index
+// CHECK:           [[VAR_30_:%.+]] = shape.from_extents [[VAR_26_]], [[VAR_27_]], [[VAR_28_]], [[VAR_29_]] : index, index, index, index
+// CHECK:           [[VAR_31_:%.+]] = shape.to_extent_tensor [[VAR_30_]] : !shape.shape -> tensor<4xindex>
+// CHECK:           [[VAR_32_:%.+]] = stablehlo.dynamic_reshape [[PARAM_0_]], [[VAR_31_]] : (tensor<5x5x1x32xf32>, tensor<4xindex>) -> tensor<?x?x?x?xf32>
+// CHECK:           return [[VAR_32_]] : tensor<?x?x?x?xf32>
 // CHECK:         }
 
 // -----