diff --git a/cmake/onnxruntime_unittests.cmake b/cmake/onnxruntime_unittests.cmake index 5e63298a4a..7f11111824 100644 --- a/cmake/onnxruntime_unittests.cmake +++ b/cmake/onnxruntime_unittests.cmake @@ -20,21 +20,30 @@ endif() set(disabled_warnings) function(AddTest) cmake_parse_arguments(_UT "DYN" "TARGET" "LIBS;SOURCES;DEPENDS" ${ARGN}) - if(_UT_LIBS) - list(REMOVE_DUPLICATES _UT_LIBS) - endif() list(REMOVE_DUPLICATES _UT_SOURCES) - if (_UT_DEPENDS) - list(REMOVE_DUPLICATES _UT_DEPENDS) - endif(_UT_DEPENDS) - if (${CMAKE_SYSTEM_NAME} STREQUAL "iOS") add_executable(${_UT_TARGET} ${TEST_SRC_DIR}/xctest/orttestmain.m) else() add_executable(${_UT_TARGET} ${_UT_SOURCES}) endif() + if (_UT_DEPENDS) + list(REMOVE_DUPLICATES _UT_DEPENDS) + endif(_UT_DEPENDS) + + if(_UT_LIBS) + list(REMOVE_DUPLICATES _UT_LIBS) + list (FIND _UT_LIBS "cudart" _index) + if (${_index} GREATER -1) + if(WIN32) + target_link_directories(${_UT_TARGET} PRIVATE ${onnxruntime_CUDA_HOME}/x64/lib64) + else() + target_link_directories(${_UT_TARGET} PRIVATE ${onnxruntime_CUDA_HOME}/lib64) + endif() + endif() + endif() + source_group(TREE ${REPO_ROOT} FILES ${_UT_SOURCES}) if (MSVC AND NOT CMAKE_SIZEOF_VOID_P EQUAL 8) @@ -59,7 +68,7 @@ function(AddTest) onnxruntime_add_include_to_target(${_UT_TARGET} date_interface flatbuffers) target_include_directories(${_UT_TARGET} PRIVATE ${TEST_INC_DIR}) if (onnxruntime_USE_CUDA) - target_include_directories(${_UT_TARGET} PRIVATE ${CUDA_INCLUDE_DIRS} ${onnxruntime_CUDNN_HOME}/include) + target_include_directories(${_UT_TARGET} PRIVATE ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES} ${onnxruntime_CUDNN_HOME}/include) endif() if(MSVC) target_compile_options(${_UT_TARGET} PRIVATE "$<$:SHELL:--compiler-options /utf-8>" @@ -849,6 +858,9 @@ if (onnxruntime_BUILD_SHARED_LIB) if(NOT WIN32) list(APPEND onnxruntime_shared_lib_test_LIBS nsync_cpp) endif() + if (onnxruntime_USE_CUDA) + list(APPEND onnxruntime_shared_lib_test_LIBS cudart) + endif() if (CMAKE_SYSTEM_NAME STREQUAL "Android") list(APPEND onnxruntime_shared_lib_test_LIBS ${android_shared_libs}) endif() diff --git a/onnxruntime/test/shared_lib/test_inference.cc b/onnxruntime/test/shared_lib/test_inference.cc index 9b99f3b448..db0094b78f 100644 --- a/onnxruntime/test/shared_lib/test_inference.cc +++ b/onnxruntime/test/shared_lib/test_inference.cc @@ -23,6 +23,10 @@ #include #endif +#ifdef USE_CUDA +#include +#endif + struct Input { const char* name = nullptr; std::vector dims; @@ -601,6 +605,113 @@ TEST(CApiTest, io_binding) { binding.ClearBoundOutputs(); } +#ifdef USE_CUDA +TEST(CApiTest, io_binding_cuda) { + struct CudaMemoryDeleter { + explicit CudaMemoryDeleter(const Ort::Allocator* alloc) { + alloc_ = alloc; + } + void operator()(void* ptr) const { + alloc_->Free(ptr); + } + + const Ort::Allocator* alloc_; + }; + + Ort::SessionOptions session_options; + Ort::ThrowOnError(OrtSessionOptionsAppendExecutionProvider_CUDA(session_options, 0)); + Ort::Session session(*ort_env, MODEL_URI, session_options); + + Ort::MemoryInfo info_cuda("Cuda", OrtAllocatorType::OrtArenaAllocator, 0, OrtMemTypeDefault); + Ort::Allocator cuda_allocator(session, info_cuda); + auto allocator_info = cuda_allocator.GetInfo(); + ASSERT_TRUE(info_cuda == allocator_info); + + const std::array x_shape = {3, 2}; + std::array x_values = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f}; + auto input_data = std::unique_ptr(cuda_allocator.Alloc(x_values.size() * sizeof(float)), + CudaMemoryDeleter(&cuda_allocator)); + ASSERT_NE(input_data.get(), nullptr); + cudaMemcpy(input_data.get(), x_values.data(), sizeof(float) * x_values.size(), cudaMemcpyHostToDevice); + + // Create an OrtValue tensor backed by data on CUDA memory + Ort::Value bound_x = Ort::Value::CreateTensor(info_cuda, reinterpret_cast(input_data.get()), x_values.size(), + x_shape.data(), x_shape.size()); + + const std::array expected_y_shape = {3, 2}; + const std::array expected_y = {1.0f, 4.0f, 9.0f, 16.0f, 25.0f, 36.0f}; + auto output_data = std::unique_ptr(cuda_allocator.Alloc(expected_y.size() * sizeof(float)), + CudaMemoryDeleter(&cuda_allocator)); + ASSERT_NE(output_data.get(), nullptr); + + // Create an OrtValue tensor backed by data on CUDA memory + Ort::Value bound_y = Ort::Value::CreateTensor(info_cuda, reinterpret_cast(output_data.get()), expected_y.size(), + expected_y_shape.data(), expected_y_shape.size()); + + Ort::IoBinding binding(session); + binding.BindInput("X", bound_x); + binding.BindOutput("Y", bound_y); + + session.Run(Ort::RunOptions(), binding); + + // Check the values against the bound raw memory (needs copying from device to host first) + std::array y_values_0; + cudaMemcpy(y_values_0.data(), output_data.get(), sizeof(float) * y_values_0.size(), cudaMemcpyDeviceToHost); + ASSERT_TRUE(std::equal(std::begin(y_values_0), std::end(y_values_0), std::begin(expected_y))); + + // Now compare values via GetOutputValues + { + std::vector output_values = binding.GetOutputValues(); + ASSERT_EQ(output_values.size(), 1U); + const Ort::Value& Y_value = output_values[0]; + ASSERT_TRUE(Y_value.IsTensor()); + Ort::TensorTypeAndShapeInfo type_info = Y_value.GetTensorTypeAndShapeInfo(); + ASSERT_EQ(ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT, type_info.GetElementType()); + auto count = type_info.GetElementCount(); + ASSERT_EQ(expected_y.size(), count); + const float* values = Y_value.GetTensorData(); + + std::array y_values_1; + cudaMemcpy(y_values_1.data(), values, sizeof(float) * y_values_1.size(), cudaMemcpyDeviceToHost); + ASSERT_TRUE(std::equal(std::begin(y_values_1), std::end(y_values_1), std::begin(expected_y))); + } + + { + std::vector output_names = binding.GetOutputNames(); + ASSERT_EQ(1U, output_names.size()); + ASSERT_EQ(output_names[0].compare("Y"), 0); + } + + // Now replace binding of Y with an on device binding instead of pre-allocated memory. + // This is when we can not allocate an OrtValue due to unknown dimensions + { + binding.BindOutput("Y", info_cuda); + session.Run(Ort::RunOptions(), binding); + } + + // Check the output value allocated based on the device binding. + { + std::vector output_values = binding.GetOutputValues(); + ASSERT_EQ(output_values.size(), 1U); + const Ort::Value& Y_value = output_values[0]; + ASSERT_TRUE(Y_value.IsTensor()); + Ort::TensorTypeAndShapeInfo type_info = Y_value.GetTensorTypeAndShapeInfo(); + ASSERT_EQ(ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT, type_info.GetElementType()); + auto count = type_info.GetElementCount(); + ASSERT_EQ(expected_y.size(), count); + const float* values = Y_value.GetTensorData(); + + std::array y_values_2; + cudaMemcpy(y_values_2.data(), values, sizeof(float) * y_values_2.size(), cudaMemcpyDeviceToHost); + ASSERT_TRUE(std::equal(std::begin(y_values_2), std::end(y_values_2), std::begin(expected_y))); + } + + // Clean up + binding.ClearBoundInputs(); + binding.ClearBoundOutputs(); +} +#endif + TEST(CApiTest, create_tensor) { const char* s[] = {"abc", "kmp"}; int64_t expected_len = 2;