From f3e251b9682f93f17471055e551732529330a31a Mon Sep 17 00:00:00 2001 From: cloudhan Date: Mon, 9 Oct 2023 23:50:31 +0800 Subject: [PATCH] Enable tests for Loop and Scan --- .../test/providers/cpu/controlflow/loop_test.cc | 8 ++++++-- .../test/providers/cpu/controlflow/scan_test.cc | 10 +++++++++- 2 files changed, 15 insertions(+), 3 deletions(-) diff --git a/onnxruntime/test/providers/cpu/controlflow/loop_test.cc b/onnxruntime/test/providers/cpu/controlflow/loop_test.cc index 8dcf632192249..9c0b779870c70 100644 --- a/onnxruntime/test/providers/cpu/controlflow/loop_test.cc +++ b/onnxruntime/test/providers/cpu/controlflow/loop_test.cc @@ -358,7 +358,11 @@ void RunTest(int64_t max_iterations, // we want the CUDA provider to be first, and the CPU provider second. all except the Loop node should run on // CUDA given that, which creates the scenario where we need to copy to/from CPU to execute the Loop node correctly. std::vector> execution_providers; +#if defined(USE_CUDA) execution_providers.push_back(DefaultCudaExecutionProvider()); +#elif defined(USE_ROCM) + execution_providers.push_back(DefaultRocmExecutionProvider()); +#endif execution_providers.push_back(DefaultCpuExecutionProvider()); test.Run(expect_result, failure_message, {kTensorrtExecutionProvider}, nullptr, &execution_providers); @@ -1038,8 +1042,8 @@ TEST(Loop, IterationCountAsOutput) { test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider}); } -#ifdef USE_CUDA -// test that when part of the subgraph run on CUDA it executes successfully +#if defined(USE_CUDA) || defined(USE_ROCM) +// test that when part of the subgraph run on CUDA/ROCm it executes successfully TEST(Loop, MixedExecutionProviders) { RunOptions options{}; options.mixed_execution_providers = true; diff --git a/onnxruntime/test/providers/cpu/controlflow/scan_test.cc b/onnxruntime/test/providers/cpu/controlflow/scan_test.cc index 8008fd129c19b..3d46893cdb82d 100644 --- a/onnxruntime/test/providers/cpu/controlflow/scan_test.cc +++ b/onnxruntime/test/providers/cpu/controlflow/scan_test.cc @@ -411,7 +411,11 @@ static void RunTest_v9(const std::string test_name, int64_t sequence_len, int64_ // we want the CUDA provider to be first, and the CPU provider second. all except the Scan node should run on // CUDA given that, which creates the scenario where we need to copy to/from CPU to execute the Scan node correctly. std::vector> execution_providers; +#if defined(USE_CUDA) execution_providers.push_back(DefaultCudaExecutionProvider()); +#elif defined(USE_ROCM) + execution_providers.push_back(DefaultRocmExecutionProvider()); +#endif execution_providers.push_back(DefaultCpuExecutionProvider()); test.Run(expect_result, failure_message, options.excluded_provider_types, nullptr, &execution_providers); @@ -1162,7 +1166,11 @@ void UnknownDimInSubgraphOutput(bool is_v8, bool mixed_execution_providers = fal // we want the CUDA provider to be first, and the CPU provider second. all except the Scan node should run on // CUDA given that, which creates the scenario where we need to copy to/from CPU to execute the Scan node correctly. std::vector> execution_providers; +#if defined(USE_CUDA) execution_providers.push_back(DefaultCudaExecutionProvider()); +#elif defined(USE_ROCM) + execution_providers.push_back(DefaultRocmExecutionProvider()); +#endif execution_providers.push_back(DefaultCpuExecutionProvider()); test.Run(OpTester::ExpectResult::kExpectSuccess, "", RunOptions().excluded_provider_types, nullptr, @@ -1174,7 +1182,7 @@ void UnknownDimInSubgraphOutput(bool is_v8, bool mixed_execution_providers = fal TEST_8_AND_9(UnknownDimInSubgraphOutput); -#ifdef USE_CUDA +#if defined(USE_CUDA) || defined(USE_ROCM) TEST(Scan, MixedExecutionProviders) { RunOptions options{}; options.is_v8 = false;