file(
  GLOB TEST_INFERENCE_IR_PASSES
  RELATIVE "${CMAKE_CURRENT_SOURCE_DIR}"
  "test_*.py")
string(REPLACE ".py" "" TEST_INFERENCE_IR_PASSES "${TEST_INFERENCE_IR_PASSES}")

file(
  GLOB TEST_TRT_IR_PASSES
  RELATIVE "${CMAKE_CURRENT_SOURCE_DIR}"
  "test_trt_*.py")
string(REPLACE ".py" "" TEST_TRT_IR_PASSES "${TEST_TRT_IR_PASSES}")

file(
  GLOB TEST_TRT_CONVERTER
  RELATIVE "${CMAKE_CURRENT_SOURCE_DIR}"
  "test_trt_convert_*.py")
string(REPLACE ".py" "" TEST_TRT_CONVERTER "${TEST_TRT_CONVERTER}")

if(NOT WITH_DISTRIBUTE)
  list(REMOVE_ITEM TEST_INFERENCE_IR_PASSES "test_delete_c_identity_op_pass")
  list(REMOVE_ITEM TEST_INFERENCE_IR_PASSES "test_trt_convert_c_allreduce")
  list(REMOVE_ITEM TEST_TRT_IR_PASSES "test_trt_convert_c_allreduce")
  list(REMOVE_ITEM TEST_TRT_CONVERTER "test_trt_convert_c_allreduce")
endif()

if(WIN32)
  list(REMOVE_ITEM TEST_INFERENCE_IR_PASSES
       "test_trt_convert_fused_token_prune")
  list(REMOVE_ITEM TEST_INFERENCE_IR_PASSES
       "test_preln_groupnorm_act_fuse_pass")
  list(REMOVE_ITEM TEST_INFERENCE_IR_PASSES
       "test_element_groupnorm_act_fuse_pass")
  list(REMOVE_ITEM TEST_INFERENCE_IR_PASSES "test_groupnorm_act_pass_fuse_pass")
  list(REMOVE_ITEM TEST_TRT_IR_PASSES "test_trt_convert_trans_layernorm")
  list(REMOVE_ITEM TEST_TRT_CONVERTER "test_trt_convert_trans_layernorm")
  list(REMOVE_ITEM TEST_INFERENCE_IR_PASSES "test_trt_convert_trans_layernorm")
  list(REMOVE_ITEM TEST_TRT_IR_PASSES "test_trt_convert_fused_token_prune")
  list(REMOVE_ITEM TEST_TRT_CONVERTER "test_trt_convert_fused_token_prune")
  list(REMOVE_ITEM TEST_TRT_IR_PASSES
       "test_trt_convert_elementwiseadd_transpose")
  list(REMOVE_ITEM TEST_INFERENCE_IR_PASSES
       "test_trt_convert_elementwiseadd_transpose")
  list(REMOVE_ITEM TEST_TRT_CONVERTER
       "test_trt_convert_elementwiseadd_transpose")
  list(REMOVE_ITEM TEST_INFERENCE_IR_PASSES
       "test_split_layernorm_to_math_ops_pass")
endif()

# Only for cpu(mkl + openblas)
set(TEST_INFERENCE_CPU_UT "test_mul_lstm_fuse_pass" "test_mul_gru_fuse_pass")

foreach(CPU_UT ${TEST_INFERENCE_CPU_UT})
  list(REMOVE_ITEM TEST_INFERENCE_IR_PASSES ${CPU_UT})
endforeach()

foreach(TEST_INFERENCE_IR_PASS ${TEST_TRT_IR_PASSES})
  list(REMOVE_ITEM TEST_INFERENCE_IR_PASSES ${TEST_INFERENCE_IR_PASS})
endforeach()

if(WITH_GPU AND TENSORRT_FOUND)
  list(REMOVE_ITEM TEST_TRT_IR_PASSES test_trt_multiclass_nms_op)

  foreach(TRT_CONVERT ${TEST_TRT_CONVERTER})
    list(REMOVE_ITEM TEST_TRT_IR_PASSES ${TRT_CONVERT})
  endforeach()

  foreach(target ${TEST_TRT_IR_PASSES})
    if(${target} STREQUAL "test_trt_slice_dynamic_plugin")
      if("${TENSORRT_MAJOR_VERSION}.${TENSORRT_MINOR_VERSION}" VERSION_GREATER
         "7.1")
        py_test_modules(${target} MODULES ${target})
        set_tests_properties(${target} PROPERTIES TIMEOUT 60)
        set_tests_properties(${target} PROPERTIES LABELS "RUN_TYPE=INFER")
      endif()
    else()
      py_test_modules(${target} MODULES ${target})
      set_tests_properties(${target} PROPERTIES LABELS "RUN_TYPE=INFER")
    endif()
  endforeach()

  foreach(target ${TEST_TRT_CONVERTER})
    py_test_modules(${target} MODULES ${target})
    set_tests_properties(${target} PROPERTIES TIMEOUT 300)
    set_tests_properties(${target} PROPERTIES LABELS "RUN_TYPE=INFER")
  endforeach()
