fix some inference UT in pir (#67105)

Yuan Xiaolan · web-flow · commit b1657b626f83 · 2024-08-12T19:35:59.000+08:00
diff --git a/test/ir/inference/CMakeLists.txt b/test/ir/inference/CMakeLists.txt
@@ -308,17 +308,87 @@ if(WITH_GPU AND TENSORRT_FOUND)
   endif()
 endif()
 
+if(WIN32 AND WIN_UNITTEST_LEVEL LESS 2)
+  message(STATUS "Skip tests unrelated to CUDA/TRT")
+elseif(WITH_ONEDNN)
+
+  set(PIR_COVERAGE_MKLDNN_TESTS
+      test_mkldnn_conv_affine_channel_fuse_pass
+      test_mkldnn_conv_gelu_fuse_pass
+      test_mkldnn_conv_hard_sigmoid_fuse_pass
+      test_mkldnn_conv_hard_swish_fuse_pass
+      test_mkldnn_conv_mish_fuse_pass
+      test_mkldnn_conv_transpose_bias_fuse_pass
+      test_mkldnn_conv3d_op
+      test_mkldnn_depthwise_conv_pass
+      test_mkldnn_shape_op
+      test_mkldnn_shuffle_channel_op)
+  foreach(target ${PIR_COVERAGE_MKLDNN_TESTS})
+    py_test_modules(${target}_pir MODULES ${target} ENVS FLAGS_enable_pir_api=1)
+    set_tests_properties(${target} PROPERTIES LABELS "RUN_TYPE=INFER")
+    message(STATUS "PIR Copied Test: ${target}_pir in inference test")
+  endforeach()
+
+  set(PIR_COVERAGE_ONEDNN_TESTS
+      test_onednn_batch_norm_act_fuse_pass
+      test_onednn_conv_bias_fuse_pass
+      test_onednn_conv_bn_fuse_pass
+      test_onednn_conv_concat_activation_fuse_pass
+      test_onednn_conv_elementwise_add_fuse_pass
+      test_onednn_elementwise_add_activation_fuse_pass
+      test_onednn_fc_activation_fuse_pass
+      test_onednn_fc_gru_fuse_pass
+      test_onednn_fc_lstm_fuse_pass)
+  foreach(target ${PIR_COVERAGE_ONEDNN_TESTS})
+    py_test_modules(${target}_pir MODULES ${target} ENVS FLAGS_enable_pir_api=1)
+    set_tests_properties(${target} PROPERTIES LABELS "RUN_TYPE=INFER")
+    message(STATUS "PIR Copied Test: ${target}_pir in inference test")
+  endforeach()
+
+  set_tests_properties(test_mkldnn_shape_op_pir PROPERTIES TIMEOUT 300)
+  set_tests_properties(test_mkldnn_shuffle_channel_op_pir PROPERTIES TIMEOUT
+                                                                     300)
+  set_tests_properties(test_onednn_conv_bias_fuse_pass_pir PROPERTIES TIMEOUT
+                                                                      300)
+  set_tests_properties(test_onednn_conv_concat_activation_fuse_pass_pir
+                       PROPERTIES TIMEOUT 300)
+
+endif()
+
 if(WITH_GPU AND TENSORRT_FOUND)
   if(WITH_ONEDNN)
     py_test_modules(test_matmul_scale_fuse_pass_pir MODULES
                     test_matmul_scale_fuse_pass ENVS FLAGS_enable_pir_api=1)
+    py_test_modules(
+      test_multihead_matmul_fuse_pass_v3_pir MODULES
+      test_multihead_matmul_fuse_pass_v3 ENVS FLAGS_enable_pir_api=1)
+    py_test_modules(
+      test_multihead_matmul_roformer_fuse_pass_pir MODULES
+      test_multihead_matmul_roformer_fuse_pass ENVS FLAGS_enable_pir_api=1)
+    py_test_modules(test_matmul_v2_scale_fuse_pass_pir MODULES
+                    test_matmul_v2_scale_fuse_pass ENVS FLAGS_enable_pir_api=1)
+    set_tests_properties(test_multihead_matmul_fuse_pass_v3_pir
+                         PROPERTIES LABELS "RUN_TYPE=INFER")
+    set_tests_properties(test_multihead_matmul_roformer_fuse_pass_pir
+                         PROPERTIES LABELS "RUN_TYPE=INFER")
+    set_tests_properties(test_multihead_matmul_fuse_pass_v3_pir
+                         PROPERTIES TIMEOUT 300)
+    set_tests_properties(test_multihead_matmul_roformer_fuse_pass_pir
+                         PROPERTIES TIMEOUT 300)
+
     if(WIN32)
       set_tests_properties(test_matmul_scale_fuse_pass_pir PROPERTIES TIMEOUT
                                                                       300)
+      set_tests_properties(test_matmul_v2_scale_fuse_pass_pir PROPERTIES TIMEOUT
+                                                                         300)
+
+      if(WIN_UNITTEST_LEVEL EQUAL 2)
+        set_tests_properties(test_onednn_conv_bias_fuse_pass_pir
+                             PROPERTIES TIMEOUT 600)
+        set_tests_properties(test_onednn_conv_concat_activation_fuse_pass_pir
+                             PROPERTIES TIMEOUT 600)
+      endif()
     else()
-      py_test_modules(
-        test_matmul_v2_scale_fuse_pass_pir MODULES
-        test_matmul_v2_scale_fuse_pass ENVS FLAGS_enable_pir_api=1)
       set_tests_properties(test_matmul_scale_fuse_pass_pir PROPERTIES TIMEOUT
                                                                       60)
       set_tests_properties(test_matmul_v2_scale_fuse_pass_pir PROPERTIES TIMEOUT
@@ -327,21 +397,6 @@ if(WITH_GPU AND TENSORRT_FOUND)
     if(WIN32 AND WIN_UNITTEST_LEVEL LESS 2)
       message(STATUS "Skip tests unrelated to CUDA/TRT")
     else()
-      set(PIR_COVERAGE_MKLDNN_TESTS
-          test_mkldnn_conv_affine_channel_fuse_pass
-          test_mkldnn_conv_gelu_fuse_pass
-          test_mkldnn_conv_hard_sigmoid_fuse_pass
-          test_mkldnn_conv_hard_swish_fuse_pass
-          test_mkldnn_conv_mish_fuse_pass
-          test_mkldnn_conv_transpose_bias_fuse_pass
-          test_mkldnn_conv3d_op
-          test_mkldnn_depthwise_conv_pass)
-      foreach(target ${PIR_COVERAGE_MKLDNN_TESTS})
-        py_test_modules(${target}_pir MODULES ${target} ENVS
-                        FLAGS_enable_pir_api=1)
-        set_tests_properties(${target} PROPERTIES LABELS "RUN_TYPE=INFER")
-        message(STATUS "PIR Copied Test: ${target}_pir in inference test")
-      endforeach()
 
       set_tests_properties(test_mkldnn_conv_affine_channel_fuse_pass_pir
                            PROPERTIES TIMEOUT 120)
@@ -358,6 +413,21 @@ if(WITH_GPU AND TENSORRT_FOUND)
       set_tests_properties(test_mkldnn_conv3d_op_pir PROPERTIES TIMEOUT 300)
       set_tests_properties(test_mkldnn_depthwise_conv_pass_pir
                            PROPERTIES TIMEOUT 120)
+
+      set_tests_properties(test_onednn_conv_bn_fuse_pass_pir PROPERTIES TIMEOUT
+                                                                        120)
+      set_tests_properties(test_onednn_conv_elementwise_add_fuse_pass_pir
+                           PROPERTIES TIMEOUT 120)
+      set_tests_properties(test_onednn_batch_norm_act_fuse_pass_pir
+                           PROPERTIES TIMEOUT 100)
+      set_tests_properties(test_onednn_fc_activation_fuse_pass_pir
+                           PROPERTIES TIMEOUT 300)
+      set_tests_properties(test_onednn_elementwise_add_activation_fuse_pass_pir
+                           PROPERTIES TIMEOUT 300)
+      set_tests_properties(test_onednn_fc_gru_fuse_pass_pir PROPERTIES TIMEOUT
+                                                                       300)
+      set_tests_properties(test_onednn_fc_lstm_fuse_pass_pir PROPERTIES TIMEOUT
+                                                                        300)
     endif()
   endif()
 endif()