endif()

file(
  GLOB TEST_MKLDNN_IR_PASSES
  RELATIVE "${CMAKE_CURRENT_SOURCE_DIR}"
  "test_mkldnn_*.py")
string(REPLACE ".py" "" TEST_MKLDNN_IR_PASSES "${TEST_MKLDNN_IR_PASSES}")
foreach(TEST_INFERENCE_IR_PASS ${TEST_MKLDNN_IR_PASSES})
  list(REMOVE_ITEM TEST_INFERENCE_IR_PASSES ${TEST_INFERENCE_IR_PASS})
endforeach()

file(
  GLOB TEST_ONEDNN_IR_PASSES
  RELATIVE "${CMAKE_CURRENT_SOURCE_DIR}"
  "test_onednn_*.py")
string(REPLACE ".py" "" TEST_ONEDNN_IR_PASSES "${TEST_ONEDNN_IR_PASSES}")
foreach(TEST_INFERENCE_IR_PASS ${TEST_ONEDNN_IR_PASSES})
  list(REMOVE_ITEM TEST_INFERENCE_IR_PASSES ${TEST_INFERENCE_IR_PASS})
endforeach()

if(WITH_MKLDNN)
  foreach(target ${TEST_MKLDNN_IR_PASSES})
    py_test_modules(${target} MODULES ${target})
    set_tests_properties(${target} PROPERTIES LABELS "RUN_TYPE=INFER")
  endforeach()

  foreach(target ${TEST_ONEDNN_IR_PASSES})
    py_test_modules(${target} MODULES ${target})
    set_tests_properties(${target} PROPERTIES LABELS "RUN_TYPE=INFER")
  endforeach()
endif()

file(
  GLOB TEST_XPU_IR_PASSES
  RELATIVE "${CMAKE_CURRENT_SOURCE_DIR}"
  "test_xpu_*.py")
string(REPLACE ".py" "" TEST_XPU_IR_PASSES "${TEST_XPU_IR_PASSES}")
foreach(TEST_XPU_IR_PASS ${TEST_XPU_IR_PASSES})
  list(REMOVE_ITEM TEST_INFERENCE_IR_PASSES ${TEST_XPU_IR_PASS})
endforeach()

if(WITH_XPU)
  foreach(target ${TEST_XPU_IR_PASSES})
    py_test_modules(${target} MODULES ${target})
    set_tests_properties(${target} PROPERTIES LABELS "RUN_TYPE=INFER")
  endforeach()
endif()

# below are cutlass unitests
file(
  GLOB TEST_CUTLASS
  RELATIVE "${CMAKE_CURRENT_SOURCE_DIR}"
  "test_cutlass_*.py")
string(REPLACE ".py" "" TEST_CUTLASS "${TEST_CUTLASS}")
list(REMOVE_ITEM TEST_INFERENCE_IR_PASSES ${TEST_CUTLASS})
if(WITH_CUTLASS)
  foreach(target ${TEST_CUTLASS})
    py_test_modules(${target} MODULES ${target})
  endforeach()
endif()

if(WITH_MKLDNN
   AND TENSORRT_FOUND
   AND WITH_GPU)
  foreach(target ${TEST_INFERENCE_IR_PASSES})
    py_test_modules(${target} MODULES ${target})
    set_tests_properties(${target} PROPERTIES LABELS "RUN_TYPE=INFER")
  endforeach()
endif()

if(NOT WITH_MKLDNN
   AND NOT TENSORRT_FOUND
   AND NOT WITH_GPU)
  foreach(target ${TEST_INFERENCE_CPU_UT})
    py_test_modules(${target} MODULES ${target})
    set_tests_properties(${target} PROPERTIES LABELS "RUN_TYPE=INFER")
  endforeach()

  set_tests_properties(test_mul_lstm_fuse_pass PROPERTIES TIMEOUT 300)
  set_tests_properties(test_mul_gru_fuse_pass PROPERTIES TIMEOUT 300)
endif()

if(WITH_GPU AND TENSORRT_FOUND)
  set_tests_properties(test_trt_subgraph_pass PROPERTIES TIMEOUT 120)
  # longer timeout for trt_activation_pass for longer trt optimization time in trt 8
  set_tests_properties(test_trt_conv_pass PROPERTIES TIMEOUT 120)
  #set_tests_properties(test_trt_multiclass_nms_op PROPERTIES TIMEOUT 200)
  set_tests_properties(test_trt_dynamic_shape PROPERTIES TIMEOUT 120)
  set_tests_properties(test_trt_inspector PROPERTIES TIMEOUT 60)
  set_tests_properties(test_trt_inference_predictor PROPERTIES TIMEOUT 60)
  set_tests_properties(test_trt_inference_fp16_io PROPERTIES TIMEOUT 300)
  set_tests_properties(test_save_optimized_model_pass PROPERTIES TIMEOUT 300)

  if(WITH_NV_JETSON)
    set_tests_properties(
      test_trt_pool_op
      PROPERTIES ENVIRONMENT FLAGS_fraction_of_gpu_memory_to_use=0.1 TIMEOUT
                 450)
    set_tests_properties(
      test_trt_pool3d_op
      PROPERTIES ENVIRONMENT FLAGS_fraction_of_gpu_memory_to_use=0.1 TIMEOUT
                 450)
  else()
    set_tests_properties(
      test_trt_pool_op
      PROPERTIES ENVIRONMENT FLAGS_fraction_of_gpu_memory_to_use=0.1 TIMEOUT
                 300)
    set_tests_properties(
      test_trt_pool3d_op
      PROPERTIES ENVIRONMENT FLAGS_fraction_of_gpu_memory_to_use=0.1 TIMEOUT 45)
  endif()
  set_tests_properties(test_trt_tile_op PROPERTIES TIMEOUT 60)
  set_tests_properties(test_trt_fc_fuse_quant_dequant_pass PROPERTIES TIMEOUT
                                                                      100)
  set_tests_properties(test_trt_conv_quant_dequant_pass PROPERTIES TIMEOUT 180)
  set_tests_properties(test_trt_matmul_quant_dequant PROPERTIES TIMEOUT 450)
  set_tests_properties(test_trt_conv3d_op PROPERTIES TIMEOUT 60)
  set_tests_properties(test_trt_conv3d_transpose_op PROPERTIES TIMEOUT 60)
  set_tests_properties(test_trt_nearest_interp_v2_op PROPERTIES TIMEOUT 30)
  set_tests_properties(test_trt_multiclass_nms3_op PROPERTIES TIMEOUT 60)

  if(WITH_MKLDNN
     AND TENSORRT_FOUND
     AND WITH_GPU)
    set_tests_properties(test_merge_layernorm_fuse_pass PROPERTIES TIMEOUT 180)
    set_tests_properties(test_skip_merge_layernorm_fuse_pass PROPERTIES TIMEOUT
                                                                        180)
    set_tests_properties(test_trt_emb_eltwise_layernorm_fuse_pass
                         PROPERTIES TIMEOUT 180)

    set_tests_properties(test_fc_fuse_pass PROPERTIES TIMEOUT 240)
    set_tests_properties(test_reverse_roll_fuse_pass PROPERTIES TIMEOUT 120)
    set_tests_properties(test_inplace_op_pass PROPERTIES TIMEOUT 120)
    set_tests_properties(test_transfer_layout_elim_pass PROPERTIES TIMEOUT 300)

    set_tests_properties(test_simplify_with_basic_ops_pass_autoscan
                         PROPERTIES TIMEOUT 240)
    set_tests_properties(test_adaptive_pool2d_convert_global_pass_autoscan
                         PROPERTIES TIMEOUT 100)
    set_tests_properties(test_conv_act_onednn_fuse_pass PROPERTIES TIMEOUT 120)
    set_tests_properties(test_conv_elementwise_add2_act_fuse_pass
                         PROPERTIES TIMEOUT 120)
    set_tests_properties(test_conv_elementwise_add_act_fuse_pass
                         PROPERTIES TIMEOUT 120)
    set_tests_properties(test_onednn_conv_bias_fuse_pass PROPERTIES TIMEOUT 300)
    set_tests_properties(test_onednn_conv_concat_activation_fuse_pass
                         PROPERTIES TIMEOUT 300)
    set_tests_properties(test_onednn_multi_gru_fuse_pass PROPERTIES TIMEOUT 120)
    set_tests_properties(test_onednn_multi_gru_seq_fuse_pass PROPERTIES TIMEOUT
                                                                        120)
    set_tests_properties(test_flatten2_matmul_fuse_pass PROPERTIES TIMEOUT 240)
    set_tests_properties(test_squeeze2_matmul_fuse_pass PROPERTIES TIMEOUT 240)
    set_tests_properties(test_reshape2_matmul_fuse_pass PROPERTIES TIMEOUT 240)
    set_tests_properties(test_preln_layernorm_x_fuse_pass PROPERTIES TIMEOUT
                                                                     240)
    set_tests_properties(test_shuffle_channel_detect_pass PROPERTIES TIMEOUT
                                                                     120)
    if(WIN32)
      set_tests_properties(test_matmul_scale_fuse_pass PROPERTIES TIMEOUT 300)
      set_tests_properties(test_matmul_v2_scale_fuse_pass PROPERTIES TIMEOUT
                                                                     300)
      set_tests_properties(test_map_matmul_v2_to_matmul_pass PROPERTIES TIMEOUT
                                                                        360)
      set_tests_properties(test_map_matmul_v2_to_mul_pass PROPERTIES TIMEOUT
                                                                     360)
      set_tests_properties(test_map_matmul_to_mul_pass PROPERTIES TIMEOUT 360)
      set_tests_properties(test_layernorm_shift_partition_pass
                           PROPERTIES TIMEOUT 360)
    else()
      set_tests_properties(test_matmul_scale_fuse_pass PROPERTIES TIMEOUT 60)
      set_tests_properties(test_matmul_v2_scale_fuse_pass PROPERTIES TIMEOUT 60)
      set_tests_properties(test_map_matmul_v2_to_matmul_pass PROPERTIES TIMEOUT
                                                                        120)
      set_tests_properties(test_map_matmul_v2_to_mul_pass PROPERTIES TIMEOUT
                                                                     120)
      set_tests_properties(test_map_matmul_to_mul_pass PROPERTIES TIMEOUT 120)
      set_tests_properties(test_element_groupnorm_act_fuse_pass
                           PROPERTIES TIMEOUT 120)
      set_tests_properties(test_preln_groupnorm_act_fuse_pass PROPERTIES TIMEOUT
                                                                         120)
      set_tests_properties(test_groupnorm_act_pass_fuse_pass PROPERTIES TIMEOUT
                                                                        120)
      set_tests_properties(test_split_layernorm_to_math_ops_pass
                           PROPERTIES TIMEOUT 240)
    endif()
  endif()

  if(WITH_MKLDNN)
    set_tests_properties(test_onednn_conv_bn_fuse_pass PROPERTIES TIMEOUT 120)
    set_tests_properties(test_onednn_conv_elementwise_add_fuse_pass
                         PROPERTIES TIMEOUT 120)
    set_tests_properties(test_mkldnn_depthwise_conv_pass PROPERTIES TIMEOUT 120)
    set_tests_properties(test_onednn_reshape_transpose_matmul_fuse_pass
                         PROPERTIES TIMEOUT 100)
    set_tests_properties(test_mkldnn_mish_op PROPERTIES TIMEOUT 300)
    set_tests_properties(test_mkldnn_conv3d_op PROPERTIES TIMEOUT 300)
    set_tests_properties(test_mkldnn_prelu_op PROPERTIES TIMEOUT 300)
    set_tests_properties(test_conv_act_onednn_fuse_pass PROPERTIES TIMEOUT 120)
    set_tests_properties(test_conv_transpose_eltwiseadd_bn_fuse_pass
                         PROPERTIES TIMEOUT 250)
    set_tests_properties(test_onednn_matmul_transpose_reshape_fuse_pass
                         PROPERTIES TIMEOUT 100)
    set_tests_properties(test_conv_transpose_bn_fuse_pass PROPERTIES TIMEOUT
                                                                     300)
    set_tests_properties(test_mkldnn_conv_hard_sigmoid_fuse_pass
                         PROPERTIES TIMEOUT 300)
    set_tests_properties(test_mkldnn_conv_hard_swish_fuse_pass
                         PROPERTIES TIMEOUT 300)
    set_tests_properties(test_onednn_batch_norm_act_fuse_pass PROPERTIES TIMEOUT
                                                                         100)
    set_tests_properties(test_mkldnn_matmul_v2_transpose_reshape_fuse_pass
                         PROPERTIES TIMEOUT 100)
    set_tests_properties(test_mkldnn_conv_transpose_bias_fuse_pass
                         PROPERTIES TIMEOUT 100)
    set_tests_properties(test_conv_eltwiseadd_bn_fuse_pass PROPERTIES TIMEOUT
                                                                      300)
    set_tests_properties(test_mkldnn_conv_mish_fuse_pass PROPERTIES TIMEOUT 300)
    set_tests_properties(test_onednn_fc_activation_fuse_pass PROPERTIES TIMEOUT
                                                                        300)
    set_tests_properties(test_mkldnn_conv_affine_channel_fuse_pass
                         PROPERTIES TIMEOUT 60)
  endif()
endif()