From 8b9656717b4356732096d2e9a40f2d31f771a04a Mon Sep 17 00:00:00 2001 From: Changming Sun Date: Wed, 19 Jun 2024 14:44:39 -0700 Subject: [PATCH 01/52] Fix a perm issue in Windows Static Analysis pipeline (#21100) ### Description Due to a security setting change, now we need to explicitly set the permissions. I forgot doing that when bringing the old change back. ### Motivation and Context Now the pipeline cannot publish scanning result to Github --- .github/workflows/sca.yml | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/.github/workflows/sca.yml b/.github/workflows/sca.yml index eb35f6a814987..0867d4c343e91 100644 --- a/.github/workflows/sca.yml +++ b/.github/workflows/sca.yml @@ -16,6 +16,8 @@ env: jobs: Onnxruntime-SCA-training-CUDA: + permissions: + security-events: write runs-on: ["self-hosted", "1ES.Pool=onnxruntime-github-vs2022-mms"] steps: - uses: actions/checkout@v4 @@ -57,6 +59,8 @@ jobs: # No python Onnxruntime-SCA-win32-WINML-x64: + permissions: + security-events: write runs-on: ["self-hosted", "1ES.Pool=onnxruntime-github-vs2022-mms"] steps: - uses: actions/checkout@v4 @@ -95,6 +99,8 @@ jobs: # No java, No python Onnxruntime-SCA-win32-WINML-x86: + permissions: + security-events: write runs-on: ["self-hosted", "1ES.Pool=onnxruntime-github-vs2022-mms"] steps: - uses: actions/checkout@v4 From 8ab8e649a70e6f8a237bc067c8e6c89cdcb2b262 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Cl=C3=A9ment=20P=C3=A9ron?= Date: Thu, 20 Jun 2024 01:14:58 +0200 Subject: [PATCH 02/52] tools: build: fix typo (#21052) ### Description Typo in the python build script --- cmake/adjust_global_compile_flags.cmake | 2 +- tools/ci_build/build.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/cmake/adjust_global_compile_flags.cmake b/cmake/adjust_global_compile_flags.cmake index 690b6d4e66154..6eb784a4063ed 100644 --- a/cmake/adjust_global_compile_flags.cmake +++ b/cmake/adjust_global_compile_flags.cmake @@ -321,7 +321,7 @@ else() string(APPEND CMAKE_CXX_FLAGS " -mavx512f -mavx512cd -mavx512bw -mavx512dq -mavx512vl") string(APPEND CMAKE_C_FLAGS " -mavx512f -mavx512cd -mavx512bw -mavx512dq -mavx512vl") endif() - if (CMAKE_SYSTEM_NAME STREQUAL "Android" AND Onnxruntime_GCOV_COVERAGE) + if (CMAKE_SYSTEM_NAME STREQUAL "Android" AND onnxruntime_GCOV_COVERAGE) string(APPEND CMAKE_CXX_FLAGS " -g -O0 --coverage ") string(APPEND CMAKE_C_FLAGS " -g -O0 --coverage ") endif() diff --git a/tools/ci_build/build.py b/tools/ci_build/build.py index 3e1b6528440b5..1145509eef261 100644 --- a/tools/ci_build/build.py +++ b/tools/ci_build/build.py @@ -1079,7 +1079,7 @@ def generate_build_tree( "-Donnxruntime_USE_NCCL=" + ("ON" if args.enable_nccl else "OFF"), "-Donnxruntime_BUILD_BENCHMARKS=" + ("ON" if args.build_micro_benchmarks else "OFF"), "-Donnxruntime_USE_ROCM=" + ("ON" if args.use_rocm else "OFF"), - "-DOnnxruntime_GCOV_COVERAGE=" + ("ON" if args.code_coverage else "OFF"), + "-Donnxruntime_GCOV_COVERAGE=" + ("ON" if args.code_coverage else "OFF"), "-Donnxruntime_USE_MPI=" + ("ON" if args.use_mpi else "OFF"), "-Donnxruntime_ENABLE_MEMORY_PROFILE=" + ("ON" if args.enable_memory_profile else "OFF"), "-Donnxruntime_ENABLE_CUDA_LINE_NUMBER_INFO=" + ("ON" if args.enable_cuda_line_info else "OFF"), From be423747b19c9011f0ef0544d840fcbc7d769015 Mon Sep 17 00:00:00 2001 From: Changming Sun Date: Wed, 19 Jun 2024 16:21:33 -0700 Subject: [PATCH 03/52] Delete pyop (#21094) ### Description Remove the "--enable_language_interop_ops" build flag, because the code is incompatible with the latest numpy, and the build flag is not used anywhere except a macOS CI pipeline. It does not seem to have a ship plan. ### Motivation and Context The build error was: ``` onnxruntime/core/language_interop_ops/pyop/pyop.cc:122:85: error: no member named 'elsize' in '_PyArray_Descr' static_cast(PyArray_DescrFromType(type)->elsize), ~~~~~~~~~~~~~~~~~~~~~~~~~~~ ^ ``` --- cmake/CMakeLists.txt | 10 +- cmake/onnxruntime.cmake | 7 - cmake/onnxruntime_language_interop_ops.cmake | 8 - cmake/onnxruntime_pyop.cmake | 12 - cmake/onnxruntime_python.cmake | 7 - cmake/onnxruntime_training.cmake | 4 - cmake/onnxruntime_unittests.cmake | 20 - .../language_interop_ops.cc | 65 --- .../language_interop_ops.h | 16 - .../core/language_interop_ops/pyop/pyop.cc | 399 ------------------ .../core/language_interop_ops/pyop/pyop.h | 101 ----- tools/ci_build/build.py | 6 - .../azure-pipelines/mac-ci-pipeline.yml | 2 +- 13 files changed, 2 insertions(+), 655 deletions(-) delete mode 100644 cmake/onnxruntime_language_interop_ops.cmake delete mode 100644 cmake/onnxruntime_pyop.cmake delete mode 100644 onnxruntime/core/language_interop_ops/language_interop_ops.cc delete mode 100644 onnxruntime/core/language_interop_ops/language_interop_ops.h delete mode 100644 onnxruntime/core/language_interop_ops/pyop/pyop.cc delete mode 100644 onnxruntime/core/language_interop_ops/pyop/pyop.h diff --git a/cmake/CMakeLists.txt b/cmake/CMakeLists.txt index 00097c25d2ba5..ce22def914851 100644 --- a/cmake/CMakeLists.txt +++ b/cmake/CMakeLists.txt @@ -123,7 +123,6 @@ option(onnxruntime_GCOV_COVERAGE "Compile with options necessary to run code cov option(onnxruntime_DONT_VECTORIZE "Do not vectorize operations in Eigen" OFF) option(onnxruntime_USE_FULL_PROTOBUF "Link to libprotobuf instead of libprotobuf-lite when this option is ON" OFF) -option(onnxruntime_ENABLE_LANGUAGE_INTEROP_OPS "Enable operator implemented in language other than cpp" OFF) option(onnxruntime_DEBUG_NODE_INPUTS_OUTPUTS "Dump debug information about node inputs and outputs when executing the model." OFF) cmake_dependent_option(onnxruntime_DEBUG_NODE_INPUTS_OUTPUTS_ENABLE_DUMP_TO_SQLDB "Build dump debug information about node inputs and outputs with support for sql database." OFF "onnxruntime_DEBUG_NODE_INPUTS_OUTPUTS" OFF) option(onnxruntime_USE_DML "Build with DirectML support" OFF) @@ -439,13 +438,6 @@ if (onnxruntime_ENABLE_MEMORY_PROFILE) endif() set(ONNX_ML 1) -if (NOT onnxruntime_ENABLE_PYTHON) - set(onnxruntime_ENABLE_LANGUAGE_INTEROP_OPS OFF) -endif() - -if (onnxruntime_ENABLE_LANGUAGE_INTEROP_OPS) - add_compile_definitions(ENABLE_LANGUAGE_INTEROP_OPS) -endif() if (NOT (UNIX AND onnxruntime_ENABLE_PYTHON AND onnxruntime_ENABLE_TRAINING AND (NOT onnxruntime_BUILD_SHARED_LIB))) if (onnxruntime_ENABLE_TRAINING_TORCH_INTEROP) @@ -578,7 +570,7 @@ endif() #Need python to generate def file if (onnxruntime_BUILD_SHARED_LIB OR onnxruntime_ENABLE_PYTHON) if (onnxruntime_ENABLE_PYTHON) - if (onnxruntime_ENABLE_LANGUAGE_INTEROP_OPS OR onnxruntime_REQUIRE_PYTHON_EMBED_LIB) + if (onnxruntime_REQUIRE_PYTHON_EMBED_LIB) find_package(Python 3.8 COMPONENTS Interpreter Development NumPy) else() find_package(Python 3.8 COMPONENTS Interpreter Development.Module NumPy) diff --git a/cmake/onnxruntime.cmake b/cmake/onnxruntime.cmake index e15c8a046dc20..977aa44b0e8d7 100644 --- a/cmake/onnxruntime.cmake +++ b/cmake/onnxruntime.cmake @@ -205,13 +205,6 @@ set(onnxruntime_INTERNAL_LIBRARIES onnxruntime_flatbuffers ) -if (onnxruntime_ENABLE_LANGUAGE_INTEROP_OPS) - list(APPEND onnxruntime_INTERNAL_LIBRARIES - onnxruntime_language_interop - onnxruntime_pyop - ) -endif() - if (onnxruntime_USE_EXTENSIONS) list(APPEND onnxruntime_INTERNAL_LIBRARIES onnxruntime_extensions diff --git a/cmake/onnxruntime_language_interop_ops.cmake b/cmake/onnxruntime_language_interop_ops.cmake deleted file mode 100644 index 5d88332eb2e9a..0000000000000 --- a/cmake/onnxruntime_language_interop_ops.cmake +++ /dev/null @@ -1,8 +0,0 @@ -# Copyright (c) Microsoft Corporation. All rights reserved. -# Licensed under the MIT License. -include(onnxruntime_pyop.cmake) -file (GLOB onnxruntime_language_interop_ops_src "${ONNXRUNTIME_ROOT}/core/language_interop_ops/language_interop_ops.cc") -onnxruntime_add_static_library(onnxruntime_language_interop ${onnxruntime_language_interop_ops_src}) -add_dependencies(onnxruntime_language_interop onnxruntime_pyop) -onnxruntime_add_include_to_target(onnxruntime_language_interop onnxruntime_common onnxruntime_graph onnxruntime_framework onnxruntime_pyop onnx onnx_proto ${PROTOBUF_LIB} flatbuffers::flatbuffers safeint_interface Boost::mp11) -target_include_directories(onnxruntime_language_interop PRIVATE ${ONNXRUNTIME_ROOT} ${eigen_INCLUDE_DIRS}) \ No newline at end of file diff --git a/cmake/onnxruntime_pyop.cmake b/cmake/onnxruntime_pyop.cmake deleted file mode 100644 index f7583690945a1..0000000000000 --- a/cmake/onnxruntime_pyop.cmake +++ /dev/null @@ -1,12 +0,0 @@ -# Copyright (c) Microsoft Corporation. All rights reserved. -# Licensed under the MIT License. -onnxruntime_add_static_library(onnxruntime_pyop "${ONNXRUNTIME_ROOT}/core/language_interop_ops/pyop/pyop.cc") -add_dependencies(onnxruntime_pyop ${onnxruntime_EXTERNAL_DEPENDENCIES}) -onnxruntime_add_include_to_target(onnxruntime_pyop onnxruntime_common onnxruntime_graph onnxruntime_framework onnx onnx_proto ${PROTOBUF_LIB} flatbuffers::flatbuffers ${GSL_TARGET} Boost::mp11) -target_include_directories(onnxruntime_pyop PRIVATE ${ONNXRUNTIME_ROOT} ${eigen_INCLUDE_DIRS}) -onnxruntime_add_include_to_target(onnxruntime_pyop Python::Module Python::NumPy) -if (TARGET Python::Python) - target_link_libraries(onnxruntime_pyop PRIVATE Python::Python) -else() - target_link_libraries(onnxruntime_pyop PRIVATE Python::Module) -endif() diff --git a/cmake/onnxruntime_python.cmake b/cmake/onnxruntime_python.cmake index da5732e484d6b..3c2833d87d652 100644 --- a/cmake/onnxruntime_python.cmake +++ b/cmake/onnxruntime_python.cmake @@ -193,10 +193,6 @@ target_link_libraries(onnxruntime_pybind11_state PRIVATE ${pybind11_lib} ) -if (onnxruntime_ENABLE_LANGUAGE_INTEROP_OPS) - target_link_libraries(onnxruntime_pybind11_state PRIVATE onnxruntime_language_interop onnxruntime_pyop) -endif() - set(onnxruntime_pybind11_state_dependencies ${onnxruntime_EXTERNAL_DEPENDENCIES} ${pybind11_dep} @@ -1027,6 +1023,3 @@ if (onnxruntime_USE_QNN) endif() endif() -if (onnxruntime_ENABLE_LANGUAGE_INTEROP_OPS) - include(onnxruntime_language_interop_ops.cmake) -endif() diff --git a/cmake/onnxruntime_training.cmake b/cmake/onnxruntime_training.cmake index f9ba2b341f741..01590a431205c 100644 --- a/cmake/onnxruntime_training.cmake +++ b/cmake/onnxruntime_training.cmake @@ -141,10 +141,6 @@ if (onnxruntime_BUILD_UNIT_TESTS) Boost::mp11 safeint_interface ) - if (onnxruntime_ENABLE_LANGUAGE_INTEROP_OPS) - list(APPEND ONNXRUNTIME_LIBS onnxruntime_language_interop onnxruntime_pyop) - endif() - if(UNIX AND NOT APPLE) if (HAS_NO_MAYBE_UNINITIALIZED) target_compile_options(onnxruntime_training_mnist PUBLIC "-Wno-maybe-uninitialized") diff --git a/cmake/onnxruntime_unittests.cmake b/cmake/onnxruntime_unittests.cmake index fb96fd1cad39d..ed71e7a57a500 100644 --- a/cmake/onnxruntime_unittests.cmake +++ b/cmake/onnxruntime_unittests.cmake @@ -583,10 +583,6 @@ if(onnxruntime_USE_ARMNN) list(APPEND onnxruntime_test_providers_dependencies onnxruntime_providers_armnn) endif() -if (onnxruntime_ENABLE_LANGUAGE_INTEROP_OPS) - set(ONNXRUNTIME_INTEROP_TEST_LIBS PRIVATE onnxruntime_language_interop onnxruntime_pyop) -endif() - set(ONNXRUNTIME_TEST_LIBS onnxruntime_session ${ONNXRUNTIME_INTEROP_TEST_LIBS} @@ -916,10 +912,6 @@ endif() if (onnxruntime_DEBUG_NODE_INPUTS_OUTPUTS) target_compile_definitions(onnxruntime_test_all PRIVATE DEBUG_NODE_INPUTS_OUTPUTS) endif() - -if (onnxruntime_ENABLE_LANGUAGE_INTEROP_OPS) - target_link_libraries(onnxruntime_test_all PRIVATE onnxruntime_language_interop onnxruntime_pyop) -endif() if (onnxruntime_USE_ROCM) if (onnxruntime_USE_COMPOSABLE_KERNEL) target_compile_definitions(onnxruntime_test_all PRIVATE USE_COMPOSABLE_KERNEL) @@ -1057,10 +1049,6 @@ set(onnx_test_libs onnx_test_data_proto ${onnxruntime_EXTERNAL_LIBRARIES}) -if (onnxruntime_ENABLE_LANGUAGE_INTEROP_OPS) - list(APPEND onnx_test_libs onnxruntime_language_interop onnxruntime_pyop) -endif() - if (NOT IOS) onnxruntime_add_executable(onnx_test_runner ${onnx_test_runner_src_dir}/main.cc) if(MSVC) @@ -1241,10 +1229,6 @@ if (NOT onnxruntime_ENABLE_TRAINING_TORCH_INTEROP) endif() set_target_properties(onnxruntime_perf_test PROPERTIES FOLDER "ONNXRuntimeTest") - if (onnxruntime_ENABLE_LANGUAGE_INTEROP_OPS AND NOT onnxruntime_BUILD_SHARED_LIB) - target_link_libraries(onnxruntime_perf_test PRIVATE onnxruntime_language_interop onnxruntime_pyop) - endif() - if (onnxruntime_USE_TVM) if (WIN32) target_link_options(onnxruntime_perf_test PRIVATE "/STACK:4000000") @@ -1474,10 +1458,6 @@ endif() onnxruntime_flatbuffers ) - if (onnxruntime_ENABLE_LANGUAGE_INTEROP_OPS) - list(APPEND ONNXRUNTIME_TEST_LIBS onnxruntime_language_interop onnxruntime_pyop) - endif() - target_link_libraries(onnxruntime_test_trainer PRIVATE ${ONNXRUNTIME_TEST_LIBS} ${onnxruntime_EXTERNAL_LIBRARIES} diff --git a/onnxruntime/core/language_interop_ops/language_interop_ops.cc b/onnxruntime/core/language_interop_ops/language_interop_ops.cc deleted file mode 100644 index b40ee08479055..0000000000000 --- a/onnxruntime/core/language_interop_ops/language_interop_ops.cc +++ /dev/null @@ -1,65 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT License. - -#include "language_interop_ops.h" -#include "core/framework/tensorprotoutils.h" -#include "core/platform/env.h" -#include "core/session/inference_session.h" -#include "pyop/pyop.h" -#include - -namespace onnxruntime { - -void LoadInterOp(const std::basic_string& model_uri, InterOpDomains& domains, const InterOpLogFunc& log_func) { - int fd; - - // match the error message from model.cc to keep the nodejs tests happy. - // as this is deprecated just cut-and-paste equivalent code for now. - auto status = Env::Default().FileOpenRd(model_uri, fd); - if (!status.IsOK()) { - if (status.Category() == common::SYSTEM) { - switch (status.Code()) { - case ENOENT: - status = ORT_MAKE_STATUS(ONNXRUNTIME, NO_SUCHFILE, "Load model ", ToUTF8String(model_uri), - " failed. File doesn't exist"); - break; - case EINVAL: - status = ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Load model ", ToUTF8String(model_uri), " failed"); - break; - default: - status = ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "system error number ", status.Code()); - } - } - } - - ORT_ENFORCE(status.IsOK(), status.ErrorMessage()); - - google::protobuf::io::FileInputStream f(fd); - f.SetCloseOnDelete(true); - ONNX_NAMESPACE::ModelProto model_proto; - ORT_ENFORCE(model_proto.ParseFromZeroCopyStream(&f), "Failed to parse model proto"); - LoadInterOp(model_proto, domains, log_func); -} - -void LoadInterOp(const ONNX_NAMESPACE::ModelProto& model_proto, InterOpDomains& domains, const InterOpLogFunc& log_func) { - LoadInterOp(model_proto.graph(), domains, log_func); -} - -void LoadInterOp(const ONNX_NAMESPACE::GraphProto& graph_proto, InterOpDomains& domains, const InterOpLogFunc& log_func) { - for (int i = 0; i < graph_proto.node_size(); ++i) { - const auto& node_proto = graph_proto.node(i); - if (node_proto.op_type() == "PyOp") { - auto pyop_domain = Ort::CustomOpDomain(node_proto.domain().c_str()); - pyop_domain.Add(LoadPyOp(node_proto, log_func)); - domains.push_back(std::move(pyop_domain)); - } else { - for (int j = 0, limit = node_proto.attribute_size(); j < limit; ++j) { - const auto& attr = node_proto.attribute(j); - if (utils::HasGraph(attr)) { - LoadInterOp(attr.g(), domains, log_func); // load pyop in subgraph - } - } // for - } // else - } // for -} -} // namespace onnxruntime diff --git a/onnxruntime/core/language_interop_ops/language_interop_ops.h b/onnxruntime/core/language_interop_ops/language_interop_ops.h deleted file mode 100644 index 2ab5945b17bc2..0000000000000 --- a/onnxruntime/core/language_interop_ops/language_interop_ops.h +++ /dev/null @@ -1,16 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT License. -#pragma once -#include -#include -#include -#include "core/graph/graph.h" -#include "core/session/onnxruntime_cxx_api.h" - -namespace onnxruntime { -using InterOpLogFunc = std::function; -using InterOpDomains = std::vector; -void LoadInterOp(const std::basic_string& model_uri, InterOpDomains& domains, const InterOpLogFunc& log_func); -void LoadInterOp(const ONNX_NAMESPACE::ModelProto& model_proto, InterOpDomains& domains, const InterOpLogFunc& log_func); -void LoadInterOp(const ONNX_NAMESPACE::GraphProto& graph_proto, InterOpDomains& domains, const InterOpLogFunc& log_func); -} // namespace onnxruntime diff --git a/onnxruntime/core/language_interop_ops/pyop/pyop.cc b/onnxruntime/core/language_interop_ops/pyop/pyop.cc deleted file mode 100644 index ccbe4c0d83006..0000000000000 --- a/onnxruntime/core/language_interop_ops/pyop/pyop.cc +++ /dev/null @@ -1,399 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT License. - -#include "pyop.h" - -#ifdef _WIN32 -#define LIB_PYOP "onnxruntime_pywrapper.dll" -#define LOAD_PYOP_LIB(n, v, m) ORT_ENFORCE((v = LoadLibraryA(n)) != nullptr, m) -#else -#ifdef __APPLE__ -#define LIB_PYOP "./libonnxruntime_pywrapper.dylib" -#else -#define LIB_PYOP "./libonnxruntime_pywrapper.so" -#endif -#define LOAD_PYOP_LIB(n, v, m) ORT_ENFORCE((v = dlopen(n, RTLD_NOW | RTLD_GLOBAL)) != nullptr, m) -#include "dlfcn.h" -#endif - -#include "core/framework/tensorprotoutils.h" -#include "core/platform/env.h" -#ifdef _DEBUG -#undef _DEBUG -#include -#define _DEBUG -#else -#include -#endif - -#define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION -#include "numpy/arrayobject.h" -#include -#include -#include -#include -#include -#include -#include -#include -#include - -namespace onnxruntime { - -PyOpLibProxy& PyOpLibProxy::GetInstance() { - static PyOpLibProxy proxy; - return proxy; -} - -class Scope { - public: - Scope(const std::vector& objs = {}) : objs_(objs) { - mtx_.lock(); - } - ~Scope() { - for (auto obj : objs_) { - Py_XDECREF(obj); - } - mtx_.unlock(); - } - void Add(PyObject* obj) { - objs_.push_back(obj); - } - - private: - static std::mutex mtx_; - std::vector objs_; -}; - -PyOpLibProxy::PyOpLibProxy() { - Scope scope; - Py_Initialize(); - if (_import_array() < 0) { - return; - } - auto path_list = PySys_GetObject("path"); // do not release it - if (nullptr == path_list || !PyList_Check(path_list) || - PyList_Append(path_list, PyUnicode_FromString(".")) != 0) { - return; - } - initialized_ = true; -} - -PyOpLibProxy::~PyOpLibProxy() { - if (initialized_) { - Py_Finalize(); - } -} - -std::mutex Scope::mtx_; - -const char* PyOpLibProxy::GetLastErrorMessage(std::string& err) { - Scope scope; - if (PyErr_Occurred()) { - PyObject *type, *value, *trace; - PyErr_Fetch(&type, &value, &trace); - if (nullptr != value) { - auto pyVal = PyObject_Repr(value); - scope.Add(pyVal); - auto pyStr = PyUnicode_AsEncodedString(pyVal, "utf-8", "Error ~"); - scope.Add(pyStr); - err = PyBytes_AS_STRING(pyStr); - } - PyErr_Restore(type, value, trace); - } - return err.c_str(); -} - -int32_t PyOpLibProxy::GetGil() const { - return PyGILState_Ensure(); -} - -void PyOpLibProxy::PutGil(int32_t state) const { - PyGILState_Release((PyGILState_STATE)state); -} - -PyObject* MakePyObj(const void* data, int32_t type, const std::vector& dim) { - std::vector np_dim; - for (auto d : dim) { - np_dim.push_back(static_cast(d)); - } - auto pyObj = static_cast(PyArray_EMPTY(static_cast(np_dim.size()), np_dim.data(), type, 0)); - auto data_len = std::accumulate(begin(np_dim), end(np_dim), - static_cast(PyArray_DescrFromType(type)->elsize), - std::multiplies()); - auto np_array = reinterpret_cast(pyObj); - memcpy(PyArray_DATA(np_array), data, data_len); - return pyObj; -} - -bool ExtractOutput(PyObject* pyObj, - std::vector>& outputs, - std::vector& outputs_elem_size, - std::vector>& outputs_dim) { - if (!PyArray_Check(pyObj)) { - return false; - } - - outputs_dim.push_back({}); - auto np_array = reinterpret_cast(pyObj); - outputs_elem_size.push_back(static_cast(PyArray_ITEMSIZE(np_array))); - - for (int i = 0; i < PyArray_NDIM(np_array); ++i) { - outputs_dim.back().push_back(PyArray_SHAPE(np_array)[i]); - } - - auto data_len = std::accumulate(begin(outputs_dim.back()), - end(outputs_dim.back()), - static_cast(outputs_elem_size.back()), - std::multiplies()); - - outputs.push_back(std::unique_ptr(new char[data_len])); - memcpy(static_cast(outputs.back().get()), PyArray_DATA(np_array), data_len); - return true; -} - -void* PyOpLibProxy::NewInstance(const char* module, const char* class_name, - const std::unordered_map& args) { - Scope scope; - auto pyModule = PyImport_ImportModule(module); - if (nullptr == pyModule) { - return nullptr; - } - - scope.Add(pyModule); - auto pyClass = PyObject_GetAttrString(pyModule, class_name); - if (nullptr == pyClass) { - return nullptr; - } - - scope.Add(pyClass); - auto empty_args = PyTuple_New(0); - scope.Add(empty_args); - auto named_args = PyDict_New(); - scope.Add(named_args); - for (const auto& iter : args) { - PyDict_SetItemString(named_args, iter.first.c_str(), PyUnicode_FromString(iter.second.c_str())); - } - - return PyObject_Call(pyClass, empty_args, named_args); -} - -void PyOpLibProxy::ReleaseInstance(void* instance) { - Scope scope({static_cast(instance)}); -} - -bool PyOpLibProxy::InvokePythonFunc(void* raw_inst, - const char* function, - const std::vector& inputs, - const std::vector& inputs_type, - const std::vector>& inputs_dim, - std::vector>& outputs, - std::vector& outputs_elem_size, - std::vector>& outputs_dim, - std::function logging_func) { - Scope scope; - auto instance = static_cast(raw_inst); - if (nullptr == instance || nullptr == function) { - logging_func("InvokePythonFunc: found invalid instance or function"); - return false; - } - - auto pyFunc = PyObject_GetAttrString(instance, function); - if (nullptr == pyFunc) { - logging_func("InvokePythonFunc: failed to create function object"); - return false; - } - - scope.Add(pyFunc); - auto pyArgs = PyTuple_New(inputs.size()); - for (size_t i = 0; i < inputs.size(); ++i) { - PyTuple_SetItem(pyArgs, i, MakePyObj(inputs[i], inputs_type[i], inputs_dim[i])); - } - - scope.Add(pyArgs); - auto pyResult = PyObject_CallObject(pyFunc, pyArgs); - if (nullptr == pyResult) { - logging_func("InvokePythonFunc: no result"); - return false; - } - - scope.Add(pyResult); - if (PyArray_Check(pyResult)) { - ExtractOutput(pyResult, outputs, outputs_elem_size, outputs_dim); - } else if (PyTuple_Check(pyResult)) { - for (int32_t i = 0; i < PyTuple_Size(pyResult); ++i) { - if (!ExtractOutput(PyTuple_GetItem(pyResult, i), outputs, outputs_elem_size, outputs_dim)) { - logging_func("InvokePythonFunc: failed to extract output"); - return false; - } - } - } else { - logging_func("InvokePythonFunc: returned value must be numpy(s)"); - return false; - } - return true; -} // bool InvokePythonFunc - -PyCustomKernel::PyCustomKernel(const OnnxAttrs& attrs, - const std::string& module, - const std::string& class_name, - const std::string& compute, - PyOpLogFunc logging_func) : attrs_(attrs), module_(module), class_name_(class_name), compute_(compute), logging_func_(logging_func) { - std::string err; - auto state = PyOpLibProxy::GetInstance().GetGil(); - ORT_ENFORCE(PyOpLibProxy::GetInstance().Initialized(), "Py library not properly initialized."); - instance_ = PyOpLibProxy::GetInstance().NewInstance(module.c_str(), class_name_.c_str(), attrs_); - PyOpLibProxy::GetInstance().PutGil(state); - ORT_ENFORCE(nullptr != instance_, PyOpLibProxy::GetInstance().GetLastErrorMessage(err)); -} - -PyCustomKernel::~PyCustomKernel() { - if (nullptr != instance_) { - auto state = PyOpLibProxy::GetInstance().GetGil(); - PyOpLibProxy::GetInstance().ReleaseInstance(instance_); - PyOpLibProxy::GetInstance().PutGil(state); - instance_ = nullptr; - } -} - -// Do nothing since Custom Op does not trigger shape inference -void PyCustomKernel::GetOutputShape(OrtKernelContext*, size_t, OrtTensorTypeAndShapeInfo*) {} - -void PyCustomKernel::Compute(OrtKernelContext* context) { - ORT_ENFORCE(nullptr != context); - - Ort::KernelContext ctx(context); - const auto inputs_count = ctx.GetInputCount(); - - std::vector inputs; - std::vector> outputs; - std::vector inputs_type, outputs_elem_size; - std::vector> inputs_dim, outputs_dim; - - inputs.reserve(inputs_count); - inputs_dim.reserve(inputs_count); - for (size_t i = 0; i < inputs_count; ++i) { - auto value = ctx.GetInput(i); - ORT_ENFORCE(value.IsTensor(), "input must be a tensor"); - - inputs.push_back(value.GetTensorRawData()); - - auto type_and_shape = value.GetTensorTypeAndShapeInfo(); - inputs_type.push_back(GetNumpyType(type_and_shape.GetElementType())); - auto shape = type_and_shape.GetShape(); - inputs_dim.push_back(std::move(shape)); - } - - std::string err; - auto state = PyOpLibProxy::GetInstance().GetGil(); - ORT_ENFORCE(PyOpLibProxy::GetInstance().InvokePythonFunc(instance_, compute_.c_str(), inputs, inputs_type, - inputs_dim, outputs, outputs_elem_size, - outputs_dim, logging_func_), - PyOpLibProxy::GetInstance().GetLastErrorMessage(err)); // ORT_ENFORCE - PyOpLibProxy::GetInstance().PutGil(state); - - for (size_t i = 0; i < outputs.size(); ++i) { - auto ort_output = ctx.GetOutput(i, outputs_dim[i].data(), outputs_dim[i].size()); - auto output_mem_addr = ort_output.GetTensorMutableData(); - auto output_len = std::accumulate(begin(outputs_dim[i]), end(outputs_dim[i]), static_cast(outputs_elem_size[i]), std::multiplies()); - memcpy(output_mem_addr, outputs[i].get(), output_len); - } -} - -int32_t PyCustomKernel::GetNumpyType(int32_t elem_type) const { - int32_t numpy_type; - namespace on = ONNX_NAMESPACE; - switch (elem_type) { - case on::TensorProto_DataType_BOOL: - numpy_type = 0; - break; - case on::TensorProto_DataType_INT8: - numpy_type = 1; - break; - case on::TensorProto_DataType_UINT8: - numpy_type = 2; - break; - case on::TensorProto_DataType_INT16: - numpy_type = 3; - break; - case on::TensorProto_DataType_UINT16: - numpy_type = 4; - break; - case on::TensorProto_DataType_INT32: - numpy_type = 5; - break; - case on::TensorProto_DataType_UINT32: - numpy_type = 6; - break; - case on::TensorProto_DataType_INT64: - numpy_type = 9; - break; - case on::TensorProto_DataType_UINT64: - numpy_type = 10; - break; - case on::TensorProto_DataType_FLOAT: - numpy_type = 11; - break; - case on::TensorProto_DataType_DOUBLE: - numpy_type = 12; - break; - default: - ORT_THROW("Input primitive type not supported: ", elem_type); - } - return numpy_type; -} - -PyCustomOp::PyCustomOp(const OnnxAttrs& attrs, - const OnnxTypes& inputs_type, - const OnnxTypes& outputs_type, - const std::string& module, - const std::string& class_name, - const std::string& compute, - PyOpLogFunc logging_func) : attrs_(attrs), inputs_type_(inputs_type), outputs_type_(outputs_type), module_(module), class_name_(class_name), compute_(compute), logging_func_(logging_func) { OrtCustomOp::version = ORT_API_VERSION; } - -void* PyCustomOp::CreateKernel(const OrtApi&, const OrtKernelInfo*) const { - return new PyCustomKernel(attrs_, module_, class_name_, compute_, logging_func_); -} - -const char* PyCustomOp::GetName() const { return "PyOp"; } - -size_t PyCustomOp::GetInputTypeCount() const { return inputs_type_.size(); } -ONNXTensorElementDataType PyCustomOp::GetInputType(size_t index) const { return inputs_type_[index]; } - -size_t PyCustomOp::GetOutputTypeCount() const { return outputs_type_.size(); } -ONNXTensorElementDataType PyCustomOp::GetOutputType(size_t index) const { return outputs_type_[index]; } - -PyCustomOp* LoadPyOp(const ONNX_NAMESPACE::NodeProto& node_proto, PyOpLogFunc log_func) { - OnnxAttrs onnx_attrs; - OnnxTypes input_types, output_types; - std::string module, class_name, compute = "compute"; - for (int j = 0; j < node_proto.attribute_size(); ++j) { - const auto& attr = node_proto.attribute(j); - if (utils::HasString(attr)) { - if (attr.name() == "module") - module = attr.s(); - else if (attr.name() == "class_name") - class_name = attr.s(); - else if (attr.name() == "compute") - compute = attr.s(); - else - onnx_attrs[attr.name()] = attr.s(); - } else if (attr.ints_size() > 0) { - if (attr.name() == "input_types") { - for (int k = 0; k < attr.ints_size(); ++k) { - input_types.push_back(static_cast(attr.ints(k))); - } - } else if (attr.name() == "output_types") { - for (int k = 0; k < attr.ints_size(); ++k) { - output_types.push_back(static_cast(attr.ints(k))); - } - } - } - } // for - ORT_ENFORCE(module != "", "PyOp module not specified"); - ORT_ENFORCE(class_name != "", "PyOp class name not specified"); - ORT_ENFORCE(!input_types.empty(), "PyOp node inputs not specified"); - ORT_ENFORCE(!output_types.empty(), "PyOp node outputs not specified"); - return new PyCustomOp(onnx_attrs, input_types, output_types, module, class_name, compute, log_func); -} -} // namespace onnxruntime diff --git a/onnxruntime/core/language_interop_ops/pyop/pyop.h b/onnxruntime/core/language_interop_ops/pyop/pyop.h deleted file mode 100644 index 049a247aec469..0000000000000 --- a/onnxruntime/core/language_interop_ops/pyop/pyop.h +++ /dev/null @@ -1,101 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT License. - -#pragma once - -#include "core/platform/env.h" -#define LOAD_PYOP_SYM(n, v, m) ORT_ENFORCE(Env::Default().GetSymbolFromLibrary(handle_, n, reinterpret_cast(&v)) == Status::OK(), m) - -#include "core/session/onnxruntime_cxx_api.h" -#include -#include -#include -#ifdef _WIN32 -#include -#else -#define HMODULE void* -#endif - -namespace ONNX_NAMESPACE { -class NodeProto; -} - -namespace onnxruntime { - -using OnnxTypes = std::vector; -using OnnxAttrs = std::unordered_map; -using PyOpLogFunc = std::function; - -class PyOpLibProxy { - public: - static PyOpLibProxy& GetInstance(); - void ReleaseInstance(void*); - bool InvokePythonFunc(void*, - const char*, - const std::vector&, - const std::vector&, - const std::vector>&, - std::vector>&, - std::vector&, - std::vector>&, - std::function); - const char* GetLastErrorMessage(std::string&); - void* NewInstance(const char*, const char*, const OnnxAttrs&); - bool Initialized() const { return initialized_; } - int32_t GetGil() const; - void PutGil(int32_t) const; - - private: - PyOpLibProxy(); - ~PyOpLibProxy(); - bool initialized_ = false; -}; - -struct PyCustomKernel { - PyCustomKernel(const OnnxAttrs& attrs, - const std::string& module, - const std::string& class_name, - const std::string& compute, - PyOpLogFunc logging_func); - ~PyCustomKernel(); - void GetOutputShape(OrtKernelContext*, size_t, OrtTensorTypeAndShapeInfo*); - void Compute(OrtKernelContext* context); - int32_t GetNumpyType(int32_t elem_type) const; - - private: - OnnxAttrs attrs_; - std::string module_; - std::string class_name_; - std::string compute_; - void* instance_ = nullptr; - PyOpLogFunc logging_func_; -}; - -struct PyCustomOp : Ort::CustomOpBase { - PyCustomOp( - const OnnxAttrs& attrs, - const OnnxTypes& inputs_type, - const OnnxTypes& outputs_type, - const std::string& module, - const std::string& class_name, - const std::string& compute = "compute", - PyOpLogFunc logging_func = [](const char*) {}); - void* CreateKernel(const OrtApi&, const OrtKernelInfo*) const; - const char* GetName() const; - size_t GetInputTypeCount() const; - ONNXTensorElementDataType GetInputType(size_t index) const; - size_t GetOutputTypeCount() const; - ONNXTensorElementDataType GetOutputType(size_t index) const; - - private: - OnnxAttrs attrs_; - OnnxTypes inputs_type_; - OnnxTypes outputs_type_; - std::string module_; - std::string class_name_; - std::string compute_; - PyOpLogFunc logging_func_; -}; // struct PyCustomOp - -PyCustomOp* LoadPyOp(const ONNX_NAMESPACE::NodeProto& node_proto, PyOpLogFunc log_func); -} // namespace onnxruntime diff --git a/tools/ci_build/build.py b/tools/ci_build/build.py index 1145509eef261..3e587e9b56e2e 100644 --- a/tools/ci_build/build.py +++ b/tools/ci_build/build.py @@ -605,11 +605,6 @@ def convert_arg_line_to_args(self, arg_line): parser.add_argument( "--enable_msvc_static_runtime", action="store_true", help="Enable static linking of MSVC runtimes." ) - parser.add_argument( - "--enable_language_interop_ops", - action="store_true", - help="Enable operator implemented in language other than cpp", - ) parser.add_argument( "--cmake_generator", choices=[ @@ -1053,7 +1048,6 @@ def generate_build_tree( else "OFF" ), "-Donnxruntime_REDUCED_OPS_BUILD=" + ("ON" if is_reduced_ops_build(args) else "OFF"), - "-Donnxruntime_ENABLE_LANGUAGE_INTEROP_OPS=" + ("ON" if args.enable_language_interop_ops else "OFF"), "-Donnxruntime_USE_DML=" + ("ON" if args.use_dml else "OFF"), "-Donnxruntime_USE_WINML=" + ("ON" if args.use_winml else "OFF"), "-Donnxruntime_BUILD_MS_EXPERIMENTAL_OPS=" + ("ON" if args.ms_experimental else "OFF"), diff --git a/tools/ci_build/github/azure-pipelines/mac-ci-pipeline.yml b/tools/ci_build/github/azure-pipelines/mac-ci-pipeline.yml index 5894631739ac8..8fa2b805dadc9 100644 --- a/tools/ci_build/github/azure-pipelines/mac-ci-pipeline.yml +++ b/tools/ci_build/github/azure-pipelines/mac-ci-pipeline.yml @@ -32,5 +32,5 @@ stages: parameters: AllowReleasedOpsetOnly: 0 BuildForAllArchs: false - AdditionalBuildFlags: --build_objc --enable_language_interop_ops --build_wheel --use_xnnpack + AdditionalBuildFlags: --build_objc --build_wheel --use_xnnpack WithCache: true From 8448f31d9042b6129886fb6cebfb23d6fe3293b8 Mon Sep 17 00:00:00 2001 From: Jian Chen Date: Wed, 19 Jun 2024 16:23:47 -0700 Subject: [PATCH 04/52] change is_pod tp is_trivial (#21071) ### Description change is_pod tp is_trivial ### Motivation and Context This is commonnly needed for both linux and win c++20 upgrade. is_trivial was introduced backed in C++11 --- onnxruntime/contrib_ops/cuda/math/cufft_plan_cache.h | 12 ++++++------ .../orttraining/training_ops/cuda/nn/conv_shared.h | 8 ++++---- .../orttraining/training_ops/rocm/nn/conv_grad.cc | 8 ++++---- 3 files changed, 14 insertions(+), 14 deletions(-) diff --git a/onnxruntime/contrib_ops/cuda/math/cufft_plan_cache.h b/onnxruntime/contrib_ops/cuda/math/cufft_plan_cache.h index 9cd96d5c62c94..3d21b12f9b55c 100644 --- a/onnxruntime/contrib_ops/cuda/math/cufft_plan_cache.h +++ b/onnxruntime/contrib_ops/cuda/math/cufft_plan_cache.h @@ -33,10 +33,10 @@ struct CufftPlanInfo { // see https://en.wikipedia.org/wiki/Fowler%E2%80%93Noll%E2%80%93Vo_hash_function template struct ParamsHash { - // Params must be a POD because we read out its memory - // contenst as char* when hashing + // Params must be a trivial type because we read out its memory + // contents as char* when hashing - static_assert(std::is_pod::value, "Params is not POD"); + static_assert(std::is_trivial::value, "Params is not a trivial type"); size_t operator()(const T& params) const { auto ptr = reinterpret_cast(¶ms); uint32_t value = 0x811C9DC5; @@ -50,10 +50,10 @@ struct ParamsHash { template struct ParamsEqual { - // Params must be a POD because we read out its memory - // contenst as char* when comparing + // Params must be a trivial type because we read out its memory + // contents as char* when comparing - static_assert(std::is_pod::value, "Params is not POD"); + static_assert(std::is_trivial::value, "Params is not a trivial type"); bool operator()(const T& a, const T& b) const { auto ptr1 = reinterpret_cast(&a); diff --git a/orttraining/orttraining/training_ops/cuda/nn/conv_shared.h b/orttraining/orttraining/training_ops/cuda/nn/conv_shared.h index 3fdb4306bfbbb..ec227c4641766 100644 --- a/orttraining/orttraining/training_ops/cuda/nn/conv_shared.h +++ b/orttraining/orttraining/training_ops/cuda/nn/conv_shared.h @@ -54,15 +54,15 @@ struct ConvArgs { }; struct ConvParamsHash { - // ConvParams must be a POD because we read out its memory constant as char* when hashing. - static_assert(std::is_pod::value, "ConvParams is not POD"); + // ConvParams must be a trivial type because we read out its memory contents as char* when hashing. + static_assert(std::is_trivial::value, "ConvParams is not a trivial type"); size_t operator()(const ConvParams& conv_params) const; }; struct ConvParamsEqual { - // ConvParams must be a POD because we read out its memory constant as char* when hashing. - static_assert(std::is_pod::value, "ConvParams is not POD"); + // ConvParams must be a trivial type because we read out its memory contents as char* when hashing. + static_assert(std::is_trivial::value, "ConvParams is not a trivial type"); bool operator()(const ConvParams& a, const ConvParams& b) const; }; diff --git a/orttraining/orttraining/training_ops/rocm/nn/conv_grad.cc b/orttraining/orttraining/training_ops/rocm/nn/conv_grad.cc index 9da0725274641..22fa5b6f55a5d 100644 --- a/orttraining/orttraining/training_ops/rocm/nn/conv_grad.cc +++ b/orttraining/orttraining/training_ops/rocm/nn/conv_grad.cc @@ -71,8 +71,8 @@ std::vector GetValidAlgorithms(const T_Perf* perf_results, int n_algo) { } struct ConvParamsHash { - // ConvParams must be a POD because we read out its memory constant as char* when hashing. - static_assert(std::is_pod::value, "ConvParams is not POD"); + // ConvParams must be a trivial type because we read out its memory contents as char* when hashing. + static_assert(std::is_trivial::value, "ConvParams is not a trivial type"); size_t operator()(const ConvParams& conv_params) const { auto ptr = reinterpret_cast(&conv_params); uint32_t value = 0x811C9DC5; @@ -85,8 +85,8 @@ struct ConvParamsHash { }; struct ConvParamsEqual { - // ConvParams must be a POD because we read out its memory constant as char* when hashing. - static_assert(std::is_pod::value, "ConvParams is not POD"); + // ConvParams must be a trivial type because we read out its memory contents as char* when hashing. + static_assert(std::is_trivial::value, "ConvParams is not a trivial type"); bool operator()(const ConvParams& a, const ConvParams& b) const { auto ptr1 = reinterpret_cast(&a); auto ptr2 = reinterpret_cast(&b); From 6817b013b9d19b8e557da3ca0f593a6cd53c2266 Mon Sep 17 00:00:00 2001 From: Jing Fang <126209182+fajin-corp@users.noreply.github.com> Date: Wed, 19 Jun 2024 17:15:45 -0700 Subject: [PATCH 05/52] [MLAS] add q4 quantize and transpose kernel to support MatMulNBits QDQ fuse (#21054) ### Description 1. added kernel to quantize matmul B tensor to q4, and store in the same shape as original tensor. scales and zero points are calculated as well. scales and zero points have the same shape. 2. added kernel to transpose q4 B tensor to B tensor in MatMulNBits. Scales and zero points are transposed as well. #### Benchmark <1024 x 4096 input, 64 quant block, 8 threads>: - quantize: 23035923 ns - transpose: 718635 ns <1024 x 4095 input, 64 quant block, 8 threads>: - quantize: 26759319 ns - transpose: 1279064 ns ### Motivation and Context The MatMulNbits tool chain current only supports converting a MatMul op direct to MatMulNBits op. MatMulNbits op is not an ONNX standard op. Therefore, we need the tool chain to support converting MatMul to Q/DQ format, and later in the transform step converts DQ + MatMul to MatMulNBits. The tensors stored in DQ are the quantized constants and will be stored in the MatMulNBits. --- onnxruntime/core/mlas/inc/mlas_q4.h | 67 ++ onnxruntime/core/mlas/lib/q4_dq.cpp | 775 +++++++++++++++++- onnxruntime/core/util/qmath.h | 320 +++++--- .../python/onnxruntime_pybind_quant.cc | 33 + onnxruntime/test/mlas/bench/bench_q4dq.cpp | 118 +++ .../test/mlas/unittest/test_blockq4.cpp | 51 ++ .../test/onnx/microbenchmark/quantize.cc | 40 +- 7 files changed, 1258 insertions(+), 146 deletions(-) create mode 100644 onnxruntime/test/mlas/bench/bench_q4dq.cpp diff --git a/onnxruntime/core/mlas/inc/mlas_q4.h b/onnxruntime/core/mlas/inc/mlas_q4.h index 316344ad8c214..898fb23cf3e4f 100644 --- a/onnxruntime/core/mlas/inc/mlas_q4.h +++ b/onnxruntime/core/mlas/inc/mlas_q4.h @@ -358,3 +358,70 @@ MlasDequantizeBlockwise( int columns, MLAS_THREADPOOL* thread_pool ); + +/** + * @brief Blockwise 2 bits or 4 bits quantization. After quantization, the weights and zero points + * are packed row-wise. In terms of the qbits type, dst and src have the same shape, and + * scales and zero_points have the same shape. + * columns must be multiple of 8 / qbits. + * @tparam Tin + * @tparam qbits number of bits used for quantization, 2 or 4 + * @param src points to the floating point matrix, to be quantized, row major shape [rows, columns] + * @param scales points to the scales matrix, row major + * @param zero_points points to the zero_points matrix, row major + * @param dst points to the quantized matrix, shape [rows, columns] row major in qbits type. + * In uint8_t type, shape is [rows, columns * qbits / 8]. + * @param columnwise true when quantize elements in a column, false when quantize elements in a row. + * @param rows + * @param columns + * @param quant_block_size number of elements in a quantize block + * @param thread_pool + */ +template +void +MlasQDQQuantizeBlockwise( + const Tin* src, + Tin* scales, + uint8_t* zero_points, + uint8_t* dst, + bool columnwise, + int rows, + int columns, + int quant_block_size, + MLAS_THREADPOOL* thread_pool +); + +/** + * @brief Transpose blockwise quantized tensors. The src tensors are row major. src weights and zero + * points are packed row-wise. The dst tensors are column major. dst weights and zero points + * are packed column-wise. + * @tparam Tin + * @tparam qbits number of bits used for quantization, 2 or 4 + * @param src_weights points to the quantized matrix, row major, shape [rows, columns] in qbits type. + * In uint8_t type, shape is [rows, columns * qbits / 8]. + * @param src_scales points to the scales matrix, row major + * @param src_zero_points points to the zero_points matrix, row major. Packed row-wise. + * @param dst_weights points to the quantized matrix, column major. Packed column-wise. + * @param dst_scales points to the scales matrix, column major + * @param dst_zero_points points to the zero_points matrix, column major. Packed column-wise. + * @param columnwise true when quantize elements in a column, false when quantize elements in a row. + * @param rows + * @param columns + * @param quant_block_size number of elements in a quantize block + * @param thread_pool + */ +template +void +MlasQDQTransposeBlockwiseQuantized( + const uint8_t* src_weights, + const Tin* src_scales, + const uint8_t* src_zero_points, + uint8_t* dst_weights, + Tin* dst_scales, + uint8_t* dst_zero_points, + bool columnwise, + int rows, + int columns, + int quant_block_size, + MLAS_THREADPOOL* thread_pool +); diff --git a/onnxruntime/core/mlas/lib/q4_dq.cpp b/onnxruntime/core/mlas/lib/q4_dq.cpp index b5784ecb56d01..62fe58ca333de 100644 --- a/onnxruntime/core/mlas/lib/q4_dq.cpp +++ b/onnxruntime/core/mlas/lib/q4_dq.cpp @@ -638,6 +638,669 @@ struct BlockwiseQuantizer { } }; +/** + * @brief Blockwise quantization methods for QDQ format. Input tensor is quantized along column + * or row. Scales and zeros are calculated. Based on qbits, consecutive quantized elements + * in memory are packed together, which means the packing is along the row. Quantized data + * are stored in row major, so the output tensor reserves same shape, in terms of qbits type, + * as the input tensor. + * @tparam Tin source data type, e.g. fp32/fp16 + * @tparam qbits number of bits in each quantized element + */ +template +struct BlockwiseQDQQuantizer; + +template +struct BlockwiseQDQQuantizer { + static MLAS_FORCEINLINE uint8_t GetElem(uint8_t val, int32_t idx) + { + return (val >> (idx << 2)) & 0xF; + } + + static MLAS_FORCEINLINE uint8_t SetElem(uint8_t val, int32_t idx, uint8_t dst) + { + auto shift = idx << 2; + return ((val & 0xF) << shift) | (dst & (~(0xF << shift))); + } + + static MLAS_FORCEINLINE uint8_t Pack(uint8_t v0, uint8_t v1) + { + return (v0 & 0xF) | ((v1 & 0xF) << 4); + } + + // If src is row major, then dst is column major. Transpose: + // | src0: low 4 bit | src0: high 4 bit | + // | src1: low 4 bit | src1: high 4 bit | + // --> + // | dst0: low 4 bit | dst1: low 4 bit | + // | dst0: high 4 bit| dst1: high 4 bit | + // If src is column major, then dst is row major. Transpose: + // | src0: low 4 bit | src1: low 4 bit | + // | src0: high 4 bit| src1: high 4 bit | + // --> + // | dst0: low 4 bit | dst0: high 4 bit | + // | dst1: low 4 bit | dst1: high 4 bit | + static MLAS_FORCEINLINE void Transpose(uint8_t src0, uint8_t src1, uint8_t& dst0, uint8_t& dst1) + { + dst0 = (src0 & 0xF) | ((src1 & 0xF) << 4); + dst1 = ((src0 & 0xF0) >> 4) | (src1 & 0xF0); + } + + static MLAS_FORCEINLINE uint8_t QuantizeV(Tin src, float reciprocal_scale, uint8_t zero_point) + { + return static_cast( + std::clamp( + static_cast( + std::roundf(static_cast(src) * reciprocal_scale) + ) + static_cast(zero_point), + 0, + BitsTraits<4>::kMax + ) + ); + } + + /** + * @brief Quantize a matrix shape [rows, columns] row-wise. Scales and zero points are calculated. + * Quantized data are packed row-wise based on qbits. Quantized data are stored in row + * major, so the output tensor reserves the shape, in terms output type. + * Thread block is [1, quant_block_size * 2]. + * @param src the source matrix, row major: [rows * columns] + * @param scales the scales of quantized blocks, row major layout with shape: + * [rows * ceil(columns / quant_block_size)] + * @param zero_points the zero points of quantized blocks, packed. Same shape as scales + * in terms of output type. In terms of uint8_t, the shape is: + * [ceil(rows * ceil(columns / quant_block_size) * qbits / 8)] + * @param dst the quantized weights, row major: [rows * columns] in terms of + * output type. In terms of uint8_t, the shape is: [ceil(rows * columns * qbits / 8] + * @param rows number of rows in the source matrix + * @param columns number of columns in the source matrix, must satisfy + * ceil(columns / quant_block_size) % 2 == 0, so in each thread block, + * zero points are packed into one byte. + * @param quant_block_size number of elements quantized together. + * @param thread_pool thread pool for parallel processing + */ + static void QuantizeRowWise( + const Tin* src, + Tin* scales, + uint8_t* zero_points, + uint8_t* dst, + int32_t rows, + int32_t columns, + int32_t quant_block_size, + MLAS_THREADPOOL* thread_pool + ) + { + MLAS_UNREFERENCED_PARAMETER(src); + MLAS_UNREFERENCED_PARAMETER(scales); + MLAS_UNREFERENCED_PARAMETER(zero_points); + MLAS_UNREFERENCED_PARAMETER(dst); + MLAS_UNREFERENCED_PARAMETER(rows); + MLAS_UNREFERENCED_PARAMETER(columns); + MLAS_UNREFERENCED_PARAMETER(quant_block_size); + MLAS_UNREFERENCED_PARAMETER(thread_pool); + ORT_THROW("BlockwiseQDQQuantizer::BlockwiseQDQQuantizer is not implemented"); + } + + /** + * @brief Quantize a matrix shape [rows, columns] column-wise. Scales and zero points are calculated. + * Quantized data are packed row-wise based on qbits. Quantized data are stored in row major + * so the output tensor reserves the shape, in terms output type. + * @param src the source matrix, row major: [rows * columns] + * @param scales the scales of quantized blocks, row major with shape: + * [ceil(rows/quant_block_size) * columns] + * @param zero_points the zero points of quantized blocks, packed. Same shape as scales in terms + * of output type. In uint8_t, the shape is: + * [ceil(columns * ceil(rows / quant_block_size) * qbits / 8)] + * @param dst the quantized weights, row major: [rows * columns] in terms of output type. + * In uint8_t, the shape is: [ceil(rows * columns * qbits / 8] + * @param rows number of rows in the source matrix + * @param columns number of columns in the source matrix. + * @param quant_block_size number of rows/columns quantized together + * @param thread_pool thread pool for parallel processing + */ + static void QuantizeColumnWise( + const Tin* src, + Tin* scales, + uint8_t* zero_points, + uint8_t* dst, + int32_t rows, + int32_t columns, + int32_t quant_block_size, + MLAS_THREADPOOL* thread_pool + ) + { + // Must avoid multiple thread write to a single byte, which means the starting index + // of a thread block must be even. To achieve that, we need to customize the thread + // block size based on the parity of columns. + if (columns & 1) { + QuantizeColumnWisePackUnaligned( + src, scales, zero_points, dst, rows, columns, quant_block_size, thread_pool + ); + } else { + QuantizeColumnWisePackAligned( + src, scales, zero_points, dst, rows, columns, quant_block_size, thread_pool + ); + } + } + + + /** + * @brief Transpose quantized tensors, which has been column-wise quantized, for use in MatMulNbits. + * Since both src tensor and dst tensor are packed, it's not needed to consider sign + * during the unpacking/packing in transpose. + * @param src_weights The quantized weights, row major: [rows, columns] in qbits type. + * In uint8_t, size of [ceil(rows * columns * qbits / 8)]. + * @param src_scales [ceil(rows / quant_block_size), columns] + * @param src_zero_points [ceil(rows / quant_block_size), columns] in qbits type. In uint8_t, size of + * [ceil(ceil(rows / quant_block_size) * columns * qbits / 8 )]. + * @param dst_weights the transposed quantized weights, column major. In uint8_t, the shape is + * [columns, ceil(rows / quant_block_size), ceil(quant_block_size * qbits / 8)] + * @param dst_scales [columns, ceil(rows / quant_block_size)] + * @param dst_zero_points [columns, ceil(ceil(rows / quant_block_size) * qbits / 8)] in uint8_t. + * @param rows number of src rows in qbits type. + * @param columns number of src columns in qbits type. + * @param quant_block_size number of elements quantized together + * @param thread_pool thread pool for parallel processing + */ + static void TransposeColumnWiseQuantized( + const uint8_t* src_weights, + const Tin* src_scales, + const uint8_t* src_zero_points, + uint8_t* dst_weights, + Tin* dst_scales, + uint8_t* dst_zero_points, + int32_t rows, + int32_t columns, + int32_t quant_block_size, + MLAS_THREADPOOL* thread_pool + ) + { + // Must avoid multiple thread write to a single byte, which means the starting index + // of a thread block must be even. To achieve that, we need to customize the thread + // block size based on the parity of columns. + if (columns & 1) { + TransposeColumnWiseQuantizedPackUnaligned( + src_weights, src_scales, src_zero_points, + dst_weights, dst_scales, dst_zero_points, + rows, columns, quant_block_size, thread_pool + ); + } else { + TransposeColumnWiseQuantizedPackAligned( + src_weights, src_scales, src_zero_points, + dst_weights, dst_scales, dst_zero_points, + rows, columns, quant_block_size, thread_pool + ); + } + } + +private: + static void QuantizeColumnWisePackAligned( + const Tin* src, + Tin* scales, + uint8_t* zero_points, + uint8_t* dst, + int32_t rows, + int32_t columns, + int32_t quant_block_size, + MLAS_THREADPOOL* thread_pool + ) + { + ORT_ENFORCE(columns % 2 == 0, "Columns must be multiple of 2."); + // Thread block is [quant_block_size, thread_blk_size]. thread_blk_size % 2 == 0. + constexpr int32_t thread_blk_size = 128; + const auto num_row_thread_blk = (rows + quant_block_size - 1) / quant_block_size; + const auto num_col_thread_blk = (columns + thread_blk_size - 1) / thread_blk_size; + const auto num_thread_blk = num_row_thread_blk * num_col_thread_blk; + constexpr auto minf = std::numeric_limits::lowest(); + constexpr auto maxf = std::numeric_limits::max(); + + MlasTryBatchParallel( + thread_pool, static_cast(num_thread_blk), + [&](ptrdiff_t thread_blk_idx) { + // !!warning!!: buffering the whole thread block + constexpr int32_t buffer_size = 128; + ORT_ENFORCE(buffer_size == thread_blk_size, "buffer size must be equal to thread block size."); + float reciprocal_scale_t[buffer_size]; + uint8_t zp_t[buffer_size]; + float vmin_t[buffer_size]; + float vmax_t[buffer_size]; + + const int32_t row_thread_blk_idx = static_cast(thread_blk_idx / num_col_thread_blk); + const int32_t col_thread_blk_idx = static_cast(thread_blk_idx % num_col_thread_blk); + const int32_t row_idx = row_thread_blk_idx * quant_block_size; + const int32_t col_idx = col_thread_blk_idx * buffer_size; + const int32_t row_size = std::min(quant_block_size, rows - row_idx); + const int32_t col_size = std::min(buffer_size, columns - col_idx); + // input_idx, scale_idx, col_size are aligned to 2 + auto input_idx = row_idx * columns + col_idx; + auto scale_idx = row_thread_blk_idx * columns + col_idx; + + Tin scale0_tt, scale1_tt; + uint8_t v0_tt, v1_tt; + + std::fill_n(vmin_t, buffer_size, maxf); + std::fill_n(vmax_t, buffer_size, minf); + + // calculate min/max + for (int32_t j = 0, input_idx_t = input_idx; j < row_size; ++j, input_idx_t += columns) { + // TODO(fajin): use SIMD + for (int32_t i = 0; i < col_size; i += 2) { + auto v0 = static_cast(src[input_idx_t + i]); + auto v1 = static_cast(src[input_idx_t + i + 1]); + vmin_t[i] = std::min(vmin_t[i], v0); + vmax_t[i] = std::max(vmax_t[i], v0); + vmin_t[i + 1] = std::min(vmin_t[i + 1], v1); + vmax_t[i + 1] = std::max(vmax_t[i + 1], v1); + } + } + + // calculate scale and zero point, and store + for (int32_t i = 0; i < col_size; i += 2) { + v0_tt = v1_tt = BitsTraits<4>::kMid; + + if (zero_points) { + range2scalezp(vmin_t[i], vmax_t[i], scale0_tt, v0_tt); + range2scalezp(vmin_t[i + 1], vmax_t[i + 1], scale1_tt, v1_tt); + zero_points[(scale_idx + i) >> 1] = Pack(v0_tt, v1_tt); + } else { + range2scale(vmin_t[i], vmax_t[i], scale0_tt); + range2scale(vmin_t[i + 1], vmax_t[i + 1], scale1_tt); + } + + scales[scale_idx + i] = scale0_tt; + scales[scale_idx + i + 1] = scale1_tt; + + float scalef0 = static_cast(scale0_tt); + reciprocal_scale_t[i] = scalef0 ? 1.0f / scalef0 : 0.0f; + zp_t[i] = v0_tt; + + float scalef1 = static_cast(scale1_tt); + reciprocal_scale_t[i + 1] = scalef1 ? 1.0f / scalef1 : 0.0f; + zp_t[i + 1] = v1_tt; + } + + // quantize and pack + for (int32_t j = 0, input_idx_t = input_idx; j < row_size; ++j, input_idx_t += columns) { + // TODO(fajin): use SIMD + for (int32_t i = 0; i < col_size; i += 2) { + v0_tt = QuantizeV(src[input_idx_t + i], reciprocal_scale_t[i], zp_t[i]); + v1_tt = QuantizeV(src[input_idx_t + i + 1], reciprocal_scale_t[i + 1], zp_t[i + 1]); + dst[(input_idx_t + i) >> 1] = Pack(v0_tt, v1_tt); + } + } + } + ); + } + + static void QuantizeColumnWisePackUnaligned( + const Tin* src, + Tin* scales, + uint8_t* zero_points, + uint8_t* dst, + int32_t rows, + int32_t columns, + int32_t quant_block_size, + MLAS_THREADPOOL* thread_pool + ) + { + // Thread block is [quant_block_size * 2, columns], so the packed bytes do not cross threads. + constexpr auto minf = std::numeric_limits::lowest(); + constexpr auto maxf = std::numeric_limits::max(); + auto row_thread_blk_size = quant_block_size * 2; + auto num_row_thread_blk = (rows + row_thread_blk_size - 1) / (row_thread_blk_size); + + MlasTryBatchParallel( + thread_pool, static_cast(num_row_thread_blk), + [&](ptrdiff_t thread_blk_idx) { + constexpr int32_t buffer_size = 128; + float reciprocal_scale_t[buffer_size]; + uint8_t zp_t[buffer_size]; + float vmin_t[buffer_size]; + float vmax_t[buffer_size]; + + auto row_thread_blk_idx = static_cast(thread_blk_idx); + int32_t row_idx = row_thread_blk_idx * row_thread_blk_size; + int32_t row_idx_end = std::min(row_thread_blk_size + row_idx, rows); + auto input_idx = row_idx * columns; + auto scale_idx = row_thread_blk_idx * 2 * columns; + Tin scale0_tt, scale1_tt; + uint8_t v0_tt, v1_tt; + + for (; row_idx < row_idx_end; row_idx += quant_block_size) { + // per quant block row + auto quant_row_size = std::min(quant_block_size, row_idx_end - row_idx); + auto input_buffer_idx = input_idx; + auto scale_buffer_idx = scale_idx; + for (int32_t buffer_idx = 0; buffer_idx < columns; buffer_idx += buffer_size) { + // per buffer column + auto buffer_col_size = std::min(buffer_size, columns - buffer_idx); + + std::fill_n(vmin_t, buffer_size, maxf); + std::fill_n(vmax_t, buffer_size, minf); + // calculate min/max of [quant block, buffer] + auto input_idx_t = input_buffer_idx; + for (int32_t j = 0; j < quant_row_size; ++j, input_idx_t += columns) { + // TODO(fajin): use SIMD + for (int32_t i = 0; i < buffer_col_size; ++i) { + auto v = static_cast(src[input_idx_t + i]); + vmin_t[i] = std::min(vmin_t[i], v); + vmax_t[i] = std::max(vmax_t[i], v); + } + } + + // calculate scale and zero point + auto scale_buffer_idx_end = scale_buffer_idx + buffer_col_size; + int32_t col_idx = 0; + // leading unailgned zero points + if (scale_buffer_idx & 1) { + v0_tt = BitsTraits<4>::kMid; + if (zero_points) { + range2scalezp(vmin_t[0], vmax_t[0], scale0_tt, v0_tt); + zero_points[scale_buffer_idx >> 1] = SetElem( + v0_tt, 1, zero_points[scale_buffer_idx >> 1] + ); + } else { + range2scale(vmin_t[0], vmax_t[0], scale0_tt); + } + + scales[scale_buffer_idx] = scale0_tt; + + float scalef = static_cast(scale0_tt); + reciprocal_scale_t[0] = scalef ? 1.0f / scalef : 0.0f; + zp_t[0] = v0_tt; + + ++col_idx; + ++scale_buffer_idx; + } + // aligned zero points + for (; scale_buffer_idx < scale_buffer_idx_end - 1; col_idx += 2, scale_buffer_idx += 2) { + v0_tt = v1_tt = BitsTraits<4>::kMid; + if (zero_points) { + range2scalezp(vmin_t[col_idx], vmax_t[col_idx], scale0_tt, v0_tt); + range2scalezp(vmin_t[col_idx + 1], vmax_t[col_idx + 1], scale1_tt, v1_tt); + zero_points[scale_buffer_idx >> 1] = Pack(v0_tt, v1_tt); + } else { + range2scale(vmin_t[col_idx], vmax_t[col_idx], scale0_tt); + range2scale(vmin_t[col_idx + 1], vmax_t[col_idx + 1], scale1_tt); + } + + scales[scale_buffer_idx] = scale0_tt; + scales[scale_buffer_idx + 1] = scale1_tt; + + float scalef0 = static_cast(scale0_tt); + reciprocal_scale_t[col_idx] = scalef0 ? 1.0f / scalef0 : 0.0f; + zp_t[col_idx] = v0_tt; + + float scalef1 = static_cast(scale1_tt); + reciprocal_scale_t[col_idx + 1] = scalef1 ? 1.0f / scalef1 : 0.0f; + zp_t[col_idx + 1] = v1_tt; + } + // tailing unaligned elements + if (scale_buffer_idx < scale_buffer_idx_end) { + v0_tt = BitsTraits<4>::kMid; + if (zero_points) { + range2scalezp(vmin_t[col_idx], vmax_t[col_idx], scale0_tt, v0_tt); + zero_points[scale_buffer_idx >> 1] = SetElem( + v0_tt, 0, zero_points[scale_buffer_idx >> 1] + ); + } else { + range2scale(vmin_t[col_idx], vmax_t[col_idx], scale0_tt); + } + + scales[scale_buffer_idx] = scale0_tt; + + float scalef = static_cast(scale0_tt); + reciprocal_scale_t[col_idx] = scalef ? 1.0f / scalef : 0.0f; + zp_t[col_idx] = v0_tt; + + ++scale_buffer_idx; + } + + // quantize and pack + input_idx_t = input_buffer_idx; + for (int32_t j = 0; j < quant_row_size; ++j, input_idx_t += columns) { + auto input_idx_t_start = input_idx_t; + auto input_idx_t_end = input_idx_t + buffer_col_size; + col_idx = 0; + // leading unaligned output + if (input_idx_t_start & 1) { + v1_tt = QuantizeV(src[input_idx_t_start], reciprocal_scale_t[col_idx], zp_t[col_idx]); + dst[input_idx_t_start >> 1] = SetElem(v1_tt, 1, dst[input_idx_t_start >> 1]); + + ++col_idx; + ++input_idx_t_start; + } + // aligned output + // TODO(fajin): use SIMD + for (; input_idx_t_start < input_idx_t_end - 1; col_idx += 2, input_idx_t_start += 2) { + v0_tt = QuantizeV(src[input_idx_t_start], reciprocal_scale_t[col_idx], zp_t[col_idx]); + v1_tt = QuantizeV( + src[input_idx_t_start + 1], reciprocal_scale_t[col_idx + 1], zp_t[col_idx + 1] + ); + + dst[input_idx_t_start >> 1] = Pack(v0_tt, v1_tt); + } + // tailing unaligned output + if (input_idx_t_start < input_idx_t_end) { + v0_tt = QuantizeV(src[input_idx_t_start], reciprocal_scale_t[col_idx], zp_t[col_idx]); + dst[input_idx_t_start >> 1] = SetElem(v0_tt, 0, dst[input_idx_t_start >> 1]); + } + } + + input_buffer_idx += buffer_size; + } + + input_idx += quant_block_size * columns; + scale_idx += columns; + } + } + ); + } + + static void TransposeColumnWiseQuantizedPackAligned( + const uint8_t* src_weights, // [rows, columns / 2] + const Tin* src_scales, // [ceil(rows / quant_block_size), columns] + const uint8_t* src_zero_points, // [ceil(rows / quant_block_size), columns / 2] + uint8_t* dst_weights, // [columns, ceil(rows / quant_block_size), ceil(quant_block_size / 2)] + Tin* dst_scales, // [columns, ceil(rows / quant_block_size)] + uint8_t* dst_zero_points, // [columns, ceil(ceil(rows / quant_block_size) / 2)] + int32_t rows, + int32_t columns, + int32_t quant_block_size, + MLAS_THREADPOOL* thread_pool + ) + { + ORT_ENFORCE(columns % 2 == 0, "Columns must be multiple of 2"); + + auto row_quant_blk_num = (rows + quant_block_size - 1) / quant_block_size; + auto dst_bytes_per_quant_blk = (quant_block_size * 4 + 7) / 8; + // number of rows in transposed dst + auto dstT_num_row = row_quant_blk_num * dst_bytes_per_quant_blk; + auto packed_col_size = columns / 2; + + // weight transpose thread block is [dst_bytes_per_quant_blk, 2] on dst_Transpose. + // Map to src it is [quant_block_size, 1]. Both in uint8_t. + auto num_thread_blk = row_quant_blk_num * packed_col_size; + MlasTryBatchParallel( + thread_pool, static_cast(num_thread_blk), + [&](ptrdiff_t thread_blk_idx) { + uint8_t src0_t, src1_t; + uint8_t dst0_t, dst1_t; + + auto row_thread_blk_idx = static_cast(thread_blk_idx / packed_col_size); + auto col_thread_blk_idx = static_cast(thread_blk_idx % packed_col_size); + + auto dstT_row_idx = row_thread_blk_idx * dst_bytes_per_quant_blk; + auto dstT_col_idx = col_thread_blk_idx * 2; + auto dst_idx = dstT_col_idx * dstT_num_row + dstT_row_idx; + + auto src_row_idx = row_thread_blk_idx * quant_block_size; + auto src_row_end_idx = std::min(src_row_idx + quant_block_size, rows); + auto src_col_idx = col_thread_blk_idx; + auto src_idx = src_row_idx * packed_col_size + src_col_idx; + auto src_end_idx = src_row_end_idx * packed_col_size + src_col_idx; + + for (; src_idx < src_end_idx - packed_col_size; ++dst_idx) { + src0_t = src_weights[src_idx]; + src1_t = src_weights[src_idx + packed_col_size]; + src_idx += packed_col_size + packed_col_size; + Transpose(src0_t, src1_t, dst0_t, dst1_t); + dst_weights[dst_idx] = dst0_t; + dst_weights[dst_idx + dstT_num_row] = dst1_t; + } + + if (src_idx < src_end_idx) { + src0_t = src_weights[src_idx]; + src1_t = 0; + Transpose(src0_t, src1_t, dst0_t, dst1_t); + dst_weights[dst_idx] = dst0_t; + dst_weights[dst_idx + dstT_num_row] = dst1_t; + } + } + ); + + // Transpose scales. Thread block is [row_quant_blk_num, 1] on dst_Transpose. + MlasTryBatchParallel( + thread_pool, static_cast(columns), + [&](ptrdiff_t thread_blk_idx) { + auto col_thread_blk_idx = static_cast(thread_blk_idx); + auto src_idx = col_thread_blk_idx; + auto dst_idx = col_thread_blk_idx * row_quant_blk_num; + for (int32_t i = 0; i < row_quant_blk_num; ++i, ++dst_idx, src_idx += columns) { + dst_scales[dst_idx] = src_scales[src_idx]; + } + } + ); + + if (src_zero_points) { + // Transpose zero points. Thread block is [ceil(row_quant_blk_num / 2), 2] + // on dst_Transpose. Map to src it is [row_quant_blk_num, 1]. Both in uint8_t. + auto dst_zp_row_num = (row_quant_blk_num + 1) / 2; + MlasTryBatchParallel( + thread_pool, static_cast(packed_col_size), + [&](ptrdiff_t thread_blk_idx) { + uint8_t src0_t, src1_t; + uint8_t dst0_t, dst1_t; + + auto col_thread_blk_idx = static_cast(thread_blk_idx); + auto src_idx = col_thread_blk_idx; + auto src_end_idx = row_quant_blk_num * packed_col_size + col_thread_blk_idx; + auto dst_idx = col_thread_blk_idx * 2 * dst_zp_row_num; + + for (; src_idx < src_end_idx - packed_col_size; ++dst_idx) { + src0_t = src_zero_points[src_idx]; + src1_t = src_zero_points[src_idx + packed_col_size]; + Transpose(src0_t, src1_t, dst0_t, dst1_t); + dst_zero_points[dst_idx] = dst0_t; + dst_zero_points[dst_idx + dst_zp_row_num] = dst1_t; + src_idx += packed_col_size + packed_col_size; + } + + if (src_idx < src_end_idx) { + src0_t = src_zero_points[src_idx]; + src1_t = 0; + Transpose(src0_t, src1_t, dst0_t, dst1_t); + dst_zero_points[dst_idx] = dst0_t; + dst_zero_points[dst_idx + dst_zp_row_num] = dst1_t; + } + } + ); + } + } + + static void TransposeColumnWiseQuantizedPackUnaligned( + const uint8_t* src_weights, // size of [ceil(rows * columns / 2)] + const Tin* src_scales, // [ceil(rows / quant_block_size), columns] + const uint8_t* src_zero_points, // size of [ceil(ceil(rows / quant_block_size) * columns / 2)] + uint8_t *dst_weights, // [columns, ceil(rows / quant_block_size), ceil(quant_block_size / 2)] + Tin* dst_scales, // [columns, ceil(rows / quant_block_size)] + uint8_t* dst_zero_points, // [columns, ceil(ceil(rows / quant_block_size) / 2)] + int32_t rows, + int32_t columns, + int32_t quant_block_size, + MLAS_THREADPOOL* thread_pool) + { + auto row_quant_blk_num = (rows + quant_block_size - 1) / quant_block_size; + auto dst_bytes_per_quant_blk = (quant_block_size * 4 + 7) / 8; + // number of rows in transposed dst + auto dstT_num_row = row_quant_blk_num * dst_bytes_per_quant_blk; + + // weight transpose thread block is [dst_bytes_per_quant_blk, 1] on dst_Transpose in uint8_t. + // Map to src it is [quant_block_size, 1] in int4. + auto num_thread_blk = row_quant_blk_num * columns; + MlasTryBatchParallel( + thread_pool, static_cast(num_thread_blk), + [&](ptrdiff_t thread_blk_idx) { + uint8_t src0_t, src1_t; + + auto row_thread_blk_idx = static_cast(thread_blk_idx / columns); + auto col_thread_blk_idx = static_cast(thread_blk_idx % columns); + + auto dstT_row_idx = row_thread_blk_idx * dst_bytes_per_quant_blk; + auto dst_idx = col_thread_blk_idx * dstT_num_row + dstT_row_idx; + + auto src_row_idx = row_thread_blk_idx * quant_block_size; + auto src_row_end_idx = std::min(src_row_idx + quant_block_size, rows); + auto src_idx = src_row_idx * columns + col_thread_blk_idx; + auto src_end_idx = src_row_end_idx * columns + col_thread_blk_idx; + + for (; src_idx < src_end_idx - columns; ++dst_idx) { + src0_t = GetElem(src_weights[src_idx >> 1], src_idx & 1); + src1_t = GetElem(src_weights[(src_idx + columns) >> 1], (src_idx + columns) & 1); + dst_weights[dst_idx] = (src0_t & 0xf) | ((src1_t & 0xf) << 4); + src_idx += columns + columns; + } + + if (src_idx < src_end_idx) { + src0_t = GetElem(src_weights[src_idx >> 1], src_idx & 1); + dst_weights[dst_idx] = src0_t & 0xf; + } + } + ); + + // Transpose scales. Thread block is [row_quant_blk_num, 1] on dst_Transpose. + MlasTryBatchParallel( + thread_pool, static_cast(columns), + [&](ptrdiff_t thread_blk_idx) { + auto col_thread_blk_idx = static_cast(thread_blk_idx); + auto src_idx = col_thread_blk_idx; + auto dst_idx = col_thread_blk_idx * row_quant_blk_num; + for (int32_t i = 0; i < row_quant_blk_num; ++i, ++dst_idx, src_idx += columns) { + dst_scales[dst_idx] = src_scales[src_idx]; + } + } + ); + + if (src_zero_points) { + // Transpose zero points. Thread block is [ceil(row_quant_blk_num / 2), 1] on dst_Transpose in uint8_t. + // Map to src it is [row_quant_blk_num, 1] in int4. + auto dst_zp_row_num = (row_quant_blk_num + 1) / 2; + MlasTryBatchParallel( + thread_pool, static_cast(columns), + [&](ptrdiff_t thread_blk_idx) { + uint8_t src0_t, src1_t; + + auto col_thread_blk_idx = static_cast(thread_blk_idx); + auto src_idx = col_thread_blk_idx; + auto src_end_idx = row_quant_blk_num * columns + col_thread_blk_idx; + auto dst_idx = col_thread_blk_idx * dst_zp_row_num; + + for (; src_idx < src_end_idx - columns; ++dst_idx) { + src0_t = GetElem(src_zero_points[src_idx >> 1], src_idx & 1); + src1_t = GetElem(src_zero_points[(src_idx + columns) >> 1], (src_idx + columns) & 1); + dst_zero_points[dst_idx] = (src0_t & 0xf) | ((src1_t & 0xf) << 4); + src_idx += columns + columns; + } + + if (src_idx < src_end_idx) { + src0_t = GetElem(src_zero_points[src_idx >> 1], src_idx & 1); + dst_zero_points[dst_idx] = src0_t & 0xf; + } + } + ); + } + } +}; template void @@ -1068,8 +1731,7 @@ MlasDequantizeBlockwise( } } -template -void +template void MlasDequantizeBlockwise( float* dst, const uint8_t* src, @@ -1080,4 +1742,111 @@ MlasDequantizeBlockwise( int rows, int columns, MLAS_THREADPOOL* thread_pool - ); +); + +template +void +MlasQDQQuantizeBlockwise( + const Tin* src, + Tin* scales, + uint8_t* zero_points, + uint8_t* dst, + bool columnwise, + int rows, + int columns, + int quant_block_size, + MLAS_THREADPOOL* thread_pool +) +{ + if (columnwise) { + BlockwiseQDQQuantizer::QuantizeColumnWise( + src, scales, zero_points, dst, rows, columns, quant_block_size, thread_pool + ); + } else { + BlockwiseQDQQuantizer::QuantizeRowWise( + src, scales, zero_points, dst, rows, columns, quant_block_size, thread_pool + ); + } +} + +template void +MlasQDQQuantizeBlockwise( + const float* src, + float* scales, + uint8_t* zero_points, + uint8_t* dst, + bool columnwise, + int rows, + int columns, + int quant_block_size, + MLAS_THREADPOOL* thread_pool +); + +template void +MlasQDQQuantizeBlockwise( + const MLAS_FP16* src, + MLAS_FP16* scales, + uint8_t* zero_points, + uint8_t* dst, + bool columnwise, + int rows, + int columns, + int quant_block_size, + MLAS_THREADPOOL* thread_pool +); + +template +void +MlasQDQTransposeBlockwiseQuantized( + const uint8_t* src_weights, + const Tin* src_scales, + const uint8_t* src_zero_points, + uint8_t* dst_weights, + Tin* dst_scales, + uint8_t* dst_zero_points, + bool columnwise, + int rows, + int columns, + int quant_block_size, + MLAS_THREADPOOL* thread_pool +) +{ + if (columnwise) { + BlockwiseQDQQuantizer::TransposeColumnWiseQuantized( + src_weights, src_scales, src_zero_points, dst_weights, dst_scales, dst_zero_points, + rows, columns, quant_block_size, thread_pool + ); + } else { + ORT_THROW("Row-wise MlasQDQTransposeBlockwiseQuantized is not implemented"); + } +} + +template void +MlasQDQTransposeBlockwiseQuantized( + const uint8_t* src_weights, + const float* src_scales, + const uint8_t* src_zero_points, + uint8_t* dst_weights, + float* dst_scales, + uint8_t* dst_zero_points, + bool columnwise, + int rows, + int columns, + int quant_block_size, + MLAS_THREADPOOL* thread_pool +); + +template void +MlasQDQTransposeBlockwiseQuantized( + const uint8_t* src_weights, + const MLAS_FP16* src_scales, + const uint8_t* src_zero_points, + uint8_t* dst_weights, + MLAS_FP16* dst_scales, + uint8_t* dst_zero_points, + bool columnwise, + int rows, + int columns, + int quant_block_size, + MLAS_THREADPOOL* thread_pool +); diff --git a/onnxruntime/core/util/qmath.h b/onnxruntime/core/util/qmath.h index fcd1db31f95ef..c982a7aa2e7e0 100644 --- a/onnxruntime/core/util/qmath.h +++ b/onnxruntime/core/util/qmath.h @@ -552,53 +552,80 @@ struct BlockedQuantizeLinear { std::ptrdiff_t N, const std::ptrdiff_t quant_block_size, const std::ptrdiff_t thread_block_size, bool saturate) { ORT_UNUSED_PARAMETER(saturate); + // to avoid a byte being writen from mutiple threads, use 2 * N as thread block + ORT_UNUSED_PARAMETER(thread_block_size); constexpr auto low = static_cast(TOut::min_val); constexpr auto high = static_cast(TOut::max_val); - const auto num_thread_block_N = (N + thread_block_size - 1) / thread_block_size; - const auto num_thread_block = M * K * num_thread_block_N; - const TensorOpCost unit_cost{static_cast(thread_block_size * sizeof(float) * 2), - static_cast(thread_block_size * sizeof(typename TOut::UnpackedType)), - static_cast(thread_block_size) * 2.0}; - auto KN = K * N; - auto num_quant_block_KN = (K + quant_block_size - 1) / quant_block_size * N; - const auto num_thread_block_KN = K * num_thread_block_N; + auto size_thread_block = 2 * N; + auto num_thread_block = (M * K + 1) / 2; + auto num_quant_block_K = (K + quant_block_size - 1) / quant_block_size; + auto num_quant_block_KN = num_quant_block_K * N; + auto MK = M * K; + const TensorOpCost unit_cost{static_cast(size_thread_block * sizeof(float) * 2), + static_cast(size_thread_block * sizeof(typename TOut::UnpackedType)), + static_cast(size_thread_block) * 2.0}; concurrency::ThreadPool::TryParallelFor( thread_pool, num_thread_block, unit_cost, [&](std::ptrdiff_t begin, std::ptrdiff_t end) { - auto m = begin / num_thread_block_KN, k = begin % num_thread_block_KN / num_thread_block_N; - auto n_blk = begin % num_thread_block_N, n = n_blk * thread_block_size; - auto output_idx = m * KN + k * N + n; - auto quant_param_idx = m * num_quant_block_KN + k / quant_block_size * N; - auto quant_param_idx_t = quant_param_idx + n; + begin <<= 1, end = std::min(end << 1, MK); + auto output_idx = begin * N; + auto m = begin / K, k = begin % K; + auto zp_idx = m * num_quant_block_KN + k / quant_block_size * N; for (; begin < end; ++begin) { - auto n_end = std::min(N, n + thread_block_size); - // TODO(fajin): 1> use SIMD, 2> set block to quant_block_size * thread_block_size - // TODO(fajin): process 2 elements at a time - for (; n < n_end; ++n, ++output_idx, ++quant_param_idx_t) { - // TODO(fajin): perf difference + auto zp_idx_t = zp_idx; + auto output_idx_end = output_idx + N; + + // leading unaligned output + if (output_idx & 1) { auto zp = zero_point - ? static_cast(zero_point[quant_param_idx_t >> 1].GetElem(quant_param_idx_t & 1)) + ? static_cast(zero_point[zp_idx_t >> 1].GetElem(zp_idx_t & 1)) : 0; - auto sc = scale[quant_param_idx_t]; + auto sc = scale[zp_idx_t]; auto v = std::clamp(static_cast(std::nearbyint(input[output_idx] / sc)) + zp, low, high); - output[output_idx >> 1].SetElem(output_idx & 1, static_cast(v)); + output[output_idx >> 1].SetElem(1, static_cast(v)); + ++output_idx; + ++zp_idx_t; } - if (n == N) { - n = 0; - ++k; - if (k == K) { - k = 0; - quant_param_idx += N; - } else if (k % quant_block_size == 0) { - quant_param_idx += N; - } + // TODO(fajin): use SIMD + // aligned output + auto output_t = reinterpret_cast(output); + for (; output_idx < output_idx_end - 1; output_idx += 2, zp_idx_t += 2) { + auto zp0 = zero_point + ? static_cast(zero_point[zp_idx_t >> 1].GetElem(zp_idx_t & 1)) + : 0; + auto zp1 = zero_point + ? static_cast(zero_point[(zp_idx_t + 1) >> 1].GetElem((zp_idx_t + 1) & 1)) + : 0; + auto sc0 = scale[zp_idx_t]; + auto sc1 = scale[zp_idx_t + 1]; + auto v0 = std::clamp(static_cast(std::nearbyint(input[output_idx] / sc0)) + zp0, low, high); + auto v1 = std::clamp(static_cast(std::nearbyint(input[output_idx + 1] / sc1)) + zp1, low, high); + output_t[output_idx >> 1] = static_cast((v0 & 0xF) | ((v1 & 0xF) << 4)); + } - quant_param_idx_t = quant_param_idx; + // tailing unaligned output + if (output_idx < output_idx_end) { + auto zp = zero_point + ? static_cast(zero_point[zp_idx_t >> 1].GetElem(zp_idx_t & 1)) + : 0; + auto sc = scale[zp_idx_t]; + auto v = std::clamp(static_cast(std::nearbyint(input[output_idx] / sc)) + zp, low, high); + output[output_idx >> 1].SetElem(0, static_cast(v)); + + ++output_idx; + } + + ++k; + if (k == K) { + k = 0; + zp_idx += N; + } else if (k % quant_block_size == 0) { + zp_idx += N; } } }); @@ -610,53 +637,59 @@ struct BlockedQuantizeLinear { ORT_UNUSED_PARAMETER(saturate); constexpr auto low = static_cast(TOut::min_val); constexpr auto high = static_cast(TOut::max_val); - // quant block size is used as thread block size - const auto num_thread_block_K = (K + quant_block_size - 1) / quant_block_size; - const auto num_thread_block = num_thread_block_K * M; - const TensorOpCost unit_cost{static_cast(quant_block_size * sizeof(float)), - static_cast(quant_block_size * sizeof(typename TOut ::UnpackedType)), - static_cast(quant_block_size) * 2.0}; + // to avoid a byte being writen from mutiple threads, use 2 * K as thread block + auto size_thread_block = 2 * K; + auto quant_block_num_K = (K + quant_block_size - 1) / quant_block_size; + auto num_thread_block = (M + 1) / 2; + TensorOpCost unit_cost{static_cast(size_thread_block * sizeof(float)), + static_cast(size_thread_block * sizeof(typename TOut ::UnpackedType)), + static_cast(size_thread_block) * 2.0}; concurrency::ThreadPool::TryParallelFor( thread_pool, num_thread_block, unit_cost, [&](std::ptrdiff_t begin, std::ptrdiff_t end) { - auto m = begin / num_thread_block_K, k_blk = begin % num_thread_block_K, k = k_blk * quant_block_size; - auto output_idx = m * K + k; - - for (; begin < end; ++begin) { - auto zp = zero_point ? static_cast(zero_point[begin >> 1].GetElem(begin & 1)) : 0; - auto sc = scale[begin]; - size_t output_idx_end = std::min(K - k, quant_block_size) + output_idx; - size_t out_start = output_idx, out_end = output_idx_end; - - if (out_start & 1) { - auto v = std::clamp(static_cast(std::nearbyint(input[out_start] / sc)) + zp, low, high); - output[out_start >> 1].SetElem(1, static_cast(v)); - ++out_start; - } + begin <<= 1, end = std::min(end << 1, M); + auto output_idx = begin * K; + auto zp_idx = begin * quant_block_num_K; + + for (; begin < end; ++begin, output_idx += K) { + auto output_row_idx_start = output_idx; + auto output_row_idx_end = output_row_idx_start + K; + + for (; output_row_idx_start < output_row_idx_end; output_row_idx_start += quant_block_size, ++zp_idx) { + auto zp = zero_point ? static_cast(zero_point[zp_idx >> 1].GetElem(zp_idx & 1)) : 0; + auto sc = scale[zp_idx]; + size_t out_start = output_row_idx_start; + size_t out_end = std::min(output_row_idx_start + quant_block_size, output_row_idx_end); + + if (out_start & 1) { + auto v = std::clamp(static_cast(std::nearbyint(input[out_start] / sc)) + zp, low, high); + output[out_start >> 1].SetElem(1, static_cast(v)); + ++out_start; + } - if (out_end & 1) { - --out_end; - auto v = std::clamp(static_cast(std::nearbyint(input[out_end] / sc)) + zp, low, high); - output[out_end >> 1].SetElem(0, static_cast(v)); - } + if (out_end & 1) { + --out_end; + auto v = std::clamp(static_cast(std::nearbyint(input[out_end] / sc)) + zp, low, high); + output[out_end >> 1].SetElem(0, static_cast(v)); + } - if constexpr (std::is_same::value) { - MlasQuantizeLinearS4(input + out_start, reinterpret_cast(&(output[out_start >> 1])), - out_end - out_start, sc, static_cast(zp)); - } else { - MlasQuantizeLinearU4(input + out_start, reinterpret_cast(&(output[out_start >> 1])), - out_end - out_start, sc, static_cast(zp)); + if constexpr (std::is_same::value) { + MlasQuantizeLinearS4(input + out_start, reinterpret_cast(&(output[out_start >> 1])), + out_end - out_start, sc, static_cast(zp)); + } else { + MlasQuantizeLinearU4(input + out_start, reinterpret_cast(&(output[out_start >> 1])), + out_end - out_start, sc, static_cast(zp)); + } } - - output_idx = output_idx_end; - k = output_idx % K; } }); } }; +// Bug(fajin): the same byte in output / zero_point must not be written by different threads, otherwise +// the result is undefined. This is not handled in the current implementation. template struct BlockedQuantizeLinear { static void opNotLastAxis(concurrency::ThreadPool* thread_pool, const MLFloat16* input, const MLFloat16* scale, @@ -664,54 +697,84 @@ struct BlockedQuantizeLinear { std::ptrdiff_t N, const std::ptrdiff_t quant_block_size, const std::ptrdiff_t thread_block_size, bool saturate) { ORT_UNUSED_PARAMETER(saturate); + // to avoid a byte being writen from mutiple threads, use 2 * N as thread block + ORT_UNUSED_PARAMETER(thread_block_size); constexpr auto low = static_cast(TOut::min_val); constexpr auto high = static_cast(TOut::max_val); - const auto num_thread_block_N = (N + thread_block_size - 1) / thread_block_size; - const auto num_thread_block = M * K * num_thread_block_N; - const TensorOpCost unit_cost{static_cast(thread_block_size * sizeof(MLFloat16) * 2), - static_cast(thread_block_size * sizeof(typename TOut::UnpackedType)), - static_cast(thread_block_size) * 2.0}; - auto KN = K * N; - auto num_quant_block_KN = (K + quant_block_size - 1) / quant_block_size * N; - const auto num_thread_block_KN = K * num_thread_block_N; + auto size_thread_block = 2 * N; + auto num_thread_block = (M * K + 1) / 2; + auto num_quant_block_K = (K + quant_block_size - 1) / quant_block_size; + auto num_quant_block_KN = num_quant_block_K * N; + auto MK = M * K; + const TensorOpCost unit_cost{static_cast(size_thread_block * sizeof(float) * 2), + static_cast(size_thread_block * sizeof(typename TOut::UnpackedType)), + static_cast(size_thread_block) * 2.0}; concurrency::ThreadPool::TryParallelFor( thread_pool, num_thread_block, unit_cost, [&](std::ptrdiff_t begin, std::ptrdiff_t end) { - auto m = begin / num_thread_block_KN, k = begin % num_thread_block_KN / num_thread_block_N; - auto n_blk = begin % num_thread_block_N, n = n_blk * thread_block_size; - auto output_idx = m * KN + k * N + n; - auto quant_param_idx = m * num_quant_block_KN + k / quant_block_size * N; - auto quant_param_idx_t = quant_param_idx + n; + begin <<= 1, end = std::min(end << 1, MK); + auto output_idx = begin * N; + auto m = begin / K, k = begin % K; + auto zp_idx = m * num_quant_block_KN + k / quant_block_size * N; for (; begin < end; ++begin) { - auto n_end = std::min(N, n + thread_block_size); - // TODO(fajin): 1> use SIMD, 2> set block to quant_block_size * thread_block_size - // TODO(fajin): process 2 elements at a time - for (; n < n_end; ++n, ++output_idx, ++quant_param_idx_t) { - // TODO(fajin): perf difference + auto zp_idx_t = zp_idx; + auto output_idx_end = output_idx + N; + + // leading unaligned output + if (output_idx & 1) { auto zp = zero_point - ? static_cast(zero_point[quant_param_idx_t >> 1].GetElem(quant_param_idx_t & 1)) + ? static_cast(zero_point[zp_idx_t >> 1].GetElem(zp_idx_t & 1)) : 0; - auto sc = scale[quant_param_idx_t].ToFloat(); - auto v = std::clamp(static_cast(std::nearbyint(input[output_idx].ToFloat() / sc)) + zp, - low, high); - output[output_idx >> 1].SetElem(output_idx & 1, static_cast(v)); + auto sc = scale[zp_idx_t].ToFloat(); + auto v = std::clamp( + static_cast(std::nearbyint(input[output_idx].ToFloat() / sc)) + zp, low, high); + output[output_idx >> 1].SetElem(1, static_cast(v)); + ++output_idx; + ++zp_idx_t; } - if (n == N) { - n = 0; - ++k; - if (k == K) { - k = 0; - quant_param_idx += N; - } else if (k % quant_block_size == 0) { - quant_param_idx += N; - } + // TODO(fajin): use SIMD + // aligned output + auto output_t = reinterpret_cast(output); + for (; output_idx < output_idx_end - 1; output_idx += 2, zp_idx_t += 2) { + auto zp0 = zero_point + ? static_cast(zero_point[zp_idx_t >> 1].GetElem(zp_idx_t & 1)) + : 0; + auto zp1 = zero_point + ? static_cast(zero_point[(zp_idx_t + 1) >> 1].GetElem((zp_idx_t + 1) & 1)) + : 0; + auto sc0 = scale[zp_idx_t].ToFloat(); + auto sc1 = scale[zp_idx_t + 1].ToFloat(); + auto v0 = std::clamp( + static_cast(std::nearbyint(input[output_idx].ToFloat() / sc0)) + zp0, low, high); + auto v1 = std::clamp( + static_cast(std::nearbyint(input[output_idx + 1].ToFloat() / sc1)) + zp1, low, high); + output_t[output_idx >> 1] = static_cast((v0 & 0xF) | ((v1 & 0xF) << 4)); + } - quant_param_idx_t = quant_param_idx; + // tailing unaligned output + if (output_idx < output_idx_end) { + auto zp = zero_point + ? static_cast(zero_point[zp_idx_t >> 1].GetElem(zp_idx_t & 1)) + : 0; + auto sc = scale[zp_idx_t].ToFloat(); + auto v = std::clamp( + static_cast(std::nearbyint(input[output_idx].ToFloat() / sc)) + zp, low, high); + output[output_idx >> 1].SetElem(0, static_cast(v)); + + ++output_idx; + } + + ++k; + if (k == K) { + k = 0; + zp_idx += N; + } else if (k % quant_block_size == 0) { + zp_idx += N; } } }); @@ -723,32 +786,55 @@ struct BlockedQuantizeLinear { ORT_UNUSED_PARAMETER(saturate); constexpr auto low = static_cast(TOut::min_val); constexpr auto high = static_cast(TOut::max_val); - // quant block size is used as thread block size - const auto num_thread_block_K = (K + quant_block_size - 1) / quant_block_size; - const auto num_thread_block = num_thread_block_K * M; - const TensorOpCost unit_cost{static_cast(quant_block_size * sizeof(MLFloat16)), - static_cast(quant_block_size * sizeof(typename TOut::UnpackedType)), - static_cast(quant_block_size) * 2.0}; + // to avoid a byte being writen from mutiple threads, use 2 * K as thread block + auto size_thread_block = 2 * K; + auto quant_block_num_K = (K + quant_block_size - 1) / quant_block_size; + auto num_thread_block = (M + 1) / 2; + TensorOpCost unit_cost{static_cast(size_thread_block * sizeof(float)), + static_cast(size_thread_block * sizeof(typename TOut ::UnpackedType)), + static_cast(size_thread_block) * 2.0}; concurrency::ThreadPool::TryParallelFor( thread_pool, num_thread_block, unit_cost, [&](std::ptrdiff_t begin, std::ptrdiff_t end) { - auto m = begin / num_thread_block_K, k_blk = begin % num_thread_block_K, k = k_blk * quant_block_size; - auto output_idx = m * K + k; + begin <<= 1, end = std::min(end << 1, M); + auto output_idx = begin * K; + auto zp_idx = begin * quant_block_num_K; + + for (; begin < end; ++begin, output_idx += K) { + auto output_row_idx_start = output_idx; + auto output_row_idx_end = output_row_idx_start + K; + + for (; output_row_idx_start < output_row_idx_end; output_row_idx_start += quant_block_size, ++zp_idx) { + auto zp = zero_point ? static_cast(zero_point[zp_idx >> 1].GetElem(zp_idx & 1)) : 0; + auto sc = scale[zp_idx].ToFloat(); + size_t out_start = output_row_idx_start; + size_t out_end = std::min(output_row_idx_start + quant_block_size, output_row_idx_end); + + if (out_start & 1) { + auto v = std::clamp( + static_cast(std::nearbyint(input[out_start].ToFloat() / sc)) + zp, low, high); + output[out_start >> 1].SetElem(1, static_cast(v)); + ++out_start; + } - for (; begin < end; ++begin) { - // each thread block is also a quantization block - auto zp = zero_point ? static_cast(zero_point[begin >> 1].GetElem(begin & 1)) : 0; - auto sc = scale[begin].ToFloat(); - auto output_idx_end = std::min(K - k, quant_block_size) + output_idx; - for (; output_idx < output_idx_end; ++output_idx) { - auto v = std::clamp(static_cast(std::nearbyint(input[output_idx].ToFloat() / sc)) + zp, - low, high); - output[output_idx >> 1].SetElem(output_idx & 1, static_cast(v)); - } + if (out_end & 1) { + --out_end; + auto v = std::clamp( + static_cast(std::nearbyint(input[out_end].ToFloat() / sc)) + zp, low, high); + output[out_end >> 1].SetElem(0, static_cast(v)); + } - k = output_idx % K; + auto output_t = reinterpret_cast(output); + for (; out_start < out_end; out_start += 2) { + auto v0 = std::clamp( + static_cast(std::nearbyint(input[out_start].ToFloat() / sc)) + zp, low, high); + auto v1 = std::clamp( + static_cast(std::nearbyint(input[out_start + 1].ToFloat() / sc)) + zp, low, high); + output_t[out_start >> 1] = static_cast((v0 & 0xF) | ((v1 & 0xF) << 4)); + } + } } }); } diff --git a/onnxruntime/python/onnxruntime_pybind_quant.cc b/onnxruntime/python/onnxruntime_pybind_quant.cc index ff76887e917cd..5e8e5c1a2a2fc 100644 --- a/onnxruntime/python/onnxruntime_pybind_quant.cc +++ b/onnxruntime/python/onnxruntime_pybind_quant.cc @@ -66,6 +66,37 @@ void QuantizeMatMul4BitsBlockwise( tp.get()); } +template +void QuantizeQDQMatMul4BitsBlockwise( + py::array_t dst, // shape: [K, N / 2] + py::array_t src, // shape: [K, N] + py::array_t scale, // shape: [block_per_K, N] + py::array_t zero_points, // shape: [block_per_K, N / 2] + int32_t quant_block_size, + int32_t N, + int32_t K, + bool is_symmetric) { + OrtThreadPoolParams to; + auto tp = concurrency::CreateThreadPool(&onnxruntime::Env::Default(), to, + concurrency::ThreadPoolType::INTRA_OP); + + py::buffer_info dst_buf = dst.request(); + py::buffer_info src_buf = src.request(); + py::buffer_info scale_buf = scale.request(); + py::buffer_info zp_buf = zero_points.request(); + + MlasQDQQuantizeBlockwise( + reinterpret_cast(src_buf.ptr), + reinterpret_cast(scale_buf.ptr), + is_symmetric ? nullptr : reinterpret_cast(zp_buf.ptr), + reinterpret_cast(dst_buf.ptr), + true, + K, + N, + quant_block_size, + tp.get()); +} + template void QuantizeMatMulBnb4Blockwise( py::array_t dst, @@ -99,6 +130,8 @@ void CreateQuantPybindModule(py::module& m) { m.def("quantize_matmul_4bits", &QuantizeMatMul4BitsBlockwise); m.def("quantize_matmul_bnb4", &QuantizeMatMulBnb4Blockwise); m.def("quantize_matmul_bnb4", &QuantizeMatMulBnb4Blockwise); + m.def("quantize_qdq_matmul_4bits", &QuantizeQDQMatMul4BitsBlockwise); + m.def("quantize_qdq_matmul_4bits", &QuantizeQDQMatMul4BitsBlockwise); } } // namespace python diff --git a/onnxruntime/test/mlas/bench/bench_q4dq.cpp b/onnxruntime/test/mlas/bench/bench_q4dq.cpp new file mode 100644 index 0000000000000..00234ecfd2ce2 --- /dev/null +++ b/onnxruntime/test/mlas/bench/bench_q4dq.cpp @@ -0,0 +1,118 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include +#include + +#include "core/mlas/inc/mlas_q4.h" +#include "test/mlas/bench/bench_util.h" +#include "core/util/thread_utils.h" + +static void BM_QDQBlockwiseQuantizer_QuantizeColumnwise(benchmark::State& state) { + int M = state.range(0); + int N = state.range(1); + int quant_block_size = state.range(2); + int threads = state.range(3); + size_t scale_size = (M + quant_block_size - 1) / quant_block_size * N; + + auto src = RandomVectorUniform(M * N, -16.0f, 14.0f); + auto scales = std::vector(scale_size); + auto zero_points = std::vector((scale_size + 1) / 2); + auto dst = std::vector((M * N + 1) / 2); + + OrtThreadPoolParams tpo; + tpo.thread_pool_size = static_cast(threads); + tpo.auto_set_affinity = true; + std::unique_ptr tp( + onnxruntime::concurrency::CreateThreadPool(&onnxruntime::Env::Default(), + tpo, onnxruntime::concurrency::ThreadPoolType::INTRA_OP)); + + for (auto _ : state) { + benchmark::DoNotOptimize(dst.data()); + MlasQDQQuantizeBlockwise( + src.data(), scales.data(), zero_points.data(), dst.data(), + true, M, N, quant_block_size, tp.get()); + benchmark::ClobberMemory(); + } +} + +static void BM_MlasQuantizeBlockwise(benchmark::State& state) { + int M = state.range(0); + int N = state.range(1); + int quant_block_size = state.range(2); + int threads = state.range(3); + size_t scale_size = (M + quant_block_size - 1) / quant_block_size * N; + + auto src = RandomVectorUniform(M * N, -16.0f, 14.0f); + auto scales = std::vector(scale_size); + auto zero_points = std::vector((scale_size + 1) / 2); + auto dst = std::vector((M * N + 1) / 2); + + OrtThreadPoolParams tpo; + tpo.thread_pool_size = static_cast(threads); + tpo.auto_set_affinity = true; + std::unique_ptr tp( + onnxruntime::concurrency::CreateThreadPool(&onnxruntime::Env::Default(), + tpo, onnxruntime::concurrency::ThreadPoolType::INTRA_OP)); + + for (auto _ : state) { + benchmark::DoNotOptimize(dst.data()); + MlasQuantizeBlockwise( + dst.data(), scales.data(), zero_points.data(), src.data(), + quant_block_size, true, M, N, N, tp.get()); + benchmark::ClobberMemory(); + } +} + +static void BM_QDQBlockwiseQuantizer_TransposeColumnwise(benchmark::State& state) { + int M = state.range(0); + int N = state.range(1); + int quant_block_size = state.range(2); + int threads = state.range(3); + int quant_num_M = (M + quant_block_size - 1) / quant_block_size; + int blob_size = (quant_block_size + 1) / 2; + size_t scale_size = quant_num_M * N; + + auto scales = RandomVectorUniform(scale_size, -16.0f, 14.0f); + auto zero_points = RandomVectorUniform(static_cast((scale_size + 1) / 2), 0, 255); + auto dst = RandomVectorUniform(static_cast((M * N + 1) / 2), 0, 255); + auto scales_T = std::vector(scale_size); + auto zero_points_T = std::vector(((quant_num_M + 1) / 2) * N); + auto dst_T = std::vector(quant_num_M * blob_size * N); + + OrtThreadPoolParams tpo; + tpo.thread_pool_size = static_cast(threads); + tpo.auto_set_affinity = true; + std::unique_ptr tp( + onnxruntime::concurrency::CreateThreadPool(&onnxruntime::Env::Default(), + tpo, onnxruntime::concurrency::ThreadPoolType::INTRA_OP)); + + for (auto _ : state) { + benchmark::DoNotOptimize(dst.data()); + MlasQDQTransposeBlockwiseQuantized( + dst.data(), scales.data(), zero_points.data(), dst_T.data(), scales_T.data(), zero_points_T.data(), + true, M, N, quant_block_size, tp.get()); + benchmark::ClobberMemory(); + } +} + +BENCHMARK(BM_QDQBlockwiseQuantizer_QuantizeColumnwise) + ->UseRealTime() + ->Apply([](benchmark::internal::Benchmark* b) { + b->ArgNames({"M", "N", "quant_block_size", "threads"}); + b->ArgsProduct({{1024, 4096}, {4096, 4095}, {64, 128}, {8}}); + }); + +BENCHMARK(BM_MlasQuantizeBlockwise) + ->UseRealTime() + ->Apply([](benchmark::internal::Benchmark* b) { + b->ArgNames({"M", "N", "quant_block_size", "threads"}); + b->ArgsProduct({{1024, 4096}, {4096, 4095}, {64, 128}, {8}}); + }); + +BENCHMARK(BM_QDQBlockwiseQuantizer_TransposeColumnwise) + ->UseRealTime() + ->Apply([](benchmark::internal::Benchmark* b) { + b->ArgNames({"M", "N", "quant_block_size", "threads"}); + b->ArgsProduct({{1024, 4096}, {4096, 4095}, {64, 128}, {2, 8, 16}}); + }); diff --git a/onnxruntime/test/mlas/unittest/test_blockq4.cpp b/onnxruntime/test/mlas/unittest/test_blockq4.cpp index 07f0748fb7ed1..b466e883059f4 100644 --- a/onnxruntime/test/mlas/unittest/test_blockq4.cpp +++ b/onnxruntime/test/mlas/unittest/test_blockq4.cpp @@ -29,10 +29,18 @@ class MlasBlockwiseQdqTest : public MlasTestBase { MatrixGuardBuffer OutputElements; MatrixGuardBuffer OutputScales; MatrixGuardBuffer OutputOffsets; + MatrixGuardBuffer QDQOutputElements; + MatrixGuardBuffer QDQOutputScales; + MatrixGuardBuffer QDQOutputOffsets; + MatrixGuardBuffer QDQTransposedOutputElements; + MatrixGuardBuffer QDQTransposedOutputScales; + MatrixGuardBuffer QDQTransposedOutputOffsets; void Test(int rows, int columns, int block_size, bool columnwise, bool symmetric) { float* dequant_buf = FpBuf.GetBuffer(rows * columns, true); float* transposed = FpBuf2.GetBuffer(rows * columns, true); + size_t scale_size = (rows + block_size - 1) / block_size * columns; + size_t zp_size = (scale_size + 1) / 2; MLAS_THREADPOOL* threadpool_ptr = GetMlasThreadPool(); @@ -49,6 +57,8 @@ class MlasBlockwiseQdqTest : public MlasTestBase { q_data_size_in_bytes, q_scale_size, &q_zp_size_in_bytes); uint8_t* elements = InputElements.GetBuffer(q_data_size_in_bytes, true); + uint8_t* qdq_weights = QDQOutputElements.GetBuffer((rows * columns + 1) / 2, true); + uint8_t* qdq_weights_T = QDQTransposedOutputElements.GetBuffer(q_data_size_in_bytes, true); int v = 7; for (int c = 0; c < columns; c++) { @@ -75,7 +85,11 @@ class MlasBlockwiseQdqTest : public MlasTestBase { } float* scales = InputScales.GetBuffer(q_scale_size); + float* qdq_scales = QDQOutputScales.GetBuffer(scale_size); + float* qdq_scales_T = QDQTransposedOutputScales.GetBuffer(q_scale_size); uint8_t* zp = symmetric ? nullptr : InputOffsets.GetBuffer(q_zp_size_in_bytes, true); + uint8_t* qdq_zp = symmetric ? nullptr : QDQOutputOffsets.GetBuffer(zp_size, true); + uint8_t* qdq_zp_T = symmetric ? nullptr : QDQTransposedOutputOffsets.GetBuffer(q_zp_size_in_bytes, true); if (zp) { for (int c = 0; c < meta_cols; c++) { for (int r = 0; r < meta_rows; r += 2) { @@ -112,16 +126,37 @@ class MlasBlockwiseQdqTest : public MlasTestBase { MlasQuantizeBlockwise(o_elements, o_scales, o_zp, transposed, block_size, columnwise, rows, columns, columns, threadpool_ptr); + if (columnwise) { + MlasQDQQuantizeBlockwise( + transposed, qdq_scales, qdq_zp, qdq_weights, + true, rows, columns, block_size, threadpool_ptr); + + MlasQDQTransposeBlockwiseQuantized( + qdq_weights, qdq_scales, qdq_zp, qdq_weights_T, qdq_scales_T, qdq_zp_T, + true, rows, columns, block_size, threadpool_ptr); + } + for (int c = 0; c < columns; c++) { for (int r = 0; r < rows; r += 2) { int idx = c * q_rows + r / 2; ASSERT_EQ(o_elements[idx] & 0xf, elements[idx] & 0xf) << ", index=[" << r << "x" << c << "], shape=[" << rows << "x" << columns << "] block: " << block_size << ", symmetric: " << symmetric << ", columnwise: " << columnwise; + if (columnwise) { + ASSERT_EQ(qdq_weights_T[idx] & 0xf, elements[idx] & 0xf) + << ", index=[" << r << "x" << c << "], shape=[" << rows << "x" << columns + << "] block: " << block_size << ", symmetric: " << symmetric << ", columnwise: " << columnwise; + } + if (r + 1 < rows) { ASSERT_EQ(o_elements[idx] >> 4, elements[idx] >> 4) << ", index=[" << r + 1 << "x" << c << "], shape=[" << rows << "x" << columns << "] block: " << block_size << ", symmetric: " << symmetric << ", columnwise: " << columnwise; + if (columnwise) { + ASSERT_EQ(qdq_weights_T[idx] >> 4, elements[idx] >> 4) + << ", index=[" << r + 1 << "x" << c << "], shape=[" << rows << "x" << columns + << "] block: " << block_size << ", symmetric: " << symmetric << ", columnwise: " << columnwise; + } } } } @@ -132,6 +167,12 @@ class MlasBlockwiseQdqTest : public MlasTestBase { ASSERT_EQ(o_scales[idx], scales[idx]) << ", index=" << r << "x" << c << ", shape=[" << rows << "x" << columns << "] block: " << block_size << ", symmetric: " << symmetric << ", columnwise: " << columnwise; + + if (columnwise) { + ASSERT_EQ(qdq_scales_T[idx], scales[idx]) + << ", index=" << r << "x" << c << ", shape=[" << rows << "x" << columns + << "] block: " << block_size << ", symmetric: " << symmetric << ", columnwise: " << columnwise; + } } } @@ -142,10 +183,20 @@ class MlasBlockwiseQdqTest : public MlasTestBase { ASSERT_EQ(o_zp[idx] & 0xf, zp[idx] & 0xf) << ", index=" << r << "x" << c << ", shape=[" << rows << "x" << columns << "] block: " << block_size << ", symmetric: " << symmetric << ", columnwise: " << columnwise; + if (columnwise) { + ASSERT_EQ(qdq_zp_T[idx] & 0xf, zp[idx] & 0xf) + << ", index=" << r << "x" << c << ", shape=[" << rows << "x" << columns + << "] block: " << block_size << ", symmetric: " << symmetric << ", columnwise: " << columnwise; + } if (r + 1 < meta_rows) { ASSERT_EQ(o_zp[idx] >> 4, zp[idx] >> 4) << ", index=" << r + 1 << "x" << c << ", shape=[" << rows << "x" << columns << "] block: " << block_size << ", symmetric: " << symmetric << ", columnwise: " << columnwise; + if (columnwise) { + ASSERT_EQ(qdq_zp_T[idx] >> 4, zp[idx] >> 4) + << ", index=" << r + 1 << "x" << c << ", shape=[" << rows << "x" << columns + << "] block: " << block_size << ", symmetric: " << symmetric << ", columnwise: " << columnwise; + } } } } diff --git a/onnxruntime/test/onnx/microbenchmark/quantize.cc b/onnxruntime/test/onnx/microbenchmark/quantize.cc index fda4324c0e83d..a6ab8484231b8 100644 --- a/onnxruntime/test/onnx/microbenchmark/quantize.cc +++ b/onnxruntime/test/onnx/microbenchmark/quantize.cc @@ -82,12 +82,11 @@ BENCHMARK(BM_Quantize) static void BM_BlockedQuantize_NotLastAxis(benchmark::State& state) { using Int4 = onnxruntime::Int4x2; using UnpackedType = Int4::UnpackedType; - const std::ptrdiff_t M[] = {96, 192, 192}; - const std::ptrdiff_t N[] = {2048, 2048, 4096}; - const int64_t size_idx = state.range(0); - const int64_t threads = state.range(1); + const int64_t M = state.range(0); + const int64_t N = state.range(1); const int64_t block_size = state.range(2); - size_t batch_size = M[size_idx] * N[size_idx]; + const int64_t threads = state.range(3); + size_t batch_size = M * N; size_t quant_block_size = 64; size_t scale_size = batch_size / quant_block_size; @@ -108,7 +107,7 @@ static void BM_BlockedQuantize_NotLastAxis(benchmark::State& state) { benchmark::DoNotOptimize(a_data_quant); onnxruntime::BlockedQuantizeLinear::opNotLastAxis( tp.get(), a_data, scale, reinterpret_cast(zero_point), reinterpret_cast(a_data_quant), - 1, M[size_idx], N[size_idx], static_cast(quant_block_size), + 1, M, N, static_cast(quant_block_size), static_cast(block_size), true); benchmark::ClobberMemory(); } @@ -121,12 +120,11 @@ static void BM_BlockedQuantize_NotLastAxis(benchmark::State& state) { static void BM_BlockedQuantize_LastAxis(benchmark::State& state) { using Int4 = onnxruntime::Int4x2; using UnpackedType = Int4::UnpackedType; - const std::ptrdiff_t M[] = {96, 192, 192}; - const std::ptrdiff_t N[] = {2048, 2048, 4096}; - const int64_t size_idx = state.range(0); - const int64_t threads = state.range(1); + const int64_t M = state.range(0); + const int64_t N = state.range(1); const int64_t quant_block_size = state.range(2); - size_t batch_size = M[size_idx] * N[size_idx]; + const int64_t threads = state.range(3); + size_t batch_size = M * N; size_t scale_size = batch_size / quant_block_size; float* a_data = GenerateArrayWithRandomValue(batch_size, -16, 14); @@ -146,7 +144,7 @@ static void BM_BlockedQuantize_LastAxis(benchmark::State& state) { benchmark::DoNotOptimize(a_data_quant); onnxruntime::BlockedQuantizeLinear::opLastAxis( tp.get(), a_data, scale, reinterpret_cast(zero_point), reinterpret_cast(a_data_quant), - M[size_idx], N[size_idx], static_cast(quant_block_size), true); + M, N, static_cast(quant_block_size), true); benchmark::ClobberMemory(); } aligned_free(a_data_quant); @@ -159,24 +157,14 @@ BENCHMARK(BM_BlockedQuantize_NotLastAxis) ->UseRealTime() ->Unit(benchmark::TimeUnit::kNanosecond) ->Apply([](benchmark::internal::Benchmark* b) { - for (int size_idx : {0, 1, 2}) { - for (int thread : {2, 4, 8}) { - for (int block_size : {64, 128}) { - b->Args({size_idx, thread, block_size}); - } - } - } + b->ArgNames({"M", "N", "block_size", "threads"}); + b->ArgsProduct({{1024, 4096}, {4096}, {128}, {2, 8}}); }); BENCHMARK(BM_BlockedQuantize_LastAxis) ->UseRealTime() ->Unit(benchmark::TimeUnit::kNanosecond) ->Apply([](benchmark::internal::Benchmark* b) { - for (int size_idx : {0, 1, 2}) { - for (int thread : {2, 4, 8}) { - for (int quant_block_size : {16, 64, 256}) { - b->Args({size_idx, thread, quant_block_size}); - } - } - } + b->ArgNames({"M", "N", "quant_block_size", "threads"}); + b->ArgsProduct({{1024, 4096}, {4096}, {64, 128}, {2, 8}}); }); From e73754786264e63d7ccbc02d0607d41ec6dfaa3a Mon Sep 17 00:00:00 2001 From: Chi Lo <54722500+chilo-ms@users.noreply.github.com> Date: Wed, 19 Jun 2024 18:36:26 -0700 Subject: [PATCH 06/52] Add support for INT64 types in TensorRT constant layer calibration (#21101) This PR is a duplicate of the https://github.com/microsoft/onnxruntime/pull/21041 Create this PR in case the original one can't be updated for patch release timeline. --- .../providers/tensorrt/tensorrt_execution_provider.cc | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc b/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc index 9c2db494f0e41..13316d6cbc749 100644 --- a/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc +++ b/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc @@ -71,6 +71,7 @@ bool SetDynamicRange(nvinfer1::INetworkDefinition& network, std::unordered_mapsetDynamicRange(-dynamic_range_iter->second, dynamic_range_iter->second)) { + LOGS_DEFAULT(ERROR) << "Failed to set dynamic range for network input " << tensor_name; return false; } } @@ -84,10 +85,12 @@ bool SetDynamicRange(nvinfer1::INetworkDefinition& network, std::unordered_mapgetOutput(j)->setDynamicRange(-dynamic_range_iter->second, dynamic_range_iter->second)) { + LOGS_DEFAULT(ERROR) << "Failed to set dynamic range for tensor " << tensor_name; return false; } } else if (trt_layer->getType() == nvinfer1::LayerType::kCONSTANT) { nvinfer1::IConstantLayer* const_layer = static_cast(trt_layer); + const std::string const_layer_name = const_layer->getName(); auto trt_weights = const_layer->getWeights(); double max_weight = std::numeric_limits::min(); for (int64_t k = 0, end = trt_weights.count; k < end; ++k) { @@ -108,13 +111,19 @@ bool SetDynamicRange(nvinfer1::INetworkDefinition& network, std::unordered_map(trt_weights.values)[k]; break; +#if NV_TENSORRT_MAJOR >= 10 + case nvinfer1::DataType::kINT64: + weight = static_cast(static_cast(trt_weights.values)[k]); + break; +#endif // NV_TENSORRT_MAJOR >= 10 default: - LOGS_DEFAULT(ERROR) << "Found unsupported datatype!"; + LOGS_DEFAULT(ERROR) << "Found unsupported datatype for layer " << const_layer_name; return false; } max_weight = std::max(max_weight, std::abs(weight)); } if (!trt_layer->getOutput(j)->setDynamicRange(static_cast(-max_weight), static_cast(max_weight))) { + LOGS_DEFAULT(ERROR) << "Failed to set dynamic range for layer " << const_layer_name; return false; } } From 27f3ac78d4e7ad6ba17fa9f948e447b0d21b986e Mon Sep 17 00:00:00 2001 From: Changming Sun Date: Wed, 19 Jun 2024 20:11:15 -0700 Subject: [PATCH 07/52] Delete RoslynAnalyzers (#21104) ### Description Delete RoslynAnalyzers. Use CodeQL instead. ### Motivation and Context Now we already have CodeQL which is modern and also covers C# code. The RoslynAnalyzers one is not in our pull request pipelines. The "RoslynAnalyzers@2" task is outdated and needs be upgraded. I will delete it for now since we already have CodeQL. --- .../c-api-noopenmp-packaging-pipelines.yml | 16 ---------------- .../stages/nuget-cuda-packaging-stage.yml | 7 ------- .../azure-pipelines/templates/c-api-cpu.yml | 7 ------- .../ondevice-training-cpu-packaging-pipeline.yml | 7 ------- 4 files changed, 37 deletions(-) diff --git a/tools/ci_build/github/azure-pipelines/c-api-noopenmp-packaging-pipelines.yml b/tools/ci_build/github/azure-pipelines/c-api-noopenmp-packaging-pipelines.yml index fbb03abb70e88..186c4fccb1045 100644 --- a/tools/ci_build/github/azure-pipelines/c-api-noopenmp-packaging-pipelines.yml +++ b/tools/ci_build/github/azure-pipelines/c-api-noopenmp-packaging-pipelines.yml @@ -408,22 +408,6 @@ stages: msbuildArguments: '-t:Clean -p:OnnxRuntimeBuildDirectory="$(Build.BinariesDirectory)" -p:OrtPackageId=Microsoft.ML.OnnxRuntime.ROCm' workingDirectory: '$(Build.SourcesDirectory)\csharp' - - task: RoslynAnalyzers@2 - displayName: 'Run Roslyn Analyzers' - inputs: - userProvideBuildInfo: msBuildInfo - msBuildCommandline: > - "C:\Program Files\Microsoft Visual Studio\2022\Enterprise\MSBuild\Current\Bin\msbuild.exe" - $(Build.SourcesDirectory)\csharp\OnnxRuntime.CSharp.sln - -p:configuration="RelWithDebInfo" - -p:Platform="Any CPU" - -p:OnnxRuntimeBuildDirectory="$(Build.BinariesDirectory)" - -p:OrtPackageId=Microsoft.ML.OnnxRuntime.ROCm - -p:IsLinuxBuild=true - -p:IsWindowsBuild=false - -p:IsMacOSBuild=false - condition: and(succeeded(), eq('${{ parameters.DoCompliance }}', true)) - - template: templates/component-governance-component-detection-steps.yml parameters : condition : 'succeeded' diff --git a/tools/ci_build/github/azure-pipelines/stages/nuget-cuda-packaging-stage.yml b/tools/ci_build/github/azure-pipelines/stages/nuget-cuda-packaging-stage.yml index 02df814eb2a93..21aa75d369010 100644 --- a/tools/ci_build/github/azure-pipelines/stages/nuget-cuda-packaging-stage.yml +++ b/tools/ci_build/github/azure-pipelines/stages/nuget-cuda-packaging-stage.yml @@ -193,13 +193,6 @@ stages: msbuildArguments: '-t:Clean -p:OnnxRuntimeBuildDirectory="$(Build.BinariesDirectory)" -p:OrtPackageId=Microsoft.ML.OnnxRuntime.Gpu' workingDirectory: '$(Build.SourcesDirectory)\csharp' - - task: RoslynAnalyzers@2 - displayName: 'Run Roslyn Analyzers' - inputs: - userProvideBuildInfo: msBuildInfo - msBuildCommandline: '"C:\Program Files\Microsoft Visual Studio\2022\Enterprise\MSBuild\Current\Bin\msbuild.exe" $(Build.SourcesDirectory)\csharp\OnnxRuntime.CSharp.sln -p:configuration="RelWithDebInfo" -p:Platform="Any CPU" -p:OnnxRuntimeBuildDirectory="$(Build.BinariesDirectory)" -p:OrtPackageId=Microsoft.ML.OnnxRuntime.Gpu' - condition: and(succeeded(), eq('${{ parameters.DoCompliance }}', true)) - - template: ../templates/component-governance-component-detection-steps.yml parameters: condition: 'succeeded' diff --git a/tools/ci_build/github/azure-pipelines/templates/c-api-cpu.yml b/tools/ci_build/github/azure-pipelines/templates/c-api-cpu.yml index d7aecae18968e..c41f9589d8469 100644 --- a/tools/ci_build/github/azure-pipelines/templates/c-api-cpu.yml +++ b/tools/ci_build/github/azure-pipelines/templates/c-api-cpu.yml @@ -478,13 +478,6 @@ stages: msbuildArguments: '-t:Clean -p:OnnxRuntimeBuildDirectory="$(Build.BinariesDirectory)" -p:OrtPackageId=$(OrtPackageId)' workingDirectory: '$(Build.SourcesDirectory)\csharp' - - task: RoslynAnalyzers@2 - displayName: 'Run Roslyn Analyzers' - inputs: - userProvideBuildInfo: msBuildInfo - msBuildCommandline: '"C:\Program Files\Microsoft Visual Studio\2022\Enterprise\MSBuild\Current\Bin\msbuild.exe" $(Build.SourcesDirectory)\csharp\OnnxRuntime.CSharp.sln -p:configuration="RelWithDebInfo" -p:Platform="Any CPU" -p:OnnxRuntimeBuildDirectory="$(Build.BinariesDirectory)" -p:OrtPackageId=$(OrtPackageId)' - condition: and(succeeded(), eq('${{ parameters.DoCompliance }}', true)) - - template: component-governance-component-detection-steps.yml parameters : condition : 'succeeded' diff --git a/tools/ci_build/github/azure-pipelines/templates/ondevice-training-cpu-packaging-pipeline.yml b/tools/ci_build/github/azure-pipelines/templates/ondevice-training-cpu-packaging-pipeline.yml index 61cf5de8bdf17..6fead4181ffa2 100644 --- a/tools/ci_build/github/azure-pipelines/templates/ondevice-training-cpu-packaging-pipeline.yml +++ b/tools/ci_build/github/azure-pipelines/templates/ondevice-training-cpu-packaging-pipeline.yml @@ -288,13 +288,6 @@ stages: msbuildArguments: '-t:Clean -p:OnnxRuntimeBuildDirectory="$(Build.BinariesDirectory)" -p:OrtPackageId=$(OrtPackageId)' workingDirectory: '$(Build.SourcesDirectory)\csharp' - - task: RoslynAnalyzers@2 - displayName: 'Run Roslyn Analyzers' - inputs: - userProvideBuildInfo: msBuildInfo - msBuildCommandline: '"C:\Program Files\Microsoft Visual Studio\2022\Enterprise\MSBuild\Current\Bin\msbuild.exe" $(Build.SourcesDirectory)\csharp\OnnxRuntime.CSharp.sln -p:configuration="RelWithDebInfo" -p:Platform="Any CPU" -p:OnnxRuntimeBuildDirectory="$(Build.BinariesDirectory)" -p:OrtPackageId=$(OrtPackageId)' - condition: and(succeeded(), eq('${{ parameters.DoCompliance }}', true)) - - template: component-governance-component-detection-steps.yml parameters : condition : 'succeeded' From bd3a9ee99df665202e214e781c20ed1be9f7ea6f Mon Sep 17 00:00:00 2001 From: Changming Sun Date: Wed, 19 Jun 2024 20:47:21 -0700 Subject: [PATCH 08/52] Add UsePythonVersion (#21109) ### Description The machine has multiple python installations and none of them is in PATH. Therefore we should explicitly set python version via this task to avoid having surprises. ### Motivation and Context Similar to #21095 --- .../azure-pipelines/stages/nuget-cuda-packaging-stage.yml | 5 +++++ .../templates/ondevice-training-cpu-packaging-pipeline.yml | 5 +++++ 2 files changed, 10 insertions(+) diff --git a/tools/ci_build/github/azure-pipelines/stages/nuget-cuda-packaging-stage.yml b/tools/ci_build/github/azure-pipelines/stages/nuget-cuda-packaging-stage.yml index 21aa75d369010..18615b6ca18b1 100644 --- a/tools/ci_build/github/azure-pipelines/stages/nuget-cuda-packaging-stage.yml +++ b/tools/ci_build/github/azure-pipelines/stages/nuget-cuda-packaging-stage.yml @@ -116,6 +116,11 @@ stages: DisplayName: 'ESRP - Sign C# dlls' DoEsrp: ${{ parameters.DoEsrp }} + - task: UsePythonVersion@0 + displayName: 'Use Python' + inputs: + versionSpec: 3.8 + - task: MSBuild@1 displayName: 'Build Nuget Packages' inputs: diff --git a/tools/ci_build/github/azure-pipelines/templates/ondevice-training-cpu-packaging-pipeline.yml b/tools/ci_build/github/azure-pipelines/templates/ondevice-training-cpu-packaging-pipeline.yml index 6fead4181ffa2..bc75a115326f6 100644 --- a/tools/ci_build/github/azure-pipelines/templates/ondevice-training-cpu-packaging-pipeline.yml +++ b/tools/ci_build/github/azure-pipelines/templates/ondevice-training-cpu-packaging-pipeline.yml @@ -228,6 +228,11 @@ stages: DisplayName: 'ESRP - Sign C# dlls' DoEsrp: ${{ parameters.DoEsrp }} + - task: UsePythonVersion@0 + displayName: 'Use Python' + inputs: + versionSpec: 3.8 + - task: MSBuild@1 displayName: 'Build Nuget Packages' inputs: From 55f7f9d7a9b88c4e7f0eb7cf4d7f31004761f5cb Mon Sep 17 00:00:00 2001 From: ivberg Date: Thu, 20 Jun 2024 06:45:45 -0700 Subject: [PATCH 09/52] Fix Crash When Enabling and Disabling ETW with Old Callbacks (#21086) ### Description Under certain conditions with enabling & disabling ETW continuously, we got a crash report. Allows ETW callbacks to be de-registered upon class destructor. Related to #20537 ### Motivation and Context Fixes crash ### Callstack We see it crash in [0x0] onnxruntime!::operator()+0x34 0x12941ff570 0x7ffa994f0a04 [0x1] onnxruntime!std::_Func_class::operator()+0x54 0x12941ff7b0 0x7ffa994f0d64 [0x2] onnxruntime!onnxruntime::logging::EtwRegistrationManager::InvokeCallbacks+0xcc 0x12941ff7b0 0x7ffa994f0d64 [0x3] onnxruntime!onnxruntime::logging::EtwRegistrationManager::ORT_TL_EtwEnableCallback+0x94 0x12941ff860 0x7ffa98d19628 and seems to us that the this pointer captured in etwRegistrationManager.RegisterInternalCallback( [&etwRegistrationManager, this]( ... is no longer valid when the callback is called. --- .../core/framework/execution_providers.h | 58 ++++++++++++------- .../core/platform/windows/logging/etw_sink.cc | 15 ++++- .../core/platform/windows/logging/etw_sink.h | 4 +- .../core/platform/windows/telemetry.cc | 22 +++++-- onnxruntime/core/platform/windows/telemetry.h | 4 +- .../providers/qnn/qnn_execution_provider.cc | 8 ++- .../providers/qnn/qnn_execution_provider.h | 6 ++ onnxruntime/core/session/inference_session.cc | 45 +++++++------- onnxruntime/core/session/inference_session.h | 6 ++ 9 files changed, 116 insertions(+), 52 deletions(-) diff --git a/onnxruntime/core/framework/execution_providers.h b/onnxruntime/core/framework/execution_providers.h index dc45cad692b6e..43fe92edc9dfe 100644 --- a/onnxruntime/core/framework/execution_providers.h +++ b/onnxruntime/core/framework/execution_providers.h @@ -25,30 +25,10 @@ Class for managing lookup of the execution providers in a session. */ class ExecutionProviders { public: - ExecutionProviders() = default; - - common::Status Add(const std::string& provider_id, const std::shared_ptr& p_exec_provider) { - // make sure there are no issues before we change any internal data structures - if (provider_idx_map_.find(provider_id) != provider_idx_map_.end()) { - auto status = ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Provider ", provider_id, " has already been registered."); - LOGS_DEFAULT(ERROR) << status.ErrorMessage(); - return status; - } - - // index that provider will have after insertion - auto new_provider_idx = exec_providers_.size(); - - ORT_IGNORE_RETURN_VALUE(provider_idx_map_.insert({provider_id, new_provider_idx})); - - // update execution provider options - auto providerOptions = p_exec_provider->GetProviderOptions(); - exec_provider_options_[provider_id] = providerOptions; - + ExecutionProviders() { #ifdef _WIN32 - LogProviderOptions(provider_id, providerOptions, false); - // Register callback for ETW capture state (rundown) - WindowsTelemetry::RegisterInternalCallback( + etw_callback_ = onnxruntime::WindowsTelemetry::EtwInternalCallback( [this]( LPCGUID SourceId, ULONG IsEnabled, @@ -79,6 +59,36 @@ class ExecutionProviders { } } }); + WindowsTelemetry::RegisterInternalCallback(etw_callback_); +#endif + } + + ~ExecutionProviders() { +#ifdef _WIN32 + WindowsTelemetry ::UnregisterInternalCallback(etw_callback_); +#endif + } + + common::Status + Add(const std::string& provider_id, const std::shared_ptr& p_exec_provider) { + // make sure there are no issues before we change any internal data structures + if (provider_idx_map_.find(provider_id) != provider_idx_map_.end()) { + auto status = ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Provider ", provider_id, " has already been registered."); + LOGS_DEFAULT(ERROR) << status.ErrorMessage(); + return status; + } + + // index that provider will have after insertion + auto new_provider_idx = exec_providers_.size(); + + ORT_IGNORE_RETURN_VALUE(provider_idx_map_.insert({provider_id, new_provider_idx})); + + // update execution provider options + auto providerOptions = p_exec_provider->GetProviderOptions(); + exec_provider_options_[provider_id] = providerOptions; + +#ifdef _WIN32 + LogProviderOptions(provider_id, providerOptions, false); #endif exec_provider_ids_.push_back(provider_id); @@ -156,5 +166,9 @@ class ExecutionProviders { // Whether the CPU provider was implicitly added to a session for fallback (true), // or whether it was explicitly added by the caller. bool cpu_execution_provider_was_implicitly_added_ = false; + +#ifdef _WIN32 + WindowsTelemetry::EtwInternalCallback etw_callback_; +#endif }; } // namespace onnxruntime diff --git a/onnxruntime/core/platform/windows/logging/etw_sink.cc b/onnxruntime/core/platform/windows/logging/etw_sink.cc index 5fb7f7a65161d..b0f9eaf4f62d2 100644 --- a/onnxruntime/core/platform/windows/logging/etw_sink.cc +++ b/onnxruntime/core/platform/windows/logging/etw_sink.cc @@ -104,7 +104,16 @@ HRESULT EtwRegistrationManager::Status() const { void EtwRegistrationManager::RegisterInternalCallback(const EtwInternalCallback& callback) { std::lock_guard lock(callbacks_mutex_); - callbacks_.push_back(callback); + callbacks_.push_back(&callback); +} + +void EtwRegistrationManager::UnregisterInternalCallback(const EtwInternalCallback& callback) { + std::lock_guard lock(callbacks_mutex_); + auto new_end = std::remove_if(callbacks_.begin(), callbacks_.end(), + [&callback](const EtwInternalCallback* ptr) { + return ptr == &callback; + }); + callbacks_.erase(new_end, callbacks_.end()); } void NTAPI EtwRegistrationManager::ORT_TL_EtwEnableCallback( @@ -126,6 +135,8 @@ void NTAPI EtwRegistrationManager::ORT_TL_EtwEnableCallback( } EtwRegistrationManager::~EtwRegistrationManager() { + std::lock_guard lock(callbacks_mutex_); + callbacks_.clear(); ::TraceLoggingUnregister(etw_provider_handle); } @@ -150,7 +161,7 @@ void EtwRegistrationManager::InvokeCallbacks(LPCGUID SourceId, ULONG IsEnabled, PVOID CallbackContext) { std::lock_guard lock(callbacks_mutex_); for (const auto& callback : callbacks_) { - callback(SourceId, IsEnabled, Level, MatchAnyKeyword, MatchAllKeyword, FilterData, CallbackContext); + (*callback)(SourceId, IsEnabled, Level, MatchAnyKeyword, MatchAllKeyword, FilterData, CallbackContext); } } diff --git a/onnxruntime/core/platform/windows/logging/etw_sink.h b/onnxruntime/core/platform/windows/logging/etw_sink.h index 5d35d101f1242..3af45b813a625 100644 --- a/onnxruntime/core/platform/windows/logging/etw_sink.h +++ b/onnxruntime/core/platform/windows/logging/etw_sink.h @@ -71,6 +71,8 @@ class EtwRegistrationManager { void RegisterInternalCallback(const EtwInternalCallback& callback); + void UnregisterInternalCallback(const EtwInternalCallback& callback); + private: EtwRegistrationManager(); ~EtwRegistrationManager(); @@ -90,7 +92,7 @@ class EtwRegistrationManager { _In_opt_ PEVENT_FILTER_DESCRIPTOR FilterData, _In_opt_ PVOID CallbackContext); - std::vector callbacks_; + std::vector callbacks_; OrtMutex callbacks_mutex_; mutable OrtMutex provider_change_mutex_; OrtMutex init_mutex_; diff --git a/onnxruntime/core/platform/windows/telemetry.cc b/onnxruntime/core/platform/windows/telemetry.cc index 850f40e846248..86067d377205b 100644 --- a/onnxruntime/core/platform/windows/telemetry.cc +++ b/onnxruntime/core/platform/windows/telemetry.cc @@ -64,7 +64,7 @@ bool WindowsTelemetry::enabled_ = true; uint32_t WindowsTelemetry::projection_ = 0; UCHAR WindowsTelemetry::level_ = 0; UINT64 WindowsTelemetry::keyword_ = 0; -std::vector WindowsTelemetry::callbacks_; +std::vector WindowsTelemetry::callbacks_; OrtMutex WindowsTelemetry::callbacks_mutex_; WindowsTelemetry::WindowsTelemetry() { @@ -86,6 +86,9 @@ WindowsTelemetry::~WindowsTelemetry() { TraceLoggingUnregister(telemetry_provider_handle); } } + + std::lock_guard lock_callbacks(callbacks_mutex_); + callbacks_.clear(); } bool WindowsTelemetry::IsEnabled() const { @@ -108,8 +111,17 @@ UINT64 WindowsTelemetry::Keyword() const { // } void WindowsTelemetry::RegisterInternalCallback(const EtwInternalCallback& callback) { - std::lock_guard lock(callbacks_mutex_); - callbacks_.push_back(callback); + std::lock_guard lock_callbacks(callbacks_mutex_); + callbacks_.push_back(&callback); +} + +void WindowsTelemetry::UnregisterInternalCallback(const EtwInternalCallback& callback) { + std::lock_guard lock_callbacks(callbacks_mutex_); + auto new_end = std::remove_if(callbacks_.begin(), callbacks_.end(), + [&callback](const EtwInternalCallback* ptr) { + return ptr == &callback; + }); + callbacks_.erase(new_end, callbacks_.end()); } void NTAPI WindowsTelemetry::ORT_TL_EtwEnableCallback( @@ -131,9 +143,9 @@ void NTAPI WindowsTelemetry::ORT_TL_EtwEnableCallback( void WindowsTelemetry::InvokeCallbacks(LPCGUID SourceId, ULONG IsEnabled, UCHAR Level, ULONGLONG MatchAnyKeyword, ULONGLONG MatchAllKeyword, PEVENT_FILTER_DESCRIPTOR FilterData, PVOID CallbackContext) { - std::lock_guard lock(callbacks_mutex_); + std::lock_guard lock_callbacks(callbacks_mutex_); for (const auto& callback : callbacks_) { - callback(SourceId, IsEnabled, Level, MatchAnyKeyword, MatchAllKeyword, FilterData, CallbackContext); + (*callback)(SourceId, IsEnabled, Level, MatchAnyKeyword, MatchAllKeyword, FilterData, CallbackContext); } } diff --git a/onnxruntime/core/platform/windows/telemetry.h b/onnxruntime/core/platform/windows/telemetry.h index 27cd20c2d21d1..ed80f13e633ac 100644 --- a/onnxruntime/core/platform/windows/telemetry.h +++ b/onnxruntime/core/platform/windows/telemetry.h @@ -66,13 +66,15 @@ class WindowsTelemetry : public Telemetry { static void RegisterInternalCallback(const EtwInternalCallback& callback); + static void UnregisterInternalCallback(const EtwInternalCallback& callback); + private: static OrtMutex mutex_; static uint32_t global_register_count_; static bool enabled_; static uint32_t projection_; - static std::vector callbacks_; + static std::vector callbacks_; static OrtMutex callbacks_mutex_; static OrtMutex provider_change_mutex_; static UCHAR level_; diff --git a/onnxruntime/core/providers/qnn/qnn_execution_provider.cc b/onnxruntime/core/providers/qnn/qnn_execution_provider.cc index d4616e0cefbf2..c159730d46cf1 100644 --- a/onnxruntime/core/providers/qnn/qnn_execution_provider.cc +++ b/onnxruntime/core/providers/qnn/qnn_execution_provider.cc @@ -233,7 +233,7 @@ QNNExecutionProvider::QNNExecutionProvider(const ProviderOptions& provider_optio #ifdef _WIN32 auto& etwRegistrationManager = logging::EtwRegistrationManager::Instance(); // Register callback for ETW capture state (rundown) - etwRegistrationManager.RegisterInternalCallback( + callback_ETWSink_provider_ = onnxruntime::logging::EtwRegistrationManager::EtwInternalCallback( [&etwRegistrationManager, this]( LPCGUID SourceId, ULONG IsEnabled, @@ -270,6 +270,7 @@ QNNExecutionProvider::QNNExecutionProvider(const ProviderOptions& provider_optio (void)qnn_backend_manager_->ResetQnnLogLevel(); } }); + etwRegistrationManager.RegisterInternalCallback(callback_ETWSink_provider_); #endif // In case ETW gets disabled later @@ -397,6 +398,11 @@ QNNExecutionProvider::~QNNExecutionProvider() { if (!cache) continue; ORT_IGNORE_RETURN_VALUE(cache->erase(this)); } + + // Unregister the ETW callback +#ifdef _WIN32 + logging::EtwRegistrationManager::Instance().UnregisterInternalCallback(callback_ETWSink_provider_); +#endif } bool QNNExecutionProvider::IsNodeSupported(qnn::QnnModelWrapper& qnn_model_wrapper, const NodeUnit& node_unit, diff --git a/onnxruntime/core/providers/qnn/qnn_execution_provider.h b/onnxruntime/core/providers/qnn/qnn_execution_provider.h index c5d3098f87b3a..e7419dabb14d1 100644 --- a/onnxruntime/core/providers/qnn/qnn_execution_provider.h +++ b/onnxruntime/core/providers/qnn/qnn_execution_provider.h @@ -15,6 +15,9 @@ #include #include #include +#ifdef _WIN32 +#include "core/platform/windows/logging/etw_sink.h" +#endif namespace onnxruntime { @@ -86,6 +89,9 @@ class QNNExecutionProvider : public IExecutionProvider { qnn::HtpPerformanceMode default_htp_performance_mode_ = qnn::HtpPerformanceMode::kHtpDefault; uint32_t default_rpc_control_latency_ = 0; bool enable_HTP_FP16_precision_ = false; +#ifdef _WIN32 + onnxruntime::logging::EtwRegistrationManager::EtwInternalCallback callback_ETWSink_provider_; +#endif class PerThreadContext final { public: diff --git a/onnxruntime/core/session/inference_session.cc b/onnxruntime/core/session/inference_session.cc index e8d33bc154b0c..03049c4b51c9c 100644 --- a/onnxruntime/core/session/inference_session.cc +++ b/onnxruntime/core/session/inference_session.cc @@ -250,6 +250,7 @@ std::atomic InferenceSession::global_session_id_{1}; std::map InferenceSession::active_sessions_; #ifdef _WIN32 OrtMutex InferenceSession::active_sessions_mutex_; // Protects access to active_sessions_ +onnxruntime::WindowsTelemetry::EtwInternalCallback InferenceSession::callback_ML_ORT_provider_; #endif static Status FinalizeSessionOptions(const SessionOptions& user_provided_session_options, @@ -374,15 +375,14 @@ void InferenceSession::ConstructorCommon(const SessionOptions& session_options, active_sessions_[global_session_id_++] = this; // Register callback for ETW capture state (rundown) for Microsoft.ML.ONNXRuntime provider - WindowsTelemetry::RegisterInternalCallback( - [this]( - LPCGUID SourceId, - ULONG IsEnabled, - UCHAR Level, - ULONGLONG MatchAnyKeyword, - ULONGLONG MatchAllKeyword, - PEVENT_FILTER_DESCRIPTOR FilterData, - PVOID CallbackContext) { + callback_ML_ORT_provider_ = onnxruntime::WindowsTelemetry::EtwInternalCallback( + [this](LPCGUID SourceId, + ULONG IsEnabled, + UCHAR Level, + ULONGLONG MatchAnyKeyword, + ULONGLONG MatchAllKeyword, + PEVENT_FILTER_DESCRIPTOR FilterData, + PVOID CallbackContext) { (void)SourceId; (void)Level; (void)MatchAnyKeyword; @@ -396,19 +396,18 @@ void InferenceSession::ConstructorCommon(const SessionOptions& session_options, LogAllSessions(); } }); + WindowsTelemetry::RegisterInternalCallback(callback_ML_ORT_provider_); // Register callback for ETW start / stop so that LOGS tracing can be adjusted dynamically after session start auto& etwRegistrationManager = logging::EtwRegistrationManager::Instance(); - // Register callback for ETW capture state (rundown) - etwRegistrationManager.RegisterInternalCallback( - [&etwRegistrationManager, this]( - LPCGUID SourceId, - ULONG IsEnabled, - UCHAR Level, - ULONGLONG MatchAnyKeyword, - ULONGLONG MatchAllKeyword, - PEVENT_FILTER_DESCRIPTOR FilterData, - PVOID CallbackContext) { + callback_ETWSink_provider_ = onnxruntime::logging::EtwRegistrationManager::EtwInternalCallback( + [&etwRegistrationManager, this](LPCGUID SourceId, + ULONG IsEnabled, + UCHAR Level, + ULONGLONG MatchAnyKeyword, + ULONGLONG MatchAllKeyword, + PEVENT_FILTER_DESCRIPTOR FilterData, + PVOID CallbackContext) { (void)SourceId; (void)Level; (void)MatchAnyKeyword; @@ -439,6 +438,10 @@ void InferenceSession::ConstructorCommon(const SessionOptions& session_options, } } }); + + // Register callback for ETW capture state (rundown) + etwRegistrationManager.RegisterInternalCallback(callback_ETWSink_provider_); + #endif SetLoggingManager(session_options, session_env); @@ -720,9 +723,11 @@ InferenceSession::~InferenceSession() { } } - // Unregister the session + // Unregister the session and ETW callbacks #ifdef _WIN32 std::lock_guard lock(active_sessions_mutex_); + WindowsTelemetry::UnregisterInternalCallback(callback_ML_ORT_provider_); + logging::EtwRegistrationManager::Instance().UnregisterInternalCallback(callback_ETWSink_provider_); #endif active_sessions_.erase(global_session_id_); diff --git a/onnxruntime/core/session/inference_session.h b/onnxruntime/core/session/inference_session.h index 204b24974ff50..77fba90b56b1e 100644 --- a/onnxruntime/core/session/inference_session.h +++ b/onnxruntime/core/session/inference_session.h @@ -35,6 +35,10 @@ #include "core/platform/tracing.h" #include #endif +#ifdef _WIN32 +#include "core/platform/windows/logging/etw_sink.h" +#include "core/platform/windows/telemetry.h" +#endif namespace ONNX_NAMESPACE { class ModelProto; @@ -124,6 +128,8 @@ class InferenceSession { static std::map active_sessions_; #ifdef _WIN32 static OrtMutex active_sessions_mutex_; // Protects access to active_sessions_ + static onnxruntime::WindowsTelemetry::EtwInternalCallback callback_ML_ORT_provider_; + onnxruntime::logging::EtwRegistrationManager::EtwInternalCallback callback_ETWSink_provider_; #endif public: From 0c80cd2157ae0a74df2280a0fdde458430ee8db1 Mon Sep 17 00:00:00 2001 From: Wanming Lin Date: Fri, 21 Jun 2024 02:04:01 +0800 Subject: [PATCH 10/52] [WebNN EP] Update Prelu restriction for CPU backend (#20878) --- js/web/docs/webnn-operators.md | 2 +- .../webnn/builders/impl/binary_op_builder.cc | 19 ++++++++++++++----- 2 files changed, 15 insertions(+), 6 deletions(-) diff --git a/js/web/docs/webnn-operators.md b/js/web/docs/webnn-operators.md index 4c6dab84fa973..508f85377a3a2 100644 --- a/js/web/docs/webnn-operators.md +++ b/js/web/docs/webnn-operators.md @@ -59,7 +59,7 @@ operators and the supported opset domain/versions in **WebNN EP** by ONNX Runtim | Not | ai.onnx(7+) | logicalnot | ✓ | ✓ | | | Pad | ai.onnx(7-10, 11-12, 13-17, 18, 19-20, 21+) | pad | ✓ | ✓ | modes == 'wrap' is not supported | | Pow | ai.onnx(7-11, 12, 13-14, 15+) | pow | ✓ | ✓ | | -| PRelu | ai.onnx(7-8, 9-15, 16+) | prelu | ✓ | ✓ | WebNN CPU restricts slope to be a static value | +| PRelu | ai.onnx(7-8, 9-15, 16+) | prelu | ✓ | ✓ | WebNN CPU backend restricts the last dimension of input and slope to be same (Chromium issue: https://issues.chromium.org/issues/335517470) | | Reciprocal | ai.onnx(7-12, 13+) | reciprocal | ✗ | ✓ | | | ReduceL1 | ai.onnx(7-10, 11-12, 13-17, 18+) | reduceL1 | ✗ | ✓ | Input 'axes' if present should be a constant | | ReduceL2 | ai.onnx(7-10, 11-12, 13-17, 18+) | reduceL2 | ✗ | ✓ | Input 'axes' if present should be a constant | diff --git a/onnxruntime/core/providers/webnn/builders/impl/binary_op_builder.cc b/onnxruntime/core/providers/webnn/builders/impl/binary_op_builder.cc index 2c97ef490f4a3..23e19d5943144 100644 --- a/onnxruntime/core/providers/webnn/builders/impl/binary_op_builder.cc +++ b/onnxruntime/core/providers/webnn/builders/impl/binary_op_builder.cc @@ -63,14 +63,23 @@ bool BinaryOpBuilder::IsOpSupportedImpl(const InitializedTensorSet& initializers const auto& input_defs = node.InputDefs(); const auto& op_type = node.OpType(); - // XNNPACK prelu operator expects slope to be a static value. - // https://github.com/google/XNNPACK/issues/4692 - // TODO: Remove this check after it is solved. - if (op_type == "PRelu" && !Contains(initializers, input_defs[1]->Name()) && device_type == WebnnDeviceType::CPU) { - LOGS(logger, VERBOSE) << "The second input (slope) for PRelu must be a constant initializer for WebNN CPU backend."; + std::vector input0_shape; + std::vector input1_shape; + if (!GetShape(*input_defs[0], input0_shape, logger) || + !GetShape(*input_defs[1], input1_shape, logger)) { return false; } + // 'prelu' op in WebNN CPU backend restricts the last dimension of input and slope to be same. + // TODO: Remove this workaround once the associated issue is resolved in Chromium: + // https://issues.chromium.org/issues/335517470. + if (op_type == "PRelu" && device_type == WebnnDeviceType::CPU) { + if (input0_shape.back() != input1_shape.back()) { + LOGS(logger, VERBOSE) << "The last dimension of input and slope for PRelu must be same for WebNN CPU backend."; + return false; + } + } + return true; } From 00c713088dfa5cd8335e5104186cf409c0df7511 Mon Sep 17 00:00:00 2001 From: Yi-Hong Lyu Date: Thu, 20 Jun 2024 11:21:32 -0700 Subject: [PATCH 11/52] Adpot QDQFinalCleanupTransformer for Q->DQs/DQ->Qs cases (#21018) ### Description ### Motivation and Context --- .../qdq_transformer/qdq_final_cleanup.cc | 147 +++++----- .../optimizer/qdq_transformer/qdq_util.cc | 9 +- .../core/optimizer/qdq_transformer/qdq_util.h | 3 +- .../test/optimizer/qdq_transformer_test.cc | 252 +++++++++++++++++- 4 files changed, 335 insertions(+), 76 deletions(-) diff --git a/onnxruntime/core/optimizer/qdq_transformer/qdq_final_cleanup.cc b/onnxruntime/core/optimizer/qdq_transformer/qdq_final_cleanup.cc index 49ce06e47c98a..507bc71709b2f 100644 --- a/onnxruntime/core/optimizer/qdq_transformer/qdq_final_cleanup.cc +++ b/onnxruntime/core/optimizer/qdq_transformer/qdq_final_cleanup.cc @@ -1,10 +1,11 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -#include "core/optimizer/qdq_transformer/qdq_final_cleanup.h" +#include #include "core/graph/graph_utils.h" #include "core/optimizer/initializer.h" +#include "core/optimizer/qdq_transformer/qdq_final_cleanup.h" #include "core/optimizer/qdq_transformer/qdq_util.h" #include "core/optimizer/selectors_actions/actions.h" #include "core/optimizer/utils.h" @@ -31,98 +32,106 @@ bool CleanUpNodeSequence(NodeSequence node_sequence_type, Graph& graph, NodeInde if (!match_first(first_node) || // not filtering on provider currently // !graph_utils::IsSupportedProvider(first_node, compatible_execution_providers) || - !optimizer_utils::CheckOutputEdges(graph, first_node, 1)) { + !(first_node.GetOutputEdgesCount() >= 1)) { return false; } - Node& second_node = *graph.GetNode(first_node.OutputNodesBegin()->Index()); - if (!match_second(second_node) - // not filtering on provider currently - // || !graph_utils::IsSupportedProvider(second_node, compatible_execution_providers) - ) { - return false; + std::vector second_node_ptrs; + for (auto node_it = first_node.OutputNodesBegin(); node_it != first_node.OutputNodesEnd(); ++node_it) { + second_node_ptrs.push_back(graph.GetNode(node_it->Index())); } - if (node_sequence_type == NodeSequence::DQ_Q) { - // for DQ -> Q, check for constant, matching scale/ZP values + for (auto second_node_ptr : second_node_ptrs) { + // check for constant, matching scale/ZP values const auto get_constant_initializer = [&graph](const std::string& initializer_name) { return graph.GetConstantInitializer(initializer_name, true); }; - if (!QDQ::IsQDQPairSupported(second_node, first_node, get_constant_initializer, graph.ModelPath())) { + const bool produces_graph_output = graph.NodeProducesGraphOutput(*second_node_ptr); + const auto output_edges_count = second_node_ptr->GetOutputEdgesCount(); + + if (!match_second(*second_node_ptr) || + !QDQ::IsQDQPairSupported(first_node, *second_node_ptr, get_constant_initializer, graph.ModelPath(), false) || + (produces_graph_output && output_edges_count != 0) || + (!produces_graph_output && output_edges_count != 1)) { return false; } } // we have a node sequence to clean up - - // we support a second_node that produces a graph output if it has no output edges, or a second_node with one output edge. - const bool produces_graph_output = graph.NodeProducesGraphOutput(second_node); - const auto output_edges_count = second_node.GetOutputEdgesCount(); - - if ((produces_graph_output && output_edges_count != 0) || - (!produces_graph_output && output_edges_count != 1)) { - return false; + if (logger.GetSeverity() == logging::Severity::kVERBOSE) { + LOGS(logger, VERBOSE) << "Found back-to-back nodes: " + << first_node.OpType() << " with name \"" << first_node.Name() << "\""; + for (auto& second_node_ptr : second_node_ptrs) { + LOGS(logger, VERBOSE) << ", " << second_node_ptr->OpType() << " with name \"" << second_node_ptr->Name() << "\""; + } } - LOGS(logger, VERBOSE) << "Cleaning up back-to-back nodes: " - << first_node.OpType() << " with name \"" << first_node.Name() << "\" and " - << second_node.OpType() << " with name \"" << second_node.Name() << "\""; - - // src node or graph input/initializer -> first_node -> second_node -> downstream node or graph output - NodeIndex src_node_idx = 0; - int src_arg_idx = -1; - NodeIndex downstream_node_idx = 0; - int downstream_arg_idx = -1; - - // input could be node or initializer/graph input so need to handle both. - // if it's a node we need to replace the edge, so need info on which output idx it was attached to on the src node. - const Node::EdgeEnd* input_edge = nullptr; - if (first_node.GetInputEdgesCount() == 1) { - input_edge = &*first_node.InputEdgesBegin(); - src_node_idx = input_edge->GetNode().Index(); - src_arg_idx = input_edge->GetSrcArgIndex(); - // remove edge from src to first_node. dest arg idx is 0 as first_node (Q or DQ) only has one input - graph.RemoveEdge(src_node_idx, first_node.Index(), src_arg_idx, 0); - } + for (auto second_node_ptr : second_node_ptrs) { + Node& second_node = *graph.GetNode(second_node_ptr->Index()); + // we support a second_node that produces a graph output if it has no output edges, or a second_node with one + // output edge. + const bool produces_graph_output = graph.NodeProducesGraphOutput(second_node); + + // src node or graph input/initializer -> first_node -> second_node -> downstream node or graph output + NodeIndex src_node_idx = 0; + int src_arg_idx = -1; + NodeIndex downstream_node_idx = 0; + int downstream_arg_idx = -1; + + // input could be node or initializer/graph input so need to handle both. + // if it's a node we need to replace the edge, so need info on which output idx it was attached to on the src node. + const Node::EdgeEnd* input_edge = nullptr; + if (first_node.GetInputEdgesCount() == 1) { + input_edge = &*first_node.InputEdgesBegin(); + src_node_idx = input_edge->GetNode().Index(); + src_arg_idx = input_edge->GetSrcArgIndex(); + // remove edge from src to first_node. dest arg idx is 0 as first_node (Q or DQ) only has one input + if (second_node_ptr == second_node_ptrs.back()) { + graph.RemoveEdge(src_node_idx, first_node.Index(), src_arg_idx, 0); + } + } - // remove edge between pair we're removing - // both DQ and Q are single input single output so src idx and dest idx must be 0 - graph.RemoveEdge(first_node.Index(), second_node.Index(), 0, 0); + // remove edge between pair we're removing + // both DQ and Q are single input single output so src idx and dest idx must be 0 + graph.RemoveEdge(first_node.Index(), second_node.Index(), 0, 0); - if (!produces_graph_output) { - // remove edge to downstream node - const Node::EdgeEnd& output_edge = *second_node.OutputEdgesBegin(); - downstream_node_idx = output_edge.GetNode().Index(); - downstream_arg_idx = output_edge.GetDstArgIndex(); + if (!produces_graph_output) { + // remove edge to downstream node + const Node::EdgeEnd& output_edge = *second_node.OutputEdgesBegin(); + downstream_node_idx = output_edge.GetNode().Index(); + downstream_arg_idx = output_edge.GetDstArgIndex(); - // source arg idx is 0 as Q/DQ only has one output - graph.RemoveEdge(second_node.Index(), downstream_node_idx, 0, downstream_arg_idx); + // source arg idx is 0 as Q/DQ only has one output + graph.RemoveEdge(second_node.Index(), downstream_node_idx, 0, downstream_arg_idx); - // replace input on downstream node - Node& downstream_node = *graph.GetNode(downstream_node_idx); - downstream_node.MutableInputDefs()[downstream_arg_idx] = first_node.MutableInputDefs()[0]; + // replace input on downstream node + Node& downstream_node = *graph.GetNode(downstream_node_idx); + downstream_node.MutableInputDefs()[downstream_arg_idx] = first_node.MutableInputDefs()[0]; - // create edge between src_node (if available) and downstream node - if (input_edge) { - graph.AddEdge(src_node_idx, downstream_node_idx, src_arg_idx, downstream_arg_idx); - } - } else { - NodeArg* graph_output_nodearg = second_node.MutableOutputDefs()[0]; - if (src_arg_idx >= 0) { - // update the src node to produce the graph output that was being provided by second_node - Node& src_node = *graph.GetNode(src_node_idx); - src_node.MutableOutputDefs()[src_arg_idx] = graph_output_nodearg; + // create edge between src_node (if available) and downstream node + if (input_edge) { + graph.AddEdge(src_node_idx, downstream_node_idx, src_arg_idx, downstream_arg_idx); + } } else { - // add Identity node to connect the graph input or initializer to the graph output. - Node& id_node = graph.AddNode(graph.GenerateNodeName("QDQFinalCleanupTransformer"), - "Identity", "", {first_node.MutableInputDefs()[0]}, {graph_output_nodearg}); - id_node.SetExecutionProviderType(second_node.GetExecutionProviderType()); + NodeArg* graph_output_nodearg = second_node.MutableOutputDefs()[0]; + if (src_arg_idx >= 0 && second_node_ptrs.size() == 1) { + // update the src node to produce the graph output that was being provided by second_node + Node& src_node = *graph.GetNode(src_node_idx); + src_node.MutableOutputDefs()[src_arg_idx] = graph_output_nodearg; + } else { + // add Identity node to connect the graph input or initializer to the graph output. + Node& id_node = graph.AddNode(graph.GenerateNodeName("QDQFinalCleanupTransformer"), + "Identity", "", {first_node.MutableInputDefs()[0]}, {graph_output_nodearg}); + id_node.SetExecutionProviderType(second_node.GetExecutionProviderType()); + } } - } - graph.RemoveNode(first_node.Index()); - graph.RemoveNode(second_node.Index()); + if (second_node_ptr == second_node_ptrs.back()) { + graph.RemoveNode(first_node.Index()); + } + graph.RemoveNode(second_node.Index()); + } return true; } diff --git a/onnxruntime/core/optimizer/qdq_transformer/qdq_util.cc b/onnxruntime/core/optimizer/qdq_transformer/qdq_util.cc index e245636ce9a84..5a6c47a8d8454 100644 --- a/onnxruntime/core/optimizer/qdq_transformer/qdq_util.cc +++ b/onnxruntime/core/optimizer/qdq_transformer/qdq_util.cc @@ -17,7 +17,14 @@ namespace onnxruntime::QDQ { bool IsQDQPairSupported( const Node& q_node, const Node& dq_node, const GetConstantInitializerFn& get_const_initializer, - const Path& model_path) { + const Path& model_path, + bool check_op_type) { + if (check_op_type) { + if (!MatchQNode(q_node) || !MatchDQNode(dq_node)) { + return false; + } + } + ConstPointerContainer> dq_input_defs = dq_node.InputDefs(); ConstPointerContainer> q_input_defs = q_node.InputDefs(); diff --git a/onnxruntime/core/optimizer/qdq_transformer/qdq_util.h b/onnxruntime/core/optimizer/qdq_transformer/qdq_util.h index 8333168b0093f..c5f7cd601a2f0 100644 --- a/onnxruntime/core/optimizer/qdq_transformer/qdq_util.h +++ b/onnxruntime/core/optimizer/qdq_transformer/qdq_util.h @@ -36,7 +36,8 @@ using GetConstantInitializerFn = std::function Q sequence represents a conversion in quantization data type. // Example of uint8 to uint16: diff --git a/onnxruntime/test/optimizer/qdq_transformer_test.cc b/onnxruntime/test/optimizer/qdq_transformer_test.cc index c338d542b0b79..1c77121ba9df1 100644 --- a/onnxruntime/test/optimizer/qdq_transformer_test.cc +++ b/onnxruntime/test/optimizer/qdq_transformer_test.cc @@ -3675,7 +3675,9 @@ TEST(QDQTransformerTests, QDQFinalCleanupTransformer_BasicQDQCleanup) { }; // if we block removal of the DQ node the Q node in the pair will not be removed either - const int expected_qdq_count = 0 + (block_removal_of_first_dq ? 1 : 0) + (block_removal_of_last_dq ? 1 : 0); + // TODO(yilyu): block_removal_of_first_dq is not functional, need to fix it + // TODO(yilyu): block_removal_of_last_dq is not functional, need to fix it + const int expected_qdq_count = 0; // + (block_removal_of_first_dq ? 1 : 0) + (block_removal_of_last_dq ? 1 : 0); // blocking removal of DQ by adding an additional edge will cause EnsureUniqueDQForNodeUnit to duplicate the DQ, // so we expect twice as many DQ's as original QDQ pairs const int expected_dq_count = expected_qdq_count * 2; @@ -3725,19 +3727,136 @@ TEST(QDQTransformerTests, QDQFinalCleanupTransformer_BasicQDQCleanup) { }; test_case({{1, 2, 4}, {1, 3, 4}}, false, false); // Do not block removal - test_case({{1, 2, 4}, {1, 3, 4}}, true, false); // Block removal of first dq - test_case({{1, 2, 4}, {1, 3, 4}}, false, true); // Block removal of last dq + test_case({{1, 2, 4}, {1, 3, 4}}, true, false); // Block removal of last dq + test_case({{1, 2, 4}, {1, 3, 4}}, false, true); // Block removal of first dq test_case({{1, 2, 4}, {1, 3, 4}}, true, true); // Block removal of first and last dq #if !defined(DISABLE_CONTRIB_OPS) // Use contrib QDQ ops test_case({{1, 2, 4}, {1, 3, 4}}, false, false, true); // Do not block removal - test_case({{1, 2, 4}, {1, 3, 4}}, true, false, true); // Block removal of first dq - test_case({{1, 2, 4}, {1, 3, 4}}, false, true, true); // Block removal of last dq + test_case({{1, 2, 4}, {1, 3, 4}}, true, false, true); // Block removal of last dq + test_case({{1, 2, 4}, {1, 3, 4}}, false, true, true); // Block removal of first dq test_case({{1, 2, 4}, {1, 3, 4}}, true, true, true); // Block removal of first and last dq #endif } +// test removal of Q->DQs pairs by QDQFinalCleanupTransformer +TEST(QDQTransformerTests, QDQFinalCleanupTransformer_BasicQDQsCleanup) { + auto test_case = [&](bool is_input_q, + bool is_dq1_output, + bool is_dq2_output, + bool use_contrib_qdq) { + // create model with float Input -> (Transpose1 ->) Q -> DQ1 -> (Transpose2 ->) Output1 + // -> DQ2 -> (Transpose3 ->) Output2 + // If we enable cleanup and don't run the QDQ transformer we should drop all the Q->DQ pairs + auto build_test_case = [&](ModelTestBuilder& builder) { + NodeArg* input_args = builder.MakeInput({1, 2, 4}, -1.f, 1.f); + + // Add Input -> (Transpose1 ->) + NodeArg* q_input = nullptr; + if (is_input_q) { + q_input = input_args; + } else { + auto* transpose_output = builder.MakeIntermediate(); + builder.AddNode("Transpose", {input_args}, {transpose_output}); + q_input = transpose_output; + } + + // Add Q -> + auto* q_output = builder.MakeIntermediate(); + builder.AddQuantizeLinearNode(q_input, 0.05f, 128, q_output, use_contrib_qdq); + + // Add DQ1 -> (Transpose2 ->) Output1 + if (is_dq1_output) { + auto* dq1_output = builder.MakeOutput(); + builder.AddDequantizeLinearNode(q_output, 0.05f, 128, dq1_output, use_contrib_qdq); + } else { + auto* dq1_output = builder.MakeIntermediate(); + builder.AddDequantizeLinearNode(q_output, 0.05f, 128, dq1_output, use_contrib_qdq); + auto* output1 = builder.MakeOutput(); + builder.AddNode("Transpose", {dq1_output}, {output1}); + } + + // Add DQ2 -> (Transpose3 ->) Output2 + if (is_dq2_output) { + auto* dq2_output = builder.MakeOutput(); + builder.AddDequantizeLinearNode(q_output, 0.05f, 128, dq2_output, use_contrib_qdq); + } else { + auto* dq2_output = builder.MakeIntermediate(); + builder.AddDequantizeLinearNode(q_output, 0.05f, 128, dq2_output, use_contrib_qdq); + auto* output2 = builder.MakeOutput(); + builder.AddNode("Transpose", {dq2_output}, {output2}); + } + }; + + const int expected_transpose_count = (is_input_q ? 0 : 1) + (is_dq1_output ? 0 : 1) + (is_dq2_output ? 0 : 1); + + auto check_graph = [expected_transpose_count, + use_contrib_qdq](InferenceSessionWrapper& session) { + auto op_to_count = CountOpsInGraph(session.GetGraph()); + const QDQOpKeys qdq_keys = GetQDQOpKeys(use_contrib_qdq); + EXPECT_EQ(op_to_count[qdq_keys.quantize_linear], 0); + EXPECT_EQ(op_to_count[qdq_keys.dequantize_linear], 0); + EXPECT_EQ(op_to_count["Transpose"], expected_transpose_count); + }; + + auto add_session_options = [](SessionOptions& so) { + ASSERT_STATUS_OK(so.config_options.AddConfigEntry(kOrtSessionOptionsEnableQuantQDQCleanup, "1")); + }; + + // we increase the tolerance as removing the QDQ nodes means there's no round-trip to 8-bit and back + // essentially rounding the input values. + TransformerTester(build_test_case, + check_graph, + TransformerLevel::Level1, + TransformerLevel::Level2, + 12 /*opset_version*/, + 0.025f /*per_sample_tolerance*/, + 0.01f /*relative_per_sample_tolerance*/, + std::make_unique(true /*enable_q_dq_cleanup*/), + add_session_options); + TransformerTester(build_test_case, + check_graph, + TransformerLevel::Level1, + TransformerLevel::Level2, + 18 /*opset_version*/, + 0.025f /*per_sample_tolerance*/, + 0.01f /*relative_per_sample_tolerance*/, + std::make_unique(true /*enable_q_dq_cleanup*/), + add_session_options); + TransformerTester(build_test_case, + check_graph, + TransformerLevel::Level1, + TransformerLevel::Level2, + 19 /*opset_version*/, + 0.025f /*per_sample_tolerance*/, + 0.01f /*relative_per_sample_tolerance*/, + std::make_unique(true /*enable_q_dq_cleanup*/), + add_session_options); + }; + + test_case(true, true, true, false); // is_input_q, is_dq1_output, is_dq2_output + test_case(false, true, true, false); // is_dq1_output, is_dq2_output + test_case(true, false, true, false); // is_input_q, is_dq2_output + test_case(false, false, true, false); // is_dq2_output + test_case(true, true, false, false); // is_input_q, is_dq1_output + test_case(false, true, false, false); // is_dq1_output + test_case(true, false, false, false); // is_input_q + test_case(false, false, false, false); + +#if !defined(DISABLE_CONTRIB_OPS) + // Use contrib QDQ ops + test_case(true, true, true, true); // is_input_q, is_dq1_output, is_dq2_output + test_case(false, true, true, true); // is_dq1_output, is_dq2_output + test_case(true, false, true, true); // is_input_q, is_dq2_output + test_case(false, false, true, true); // is_dq2_output + test_case(true, true, false, true); // is_input_q, is_dq1_output + test_case(false, true, false, true); // is_dq1_output + test_case(true, false, false, true); // is_input_q + test_case(false, false, false, true); +#endif +} + TEST(QDQTransformerTests, QDQFinalCleanupTransformer_BasicDQQCleanUp) { auto test_case = [](bool use_matching_qdq_params, bool use_contrib_qdq) { // input -> Q -> DQ -> Q -> DQ -> output @@ -3801,6 +3920,129 @@ TEST(QDQTransformerTests, QDQFinalCleanupTransformer_BasicDQQCleanUp) { #endif } +// test removal of DQ->Qs pairs by QDQFinalCleanupTransformer +TEST(QDQTransformerTests, QDQFinalCleanupTransformer_BasicDQQsCleanup) { + auto test_case = [&](bool is_input_dq, + bool is_q1_output, + bool is_q2_output, + bool use_contrib_qdq) { + // create model with float Input -> (Transpose1 ->) DQ -> Q1 -> (Transpose2 ->) Output1 + // -> Q2 -> (Transpose3 ->) Output2 + // If we enable cleanup and don't run the QDQ transformer we should drop all the DQ->Q pairs + auto build_test_case = [&](ModelTestBuilder& builder) { + NodeArg* input_args = builder.MakeInput({1, 2, 4}, + std::numeric_limits::min(), + std::numeric_limits::max()); + + // Add Input -> (Transpose1 ->) + NodeArg* dq_input = nullptr; + if (is_input_dq) { + dq_input = input_args; + } else { + auto* transpose_output = builder.MakeIntermediate(); + builder.AddNode("Transpose", {input_args}, {transpose_output}); + dq_input = transpose_output; + } + + // Add DQ -> + auto* dq_output = builder.MakeIntermediate(); + builder.AddDequantizeLinearNode(dq_input, 0.05f, 128, dq_output, use_contrib_qdq); + + // Add Q1 -> (Transpose2 ->) Output1 + if (is_q1_output) { + auto* q1_output = builder.MakeOutput(); + builder.AddQuantizeLinearNode(dq_output, 0.05f, 128, q1_output, use_contrib_qdq); + } else { + auto* q1_output = builder.MakeIntermediate(); + builder.AddQuantizeLinearNode(dq_output, 0.05f, 128, q1_output, use_contrib_qdq); + auto* output1 = builder.MakeOutput(); + builder.AddNode("Transpose", {q1_output}, {output1}); + } + + // Add Q2 -> (Transpose3 ->) Output2 + if (is_q2_output) { + auto* q2_output = builder.MakeOutput(); + builder.AddQuantizeLinearNode(dq_output, 0.05f, 128, q2_output, use_contrib_qdq); + } else { + auto* q2_output = builder.MakeIntermediate(); + builder.AddQuantizeLinearNode(dq_output, 0.05f, 128, q2_output, use_contrib_qdq); + auto* output2 = builder.MakeOutput(); + builder.AddNode("Transpose", {q2_output}, {output2}); + } + }; + + const int expected_transpose_count = (is_input_dq ? 0 : 1) + (is_q1_output ? 0 : 1) + (is_q2_output ? 0 : 1); + + auto check_graph = [expected_transpose_count, + use_contrib_qdq](InferenceSessionWrapper& session) { + auto op_to_count = CountOpsInGraph(session.GetGraph()); + const QDQOpKeys qdq_keys = GetQDQOpKeys(use_contrib_qdq); + EXPECT_EQ(op_to_count[qdq_keys.quantize_linear], 0); + EXPECT_EQ(op_to_count[qdq_keys.dequantize_linear], 0); + EXPECT_EQ(op_to_count["Transpose"], expected_transpose_count); + }; + + auto add_session_options = [](SessionOptions& so) { + // The function EnsureUniqueDQForEachExplicitOutputEdge does not account for this particular case. Disable it to + // prevent test failures. + ASSERT_STATUS_OK(so.config_options.AddConfigEntry(kOrtSessionOptionsDisableQuantQDQ, "1")); + + ASSERT_STATUS_OK(so.config_options.AddConfigEntry(kOrtSessionOptionsEnableQuantQDQCleanup, "1")); + }; + + // we increase the tolerance as removing the QDQ nodes means there's no round-trip to 8-bit and back + // essentially rounding the input values. + TransformerTester(build_test_case, + check_graph, + TransformerLevel::Level1, + TransformerLevel::Level2, + 12 /*opset_version*/, + 0.025f /*per_sample_tolerance*/, + 0.01f /*relative_per_sample_tolerance*/, + std::make_unique(true /*enable_q_dq_cleanup*/), + add_session_options); + TransformerTester(build_test_case, + check_graph, + TransformerLevel::Level1, + TransformerLevel::Level2, + 18 /*opset_version*/, + 0.025f /*per_sample_tolerance*/, + 0.01f /*relative_per_sample_tolerance*/, + std::make_unique(true /*enable_q_dq_cleanup*/), + add_session_options); + TransformerTester(build_test_case, + check_graph, + TransformerLevel::Level1, + TransformerLevel::Level2, + 19 /*opset_version*/, + 0.025f /*per_sample_tolerance*/, + 0.01f /*relative_per_sample_tolerance*/, + std::make_unique(true /*enable_q_dq_cleanup*/), + add_session_options); + }; + + test_case(true, true, true, false); // is_input_dq, is_q1_output, is_q2_output + test_case(false, true, true, false); // is_q1_output, is_q2_output + test_case(true, false, true, false); // is_input_dq, is_q2_output + test_case(false, false, true, false); // is_q2_output + test_case(true, true, false, false); // is_input_dq, is_q1_output + test_case(false, true, false, false); // is_q1_output + test_case(true, false, false, false); // is_input_dq + test_case(false, false, false, false); + +#if !defined(DISABLE_CONTRIB_OPS) + // Use contrib QDQ ops + test_case(true, true, true, true); // is_input_dq, is_q1_output, is_q2_output + test_case(false, true, true, true); // is_q1_output, is_q2_output + test_case(true, false, true, true); // is_input_dq, is_q2_output + test_case(false, false, true, true); // is_q2_output + test_case(true, true, false, true); // is_input_dq, is_q1_output + test_case(false, true, false, true); // is_q1_output + test_case(true, false, false, true); // is_input_dq + test_case(false, false, false, true); +#endif +} + // test removal when we have graph input -> Q/DQ pair -> graph output TEST(QDQTransformerTests, QDQFinalCleanupTransformer_GraphInputToOutput) { auto test_case = [](bool is_q_dq, bool use_contrib_qdq) { From efcaa835b1d9b9678046d9e67196051274145acb Mon Sep 17 00:00:00 2001 From: Changming Sun Date: Thu, 20 Jun 2024 16:13:31 -0700 Subject: [PATCH 12/52] Update generate_nuspec_for_native_nuget.py for training (#21112) ### Description Similar to #21096 , but this one is for ORT training nuget package. --- .../c-api-noopenmp-packaging-pipelines.yml | 5 +++ .../nuget/generate_nuspec_for_native_nuget.py | 40 ++++--------------- 2 files changed, 12 insertions(+), 33 deletions(-) diff --git a/tools/ci_build/github/azure-pipelines/c-api-noopenmp-packaging-pipelines.yml b/tools/ci_build/github/azure-pipelines/c-api-noopenmp-packaging-pipelines.yml index 186c4fccb1045..990d0c253c789 100644 --- a/tools/ci_build/github/azure-pipelines/c-api-noopenmp-packaging-pipelines.yml +++ b/tools/ci_build/github/azure-pipelines/c-api-noopenmp-packaging-pipelines.yml @@ -342,6 +342,11 @@ stages: DisplayName: 'ESRP - Sign C# dlls' DoEsrp: ${{ parameters.DoEsrp }} + - task: UsePythonVersion@0 + displayName: 'Use Python' + inputs: + versionSpec: 3.8 + - task: MSBuild@1 displayName: 'Build Nuget Packages' inputs: diff --git a/tools/nuget/generate_nuspec_for_native_nuget.py b/tools/nuget/generate_nuspec_for_native_nuget.py index d200a2f666939..88d1cebc84f8d 100644 --- a/tools/nuget/generate_nuspec_for_native_nuget.py +++ b/tools/nuget/generate_nuspec_for_native_nuget.py @@ -1015,57 +1015,31 @@ def generate_files(line_list, args): # Process Training specific targets and props if args.package_name == "Microsoft.ML.OnnxRuntime.Training": - monoandroid_source_targets = os.path.join( - args.sources_path, - "csharp", - "src", - "Microsoft.ML.OnnxRuntime", - "targets", - "monoandroid11.0", - "targets.xml", - ) - monoandroid_target_targets = os.path.join( - args.sources_path, - "csharp", - "src", - "Microsoft.ML.OnnxRuntime", - "targets", - "monoandroid11.0", - args.package_name + ".targets", - ) - - net6_android_source_targets = os.path.join( + net8_android_source_targets = os.path.join( args.sources_path, "csharp", "src", "Microsoft.ML.OnnxRuntime", "targets", - "net6.0-android", + "net8.0-android", "targets.xml", ) - net6_android_target_targets = os.path.join( + net8_android_target_targets = os.path.join( args.sources_path, "csharp", "src", "Microsoft.ML.OnnxRuntime", "targets", - "net6.0-android", + "net8.0-android", args.package_name + ".targets", ) - os.system(copy_command + " " + monoandroid_source_targets + " " + monoandroid_target_targets) - os.system(copy_command + " " + net6_android_source_targets + " " + net6_android_target_targets) - - files_list.append("') - files_list.append( - "' - ) - + os.system(copy_command + " " + net8_android_source_targets + " " + net8_android_target_targets) files_list.append( - "' + "' ) files_list.append( - "' + "' ) # README From 1d7bf5694779f16fe65af6e0e7029a02e3a4b05c Mon Sep 17 00:00:00 2001 From: Ted Themistokleous <107195283+TedThemistokleous@users.noreply.github.com> Date: Thu, 20 Jun 2024 19:21:11 -0400 Subject: [PATCH 13/52] [MIGraphX EP] enable compilation and execution on Windows (#36) (#21084) --- cmake/CMakeLists.txt | 3 - cmake/onnxruntime_providers_migraphx.cmake | 56 +++--- .../providers/migraphx/gpu_data_transfer.cc | 10 - ...hip_allocator.cc => migraphx_allocator.cc} | 14 +- .../{hip_allocator.h => migraphx_allocator.h} | 14 +- .../core/providers/migraphx/migraphx_call.cc | 25 +-- .../core/providers/migraphx/migraphx_call.h | 2 - .../migraphx/migraphx_execution_provider.cc | 95 ++++------ .../migraphx/migraphx_execution_provider.h | 16 +- .../migraphx_execution_provider_info.h | 2 +- .../migraphx_execution_provider_utils.h | 2 +- .../core/providers/migraphx/migraphx_inc.h | 2 +- .../migraphx/migraphx_provider_factory.cc | 19 +- .../migraphx/migraphx_provider_factory.h | 9 + .../migraphx/migraphx_stream_handle.cc | 171 ++++++++++++++++++ .../migraphx/migraphx_stream_handle.h | 48 +++++ .../providers/shared_library/provider_api.h | 3 + .../provider_bridge_provider.cc | 12 +- .../shared_library/provider_interfaces.h | 5 + .../core/session/provider_bridge_ort.cc | 21 +++ setup.py | 63 +++++-- tools/ci_build/build.py | 11 +- 22 files changed, 430 insertions(+), 173 deletions(-) rename onnxruntime/core/providers/migraphx/{hip_allocator.cc => migraphx_allocator.cc} (83%) rename onnxruntime/core/providers/migraphx/{hip_allocator.h => migraphx_allocator.h} (78%) create mode 100644 onnxruntime/core/providers/migraphx/migraphx_stream_handle.cc create mode 100644 onnxruntime/core/providers/migraphx/migraphx_stream_handle.h diff --git a/cmake/CMakeLists.txt b/cmake/CMakeLists.txt index ce22def914851..575678029e25e 100644 --- a/cmake/CMakeLists.txt +++ b/cmake/CMakeLists.txt @@ -1472,9 +1472,6 @@ if (onnxruntime_USE_CUDA) endif() if (onnxruntime_USE_MIGRAPHX) - if (WIN32) - message(FATAL_ERROR "MIGraphX does not support build in Windows!") - endif() set(AMD_MIGRAPHX_HOME ${onnxruntime_MIGRAPHX_HOME}) endif() diff --git a/cmake/onnxruntime_providers_migraphx.cmake b/cmake/onnxruntime_providers_migraphx.cmake index 01c4f8b2c8719..3a7492ebbb0b8 100644 --- a/cmake/onnxruntime_providers_migraphx.cmake +++ b/cmake/onnxruntime_providers_migraphx.cmake @@ -19,23 +19,25 @@ endif() # Add search paths for default rocm installation - list(APPEND CMAKE_PREFIX_PATH /opt/rocm/hcc /opt/rocm/hip /opt/rocm) + list(APPEND CMAKE_PREFIX_PATH /opt/rocm/hcc /opt/rocm/hip /opt/rocm $ENV{HIP_PATH}) - find_package(hip) - find_package(migraphx PATHS ${AMD_MIGRAPHX_HOME}) + # Suppress the warning about the small capitals of the package name + cmake_policy(SET CMP0144 NEW) - find_package(miopen) - find_package(rocblas) + if(WIN32 AND NOT HIP_PLATFORM) + set(HIP_PLATFORM "amd") + endif() + + find_package(hip REQUIRED) + find_package(migraphx REQUIRED PATHS ${AMD_MIGRAPHX_HOME}) - set(migraphx_libs migraphx::c hip::host MIOpen roc::rocblas) + set(migraphx_libs migraphx::c hip::host) file(GLOB_RECURSE onnxruntime_providers_migraphx_cc_srcs CONFIGURE_DEPENDS "${ONNXRUNTIME_ROOT}/core/providers/migraphx/*.h" "${ONNXRUNTIME_ROOT}/core/providers/migraphx/*.cc" "${ONNXRUNTIME_ROOT}/core/providers/shared_library/*.h" "${ONNXRUNTIME_ROOT}/core/providers/shared_library/*.cc" - "${ONNXRUNTIME_ROOT}/core/providers/rocm/rocm_stream_handle.h" - "${ONNXRUNTIME_ROOT}/core/providers/rocm/rocm_stream_handle.cc" ) source_group(TREE ${ONNXRUNTIME_ROOT}/core FILES ${onnxruntime_providers_migraphx_cc_srcs}) onnxruntime_add_shared_library_module(onnxruntime_providers_migraphx ${onnxruntime_providers_migraphx_cc_srcs}) @@ -46,18 +48,16 @@ set_target_properties(onnxruntime_providers_migraphx PROPERTIES LINKER_LANGUAGE CXX) set_target_properties(onnxruntime_providers_migraphx PROPERTIES FOLDER "ONNXRuntime") target_compile_definitions(onnxruntime_providers_migraphx PRIVATE ONNXIFI_BUILD_LIBRARY=1) - target_compile_options(onnxruntime_providers_migraphx PRIVATE -Wno-error=sign-compare) - set_property(TARGET onnxruntime_providers_migraphx APPEND_STRING PROPERTY COMPILE_FLAGS "-Wno-deprecated-declarations") - set_property(TARGET onnxruntime_providers_migraphx APPEND_STRING PROPERTY LINK_FLAGS "-Xlinker --version-script=${ONNXRUNTIME_ROOT}/core/providers/migraphx/version_script.lds -Xlinker --gc-sections") - target_link_libraries(onnxruntime_providers_migraphx PRIVATE nsync::nsync_cpp) - - include(CheckLibraryExists) - check_library_exists(migraphx::c "migraphx_program_run_async" "/opt/rocm/migraphx/lib" HAS_STREAM_SYNC) - if(HAS_STREAM_SYNC) - target_compile_definitions(onnxruntime_providers_migraphx PRIVATE -DMIGRAPHX_STREAM_SYNC) - message(STATUS "MIGRAPHX GPU STREAM SYNC is ENABLED") + if(MSVC) + set_property(TARGET onnxruntime_providers_migraphx APPEND_STRING PROPERTY LINK_FLAGS /DEF:${ONNXRUNTIME_ROOT}/core/providers/migraphx/symbols.def) + target_link_libraries(onnxruntime_providers_migraphx PRIVATE ws2_32) else() - message(STATUS "MIGRAPHX GPU STREAM SYNC is DISABLED") + target_compile_options(onnxruntime_providers_migraphx PRIVATE -Wno-error=sign-compare) + set_property(TARGET onnxruntime_providers_migraphx APPEND_STRING PROPERTY COMPILE_FLAGS "-Wno-deprecated-declarations") + endif() + if(UNIX) + set_property(TARGET onnxruntime_providers_migraphx APPEND_STRING PROPERTY LINK_FLAGS "-Xlinker --version-script=${ONNXRUNTIME_ROOT}/core/providers/migraphx/version_script.lds -Xlinker --gc-sections") + target_link_libraries(onnxruntime_providers_migraphx PRIVATE nsync::nsync_cpp stdc++fs) endif() if (onnxruntime_ENABLE_TRAINING_OPS) @@ -68,8 +68,16 @@ endif() endif() - install(TARGETS onnxruntime_providers_migraphx - ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR} - LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR} - RUNTIME DESTINATION ${CMAKE_INSTALL_BINDIR} - ) + if(CMAKE_SYSTEM_NAME STREQUAL "Windows") + install(TARGETS onnxruntime_providers_migraphx + ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR} + LIBRARY DESTINATION ${CMAKE_INSTALL_BINDIR} + RUNTIME DESTINATION ${CMAKE_INSTALL_BINDIR} + ) + else() + install(TARGETS onnxruntime_providers_migraphx + ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR} + LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR} + RUNTIME DESTINATION ${CMAKE_INSTALL_BINDIR} + ) + endif() diff --git a/onnxruntime/core/providers/migraphx/gpu_data_transfer.cc b/onnxruntime/core/providers/migraphx/gpu_data_transfer.cc index 72193ef6268c1..94480c308b99f 100644 --- a/onnxruntime/core/providers/migraphx/gpu_data_transfer.cc +++ b/onnxruntime/core/providers/migraphx/gpu_data_transfer.cc @@ -60,17 +60,7 @@ common::Status GPUDataTransfer::CopyTensorAsync(const Tensor& src, Tensor& dst, HIP_CALL_THROW(hipMemcpy(dst_data, src_data, bytes, hipMemcpyHostToDevice)); } } else if (src_device.Type() == OrtDevice::GPU) { -#ifndef MIGRAPHX_STREAM_SYNC - if (dst_device.Type() == OrtDevice::CPU && dst_device.MemType() == OrtDevice::MemType::HIP_PINNED) { - // copying from GPU to pinned memory, this is non-blocking - HIP_CALL_THROW(hipMemcpyAsync(dst_data, src_data, bytes, hipMemcpyDeviceToHost, static_cast(stream.GetHandle()))); - } else { - // copying from GPU to CPU memory, this is blocking - HIP_CALL_THROW(hipMemcpy(dst_data, src_data, bytes, hipMemcpyDeviceToHost)); - } -#else HIP_CALL_THROW(hipMemcpyAsync(dst_data, src_data, bytes, hipMemcpyDeviceToHost, static_cast(stream.GetHandle()))); -#endif } else { // copying between cpu memory memcpy(dst_data, src_data, bytes); diff --git a/onnxruntime/core/providers/migraphx/hip_allocator.cc b/onnxruntime/core/providers/migraphx/migraphx_allocator.cc similarity index 83% rename from onnxruntime/core/providers/migraphx/hip_allocator.cc rename to onnxruntime/core/providers/migraphx/migraphx_allocator.cc index 53f10e318e65f..0693eea056416 100644 --- a/onnxruntime/core/providers/migraphx/hip_allocator.cc +++ b/onnxruntime/core/providers/migraphx/migraphx_allocator.cc @@ -3,7 +3,7 @@ #include "core/providers/shared_library/provider_api.h" #include "migraphx_call.h" -#include "hip_allocator.h" +#include "migraphx_allocator.h" #include "core/common/status.h" #include "core/framework/float16.h" #include "core/common/status.h" @@ -11,7 +11,7 @@ namespace onnxruntime { -void HIPAllocator::CheckDevice() const { +void MIGraphXAllocator::CheckDevice() const { #ifndef NDEBUG // check device to match at debug build // if it's expected to change, call hipSetDevice instead of the check @@ -23,7 +23,7 @@ void HIPAllocator::CheckDevice() const { #endif } -void* HIPAllocator::Alloc(size_t size) { +void* MIGraphXAllocator::Alloc(size_t size) { CheckDevice(); void* p = nullptr; if (size > 0) { @@ -32,12 +32,12 @@ void* HIPAllocator::Alloc(size_t size) { return p; } -void HIPAllocator::Free(void* p) { +void MIGraphXAllocator::Free(void* p) { CheckDevice(); (void)hipFree(p); // do not throw error since it's OK for hipFree to fail during shutdown } -void* HIPExternalAllocator::Alloc(size_t size) { +void* MIGraphXExternalAllocator::Alloc(size_t size) { void* p = nullptr; if (size > 0) { p = alloc_(size); @@ -49,7 +49,7 @@ void* HIPExternalAllocator::Alloc(size_t size) { return p; } -void HIPExternalAllocator::Free(void* p) { +void MIGraphXExternalAllocator::Free(void* p) { free_(p); std::lock_guard lock(lock_); auto it = reserved_.find(p); @@ -59,7 +59,7 @@ void HIPExternalAllocator::Free(void* p) { } } -void* HIPExternalAllocator::Reserve(size_t size) { +void* MIGraphXExternalAllocator::Reserve(size_t size) { void* p = Alloc(size); if (!p) return nullptr; std::lock_guard lock(lock_); diff --git a/onnxruntime/core/providers/migraphx/hip_allocator.h b/onnxruntime/core/providers/migraphx/migraphx_allocator.h similarity index 78% rename from onnxruntime/core/providers/migraphx/hip_allocator.h rename to onnxruntime/core/providers/migraphx/migraphx_allocator.h index 3244f9f04ea70..64da844e8c714 100644 --- a/onnxruntime/core/providers/migraphx/hip_allocator.h +++ b/onnxruntime/core/providers/migraphx/migraphx_allocator.h @@ -9,12 +9,12 @@ namespace onnxruntime { -class HIPAllocator : public IAllocator { +class MIGraphXAllocator : public IAllocator { public: - HIPAllocator(int device_id, const char* name) + MIGraphXAllocator(int device_id, const char* name) : IAllocator( OrtMemoryInfo(name, OrtAllocatorType::OrtDeviceAllocator, - OrtDevice(OrtDevice::GPU, OrtDevice::MemType::DEFAULT, device_id), + OrtDevice(OrtDevice::GPU, OrtDevice::MemType::DEFAULT, static_cast(device_id)), device_id, OrtMemTypeDefault)) {} virtual void* Alloc(size_t size) override; @@ -24,14 +24,14 @@ class HIPAllocator : public IAllocator { void CheckDevice() const; }; -class HIPExternalAllocator : public HIPAllocator { +class MIGraphXExternalAllocator : public MIGraphXAllocator { typedef void* (*ExternalAlloc)(size_t size); typedef void (*ExternalFree)(void* p); typedef void (*ExternalEmptyCache)(); public: - HIPExternalAllocator(OrtDevice::DeviceId device_id, const char* name, void* alloc, void* free, void* empty_cache) - : HIPAllocator(device_id, name) { + MIGraphXExternalAllocator(OrtDevice::DeviceId device_id, const char* name, void* alloc, void* free, void* empty_cache) + : MIGraphXAllocator(device_id, name) { alloc_ = reinterpret_cast(alloc); free_ = reinterpret_cast(free); empty_cache_ = reinterpret_cast(empty_cache); @@ -55,7 +55,7 @@ class HIPPinnedAllocator : public IAllocator { HIPPinnedAllocator(int device_id, const char* name) : IAllocator( OrtMemoryInfo(name, OrtAllocatorType::OrtDeviceAllocator, - OrtDevice(OrtDevice::CPU, OrtDevice::MemType::HIP_PINNED, device_id), + OrtDevice(OrtDevice::CPU, OrtDevice::MemType::HIP_PINNED, static_cast(device_id)), device_id, OrtMemTypeCPUOutput)) {} virtual void* Alloc(size_t size) override; diff --git a/onnxruntime/core/providers/migraphx/migraphx_call.cc b/onnxruntime/core/providers/migraphx/migraphx_call.cc index 5248ac2f39214..9807cd646e51c 100644 --- a/onnxruntime/core/providers/migraphx/migraphx_call.cc +++ b/onnxruntime/core/providers/migraphx/migraphx_call.cc @@ -1,10 +1,13 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. +#ifdef _WIN32 +#include +#else #include -#include -#include -#include +#endif + +#include #include "core/common/common.h" #include "core/common/status.h" #include "core/providers/shared_library/provider_api.h" @@ -34,16 +37,20 @@ std::conditional_t RocmCall( ERRTYPE retCode, const char* exprString, const char* libName, ERRTYPE successCode, const char* msg, const char* file, const int line) { if (retCode != successCode) { try { - char hostname[HOST_NAME_MAX]; - if (gethostname(hostname, HOST_NAME_MAX) != 0) - strcpy(hostname, "?"); +#ifdef _WIN32 + // According to the POSIX spec, 255 is the safe minimum value. + static constexpr int HOST_NAME_MAX = 255; +#endif + std::string hostname(HOST_NAME_MAX, 0); + if (gethostname(hostname.data(), HOST_NAME_MAX) != 0) + hostname = "?"; int currentHipDevice; (void)hipGetDevice(¤tHipDevice); (void)hipGetLastError(); // clear last HIP error static char str[1024]; snprintf(str, 1024, "%s failure %d: %s ; GPU=%d ; hostname=%s ; file=%s ; line=%d ; expr=%s; %s", libName, (int)retCode, RocmErrString(retCode), currentHipDevice, - hostname, + hostname.c_str(), file, line, exprString, msg); if constexpr (THRW) { // throw an exception with the error info @@ -68,9 +75,5 @@ std::conditional_t RocmCall( template Status RocmCall(hipError_t retCode, const char* exprString, const char* libName, hipError_t successCode, const char* msg, const char* file, const int line); template void RocmCall(hipError_t retCode, const char* exprString, const char* libName, hipError_t successCode, const char* msg, const char* file, const int line); -template Status RocmCall(rocblas_status retCode, const char* exprString, const char* libName, rocblas_status successCode, const char* msg, const char* file, const int line); -template void RocmCall(rocblas_status retCode, const char* exprString, const char* libName, rocblas_status successCode, const char* msg, const char* file, const int line); -template Status RocmCall(miopenStatus_t retCode, const char* exprString, const char* libName, miopenStatus_t successCode, const char* msg, const char* file, const int line); -template void RocmCall(miopenStatus_t retCode, const char* exprString, const char* libName, miopenStatus_t successCode, const char* msg, const char* file, const int line); } // namespace onnxruntime diff --git a/onnxruntime/core/providers/migraphx/migraphx_call.h b/onnxruntime/core/providers/migraphx/migraphx_call.h index 15d385a636b76..f6a95cebf34b5 100644 --- a/onnxruntime/core/providers/migraphx/migraphx_call.h +++ b/onnxruntime/core/providers/migraphx/migraphx_call.h @@ -4,8 +4,6 @@ #pragma once #include "migraphx_inc.h" -#pragma once - namespace onnxruntime { // ----------------------------------------------------------------------- diff --git a/onnxruntime/core/providers/migraphx/migraphx_execution_provider.cc b/onnxruntime/core/providers/migraphx/migraphx_execution_provider.cc index 6ee85c3a4c047..097b16ecde536 100644 --- a/onnxruntime/core/providers/migraphx/migraphx_execution_provider.cc +++ b/onnxruntime/core/providers/migraphx/migraphx_execution_provider.cc @@ -13,12 +13,11 @@ #include "core/common/logging/severity.h" #include "migraphx_execution_provider.h" #include "migraphx_execution_provider_utils.h" -#include "hip_allocator.h" +#include "migraphx_allocator.h" #include "gpu_data_transfer.h" #include "migraphx_inc.h" -// TODO: find a better way to share this -#include "core/providers/rocm/rocm_stream_handle.h" +#include "migraphx_stream_handle.h" #if defined(_MSC_VER) #pragma warning(disable : 4244 4245) @@ -102,10 +101,10 @@ std::shared_ptr MIGraphXExecutionProvider::GetKernelRegistry() c } MIGraphXExecutionProvider::MIGraphXExecutionProvider(const MIGraphXExecutionProviderInfo& info) - : IExecutionProvider{onnxruntime::kMIGraphXExecutionProvider, OrtDevice(OrtDevice::GPU, OrtDevice::MemType::DEFAULT, info.device_id)}, device_id_(info.device_id) { + : IExecutionProvider{onnxruntime::kMIGraphXExecutionProvider, OrtDevice(OrtDevice::GPU, OrtDevice::MemType::DEFAULT, info.device_id)}, info_(info) { InitProviderOrtApi(); // Set GPU device to be used - HIP_CALL_THROW(hipSetDevice(device_id_)); + HIP_CALL_THROW(hipSetDevice(info_.device_id)); t_ = migraphx::target(info.target_device.c_str()); // whether fp16 is enable @@ -181,16 +180,10 @@ MIGraphXExecutionProvider::MIGraphXExecutionProvider(const MIGraphXExecutionProv dump_model_ops_ = (std::stoi(dump_model_ops_env) == 0 ? false : true); } - ROCBLAS_CALL_THROW(rocblas_create_handle(&external_rocblas_handle_)); - ROCBLAS_CALL_THROW(rocblas_set_stream(external_rocblas_handle_, stream_)); - - MIOPEN_CALL_THROW(miopenCreate(&external_miopen_handle_)); - MIOPEN_CALL_THROW(miopenSetStream(external_miopen_handle_, stream_)); - metadef_id_generator_ = ModelMetadefIdGenerator::Create(); LOGS_DEFAULT(VERBOSE) << "[MIGraphX EP] MIGraphX provider options: " - << "device_id: " << device_id_ + << "device_id: " << info_.device_id << ", migraphx_fp16_enable: " << fp16_enable_ << ", migraphx_int8_enable: " << int8_enable_ << ", migraphx_int8_enable: " << int8_enable_ @@ -205,17 +198,14 @@ MIGraphXExecutionProvider::MIGraphXExecutionProvider(const MIGraphXExecutionProv } MIGraphXExecutionProvider::~MIGraphXExecutionProvider() { - ORT_IGNORE_RETURN_VALUE(ROCBLAS_CALL(rocblas_destroy_handle(external_rocblas_handle_))); - ORT_IGNORE_RETURN_VALUE(MIOPEN_CALL(miopenDestroy(external_miopen_handle_))); } std::vector MIGraphXExecutionProvider::CreatePreferredAllocators() { AllocatorCreationInfo default_memory_info( - [](OrtDevice::DeviceId device_id) { return CreateROCMAllocator(device_id, onnxruntime::CUDA); }, device_id_); + [](OrtDevice::DeviceId device_id) { return CreateMIGraphXAllocator(device_id, onnxruntime::CUDA); }, info_.device_id); AllocatorCreationInfo pinned_allocator_info( [](OrtDevice::DeviceId device_id) { - ORT_UNUSED_PARAMETER(device_id); - return CreateROCMPinnedAllocator(onnxruntime::CUDA_PINNED); + return CreateMIGraphXPinnedAllocator(device_id, onnxruntime::CUDA_PINNED); }, 0); return std::vector{CreateAllocator(default_memory_info), CreateAllocator(pinned_allocator_info)}; @@ -254,40 +244,40 @@ static bool getMIGraphXType(ONNXTensorElementDataType type, migraphx_shape_datatype_t& mgx_type) { mgx_type = migraphx_shape_float_type; switch (type) { - case ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_FLOAT16: + case ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16: mgx_type = migraphx_shape_half_type; break; - case ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_FLOAT: + case ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT: mgx_type = migraphx_shape_float_type; break; - case ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_DOUBLE: + case ONNX_TENSOR_ELEMENT_DATA_TYPE_DOUBLE: mgx_type = migraphx_shape_double_type; break; - case ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_INT8: + case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT8: mgx_type = migraphx_shape_int8_type; break; - case ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_INT16: + case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT16: mgx_type = migraphx_shape_int16_type; break; - case ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_INT32: + case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32: mgx_type = migraphx_shape_int32_type; break; - case ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_INT64: + case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64: mgx_type = migraphx_shape_int64_type; break; - case ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_UINT8: + case ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT8: mgx_type = migraphx_shape_uint8_type; break; - case ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_UINT16: + case ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT16: mgx_type = migraphx_shape_uint16_type; break; - case ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_UINT32: + case ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT32: mgx_type = migraphx_shape_uint32_type; break; - case ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_UINT64: + case ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT64: mgx_type = migraphx_shape_uint64_type; break; - case ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_BOOL: + case ONNX_TENSOR_ELEMENT_DATA_TYPE_BOOL: mgx_type = migraphx_shape_bool_type; break; default: @@ -303,7 +293,7 @@ std::vector toVector(const ONNX_NAMESPACE::int64s& nums) { std::vector result; int num = nums.size(); for (int i = 0; i < num; ++i) { - result.push_back(nums[i]); + result.push_back(static_cast(nums[i])); } return result; @@ -501,16 +491,9 @@ static bool IsUnsupportedOpMode(const onnxruntime::GraphViewer& graph_viewer, co if (arg_s != nullptr) { const auto& tensor_dims = arg_s->dim(); std::vector dims; - std::transform(tensor_dims.begin(), - tensor_dims.end(), - std::back_inserter(dims), - [&](auto&& d) -> std::size_t { - if (d.has_dim_value()) { - return d.dim_value(); - } else { - return 0; - } - }); + for (auto&& dim : tensor_dims) { + dims.emplace_back(dim.has_dim_value() ? dim.dim_value() : 0); + } if (dims == std::vector{0}) { return true; } @@ -546,8 +529,8 @@ static bool IsUnsupportedOpMode(const onnxruntime::GraphViewer& graph_viewer, co } void SubgraphPostProcessing(const onnxruntime::GraphViewer& graph_viewer, std::vector>& clusters, - const logging::Logger& logger) { - // Then check whether a subgraph should fallback to CPU + [[maybe_unused]] const logging::Logger& logger) { + // Then check whether a subgraph should fall back to CPU // 1. Check whether a subgraph contains a RNN operator std::unordered_set rnn_names = {"RNN", "GRU", "LSTM"}; std::unordered_set op_names = {"AveragePool", "Conv", "Gemm", "LRN", "MatMul", "MaxPool"}; @@ -591,17 +574,10 @@ void SubgraphPostProcessing(const onnxruntime::GraphViewer& graph_viewer, std::v if (arg_s == nullptr) return false; const auto& tensor_dims = arg_s->dim(); std::vector dims; - std::transform(tensor_dims.begin(), - tensor_dims.end(), - std::back_inserter(dims), - [&](auto&& d) -> std::size_t { - if (d.has_dim_value()) { - return d.dim_value(); - } else { - return 1; - } - }); - return (std::accumulate(dims.begin(), dims.end(), 1, std::multiplies{}) > 300); + for (auto&& dim : tensor_dims) { + dims.emplace_back(dim.has_dim_value() ? dim.dim_value() : 1); + } + return (std::accumulate(dims.begin(), dims.end(), 1ULL, std::multiplies{}) > 300); })) { return false; } @@ -623,7 +599,7 @@ void SubgraphPostProcessing(const onnxruntime::GraphViewer& graph_viewer, std::v static bool IsNodeSupported(const std::set& op_set, const onnxruntime::GraphViewer& graph_viewer, const NodeIndex node_idx, - const logging::Logger& logger) { + [[maybe_unused]] const logging::Logger& logger) { const auto& node = graph_viewer.GetNode(node_idx); const auto& optype = node->OpType(); const auto& domain = node->Domain(); @@ -1442,14 +1418,10 @@ Status MIGraphXExecutionProvider::Compile(const std::vector& // lock to avoid race condition std::lock_guard lock(*(mgx_state->mgx_mu_ptr)); -#ifdef MIGRAPHX_STREAM_SYNC void* rocm_stream; Ort::ThrowOnError(api->KernelContext_GetGPUComputeStream(context, &rocm_stream)); auto prog_outputs = prog.run_async(m, static_cast(rocm_stream)); -#else - auto prog_outputs = prog.eval(m); - HIP_CALL_THROW(hipDeviceSynchronize()); -#endif + // In case of input parameters are reused as output parameter call hipMemcpy auto output_num = prog_outputs.size(); if (prog_output_indices.size() < output_num) { @@ -1478,8 +1450,7 @@ Status MIGraphXExecutionProvider::Compile(const std::vector& void MIGraphXExecutionProvider::RegisterStreamHandlers(IStreamCommandHandleRegistry& stream_handle_registry, AllocatorMap& allocators) const { auto allocator = allocators[GetOrtDeviceByMemType(OrtMemTypeCPU)]; - RegisterRocmStreamHandles(stream_handle_registry, OrtDevice::GPU, allocator, true, stream_, - false /*TODO:external_stream_*/, external_miopen_handle_, external_rocblas_handle_); + RegisterMIGraphXStreamHandles(stream_handle_registry, OrtDevice::GPU, allocator, true, stream_, false /*TODO:external_stream_*/); } OrtDevice MIGraphXExecutionProvider::GetOrtDeviceByMemType(OrtMemType mem_type) const { @@ -1487,7 +1458,6 @@ OrtDevice MIGraphXExecutionProvider::GetOrtDeviceByMemType(OrtMemType mem_type) if (mem_type == OrtMemTypeCPUOutput) return OrtDevice(OrtDevice::CPU, OrtDevice::MemType::HIP_PINNED, 0 /*CPU device id always be 0*/); return default_device_; } -#ifdef MIGRAPHX_STREAM_SYNC Status MIGraphXExecutionProvider::Sync() const { HIP_CALL_THROW(hipStreamSynchronize(static_cast(nullptr))); @@ -1512,5 +1482,4 @@ Status MIGraphXExecutionProvider::OnRunEnd(bool /*sync_stream*/, const onnxrunti return Status::OK(); } -#endif } // namespace onnxruntime diff --git a/onnxruntime/core/providers/migraphx/migraphx_execution_provider.h b/onnxruntime/core/providers/migraphx/migraphx_execution_provider.h index 1977f71b8b1cf..f34ca320d0a5a 100644 --- a/onnxruntime/core/providers/migraphx/migraphx_execution_provider.h +++ b/onnxruntime/core/providers/migraphx/migraphx_execution_provider.h @@ -3,9 +3,6 @@ #pragma once -#include -#include - #include "core/framework/arena_extend_strategy.h" #include "core/framework/execution_provider.h" #include "core/platform/ort_mutex.h" @@ -14,8 +11,6 @@ #include #include -// TODO: find a better way to share this -// #include "core/providers/cuda/rocm_stream_handle.h" namespace onnxruntime { @@ -62,13 +57,11 @@ class MIGraphXExecutionProvider : public IExecutionProvider { explicit MIGraphXExecutionProvider(const MIGraphXExecutionProviderInfo& info); ~MIGraphXExecutionProvider(); -#ifdef MIGRAPHX_STREAM_SYNC Status Sync() const override; Status OnRunStart(const onnxruntime::RunOptions& run_options) override; Status OnRunEnd(bool sync_stream, const onnxruntime::RunOptions& run_options) override; -#endif std::vector> GetCapability(const onnxruntime::GraphViewer& graph_viewer, @@ -85,7 +78,13 @@ class MIGraphXExecutionProvider : public IExecutionProvider { OrtDevice GetOrtDeviceByMemType(OrtMemType mem_type) const override; std::vector CreatePreferredAllocators() override; + int GetDeviceId() const override { return info_.device_id; } + ProviderOptions GetProviderOptions() const override { + return MIGraphXExecutionProviderInfo::ToProviderOptions(info_); + } + private: + MIGraphXExecutionProviderInfo info_; bool fp16_enable_ = false; bool int8_enable_ = false; std::string int8_calibration_cache_name_; @@ -98,7 +97,6 @@ class MIGraphXExecutionProvider : public IExecutionProvider { bool load_compiled_model_ = false; std::string load_compiled_path_; bool dump_model_ops_ = false; - int device_id_; migraphx::target t_; OrtMutex mgx_mu_; hipStream_t stream_ = nullptr; @@ -109,8 +107,6 @@ class MIGraphXExecutionProvider : public IExecutionProvider { std::unordered_map map_no_input_shape_; AllocatorPtr allocator_; - miopenHandle_t external_miopen_handle_ = nullptr; - rocblas_handle external_rocblas_handle_ = nullptr; std::unique_ptr metadef_id_generator_; }; diff --git a/onnxruntime/core/providers/migraphx/migraphx_execution_provider_info.h b/onnxruntime/core/providers/migraphx/migraphx_execution_provider_info.h index 8411e3eef096b..68d5d9af98ea4 100644 --- a/onnxruntime/core/providers/migraphx/migraphx_execution_provider_info.h +++ b/onnxruntime/core/providers/migraphx/migraphx_execution_provider_info.h @@ -14,7 +14,7 @@ namespace onnxruntime { // Information needed to construct trt execution providers. struct MIGraphXExecutionProviderInfo { std::string target_device; - int device_id{0}; + OrtDevice::DeviceId device_id{0}; bool fp16_enable{false}; bool int8_enable{false}; std::string int8_calibration_table_name{""}; diff --git a/onnxruntime/core/providers/migraphx/migraphx_execution_provider_utils.h b/onnxruntime/core/providers/migraphx/migraphx_execution_provider_utils.h index 071070e92a209..9274b5696185c 100644 --- a/onnxruntime/core/providers/migraphx/migraphx_execution_provider_utils.h +++ b/onnxruntime/core/providers/migraphx/migraphx_execution_provider_utils.h @@ -28,7 +28,7 @@ bool IsGraphInput(const GraphViewer& graph, const std::string& name) { return (std::find(input_names.begin(), input_names.end(), name) != input_names.end()); } -bool IsGraphInitializer(const GraphViewer& graph, const std::string& name, bool check_outer_scope = true) { +bool IsGraphInitializer(const GraphViewer& graph, const std::string& name, [[maybe_unused]] bool check_outer_scope = true) { const ONNX_NAMESPACE::TensorProto* initializer = nullptr; return graph.GetInitializedTensor(name, initializer); } diff --git a/onnxruntime/core/providers/migraphx/migraphx_inc.h b/onnxruntime/core/providers/migraphx/migraphx_inc.h index 96b24051ace76..2b035b20f619f 100644 --- a/onnxruntime/core/providers/migraphx/migraphx_inc.h +++ b/onnxruntime/core/providers/migraphx/migraphx_inc.h @@ -4,5 +4,5 @@ #pragma once #include -#include +#include #include diff --git a/onnxruntime/core/providers/migraphx/migraphx_provider_factory.cc b/onnxruntime/core/providers/migraphx/migraphx_provider_factory.cc index dd24dbdc76d2f..6d199930116e8 100644 --- a/onnxruntime/core/providers/migraphx/migraphx_provider_factory.cc +++ b/onnxruntime/core/providers/migraphx/migraphx_provider_factory.cc @@ -6,7 +6,7 @@ #include "core/providers/migraphx/migraphx_provider_factory.h" #include "migraphx_execution_provider.h" #include "migraphx_provider_factory_creator.h" -#include "hip_allocator.h" +#include "migraphx_allocator.h" #include "gpu_data_transfer.h" #include "core/framework/provider_options.h" @@ -33,10 +33,23 @@ std::unique_ptr MIGraphXProviderFactory::CreateProvider() { return std::make_unique(info_); } +struct ProviderInfo_MIGraphX_Impl final : ProviderInfo_MIGraphX { + std::unique_ptr CreateMIGraphXAllocator(int16_t device_id, const char* name) override { + return std::make_unique(device_id, name); + } + + std::unique_ptr CreateMIGraphXPinnedAllocator(int16_t device_id, const char* name) override { + return std::make_unique(device_id, name); + } + +} g_info; + struct MIGraphX_Provider : Provider { + void* GetInfo() override { return &g_info; } + std::shared_ptr CreateExecutionProviderFactory(int device_id) override { MIGraphXExecutionProviderInfo info; - info.device_id = device_id; + info.device_id = static_cast(device_id); info.target_device = "gpu"; return std::make_shared(info); } @@ -44,7 +57,7 @@ struct MIGraphX_Provider : Provider { std::shared_ptr CreateExecutionProviderFactory(const void* provider_options) override { auto& options = *reinterpret_cast(provider_options); MIGraphXExecutionProviderInfo info; - info.device_id = options.device_id; + info.device_id = static_cast(options.device_id); info.target_device = "gpu"; info.fp16_enable = options.migraphx_fp16_enable; info.int8_enable = options.migraphx_int8_enable; diff --git a/onnxruntime/core/providers/migraphx/migraphx_provider_factory.h b/onnxruntime/core/providers/migraphx/migraphx_provider_factory.h index ac9834e64942a..b257a4318dc0e 100644 --- a/onnxruntime/core/providers/migraphx/migraphx_provider_factory.h +++ b/onnxruntime/core/providers/migraphx/migraphx_provider_factory.h @@ -10,4 +10,13 @@ struct IExecutionProviderFactory; struct MIGraphXExecutionProviderInfo; enum class ArenaExtendStrategy : int32_t; struct MIGraphXExecutionProviderExternalAllocatorInfo; + +struct ProviderInfo_MIGraphX { + virtual std::unique_ptr CreateMIGraphXAllocator(int16_t device_id, const char* name) = 0; + virtual std::unique_ptr CreateMIGraphXPinnedAllocator(int16_t device_id, const char* name) = 0; + + protected: + ~ProviderInfo_MIGraphX() = default; // Can only be destroyed through a subclass instance +}; + } // namespace onnxruntime diff --git a/onnxruntime/core/providers/migraphx/migraphx_stream_handle.cc b/onnxruntime/core/providers/migraphx/migraphx_stream_handle.cc new file mode 100644 index 0000000000000..9c5bb4ecf5c97 --- /dev/null +++ b/onnxruntime/core/providers/migraphx/migraphx_stream_handle.cc @@ -0,0 +1,171 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include +#include "migraphx_stream_handle.h" + +namespace onnxruntime { + +struct MIGraphXNotification : public synchronize::Notification { + MIGraphXNotification(Stream& s) : Notification(s) { + HIP_CALL_THROW(hipEventCreateWithFlags(&event_, hipEventDisableTiming)); + } + + ~MIGraphXNotification() { + if (event_) + HIP_CALL_THROW(hipEventDestroy(event_)); + } + + void Activate() override { + // record event with hipEventBlockingSync so we can support sync on host without busy wait. + HIP_CALL_THROW(hipEventRecord(event_, static_cast(stream_.GetHandle()))); + } + + void wait_on_device(Stream& device_stream) { + ORT_ENFORCE(device_stream.GetDevice().Type() == OrtDevice::GPU, "Unexpected device:", device_stream.GetDevice().ToString()); + // launch a wait command to the migraphx stream + HIP_CALL_THROW(hipStreamWaitEvent(static_cast(device_stream.GetHandle()), event_, 0)); + }; + + void wait_on_host() { + // CUDA_CALL_THROW(cudaStreamSynchronize(stream_)); + HIP_CALL_THROW(hipEventSynchronize(event_)); + } + + hipEvent_t event_; +}; + +MIGraphXStream::MIGraphXStream(hipStream_t stream, + const OrtDevice& device, + AllocatorPtr cpu_allocator, + bool release_cpu_buffer_on_migraphx_stream) + : Stream(stream, device), + cpu_allocator_(cpu_allocator), + release_cpu_buffer_on_migraphx_stream_(release_cpu_buffer_on_migraphx_stream) { +} + +MIGraphXStream::~MIGraphXStream() { + ORT_IGNORE_RETURN_VALUE(CleanUpOnRunEnd()); + if (own_stream_) { + auto* handle = GetHandle(); + if (handle) + HIP_CALL_THROW(hipStreamDestroy(static_cast(handle))); + } +} + +std::unique_ptr MIGraphXStream::CreateNotification(size_t /*num_consumers*/) { + return std::make_unique(*this); +} + +void MIGraphXStream::Flush() { + if (own_stream_) + HIP_CALL_THROW(hipStreamSynchronize(static_cast(GetHandle()))); +} + +void MIGraphXStream::EnqueDeferredCPUBuffer(void* cpu_buffer) { + // stream is per thread, so don't need lock + deferred_cpu_buffers_.push_back(cpu_buffer); +} + +struct CpuBuffersInfo { + // This struct stores the information needed + // to release CPU buffers allocated for GPU kernels. + // It's used to enqueue their release after + // associated GPU kernels in a MIGraphX stream. + + // This is a CPU allocator in MIGraphX EP. + // It must be the one used to allocate the + // following pointers. + AllocatorPtr allocator; + // buffers[i] is the i-th pointer added by + // AddDeferredReleaseCPUPtr for a specific + // MIGraphX stream. For example, this fields + // should contain all values in + // deferred_release_buffer_pool_[my_stream] + // when release my_stream's buffers. + std::unique_ptr buffers; + // CPU buffer buffers[i]. + // Number of buffer points in "buffers". + size_t n_buffers; +}; + +static void ReleaseCpuBufferCallback(void* raw_info) { + std::unique_ptr info = std::make_unique(); + info.reset(reinterpret_cast(raw_info)); + for (size_t i = 0; i < info->n_buffers; ++i) { + info->allocator->Free(info->buffers[i]); + } +} + +Status MIGraphXStream::CleanUpOnRunEnd() { + if (deferred_cpu_buffers_.empty()) + return Status::OK(); + // Release the ownership of cpu_buffers_info so that the underlying + // object will keep alive until the end of ReleaseCpuBufferCallback. + if (release_cpu_buffer_on_migraphx_stream_ && cpu_allocator_->Info().alloc_type == OrtArenaAllocator) { + std::unique_ptr cpu_buffers_info = std::make_unique(); + cpu_buffers_info->allocator = cpu_allocator_; + cpu_buffers_info->buffers = std::make_unique(deferred_cpu_buffers_.size()); + for (size_t i = 0; i < deferred_cpu_buffers_.size(); ++i) { + cpu_buffers_info->buffers[i] = deferred_cpu_buffers_.at(i); + } + cpu_buffers_info->n_buffers = deferred_cpu_buffers_.size(); + HIP_RETURN_IF_ERROR(hipLaunchHostFunc(static_cast(GetHandle()), ReleaseCpuBufferCallback, cpu_buffers_info.release())); + } else { + HIP_RETURN_IF_ERROR(hipStreamSynchronize(static_cast(GetHandle()))); + for (auto* buffer : deferred_cpu_buffers_) { + cpu_allocator_->Free(buffer); + } + } + + deferred_cpu_buffers_.clear(); + return Status::OK(); +} + +void* MIGraphXStream::GetResource(int version, int id) const { + ORT_ENFORCE(version <= ORT_ROCM_RESOUCE_VERSION, "resource version unsupported!"); + void* resource{}; + switch (id) { + case RocmResource::hip_stream_t: + return reinterpret_cast(GetHandle()); + default: + break; + } + return resource; +} + +// CPU Stream command handles +void WaitMIGraphXNotificationOnDevice(Stream& stream, synchronize::Notification& notification) { + static_cast(¬ification)->wait_on_device(stream); +} + +void WaitMIGraphXNotificationOnHost(Stream& /*stream*/, synchronize::Notification& notification) { + static_cast(¬ification)->wait_on_host(); +} + +void RegisterMIGraphXStreamHandles(IStreamCommandHandleRegistry& stream_handle_registry, + const OrtDevice::DeviceType device_type, + AllocatorPtr cpu_allocator, + bool release_cpu_buffer_on_migraphx_stream, + hipStream_t external_stream, + bool use_existing_stream) { + // wait migraphx notification on migraphx ep + stream_handle_registry.RegisterWaitFn(device_type, device_type, WaitMIGraphXNotificationOnDevice); + // wait migraphx notification on cpu ep + stream_handle_registry.RegisterWaitFn(device_type, OrtDevice::CPU, WaitMIGraphXNotificationOnHost); + if (!use_existing_stream) + stream_handle_registry.RegisterCreateStreamFn(device_type, [cpu_allocator, release_cpu_buffer_on_migraphx_stream](const OrtDevice& device) { + HIP_CALL_THROW(hipSetDevice(device.Id())); + hipStream_t stream = nullptr; + HIP_CALL_THROW(hipStreamCreateWithFlags(&stream, hipStreamNonBlocking)); + return std::make_unique(stream, device, cpu_allocator, release_cpu_buffer_on_migraphx_stream); + }); + else + stream_handle_registry.RegisterCreateStreamFn(device_type, [cpu_allocator, + release_cpu_buffer_on_migraphx_stream, + external_stream](const OrtDevice& device) { + return std::make_unique(external_stream, device, cpu_allocator, release_cpu_buffer_on_migraphx_stream); + }); +} + +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/migraphx/migraphx_stream_handle.h b/onnxruntime/core/providers/migraphx/migraphx_stream_handle.h new file mode 100644 index 0000000000000..03a7c1607e3ad --- /dev/null +++ b/onnxruntime/core/providers/migraphx/migraphx_stream_handle.h @@ -0,0 +1,48 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once +#include "core/framework/stream_handles.h" +#include "migraphx_inc.h" +#include "migraphx_call.h" + +#define HIP_RETURN_IF_ERROR(expr) ORT_RETURN_IF_ERROR(HIP_CALL(expr)) + +namespace onnxruntime { +void WaitMIGraphXNotificationOnDevice(Stream& stream, synchronize::Notification& notification); + +struct MIGraphXStream : Stream { + MIGraphXStream(hipStream_t stream, + const OrtDevice& device, + AllocatorPtr cpu_allocator, + bool release_cpu_buffer_on_migraphx_stream); + + ~MIGraphXStream(); + + std::unique_ptr CreateNotification(size_t /*num_consumers*/) override; + + void Flush() override; + + Status CleanUpOnRunEnd() override; + + void EnqueDeferredCPUBuffer(void* cpu_buffer); + + bool own_stream_{true}; + + virtual void* GetResource(int version, int id) const; + + virtual WaitNotificationFn GetWaitNotificationFn() const { return WaitMIGraphXNotificationOnDevice; } + + private: + std::vector deferred_cpu_buffers_; + AllocatorPtr cpu_allocator_; + bool release_cpu_buffer_on_migraphx_stream_{true}; +}; + +void RegisterMIGraphXStreamHandles(IStreamCommandHandleRegistry& stream_handle_registry, + const OrtDevice::DeviceType device_type, + AllocatorPtr cpu_allocator, + bool release_cpu_buffer_on_migraphx_stream, + hipStream_t external_stream, + bool use_existing_stream); +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/shared_library/provider_api.h b/onnxruntime/core/providers/shared_library/provider_api.h index 7cdfb0ffc19f2..8917bb7fd9bb6 100644 --- a/onnxruntime/core/providers/shared_library/provider_api.h +++ b/onnxruntime/core/providers/shared_library/provider_api.h @@ -279,6 +279,9 @@ std::unique_ptr CreateCPUAllocator(const OrtMemoryInfo& memory_info) std::unique_ptr CreateCUDAAllocator(int16_t device_id, const char* name); std::unique_ptr CreateCUDAPinnedAllocator(const char* name); +std::unique_ptr CreateMIGraphXAllocator(int16_t device_id, const char* name); +std::unique_ptr CreateMIGraphXPinnedAllocator(int16_t device_id, const char* name); + std::unique_ptr CreateROCMAllocator(int16_t device_id, const char* name); std::unique_ptr CreateROCMPinnedAllocator(const char* name); diff --git a/onnxruntime/core/providers/shared_library/provider_bridge_provider.cc b/onnxruntime/core/providers/shared_library/provider_bridge_provider.cc index 27d8a0f06f565..540f671d67f8d 100644 --- a/onnxruntime/core/providers/shared_library/provider_bridge_provider.cc +++ b/onnxruntime/core/providers/shared_library/provider_bridge_provider.cc @@ -353,16 +353,12 @@ std::unique_ptr CreateGPUDataTransfer() { #endif #ifdef USE_MIGRAPHX -std::unique_ptr CreateROCMAllocator(int16_t device_id, const char* name) { - return g_host->CreateROCMAllocator(device_id, name); +std::unique_ptr CreateMIGraphXAllocator(int16_t device_id, const char* name) { + return g_host->CreateMIGraphXAllocator(device_id, name); } -std::unique_ptr CreateROCMPinnedAllocator(const char* name) { - return g_host->CreateROCMPinnedAllocator(name); -} - -std::unique_ptr CreateGPUDataTransfer() { - return g_host->CreateGPUDataTransfer(); +std::unique_ptr CreateMIGraphXPinnedAllocator(int16_t device_id, const char* name) { + return g_host->CreateMIGraphXPinnedAllocator(device_id, name); } #endif diff --git a/onnxruntime/core/providers/shared_library/provider_interfaces.h b/onnxruntime/core/providers/shared_library/provider_interfaces.h index cc3b13f696a96..f1a778e8b8f80 100644 --- a/onnxruntime/core/providers/shared_library/provider_interfaces.h +++ b/onnxruntime/core/providers/shared_library/provider_interfaces.h @@ -178,6 +178,11 @@ struct ProviderHost { virtual void CudaCall_true(int retCode, const char* exprString, const char* libName, int successCode, const char* msg, const char* file, const int line) = 0; #endif +#ifdef USE_MIGRAPHX + virtual std::unique_ptr CreateMIGraphXAllocator(int16_t device_id, const char* name) = 0; + virtual std::unique_ptr CreateMIGraphXPinnedAllocator(int16_t device_id, const char* name) = 0; +#endif + #ifdef USE_ROCM virtual std::unique_ptr CreateROCMAllocator(int16_t device_id, const char* name) = 0; virtual std::unique_ptr CreateROCMPinnedAllocator(const char* name) = 0; diff --git a/onnxruntime/core/session/provider_bridge_ort.cc b/onnxruntime/core/session/provider_bridge_ort.cc index 7f7ed5e436afe..a4959399990c5 100644 --- a/onnxruntime/core/session/provider_bridge_ort.cc +++ b/onnxruntime/core/session/provider_bridge_ort.cc @@ -130,6 +130,8 @@ ProviderInfo_Dnnl& GetProviderInfo_Dnnl(); ProviderInfo_ROCM* TryGetProviderInfo_ROCM(); ProviderInfo_ROCM& GetProviderInfo_ROCM(); ProviderHostCPU& GetProviderHostCPU(); +ProviderInfo_MIGraphX* TryGetProviderInfo_MIGraphX(); +ProviderInfo_MIGraphX& GetProviderInfo_MIGraphX(); ONNX_NAMESPACE::OpSchema CreateSchema(const std::string& domain, const std::vector& ops); struct TensorShapeProto_Dimension_Iterator_Impl : TensorShapeProto_Dimension_Iterator { TensorShapeProto_Dimension_Iterator_Impl(google::protobuf::internal::RepeatedPtrIterator&& v) : v_{std::move(v)} {} @@ -241,6 +243,11 @@ struct ProviderHostImpl : ProviderHost { void CudaCall_true(int retCode, const char* exprString, const char* libName, int successCode, const char* msg, const char* file, const int line) override { GetProviderInfo_CUDA().CudaCall_true(retCode, exprString, libName, successCode, msg, file, line); } #endif +#ifdef USE_MIGRAPHX + std::unique_ptr CreateMIGraphXAllocator(int16_t device_id, const char* name) override { return GetProviderInfo_MIGraphX().CreateMIGraphXAllocator(device_id, name); } + std::unique_ptr CreateMIGraphXPinnedAllocator(int16_t device_id, const char* name) override { return GetProviderInfo_MIGraphX().CreateMIGraphXPinnedAllocator(device_id, name); } +#endif + #ifdef USE_ROCM std::unique_ptr CreateROCMAllocator(int16_t device_id, const char* name) override { return GetProviderInfo_ROCM().CreateROCMAllocator(device_id, name); } std::unique_ptr CreateROCMPinnedAllocator(const char* name) override { return GetProviderInfo_ROCM().CreateROCMPinnedAllocator(name); } @@ -1900,6 +1907,20 @@ ProviderInfo_ROCM& GetProviderInfo_ROCM() { ORT_THROW("ROCM Provider not available, can't get interface for it"); } +ProviderInfo_MIGraphX* TryGetProviderInfo_MIGraphX() try { + return reinterpret_cast(s_library_migraphx.Get().GetInfo()); +} catch (const std::exception& exception) { + LOGS_DEFAULT(ERROR) << exception.what(); + return nullptr; +} + +ProviderInfo_MIGraphX& GetProviderInfo_MIGraphX() { + if (auto* info = TryGetProviderInfo_MIGraphX()) + return *info; + + ORT_THROW("MIGraphX Provider not available, can't get interface for it"); +} + void CopyGpuToCpu( void* dst_ptr, const void* src_ptr, diff --git a/setup.py b/setup.py index 3203993e0c4d4..baab399872b0f 100644 --- a/setup.py +++ b/setup.py @@ -56,6 +56,7 @@ def parse_arg_remove_string(argv, arg_name_equal): cuda_version = None rocm_version = None +is_migraphx = False is_rocm = False is_openvino = False # The following arguments are mutually exclusive @@ -64,8 +65,9 @@ def parse_arg_remove_string(argv, arg_name_equal): cuda_version = parse_arg_remove_string(sys.argv, "--cuda_version=") elif parse_arg_remove_boolean(sys.argv, "--use_rocm"): is_rocm = True - package_name = "onnxruntime-rocm" if not nightly_build else "ort-rocm-nightly" rocm_version = parse_arg_remove_string(sys.argv, "--rocm_version=") +elif parse_arg_remove_boolean(sys.argv, "--use_migraphx"): + is_migraphx = True elif parse_arg_remove_boolean(sys.argv, "--use_openvino"): is_openvino = True package_name = "onnxruntime-openvino" @@ -87,6 +89,9 @@ def parse_arg_remove_string(argv, arg_name_equal): elif parse_arg_remove_boolean(sys.argv, "--use_qnn"): package_name = "onnxruntime-qnn" +if is_rocm or is_migraphx: + package_name = "onnxruntime-rocm" if not nightly_build else "ort-rocm-nightly" + # PEP 513 defined manylinux1_x86_64 and manylinux1_i686 # PEP 571 defined manylinux2010_x86_64 and manylinux2010_i686 # PEP 599 defines the following platform tags: @@ -280,10 +285,21 @@ def finalize_options(self): return ret -providers_cuda_or_rocm = "libonnxruntime_providers_" + ("rocm.so" if is_rocm else "cuda.so") -providers_tensorrt_or_migraphx = "libonnxruntime_providers_" + ("migraphx.so" if is_rocm else "tensorrt.so") -providers_openvino = "libonnxruntime_providers_openvino.so" -providers_cann = "libonnxruntime_providers_cann.so" +providers_cuda_or_rocm = "onnxruntime_providers_" + ("rocm" if is_rocm else "cuda") +providers_tensorrt_or_migraphx = "onnxruntime_providers_" + ("migraphx" if is_migraphx else "tensorrt") +providers_openvino = "onnxruntime_providers_openvino" +providers_cann = "onnxruntime_providers_cann" + +if platform.system() == "Linux": + providers_cuda_or_rocm = "lib" + providers_cuda_or_rocm + ".so" + providers_tensorrt_or_migraphx = "lib" + providers_tensorrt_or_migraphx + ".so" + providers_openvino = "lib" + providers_openvino + ".so" + providers_cann = "lib" + providers_cann + ".so" +elif platform.system() == "Windows": + providers_cuda_or_rocm = providers_cuda_or_rocm + ".dll" + providers_tensorrt_or_migraphx = providers_tensorrt_or_migraphx + ".dll" + providers_openvino = providers_openvino + ".dll" + providers_cann = providers_cann + ".dll" # Additional binaries dl_libs = [] @@ -297,19 +313,22 @@ def finalize_options(self): "libmklml_gnu.so", "libiomp5.so", "mimalloc.so", + # DNNL, TensorRT & OpenVINO EPs are built as shared libs + "libonnxruntime_providers_shared.so", + "libonnxruntime_providers_dnnl.so", + "libonnxruntime_providers_openvino.so", + "libonnxruntime_providers_vitisai.so", + providers_cuda_or_rocm, + providers_tensorrt_or_migraphx, + providers_cann, ] - dl_libs = ["libonnxruntime_providers_shared.so"] - dl_libs.append(providers_cuda_or_rocm) - dl_libs.append(providers_tensorrt_or_migraphx) - dl_libs.append(providers_cann) - # DNNL, TensorRT & OpenVINO EPs are built as shared libs - libs.extend(["libonnxruntime_providers_shared.so"]) - libs.extend(["libonnxruntime_providers_dnnl.so"]) - libs.extend(["libonnxruntime_providers_openvino.so"]) - libs.extend(["libonnxruntime_providers_vitisai.so"]) - libs.append(providers_cuda_or_rocm) - libs.append(providers_tensorrt_or_migraphx) - libs.append(providers_cann) + dl_libs = [ + "libonnxruntime_providers_shared.so", + providers_cuda_or_rocm, + providers_tensorrt_or_migraphx, + providers_cann, + ] + if nightly_build: libs.extend(["libonnxruntime_pywrapper.so"]) elif platform.system() == "Darwin": @@ -323,7 +342,15 @@ def finalize_options(self): if nightly_build: libs.extend(["libonnxruntime_pywrapper.dylib"]) else: - libs = ["onnxruntime_pybind11_state.pyd", "dnnl.dll", "mklml.dll", "libiomp5md.dll"] + libs = [ + "onnxruntime_pybind11_state.pyd", + "dnnl.dll", + "mklml.dll", + "libiomp5md.dll", + providers_cuda_or_rocm, + providers_tensorrt_or_migraphx, + providers_cann, + ] # DNNL, TensorRT & OpenVINO EPs are built as shared libs libs.extend(["onnxruntime_providers_shared.dll"]) libs.extend(["onnxruntime_providers_dnnl.dll"]) diff --git a/tools/ci_build/build.py b/tools/ci_build/build.py index 3e587e9b56e2e..6159f078828f3 100644 --- a/tools/ci_build/build.py +++ b/tools/ci_build/build.py @@ -611,6 +611,7 @@ def convert_arg_line_to_args(self, arg_line): "MinGW Makefiles", "Ninja", "NMake Makefiles", + "NMake Makefiles JOM", "Unix Makefiles", "Visual Studio 17 2022", "Xcode", @@ -2207,6 +2208,7 @@ def build_python_wheel( use_cuda, cuda_version, use_rocm, + use_migraphx, rocm_version, use_dnnl, use_tensorrt, @@ -2258,6 +2260,8 @@ def build_python_wheel( args.append("--use_rocm") if rocm_version: args.append(f"--rocm_version={rocm_version}") + elif use_migraphx: + args.append("--use_migraphx") elif use_openvino: args.append("--use_openvino") elif use_dnnl: @@ -2583,9 +2587,6 @@ def main(): if args.use_tensorrt: args.use_cuda = True - if args.use_migraphx: - args.use_rocm = True - if args.build_wheel or args.gen_doc or args.use_tvm or args.enable_training: args.enable_pybind = True @@ -2872,7 +2873,8 @@ def main(): # fail unexpectedly. Similar, if your packaging step forgot to copy a file into the package, we don't know it # either. if args.build: - # TODO: find asan DLL and copy it to onnxruntime/capi folder when args.enable_address_sanitizer is True and the target OS is Windows + # TODO: find asan DLL and copy it to onnxruntime/capi folder when args.enable_address_sanitizer is True and + # the target OS is Windows if args.build_wheel: nightly_build = bool(os.getenv("NIGHTLY_BUILD") == "1") default_training_package_device = bool(os.getenv("DEFAULT_TRAINING_PACKAGE_DEVICE") == "1") @@ -2883,6 +2885,7 @@ def main(): args.use_cuda, args.cuda_version, args.use_rocm, + args.use_migraphx, args.rocm_version, args.use_dnnl, args.use_tensorrt, From b9eb1dc21efb8ee8739a34f403d7f21dc9e88cee Mon Sep 17 00:00:00 2001 From: Jake Mathern Date: Thu, 20 Jun 2024 16:28:15 -0700 Subject: [PATCH 14/52] Update protobuf_cmake.patch to allow extra disablements configurable by projects that build ORT (#20875) ### Description Update protobuf_cmake.patch to allow extra disablements. ORT repo already patches protobuf to not disable the warning 4996. ### Motivation and Context To meet SDL requirements, Microsoft repos have to fail build if there is warning 4996 Binskim also gives errors if warning 4996 is disabled. We can suppress the Binskim issues, but we need a way to disable the warnings for the minimal set of code that has them. Right now, WindowsAI disables 4996 for entirety of ORT, but it should only be disabled for protobuf. --- cmake/patches/protobuf/protobuf_cmake.patch | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/cmake/patches/protobuf/protobuf_cmake.patch b/cmake/patches/protobuf/protobuf_cmake.patch index 00e1f6f81d6f0..e315420e4defe 100644 --- a/cmake/patches/protobuf/protobuf_cmake.patch +++ b/cmake/patches/protobuf/protobuf_cmake.patch @@ -12,11 +12,12 @@ index 04cb3303a..c023001de 100644 /wd4305 # 'identifier' : truncation from 'type1' to 'type2' /wd4307 # 'operator' : integral constant overflow /wd4309 # 'conversion' : truncation of constant value -@@ -259,7 +257,6 @@ if (MSVC) +@@ -259,7 +257,7 @@ if (MSVC) /wd4355 # 'this' : used in base member initializer list /wd4506 # no definition for inline function 'function' /wd4800 # 'type' : forcing value to bool 'true' or 'false' (performance warning) - /wd4996 # The compiler encountered a deprecated declaration. ++ ${onnxruntime_PROTOBUF_EXTRA_WARNING_DISABLEMENT} ) # Allow big object add_definitions(/bigobj) From 69d522f4e902ae412c5656c6348b4d81e7e3fe42 Mon Sep 17 00:00:00 2001 From: Yi Zhang Date: Fri, 21 Jun 2024 12:41:06 +0800 Subject: [PATCH 15/52] [Fix] use cmdline in Final Jar Testing Stage for new managed Windows Image (#21130) ### Description No bash command in Managed Windows image. Use CmdlLine step instead. ### Verified Link https://dev.azure.com/aiinfra/Lotus/_build/results?buildId=491902&view=logs&j=f1f8e11e-a9fa-53e5-cd29-3ba2c1988550 --- .../templates/final-jar-testing.yml | 47 ++++++++++++------- 1 file changed, 29 insertions(+), 18 deletions(-) diff --git a/tools/ci_build/github/azure-pipelines/templates/final-jar-testing.yml b/tools/ci_build/github/azure-pipelines/templates/final-jar-testing.yml index 31519a2cef376..c9b7c01146981 100644 --- a/tools/ci_build/github/azure-pipelines/templates/final-jar-testing.yml +++ b/tools/ci_build/github/azure-pipelines/templates/final-jar-testing.yml @@ -53,24 +53,35 @@ stages: SpecificArtifact: ${{ parameters.SpecificArtifact }} BuildId: ${{ parameters.BuildId }} - - task: Bash@3 - inputs: - targetType: 'inline' - script: | - echo "Java Version" - java --version - mkdir test - pushd test - jar xf '$(Build.BinariesDirectory)/final-jar/testing.jar' - popd - # if you want to run the tests in the power shell, you need to replace ':' to ';', that is, "-cp .;.\test;protobuf-java-3.21.7.jar;onnxruntime-$(OnnxRuntimeVersion).jar" - java -jar ./junit-platform-console-standalone-1.6.2.jar -cp .:./test:./protobuf-java-3.21.7.jar:./onnxruntime-$(OnnxRuntimeVersion).jar --scan-class-path --fail-if-no-tests --disable-banner - workingDirectory: '$(Build.BinariesDirectory)/final-jar' - env: - ${{ if eq(parameters.OS, 'MacOS') }}: - DYLD_LIBRARY_PATH: '$(Build.BinariesDirectory)/final-jar/test:$(DYLD_LIBRARY_PATH)' - ${{ if eq(parameters.OS, 'Linux') }}: - LD_LIBRARY_PATH: '$(Build.BinariesDirectory)/final-jar/test:$(LD_LIBRARY_PATH)' + - ${{ if eq(parameters.OS, 'Windows') }}: + - task: CmdLine@2 + inputs: + script: | + mkdir test + pushd test + jar xf $(Build.BinariesDirectory)\final-jar\testing.jar + popd + java -jar junit-platform-console-standalone-1.6.2.jar -cp .;.\test;protobuf-java-3.21.7.jar;onnxruntime-$(OnnxRuntimeVersion).jar --scan-class-path --fail-if-no-tests --disable-banner + workingDirectory: '$(Build.BinariesDirectory)\final-jar' + - ${{ else }}: + - task: Bash@3 + inputs: + targetType: 'inline' + script: | + echo "Java Version" + java --version + mkdir test + pushd test + jar xf '$(Build.BinariesDirectory)/final-jar/testing.jar' + popd + # if you want to run the tests in the power shell, you need to replace ':' to ';', that is, "-cp .;.\test;protobuf-java-3.21.7.jar;onnxruntime-$(OnnxRuntimeVersion).jar" + java -jar ./junit-platform-console-standalone-1.6.2.jar -cp .:./test:./protobuf-java-3.21.7.jar:./onnxruntime-$(OnnxRuntimeVersion).jar --scan-class-path --fail-if-no-tests --disable-banner + workingDirectory: '$(Build.BinariesDirectory)/final-jar' + env: + ${{ if eq(parameters.OS, 'MacOS') }}: + DYLD_LIBRARY_PATH: '$(Build.BinariesDirectory)/final-jar/test:$(DYLD_LIBRARY_PATH)' + ${{ if eq(parameters.OS, 'Linux') }}: + LD_LIBRARY_PATH: '$(Build.BinariesDirectory)/final-jar/test:$(LD_LIBRARY_PATH)' - ${{ if eq(parameters['OS'], 'MacOS') }}: - template: use-xcode-version.yml From f5625b8858b5ab5bad27efff9d671ca87b4263c7 Mon Sep 17 00:00:00 2001 From: Changming Sun Date: Fri, 21 Jun 2024 01:01:07 -0700 Subject: [PATCH 16/52] Revert "[MIGraphX EP] enable compilation and execution on Windows (21084)" (#21132) ### Description This reverts commit 1d7bf5694779f16fe65af6e0e7029a02e3a4b05c because it broken the AMD GPU CI pipeline. Sorry when I reviewed the PR I forgot to run the AMD GPU CI pipeline. Will revert the PR first then ask the author to fix the issue. --- cmake/CMakeLists.txt | 3 + cmake/onnxruntime_providers_migraphx.cmake | 56 +++--- .../providers/migraphx/gpu_data_transfer.cc | 10 + ...migraphx_allocator.cc => hip_allocator.cc} | 14 +- .../{migraphx_allocator.h => hip_allocator.h} | 14 +- .../core/providers/migraphx/migraphx_call.cc | 25 ++- .../core/providers/migraphx/migraphx_call.h | 2 + .../migraphx/migraphx_execution_provider.cc | 95 ++++++---- .../migraphx/migraphx_execution_provider.h | 16 +- .../migraphx_execution_provider_info.h | 2 +- .../migraphx_execution_provider_utils.h | 2 +- .../core/providers/migraphx/migraphx_inc.h | 2 +- .../migraphx/migraphx_provider_factory.cc | 19 +- .../migraphx/migraphx_provider_factory.h | 9 - .../migraphx/migraphx_stream_handle.cc | 171 ------------------ .../migraphx/migraphx_stream_handle.h | 48 ----- .../providers/shared_library/provider_api.h | 3 - .../provider_bridge_provider.cc | 12 +- .../shared_library/provider_interfaces.h | 5 - .../core/session/provider_bridge_ort.cc | 21 --- setup.py | 63 ++----- tools/ci_build/build.py | 11 +- 22 files changed, 173 insertions(+), 430 deletions(-) rename onnxruntime/core/providers/migraphx/{migraphx_allocator.cc => hip_allocator.cc} (83%) rename onnxruntime/core/providers/migraphx/{migraphx_allocator.h => hip_allocator.h} (78%) delete mode 100644 onnxruntime/core/providers/migraphx/migraphx_stream_handle.cc delete mode 100644 onnxruntime/core/providers/migraphx/migraphx_stream_handle.h diff --git a/cmake/CMakeLists.txt b/cmake/CMakeLists.txt index 575678029e25e..ce22def914851 100644 --- a/cmake/CMakeLists.txt +++ b/cmake/CMakeLists.txt @@ -1472,6 +1472,9 @@ if (onnxruntime_USE_CUDA) endif() if (onnxruntime_USE_MIGRAPHX) + if (WIN32) + message(FATAL_ERROR "MIGraphX does not support build in Windows!") + endif() set(AMD_MIGRAPHX_HOME ${onnxruntime_MIGRAPHX_HOME}) endif() diff --git a/cmake/onnxruntime_providers_migraphx.cmake b/cmake/onnxruntime_providers_migraphx.cmake index 3a7492ebbb0b8..01c4f8b2c8719 100644 --- a/cmake/onnxruntime_providers_migraphx.cmake +++ b/cmake/onnxruntime_providers_migraphx.cmake @@ -19,25 +19,23 @@ endif() # Add search paths for default rocm installation - list(APPEND CMAKE_PREFIX_PATH /opt/rocm/hcc /opt/rocm/hip /opt/rocm $ENV{HIP_PATH}) + list(APPEND CMAKE_PREFIX_PATH /opt/rocm/hcc /opt/rocm/hip /opt/rocm) - # Suppress the warning about the small capitals of the package name - cmake_policy(SET CMP0144 NEW) + find_package(hip) + find_package(migraphx PATHS ${AMD_MIGRAPHX_HOME}) - if(WIN32 AND NOT HIP_PLATFORM) - set(HIP_PLATFORM "amd") - endif() - - find_package(hip REQUIRED) - find_package(migraphx REQUIRED PATHS ${AMD_MIGRAPHX_HOME}) + find_package(miopen) + find_package(rocblas) - set(migraphx_libs migraphx::c hip::host) + set(migraphx_libs migraphx::c hip::host MIOpen roc::rocblas) file(GLOB_RECURSE onnxruntime_providers_migraphx_cc_srcs CONFIGURE_DEPENDS "${ONNXRUNTIME_ROOT}/core/providers/migraphx/*.h" "${ONNXRUNTIME_ROOT}/core/providers/migraphx/*.cc" "${ONNXRUNTIME_ROOT}/core/providers/shared_library/*.h" "${ONNXRUNTIME_ROOT}/core/providers/shared_library/*.cc" + "${ONNXRUNTIME_ROOT}/core/providers/rocm/rocm_stream_handle.h" + "${ONNXRUNTIME_ROOT}/core/providers/rocm/rocm_stream_handle.cc" ) source_group(TREE ${ONNXRUNTIME_ROOT}/core FILES ${onnxruntime_providers_migraphx_cc_srcs}) onnxruntime_add_shared_library_module(onnxruntime_providers_migraphx ${onnxruntime_providers_migraphx_cc_srcs}) @@ -48,16 +46,18 @@ set_target_properties(onnxruntime_providers_migraphx PROPERTIES LINKER_LANGUAGE CXX) set_target_properties(onnxruntime_providers_migraphx PROPERTIES FOLDER "ONNXRuntime") target_compile_definitions(onnxruntime_providers_migraphx PRIVATE ONNXIFI_BUILD_LIBRARY=1) - if(MSVC) - set_property(TARGET onnxruntime_providers_migraphx APPEND_STRING PROPERTY LINK_FLAGS /DEF:${ONNXRUNTIME_ROOT}/core/providers/migraphx/symbols.def) - target_link_libraries(onnxruntime_providers_migraphx PRIVATE ws2_32) + target_compile_options(onnxruntime_providers_migraphx PRIVATE -Wno-error=sign-compare) + set_property(TARGET onnxruntime_providers_migraphx APPEND_STRING PROPERTY COMPILE_FLAGS "-Wno-deprecated-declarations") + set_property(TARGET onnxruntime_providers_migraphx APPEND_STRING PROPERTY LINK_FLAGS "-Xlinker --version-script=${ONNXRUNTIME_ROOT}/core/providers/migraphx/version_script.lds -Xlinker --gc-sections") + target_link_libraries(onnxruntime_providers_migraphx PRIVATE nsync::nsync_cpp) + + include(CheckLibraryExists) + check_library_exists(migraphx::c "migraphx_program_run_async" "/opt/rocm/migraphx/lib" HAS_STREAM_SYNC) + if(HAS_STREAM_SYNC) + target_compile_definitions(onnxruntime_providers_migraphx PRIVATE -DMIGRAPHX_STREAM_SYNC) + message(STATUS "MIGRAPHX GPU STREAM SYNC is ENABLED") else() - target_compile_options(onnxruntime_providers_migraphx PRIVATE -Wno-error=sign-compare) - set_property(TARGET onnxruntime_providers_migraphx APPEND_STRING PROPERTY COMPILE_FLAGS "-Wno-deprecated-declarations") - endif() - if(UNIX) - set_property(TARGET onnxruntime_providers_migraphx APPEND_STRING PROPERTY LINK_FLAGS "-Xlinker --version-script=${ONNXRUNTIME_ROOT}/core/providers/migraphx/version_script.lds -Xlinker --gc-sections") - target_link_libraries(onnxruntime_providers_migraphx PRIVATE nsync::nsync_cpp stdc++fs) + message(STATUS "MIGRAPHX GPU STREAM SYNC is DISABLED") endif() if (onnxruntime_ENABLE_TRAINING_OPS) @@ -68,16 +68,8 @@ endif() endif() - if(CMAKE_SYSTEM_NAME STREQUAL "Windows") - install(TARGETS onnxruntime_providers_migraphx - ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR} - LIBRARY DESTINATION ${CMAKE_INSTALL_BINDIR} - RUNTIME DESTINATION ${CMAKE_INSTALL_BINDIR} - ) - else() - install(TARGETS onnxruntime_providers_migraphx - ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR} - LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR} - RUNTIME DESTINATION ${CMAKE_INSTALL_BINDIR} - ) - endif() + install(TARGETS onnxruntime_providers_migraphx + ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR} + LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR} + RUNTIME DESTINATION ${CMAKE_INSTALL_BINDIR} + ) diff --git a/onnxruntime/core/providers/migraphx/gpu_data_transfer.cc b/onnxruntime/core/providers/migraphx/gpu_data_transfer.cc index 94480c308b99f..72193ef6268c1 100644 --- a/onnxruntime/core/providers/migraphx/gpu_data_transfer.cc +++ b/onnxruntime/core/providers/migraphx/gpu_data_transfer.cc @@ -60,7 +60,17 @@ common::Status GPUDataTransfer::CopyTensorAsync(const Tensor& src, Tensor& dst, HIP_CALL_THROW(hipMemcpy(dst_data, src_data, bytes, hipMemcpyHostToDevice)); } } else if (src_device.Type() == OrtDevice::GPU) { +#ifndef MIGRAPHX_STREAM_SYNC + if (dst_device.Type() == OrtDevice::CPU && dst_device.MemType() == OrtDevice::MemType::HIP_PINNED) { + // copying from GPU to pinned memory, this is non-blocking + HIP_CALL_THROW(hipMemcpyAsync(dst_data, src_data, bytes, hipMemcpyDeviceToHost, static_cast(stream.GetHandle()))); + } else { + // copying from GPU to CPU memory, this is blocking + HIP_CALL_THROW(hipMemcpy(dst_data, src_data, bytes, hipMemcpyDeviceToHost)); + } +#else HIP_CALL_THROW(hipMemcpyAsync(dst_data, src_data, bytes, hipMemcpyDeviceToHost, static_cast(stream.GetHandle()))); +#endif } else { // copying between cpu memory memcpy(dst_data, src_data, bytes); diff --git a/onnxruntime/core/providers/migraphx/migraphx_allocator.cc b/onnxruntime/core/providers/migraphx/hip_allocator.cc similarity index 83% rename from onnxruntime/core/providers/migraphx/migraphx_allocator.cc rename to onnxruntime/core/providers/migraphx/hip_allocator.cc index 0693eea056416..53f10e318e65f 100644 --- a/onnxruntime/core/providers/migraphx/migraphx_allocator.cc +++ b/onnxruntime/core/providers/migraphx/hip_allocator.cc @@ -3,7 +3,7 @@ #include "core/providers/shared_library/provider_api.h" #include "migraphx_call.h" -#include "migraphx_allocator.h" +#include "hip_allocator.h" #include "core/common/status.h" #include "core/framework/float16.h" #include "core/common/status.h" @@ -11,7 +11,7 @@ namespace onnxruntime { -void MIGraphXAllocator::CheckDevice() const { +void HIPAllocator::CheckDevice() const { #ifndef NDEBUG // check device to match at debug build // if it's expected to change, call hipSetDevice instead of the check @@ -23,7 +23,7 @@ void MIGraphXAllocator::CheckDevice() const { #endif } -void* MIGraphXAllocator::Alloc(size_t size) { +void* HIPAllocator::Alloc(size_t size) { CheckDevice(); void* p = nullptr; if (size > 0) { @@ -32,12 +32,12 @@ void* MIGraphXAllocator::Alloc(size_t size) { return p; } -void MIGraphXAllocator::Free(void* p) { +void HIPAllocator::Free(void* p) { CheckDevice(); (void)hipFree(p); // do not throw error since it's OK for hipFree to fail during shutdown } -void* MIGraphXExternalAllocator::Alloc(size_t size) { +void* HIPExternalAllocator::Alloc(size_t size) { void* p = nullptr; if (size > 0) { p = alloc_(size); @@ -49,7 +49,7 @@ void* MIGraphXExternalAllocator::Alloc(size_t size) { return p; } -void MIGraphXExternalAllocator::Free(void* p) { +void HIPExternalAllocator::Free(void* p) { free_(p); std::lock_guard lock(lock_); auto it = reserved_.find(p); @@ -59,7 +59,7 @@ void MIGraphXExternalAllocator::Free(void* p) { } } -void* MIGraphXExternalAllocator::Reserve(size_t size) { +void* HIPExternalAllocator::Reserve(size_t size) { void* p = Alloc(size); if (!p) return nullptr; std::lock_guard lock(lock_); diff --git a/onnxruntime/core/providers/migraphx/migraphx_allocator.h b/onnxruntime/core/providers/migraphx/hip_allocator.h similarity index 78% rename from onnxruntime/core/providers/migraphx/migraphx_allocator.h rename to onnxruntime/core/providers/migraphx/hip_allocator.h index 64da844e8c714..3244f9f04ea70 100644 --- a/onnxruntime/core/providers/migraphx/migraphx_allocator.h +++ b/onnxruntime/core/providers/migraphx/hip_allocator.h @@ -9,12 +9,12 @@ namespace onnxruntime { -class MIGraphXAllocator : public IAllocator { +class HIPAllocator : public IAllocator { public: - MIGraphXAllocator(int device_id, const char* name) + HIPAllocator(int device_id, const char* name) : IAllocator( OrtMemoryInfo(name, OrtAllocatorType::OrtDeviceAllocator, - OrtDevice(OrtDevice::GPU, OrtDevice::MemType::DEFAULT, static_cast(device_id)), + OrtDevice(OrtDevice::GPU, OrtDevice::MemType::DEFAULT, device_id), device_id, OrtMemTypeDefault)) {} virtual void* Alloc(size_t size) override; @@ -24,14 +24,14 @@ class MIGraphXAllocator : public IAllocator { void CheckDevice() const; }; -class MIGraphXExternalAllocator : public MIGraphXAllocator { +class HIPExternalAllocator : public HIPAllocator { typedef void* (*ExternalAlloc)(size_t size); typedef void (*ExternalFree)(void* p); typedef void (*ExternalEmptyCache)(); public: - MIGraphXExternalAllocator(OrtDevice::DeviceId device_id, const char* name, void* alloc, void* free, void* empty_cache) - : MIGraphXAllocator(device_id, name) { + HIPExternalAllocator(OrtDevice::DeviceId device_id, const char* name, void* alloc, void* free, void* empty_cache) + : HIPAllocator(device_id, name) { alloc_ = reinterpret_cast(alloc); free_ = reinterpret_cast(free); empty_cache_ = reinterpret_cast(empty_cache); @@ -55,7 +55,7 @@ class HIPPinnedAllocator : public IAllocator { HIPPinnedAllocator(int device_id, const char* name) : IAllocator( OrtMemoryInfo(name, OrtAllocatorType::OrtDeviceAllocator, - OrtDevice(OrtDevice::CPU, OrtDevice::MemType::HIP_PINNED, static_cast(device_id)), + OrtDevice(OrtDevice::CPU, OrtDevice::MemType::HIP_PINNED, device_id), device_id, OrtMemTypeCPUOutput)) {} virtual void* Alloc(size_t size) override; diff --git a/onnxruntime/core/providers/migraphx/migraphx_call.cc b/onnxruntime/core/providers/migraphx/migraphx_call.cc index 9807cd646e51c..5248ac2f39214 100644 --- a/onnxruntime/core/providers/migraphx/migraphx_call.cc +++ b/onnxruntime/core/providers/migraphx/migraphx_call.cc @@ -1,13 +1,10 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -#ifdef _WIN32 -#include -#else #include -#endif - -#include +#include +#include +#include #include "core/common/common.h" #include "core/common/status.h" #include "core/providers/shared_library/provider_api.h" @@ -37,20 +34,16 @@ std::conditional_t RocmCall( ERRTYPE retCode, const char* exprString, const char* libName, ERRTYPE successCode, const char* msg, const char* file, const int line) { if (retCode != successCode) { try { -#ifdef _WIN32 - // According to the POSIX spec, 255 is the safe minimum value. - static constexpr int HOST_NAME_MAX = 255; -#endif - std::string hostname(HOST_NAME_MAX, 0); - if (gethostname(hostname.data(), HOST_NAME_MAX) != 0) - hostname = "?"; + char hostname[HOST_NAME_MAX]; + if (gethostname(hostname, HOST_NAME_MAX) != 0) + strcpy(hostname, "?"); int currentHipDevice; (void)hipGetDevice(¤tHipDevice); (void)hipGetLastError(); // clear last HIP error static char str[1024]; snprintf(str, 1024, "%s failure %d: %s ; GPU=%d ; hostname=%s ; file=%s ; line=%d ; expr=%s; %s", libName, (int)retCode, RocmErrString(retCode), currentHipDevice, - hostname.c_str(), + hostname, file, line, exprString, msg); if constexpr (THRW) { // throw an exception with the error info @@ -75,5 +68,9 @@ std::conditional_t RocmCall( template Status RocmCall(hipError_t retCode, const char* exprString, const char* libName, hipError_t successCode, const char* msg, const char* file, const int line); template void RocmCall(hipError_t retCode, const char* exprString, const char* libName, hipError_t successCode, const char* msg, const char* file, const int line); +template Status RocmCall(rocblas_status retCode, const char* exprString, const char* libName, rocblas_status successCode, const char* msg, const char* file, const int line); +template void RocmCall(rocblas_status retCode, const char* exprString, const char* libName, rocblas_status successCode, const char* msg, const char* file, const int line); +template Status RocmCall(miopenStatus_t retCode, const char* exprString, const char* libName, miopenStatus_t successCode, const char* msg, const char* file, const int line); +template void RocmCall(miopenStatus_t retCode, const char* exprString, const char* libName, miopenStatus_t successCode, const char* msg, const char* file, const int line); } // namespace onnxruntime diff --git a/onnxruntime/core/providers/migraphx/migraphx_call.h b/onnxruntime/core/providers/migraphx/migraphx_call.h index f6a95cebf34b5..15d385a636b76 100644 --- a/onnxruntime/core/providers/migraphx/migraphx_call.h +++ b/onnxruntime/core/providers/migraphx/migraphx_call.h @@ -4,6 +4,8 @@ #pragma once #include "migraphx_inc.h" +#pragma once + namespace onnxruntime { // ----------------------------------------------------------------------- diff --git a/onnxruntime/core/providers/migraphx/migraphx_execution_provider.cc b/onnxruntime/core/providers/migraphx/migraphx_execution_provider.cc index 097b16ecde536..6ee85c3a4c047 100644 --- a/onnxruntime/core/providers/migraphx/migraphx_execution_provider.cc +++ b/onnxruntime/core/providers/migraphx/migraphx_execution_provider.cc @@ -13,11 +13,12 @@ #include "core/common/logging/severity.h" #include "migraphx_execution_provider.h" #include "migraphx_execution_provider_utils.h" -#include "migraphx_allocator.h" +#include "hip_allocator.h" #include "gpu_data_transfer.h" #include "migraphx_inc.h" -#include "migraphx_stream_handle.h" +// TODO: find a better way to share this +#include "core/providers/rocm/rocm_stream_handle.h" #if defined(_MSC_VER) #pragma warning(disable : 4244 4245) @@ -101,10 +102,10 @@ std::shared_ptr MIGraphXExecutionProvider::GetKernelRegistry() c } MIGraphXExecutionProvider::MIGraphXExecutionProvider(const MIGraphXExecutionProviderInfo& info) - : IExecutionProvider{onnxruntime::kMIGraphXExecutionProvider, OrtDevice(OrtDevice::GPU, OrtDevice::MemType::DEFAULT, info.device_id)}, info_(info) { + : IExecutionProvider{onnxruntime::kMIGraphXExecutionProvider, OrtDevice(OrtDevice::GPU, OrtDevice::MemType::DEFAULT, info.device_id)}, device_id_(info.device_id) { InitProviderOrtApi(); // Set GPU device to be used - HIP_CALL_THROW(hipSetDevice(info_.device_id)); + HIP_CALL_THROW(hipSetDevice(device_id_)); t_ = migraphx::target(info.target_device.c_str()); // whether fp16 is enable @@ -180,10 +181,16 @@ MIGraphXExecutionProvider::MIGraphXExecutionProvider(const MIGraphXExecutionProv dump_model_ops_ = (std::stoi(dump_model_ops_env) == 0 ? false : true); } + ROCBLAS_CALL_THROW(rocblas_create_handle(&external_rocblas_handle_)); + ROCBLAS_CALL_THROW(rocblas_set_stream(external_rocblas_handle_, stream_)); + + MIOPEN_CALL_THROW(miopenCreate(&external_miopen_handle_)); + MIOPEN_CALL_THROW(miopenSetStream(external_miopen_handle_, stream_)); + metadef_id_generator_ = ModelMetadefIdGenerator::Create(); LOGS_DEFAULT(VERBOSE) << "[MIGraphX EP] MIGraphX provider options: " - << "device_id: " << info_.device_id + << "device_id: " << device_id_ << ", migraphx_fp16_enable: " << fp16_enable_ << ", migraphx_int8_enable: " << int8_enable_ << ", migraphx_int8_enable: " << int8_enable_ @@ -198,14 +205,17 @@ MIGraphXExecutionProvider::MIGraphXExecutionProvider(const MIGraphXExecutionProv } MIGraphXExecutionProvider::~MIGraphXExecutionProvider() { + ORT_IGNORE_RETURN_VALUE(ROCBLAS_CALL(rocblas_destroy_handle(external_rocblas_handle_))); + ORT_IGNORE_RETURN_VALUE(MIOPEN_CALL(miopenDestroy(external_miopen_handle_))); } std::vector MIGraphXExecutionProvider::CreatePreferredAllocators() { AllocatorCreationInfo default_memory_info( - [](OrtDevice::DeviceId device_id) { return CreateMIGraphXAllocator(device_id, onnxruntime::CUDA); }, info_.device_id); + [](OrtDevice::DeviceId device_id) { return CreateROCMAllocator(device_id, onnxruntime::CUDA); }, device_id_); AllocatorCreationInfo pinned_allocator_info( [](OrtDevice::DeviceId device_id) { - return CreateMIGraphXPinnedAllocator(device_id, onnxruntime::CUDA_PINNED); + ORT_UNUSED_PARAMETER(device_id); + return CreateROCMPinnedAllocator(onnxruntime::CUDA_PINNED); }, 0); return std::vector{CreateAllocator(default_memory_info), CreateAllocator(pinned_allocator_info)}; @@ -244,40 +254,40 @@ static bool getMIGraphXType(ONNXTensorElementDataType type, migraphx_shape_datatype_t& mgx_type) { mgx_type = migraphx_shape_float_type; switch (type) { - case ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16: + case ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_FLOAT16: mgx_type = migraphx_shape_half_type; break; - case ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT: + case ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_FLOAT: mgx_type = migraphx_shape_float_type; break; - case ONNX_TENSOR_ELEMENT_DATA_TYPE_DOUBLE: + case ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_DOUBLE: mgx_type = migraphx_shape_double_type; break; - case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT8: + case ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_INT8: mgx_type = migraphx_shape_int8_type; break; - case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT16: + case ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_INT16: mgx_type = migraphx_shape_int16_type; break; - case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32: + case ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_INT32: mgx_type = migraphx_shape_int32_type; break; - case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64: + case ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_INT64: mgx_type = migraphx_shape_int64_type; break; - case ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT8: + case ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_UINT8: mgx_type = migraphx_shape_uint8_type; break; - case ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT16: + case ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_UINT16: mgx_type = migraphx_shape_uint16_type; break; - case ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT32: + case ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_UINT32: mgx_type = migraphx_shape_uint32_type; break; - case ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT64: + case ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_UINT64: mgx_type = migraphx_shape_uint64_type; break; - case ONNX_TENSOR_ELEMENT_DATA_TYPE_BOOL: + case ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_BOOL: mgx_type = migraphx_shape_bool_type; break; default: @@ -293,7 +303,7 @@ std::vector toVector(const ONNX_NAMESPACE::int64s& nums) { std::vector result; int num = nums.size(); for (int i = 0; i < num; ++i) { - result.push_back(static_cast(nums[i])); + result.push_back(nums[i]); } return result; @@ -491,9 +501,16 @@ static bool IsUnsupportedOpMode(const onnxruntime::GraphViewer& graph_viewer, co if (arg_s != nullptr) { const auto& tensor_dims = arg_s->dim(); std::vector dims; - for (auto&& dim : tensor_dims) { - dims.emplace_back(dim.has_dim_value() ? dim.dim_value() : 0); - } + std::transform(tensor_dims.begin(), + tensor_dims.end(), + std::back_inserter(dims), + [&](auto&& d) -> std::size_t { + if (d.has_dim_value()) { + return d.dim_value(); + } else { + return 0; + } + }); if (dims == std::vector{0}) { return true; } @@ -529,8 +546,8 @@ static bool IsUnsupportedOpMode(const onnxruntime::GraphViewer& graph_viewer, co } void SubgraphPostProcessing(const onnxruntime::GraphViewer& graph_viewer, std::vector>& clusters, - [[maybe_unused]] const logging::Logger& logger) { - // Then check whether a subgraph should fall back to CPU + const logging::Logger& logger) { + // Then check whether a subgraph should fallback to CPU // 1. Check whether a subgraph contains a RNN operator std::unordered_set rnn_names = {"RNN", "GRU", "LSTM"}; std::unordered_set op_names = {"AveragePool", "Conv", "Gemm", "LRN", "MatMul", "MaxPool"}; @@ -574,10 +591,17 @@ void SubgraphPostProcessing(const onnxruntime::GraphViewer& graph_viewer, std::v if (arg_s == nullptr) return false; const auto& tensor_dims = arg_s->dim(); std::vector dims; - for (auto&& dim : tensor_dims) { - dims.emplace_back(dim.has_dim_value() ? dim.dim_value() : 1); - } - return (std::accumulate(dims.begin(), dims.end(), 1ULL, std::multiplies{}) > 300); + std::transform(tensor_dims.begin(), + tensor_dims.end(), + std::back_inserter(dims), + [&](auto&& d) -> std::size_t { + if (d.has_dim_value()) { + return d.dim_value(); + } else { + return 1; + } + }); + return (std::accumulate(dims.begin(), dims.end(), 1, std::multiplies{}) > 300); })) { return false; } @@ -599,7 +623,7 @@ void SubgraphPostProcessing(const onnxruntime::GraphViewer& graph_viewer, std::v static bool IsNodeSupported(const std::set& op_set, const onnxruntime::GraphViewer& graph_viewer, const NodeIndex node_idx, - [[maybe_unused]] const logging::Logger& logger) { + const logging::Logger& logger) { const auto& node = graph_viewer.GetNode(node_idx); const auto& optype = node->OpType(); const auto& domain = node->Domain(); @@ -1418,10 +1442,14 @@ Status MIGraphXExecutionProvider::Compile(const std::vector& // lock to avoid race condition std::lock_guard lock(*(mgx_state->mgx_mu_ptr)); +#ifdef MIGRAPHX_STREAM_SYNC void* rocm_stream; Ort::ThrowOnError(api->KernelContext_GetGPUComputeStream(context, &rocm_stream)); auto prog_outputs = prog.run_async(m, static_cast(rocm_stream)); - +#else + auto prog_outputs = prog.eval(m); + HIP_CALL_THROW(hipDeviceSynchronize()); +#endif // In case of input parameters are reused as output parameter call hipMemcpy auto output_num = prog_outputs.size(); if (prog_output_indices.size() < output_num) { @@ -1450,7 +1478,8 @@ Status MIGraphXExecutionProvider::Compile(const std::vector& void MIGraphXExecutionProvider::RegisterStreamHandlers(IStreamCommandHandleRegistry& stream_handle_registry, AllocatorMap& allocators) const { auto allocator = allocators[GetOrtDeviceByMemType(OrtMemTypeCPU)]; - RegisterMIGraphXStreamHandles(stream_handle_registry, OrtDevice::GPU, allocator, true, stream_, false /*TODO:external_stream_*/); + RegisterRocmStreamHandles(stream_handle_registry, OrtDevice::GPU, allocator, true, stream_, + false /*TODO:external_stream_*/, external_miopen_handle_, external_rocblas_handle_); } OrtDevice MIGraphXExecutionProvider::GetOrtDeviceByMemType(OrtMemType mem_type) const { @@ -1458,6 +1487,7 @@ OrtDevice MIGraphXExecutionProvider::GetOrtDeviceByMemType(OrtMemType mem_type) if (mem_type == OrtMemTypeCPUOutput) return OrtDevice(OrtDevice::CPU, OrtDevice::MemType::HIP_PINNED, 0 /*CPU device id always be 0*/); return default_device_; } +#ifdef MIGRAPHX_STREAM_SYNC Status MIGraphXExecutionProvider::Sync() const { HIP_CALL_THROW(hipStreamSynchronize(static_cast(nullptr))); @@ -1482,4 +1512,5 @@ Status MIGraphXExecutionProvider::OnRunEnd(bool /*sync_stream*/, const onnxrunti return Status::OK(); } +#endif } // namespace onnxruntime diff --git a/onnxruntime/core/providers/migraphx/migraphx_execution_provider.h b/onnxruntime/core/providers/migraphx/migraphx_execution_provider.h index f34ca320d0a5a..1977f71b8b1cf 100644 --- a/onnxruntime/core/providers/migraphx/migraphx_execution_provider.h +++ b/onnxruntime/core/providers/migraphx/migraphx_execution_provider.h @@ -3,6 +3,9 @@ #pragma once +#include +#include + #include "core/framework/arena_extend_strategy.h" #include "core/framework/execution_provider.h" #include "core/platform/ort_mutex.h" @@ -11,6 +14,8 @@ #include #include +// TODO: find a better way to share this +// #include "core/providers/cuda/rocm_stream_handle.h" namespace onnxruntime { @@ -57,11 +62,13 @@ class MIGraphXExecutionProvider : public IExecutionProvider { explicit MIGraphXExecutionProvider(const MIGraphXExecutionProviderInfo& info); ~MIGraphXExecutionProvider(); +#ifdef MIGRAPHX_STREAM_SYNC Status Sync() const override; Status OnRunStart(const onnxruntime::RunOptions& run_options) override; Status OnRunEnd(bool sync_stream, const onnxruntime::RunOptions& run_options) override; +#endif std::vector> GetCapability(const onnxruntime::GraphViewer& graph_viewer, @@ -78,13 +85,7 @@ class MIGraphXExecutionProvider : public IExecutionProvider { OrtDevice GetOrtDeviceByMemType(OrtMemType mem_type) const override; std::vector CreatePreferredAllocators() override; - int GetDeviceId() const override { return info_.device_id; } - ProviderOptions GetProviderOptions() const override { - return MIGraphXExecutionProviderInfo::ToProviderOptions(info_); - } - private: - MIGraphXExecutionProviderInfo info_; bool fp16_enable_ = false; bool int8_enable_ = false; std::string int8_calibration_cache_name_; @@ -97,6 +98,7 @@ class MIGraphXExecutionProvider : public IExecutionProvider { bool load_compiled_model_ = false; std::string load_compiled_path_; bool dump_model_ops_ = false; + int device_id_; migraphx::target t_; OrtMutex mgx_mu_; hipStream_t stream_ = nullptr; @@ -107,6 +109,8 @@ class MIGraphXExecutionProvider : public IExecutionProvider { std::unordered_map map_no_input_shape_; AllocatorPtr allocator_; + miopenHandle_t external_miopen_handle_ = nullptr; + rocblas_handle external_rocblas_handle_ = nullptr; std::unique_ptr metadef_id_generator_; }; diff --git a/onnxruntime/core/providers/migraphx/migraphx_execution_provider_info.h b/onnxruntime/core/providers/migraphx/migraphx_execution_provider_info.h index 68d5d9af98ea4..8411e3eef096b 100644 --- a/onnxruntime/core/providers/migraphx/migraphx_execution_provider_info.h +++ b/onnxruntime/core/providers/migraphx/migraphx_execution_provider_info.h @@ -14,7 +14,7 @@ namespace onnxruntime { // Information needed to construct trt execution providers. struct MIGraphXExecutionProviderInfo { std::string target_device; - OrtDevice::DeviceId device_id{0}; + int device_id{0}; bool fp16_enable{false}; bool int8_enable{false}; std::string int8_calibration_table_name{""}; diff --git a/onnxruntime/core/providers/migraphx/migraphx_execution_provider_utils.h b/onnxruntime/core/providers/migraphx/migraphx_execution_provider_utils.h index 9274b5696185c..071070e92a209 100644 --- a/onnxruntime/core/providers/migraphx/migraphx_execution_provider_utils.h +++ b/onnxruntime/core/providers/migraphx/migraphx_execution_provider_utils.h @@ -28,7 +28,7 @@ bool IsGraphInput(const GraphViewer& graph, const std::string& name) { return (std::find(input_names.begin(), input_names.end(), name) != input_names.end()); } -bool IsGraphInitializer(const GraphViewer& graph, const std::string& name, [[maybe_unused]] bool check_outer_scope = true) { +bool IsGraphInitializer(const GraphViewer& graph, const std::string& name, bool check_outer_scope = true) { const ONNX_NAMESPACE::TensorProto* initializer = nullptr; return graph.GetInitializedTensor(name, initializer); } diff --git a/onnxruntime/core/providers/migraphx/migraphx_inc.h b/onnxruntime/core/providers/migraphx/migraphx_inc.h index 2b035b20f619f..96b24051ace76 100644 --- a/onnxruntime/core/providers/migraphx/migraphx_inc.h +++ b/onnxruntime/core/providers/migraphx/migraphx_inc.h @@ -4,5 +4,5 @@ #pragma once #include -#include +#include #include diff --git a/onnxruntime/core/providers/migraphx/migraphx_provider_factory.cc b/onnxruntime/core/providers/migraphx/migraphx_provider_factory.cc index 6d199930116e8..dd24dbdc76d2f 100644 --- a/onnxruntime/core/providers/migraphx/migraphx_provider_factory.cc +++ b/onnxruntime/core/providers/migraphx/migraphx_provider_factory.cc @@ -6,7 +6,7 @@ #include "core/providers/migraphx/migraphx_provider_factory.h" #include "migraphx_execution_provider.h" #include "migraphx_provider_factory_creator.h" -#include "migraphx_allocator.h" +#include "hip_allocator.h" #include "gpu_data_transfer.h" #include "core/framework/provider_options.h" @@ -33,23 +33,10 @@ std::unique_ptr MIGraphXProviderFactory::CreateProvider() { return std::make_unique(info_); } -struct ProviderInfo_MIGraphX_Impl final : ProviderInfo_MIGraphX { - std::unique_ptr CreateMIGraphXAllocator(int16_t device_id, const char* name) override { - return std::make_unique(device_id, name); - } - - std::unique_ptr CreateMIGraphXPinnedAllocator(int16_t device_id, const char* name) override { - return std::make_unique(device_id, name); - } - -} g_info; - struct MIGraphX_Provider : Provider { - void* GetInfo() override { return &g_info; } - std::shared_ptr CreateExecutionProviderFactory(int device_id) override { MIGraphXExecutionProviderInfo info; - info.device_id = static_cast(device_id); + info.device_id = device_id; info.target_device = "gpu"; return std::make_shared(info); } @@ -57,7 +44,7 @@ struct MIGraphX_Provider : Provider { std::shared_ptr CreateExecutionProviderFactory(const void* provider_options) override { auto& options = *reinterpret_cast(provider_options); MIGraphXExecutionProviderInfo info; - info.device_id = static_cast(options.device_id); + info.device_id = options.device_id; info.target_device = "gpu"; info.fp16_enable = options.migraphx_fp16_enable; info.int8_enable = options.migraphx_int8_enable; diff --git a/onnxruntime/core/providers/migraphx/migraphx_provider_factory.h b/onnxruntime/core/providers/migraphx/migraphx_provider_factory.h index b257a4318dc0e..ac9834e64942a 100644 --- a/onnxruntime/core/providers/migraphx/migraphx_provider_factory.h +++ b/onnxruntime/core/providers/migraphx/migraphx_provider_factory.h @@ -10,13 +10,4 @@ struct IExecutionProviderFactory; struct MIGraphXExecutionProviderInfo; enum class ArenaExtendStrategy : int32_t; struct MIGraphXExecutionProviderExternalAllocatorInfo; - -struct ProviderInfo_MIGraphX { - virtual std::unique_ptr CreateMIGraphXAllocator(int16_t device_id, const char* name) = 0; - virtual std::unique_ptr CreateMIGraphXPinnedAllocator(int16_t device_id, const char* name) = 0; - - protected: - ~ProviderInfo_MIGraphX() = default; // Can only be destroyed through a subclass instance -}; - } // namespace onnxruntime diff --git a/onnxruntime/core/providers/migraphx/migraphx_stream_handle.cc b/onnxruntime/core/providers/migraphx/migraphx_stream_handle.cc deleted file mode 100644 index 9c5bb4ecf5c97..0000000000000 --- a/onnxruntime/core/providers/migraphx/migraphx_stream_handle.cc +++ /dev/null @@ -1,171 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT License. - -#include -#include "migraphx_stream_handle.h" - -namespace onnxruntime { - -struct MIGraphXNotification : public synchronize::Notification { - MIGraphXNotification(Stream& s) : Notification(s) { - HIP_CALL_THROW(hipEventCreateWithFlags(&event_, hipEventDisableTiming)); - } - - ~MIGraphXNotification() { - if (event_) - HIP_CALL_THROW(hipEventDestroy(event_)); - } - - void Activate() override { - // record event with hipEventBlockingSync so we can support sync on host without busy wait. - HIP_CALL_THROW(hipEventRecord(event_, static_cast(stream_.GetHandle()))); - } - - void wait_on_device(Stream& device_stream) { - ORT_ENFORCE(device_stream.GetDevice().Type() == OrtDevice::GPU, "Unexpected device:", device_stream.GetDevice().ToString()); - // launch a wait command to the migraphx stream - HIP_CALL_THROW(hipStreamWaitEvent(static_cast(device_stream.GetHandle()), event_, 0)); - }; - - void wait_on_host() { - // CUDA_CALL_THROW(cudaStreamSynchronize(stream_)); - HIP_CALL_THROW(hipEventSynchronize(event_)); - } - - hipEvent_t event_; -}; - -MIGraphXStream::MIGraphXStream(hipStream_t stream, - const OrtDevice& device, - AllocatorPtr cpu_allocator, - bool release_cpu_buffer_on_migraphx_stream) - : Stream(stream, device), - cpu_allocator_(cpu_allocator), - release_cpu_buffer_on_migraphx_stream_(release_cpu_buffer_on_migraphx_stream) { -} - -MIGraphXStream::~MIGraphXStream() { - ORT_IGNORE_RETURN_VALUE(CleanUpOnRunEnd()); - if (own_stream_) { - auto* handle = GetHandle(); - if (handle) - HIP_CALL_THROW(hipStreamDestroy(static_cast(handle))); - } -} - -std::unique_ptr MIGraphXStream::CreateNotification(size_t /*num_consumers*/) { - return std::make_unique(*this); -} - -void MIGraphXStream::Flush() { - if (own_stream_) - HIP_CALL_THROW(hipStreamSynchronize(static_cast(GetHandle()))); -} - -void MIGraphXStream::EnqueDeferredCPUBuffer(void* cpu_buffer) { - // stream is per thread, so don't need lock - deferred_cpu_buffers_.push_back(cpu_buffer); -} - -struct CpuBuffersInfo { - // This struct stores the information needed - // to release CPU buffers allocated for GPU kernels. - // It's used to enqueue their release after - // associated GPU kernels in a MIGraphX stream. - - // This is a CPU allocator in MIGraphX EP. - // It must be the one used to allocate the - // following pointers. - AllocatorPtr allocator; - // buffers[i] is the i-th pointer added by - // AddDeferredReleaseCPUPtr for a specific - // MIGraphX stream. For example, this fields - // should contain all values in - // deferred_release_buffer_pool_[my_stream] - // when release my_stream's buffers. - std::unique_ptr buffers; - // CPU buffer buffers[i]. - // Number of buffer points in "buffers". - size_t n_buffers; -}; - -static void ReleaseCpuBufferCallback(void* raw_info) { - std::unique_ptr info = std::make_unique(); - info.reset(reinterpret_cast(raw_info)); - for (size_t i = 0; i < info->n_buffers; ++i) { - info->allocator->Free(info->buffers[i]); - } -} - -Status MIGraphXStream::CleanUpOnRunEnd() { - if (deferred_cpu_buffers_.empty()) - return Status::OK(); - // Release the ownership of cpu_buffers_info so that the underlying - // object will keep alive until the end of ReleaseCpuBufferCallback. - if (release_cpu_buffer_on_migraphx_stream_ && cpu_allocator_->Info().alloc_type == OrtArenaAllocator) { - std::unique_ptr cpu_buffers_info = std::make_unique(); - cpu_buffers_info->allocator = cpu_allocator_; - cpu_buffers_info->buffers = std::make_unique(deferred_cpu_buffers_.size()); - for (size_t i = 0; i < deferred_cpu_buffers_.size(); ++i) { - cpu_buffers_info->buffers[i] = deferred_cpu_buffers_.at(i); - } - cpu_buffers_info->n_buffers = deferred_cpu_buffers_.size(); - HIP_RETURN_IF_ERROR(hipLaunchHostFunc(static_cast(GetHandle()), ReleaseCpuBufferCallback, cpu_buffers_info.release())); - } else { - HIP_RETURN_IF_ERROR(hipStreamSynchronize(static_cast(GetHandle()))); - for (auto* buffer : deferred_cpu_buffers_) { - cpu_allocator_->Free(buffer); - } - } - - deferred_cpu_buffers_.clear(); - return Status::OK(); -} - -void* MIGraphXStream::GetResource(int version, int id) const { - ORT_ENFORCE(version <= ORT_ROCM_RESOUCE_VERSION, "resource version unsupported!"); - void* resource{}; - switch (id) { - case RocmResource::hip_stream_t: - return reinterpret_cast(GetHandle()); - default: - break; - } - return resource; -} - -// CPU Stream command handles -void WaitMIGraphXNotificationOnDevice(Stream& stream, synchronize::Notification& notification) { - static_cast(¬ification)->wait_on_device(stream); -} - -void WaitMIGraphXNotificationOnHost(Stream& /*stream*/, synchronize::Notification& notification) { - static_cast(¬ification)->wait_on_host(); -} - -void RegisterMIGraphXStreamHandles(IStreamCommandHandleRegistry& stream_handle_registry, - const OrtDevice::DeviceType device_type, - AllocatorPtr cpu_allocator, - bool release_cpu_buffer_on_migraphx_stream, - hipStream_t external_stream, - bool use_existing_stream) { - // wait migraphx notification on migraphx ep - stream_handle_registry.RegisterWaitFn(device_type, device_type, WaitMIGraphXNotificationOnDevice); - // wait migraphx notification on cpu ep - stream_handle_registry.RegisterWaitFn(device_type, OrtDevice::CPU, WaitMIGraphXNotificationOnHost); - if (!use_existing_stream) - stream_handle_registry.RegisterCreateStreamFn(device_type, [cpu_allocator, release_cpu_buffer_on_migraphx_stream](const OrtDevice& device) { - HIP_CALL_THROW(hipSetDevice(device.Id())); - hipStream_t stream = nullptr; - HIP_CALL_THROW(hipStreamCreateWithFlags(&stream, hipStreamNonBlocking)); - return std::make_unique(stream, device, cpu_allocator, release_cpu_buffer_on_migraphx_stream); - }); - else - stream_handle_registry.RegisterCreateStreamFn(device_type, [cpu_allocator, - release_cpu_buffer_on_migraphx_stream, - external_stream](const OrtDevice& device) { - return std::make_unique(external_stream, device, cpu_allocator, release_cpu_buffer_on_migraphx_stream); - }); -} - -} // namespace onnxruntime diff --git a/onnxruntime/core/providers/migraphx/migraphx_stream_handle.h b/onnxruntime/core/providers/migraphx/migraphx_stream_handle.h deleted file mode 100644 index 03a7c1607e3ad..0000000000000 --- a/onnxruntime/core/providers/migraphx/migraphx_stream_handle.h +++ /dev/null @@ -1,48 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT License. - -#pragma once -#include "core/framework/stream_handles.h" -#include "migraphx_inc.h" -#include "migraphx_call.h" - -#define HIP_RETURN_IF_ERROR(expr) ORT_RETURN_IF_ERROR(HIP_CALL(expr)) - -namespace onnxruntime { -void WaitMIGraphXNotificationOnDevice(Stream& stream, synchronize::Notification& notification); - -struct MIGraphXStream : Stream { - MIGraphXStream(hipStream_t stream, - const OrtDevice& device, - AllocatorPtr cpu_allocator, - bool release_cpu_buffer_on_migraphx_stream); - - ~MIGraphXStream(); - - std::unique_ptr CreateNotification(size_t /*num_consumers*/) override; - - void Flush() override; - - Status CleanUpOnRunEnd() override; - - void EnqueDeferredCPUBuffer(void* cpu_buffer); - - bool own_stream_{true}; - - virtual void* GetResource(int version, int id) const; - - virtual WaitNotificationFn GetWaitNotificationFn() const { return WaitMIGraphXNotificationOnDevice; } - - private: - std::vector deferred_cpu_buffers_; - AllocatorPtr cpu_allocator_; - bool release_cpu_buffer_on_migraphx_stream_{true}; -}; - -void RegisterMIGraphXStreamHandles(IStreamCommandHandleRegistry& stream_handle_registry, - const OrtDevice::DeviceType device_type, - AllocatorPtr cpu_allocator, - bool release_cpu_buffer_on_migraphx_stream, - hipStream_t external_stream, - bool use_existing_stream); -} // namespace onnxruntime diff --git a/onnxruntime/core/providers/shared_library/provider_api.h b/onnxruntime/core/providers/shared_library/provider_api.h index 8917bb7fd9bb6..7cdfb0ffc19f2 100644 --- a/onnxruntime/core/providers/shared_library/provider_api.h +++ b/onnxruntime/core/providers/shared_library/provider_api.h @@ -279,9 +279,6 @@ std::unique_ptr CreateCPUAllocator(const OrtMemoryInfo& memory_info) std::unique_ptr CreateCUDAAllocator(int16_t device_id, const char* name); std::unique_ptr CreateCUDAPinnedAllocator(const char* name); -std::unique_ptr CreateMIGraphXAllocator(int16_t device_id, const char* name); -std::unique_ptr CreateMIGraphXPinnedAllocator(int16_t device_id, const char* name); - std::unique_ptr CreateROCMAllocator(int16_t device_id, const char* name); std::unique_ptr CreateROCMPinnedAllocator(const char* name); diff --git a/onnxruntime/core/providers/shared_library/provider_bridge_provider.cc b/onnxruntime/core/providers/shared_library/provider_bridge_provider.cc index 540f671d67f8d..27d8a0f06f565 100644 --- a/onnxruntime/core/providers/shared_library/provider_bridge_provider.cc +++ b/onnxruntime/core/providers/shared_library/provider_bridge_provider.cc @@ -353,12 +353,16 @@ std::unique_ptr CreateGPUDataTransfer() { #endif #ifdef USE_MIGRAPHX -std::unique_ptr CreateMIGraphXAllocator(int16_t device_id, const char* name) { - return g_host->CreateMIGraphXAllocator(device_id, name); +std::unique_ptr CreateROCMAllocator(int16_t device_id, const char* name) { + return g_host->CreateROCMAllocator(device_id, name); } -std::unique_ptr CreateMIGraphXPinnedAllocator(int16_t device_id, const char* name) { - return g_host->CreateMIGraphXPinnedAllocator(device_id, name); +std::unique_ptr CreateROCMPinnedAllocator(const char* name) { + return g_host->CreateROCMPinnedAllocator(name); +} + +std::unique_ptr CreateGPUDataTransfer() { + return g_host->CreateGPUDataTransfer(); } #endif diff --git a/onnxruntime/core/providers/shared_library/provider_interfaces.h b/onnxruntime/core/providers/shared_library/provider_interfaces.h index f1a778e8b8f80..cc3b13f696a96 100644 --- a/onnxruntime/core/providers/shared_library/provider_interfaces.h +++ b/onnxruntime/core/providers/shared_library/provider_interfaces.h @@ -178,11 +178,6 @@ struct ProviderHost { virtual void CudaCall_true(int retCode, const char* exprString, const char* libName, int successCode, const char* msg, const char* file, const int line) = 0; #endif -#ifdef USE_MIGRAPHX - virtual std::unique_ptr CreateMIGraphXAllocator(int16_t device_id, const char* name) = 0; - virtual std::unique_ptr CreateMIGraphXPinnedAllocator(int16_t device_id, const char* name) = 0; -#endif - #ifdef USE_ROCM virtual std::unique_ptr CreateROCMAllocator(int16_t device_id, const char* name) = 0; virtual std::unique_ptr CreateROCMPinnedAllocator(const char* name) = 0; diff --git a/onnxruntime/core/session/provider_bridge_ort.cc b/onnxruntime/core/session/provider_bridge_ort.cc index a4959399990c5..7f7ed5e436afe 100644 --- a/onnxruntime/core/session/provider_bridge_ort.cc +++ b/onnxruntime/core/session/provider_bridge_ort.cc @@ -130,8 +130,6 @@ ProviderInfo_Dnnl& GetProviderInfo_Dnnl(); ProviderInfo_ROCM* TryGetProviderInfo_ROCM(); ProviderInfo_ROCM& GetProviderInfo_ROCM(); ProviderHostCPU& GetProviderHostCPU(); -ProviderInfo_MIGraphX* TryGetProviderInfo_MIGraphX(); -ProviderInfo_MIGraphX& GetProviderInfo_MIGraphX(); ONNX_NAMESPACE::OpSchema CreateSchema(const std::string& domain, const std::vector& ops); struct TensorShapeProto_Dimension_Iterator_Impl : TensorShapeProto_Dimension_Iterator { TensorShapeProto_Dimension_Iterator_Impl(google::protobuf::internal::RepeatedPtrIterator&& v) : v_{std::move(v)} {} @@ -243,11 +241,6 @@ struct ProviderHostImpl : ProviderHost { void CudaCall_true(int retCode, const char* exprString, const char* libName, int successCode, const char* msg, const char* file, const int line) override { GetProviderInfo_CUDA().CudaCall_true(retCode, exprString, libName, successCode, msg, file, line); } #endif -#ifdef USE_MIGRAPHX - std::unique_ptr CreateMIGraphXAllocator(int16_t device_id, const char* name) override { return GetProviderInfo_MIGraphX().CreateMIGraphXAllocator(device_id, name); } - std::unique_ptr CreateMIGraphXPinnedAllocator(int16_t device_id, const char* name) override { return GetProviderInfo_MIGraphX().CreateMIGraphXPinnedAllocator(device_id, name); } -#endif - #ifdef USE_ROCM std::unique_ptr CreateROCMAllocator(int16_t device_id, const char* name) override { return GetProviderInfo_ROCM().CreateROCMAllocator(device_id, name); } std::unique_ptr CreateROCMPinnedAllocator(const char* name) override { return GetProviderInfo_ROCM().CreateROCMPinnedAllocator(name); } @@ -1907,20 +1900,6 @@ ProviderInfo_ROCM& GetProviderInfo_ROCM() { ORT_THROW("ROCM Provider not available, can't get interface for it"); } -ProviderInfo_MIGraphX* TryGetProviderInfo_MIGraphX() try { - return reinterpret_cast(s_library_migraphx.Get().GetInfo()); -} catch (const std::exception& exception) { - LOGS_DEFAULT(ERROR) << exception.what(); - return nullptr; -} - -ProviderInfo_MIGraphX& GetProviderInfo_MIGraphX() { - if (auto* info = TryGetProviderInfo_MIGraphX()) - return *info; - - ORT_THROW("MIGraphX Provider not available, can't get interface for it"); -} - void CopyGpuToCpu( void* dst_ptr, const void* src_ptr, diff --git a/setup.py b/setup.py index baab399872b0f..3203993e0c4d4 100644 --- a/setup.py +++ b/setup.py @@ -56,7 +56,6 @@ def parse_arg_remove_string(argv, arg_name_equal): cuda_version = None rocm_version = None -is_migraphx = False is_rocm = False is_openvino = False # The following arguments are mutually exclusive @@ -65,9 +64,8 @@ def parse_arg_remove_string(argv, arg_name_equal): cuda_version = parse_arg_remove_string(sys.argv, "--cuda_version=") elif parse_arg_remove_boolean(sys.argv, "--use_rocm"): is_rocm = True + package_name = "onnxruntime-rocm" if not nightly_build else "ort-rocm-nightly" rocm_version = parse_arg_remove_string(sys.argv, "--rocm_version=") -elif parse_arg_remove_boolean(sys.argv, "--use_migraphx"): - is_migraphx = True elif parse_arg_remove_boolean(sys.argv, "--use_openvino"): is_openvino = True package_name = "onnxruntime-openvino" @@ -89,9 +87,6 @@ def parse_arg_remove_string(argv, arg_name_equal): elif parse_arg_remove_boolean(sys.argv, "--use_qnn"): package_name = "onnxruntime-qnn" -if is_rocm or is_migraphx: - package_name = "onnxruntime-rocm" if not nightly_build else "ort-rocm-nightly" - # PEP 513 defined manylinux1_x86_64 and manylinux1_i686 # PEP 571 defined manylinux2010_x86_64 and manylinux2010_i686 # PEP 599 defines the following platform tags: @@ -285,21 +280,10 @@ def finalize_options(self): return ret -providers_cuda_or_rocm = "onnxruntime_providers_" + ("rocm" if is_rocm else "cuda") -providers_tensorrt_or_migraphx = "onnxruntime_providers_" + ("migraphx" if is_migraphx else "tensorrt") -providers_openvino = "onnxruntime_providers_openvino" -providers_cann = "onnxruntime_providers_cann" - -if platform.system() == "Linux": - providers_cuda_or_rocm = "lib" + providers_cuda_or_rocm + ".so" - providers_tensorrt_or_migraphx = "lib" + providers_tensorrt_or_migraphx + ".so" - providers_openvino = "lib" + providers_openvino + ".so" - providers_cann = "lib" + providers_cann + ".so" -elif platform.system() == "Windows": - providers_cuda_or_rocm = providers_cuda_or_rocm + ".dll" - providers_tensorrt_or_migraphx = providers_tensorrt_or_migraphx + ".dll" - providers_openvino = providers_openvino + ".dll" - providers_cann = providers_cann + ".dll" +providers_cuda_or_rocm = "libonnxruntime_providers_" + ("rocm.so" if is_rocm else "cuda.so") +providers_tensorrt_or_migraphx = "libonnxruntime_providers_" + ("migraphx.so" if is_rocm else "tensorrt.so") +providers_openvino = "libonnxruntime_providers_openvino.so" +providers_cann = "libonnxruntime_providers_cann.so" # Additional binaries dl_libs = [] @@ -313,22 +297,19 @@ def finalize_options(self): "libmklml_gnu.so", "libiomp5.so", "mimalloc.so", - # DNNL, TensorRT & OpenVINO EPs are built as shared libs - "libonnxruntime_providers_shared.so", - "libonnxruntime_providers_dnnl.so", - "libonnxruntime_providers_openvino.so", - "libonnxruntime_providers_vitisai.so", - providers_cuda_or_rocm, - providers_tensorrt_or_migraphx, - providers_cann, ] - dl_libs = [ - "libonnxruntime_providers_shared.so", - providers_cuda_or_rocm, - providers_tensorrt_or_migraphx, - providers_cann, - ] - + dl_libs = ["libonnxruntime_providers_shared.so"] + dl_libs.append(providers_cuda_or_rocm) + dl_libs.append(providers_tensorrt_or_migraphx) + dl_libs.append(providers_cann) + # DNNL, TensorRT & OpenVINO EPs are built as shared libs + libs.extend(["libonnxruntime_providers_shared.so"]) + libs.extend(["libonnxruntime_providers_dnnl.so"]) + libs.extend(["libonnxruntime_providers_openvino.so"]) + libs.extend(["libonnxruntime_providers_vitisai.so"]) + libs.append(providers_cuda_or_rocm) + libs.append(providers_tensorrt_or_migraphx) + libs.append(providers_cann) if nightly_build: libs.extend(["libonnxruntime_pywrapper.so"]) elif platform.system() == "Darwin": @@ -342,15 +323,7 @@ def finalize_options(self): if nightly_build: libs.extend(["libonnxruntime_pywrapper.dylib"]) else: - libs = [ - "onnxruntime_pybind11_state.pyd", - "dnnl.dll", - "mklml.dll", - "libiomp5md.dll", - providers_cuda_or_rocm, - providers_tensorrt_or_migraphx, - providers_cann, - ] + libs = ["onnxruntime_pybind11_state.pyd", "dnnl.dll", "mklml.dll", "libiomp5md.dll"] # DNNL, TensorRT & OpenVINO EPs are built as shared libs libs.extend(["onnxruntime_providers_shared.dll"]) libs.extend(["onnxruntime_providers_dnnl.dll"]) diff --git a/tools/ci_build/build.py b/tools/ci_build/build.py index 6159f078828f3..3e587e9b56e2e 100644 --- a/tools/ci_build/build.py +++ b/tools/ci_build/build.py @@ -611,7 +611,6 @@ def convert_arg_line_to_args(self, arg_line): "MinGW Makefiles", "Ninja", "NMake Makefiles", - "NMake Makefiles JOM", "Unix Makefiles", "Visual Studio 17 2022", "Xcode", @@ -2208,7 +2207,6 @@ def build_python_wheel( use_cuda, cuda_version, use_rocm, - use_migraphx, rocm_version, use_dnnl, use_tensorrt, @@ -2260,8 +2258,6 @@ def build_python_wheel( args.append("--use_rocm") if rocm_version: args.append(f"--rocm_version={rocm_version}") - elif use_migraphx: - args.append("--use_migraphx") elif use_openvino: args.append("--use_openvino") elif use_dnnl: @@ -2587,6 +2583,9 @@ def main(): if args.use_tensorrt: args.use_cuda = True + if args.use_migraphx: + args.use_rocm = True + if args.build_wheel or args.gen_doc or args.use_tvm or args.enable_training: args.enable_pybind = True @@ -2873,8 +2872,7 @@ def main(): # fail unexpectedly. Similar, if your packaging step forgot to copy a file into the package, we don't know it # either. if args.build: - # TODO: find asan DLL and copy it to onnxruntime/capi folder when args.enable_address_sanitizer is True and - # the target OS is Windows + # TODO: find asan DLL and copy it to onnxruntime/capi folder when args.enable_address_sanitizer is True and the target OS is Windows if args.build_wheel: nightly_build = bool(os.getenv("NIGHTLY_BUILD") == "1") default_training_package_device = bool(os.getenv("DEFAULT_TRAINING_PACKAGE_DEVICE") == "1") @@ -2885,7 +2883,6 @@ def main(): args.use_cuda, args.cuda_version, args.use_rocm, - args.use_migraphx, args.rocm_version, args.use_dnnl, args.use_tensorrt, From 7cf9263ee741222bbc1744ca2970d5786d23e9e6 Mon Sep 17 00:00:00 2001 From: RuomeiMS Date: Fri, 21 Jun 2024 16:23:23 +0100 Subject: [PATCH 17/52] Add changes for strided calibration (#20949) Context and motivation: When quantizing large transformer models, we faced OOM issue when the number of calibration samples goes up. To resolve this, in the PR we want to add support for reading quantization data in chunck, calculating ranges for intermediate tensors, then accumulating results for the final ranges. --- .../python/tools/quantization/calibrate.py | 37 +++++- .../execution_providers/qnn/quant_config.py | 2 + .../python/tools/quantization/quantize.py | 17 ++- .../test/python/quantization/op_test_utils.py | 119 +++++++++++++++++- .../quantization/test_quantize_static.py | 56 ++++++++- 5 files changed, 222 insertions(+), 9 deletions(-) diff --git a/onnxruntime/python/tools/quantization/calibrate.py b/onnxruntime/python/tools/quantization/calibrate.py index 3f5e4e660003f..10492ae419817 100644 --- a/onnxruntime/python/tools/quantization/calibrate.py +++ b/onnxruntime/python/tools/quantization/calibrate.py @@ -128,6 +128,9 @@ def __setitem__(self, key, value): def values(self): return self.data.values() + def items(self): + return self.data.items() + class CalibrationMethod(Enum): MinMax = 0 @@ -155,6 +158,12 @@ def __next__(self): raise StopIteration return result + def __len__(self): + raise NotImplementedError + + def set_range(self, start_index: int, end_index: int): + raise NotImplementedError + class CalibraterBase: def __init__( @@ -409,13 +418,31 @@ def merge_range(self, old_range, new_range): return new_range for key, value in old_range.items(): + # Handling for structured data types with TensorData + if isinstance(value, TensorData): + old_min = value.range_value[0] + old_max = value.range_value[1] + else: + old_min, old_max = value + + if isinstance(new_range[key], TensorData): + new_min = new_range[key].range_value[0] + new_max = new_range[key].range_value[1] + else: + new_min, new_max = new_range[key] + if self.moving_average: - min_value = value[0] + self.averaging_constant * (new_range[key][0] - value[0]) - max_value = value[1] + self.averaging_constant * (new_range[key][1] - value[1]) + min_value = old_min + self.averaging_constant * (new_min - old_min) + max_value = old_max + self.averaging_constant * (new_max - old_max) + else: + min_value = min(old_min, new_min) + max_value = max(old_max, new_max) + + # If structured as TensorData, wrap the result accordingly + if isinstance(value, TensorData) or isinstance(new_range[key], TensorData): + new_range[key] = TensorData(lowest=min_value, highest=max_value) else: - min_value = min(value[0], new_range[key][0]) - max_value = max(value[1], new_range[key][1]) - new_range[key] = (min_value, max_value) + new_range[key] = (min_value, max_value) return new_range diff --git a/onnxruntime/python/tools/quantization/execution_providers/qnn/quant_config.py b/onnxruntime/python/tools/quantization/execution_providers/qnn/quant_config.py index 1ad56dc3ac455..eac5b3b78690b 100644 --- a/onnxruntime/python/tools/quantization/execution_providers/qnn/quant_config.py +++ b/onnxruntime/python/tools/quantization/execution_providers/qnn/quant_config.py @@ -52,6 +52,7 @@ def get_qnn_qdq_config( activation_symmetric: bool = False, weight_symmetric: bool | None = None, keep_removable_activations: bool = False, + stride: int | None = None, ) -> StaticQuantConfig: """ Returns a static quantization configuration suitable for running QDQ models on QNN EP. @@ -171,6 +172,7 @@ def get_qnn_qdq_config( "TensorQuantOverrides": overrides_helper.get_dict(), "ActivationSymmetric": activation_symmetric, "WeightSymmetric": weight_symmetric, + "CalibStridedMinMax": stride, } # ONNX opset < 21 does not support 16-bit quantization, so must use 'com.microsoft' domain diff --git a/onnxruntime/python/tools/quantization/quantize.py b/onnxruntime/python/tools/quantization/quantize.py index f8b74a7ae4c2e..2340c995d3d5b 100644 --- a/onnxruntime/python/tools/quantization/quantize.py +++ b/onnxruntime/python/tools/quantization/quantize.py @@ -381,6 +381,9 @@ def quantize_static( CalibTensorRangeSymmetric = True/False : Default is False. If enabled, the final range of tensor during calibration will be explicitly set to symmetric to central point "0". + CalibStridedMinMax = Optional[int] : + Default is None. If set to an integer, during calculation of the min-max, only stride amount of + data will be used and then all results will be merged in the end. CalibMovingAverage = True/False : Default is False. If enabled, the moving average of the minimum and maximum values will be computed when the calibration method selected is MinMax. @@ -522,7 +525,19 @@ def inc_dataloader(): use_external_data_format=use_external_data_format, extra_options=calib_extra_options, ) - calibrator.collect_data(calibration_data_reader) + + stride = extra_options.get("CalibStridedMinMax", None) + if stride: + total_data_size = len(calibration_data_reader) + if total_data_size % stride != 0: + raise ValueError(f"Total data size ({total_data_size}) is not divisible by stride size ({stride}).") + + for start in range(0, total_data_size, stride): + end_index = start + stride + calibration_data_reader.set_range(start_index=start, end_index=end_index) + calibrator.collect_data(calibration_data_reader) + else: + calibrator.collect_data(calibration_data_reader) tensors_range = calibrator.compute_data() if not isinstance(tensors_range, TensorsData): raise TypeError( diff --git a/onnxruntime/test/python/quantization/op_test_utils.py b/onnxruntime/test/python/quantization/op_test_utils.py index b30282f2ab41f..cf7fc292ea86b 100644 --- a/onnxruntime/test/python/quantization/op_test_utils.py +++ b/onnxruntime/test/python/quantization/op_test_utils.py @@ -217,10 +217,13 @@ def rewind(self): self.iter_next = iter(self.data_feeds) -def input_feeds_neg_one_zero_one(n, name2shape): +def input_feeds_neg_one_zero_one(n, name2shape, seed=None): """ randomize n feed according to shape, its values are from -1, 0, and 1 """ + if seed is not None: + np.random.seed(seed) + input_data_list = [] for _i in range(n): inputs = {} @@ -231,6 +234,120 @@ def input_feeds_neg_one_zero_one(n, name2shape): return dr +def input_feeds_neg_one_zero_one_list(n, name2shape, seed=None): + """ + randomize n feed according to shape, its values are from -1, 0, and 1 + """ + if seed is not None: + np.random.seed(seed) + + input_data_list = [] + for _i in range(n): + inputs = {} + for name, shape in name2shape.items(): + inputs.update({name: np.random.randint(-1, 2, shape).astype(np.float32)}) + input_data_list.extend([inputs]) + return input_data_list + + +class GenerateCalibrationData(CalibrationDataReader): + def __init__(self, data_list, input_nodes, input_shapes, no_tensor_num, in_dtypes, inputs_conv_channel_last=None): + print("Generating calibration dataset from " + str(data_list)) + print("input nodes are ", input_nodes, "input shapes are ", input_shapes) + if inputs_conv_channel_last: + print(f"Inputs that will be converted to channel last: {inputs_conv_channel_last}") + + self.enum_data_dicts = [] + self.input_nodes = input_nodes + self.input_shapes = input_shapes + self.inputs_conv_channel_last = inputs_conv_channel_last + self.calibration_dataset = data_list + + def __len__(self): + return len(self.calibration_dataset) + + def get_next(self): + feed_dict = {} + inp = next(self.calibration_dataset, None) + if inp is not None: + for i in range(len(self.input_nodes)): + input_data = inp[i].reshape(self.input_shapes[i]) + if self.inputs_conv_channel_last is not None and self.input_nodes[i] in self.inputs_conv_channel_last: + input_data = np.moveaxis(input_data, 1, -1) + dict_item = {self.input_nodes[i]: input_data} + feed_dict.update(dict_item) + return feed_dict + else: + return None + + +class StridedDataReader(GenerateCalibrationData): + def __init__( + self, + data_list, + input_nodes, + input_shapes, + no_tensor_num, + in_dtypes, + inputs_conv_channel_last=None, + stride=1, + start_index=0, + end_index=None, + ): + super().__init__(data_list, input_nodes, input_shapes, no_tensor_num, in_dtypes, inputs_conv_channel_last) + + self.stride = max(1, stride) # Ensure stride is at least 1 + self.start_index = start_index + self.end_index = ( + end_index if end_index is not None else len(self.calibration_dataset) + ) # Default to the end of the dataset + self.enum_data_dicts = iter([]) + + def get_next(self): + iter_data = next(self.enum_data_dicts, None) + if iter_data: + return iter_data + + self.enum_data_dicts = None + if self.start_index < self.end_index: + print(f"start index is {self.start_index}") + data = self.load_serial() + + self.start_index += self.stride + self.enum_data_dicts = iter(data) + + return next(self.enum_data_dicts, None) + else: + return None + + def load_serial(self): + batch_data = [] + end_loop = min(self.end_index, self.start_index + self.stride) + for i in range(self.start_index, end_loop): + print(f"debugging the load serial index {i}") + data_item = self.calibration_dataset[i] + processed_item = self.process_data_item(data_item) + batch_data.append(processed_item) + return batch_data + + def process_data_item(self, data_item): + feed_dict = {} + for _, node in enumerate(self.input_nodes): + # input_data = data_item[i].reshape(self.input_shapes[i]) + feed_dict[node] = data_item["input"] + return feed_dict + + def set_range(self, start_index, end_index=None): + self.start_index = start_index + self.end_index = end_index if end_index is not None else len(self.calibration_dataset) + self.enum_data_dicts = iter([]) + + def rewind(self): + """Rewind the data reader to the beginning of the dataset.""" + self.start_index = 0 + self.enum_data_dicts = iter([]) + + def check_op_type_order(testcase, model_to_check, ops): if isinstance(model_to_check, str): model = onnx.load(model_to_check) diff --git a/onnxruntime/test/python/quantization/test_quantize_static.py b/onnxruntime/test/python/quantization/test_quantize_static.py index 5ad5a49f00c14..01976ba633137 100644 --- a/onnxruntime/test/python/quantization/test_quantize_static.py +++ b/onnxruntime/test/python/quantization/test_quantize_static.py @@ -13,8 +13,15 @@ import numpy as np import onnx from onnx import TensorProto, helper -from op_test_utils import check_model_correctness, generate_random_initializer, input_feeds_neg_one_zero_one - +from op_test_utils import ( + StridedDataReader, + check_model_correctness, + generate_random_initializer, + input_feeds_neg_one_zero_one, + input_feeds_neg_one_zero_one_list, +) + +import onnxruntime as ort from onnxruntime.quantization import QuantType, StaticQuantConfig, quantize, quantize_static @@ -89,6 +96,51 @@ def test_save_as_external(self): check_model_correctness(self, self._model_fp32_path, quant_model_path, data_reader.get_next()) data_reader.rewind() + def run_inference(self, model_path, input_data): + session = ort.InferenceSession(model_path) + input_name = session.get_inputs()[0].name + output_name = session.get_outputs()[0].name + result = session.run([output_name], {input_name: input_data}) + return result + + def test_stride_effect_on_data_collection(self): + # Define the stride and test quantize_static with different stride values + stride = 5 + input_shapes = [1, self._channel_size, 1, 3] + data_list = input_feeds_neg_one_zero_one_list(10, {"input": [1, self._channel_size, 1, 3]}, 123) + input_nodes = ["input"] + in_dtypes = [np.float32] # Example dtype, adjust as needed + + # strided calibration + quant_model_path_1 = str(Path(self._tmp_model_dir.name) / "quant.strided.onnx") + data_reader_1 = StridedDataReader( + data_list, input_nodes, input_shapes, no_tensor_num=0, in_dtypes=in_dtypes, stride=stride + ) + quant_config_1 = StaticQuantConfig(data_reader_1, extra_options={"CalibStridedMinMax": stride}) + quantize(str(self._model_fp32_path), str(quant_model_path_1), quant_config_1) + + # non-strided calibration + quant_model_path_2 = str(Path(self._tmp_model_dir.name) / "quant.non.strided.onnx") + data_reader_2 = input_feeds_neg_one_zero_one(10, {"input": [1, self._channel_size, 1, 3]}, 123) + quant_config_2 = StaticQuantConfig(data_reader_2) + quantize(str(self._model_fp32_path), str(quant_model_path_2), quant_config_2) + + # Inference with both models and assert output closeness + np.random.seed(123) + input_data = np.random.choice([-1, 0, 1], size=[1, self._channel_size, 1, 3]).astype(np.float32) + + result_1 = self.run_inference(quant_model_path_1, input_data) + result_2 = self.run_inference(quant_model_path_2, input_data) + + # Assert that the outputs are close + np.testing.assert_allclose( + result_1, + result_2, + rtol=0.01, + atol=0.01, + err_msg="Outputs from strided and non-strided models are not close enough.", + ) + def test_static_quant_config(self): data_reader = input_feeds_neg_one_zero_one(10, {"input": [1, self._channel_size, 1, 3]}) quant_config = StaticQuantConfig(data_reader) From 6236707c640b84857f9e64c1c43fb9f304ca7749 Mon Sep 17 00:00:00 2001 From: Caroline Zhu Date: Fri, 21 Jun 2024 09:55:26 -0700 Subject: [PATCH 18/52] Enable >2GB models + allow model paths to be passed for generate_artifacts API (#20958) ### Description Alternative design from #20942 Allow users to pass in a model path for the generate_artifacts API. ### Motivation and Context - ONNX API calls such as the onnx checker + shape inference fail when given a model > 2GB, but work if a path to a model >2GB is passed in. --- .../orttraining/python/training/artifacts.py | 30 +++++-- .../python/training/onnxblock/blocks.py | 58 +++++++++++- .../python/training/onnxblock/loss/loss.py | 15 ++-- .../training/onnxblock/model_accessor.py | 27 +++++- .../python/training/onnxblock/onnxblock.py | 4 +- .../orttraining_test_ort_apis_onnxblock.py | 88 +++++++++++++++++++ 6 files changed, 202 insertions(+), 20 deletions(-) diff --git a/orttraining/orttraining/python/training/artifacts.py b/orttraining/orttraining/python/training/artifacts.py index 624b30ffdab3b..c98e5bcd97092 100644 --- a/orttraining/orttraining/python/training/artifacts.py +++ b/orttraining/orttraining/python/training/artifacts.py @@ -13,6 +13,9 @@ from onnxruntime.tools.convert_onnx_models_to_ort import OptimizationStyle, convert_onnx_models_to_ort from onnxruntime.training import onnxblock +# threshold for the size of the modelproto where you should use a path instead +USE_PATH_THRESHOLD = 2147483648 + class LossType(Enum): """Loss type to be added to the training model. @@ -37,7 +40,7 @@ class OptimType(Enum): def generate_artifacts( - model: onnx.ModelProto, + model: Union[onnx.ModelProto, str], requires_grad: Optional[List[str]] = None, frozen_params: Optional[List[str]] = None, loss: Optional[Union[LossType, onnxblock.Block]] = None, @@ -61,7 +64,8 @@ def generate_artifacts( All generated ModelProtos will use the same opsets defined by *model*. Args: - model: The base model to be used for gradient graph generation. + model: The base model or path to the base model to be used for gradient graph generation. For models >2GB, + use the path to the base model. requires_grad: List of names of model parameters that require gradient computation frozen_params: List of names of model parameters that should be frozen. loss: The loss function enum or onnxblock to be used for training. If None, no loss node is added to the graph. @@ -86,6 +90,22 @@ def generate_artifacts( RuntimeError: If the optimizer provided is not one of the supported optimizers. """ + loaded_model = None + model_path = None + + if isinstance(model, str): + loaded_model = onnx.load(model) + model_path = model + elif isinstance(model, onnx.ModelProto): + if model.ByteSize() > USE_PATH_THRESHOLD: + # infer_shapes and check_model from ONNX both require paths to be used for >2GB models. + raise RuntimeError("This model is > 2GB. Please pass in a path to the ONNX file instead.") + + loaded_model = model + model_path = None + else: + raise RuntimeError("Please pass in either a string or an ONNX ModelProto for the model.") + loss_blocks = { LossType.MSELoss: onnxblock.loss.MSELoss, LossType.CrossEntropyLoss: onnxblock.loss.CrossEntropyLoss, @@ -165,12 +185,12 @@ def build(self, *inputs_to_loss): logging.info("Custom op library provided: %s", custom_op_library) custom_op_library_path = pathlib.Path(custom_op_library) - with onnxblock.base(model), ( + with onnxblock.base(loaded_model, model_path), ( onnxblock.custom_op_library(custom_op_library_path) if custom_op_library is not None else contextlib.nullcontext() ): - _ = training_block(*[output.name for output in model.graph.output]) + _ = training_block(*[output.name for output in loaded_model.graph.output]) training_model, eval_model = training_block.to_model_proto() model_params = training_block.parameters() @@ -220,7 +240,7 @@ def _export_to_ort_format(model_path, output_dir, ort_format, custom_op_library_ return opset_version = None - for domain in model.opset_import: + for domain in loaded_model.opset_import: if domain.domain == "" or domain.domain == "ai.onnx": opset_version = domain.version break diff --git a/orttraining/orttraining/python/training/onnxblock/blocks.py b/orttraining/orttraining/python/training/onnxblock/blocks.py index 149d0a360f7d3..80f07c3738a7e 100644 --- a/orttraining/orttraining/python/training/onnxblock/blocks.py +++ b/orttraining/orttraining/python/training/onnxblock/blocks.py @@ -4,6 +4,7 @@ import contextlib import copy import logging +import os from abc import ABC, abstractmethod from typing import Any, List, Optional @@ -28,8 +29,13 @@ class Block(ABC): base (onnx.ModelProto): The base model that the subclass can manipulate. """ - def __init__(self): + def __init__(self, temp_file_name="temp.onnx"): + if os.path.isabs(temp_file_name): + raise RuntimeError("Please pass in a relative path for the temp_file_name.") self.base = None + self.temp_onnx_file_path = os.path.join(os.getcwd(), temp_file_name) + # onnx.save location parameter requires a relative path to the model path + self.temp_external_data_file_name = temp_file_name + ".data" @abstractmethod def build(self, *args, **kwargs): @@ -47,10 +53,58 @@ def __call__(self, *args, **kwargs): output = self.build(*args, **kwargs) - onnx.checker.check_model(self.base, True) + if accessor._GLOBAL_ACCESSOR.has_path: + onnx.save( + accessor._GLOBAL_ACCESSOR.model, + self.temp_onnx_file_path, + save_as_external_data=True, + all_tensors_to_one_file=True, + location=self.temp_external_data_file_name, + ) + + onnx.checker.check_model(self.temp_onnx_file_path, True) + else: + onnx.checker.check_model(self.base, True) return output + def infer_shapes_on_base(self): + """ + Performs shape inference on the global model. If a path was used, then uses the + infer_shapes_path API to support models with external data. + + Returns the shape-inferenced ModelProto. + """ + if accessor._GLOBAL_ACCESSOR.has_path: + onnx.save( + accessor._GLOBAL_ACCESSOR.model, + self.temp_onnx_file_path, + save_as_external_data=True, + all_tensors_to_one_file=True, + location=self.temp_external_data_file_name, + ) + + onnx.shape_inference.infer_shapes_path(self.temp_onnx_file_path) + # shape inferenced model is saved to original path + model = onnx.load(self.temp_onnx_file_path) + + return model + else: + return onnx.shape_inference.infer_shapes(accessor._GLOBAL_ACCESSOR.model) + + def __del__(self): + # since the ModelProto does not store the external data parameters themselves, just the metadata + # for where the external data can be found, we retain the external data files for the intermediate + # calls until the Block no longer needs to be used. + if os.path.exists(self.temp_onnx_file_path): + os.remove(self.temp_onnx_file_path) + # get absolute path for the external data file + external_data_file_path = os.path.join( + os.path.dirname(self.temp_onnx_file_path), self.temp_external_data_file_name + ) + if os.path.exists(external_data_file_path): + os.remove(external_data_file_path) + class _BinaryOp(Block): def __init__(self, op_name): diff --git a/orttraining/orttraining/python/training/onnxblock/loss/loss.py b/orttraining/orttraining/python/training/onnxblock/loss/loss.py index e719301e13f48..09429dd844187 100644 --- a/orttraining/orttraining/python/training/onnxblock/loss/loss.py +++ b/orttraining/orttraining/python/training/onnxblock/loss/loss.py @@ -93,19 +93,20 @@ def build(self, scores_input_name: str, labels_name: str = "labels"): labels_input = copy.deepcopy(_graph_utils.get_output_from_output_name(self.base, scores_input_name)) labels_input.name = labels_name labels_input.type.tensor_type.elem_type = onnx.TensorProto.INT64 - # If the predictions are (num_examples x num_classes) - # labels should be (num_examples,) - del labels_input.type.tensor_type.shape.dim[1] + # Assumes classes is the last dimension + # e.g., predictions: (num_examples, num_classes) -> labels: (num_examples,) + # or predictions: (batch_size, seq_len, vocab) -> labels: (batch_size, seq_len) + del labels_input.type.tensor_type.shape.dim[-1] self.base.graph.input.append(labels_input) loss_node_input_names = [scores_input_name, labels_name] if self._weight: loss_node_input_names.append(weight_name) + loss_node_output_name = _graph_utils.generate_graph_name("loss") - loss_node_output_names = [ - loss_node_output_name, - _graph_utils.generate_graph_name("log_prob"), - ] + log_prob_output_name = _graph_utils.generate_graph_name("log_prob") + + loss_node_output_names = [loss_node_output_name, log_prob_output_name] loss_node = onnx.helper.make_node( "SoftmaxCrossEntropyLoss", loss_node_input_names, diff --git a/orttraining/orttraining/python/training/onnxblock/model_accessor.py b/orttraining/orttraining/python/training/onnxblock/model_accessor.py index ac7a53a554e0a..302573064be6e 100644 --- a/orttraining/orttraining/python/training/onnxblock/model_accessor.py +++ b/orttraining/orttraining/python/training/onnxblock/model_accessor.py @@ -15,10 +15,12 @@ class ModelAccessor: Attributes: model: The onnx model that is manipulated by the onnx blocks. + model_path: The path to the base model. Can be None. """ - def __init__(self, model: onnx.ModelProto): + def __init__(self, model: onnx.ModelProto, model_path: str | None = None): self._model = model + self._path = model_path @property def model(self) -> onnx.ModelProto: @@ -30,6 +32,22 @@ def model(self) -> onnx.ModelProto: ) return self._model + @property + def path(self) -> str: + """ModelAccessor property that gets the path to the base model.""" + + if self._path is None: + raise RuntimeError( + "The path to the onnx model was not set. Please use the context manager onnxblock.onnx_model to create the model and pass in a string." + ) + return self._path + + @property + def has_path(self) -> bool: + """Returns True if ModelAccessor has a path to a model, False otherwise.""" + + return self._path is not None + # These variable resides in the global namespace. # Different methods can access this global model and manipulate it. @@ -39,7 +57,7 @@ def model(self) -> onnx.ModelProto: @contextmanager -def base(model: onnx.ModelProto): +def base(model: onnx.ModelProto, model_path: str | None = None): """Registers the base model to be manipulated by the onnx blocks. Example: @@ -53,6 +71,7 @@ def base(model: onnx.ModelProto): Args: model: The base model to be manipulated by the onnx blocks. + model_path: The path to the base model. None if there is no model path to pass in. Returns: ModelAccessor: The model accessor that contains the modified model. @@ -69,7 +88,7 @@ def base(model: onnx.ModelProto): "model from scratch." ) - _GLOBAL_ACCESSOR = ModelAccessor(model_clone) + _GLOBAL_ACCESSOR = ModelAccessor(model_clone, model_path) try: yield _GLOBAL_ACCESSOR finally: @@ -112,7 +131,7 @@ def empty_base(opset_version: int | None = None): ) ) - _GLOBAL_ACCESSOR = ModelAccessor(model) + _GLOBAL_ACCESSOR = ModelAccessor(model, None) try: yield _GLOBAL_ACCESSOR finally: diff --git a/orttraining/orttraining/python/training/onnxblock/onnxblock.py b/orttraining/orttraining/python/training/onnxblock/onnxblock.py index a2922353ac70e..64f7acf4dc02c 100644 --- a/orttraining/orttraining/python/training/onnxblock/onnxblock.py +++ b/orttraining/orttraining/python/training/onnxblock/onnxblock.py @@ -70,7 +70,7 @@ def __call__(self, *args, **kwargs): output = self.build(*args, **kwargs) - self._model = onnx.shape_inference.infer_shapes(accessor._GLOBAL_ACCESSOR.model) + self._model = self.infer_shapes_on_base() _graph_utils.register_graph_outputs(self._model, output) @@ -187,7 +187,7 @@ def __call__(self, *args, **kwargs): output = self.build(*args, **kwargs) - model = onnx.shape_inference.infer_shapes(accessor._GLOBAL_ACCESSOR.model) + model = self.infer_shapes_on_base() _graph_utils.register_graph_outputs(model, output) diff --git a/orttraining/orttraining/test/python/orttraining_test_ort_apis_onnxblock.py b/orttraining/orttraining/test/python/orttraining_test_ort_apis_onnxblock.py index ac49c1c2834c7..5c63be92d2b2f 100644 --- a/orttraining/orttraining/test/python/orttraining_test_ort_apis_onnxblock.py +++ b/orttraining/orttraining/test/python/orttraining_test_ort_apis_onnxblock.py @@ -1099,3 +1099,91 @@ def test_custom_optimizer_block(): for attr in node.attribute: if attr.name == "weight_decay": assert attr.f == weight_decay + + +def test_generate_artifacts_path(): + + with tempfile.TemporaryDirectory() as temp_dir: + _, simple_net = _get_models("cpu", 32, 28, 10, 10) + + requires_grad_params = ["fc1.weight", "fc1.bias", "fc2.weight", "fc2.bias"] + + onnx.save_model( + simple_net, + os.path.join(temp_dir, "simple_net.onnx"), + ) + + artifacts.generate_artifacts( + os.path.join(temp_dir, "simple_net.onnx"), + requires_grad=requires_grad_params, + loss=artifacts.LossType.CrossEntropyLoss, + optimizer=artifacts.OptimType.AdamW, + artifact_directory=temp_dir, + ) + + # generate_artifacts should have thrown if it didn't complete successfully. + # Below is a sanity check to validate that all the expected files were created. + assert os.path.exists(os.path.join(temp_dir, "training_model.onnx")) + assert os.path.exists(os.path.join(temp_dir, "eval_model.onnx")) + assert os.path.exists(os.path.join(temp_dir, "optimizer_model.onnx")) + assert os.path.exists(os.path.join(temp_dir, "checkpoint")) + + +def test_generate_artifacts_external_data_one_file(): + with tempfile.TemporaryDirectory() as temp_dir: + _, simple_net = _get_models("cpu", 32, 28, 10, 10) + + requires_grad_params = ["fc1.weight", "fc1.bias", "fc2.weight", "fc2.bias"] + + onnx.save_model( + simple_net, + os.path.join(temp_dir, "simple_net.onnx"), + save_as_external_data=True, + all_tensors_to_one_file=True, + size_threshold=0, + ) + + artifacts.generate_artifacts( + os.path.join(temp_dir, "simple_net.onnx"), + requires_grad=requires_grad_params, + loss=artifacts.LossType.CrossEntropyLoss, + optimizer=artifacts.OptimType.AdamW, + artifact_directory=temp_dir, + ) + + # generate_artifacts should have thrown if it didn't complete successfully. + # Below is a sanity check to validate that all the expected files were created. + assert os.path.exists(os.path.join(temp_dir, "training_model.onnx")) + assert os.path.exists(os.path.join(temp_dir, "eval_model.onnx")) + assert os.path.exists(os.path.join(temp_dir, "optimizer_model.onnx")) + assert os.path.exists(os.path.join(temp_dir, "checkpoint")) + + +def test_generate_artifacts_external_data_separate_files(): + with tempfile.TemporaryDirectory() as temp_dir: + _, simple_net = _get_models("cpu", 32, 28, 10, 10) + + requires_grad_params = ["fc1.weight", "fc1.bias", "fc2.weight", "fc2.bias"] + + onnx.save_model( + simple_net, + os.path.join(temp_dir, "simple_net.onnx"), + save_as_external_data=True, + all_tensors_to_one_file=False, + size_threshold=0, + ) + + artifacts.generate_artifacts( + os.path.join(temp_dir, "simple_net.onnx"), + requires_grad=requires_grad_params, + loss=artifacts.LossType.CrossEntropyLoss, + optimizer=artifacts.OptimType.AdamW, + artifact_directory=temp_dir, + ) + + # generate_artifacts should have thrown if it didn't complete successfully. + # Below is a sanity check to validate that all the expected files were created. + assert os.path.exists(os.path.join(temp_dir, "training_model.onnx")) + assert os.path.exists(os.path.join(temp_dir, "eval_model.onnx")) + assert os.path.exists(os.path.join(temp_dir, "optimizer_model.onnx")) + assert os.path.exists(os.path.join(temp_dir, "checkpoint")) From ac216267254232508ca19586937a8420d75ce64d Mon Sep 17 00:00:00 2001 From: Dwayne Robinson Date: Fri, 21 Jun 2024 11:46:16 -0700 Subject: [PATCH 19/52] DML EP EinSum make more generic to avoid EP fallback (#21114) ### Problem Newer models using more novel equations (e.g. `bhwc,hkc->bhwk` in Segment Anything's encoder or `bqc,bchw->bqhw`) cause fallback from DML to CPU, yielding performance issues. The EP had some pattern matching to map more common equations to existing DML operators, but the number of permutations was prohibitive and could not catch them all. ### Solution So, ditch the static mapping, and instead handle any 1-input or 2-input cases via remapped strides and a mini-graph of elementwise multiplication & sum reduction (as if DML had a `DML_OPERATOR_DOT_PRODUCT` that took `axes`). A subset of mappings still exist for performance (GEMM, pure reduction, transpose...), but they are identified generally rather than via a pattern table. Also... - Diagonals are supported now (e.g. iji->i). - Removes any remaining DML-specific EinSum `GTEST_SKIP` statements. - Handles any cases up to 8 unique labels (DML dimension limit is 8D). - \>= 3 inputs and arbitrary size inputs via ellipsis are not handled, but we have yet to come across a model. --- .../DmlExecutionProvider/src/DmlCommon.cpp | 30 +- .../dml/DmlExecutionProvider/src/DmlCommon.h | 6 +- .../src/Operators/DmlOperatorEinSum.cpp | 497 +++++++++++------- .../DmlExecutionProvider/src/TensorDesc.cpp | 55 +- .../dml/DmlExecutionProvider/src/TensorDesc.h | 24 +- .../OperatorAuthorHelper/OperatorHelper.cpp | 197 +++---- .../dml/OperatorAuthorHelper/OperatorHelper.h | 32 +- .../test/providers/cpu/math/einsum_test.cc | 41 +- 8 files changed, 532 insertions(+), 350 deletions(-) diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/DmlCommon.cpp b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/DmlCommon.cpp index ea66289c351ea..541254ffaf7f0 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/DmlCommon.cpp +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/DmlCommon.cpp @@ -32,7 +32,7 @@ DML_TENSOR_DATA_TYPE GetDmlDataTypeFromMlDataTypeNoThrow(MLOperatorTensorDataTyp }; } -bool IsSigned(DML_TENSOR_DATA_TYPE dataType) +bool IsSigned(DML_TENSOR_DATA_TYPE dataType) noexcept { switch (dataType) { @@ -140,7 +140,33 @@ uint32_t GetSupportedDeviceDataTypeMask(IDMLDevice* dmlDevice) return deviceTypeMask; } -void GetDescendingPackedStrides(gsl::span sizes, /*out*/ gsl::span strides) +uint32_t GetBitMaskFromIndices(gsl::span indices) noexcept +{ + uint32_t bitMask = 0; + for (auto i : indices) + { + assert(i < 32); + bitMask |= (1 << i); + } + return bitMask; +} + +uint32_t CountLeastSignificantZeros(uint32_t value) noexcept +{ + // *Use std::countr_zero instead when codebase updated to C++20. + // Use bit twiddling hack rather than for loop. + uint32_t count = 32; + value &= -int32_t(value); + if (value) count--; + if (value & 0x0000FFFF) count -= 16; + if (value & 0x00FF00FF) count -= 8; + if (value & 0x0F0F0F0F) count -= 4; + if (value & 0x33333333) count -= 2; + if (value & 0x55555555) count -= 1; + return count; +} + +void GetDescendingPackedStrides(gsl::span sizes, /*out*/ gsl::span strides) noexcept { assert(sizes.size() == strides.size()); diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/DmlCommon.h b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/DmlCommon.h index c4d260b9736df..5cd3fd0aea72c 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/DmlCommon.h +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/DmlCommon.h @@ -23,9 +23,11 @@ namespace Dml size_t ComputeByteSizeFromDimensions(gsl::span dimensions, MLOperatorTensorDataType tensorDataType); size_t ComputeByteSizeFromTensor(IMLOperatorTensor& tensor); uint32_t GetSupportedDeviceDataTypeMask(IDMLDevice* dmlDevice); - void GetDescendingPackedStrides(gsl::span sizes, /*out*/ gsl::span strides); + uint32_t GetBitMaskFromIndices(gsl::span indices) noexcept; + uint32_t CountLeastSignificantZeros(uint32_t value) noexcept; + void GetDescendingPackedStrides(gsl::span sizes, /*out*/ gsl::span strides) noexcept; - bool IsSigned(DML_TENSOR_DATA_TYPE dataType); + bool IsSigned(DML_TENSOR_DATA_TYPE dataType) noexcept; template void CastToClampedScalarUnion(DML_TENSOR_DATA_TYPE dataType, T value, DML_SCALAR_UNION* outputValue) diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperatorEinSum.cpp b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperatorEinSum.cpp index d5bf54de53c30..51b19603f5122 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperatorEinSum.cpp +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperatorEinSum.cpp @@ -3,6 +3,70 @@ #include "precomp.h" +// With a single equation, the Einstein summation operator can represent a variety of operators including: matmul, +// summation, transposition, diagonal slice, diagonal sum (trace), inner (dot) product, outer product... +// +// Parameters NumPy equivalent Description +// ------------------------------------------------------------------------------------------------------------- +// ('i', A1) A1 returns a view of A1 +// ('i->', A1) sum(A1) sums the values of A1 +// ('i,i->i', A1, B1) A1 * B1 element-wise multiplication of A1 and B1 +// ('i,i->', A1, B1) inner(A1, B1) or dot(A1, B1) inner product of A1 and B1 +// ('i,i', A1, B1) inner(A1, B1) or dot(A1, B1) inner product of A1 and B1 +// ('i,j->ij', A1, B1) outer(A1, B1) outer product of A1 and B1 +// ('ij->ij', A2) A2 returns a view of A2 +// ('ij', A2) A2 returns a view of A2 +// ('ji', A2) A2.T view transpose of A2 +// ('ji->ij', A2) A2.T view transpose of A2 +// ('ii->i', A2) diag(A2) view main diagonal of A2 +// ('ii->', A2) trace(A2) sums main diagonal of A2 +// ('ij->', A2) sum(A2) sums the values of A2 +// ('ij->j', A2) sum(A2, axis=0) sum down the columns of A2 (across rows) +// ('ij->i', A2) sum(A2, axis=1) sum horizontally along the rows of A2 +// ('ij,ij->ij', A2, B2) A2 * B2 element-wise multiplication of A2 and B2 +// ('ij,ji->ij', A2, B2) A2 * B2.transpose() element-wise multiplication of A2 and B2.T +// ('ij,jk', A2, B2) matmul(A2, B2) or dot(A2, B2) matrix multiplication of A2 and B2 +// ('ij,jk->ik', A2, B2) matmul(A2, B2) or dot(A2, B2) matrix multiplication of A2 and B2 +// ('bij,bjk->bik', A2, B2) matmul(A3, B3) matrix multiplication of A3 and B3 (a stack of 2D matrices) +// ('bij,bkj->bik', A2, B2) matmul(A3, transpose(B3)) matrix multiplication of A3 and B3 (a stack of 2D matrices) +// ('ij,kj->ik', A2, B2) inner(A2, B2) inner product of A2 and B2 +// ('ij,kj->ikj', A2, B2) A2[:, None] * B2 each row of A2 multiplied by B2 +// ('ij,kl->ijkl', A2, B2) A2[:, :, None, None] * B2 each value of A2 multiplied by B2 +// (',ij', 3, B2) Scalar times array: array([[ 0, 3, 6], [ 9, 12, 15]]) +// ("ij,j", A2, B1) matvec(A2, B1) Matrix and vector. +// ("ii,ii->i", A2, B2) A2.diag() * B2.diag() diagonals multiplied by each other +// ("ii,ii->", A2, B2) dot(A2.diag(), B2.diag()) dot product of diagonals +// +// Decomposition: +// +// Ultimately though EinSum is equivalent to an elementwise multiplication into an internal product tensor +// (given a helper function to reproject all inputs so they're shape-compatible) followed by sum reduction. +// +// 1. Determine the size of the internal product tensor by concatenating the dimensions of all inputs, +// counting each unique label once. So "bij,bjk->bik" would yield an internal product of shape [b,i,j,k]. +// 2. Project each input tensor as needed to the internal product shape (transposing and/or broadcasting). +// So an input of shape [b,i] with product shape of [b,j,i,k] would insert broadcasted j and k dimensions. +// An input of shape [a,b,c] with product shape of [b,c,a] would require a transpose. +// The input shape [a,b,a] with product shape of [a,b] would collapse the first two input 'a' dimensions. +// 3. Multiply elementwise every input tensor to compute the internal product. +// 4. Sum reduce the product tensor to the final output shape, reducing along any missing dimensions. +// So a product shape of [b,j,i,k] and output shape of [b,i,k] reduces along j. +// +// ReduceSum( +// Mul( +// ExpandTransposeCollapseAsNeeded(A, aAxesToProductAxes), +// ExpandTransposeCollapseAsNeeded(B, bAxesToProductAxes), +// ), +// reductionAxes, +// keepdims=false +// ) +// +// Notes: +// +// - DirectML has no direct EinSum operator, but common cases map to existing operators. +// - EinSum can accept a variable number of input tensors, but the DML EP only supports a limited count +// (falling back to CPU otherwise). + namespace Dml { @@ -13,30 +77,36 @@ class DmlOperatorEinSum : public DmlOperator, public EinSumHelper : DmlOperator(kernelCreationContext), EinSumHelper(kernelCreationContext, kernelCreationContext.GetTensorShapeDescription(), opsetVersion) { - ML_CHECK_VALID_ARGUMENT(static_cast(kernelCreationContext.GetInputCount()) + 1 == m_components.size(), - "EinSum input tensor count is inconsistent with the equation component count."); + ML_CHECK_VALID_ARGUMENT(kernelCreationContext.GetInputCount() >= 1, "EinSum expects at least one input tensor."); ML_CHECK_VALID_ARGUMENT(kernelCreationContext.GetOutputCount() == 1, "EinSum expects one output tensor."); + ML_CHECK_VALID_ARGUMENT( + static_cast(kernelCreationContext.GetInputCount()) + 1 == m_components.size(), + "EinSum input tensor count is inconsistent with the equation component count." + ); + assert(m_recognizedOperatorType != RecognizedOperatorType::None && "Unrecognized EinSum operators should have fallen back to CPU"); std::vector> inputIndices = {0,1,2}; std::vector> outputIndices = {0}; uint32_t bindableInputCount = kernelCreationContext.GetInputCount(); if (IsMatMulOperatorType()) { - ++bindableInputCount; // Account for the optional C tensor. + ++bindableInputCount; // Account for the optional C tensor. } inputIndices.resize(bindableInputCount); - constexpr uint32_t dimCount = 2; - DmlOperator::Initialize(kernelCreationContext, inputIndices, outputIndices, std::nullopt, std::nullopt, dimCount); + uint32_t minimumDimensionCount = 1; + DmlOperator::Initialize(kernelCreationContext, inputIndices, outputIndices, std::nullopt, std::nullopt, minimumDimensionCount); std::vector inputDescs = GetDmlInputDescs(); std::vector outputDescs = GetDmlOutputDescs(); - static_assert(RecognizedOperatorType::Total == static_cast(12), "Update this switch."); + static_assert(RecognizedOperatorType::Total == static_cast(6), "Update this switch statement."); switch (m_recognizedOperatorType) { case RecognizedOperatorType::Multiply: { + ReprojectTensorDescsToProductTensor(); + DML_ELEMENT_WISE_MULTIPLY_OPERATOR_DESC operatorDesc = {}; operatorDesc.ATensor = &inputDescs[0]; operatorDesc.BTensor = &inputDescs[1]; @@ -46,115 +116,53 @@ class DmlOperatorEinSum : public DmlOperator, public EinSumHelper } break; - case RecognizedOperatorType::OuterProduct: - { - std::array aSizes = {m_inputTensorDescs[0].GetSizes().back(), 1}; - TensorDesc aTensorDesc = TensorDesc(m_inputTensorDescs[0].GetDmlDataType(), aSizes); - auto aDmlTensorDesc = aTensorDesc.GetDmlDesc(); - - std::array bSizes = {1, m_inputTensorDescs[1].GetSizes().back()}; - TensorDesc bTensorDesc = TensorDesc(m_inputTensorDescs[1].GetDmlDataType(), bSizes); - auto bDmlTensorDesc = bTensorDesc.GetDmlDesc(); - - DML_GEMM_OPERATOR_DESC operatorDesc = {}; - operatorDesc.ATensor = &aDmlTensorDesc; - operatorDesc.BTensor = &bDmlTensorDesc; - operatorDesc.OutputTensor = &outputDescs[0]; - operatorDesc.Alpha = 1.0; - operatorDesc.Beta = 0.0; - operatorDesc.FusedActivation = nullptr; - - SetDmlOperatorDesc({ DML_OPERATOR_GEMM, &operatorDesc }, kernelCreationContext); - } - break; - case RecognizedOperatorType::MatMul: - case RecognizedOperatorType::MatMulTransposeA: - case RecognizedOperatorType::MatMulTransposeB: { - DML_GEMM_OPERATOR_DESC operatorDesc = {}; - operatorDesc.ATensor = &inputDescs[0]; - operatorDesc.BTensor = &inputDescs[1]; - // No operatorDesc.CTensor - operatorDesc.OutputTensor = &outputDescs[0]; - operatorDesc.TransA = (m_recognizedOperatorType == RecognizedOperatorType::MatMulTransposeA) ? DML_MATRIX_TRANSFORM_TRANSPOSE : DML_MATRIX_TRANSFORM_NONE; - operatorDesc.TransB = (m_recognizedOperatorType == RecognizedOperatorType::MatMulTransposeB) ? DML_MATRIX_TRANSFORM_TRANSPOSE : DML_MATRIX_TRANSFORM_NONE; - operatorDesc.Alpha = 1.0; - operatorDesc.Beta = 0.0; - operatorDesc.FusedActivation = nullptr; - - SetDmlOperatorDesc({ DML_OPERATOR_GEMM, &operatorDesc }, kernelCreationContext); - } - break; - case RecognizedOperatorType::MatMulNhcw: - case RecognizedOperatorType::MatMulNhcwTransposeA: - case RecognizedOperatorType::MatMulNhcwTransposeB: - { - // Transpose via input strides. The output tensor is not strided. Support only 4D for now. - assert(m_components.size() == 3); - assert(m_components[0].GetDimensionCount() == m_components[2].GetDimensionCount()); - assert(m_components[1].GetDimensionCount() == m_components[2].GetDimensionCount()); - assert(m_components[2].GetDimensionCount() == 4); - - // Remap transposed strides from NCHW to NHCW - constexpr std::array labelIndices = {0, 2, 1, 3}; - - assert(m_inputTensorDescs.size() >= 2); - for (uint32_t inputIndex = 0; inputIndex < 2; ++inputIndex) - { - TensorDesc& tensorDesc = m_inputTensorDescs[inputIndex]; - auto originalStrides = tensorDesc.GetStrides(); - std::vector inputSizes = kernelCreationContext.GetTensorShapeDescription().GetInputTensorShape(inputIndex); - std::vector inputStrides(inputSizes.size()); - - // If there were no strides, compute them based in descending packed order - // based on the input sizes. - if (originalStrides.empty()) - { - Dml::GetDescendingPackedStrides(inputSizes, /*out*/ inputStrides); - } - else // Copy the original strides. - { - assert(originalStrides.size() >= inputStrides.size()); - size_t offset = originalStrides.size() - inputStrides.size(); - inputStrides.assign(originalStrides.begin() + offset, originalStrides.end()); - } - - std::vector newStrides(inputStrides.size()); - std::vector newSizes(inputStrides.size()); - for (size_t dim = 0, dimensionCount = inputStrides.size(); dim < dimensionCount; ++dim) - { - uint32_t labelIndex = labelIndices[dim]; - assert(labelIndex < inputStrides.size()); - newSizes[dim] = inputSizes[labelIndex]; - newStrides[dim] = inputStrides[labelIndex]; - } - - // Override the initial input tensor with the new strides. - tensorDesc = TensorDesc(tensorDesc.GetDmlDataType(), newSizes, newStrides, 0); - tensorDesc.GetDmlDesc(); // Discard value, but keep side effect of refreshing the DML view. - } - - std::vector outputSizes = kernelCreationContext.GetTensorShapeDescription().GetOutputTensorShape(0); - std::vector newOutputSizes(outputSizes.size()); - assert(outputSizes.size() == labelIndices.size()); - - for (size_t dim = 0; dim < outputSizes.size(); ++dim) + assert(m_components.size() == 3 && "EinSum matmul expects 2 inputs and 1 output"); + assert(m_productDimensions.size() - 1 <= 4 && "DML Einsum matmul handles up to 4D"); + + // Generate bitmasks for each of the active axes per tensor using their labels. + const auto input0Labels = m_components[0].GetLabels(m_labelIndices); + const auto input1Labels = m_components[1].GetLabels(m_labelIndices); + const auto outputLabels = m_components[2].GetLabels(m_labelIndices); + const uint32_t input0AxesMask = GetBitMaskFromIndices(input0Labels); + const uint32_t input1AxesMask = GetBitMaskFromIndices(input1Labels); + const uint32_t outputAxesMask = GetBitMaskFromIndices(outputLabels); + + // Find each of the interesting axes, including the one being reduced, height, width, batch, and channel. + // - the reduced axis is the term missing from the output. + // - height and width are the unique axes respectively found in only input A or input B. + // - the batch (if present) is the first axis shared by both inputs, and the channel is the subsequent common one. + // If any axis is not found (say it's a 2D GEMM), then the axis value will be beyond the rank, which is + // safely handled correctly during projection as an inserted axis. + + auto findAndClearAxis = [](uint32_t& currentAxesMask, uint32_t contraintAxesMask) -> uint32_t { - uint32_t labelIndex = labelIndices[dim]; - newOutputSizes[dim] = outputSizes[labelIndex]; - } - - m_outputTensorDescs.front() = TensorDesc(m_outputTensorDescs.front().GetDmlDataType(), newOutputSizes, std::nullopt, 0); - m_outputTensorDescs.front().GetDmlDesc(); // Discard value, but keep side effect of refreshing the DML view. + uint32_t foundAxis = CountLeastSignificantZeros(currentAxesMask & ~contraintAxesMask); + currentAxesMask &= ~(1 << foundAxis); + return foundAxis; + }; + + uint32_t remainingAxesMask = ~0u; + uint32_t reductionAxis = findAndClearAxis(/*inout*/ remainingAxesMask, outputAxesMask); + uint32_t heightAxis = findAndClearAxis(/*inout*/ remainingAxesMask, input1AxesMask); + uint32_t widthAxis = findAndClearAxis(/*inout*/ remainingAxesMask, input0AxesMask); + uint32_t batchAxis = findAndClearAxis(/*inout*/ remainingAxesMask, 0); + uint32_t channelAxis = findAndClearAxis(/*inout*/ remainingAxesMask, 0); + + // Reproject all inputs and the output to the needed order pattern for DML compatibility, + // which only accepts the rightmost axis as GEMM-reducible when TransB is true. + ReprojectTensorDescToGivenAxes(/*inout*/ m_inputTensorDescs[0], input0Labels, {{batchAxis, channelAxis, heightAxis, reductionAxis}}); + ReprojectTensorDescToGivenAxes(/*inout*/ m_inputTensorDescs[1], input1Labels, {{batchAxis, channelAxis, widthAxis, reductionAxis}}); + ReprojectTensorDescToGivenAxes(/*inout*/ m_outputTensorDescs[0], outputLabels, {{batchAxis, channelAxis, heightAxis, widthAxis}}); DML_GEMM_OPERATOR_DESC operatorDesc = {}; operatorDesc.ATensor = &inputDescs[0]; operatorDesc.BTensor = &inputDescs[1]; // No operatorDesc.CTensor operatorDesc.OutputTensor = &outputDescs[0]; - operatorDesc.TransA = (m_recognizedOperatorType == RecognizedOperatorType::MatMulNhcwTransposeA) ? DML_MATRIX_TRANSFORM_TRANSPOSE : DML_MATRIX_TRANSFORM_NONE; - operatorDesc.TransB = (m_recognizedOperatorType == RecognizedOperatorType::MatMulNhcwTransposeB) ? DML_MATRIX_TRANSFORM_TRANSPOSE : DML_MATRIX_TRANSFORM_NONE; + operatorDesc.TransA = DML_MATRIX_TRANSFORM_NONE; + operatorDesc.TransB = DML_MATRIX_TRANSFORM_TRANSPOSE; operatorDesc.Alpha = 1.0; operatorDesc.Beta = 0.0; operatorDesc.FusedActivation = nullptr; @@ -165,41 +173,10 @@ class DmlOperatorEinSum : public DmlOperator, public EinSumHelper case RecognizedOperatorType::ReduceSum: { - // Get how many axes are kept in the final output, either 0 or 1 supported - // meaning full reduction or partial with one dimension left. *It could be - // generalized to support any number of output dimensions, but it would need - // to accomodate for Transposition too if the output labels are reordered. - auto keptAxes = m_components.back().GetLabels(m_labelIndices); - assert(keptAxes.size() <= 1); - - // DML expects output rank to match input rank (as if ONNX ReduceSum keepdims=1). - // So replace the existing tensor description with the input sizes, except that - // reduced dimensions have size 1. - std::vector reducedAxes; - std::vector inputSizes = kernelCreationContext.GetTensorShapeDescription().GetInputTensorShape(0); - std::vector outputSizes = inputSizes; - - // Determine which axes are being reduced by taking the opposite of those kept. - uint32_t keptAxesMask = 0; - for (auto axis : keptAxes) - { - keptAxesMask |= (1 << axis); - } - for (uint32_t axis = 0, axisCount = static_cast(outputSizes.size()); axis < axisCount; ++axis) - { - if (~keptAxesMask & (1< reducedAxes = GetReductionAxes(); operatorDesc.InputTensor = inputDescs.data(); operatorDesc.OutputTensor = outputDescs.data(); operatorDesc.Function = DML_REDUCE_FUNCTION_SUM; @@ -211,48 +188,8 @@ class DmlOperatorEinSum : public DmlOperator, public EinSumHelper break; case RecognizedOperatorType::Transpose: - case RecognizedOperatorType::Identity: { - if (m_recognizedOperatorType == RecognizedOperatorType::Transpose) - { - // Transpose via input strides. The output tensor is not strided. - assert(m_components.front().GetDimensionCount() == m_components.back().GetDimensionCount()); - auto originalStrides = m_inputTensorDescs.front().GetStrides(); - std::vector inputSizes = kernelCreationContext.GetTensorShapeDescription().GetInputTensorShape(0); - std::vector inputStrides(inputSizes.size()); - - // If there were no strides, compute them based in descending packed order - // based on the input sizes. - if (originalStrides.empty()) - { - Dml::GetDescendingPackedStrides(inputSizes, /*out*/ inputStrides); - } - else // Copy the original strides. - { - assert(originalStrides.size() >= inputStrides.size()); - size_t offset = originalStrides.size() - inputStrides.size(); - inputStrides.assign(originalStrides.begin() + offset, originalStrides.end()); - } - - // Remap transposed strides using the component labels from input to output. - auto labelIndices = m_components.back().GetLabels(m_labelIndices); - - std::vector newStrides(inputStrides.size()); - std::vector newSizes(inputStrides.size()); - for (size_t i = 0, dimensionCount = inputStrides.size(); i < dimensionCount; ++i) - { - uint32_t labelIndex = labelIndices[i]; - assert(labelIndex < inputStrides.size()); - newSizes[i] = inputSizes[labelIndex]; - newStrides[i] = inputStrides[labelIndex]; - } - - // Override the initial input tensor with the new strides. - m_inputTensorDescs.front() = TensorDesc(m_inputTensorDescs.front().GetDmlDataType(), newSizes, newStrides, 0); - m_outputTensorDescs.front() = TensorDesc(m_outputTensorDescs.front().GetDmlDataType(), newSizes, std::nullopt, 0); - m_inputTensorDescs.front().GetDmlDesc(); // Discard value, but keep side effect of refreshing the DML view. - m_outputTensorDescs.front().GetDmlDesc(); // Discard value, but keep side effect of refreshing the DML view. - } + ReprojectTensorDescsToProductTensor(); DML_ELEMENT_WISE_IDENTITY_OPERATOR_DESC operatorDesc = {}; operatorDesc.InputTensor = inputDescs.data(); @@ -262,10 +199,208 @@ class DmlOperatorEinSum : public DmlOperator, public EinSumHelper } break; + case RecognizedOperatorType::MultiplyReduceSum: + { + // DML has no generic DML_OPERATOR_DOT_PRODUCT. So construct one via a graph of mul+sumReduce. + + ReprojectTensorDescsToProductTensor(); + TensorDesc productTensorDesc(m_outputTensorDescs.front().GetDmlDataType(), m_productDimensions); + auto dmlProductTensorDesc = productTensorDesc.GetDmlDesc(); + + DML_ELEMENT_WISE_MULTIPLY_OPERATOR_DESC multiplyOperatorDesc = {}; + multiplyOperatorDesc.ATensor = &inputDescs[0]; + multiplyOperatorDesc.BTensor = &inputDescs[1]; + multiplyOperatorDesc.OutputTensor = &dmlProductTensorDesc; + DML_OPERATOR_DESC multiplyOperatorDescWithEnum = { DML_OPERATOR_ELEMENT_WISE_MULTIPLY, &multiplyOperatorDesc }; + + DML_REDUCE_OPERATOR_DESC reduceSumOperatorDesc = {}; + std::vector reducedAxes = GetReductionAxes(); + reduceSumOperatorDesc.Function = DML_REDUCE_FUNCTION_SUM; + reduceSumOperatorDesc.InputTensor = &dmlProductTensorDesc; + reduceSumOperatorDesc.OutputTensor = &outputDescs[0]; + reduceSumOperatorDesc.Axes = reducedAxes.data(); + reduceSumOperatorDesc.AxisCount = gsl::narrow_cast(reducedAxes.size()); + DML_OPERATOR_DESC reduceSumOperatorDescWithEnum = { DML_OPERATOR_REDUCE, &reduceSumOperatorDesc }; + + enum NodeIndex + { + NodeIndexMultiply, + NodeIndexReduceSum, + NodeIndexTotal, + }; + + const DML_OPERATOR_DESC* operatorDescPointers[2] = + { + &multiplyOperatorDescWithEnum, // NodeIndexMultiply + &reduceSumOperatorDescWithEnum, // NodeIndexReduceSum + }; + + DML_INPUT_GRAPH_EDGE_DESC inputEdges[2]; + DML_INTERMEDIATE_GRAPH_EDGE_DESC intermediateEdges[1]; + DML_OUTPUT_GRAPH_EDGE_DESC outputEdges[1]; + + DML_INPUT_GRAPH_EDGE_DESC& input0ToMultiplyEdge = inputEdges[0]; + input0ToMultiplyEdge.GraphInputIndex = 0; + input0ToMultiplyEdge.ToNodeIndex = NodeIndexMultiply; + input0ToMultiplyEdge.ToNodeInputIndex = 0; + + DML_INPUT_GRAPH_EDGE_DESC& input1ToMultiplyEdge = inputEdges[1]; + input1ToMultiplyEdge.GraphInputIndex = 1; + input1ToMultiplyEdge.ToNodeIndex = NodeIndexMultiply; + input1ToMultiplyEdge.ToNodeInputIndex = 1; + + DML_INTERMEDIATE_GRAPH_EDGE_DESC& multiplyToReduceSumEdge = intermediateEdges[0]; + multiplyToReduceSumEdge.FromNodeIndex = NodeIndexMultiply; + multiplyToReduceSumEdge.FromNodeOutputIndex = 0; + multiplyToReduceSumEdge.ToNodeIndex = NodeIndexReduceSum; + multiplyToReduceSumEdge.ToNodeInputIndex = 0; + + DML_OUTPUT_GRAPH_EDGE_DESC& reduceSumToOutputEdge = outputEdges[0]; + reduceSumToOutputEdge.FromNodeIndex = NodeIndexReduceSum; + reduceSumToOutputEdge.FromNodeOutputIndex = 0; + reduceSumToOutputEdge.GraphOutputIndex = 0; + + MLOperatorGraphDesc operatorGraphDesc = {}; + operatorGraphDesc.inputEdgeCount = uint32_t(std::size(inputEdges)); + operatorGraphDesc.inputEdges = std::data(inputEdges); + operatorGraphDesc.intermediateEdgeCount = uint32_t(std::size(intermediateEdges)); + operatorGraphDesc.intermediateEdges = std::data(intermediateEdges); + operatorGraphDesc.outputEdgeCount = uint32_t(std::size(outputEdges)); + operatorGraphDesc.outputEdges = std::data(outputEdges); + operatorGraphDesc.nodeCount = uint32_t(std::size(operatorDescPointers)); + operatorGraphDesc.nodes = std::data(operatorDescPointers); + SetDmlOperatorGraphDesc(std::move(operatorGraphDesc), kernelCreationContext); + } + break; + default: return; } } + + // Reproject all inputs and the output to the intermediate product tensor. + // e.g. + // + // Equation: i,j->ji + // + // [1] [4,5,6,7] [4, 8,12] + // [2] -> [5,10,15] + // [3] [6,12,18] + // [7,14,21] + // + // Expand inputs 0 and 1 to 2D via strides to be directly broadcast-compatible. + // + // [1,1,1,1] [4,5,6,7] [4, 8,12] + // [2,2,2,2] [4,5,6,7] -> [5,10,15] + // [3,3,3,3] [4,5,6,7] [6,12,18] + // [7,14,21] + // + // Transpose the output to be shape-compatible: + // + // [1,1,1,1] [4,5,6,7] [ 4, 5, 6, 7] + // [2,2,2,2] [4,5,6,7] -> [ 8,10,12,14] + // [3,3,3,3] [4,5,6,7] [12,15,18,21] + // + void ReprojectTensorDescsToProductTensor() + { + assert(!m_components.empty() && "Equation components should have already been parsed."); + assert(m_inputTensorDescs.size() + m_outputTensorDescs.size() == m_components.size()); + + for (size_t i = 0, count = m_inputTensorDescs.size(); i < count; ++i) + { + auto inputLabels = m_components[i].GetLabels(m_labelIndices); + ReprojectTensorDescToProductTensor(/*inout*/ m_inputTensorDescs[i], inputLabels, /*isReduced*/ false); + } + auto outputLabels = m_components.back().GetLabels(m_labelIndices); + ReprojectTensorDescToProductTensor(/*inout*/ m_outputTensorDescs.front(), outputLabels, /*isReduced*/ true); + } + + // Project the given tensor for shape compatibility to the internal product tensor, which may include broadcasting, + // transposition, and collapsing repeated terms (e.g. iji,i->j with 2 i's in the first term with strides summed). + // + // e.g. + // + // Axis labels: 3,0,2 // the 2 in the inputShape[0] corresponds to productDimensions[3]. + // Original tensor shape: [2,3,4] + // Original tensor strides: [12,4,1] // packed strides right-to-left + // Product tensor shape: [3,5,4,2] // transposed relative to input, with 1 more axis not in input tensor + // Reprojected shape: [3,5,4,2] // identical to product shape + // (or when isReduced) [3,1,4,2] // inserted dimension is 1 + // Reprojected strides: [4,0,1,12] // the newly inserted tensor has 0 stride for broadcasting + // + void ReprojectTensorDescToProductTensor( + /*inout*/ TensorDesc& tensorDesc, + gsl::span axisLabels, + bool isReduced // Return 1's for any missing dimensions not in axisLabels. + ) + { + assert(m_productDimensions.size() == m_uniqueLabelCount && "Product dimensions were not computed yet"); + const size_t newRank = m_productDimensions.size(); + + // Compute the default strides of the tensor (non-transposed). + tensorDesc.EnsureStridesExist(); + const auto originalSizes = tensorDesc.GetSizes(); + const auto originalStrides = tensorDesc.GetStrides(); + assert(originalSizes.size() >= axisLabels.size()); + assert(originalStrides.size() >= axisLabels.size()); + + // Set default sizes for shape compatibility with the product tensor, and + // set strides to 0's initially to broadcast any missing dimensions. + std::vector newSizes; + std::vector newStrides(newRank, 0u); // Default to 0 to broadcast missing entries. + if (isReduced) + { + newSizes.resize(newRank, 1u); // Fill with 1's initially for any missing (reduced) dimensions. + } + else + { + newSizes = m_productDimensions; // Use the product tensor shape directly. Missing axes will be broadcasted. + } + + // Scatter the original sizes and strides into the corresponding product tensor axis. + for (size_t i = 0, count = axisLabels.size(); i < count; ++i) + { + uint32_t productAxis = axisLabels[i]; + if (productAxis < newRank) + { + newSizes[productAxis] = originalSizes[i]; + newStrides[productAxis] += originalStrides[i]; // Add to combine diagonal cases like i,j,i->i,j + } + } + tensorDesc.SetDimensionsAndStrides(newSizes, newStrides); + tensorDesc.EnsureDimensionCount(1, TensorAxis::RightAligned); + } + + // Reproject a tensor to the given axis arrangement. + // The new tensor will have rank == newAxes.size(). + // e.g. + // + // product tensor shape = [2,3,4,5,6] // m_productDimensions + // newAxes = [4,2,0,1] + // new tensor shape = [6,4,2,3] + // + void ReprojectTensorDescToGivenAxes( + /*inout*/ TensorDesc& tensorDesc, + gsl::span axisLabels, + gsl::span newAxes + ) + { + // First, reproject the original dimensions up to the product tensor. + ReprojectTensorDescToProductTensor(/*inout*/ tensorDesc, axisLabels, /*isReduced*/ false); + tensorDesc.PermuteDimensions(newAxes, TensorAxis::LeftAligned); + } + + std::vector GetReductionAxes() const + { + // Determine which axes are reduced by looking for any output dimensions of size 1. + // Note this could include dimensions that are not actually being reduced and simply + // already had size 1 from the input, but such cases harmless nops either way. + + auto outputSizes = m_outputTensorDescs.front().GetSizes(); + std::vector reducedAxes; + FindValueIndices(outputSizes, 1u, /*out*/ reducedAxes); + return reducedAxes; + } }; void CALLBACK QueryEinSum(IMLOperatorSupportQueryContextPrivate* context, bool* isSupported) @@ -276,7 +411,7 @@ void CALLBACK QueryEinSum(IMLOperatorSupportQueryContextPrivate* context, bool* EinSumHelper helper(attributes); auto recognizedOperatorType = helper.GetRecognizedOperatorType(); - static_assert(EinSumHelper::RecognizedOperatorType::Total == static_cast(12), "Update this function."); + static_assert(EinSumHelper::RecognizedOperatorType::Total == static_cast(6), "Verify if this function needs updating."); *isSupported = (recognizedOperatorType != EinSumHelper::RecognizedOperatorType::None); } diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/TensorDesc.cpp b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/TensorDesc.cpp index f26a2ac6fa79a..f738e9c6626fa 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/TensorDesc.cpp +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/TensorDesc.cpp @@ -201,7 +201,7 @@ TensorDesc::TensorDesc( assert(m_bufferTensorDesc.TotalTensorSizeInBytes >= ComputeByteSizeFromDimensions(nonBroadcastDimensions, dataType)); } -gsl::span TensorDesc::GetStrides() const +gsl::span TensorDesc::GetStrides() const noexcept { if (m_bufferTensorDesc.Strides == nullptr) { @@ -212,8 +212,6 @@ gsl::span TensorDesc::GetStrides() const void TensorDesc::SetStrides(gsl::span strides) { - m_bufferTensorDesc.Strides = strides.empty() ? nullptr : strides.data(); - if (!strides.empty()) { ML_CHECK_VALID_ARGUMENT(strides.size() <= std::size(m_strides)); @@ -221,6 +219,8 @@ void TensorDesc::SetStrides(gsl::span strides) std::copy(strides.begin(), strides.end(), m_strides); } + m_bufferTensorDesc.Strides = strides.empty() ? nullptr : m_strides; + m_bufferTensorDesc.TotalTensorSizeInBytes = DMLCalcBufferTensorSize( m_bufferTensorDesc.DataType, m_bufferTensorDesc.DimensionCount, @@ -228,7 +228,7 @@ void TensorDesc::SetStrides(gsl::span strides) strides.empty() ? nullptr : m_strides); } -DML_TENSOR_DESC TensorDesc::GetDmlDesc() +DML_TENSOR_DESC TensorDesc::GetDmlDesc() noexcept { if (m_tensorType == DML_TENSOR_TYPE_INVALID) { @@ -289,6 +289,15 @@ void TensorDesc::ForceUnsignedDataType() } } +// Add additional padding 1's to ensure the count is at least that large. +void TensorDesc::EnsureDimensionCount(uint32_t newDimensionCount, TensorAxis alignment) +{ + if (m_bufferTensorDesc.DimensionCount < newDimensionCount) + { + SetDimensionCount(newDimensionCount, alignment); + } +} + void TensorDesc::SetDimensionCount(uint32_t newDimensionCount, TensorAxis alignment) { ML_CHECK_VALID_ARGUMENT(newDimensionCount <= MaximumDimensionCount); @@ -321,38 +330,48 @@ void TensorDesc::SetDimensionCount(uint32_t newDimensionCount, TensorAxis alignm m_bufferTensorDesc.DimensionCount = newDimensionCount; } -// Uses dimensionMapping to reorder m_sizes and m_strides to match specific Tensor layout +void TensorDesc::SetDimensionsAndStrides(gsl::span sizes, gsl::span strides) +{ + static_assert(sizeof(m_sizes) == sizeof(m_strides)); + ML_CHECK_VALID_ARGUMENT(sizes.size() <= std::size(m_sizes)); + ML_CHECK_VALID_ARGUMENT(strides.empty() || strides.size() == sizes.size()); + + std::copy(sizes.begin(), sizes.end(), m_sizes); + m_bufferTensorDesc.DimensionCount = static_cast(sizes.size()); + SetStrides(strides); +} + void TensorDesc::PermuteDimensions(gsl::span dimensionMapping, const TensorAxis alignment) { + const uint32_t oldRank = m_bufferTensorDesc.DimensionCount; EnsureStridesExist(); SetDimensionCount(static_cast(dimensionMapping.size()), alignment); - // Shuffle m_sizes and m_strides according to the indexes pointed by dimensionMapping - std::vector tempSizes{m_sizes, m_sizes + MaximumDimensionCount}; - std::vector tempStrides{m_strides, m_strides + MaximumDimensionCount}; + // Shuffle m_sizes and m_strides according to the indexes pointed by dimensionMapping. + // Note using MaximumDimensionCount instead of oldRank is intentional here, because the old rank could + // be smaller or larger than the new rank, but it will never be larger than MaximumDimensionCount. + std::vector oldSizes{m_sizes, m_sizes + MaximumDimensionCount}; + std::vector oldStrides{m_strides, m_strides + MaximumDimensionCount}; for (size_t i = 0; i < dimensionMapping.size(); i++) { - m_sizes[i] = tempSizes[dimensionMapping[i]]; - m_strides[i] = tempStrides[dimensionMapping[i]]; + uint32_t sourceAxis = dimensionMapping[i]; + m_sizes[i] = sourceAxis < oldRank ? oldSizes[sourceAxis] : 1; + m_strides[i] = sourceAxis < oldRank ? oldStrides[sourceAxis] : 0; } m_bufferTensorDesc.Sizes = m_sizes; m_bufferTensorDesc.Strides = m_strides; } -void TensorDesc::EnsureStridesExist() +void TensorDesc::EnsureStridesExist() noexcept { if (m_bufferTensorDesc.Strides != nullptr) { - // Strides are populated + // Strides are already populated return; } - uint32_t stride = 1; - for (uint32_t i = m_bufferTensorDesc.DimensionCount; i-- > 0;) - { - m_strides[i] = stride; - stride *= m_sizes[i]; - } + GetDescendingPackedStrides({m_sizes, m_bufferTensorDesc.DimensionCount}, {m_strides, m_bufferTensorDesc.DimensionCount}); + m_bufferTensorDesc.Strides = m_strides; } diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/TensorDesc.h b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/TensorDesc.h index 909e2084d0163..bd9f8a46600b9 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/TensorDesc.h +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/TensorDesc.h @@ -32,18 +32,28 @@ namespace Dml uint32_t guaranteedBaseOffsetAlignment ); - DML_TENSOR_DESC GetDmlDesc(); + DML_TENSOR_DESC GetDmlDesc() noexcept; - inline DML_TENSOR_DATA_TYPE GetDmlDataType() const { return m_bufferTensorDesc.DataType; } - inline MLOperatorTensorDataType GetMlOperatorDataType() const { return m_mlOperatorTensorDataType; } + inline DML_TENSOR_DATA_TYPE GetDmlDataType() const noexcept { return m_bufferTensorDesc.DataType; } + inline MLOperatorTensorDataType GetMlOperatorDataType() const noexcept { return m_mlOperatorTensorDataType; } void ForceUnsignedDataType(); - inline bool IsValid() const { return m_tensorType != DML_TENSOR_TYPE_INVALID; } + inline bool IsValid() const noexcept { return m_tensorType != DML_TENSOR_TYPE_INVALID; } inline uint32_t GetDimensionCount() const { return m_bufferTensorDesc.DimensionCount; } void SetDimensionCount(uint32_t newDimensionCount, TensorAxis alignment); - gsl::span GetSizes() const { return { m_sizes, m_sizes + m_bufferTensorDesc.DimensionCount }; } - gsl::span GetStrides() const; + void EnsureDimensionCount(uint32_t newDimensionCount, TensorAxis alignment); + + gsl::span GetSizes() const noexcept { return { m_sizes, m_sizes + m_bufferTensorDesc.DimensionCount }; } + gsl::span GetStrides() const noexcept; void SetStrides(gsl::span strides); + void EnsureStridesExist() noexcept; + + void SetDimensionsAndStrides(gsl::span sizes, gsl::span strides); + + // Rearranges existing m_sizes and m_strides by gathering axes from dimensionMapping. + // It IS legal to change the number of dimensions by adding filler, dropping entire dimensions for a new view, + // and even duplicating logical dimensions. Axes beyond the original rank will be filled by size 1 and stride 0. + // e.g. Existing sizes [2,3,4] with [2,0] yields [4,2]. void PermuteDimensions(gsl::span dimensionMapping, const TensorAxis alignment); inline uint64_t GetBufferSizeInBytes() const @@ -91,8 +101,6 @@ namespace Dml uint32_t m_sizes[MaximumDimensionCount] = {}; uint32_t m_strides[MaximumDimensionCount] = {}; DML_BUFFER_TENSOR_DESC m_bufferTensorDesc = {}; - - void EnsureStridesExist(); }; class TensorDescBuilder diff --git a/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorHelper.cpp b/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorHelper.cpp index 44c63089564d3..3a7cf28ef903e 100644 --- a/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorHelper.cpp +++ b/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorHelper.cpp @@ -68,6 +68,17 @@ namespace OperatorHelper originalValues = std::move(expanded); } + uint32_t GetBitMaskFromIndices(gsl::span indices) noexcept + { + uint32_t bitMask = 0; + for (auto i : indices) + { + assert(i < 32); + bitMask |= (1 << i); + } + return bitMask; + } + float CastFloat16ToFloat32(uint16_t input) { // Promote float16m10e5s1 to float32m23e8s1. @@ -1384,6 +1395,7 @@ namespace OperatorHelper m_recognizedOperatorType = DetermineRecognizedOperatorType(); } + // Updates: m_components, m_labelIndices, m_uniqueLabelCount. void EinSumHelper::ParseEquationComponents() { // Parse an equation like 'ij,jk->ik' into components {ij, jk, ik} mapping letters to @@ -1477,134 +1489,125 @@ namespace OperatorHelper currentComponent.labelIndexEnd = static_cast(m_labelIndices.size()); m_components.push_back(currentComponent); } + + m_uniqueLabelCount = labelMap.size(); } - EinSumHelper::RecognizedOperatorType EinSumHelper::DetermineRecognizedOperatorType() + EinSumHelper::RecognizedOperatorType EinSumHelper::DetermineRecognizedOperatorType() const { if (m_components.empty()) { - return RecognizedOperatorType::None; // Parsing may have found unsupported components - treating as unknown. + return RecognizedOperatorType::None; // Parsing may have found unsupported components - treating as unknown. } - // std::ranges::equal is not supported yet. - auto equals = [](gsl::span a, gsl::span b) + auto areIdenticalAxes = [](gsl::span a, gsl::span b) -> bool { - return std::equal(a.begin(), a.end(), b.begin(), b.end()); + return GetBitMaskFromIndices(a) == GetBitMaskFromIndices(b); }; - auto as_span = [](std::initializer_list il) { + auto as_span = [](std::initializer_list il) + { return gsl::make_span(il.begin(), il.size()); }; - std::array componentRanks; - if (m_components.size() > componentRanks.size()) + // Identify any common patterns that map to existing DirectML operators + // (identity, transpose, sum reduction, matmul, multiplication...). + + constexpr size_t maximumSupportedComponentCount = 3; // 2 inputs + 1 output + // Bail if more than 2 inputs. EinSum is generic and can handle any variable number of inputs, + // but 3-input models have yet to be seen. + if (m_components.size() > maximumSupportedComponentCount) { - // No recognized operator takes more than 2 inputs and 1 output. - // EinSum itself is generic and can handle any variable number of inputs, - // but DML's operators expect fixed counts. return RecognizedOperatorType::None; } - else if (m_components.size() == 2) + else if (m_components.size() == 2) // 1 input, 1 output { - auto inputLabels = m_components[0].GetLabels(m_labelIndices); - auto outputLabels = m_components[1].GetLabels(m_labelIndices); - if (inputLabels.size() == outputLabels.size()) + auto outputLabels = m_components.back().GetLabels(m_labelIndices); + + // Use reduction if the output has fewer dimensions than total dimension labels. + if (outputLabels.size() <= m_uniqueLabelCount) { - // Check identity. - if (equals(inputLabels, outputLabels)) - { - // Handles: "->", "i->i", "ij->ij", "ijk->ijk", "ijkl->ijkl" ... - return RecognizedOperatorType::Identity; - } - else // Transpose since a permutation exists. - { - // Handles: "ij->ji", "ijk->kji", "ijkl->lkji", "ijkl->ijkl" ... - return RecognizedOperatorType::Transpose; - } + // Handles: "ij->i", "i->", "ij->", "ijkl->jl", "ijkl->", "iji->" ... + return RecognizedOperatorType::ReduceSum; } - else if (outputLabels.empty()) // Scalar output, with all inputs reduced. + else { - // Handles: "i->", "ij->", "ijk->", "ijkl->" ... - return RecognizedOperatorType::ReduceSum; + // Handles identity: "->", "i->i", "ij->ij", "ijk->ijk", "ijkl->ijkl" ... + // Handles transpose: "ij->ji", "ijk->kji", "ijkl->lkji", "ijkl->ijkl" ... + // Handles diagional: "ii->i", "iij->ji" + return RecognizedOperatorType::Transpose; } } - else if (m_components.size() == 3) + else if (m_components.size() == 3) // 2 inputs, 1 output { - // If all components have the same size and label order, then apply elementwise multiplication. - auto inputALabels = m_components[0].GetLabels(m_labelIndices); - auto inputBLabels = m_components[1].GetLabels(m_labelIndices); + auto input0Labels = m_components[0].GetLabels(m_labelIndices); + auto input1Labels = m_components[1].GetLabels(m_labelIndices); auto outputLabels = m_components[2].GetLabels(m_labelIndices); - if (equals(inputALabels, outputLabels) && equals(inputBLabels, outputLabels)) + + // Use elementwise multiplication when no reduction occurs. + if (outputLabels.size() == m_uniqueLabelCount) { - // Handles: "i,i->i", "ij,ij->ij", "ijk,ijk->ijk", "ijkl,ijkl->ijkl" ... + // Handles: "i,i->i", "ij,ij->ij", "ijk,ijk->kji", "ijkl,klij->jilk" ... return RecognizedOperatorType::Multiply; } - } - - // Otherwise check for special cases of dedicated operators... - - struct RecognizedOperatorInfo - { - RecognizedOperatorType recognizedOperatorType; - std::initializer_list componentRanks; - std::initializer_list labelIndices; - }; - - const RecognizedOperatorInfo recognizedOperators[] = { - {RecognizedOperatorType::MatMul, {2,2,2},{0,1, 1,2, 0,2}}, // ij,jk->ik - {RecognizedOperatorType::MatMul, {3,3,3},{0,1,2, 0,2,3, 0,1,3}}, // bij,bjk->bik - {RecognizedOperatorType::MatMul, {4,4,4},{0,1,2,3, 0,1,3,4, 0,1,2,4}}, // abij,abjk->abik - {RecognizedOperatorType::OuterProduct, {1,1,2},{0, 1, 0,1}}, // i,j->ij - {RecognizedOperatorType::MatMulTransposeA, {2,2,2},{0,1, 0,2, 1,2}}, // ji,jk->ik - {RecognizedOperatorType::MatMulTransposeA, {3,3,3},{0,1,2, 0,1,3, 0,2,3}}, // bji,bjk->bik - {RecognizedOperatorType::MatMulTransposeA, {4,4,4},{0,1,2,3, 0,1,2,4, 0,1,3,4}}, // abji,abjk->abik - {RecognizedOperatorType::MatMulTransposeB, {2,2,2},{0,1, 2,1, 0,2}}, // ij,kj->ik - {RecognizedOperatorType::MatMulTransposeB, {3,3,3},{0,1,2, 0,3,2, 0,1,3}}, // bij,bkj->bik - {RecognizedOperatorType::MatMulTransposeB, {4,4,4},{0,1,2,3, 0,1,4,3, 0,1,2,4}}, // abij,abkj->abik - {RecognizedOperatorType::MatMulTransposeB, {1,1,0},{0,0,}}, // i,i-> (1D inner_prod) - {RecognizedOperatorType::MatMulNhcw, {4,4,4},{0,1,2,3, 0,3,2,4, 0,1,2,4}}, // aibj,ajbk->aibk - {RecognizedOperatorType::MatMulNhcwTransposeA, {4,4,4},{0,1,2,3, 0,1,2,4, 0,3,2,4}}, // ajbi,ajbk->aibk - {RecognizedOperatorType::MatMulNhcwTransposeB, {4,4,4},{0,1,2,3, 0,4,2,3, 0,1,2,4}}, // aibj,akbj->aibk - {RecognizedOperatorType::ReduceSum, {2,1 },{0,1, 0}}, // ij->i - {RecognizedOperatorType::ReduceSum, {2,1 },{0,1, 1}}, // ij->j - }; - - // For each recognized operator, compare the labels-per-component and label indices. - for (auto& recognizedOperator : recognizedOperators) - { - if (equals(m_labelIndices, as_span(recognizedOperator.labelIndices)) - && m_components.size() == recognizedOperator.componentRanks.size()) + // Use matrix multiplication when exactly 1 dimension is being reduced, + // the output is up to 4D (DML limit), and the inputs have distinct axes. + else if (outputLabels.size() + 1 == m_uniqueLabelCount + && outputLabels.size() <= 4 + && !areIdenticalAxes(input0Labels, input1Labels)) { - for (size_t i = 0; i < m_components.size(); ++i) - { - componentRanks[i] = m_components[i].GetDimensionCount(); - } - - if (equals(gsl::make_span(componentRanks.data(), m_components.size()), as_span(recognizedOperator.componentRanks))) - { - return recognizedOperator.recognizedOperatorType; - } + return RecognizedOperatorType::MatMul; + } + // Otherwise use an elementwise multiplication and sum reduction combo. + // This is the most generic, but it also uses more intermediate memory. + else + { + // Handles: "ij,ji->", "ijkl,abij->ab", ... + return RecognizedOperatorType::MultiplyReduceSum; } } return RecognizedOperatorType::None; } - std::vector EinSumHelper::GetOutputShapes(const MLShapeInferenceContext& shapeInfo) const + // Updates: m_productDimensions. + void EinSumHelper::ExtractLabelSizesFromTensors( + const IKernelInformationAdapter& kernelInformation, + const IShapeInformationAdapter& shapeInformation + ) { - assert(!m_components.empty()); // Should have already parsed components. + // Get the dimension length for each term label. e.g. + // + // equation = "ijk,jm->im" + // input[0].shape = [3,2,4] + // input[1].shape = [2,5] + // output[0].shape = [3,5] + // + // Yields: + // + // i = 3 // m_productDimensions[0] + // j = 2 // m_productDimensions[1] + // k = 4 // m_productDimensions[2] + // m = 5 // m_productDimensions[3] + // + // The character labels are already resolved to numeric indices in occurrence order, and so the product + // dimensions are simply an array of sizes: + // + // label sizes = [3,2,4,5] + // + assert(!m_components.empty()); // Should have already parsed components. - uint32_t inputCount = shapeInfo.GetInputCount(); - uint32_t outputCount = shapeInfo.GetOutputCount(); + uint32_t inputCount = kernelInformation.GetInputCount(); + uint32_t outputCount = kernelInformation.GetOutputCount(); ML_CHECK_VALID_ARGUMENT(inputCount + 1 == m_components.size(), "Mismatch between input tensor count and string equation component count."); ML_CHECK_VALID_ARGUMENT(outputCount == 1, "EinSum expects exactly 1 output tensor."); - std::vector labelSizes(m_labelIndices.size(), UINT_MAX); + m_productDimensions.resize(m_uniqueLabelCount, UINT_MAX); // Read every input tensor, comparing labels to ensure consistent sizes from the equation parsed earlier. for (uint32_t i = 0; i < inputCount; ++i) { - auto inputShape = shapeInfo.GetInputTensorShape(i); + auto inputShape = shapeInformation.GetInputTensorShape(i); auto& component = m_components[i]; auto labelIndices = component.GetLabels(m_labelIndices); uint32_t dimensionCount = component.GetDimensionCount(); @@ -1618,18 +1621,23 @@ namespace OperatorHelper // e.g. Given "ij,ji", both i's and both j's must match dimension sizes. uint32_t dimensionSize = inputShape[j]; uint32_t labelIndex = labelIndices[j]; - assert(labelIndex < labelSizes.size()); + assert(labelIndex < m_productDimensions.size()); - if (labelSizes[labelIndex] == UINT_MAX) + if (m_productDimensions[labelIndex] == UINT_MAX) { - labelSizes[labelIndex] = dimensionSize; + m_productDimensions[labelIndex] = dimensionSize; } else { - ML_CHECK_VALID_ARGUMENT(labelSizes[labelIndex] == dimensionSize, "All labels must have the same dimension sizes."); + ML_CHECK_VALID_ARGUMENT(m_productDimensions[labelIndex] == dimensionSize, "All labels must have the same dimension sizes."); } } } + } + + std::vector EinSumHelper::GetOutputShapes(const MLShapeInferenceContext& shapeInfo) const + { + assert(!m_components.empty()); // Should have already parsed components. // Generate output dimensions from corresponding input tensor labels. // e.g. Given ij,jk->ij with [2,3] and [3,5], the output is [2,5]. @@ -1637,7 +1645,7 @@ namespace OperatorHelper auto outputLabelIndices = m_components.back().GetLabels(m_labelIndices); for (auto labelIndex : outputLabelIndices) { - outputDimensions.push_back(labelSizes[labelIndex]); + outputDimensions.push_back(m_productDimensions[labelIndex]); } return { EdgeShapes(outputDimensions) }; @@ -1645,13 +1653,8 @@ namespace OperatorHelper bool EinSumHelper::IsMatMulOperatorType() const noexcept { - return m_recognizedOperatorType == RecognizedOperatorType::OuterProduct || - m_recognizedOperatorType == RecognizedOperatorType::MatMul || - m_recognizedOperatorType == RecognizedOperatorType::MatMulTransposeA || - m_recognizedOperatorType == RecognizedOperatorType::MatMulTransposeB || - m_recognizedOperatorType == RecognizedOperatorType::MatMulNhcw || - m_recognizedOperatorType == RecognizedOperatorType::MatMulNhcwTransposeA || - m_recognizedOperatorType == RecognizedOperatorType::MatMulNhcwTransposeB; + static_assert(RecognizedOperatorType::Total == static_cast(6), "Verify this for any potentially new matrix multiplication operators."); + return m_recognizedOperatorType == RecognizedOperatorType::MatMul; } std::vector MatMulHelperBase::GetOutputShapes(const MLShapeInferenceContext& shapeInfo) const diff --git a/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorHelper.h b/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorHelper.h index 06bed80a7c27d..b775de0b39cf4 100644 --- a/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorHelper.h +++ b/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorHelper.h @@ -762,9 +762,9 @@ class ReduceHelper : public ReduceHelperBase class EinSumHelper { -public: void Initialize(); +public: // Info_t is used to obtain attributes which will be used for calculating the output shape later. // Shape_t is used to obtain input shape which will be used for adjusting attribute value. template @@ -772,6 +772,7 @@ class EinSumHelper { m_equation = info.GetAttribute(AttrName::Equation); Initialize(); + ExtractLabelSizesFromTensors(KernelInformationAdapter(info), ShapeInformationAdapter(shape)); } EinSumHelper(const MLOperatorAttributes& info) @@ -785,17 +786,11 @@ class EinSumHelper enum class RecognizedOperatorType { None, - Identity, - Multiply, - OuterProduct, - MatMul, - MatMulTransposeA, - MatMulTransposeB, - MatMulNhcw, - MatMulNhcwTransposeA, - MatMulNhcwTransposeB, - ReduceSum, - Transpose, + Transpose, // 1 input, rearrangement or diagonal slice, no reduction + ReduceSum, // 1 input, no multiplication, just sum reduction + Multiply, // 2 inputs, elementwise multiplication, no sum reduction + MatMul, // 2 inputs, elementwise multiplication, sum reduction on 1 axis. + MultiplyReduceSum, // 2 inputs, elementwise multiplication, sum reduction on multiple axes Total, }; @@ -805,7 +800,11 @@ class EinSumHelper protected: void ParseEquationComponents(); - RecognizedOperatorType DetermineRecognizedOperatorType(); + void ExtractLabelSizesFromTensors( + const IKernelInformationAdapter& kernelInformation, + const IShapeInformationAdapter& shapeInformation + ); + RecognizedOperatorType DetermineRecognizedOperatorType() const; protected: struct Component @@ -824,9 +823,10 @@ class EinSumHelper }; std::string m_equation; - std::vector m_labelIndices; // Concatenation of all labels as rebased indices ("ij,ai" -> 0,1,2,0). - std::vector m_components; // All components in order, including inputs and output. - std::vector m_outputDimensions; + size_t m_uniqueLabelCount = 0; // e.g. ij,jk->ij has 3 unique labels. + std::vector m_labelIndices; // Concatenation of all labels as rebased indices ("ij,ai" -> 0,1,2,0). + std::vector m_components; // All components in order, including inputs and output. + std::vector m_productDimensions; // Dimensions of each unique label (size() == m_uniqueLabelCount). RecognizedOperatorType m_recognizedOperatorType = RecognizedOperatorType::None; }; diff --git a/onnxruntime/test/providers/cpu/math/einsum_test.cc b/onnxruntime/test/providers/cpu/math/einsum_test.cc index 423ea3f682f4c..73f31787e0597 100644 --- a/onnxruntime/test/providers/cpu/math/einsum_test.cc +++ b/onnxruntime/test/providers/cpu/math/einsum_test.cc @@ -274,16 +274,11 @@ TEST(Einsum, ExplicitEinsumAsMatmul_OutputTransposed) { } TEST(Einsum, ExplicitEinsumAsMatmul_2) { - // TODO: Unskip when fixed #41968513 - if (DefaultDmlExecutionProvider().get() != nullptr) { - GTEST_SKIP() << "Skipping because of the following error: MLOperatorAuthorImpl.cpp(2068): The parameter is incorrect."; - } - OpTester test("Einsum", 12, onnxruntime::kOnnxDomain); test.AddAttribute("equation", "ij,jk->ik"); test.AddInput("x", {2, 1}, {2.f, 3.f}); - test.AddInput("y", {2, 2}, {1.f, 2.f, 3.f, 4.f}); - test.AddOutput("o", {2, 2}, {8.f, 12.f, 12.f, 18.f}); + test.AddInput("y", {1, 2}, {1.f, 2.f}); + test.AddOutput("o", {2, 2}, {2.f, 4.f, 3.f, 6.f}); test.Run(); } @@ -325,16 +320,11 @@ TEST(Einsum, ImplicitEinsumAsBatchedMatmulWithBroadcasting_0) { } TEST(Einsum, ImplicitEinsumAsMatmul_2) { - // TODO: Unskip when fixed #41968513 - if (DefaultDmlExecutionProvider().get() != nullptr) { - GTEST_SKIP() << "Skipping because of the following error: MLOperatorAuthorImpl.cpp(2068): The parameter is incorrect."; - } - OpTester test("Einsum", 12, onnxruntime::kOnnxDomain); test.AddAttribute("equation", "ij,jk"); test.AddInput("x", {2, 1}, {2.f, 3.f}); - test.AddInput("y", {2, 2}, {1.f, 2.f, 3.f, 4.f}); - test.AddOutput("o", {2, 2}, {8.f, 12.f, 12.f, 18.f}); + test.AddInput("y", {1, 2}, {1.f, 2.f}); + test.AddOutput("o", {2, 2}, {2.f, 4.f, 3.f, 6.f}); test.Run(); } @@ -434,11 +424,6 @@ TEST(Einsum, ExplicitEinsumAsBatchedDiagonalOp_1) { // Implicit (Implicit diagonal ops will sum up diagonal values) TEST(Einsum, ImplicitEinsumAsDiagonalOp) { - // TODO: Unskip when fixed #41968513 - if (DefaultDmlExecutionProvider().get() != nullptr) { - GTEST_SKIP() << "Skipping because of the following error: The difference between expected[i] and output[i] is 5, which exceeds threshold"; - } - OpTester test("Einsum", 12, onnxruntime::kOnnxDomain); test.AddAttribute("equation", "ii"); test.AddInput("x", {2, 2}, {1.f, 2.f, 3.f, 4.f}); @@ -447,11 +432,6 @@ TEST(Einsum, ImplicitEinsumAsDiagonalOp) { } TEST(Einsum, ImplicitEinsumAsDiagonalOp_1) { - // TODO: Unskip when fixed #41968513 - if (DefaultDmlExecutionProvider().get() != nullptr) { - GTEST_SKIP() << "Skipping because of the following error: error: The difference between expected[i] and output[i] is 15, which exceeds threshold"; - } - OpTester test("Einsum", 12, onnxruntime::kOnnxDomain); test.AddAttribute("equation", "iii"); test.AddInput("x", {2, 2, 2}, {1.f, 2.f, 3.f, 4.f, 1.f, 2.f, 3.f, 4.f}); @@ -576,8 +556,17 @@ TEST(Einsum, ExplicitEinsumAsTensorContractionReshapeLeft) { OpTester test("Einsum", 12, onnxruntime::kOnnxDomain); test.AddAttribute("equation", "bsnh,btnh->bnts"); test.AddInput("x", {2, 1, 2, 2}, {1.f, 2.f, 1.f, 2.f, 1.f, 2.f, 1.f, 2.f}); - test.AddInput("y", {2, 2, 2, 1}, {1.f, 2.f, 3.f, 4.f, 5.f, 6.f, 7.f, 8.f}); - test.AddOutput("o", {2, 2, 2, 1}, {3.f, 9.f, 6.f, 12.f, 15.f, 21.f, 18.f, 24.f}); + test.AddInput("y", {2, 2, 2, 2}, {1.f, 2.f, 3.f, 4.f, 5.f, 6.f, 7.f, 8.f, 9.f, 10.f, 11.f, 12.f, 13.f, 14.f, 15.f, 16.f}); + test.AddOutput("o", {2, 2, 2, 1}, {5.f, 17.f, 11.f, 23.f, 29.f, 41.f, 35.f, 47.f}); + test.Run(); +} + +TEST(Einsum, ExplicitEinsumAsTensorContractionSameInput) { + OpTester test("Einsum", 12, onnxruntime::kOnnxDomain); + test.AddAttribute("equation", "nchw,nchw->nch"); + test.AddInput("x", {1, 3, 2, 4}, {1.f, 2.f, 3.f, 4.f, 5.f, 6.f, 7.f, 8.f, 2.f, 3.f, 4.f, 5.f, 6.f, 7.f, 8.f, 9.f, 3.f, 4.f, 5.f, 6.f, 7.f, 8.f, 9.f, 10.f}); + test.AddInput("y", {1, 3, 2, 4}, {1.f, 2.f, 3.f, 4.f, 5.f, 6.f, 7.f, 8.f, 2.f, 3.f, 4.f, 5.f, 6.f, 7.f, 8.f, 9.f, 3.f, 4.f, 5.f, 6.f, 7.f, 8.f, 9.f, 10.f}); + test.AddOutput("o", {1, 3, 2}, {30.f, 174.f, 54.f, 230.f, 86.f, 294.f}); test.Run(); } From b95982e588e2d433958aa9aa82557b714ec4e3de Mon Sep 17 00:00:00 2001 From: Dmitri Smirnov Date: Fri, 21 Jun 2024 13:58:21 -0700 Subject: [PATCH 20/52] Fix 2D detection bug (#21128) ### Description Should compare two leading dims for 1.f ### Motivation and Context Vulnerability scanner --- onnxruntime/core/providers/cuda/tensor/upsample.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/onnxruntime/core/providers/cuda/tensor/upsample.cc b/onnxruntime/core/providers/cuda/tensor/upsample.cc index 17533eb3d9a72..cbf745d3c7b4f 100644 --- a/onnxruntime/core/providers/cuda/tensor/upsample.cc +++ b/onnxruntime/core/providers/cuda/tensor/upsample.cc @@ -230,7 +230,7 @@ Status Upsample::BaseCompute(OpKernelContext* context, } const bool is_2D = X_dims.size() == 2; - const bool is_nchw = is_2D ? true : (scales[1] == 1.0f && scales[1] == 1.0f); + const bool is_nchw = is_2D ? true : (scales[0] == 1.0f && scales[1] == 1.0f); ORT_RETURN_IF_NOT(is_nchw, "Resize 'Cubic' mode only supports NCWH layout " From 5b5ce0bfb0183bc26a1aab58f61fc8785cdab440 Mon Sep 17 00:00:00 2001 From: Yi Zhang Date: Mon, 24 Jun 2024 13:36:13 +0800 Subject: [PATCH 21/52] Add UsePython Task in Nuget Publish workflow (#21144) ### Description Otherwise it would fail in https://github.com/microsoft/onnxruntime/blob/b95982e588e2d433958aa9aa82557b714ec4e3de/tools/ci_build/github/azure-pipelines/publish-nuget.yml#L78-L81 ### Motivation and Context The Windows CPU image is migrated to managed image ### Verification Link https://dev.azure.com/aiinfra/Lotus/_build?definitionId=1313 --- tools/ci_build/github/azure-pipelines/publish-nuget.yml | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tools/ci_build/github/azure-pipelines/publish-nuget.yml b/tools/ci_build/github/azure-pipelines/publish-nuget.yml index 8ce7915da76d1..367977ff59192 100644 --- a/tools/ci_build/github/azure-pipelines/publish-nuget.yml +++ b/tools/ci_build/github/azure-pipelines/publish-nuget.yml @@ -29,6 +29,12 @@ stages: - checkout: self submodules: false + + - task: UsePythonVersion@0 + inputs: + versionSpec: '3.9' + addToPath: true + - template: templates/set-version-number-variables-step.yml - script: mkdir "$(Build.BinariesDirectory)\nuget-artifact\final-package" From 269d9b094ff44f91e6cae9d29fbc7ce0ea2e90c1 Mon Sep 17 00:00:00 2001 From: zhijiang <43435212+zhijxu-MS@users.noreply.github.com> Date: Mon, 24 Jun 2024 16:07:39 +0800 Subject: [PATCH 22/52] Zhijxu/fix softmax cudnn bf16 (#21045) if seq >2048, ort will fallback to cudnn version, while when dtype is bf16, ort will throw exception, this PR trying to fix it. --- .../core/providers/cuda/cudnn_common.cc | 4 +++ onnxruntime/test/common/random_generator.h | 2 +- .../orttraining_test_ortmodule_onnx_ops.py | 32 +++++++++++++++++++ 3 files changed, 37 insertions(+), 1 deletion(-) diff --git a/onnxruntime/core/providers/cuda/cudnn_common.cc b/onnxruntime/core/providers/cuda/cudnn_common.cc index 9aa011c1d0ec4..914dc02a9eda4 100644 --- a/onnxruntime/core/providers/cuda/cudnn_common.cc +++ b/onnxruntime/core/providers/cuda/cudnn_common.cc @@ -174,7 +174,11 @@ cudnnDataType_t CudnnTensor::GetDataType() { template <> cudnnDataType_t CudnnTensor::GetDataType() { +#if defined(CUDNN_VERSION) && CUDNN_VERSION >= 8200 + return CUDNN_DATA_BFLOAT16; +#else ORT_THROW("cuDNN doesn't support BFloat16."); +#endif } template <> diff --git a/onnxruntime/test/common/random_generator.h b/onnxruntime/test/common/random_generator.h index cb1ce885d2d45..336e0f197fcc9 100644 --- a/onnxruntime/test/common/random_generator.h +++ b/onnxruntime/test/common/random_generator.h @@ -70,7 +70,7 @@ class RandomValueGenerator { // Random values generated are in the range [min, max). template typename std::enable_if< - std::is_same_v, + std::is_same_v || std::is_same_v, std::vector>::type Uniform(gsl::span dims, float min, float max) { std::vector val(detail::SizeFromDims(dims)); diff --git a/orttraining/orttraining/test/python/orttraining_test_ortmodule_onnx_ops.py b/orttraining/orttraining/test/python/orttraining_test_ortmodule_onnx_ops.py index df0b5f195f0b9..88735ff18515e 100644 --- a/orttraining/orttraining/test/python/orttraining_test_ortmodule_onnx_ops.py +++ b/orttraining/orttraining/test/python/orttraining_test_ortmodule_onnx_ops.py @@ -146,6 +146,38 @@ def test_onnx_ops(self): device = torch.device(device_name) self.gradient_correctness(name, device) + @unittest.skipIf(not torch.cuda.is_bf16_supported(), "Test requires CUDA and BF16 support") + def test_softmax_bf16_large(self): + if not torch.cuda.is_available(): + # only test bf16 on cuda + return + + class Model(torch.nn.Module): + def __init__(self): + super().__init__() + + def forward(self, input): + out = torch.softmax(input, dim=-1) + return out + + device = "cuda:0" + input_shape = [2, 4096] + # run torch to get the expected result + data_torch = torch.randn(size=input_shape, device=device, dtype=torch.bfloat16) + 10 + data_torch.requires_grad = True + torch_model = Model() + torch_res = torch_model(input=data_torch) + init_grad = torch.ones_like(torch_res) + torch_res.backward(gradient=init_grad) + # run ort + ort_model = ORTModule(torch_model) + data_ort = data_torch.detach().clone() + data_ort.requires_grad = True + ort_res = ort_model(input=data_ort) + ort_res.backward(gradient=init_grad) + # compara result + torch.testing.assert_close(data_torch.grad, data_ort.grad, rtol=1e-5, atol=1e-4) + if __name__ == "__main__": unittest.main() From ebd0368bb0a2ca55685810b76b096717fd734ab6 Mon Sep 17 00:00:00 2001 From: aciddelgado <139922440+aciddelgado@users.noreply.github.com> Date: Mon, 24 Jun 2024 09:43:49 -0700 Subject: [PATCH 23/52] Make Flash Attention work on Windows (#21015) ### Description Previously, Flash Attention only worked on Linux systems. This PR will make it work and enable it to be built and run on Windows. Limitations of Flash Attention in Windows: Requires CUDA 12. ### Motivation and Context This will significantly increase the performance of Windows-based LLM's with hardware sm>=80. To illustrate the improvement of Flash Attention over Memory Efficient Attention, here are some average benchmark numbers for the GQA operator, run with configurations based on several recent models (Llama, Mixtral, Phi-3). The benchmarks were obtained on RTX4090 GPU using the test script located at (onnxruntime/test/python/transformers/benchmark_gqa_windows.py). * Clarifying Note: These benchmarks are just for the GQA operator, not the entire model. ### Memory Efficient Attention Kernel Benchmarks: | Model Name | Max Sequence Length | Inference Interval (ms) | Throughput (samples/second) | |----------------------------------------|---------------------|-------------------------|-----------------------------| | Llama3-8B (Average Prompt) | 8192 | 0.19790525 | 13105.63425 | | Llama3-8B (Average Token) | 8192 | 0.207775538 | 12025.10172 | | Llama3-70B (Average Prompt) | 8192 | 0.216049167 | 11563.31185 | | Llama3-70B (Average Token) | 8192 | 0.209730731 | 12284.38149 | | Mixtral-8x22B-v0.1 (Average Prompt) | 32768 | 0.371928785 | 7031.440056 | | Mixtral-8x22B-v0.1 (Average Token) | 32768 | 0.2996659 | 7607.947159 | | Phi-3-mini-128k (Average Prompt) | 131072 | 0.183195867 | 15542.0852 | | Phi-3-mini-128k (Average Token) | 131072 | 0.198215688 | 12874.53494 | | Phi-3-small-128k (Average Prompt) | 65536 | 2.9884929 | 2332.584142 | | Phi-3-small-128k (Average Token) | 65536 | 0.845072406 | 2877.85822 | | Phi-3-medium-128K (Average Prompt) | 32768 | 0.324974429 | 8094.909517 | | Phi-3-medium-128K (Average Token) | 32768 | 0.263662567 | 8978.463687 | ### Flash Attention Kernel Benchmarks: | Model Name | Max Sequence Length | Inference Interval (ms) | Throughput (samples/second) | |--------------------------------------|---------------------|-------------------------|-----------------------------| | Llama3-8B (Average Prompt) | 8192 | 0.163566292 | 16213.69057 | | Llama3-8B (Average Token) | 8192 | 0.161643692 | 16196.14715 | | Llama3-70B (Average Prompt) | 8192 | 0.160510375 | 17448.67753 | | Llama3-70B (Average Token) | 8192 | 0.169427308 | 14702.62043 | | Mixtral-8x22B-v0.1 (Average Prompt) | 32768 | 0.164121964 | 15618.51301 | | Mixtral-8x22B-v0.1 (Average Token) | 32768 | 0.1715865 | 14524.32273 | | Phi-3-mini-128k (Average Prompt) | 131072 | 0.167527167 | 14576.725 | | Phi-3-mini-128k (Average Token) | 131072 | 0.175940594 | 15762.051 | | Phi-3-small-128k (Average Prompt) | 65536 | 0.162719733 | 17824.494 | | Phi-3-small-128k (Average Token) | 65536 | 0.14977525 | 16749.19858 | | Phi-3-medium-128K (Average Prompt) | 32768 | 0.156490786 | 17679.2513 | | Phi-3-medium-128K (Average Token) | 32768 | 0.165333833 | 14932.26079 | Flash Attention is consistently faster for every configuration we benchmarked, with improvements in our trials ranging from ~20% to ~650%. In addition to these improvements in performance, Flash Attention has better memory usage. For example, Memory Efficient Attention cannot handle a max sequence length higher than 32,768, but Flash Attention can handle max sequence lengths at least as high as 131,072. --------- Co-authored-by: Tianlei Wu --- .github/workflows/lint.yml | 1 + .lintrunner.toml | 1 + cmake/CMakeLists.txt | 5 +- .../contrib_ops/cuda/bert/attention.cc | 2 +- .../contrib_ops/cuda/bert/attention_impl.cu | 5 + .../cuda/bert/flash_attention/alibi.h | 67 ++ .../cuda/bert/flash_attention/block_info.h | 27 +- .../cuda/bert/flash_attention/flash.h | 13 +- .../cuda/bert/flash_attention/flash_api.cc | 116 +++- .../cuda/bert/flash_attention/flash_api.h | 21 +- .../flash_fwd_hdim128_bf16_sm80.cu | 4 +- .../flash_fwd_hdim160_bf16_sm80.cu | 4 +- .../flash_fwd_hdim192_bf16_sm80.cu | 4 +- .../flash_fwd_hdim224_bf16_sm80.cu | 4 +- .../flash_fwd_hdim256_bf16_sm80.cu | 4 +- .../flash_fwd_hdim32_bf16_sm80.cu | 4 +- .../flash_fwd_hdim64_bf16_sm80.cu | 4 +- .../flash_fwd_hdim96_bf16_sm80.cu | 4 +- .../bert/flash_attention/flash_fwd_kernel.h | 588 ++++++++---------- .../flash_fwd_launch_template.h | 131 ++-- .../cuda/bert/flash_attention/kernel_traits.h | 169 ++--- .../cuda/bert/flash_attention/mask.h | 208 +++++++ .../cuda/bert/flash_attention/rotary.h | 154 +++++ .../cuda/bert/flash_attention/softmax.h | 169 ++--- .../cuda/bert/flash_attention/static_switch.h | 33 +- .../cuda/bert/flash_attention/utils.h | 225 ++----- .../cuda/bert/group_query_attention.cc | 2 +- .../cuda/bert/group_query_attention_impl.cu | 2 +- .../cuda/bert/multihead_attention.cc | 2 +- .../bert/packed_multihead_attention_impl.cu | 4 +- .../transformers/benchmark_gqa_windows.py | 221 +++++++ .../transformers/test_flash_attn_cuda.py | 56 +- .../python/transformers/test_parity_moe.py | 3 + .../azure-pipelines/win-gpu-ci-pipeline.yml | 2 +- 34 files changed, 1397 insertions(+), 862 deletions(-) create mode 100644 onnxruntime/contrib_ops/cuda/bert/flash_attention/alibi.h create mode 100644 onnxruntime/contrib_ops/cuda/bert/flash_attention/mask.h create mode 100644 onnxruntime/contrib_ops/cuda/bert/flash_attention/rotary.h create mode 100644 onnxruntime/test/python/transformers/benchmark_gqa_windows.py diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml index 34911cfc7972e..fc70fa14cff36 100644 --- a/.github/workflows/lint.yml +++ b/.github/workflows/lint.yml @@ -97,6 +97,7 @@ jobs: --exclude=java/src/main/native/*.c --exclude=onnxruntime/core/mlas/inc/* --exclude=onnxruntime/core/mlas/lib/* + --exclude=onnxruntime/contrib_ops/cuda/bert/flash_attention/* filter: "-runtime/references" lint-js: diff --git a/.lintrunner.toml b/.lintrunner.toml index 08d8af1206a1c..e6d06b34726fe 100644 --- a/.lintrunner.toml +++ b/.lintrunner.toml @@ -136,6 +136,7 @@ exclude_patterns = [ 'onnxruntime/core/mickey/cutlass_ext/**', # CUTLASS based libs recommends NO automatic code formatting 'onnxruntime/core/mickey/gemm/**', # CUTLASS based libs recommends NO automatic code formatting 'winml/lib/Api.Image/shaders/**', # Contains data chunks + 'onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_fwd_launch_template.h', # Bool Switches hang Clang ] command = [ 'python', diff --git a/cmake/CMakeLists.txt b/cmake/CMakeLists.txt index ce22def914851..b2122bf56abd8 100644 --- a/cmake/CMakeLists.txt +++ b/cmake/CMakeLists.txt @@ -102,7 +102,7 @@ option(onnxruntime_USE_PREINSTALLED_EIGEN "Use pre-installed EIGEN. Need to prov option(onnxruntime_BUILD_BENCHMARKS "Build ONNXRuntime micro-benchmarks" OFF) option(onnxruntime_USE_LLVM "Build TVM with LLVM" OFF) -cmake_dependent_option(onnxruntime_USE_FLASH_ATTENTION "Build flash attention kernel for scaled dot product attention" ON "NOT WIN32; onnxruntime_USE_CUDA" OFF) +cmake_dependent_option(onnxruntime_USE_FLASH_ATTENTION "Build flash attention kernel for scaled dot product attention" ON "onnxruntime_USE_CUDA" OFF) option(onnxruntime_USE_MEMORY_EFFICIENT_ATTENTION "Build memory efficient attention kernel for scaled dot product attention" ON) option(onnxruntime_BUILD_FOR_NATIVE_MACHINE "Enable this option for turning on optimization specific to this machine" OFF) @@ -734,6 +734,9 @@ if (onnxruntime_USE_CUDA) message( STATUS "Turn off flash attention since CUDA compiler version < 11.6") set(onnxruntime_USE_FLASH_ATTENTION OFF) set(onnxruntime_USE_MEMORY_EFFICIENT_ATTENTION OFF) + elseif(WIN32 AND CMAKE_CUDA_COMPILER_VERSION VERSION_LESS 12) + message( STATUS "Flash-Attention unsupported in Windows with CUDA compiler version < 12.0") + set(onnxruntime_USE_FLASH_ATTENTION OFF) endif() if (CMAKE_CUDA_COMPILER_VERSION VERSION_LESS 11.4) message( FATAL_ERROR "Failed build due to CUDA compiler version < 11.4") diff --git a/onnxruntime/contrib_ops/cuda/bert/attention.cc b/onnxruntime/contrib_ops/cuda/bert/attention.cc index 3e6edb162360d..d9907f09121d0 100644 --- a/onnxruntime/contrib_ops/cuda/bert/attention.cc +++ b/onnxruntime/contrib_ops/cuda/bert/attention.cc @@ -145,7 +145,7 @@ Status Attention::ComputeInternal(OpKernelContext* context) const { auto [num_splits, slse_accum_bytes, o_accum_bytes] = onnxruntime::flash::get_num_splits_and_buffer_sizes( parameters.batch_size, parameters.sequence_length, parameters.kv_sequence_length, parameters.num_heads, parameters.head_size, device_prop.multiProcessorCount); - parameters.num_splits = num_splits; + parameters.num_splits = static_cast(num_splits); softmax_lse_accum_bytes = slse_accum_bytes; out_accum_bytes = o_accum_bytes; } diff --git a/onnxruntime/contrib_ops/cuda/bert/attention_impl.cu b/onnxruntime/contrib_ops/cuda/bert/attention_impl.cu index 5c13bb731ce28..150079cdf157a 100644 --- a/onnxruntime/contrib_ops/cuda/bert/attention_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/attention_impl.cu @@ -334,6 +334,11 @@ Status FlashAttention( contrib::AttentionParameters& parameters, AttentionData& data, float scale) { + ORT_UNUSED_PARAMETER(device_prop); + ORT_UNUSED_PARAMETER(stream); + ORT_UNUSED_PARAMETER(parameters); + ORT_UNUSED_PARAMETER(data); + ORT_UNUSED_PARAMETER(scale); return ORT_MAKE_STATUS(ONNXRUNTIME, StatusCode::NOT_IMPLEMENTED, "flash attention does not support float tensor"); } #endif diff --git a/onnxruntime/contrib_ops/cuda/bert/flash_attention/alibi.h b/onnxruntime/contrib_ops/cuda/bert/flash_attention/alibi.h new file mode 100644 index 0000000000000..5d94190ecbeb9 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/bert/flash_attention/alibi.h @@ -0,0 +1,67 @@ +#include +#include +#include +#include +#include "utils.h" + +namespace onnxruntime { +namespace flash { + +using namespace cute; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct Alibi { + const float alibi_slope; + const int max_seqlen_k, max_seqlen_q; + + __forceinline__ __device__ Alibi(const float alibi_slope, const int max_seqlen_k, const int max_seqlen_q) + : alibi_slope(alibi_slope), max_seqlen_k(max_seqlen_k), max_seqlen_q(max_seqlen_q){}; + + template + __forceinline__ __device__ void apply_alibi(Tensor& tensor, + const int col_idx_offset_, + const int row_idx_offset, + const int warp_row_stride) { + // tensor has shape (ncol=(2, MMA_M), nrow=(2, MMA_N)) + static_assert(Layout::rank == 2, "Only support 2D Tensor"); + const int lane_id = threadIdx.x % 32; + const int col_idx_offset = col_idx_offset_ + (lane_id % 4) * 2; + if constexpr (Is_causal) { // Simpler, we add the same bias vector to all rows +#pragma unroll + for (int nj = 0; nj < size<1, 1>(tensor); ++nj) { + const int col_idx_base = col_idx_offset + nj * 8; +#pragma unroll + for (int j = 0; j < size<1, 0>(tensor); ++j) { + const int col_idx = col_idx_base + j; +#pragma unroll + for (int mi = 0; mi < size<0>(tensor); ++mi) { + tensor(mi, make_coord(j, nj)) += alibi_slope * col_idx; + } + } + } + } else { // Bias depends on both row_idx and col_idx +#pragma unroll + for (int mi = 0; mi < size<0, 1>(tensor); ++mi) { + const int row_idx_base = row_idx_offset + mi * warp_row_stride; +#pragma unroll + for (int i = 0; i < size<0, 0>(tensor); ++i) { + const int row_idx = row_idx_base + i * 8; +#pragma unroll + for (int nj = 0; nj < size<1, 1>(tensor); ++nj) { + const int col_idx_base = col_idx_offset + nj * 8; +#pragma unroll + for (int j = 0; j < size<1, 0>(tensor); ++j) { + const int col_idx = col_idx_base + j; + tensor(make_coord(i, mi), make_coord(j, nj)) -= alibi_slope * abs(row_idx + max_seqlen_k - max_seqlen_q - col_idx); + } + } + } + } + } + } +}; + +} // namespace flash +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/bert/flash_attention/block_info.h b/onnxruntime/contrib_ops/cuda/bert/flash_attention/block_info.h index 811b1be7d4315..dde6143153e8e 100644 --- a/onnxruntime/contrib_ops/cuda/bert/flash_attention/block_info.h +++ b/onnxruntime/contrib_ops/cuda/bert/flash_attention/block_info.h @@ -12,22 +12,36 @@ struct BlockInfo { template __device__ BlockInfo(const Params& params, const int bidb) : sum_s_q(!Varlen || params.cu_seqlens_q == nullptr ? -1 : params.cu_seqlens_q[bidb]), - sum_s_k(!Varlen || params.cu_seqlens_k == nullptr || !params.is_seqlens_k_cumulative ? -1 : params.cu_seqlens_k[bidb]), - actual_seqlen_q(!Varlen || params.cu_seqlens_q == nullptr ? params.seqlen_q : params.cu_seqlens_q[bidb + 1] - sum_s_q) + sum_s_k(!Varlen || params.cu_seqlens_k == nullptr || !params.is_seqlens_k_cumulative + ? -1 + : params.cu_seqlens_k[bidb]), + actual_seqlen_q(!Varlen || params.cu_seqlens_q == nullptr + ? params.seqlen_q + : params.cu_seqlens_q[bidb + 1] - sum_s_q) // If is_seqlens_k_cumulative, then seqlen_k is cu_seqlens_k[bidb + 1] - cu_seqlens_k[bidb]. // Otherwise it's cu_seqlens_k[bidb], i.e., we use cu_seqlens_k to store the sequence lengths of K. , - seqlen_k_cache(!Varlen || params.cu_seqlens_k == nullptr ? params.seqlen_k : (params.is_seqlens_k_cumulative ? params.cu_seqlens_k[bidb + 1] - sum_s_k : params.cu_seqlens_k[bidb])), - actual_seqlen_k(seqlen_k_cache + (params.knew_ptr == nullptr ? 0 : params.seqlen_knew)) { + seqlen_k_cache(!Varlen || params.cu_seqlens_k == nullptr + ? params.seqlen_k + : (params.is_seqlens_k_cumulative + ? params.cu_seqlens_k[bidb + 1] - sum_s_k + : params.cu_seqlens_k[bidb])), + actual_seqlen_k(params.seqused_k + ? params.seqused_k[bidb] + : seqlen_k_cache + (params.knew_ptr == nullptr ? 0 : params.seqlen_knew)) { } template - inline __device__ index_t q_offset(const index_t batch_stride, const index_t row_stride, const int bidb) const { + __forceinline__ __device__ + index_t + q_offset(const index_t batch_stride, const index_t row_stride, const int bidb) const { return sum_s_q == -1 ? bidb * batch_stride : uint32_t(sum_s_q) * row_stride; } template - inline __device__ index_t k_offset(const index_t batch_stride, const index_t row_stride, const int bidb) const { + __forceinline__ __device__ + index_t + k_offset(const index_t batch_stride, const index_t row_stride, const int bidb) const { return sum_s_k == -1 ? bidb * batch_stride : uint32_t(sum_s_k) * row_stride; } @@ -41,6 +55,5 @@ struct BlockInfo { //////////////////////////////////////////////////////////////////////////////////////////////////// -//////////////////////////////////////////////////////////////////////////////////////////////////// } // namespace flash } // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/bert/flash_attention/flash.h b/onnxruntime/contrib_ops/cuda/bert/flash_attention/flash.h index cbe536c6ce45a..0463d3795b446 100644 --- a/onnxruntime/contrib_ops/cuda/bert/flash_attention/flash.h +++ b/onnxruntime/contrib_ops/cuda/bert/flash_attention/flash.h @@ -16,7 +16,7 @@ constexpr int D_DIM = 2; //////////////////////////////////////////////////////////////////////////////////////////////////// struct Qkv_params { - using index_t = uint32_t; + using index_t = int64_t; // The QKV matrices. void* __restrict__ q_ptr = nullptr; void* __restrict__ k_ptr = nullptr; @@ -79,6 +79,9 @@ struct Flash_fwd_params : public Qkv_params { int* __restrict__ cu_seqlens_q = nullptr; int* __restrict__ cu_seqlens_k = nullptr; + // If provided, the actual length of each k sequence. + int* __restrict__ seqused_k = nullptr; + int* __restrict__ blockmask = nullptr; // The K_new and V_new matrices. @@ -100,6 +103,11 @@ struct Flash_fwd_params : public Qkv_params { // The indices to index into the KV cache. int* __restrict__ cache_batch_idx = nullptr; + // Paged KV cache + int* __restrict__ block_table = nullptr; + index_t block_table_batch_stride = 0; + int page_block_size = 0; + // Local window size int window_size_left = -1; int window_size_right = -1; @@ -115,6 +123,9 @@ struct Flash_fwd_params : public Qkv_params { int num_splits = 0; // For split-KV version + void* __restrict__ alibi_slopes_ptr = nullptr; + index_t alibi_slopes_batch_stride = 0; + const cudaDeviceProp* dprops = nullptr; }; diff --git a/onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_api.cc b/onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_api.cc index 0f58a74c4d2fd..e04cdf369c6ac 100644 --- a/onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_api.cc +++ b/onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_api.cc @@ -31,6 +31,7 @@ void set_params_fprop(Flash_fwd_params& params, void* out, void* cu_seqlens_q_d, void* cu_seqlens_k_d, + void* seqused_k, void* p_d, void* softmax_lse_d, float softmax_scale, @@ -82,6 +83,7 @@ void set_params_fprop(Flash_fwd_params& params, params.cu_seqlens_q = static_cast(cu_seqlens_q_d); params.cu_seqlens_k = static_cast(cu_seqlens_k_d); + params.seqused_k = static_cast(seqused_k); // P = softmax(QK^T) params.p_ptr = p_d; @@ -123,24 +125,25 @@ void set_params_fprop(Flash_fwd_params& params, params.is_seqlens_k_cumulative = true; } -size_t get_softmax_lse_size(int seqlen, int batch_size, int num_heads) { +size_t get_softmax_lse_size(size_t seqlen, size_t batch_size, size_t num_heads) { size_t bytes = sizeof(float) * batch_size * num_heads * seqlen; return bytes; } -size_t get_softmax_lse_accum_size(int num_splits, int batch_size, int num_heads, int seqlen_q) { +size_t get_softmax_lse_accum_size(size_t num_splits, size_t batch_size, size_t num_heads, size_t seqlen_q) { size_t bytes = sizeof(float) * num_splits * batch_size * seqlen_q * num_heads; return bytes; } -size_t get_out_accum_size(int num_splits, int batch_size, int num_heads, int seqlen_q, int head_size_rounded) { +size_t get_out_accum_size(size_t num_splits, size_t batch_size, size_t num_heads, + size_t seqlen_q, size_t head_size_rounded) { size_t bytes = sizeof(float) * num_splits * batch_size * seqlen_q * num_heads * head_size_rounded; return bytes; } void run_mha_fwd(Flash_fwd_params& params, cudaStream_t stream, bool force_split_kernel = false) { FP16_SWITCH(!params.is_bf16, [&] { - FWD_HEADDIM_SWITCH(params.d, [&] { + HEADDIM_SWITCH(params.d, [&] { if (params.num_splits <= 1 && !force_split_kernel) { // If we don't set it num_splits == 0 run_mha_fwd_(params, stream); } else { @@ -156,15 +159,15 @@ void run_mha_fwd(Flash_fwd_params& params, cudaStream_t stream, bool force_split // splits as that would incur more HBM reads/writes. // So we find the best efficiency, then find the smallest number of splits that gets 85% // of the best efficiency. -int num_splits_heuristic(int batch_size, int seqlen_q, int seqlen_k, int num_heads, int head_size, int num_SMs, - int max_splits) { +size_t num_splits_heuristic(size_t batch_size, size_t seqlen_q, size_t seqlen_k, size_t num_heads, + size_t head_size, size_t num_SMs, size_t max_splits) { // This needs to match with run_mha_fwd_splitkv_dispatch - const int block_n = head_size <= 64 ? 256 : (head_size <= 128 ? 128 : 64); - const int num_n_blocks = (seqlen_k + block_n - 1) / block_n; + const size_t block_n = head_size <= 64 ? 256 : (head_size <= 128 ? 128 : 64); + const size_t num_n_blocks = (seqlen_k + block_n - 1) / block_n; // Technically kBlockM = 64 only for the splitKV kernels, not the standard kernel. // In any case we don't expect seqlen_q to be larger than 64 for inference. - const int num_m_blocks = (seqlen_q + 64 - 1) / 64; - int batch_nheads_mblocks = batch_size * num_heads * num_m_blocks; + const size_t num_m_blocks = (seqlen_q + 64 - 1) / 64; + size_t batch_nheads_mblocks = batch_size * num_heads * num_m_blocks; // If we have enough to almost fill the SMs, then just use 1 split if (batch_nheads_mblocks >= 0.8f * num_SMs) { return 1; @@ -173,15 +176,15 @@ int num_splits_heuristic(int batch_size, int seqlen_q, int seqlen_k, int num_hea float max_efficiency = 0.f; std::vector efficiency; efficiency.reserve(max_splits); - auto ceildiv = [](int a, int b) { return (a + b - 1) / b; }; + auto ceildiv = [](size_t a, size_t b) { return (a + b - 1) / b; }; // Some splits are not eligible. For example, if we have 64 blocks and choose 11 splits, // we'll have 6 * 10 + 4 blocks. If we choose 12 splits, we'll have 6 * 11 + (-2) blocks // (i.e. it's 11 splits anyway). // So we check if the number of blocks per split is the same as the previous num_splits. - auto is_split_eligible = [&ceildiv, &num_n_blocks](int num_splits) { + auto is_split_eligible = [&ceildiv, &num_n_blocks](size_t num_splits) { return num_splits == 1 || ceildiv(num_n_blocks, num_splits) != ceildiv(num_n_blocks, num_splits - 1); }; - for (int num_splits = 1; num_splits <= max_splits; num_splits++) { + for (size_t num_splits = 1; num_splits <= max_splits; num_splits++) { if (!is_split_eligible(num_splits)) { efficiency.push_back(0.f); } else { @@ -194,7 +197,7 @@ int num_splits_heuristic(int batch_size, int seqlen_q, int seqlen_k, int num_hea efficiency.push_back(eff); } } - for (int num_splits = 1; num_splits <= max_splits; num_splits++) { + for (size_t num_splits = 1; num_splits <= max_splits; num_splits++) { if (!is_split_eligible(num_splits)) { continue; } @@ -207,25 +210,39 @@ int num_splits_heuristic(int batch_size, int seqlen_q, int seqlen_k, int num_hea } // Returns (num_splits, softmax_lse_accum bytes, out_accum bytes) -std::tuple get_num_splits_and_buffer_sizes(int batch_size, int seqlen_q, int seqlen_k, int num_heads, - int head_size, int num_SMs) { - int max_splits = 128; +std::tuple get_num_splits_and_buffer_sizes(size_t batch_size, size_t seqlen_q, size_t seqlen_k, + size_t num_heads, size_t head_size, size_t num_SMs) { + size_t max_splits = 128; // split kv buffers - int num_splits = num_splits_heuristic(batch_size, seqlen_q, seqlen_k, num_heads, head_size, - num_SMs, max_splits); + size_t num_splits = num_splits_heuristic(batch_size, seqlen_q, seqlen_k, num_heads, head_size, + num_SMs, max_splits); if (num_splits > 1) { // softmax_lse_accum buffer - int softmax_lse_accum_bytes = get_softmax_lse_accum_size(num_splits, batch_size, num_heads, seqlen_q); + size_t softmax_lse_accum_bytes = get_softmax_lse_accum_size(num_splits, batch_size, num_heads, seqlen_q); // out_accum buffer - auto round_multiple = [](int x, int m) { return (x + m - 1) / m * m; }; - const int head_size_rounded = round_multiple(head_size, 32); - int out_accum_bytes = get_out_accum_size(num_splits, batch_size, num_heads, seqlen_q, head_size_rounded); + auto round_multiple = [](size_t x, size_t m) { return (x + m - 1) / m * m; }; + const size_t head_size_rounded = round_multiple(head_size, 32); + size_t out_accum_bytes = get_out_accum_size(num_splits, batch_size, num_heads, seqlen_q, head_size_rounded); return {num_splits, softmax_lse_accum_bytes, out_accum_bytes}; } else { return {0, 0, 0}; } } +// void set_params_alibi(Flash_fwd_params ¶ms, void* alibi_slopes, int batch_size, int num_heads){ +// if (alibi_slopes != nullptr) { +// // TORCH_CHECK(alibi_slopes.dtype() == torch::kFloat32, "ALiBi slopes must have dtype fp32"); +// // CHECK_DEVICE(alibi_slopes); +// // TORCH_CHECK(alibi_slopes.stride(-1) == 1, "ALiBi slopes tensor must have contiguous last dimension"); +// // TORCH_CHECK(alibi_slopes.sizes() == torch::IntArrayRef({num_heads}) +// || alibi_slopes.sizes() == torch::IntArrayRef({batch_size, num_heads})); +// params.alibi_slopes_ptr = alibi_slopes; +// params.alibi_slopes_batch_stride = alibi_slopes.dim() == 2 ? num_heads : 0; // TODO: flag for bool +// } else { +// params.alibi_slopes_ptr = nullptr; +// } +// } + Status mha_fwd(const cudaDeviceProp& dprops, cudaStream_t stream, void* q, // batch_size x seqlen_q x num_heads x head_size @@ -262,7 +279,8 @@ Status mha_fwd(const cudaDeviceProp& dprops, q, k, v, out, /*cu_seqlens_q*/ nullptr, /*cu_seqlens_k*/ nullptr, - nullptr, + /*seqused_k=*/nullptr, + /*p_ptr=*/nullptr, softmax_lse, softmax_scale, is_causal, @@ -289,6 +307,8 @@ Status mha_fwd(const cudaDeviceProp& dprops, params.oaccum_ptr = nullptr; } + params.alibi_slopes_ptr = nullptr; + run_mha_fwd(params, stream); return Status::OK(); } @@ -301,6 +321,8 @@ Status mha_varlen_fwd(const cudaDeviceProp& dprops, void* out, // half (total_q, num_heads, head_size) int* cu_seqlens_q, // int (batch_size + 1) int* cu_seqlens_k, // int (batch_size + 1) + void* seqused_k, // batch_size; If given, use this many elements of each batch element's keys. + int* block_table, // batch_size x max_num_blocks_per_seq void* softmax_lse, // float (batch_size, num_heads, max_seqlen_q) int batch_size, int num_heads, @@ -310,11 +332,14 @@ Status mha_varlen_fwd(const cudaDeviceProp& dprops, int max_seqlen_k, float softmax_scale, bool is_causal, - bool is_bf16) { + bool is_bf16, + int max_num_blocks_per_seq, + int page_block_size) { auto round_multiple = [](int x, int m) { return (x + m - 1) / m * m; }; const int head_size_rounded = round_multiple(head_size, 32); const int seqlen_q_rounded = round_multiple(max_seqlen_q, 128); const int seqlen_k_rounded = round_multiple(max_seqlen_k, 128); + const bool paged_KV = block_table != nullptr; Flash_fwd_params params; set_params_fprop(params, @@ -326,7 +351,8 @@ Status mha_varlen_fwd(const cudaDeviceProp& dprops, q, k, v, out, cu_seqlens_q, cu_seqlens_k, - nullptr, + seqused_k, + /*p_ptr=*/nullptr, softmax_lse, softmax_scale, is_causal, @@ -340,11 +366,25 @@ Status mha_varlen_fwd(const cudaDeviceProp& dprops, params.oaccum_ptr = nullptr; params.knew_ptr = nullptr; params.vnew_ptr = nullptr; + params.alibi_slopes_ptr = nullptr; + if (paged_KV) { + params.block_table = block_table; // TODO(aciddelgado): cast to int pointer + params.block_table_batch_stride = max_num_blocks_per_seq; + // params.num_blocks = num_blocks; + params.page_block_size = page_block_size; + params.k_batch_stride = page_block_size * num_heads_k * head_size; + params.v_batch_stride = page_block_size * num_heads_k * head_size; + } else { + params.block_table = nullptr; + params.block_table_batch_stride = 0; + // params.num_blocks = 0; + params.page_block_size = 1; + } run_mha_fwd(params, stream); return Status::OK(); } -bool is_supported(const cudaDeviceProp& dprops, int head_size, int num_heads, int num_heads_k) { +bool is_supported(const cudaDeviceProp& dprops, size_t head_size, size_t num_heads, size_t num_heads_k) { bool is_sm8x = dprops.major == 8 && dprops.minor >= 0; bool is_sm90 = dprops.major == 9 && dprops.minor == 0; return (is_sm8x || is_sm90) && (head_size % 8 == 0) && (head_size <= 256) && (num_heads % num_heads_k == 0); @@ -364,6 +404,7 @@ Status mha_fwd_kvcache(const cudaDeviceProp& dprops, void* seqlens_k_, // batch_size void* rotary_cos, // seqlen_ro x (rotary_dim / 2) void* rotary_sin, // seqlen_ro x (rotary_dim / 2) + int* block_table, // batch_size x max_num_blocks_per_seq int batch_size, int num_heads, int num_heads_k, @@ -381,11 +422,14 @@ Status mha_fwd_kvcache(const cudaDeviceProp& dprops, void* out_accum, // num_splits x batch_size x seqlen_q x num_heads x head_size_rounded int local_window_size, bool is_rotary_interleaved, - bool is_packed_qkv) { + bool is_packed_qkv, + int max_num_blocks_per_seq, + int page_block_size) { auto round_multiple = [](int x, int m) { return (x + m - 1) / m * m; }; const int head_size_rounded = round_multiple(head_size, 32); const int seqlen_q_rounded = round_multiple(seqlen_q, 128); const int seqlen_k_rounded = round_multiple(seqlen_k, 128); + const bool paged_KV = block_table != nullptr; // In kv-cache case, seqlen_k_max as kv sequence length Flash_fwd_params params; @@ -398,6 +442,7 @@ Status mha_fwd_kvcache(const cudaDeviceProp& dprops, q, kcache, vcache, out, /*cu_seqlens_q_d=*/nullptr, /*cu_seqlens_k_d=*/nullptr, + /*seqused_k=*/nullptr, /*p_ptr=*/nullptr, softmax_lse, softmax_scale, @@ -461,6 +506,21 @@ Status mha_fwd_kvcache(const cudaDeviceProp& dprops, params.oaccum_ptr = nullptr; } + params.alibi_slopes_ptr = nullptr; + if (paged_KV) { + params.block_table = block_table; // TODO(aciddelgado): cast to int pointer + params.block_table_batch_stride = max_num_blocks_per_seq; + // params.num_blocks = num_blocks; + params.page_block_size = page_block_size; + params.k_batch_stride = page_block_size * num_heads_k * head_size; + params.v_batch_stride = page_block_size * num_heads_k * head_size; + } else { + params.block_table = nullptr; + params.block_table_batch_stride = 0; + // params.num_blocks = 0; + params.page_block_size = 1; + } + // Only split kernel supports appending to KV cache run_mha_fwd(params, stream, /*force_split_kernel=*/k_new != nullptr); diff --git a/onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_api.h b/onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_api.h index 24891bcc4d499..4c59561449851 100644 --- a/onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_api.h +++ b/onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_api.h @@ -66,6 +66,8 @@ Status mha_varlen_fwd(const cudaDeviceProp& dprops, void* out, // half (total_q, num_heads, v_head_size) int* cu_seqlens_q, // int (batch_size + 1) int* cu_seqlens_k, // int (batch_size + 1) + void* seqused_k, // batch_size; If given, only this many elements of each batch element's keys are used. + int* block_table, // batch_size x max_num_blocks_per_seq void* softmax_lse, // float (batch_size, num_heads, max_seqlen_q) int batch_size, int num_heads, @@ -75,7 +77,9 @@ Status mha_varlen_fwd(const cudaDeviceProp& dprops, int max_seqlen_k, float softmax_scale, bool is_causal, - bool is_bf16); + bool is_bf16, + int max_num_blocks_per_seq = 0, + int page_block_size = 1); Status mha_fwd_kvcache(const cudaDeviceProp& dprops, cudaStream_t stream, @@ -87,8 +91,9 @@ Status mha_fwd_kvcache(const cudaDeviceProp& dprops, void* out, // batch_size x seqlen_q x num_heads x head_size void* softmax_lse, // batch_size x num_heads x seqlen_q void* seqlens_k_, // batch_size - void* rotary_sin, // seqlen_ro x (rotary_dim / 2) void* rotary_cos, // seqlen_ro x (rotary_dim / 2) + void* rotary_sin, // seqlen_ro x (rotary_dim / 2) + int* block_table, // batch_size x max_num_blocks_per_seq int batch_size, int num_heads, int num_heads_k, @@ -106,14 +111,16 @@ Status mha_fwd_kvcache(const cudaDeviceProp& dprops, void* out_accum = nullptr, // num_splits x batch_size x seqlen_q x num_heads x head_size_rounded int local_window_size = -1, bool is_rotary_interleaved = false, - bool is_packed_qkv = false); + bool is_packed_qkv = false, + int max_num_blocks_per_seq = 0, + int page_block_size = 1); -size_t get_softmax_lse_size(int max_seqlen_q, int batch_size, int num_heads); +size_t get_softmax_lse_size(size_t max_seqlen_q, size_t batch_size, size_t num_heads); -std::tuple get_num_splits_and_buffer_sizes(int batch_size, int seqlen_q, int seqlen_k, int num_heads, - int head_size, int num_SMs); +std::tuple get_num_splits_and_buffer_sizes(size_t batch_size, size_t seqlen_q, size_t seqlen_k, size_t num_heads, + size_t head_size, size_t num_SMs); -bool is_supported(const cudaDeviceProp& dprops, int head_size, int num_heads, int num_heads_k); +bool is_supported(const cudaDeviceProp& dprops, size_t head_size, size_t num_heads, size_t num_heads_k); } // namespace flash } // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_fwd_hdim128_bf16_sm80.cu b/onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_fwd_hdim128_bf16_sm80.cu index 431eb2bd69def..1ef1ce251ecba 100644 --- a/onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_fwd_hdim128_bf16_sm80.cu +++ b/onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_fwd_hdim128_bf16_sm80.cu @@ -8,9 +8,9 @@ namespace onnxruntime { namespace flash { -template<> +template <> void run_mha_fwd_(Flash_fwd_params& params, cudaStream_t stream) { - run_mha_fwd_hdim128(params, stream); + run_mha_fwd_hdim128(params, stream); } } // namespace flash diff --git a/onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_fwd_hdim160_bf16_sm80.cu b/onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_fwd_hdim160_bf16_sm80.cu index 0cb48272dec3f..52ff792c6edcb 100644 --- a/onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_fwd_hdim160_bf16_sm80.cu +++ b/onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_fwd_hdim160_bf16_sm80.cu @@ -8,9 +8,9 @@ namespace onnxruntime { namespace flash { -template<> +template <> void run_mha_fwd_(Flash_fwd_params& params, cudaStream_t stream) { - run_mha_fwd_hdim160(params, stream); + run_mha_fwd_hdim160(params, stream); } } // namespace flash diff --git a/onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_fwd_hdim192_bf16_sm80.cu b/onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_fwd_hdim192_bf16_sm80.cu index 142e922f71031..3bdc5e4b0443f 100644 --- a/onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_fwd_hdim192_bf16_sm80.cu +++ b/onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_fwd_hdim192_bf16_sm80.cu @@ -8,9 +8,9 @@ namespace onnxruntime { namespace flash { -template<> +template <> void run_mha_fwd_(Flash_fwd_params& params, cudaStream_t stream) { - run_mha_fwd_hdim192(params, stream); + run_mha_fwd_hdim192(params, stream); } } // namespace flash diff --git a/onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_fwd_hdim224_bf16_sm80.cu b/onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_fwd_hdim224_bf16_sm80.cu index 2142b1c343110..e4972875d0512 100644 --- a/onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_fwd_hdim224_bf16_sm80.cu +++ b/onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_fwd_hdim224_bf16_sm80.cu @@ -8,9 +8,9 @@ namespace onnxruntime { namespace flash { -template<> +template <> void run_mha_fwd_(Flash_fwd_params& params, cudaStream_t stream) { - run_mha_fwd_hdim224(params, stream); + run_mha_fwd_hdim224(params, stream); } } // namespace flash diff --git a/onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_fwd_hdim256_bf16_sm80.cu b/onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_fwd_hdim256_bf16_sm80.cu index 751363184e23a..59568b0bb03ce 100644 --- a/onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_fwd_hdim256_bf16_sm80.cu +++ b/onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_fwd_hdim256_bf16_sm80.cu @@ -8,9 +8,9 @@ namespace onnxruntime { namespace flash { -template<> +template <> void run_mha_fwd_(Flash_fwd_params& params, cudaStream_t stream) { - run_mha_fwd_hdim256(params, stream); + run_mha_fwd_hdim256(params, stream); } } // namespace flash diff --git a/onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_fwd_hdim32_bf16_sm80.cu b/onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_fwd_hdim32_bf16_sm80.cu index ebf0236435971..ad3d4df7dfc85 100644 --- a/onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_fwd_hdim32_bf16_sm80.cu +++ b/onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_fwd_hdim32_bf16_sm80.cu @@ -8,9 +8,9 @@ namespace onnxruntime { namespace flash { -template<> +template <> void run_mha_fwd_(Flash_fwd_params& params, cudaStream_t stream) { - run_mha_fwd_hdim32(params, stream); + run_mha_fwd_hdim32(params, stream); } } // namespace flash diff --git a/onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_fwd_hdim64_bf16_sm80.cu b/onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_fwd_hdim64_bf16_sm80.cu index 166bb2a0072f4..006416458c91b 100644 --- a/onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_fwd_hdim64_bf16_sm80.cu +++ b/onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_fwd_hdim64_bf16_sm80.cu @@ -8,9 +8,9 @@ namespace onnxruntime { namespace flash { -template<> +template <> void run_mha_fwd_(Flash_fwd_params& params, cudaStream_t stream) { - run_mha_fwd_hdim64(params, stream); + run_mha_fwd_hdim64(params, stream); } } // namespace flash diff --git a/onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_fwd_hdim96_bf16_sm80.cu b/onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_fwd_hdim96_bf16_sm80.cu index c8760b8168db6..d5a273a3f4163 100644 --- a/onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_fwd_hdim96_bf16_sm80.cu +++ b/onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_fwd_hdim96_bf16_sm80.cu @@ -8,9 +8,9 @@ namespace onnxruntime { namespace flash { -template<> +template <> void run_mha_fwd_(Flash_fwd_params& params, cudaStream_t stream) { - run_mha_fwd_hdim96(params, stream); + run_mha_fwd_hdim96(params, stream); } } // namespace flash diff --git a/onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_fwd_kernel.h b/onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_fwd_kernel.h index 1fac03882b4b1..1c8a93674a80b 100644 --- a/onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_fwd_kernel.h +++ b/onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_fwd_kernel.h @@ -7,20 +7,26 @@ #pragma GCC diagnostic push #pragma GCC diagnostic ignored "-Wunused-variable" #pragma GCC diagnostic ignored "-Wunused-but-set-variable" +#elif defined(_MSC_VER) +#pragma warning(push) +#pragma warning(disable : 4267) +#pragma warning(disable : 4100) +#pragma warning(disable : 4101) // equivalent to GCC's -Wunused-variable +#pragma warning(disable : 4189) // equivalent to GCC's -Wunused-but-set-variable #endif -#include #include #include #include #include -#include #include "contrib_ops/cuda/bert/flash_attention/block_info.h" #include "contrib_ops/cuda/bert/flash_attention/kernel_traits.h" #include "contrib_ops/cuda/bert/flash_attention/utils.h" #include "contrib_ops/cuda/bert/flash_attention/softmax.h" +#include "contrib_ops/cuda/bert/flash_attention/mask.h" +#include "contrib_ops/cuda/bert/flash_attention/rotary.h" namespace onnxruntime { namespace flash { @@ -28,60 +34,7 @@ using namespace cute; //////////////////////////////////////////////////////////////////////////////////////////////////// -template -inline __device__ void softmax_rescale_o(Tensor0& scores, Tensor1& scores_max, Tensor1& scores_sum, - Tensor2& acc_o, float softmax_scale_log2) { - if (Is_first) { - flash::template reduce_max(scores, scores_max); - flash::scale_apply_exp2(scores, scores_max, softmax_scale_log2); - flash::reduce_sum(scores, scores_sum); - } else { - cute::Tensor scores_max_prev = make_fragment_like(scores_max); - cute::copy(scores_max, scores_max_prev); - flash::template reduce_max(scores, scores_max); - // Reshape acc_o from (MMA=4, MMA_M, MMA_K) to (nrow=(2, MMA_M), ncol=(2, MMA_K)) - cute::Tensor acc_o_rowcol = make_tensor(acc_o.data(), flash::convert_layout_acc_rowcol(acc_o.layout())); -#pragma unroll - for (int mi = 0; mi < cute::size(scores_max); ++mi) { - float scores_max_cur = !Check_inf - ? scores_max(mi) - : (scores_max(mi) == -INFINITY ? 0.0f : scores_max(mi)); - float scores_scale = exp2f((scores_max_prev(mi) - scores_max_cur) * softmax_scale_log2); - scores_sum(mi) *= scores_scale; -#pragma unroll - for (int ni = 0; ni < cute::size<1>(acc_o_rowcol); ++ni) { - acc_o_rowcol(mi, ni) *= scores_scale; - } - } - flash::scale_apply_exp2(scores, scores_max, softmax_scale_log2); - cute::Tensor scores_sum_cur = make_fragment_like(scores_sum); - flash::reduce_sum(scores, scores_sum_cur); -#pragma unroll - for (int mi = 0; mi < cute::size(scores_sum); ++mi) { - scores_sum(mi) += scores_sum_cur(mi); - } - } -}; - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -template -inline __device__ void write_softmax_to_gmem( - cute::Tensor const& tOrP, cute::Tensor& tPgP, TiledCopy gmem_tiled_copy_P) { - // Reshape tOrP from (8, MMA_M, MMA_N) to (8, MMA_M * MMA_N) - cute::Layout l = tOrP.layout(); - cute::Tensor tPrP = make_tensor(tOrP.data(), make_layout(get<0>(l), make_layout(get<1>(l), get<2>(l)))); - CUTE_STATIC_ASSERT_V(cute::size<2>(tPgP) == _1{}); - CUTE_STATIC_ASSERT_V(cute::size<1>(tPrP) == cute::size<1>(tPgP)); -#pragma unroll - for (int mi = 0; mi < cute::size<1>(tPrP); ++mi) { - cute::copy(gmem_tiled_copy_P, tPrP(_, mi), tPgP(_, mi, 0)); - } -}; - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -template +template inline __device__ void compute_attn_1rowblock(const Params& params, const int bidb, const int bidh, const int m_block) { using Element = typename Kernel_traits::Element; using ElementAccum = typename Kernel_traits::ElementAccum; @@ -97,25 +50,30 @@ inline __device__ void compute_attn_1rowblock(const Params& params, const int bi constexpr int kBlockN = Kernel_traits::kBlockN; constexpr int kHeadDim = Kernel_traits::kHeadDim; constexpr int kNWarps = Kernel_traits::kNWarps; + // constexpr int MMA_M = kBlockM / decltype(cute::size<0>(typename Kernel_traits::TiledMma::TiledShape_MNK{}))::value; const BlockInfo binfo(params, bidb); - if (m_block * kBlockM >= binfo.actual_seqlen_q || binfo.actual_seqlen_k == 0) return; + if (m_block * kBlockM >= binfo.actual_seqlen_q) return; - const int n_block_min = !Is_local ? 0 : std::max(0, (m_block * kBlockM + binfo.actual_seqlen_k - binfo.actual_seqlen_q - params.window_size_left) / kBlockN); + const int n_block_min = !Is_local + ? 0 + : std::max(0, (m_block * kBlockM + binfo.actual_seqlen_k - binfo.actual_seqlen_q - params.window_size_left) / kBlockN); int n_block_max = cute::ceil_div(binfo.actual_seqlen_k, kBlockN); if (Is_causal || Is_local) { n_block_max = std::min(n_block_max, cute::ceil_div((m_block + 1) * kBlockM + binfo.actual_seqlen_k - binfo.actual_seqlen_q + params.window_size_right, kBlockN)); - // We exit early and write 0 to gO and gLSE. + // We exit early and write 0 to gO and gLSE. This also covers the case where actual_seqlen_k == 0. // Otherwise we might read OOB elements from gK and gV. - if (n_block_max <= n_block_min) { - const index_t row_offset_o = binfo.q_offset(params.o_batch_stride, params.o_row_stride, bidb) + m_block * kBlockM * params.o_row_stride + bidh * params.o_head_stride; - const index_t row_offset_lse = (bidb * params.h + bidh) * params.seqlen_q + m_block * kBlockM; - Tensor gO = make_tensor(make_gmem_ptr(reinterpret_cast(params.o_ptr) + row_offset_o), - Shape, Int>{}, - make_stride(params.o_row_stride, _1{})); - Tensor gLSE = make_tensor(make_gmem_ptr(reinterpret_cast(params.softmax_lse_ptr) + row_offset_lse), - Shape>{}, Stride<_1>{}); + if ((Is_causal || Is_local || !Is_even_MN) && n_block_max <= n_block_min) { + Tensor mO = make_tensor(make_gmem_ptr(reinterpret_cast(params.o_ptr) + binfo.q_offset(params.o_batch_stride, params.o_row_stride, bidb)), + make_shape(binfo.actual_seqlen_q, params.h, params.d), + make_stride(params.o_row_stride, params.o_head_stride, _1{})); + Tensor gO = local_tile(mO(_, bidh, _), Shape, Int>{}, + make_coord(m_block, 0)); // (kBlockM, kHeadDim) + Tensor mLSE = make_tensor(make_gmem_ptr(reinterpret_cast(params.softmax_lse_ptr)), + make_shape(params.b, params.h, params.seqlen_q), + make_stride(params.h * params.seqlen_q, params.seqlen_q, _1{})); + Tensor gLSE = local_tile(mLSE(bidb, bidh, _), Shape>{}, make_coord(m_block)); typename Kernel_traits::GmemTiledCopyO gmem_tiled_copy_O; auto gmem_thr_copy_O = gmem_tiled_copy_O.get_thread_slice(tidx); @@ -151,53 +109,55 @@ inline __device__ void compute_attn_1rowblock(const Params& params, const int bi // that needs masking when we read K and V from global memory. Moreover, iterating in reverse // might save us 1 register (we just need n_block instead of both n_block and n_block_max). - const index_t row_offset_q = binfo.q_offset(params.q_batch_stride, params.q_row_stride, bidb) + m_block * kBlockM * params.q_row_stride + bidh * params.q_head_stride; - // We move K and V to the last block. - const index_t row_offset_k = binfo.k_offset(params.k_batch_stride, params.k_row_stride, bidb) + (n_block_max - 1) * kBlockN * params.k_row_stride + (bidh / params.h_h_k_ratio) * params.k_head_stride; - const index_t row_offset_v = binfo.k_offset(params.v_batch_stride, params.v_row_stride, bidb) + (n_block_max - 1) * kBlockN * params.v_row_stride + (bidh / params.h_h_k_ratio) * params.v_head_stride; - const index_t row_offset_p = ((bidb * params.h + bidh) * params.seqlen_q_rounded + m_block * kBlockM) * params.seqlen_k_rounded + (n_block_max - 1) * kBlockN; - cute::Tensor gQ = make_tensor(make_gmem_ptr(reinterpret_cast(params.q_ptr) + row_offset_q), - cute::Shape, cute::Int>{}, - make_stride(params.q_row_stride, _1{})); - cute::Tensor gK = make_tensor(make_gmem_ptr(reinterpret_cast(params.k_ptr) + row_offset_k), - cute::Shape, cute::Int>{}, - make_stride(params.k_row_stride, _1{})); - cute::Tensor gV = make_tensor(make_gmem_ptr(reinterpret_cast(params.v_ptr) + row_offset_v), - cute::Shape, cute::Int>{}, - make_stride(params.v_row_stride, _1{})); - cute::Tensor gP = make_tensor(make_gmem_ptr(reinterpret_cast(params.p_ptr) + row_offset_p), - cute::Shape, cute::Int>{}, - make_stride(params.seqlen_k_rounded, _1{})); - - cute::Tensor sQ = make_tensor(make_smem_ptr(reinterpret_cast(smem_)), - typename Kernel_traits::SmemLayoutQ{}); + const index_t row_offset_p = ((bidb * params.h + bidh) * params.seqlen_q_rounded + m_block * kBlockM) * params.seqlen_k_rounded + (n_block_max - 1) * kBlockN; // We move K and V to the last block. + + Tensor mQ = make_tensor(make_gmem_ptr(reinterpret_cast(params.q_ptr) + binfo.q_offset(params.q_batch_stride, params.q_row_stride, bidb)), + make_shape(binfo.actual_seqlen_q, params.h, params.d), + make_stride(params.q_row_stride, params.q_head_stride, _1{})); + Tensor gQ = local_tile(mQ(_, bidh, _), Shape, Int>{}, + make_coord(m_block, 0)); // (kBlockM, kHeadDim) + Tensor mK = make_tensor(make_gmem_ptr(reinterpret_cast(params.k_ptr) + binfo.k_offset(params.k_batch_stride, params.k_row_stride, bidb)), + make_shape(binfo.actual_seqlen_k, params.h_k, params.d), + make_stride(params.k_row_stride, params.k_head_stride, _1{})); + Tensor gK = local_tile(mK(_, bidh / params.h_h_k_ratio, _), Shape, Int>{}, + make_coord(_, 0)); // (kBlockN, kHeadDim, nblocksN) + Tensor mV = make_tensor(make_gmem_ptr(reinterpret_cast(params.v_ptr) + binfo.k_offset(params.v_batch_stride, params.v_row_stride, bidb)), + make_shape(binfo.actual_seqlen_k, params.h_k, params.d), + make_stride(params.v_row_stride, params.v_head_stride, _1{})); + Tensor gV = local_tile(mV(_, bidh / params.h_h_k_ratio, _), Shape, Int>{}, + make_coord(_, 0)); // (kBlockN, kHeadDim, nblocksN) + Tensor gP = make_tensor(make_gmem_ptr(reinterpret_cast(params.p_ptr) + row_offset_p), + Shape, Int>{}, + make_stride(params.seqlen_k_rounded, _1{})); + + Tensor sQ = make_tensor(make_smem_ptr(reinterpret_cast(smem_)), + typename Kernel_traits::SmemLayoutQ{}); // Careful we're using the same smem for sQ and sK | sV if Share_Q_K_smem; - cute::Tensor sK = make_tensor(sQ.data() + (Kernel_traits::Share_Q_K_smem ? 0 : cute::size(sQ)), - typename Kernel_traits::SmemLayoutKV{}); - cute::Tensor sV = make_tensor(sK.data() + cute::size(sK), typename Kernel_traits::SmemLayoutKV{}); - cute::Tensor sVt = make_tensor(sV.data(), typename Kernel_traits::SmemLayoutVtransposed{}); - cute::Tensor sVtNoSwizzle = make_tensor(sV.data(), typename Kernel_traits::SmemLayoutVtransposedNoSwizzle{}); + Tensor sK = make_tensor(sQ.data() + (Kernel_traits::Share_Q_K_smem ? 0 : size(sQ)), + typename Kernel_traits::SmemLayoutKV{}); + Tensor sV = make_tensor(sK.data() + size(sK), typename Kernel_traits::SmemLayoutKV{}); + Tensor sVt = make_tensor(sV.data(), typename Kernel_traits::SmemLayoutVtransposed{}); + Tensor sVtNoSwizzle = make_tensor(sV.data(), typename Kernel_traits::SmemLayoutVtransposedNoSwizzle{}); typename Kernel_traits::GmemTiledCopyQKV gmem_tiled_copy_QKV; auto gmem_thr_copy_QKV = gmem_tiled_copy_QKV.get_thread_slice(tidx); - typename Kernel_traits::GmemTiledCopyP gmem_tiled_copy_P; - auto gmem_thr_copy_P = gmem_tiled_copy_P.get_thread_slice(tidx); - cute::Tensor tQgQ = gmem_thr_copy_QKV.partition_S(gQ); - cute::Tensor tQsQ = gmem_thr_copy_QKV.partition_D(sQ); - cute::Tensor tKgK = gmem_thr_copy_QKV.partition_S(gK); // (KCPY, KCPY_N, KCPY_K) - cute::Tensor tKsK = gmem_thr_copy_QKV.partition_D(sK); - cute::Tensor tVgV = gmem_thr_copy_QKV.partition_S(gV); // (VCPY, VCPY_N, VCPY_K) - cute::Tensor tVsV = gmem_thr_copy_QKV.partition_D(sV); - cute::Tensor tPgP = gmem_thr_copy_P.partition_D(gP); + Tensor tQgQ = gmem_thr_copy_QKV.partition_S(gQ); + Tensor tQsQ = gmem_thr_copy_QKV.partition_D(sQ); + Tensor tKgK = gmem_thr_copy_QKV.partition_S(gK); // (KCPY, KCPY_N, KCPY_K, nblocksN) + Tensor tKsK = gmem_thr_copy_QKV.partition_D(sK); + Tensor tVgV = gmem_thr_copy_QKV.partition_S(gV); // (VCPY, VCPY_N, VCPY_K, nblocksN) + Tensor tVsV = gmem_thr_copy_QKV.partition_D(sV); typename Kernel_traits::TiledMma tiled_mma; auto thr_mma = tiled_mma.get_thread_slice(tidx); - cute::Tensor tSrQ = thr_mma.partition_fragment_A(sQ); // (MMA,MMA_M,MMA_K) - cute::Tensor tSrK = thr_mma.partition_fragment_B(sK); // (MMA,MMA_N,MMA_K) - cute::Tensor tOrVt = thr_mma.partition_fragment_B(sVtNoSwizzle); // (MMA, MMA_K,MMA_N) + Tensor tSrQ = thr_mma.partition_fragment_A(sQ); // (MMA,MMA_M,MMA_K) + Tensor tSrK = thr_mma.partition_fragment_B(sK); // (MMA,MMA_N,MMA_K) + Tensor tOrVt = thr_mma.partition_fragment_B(sVtNoSwizzle); // (MMA, MMA_K,MMA_N) - cute::Tensor acc_o = partition_fragment_C(tiled_mma, cute::Shape, cute::Int>{}); // MMA, MMA_M, MMA_K + Tensor tSgS = thr_mma.partition_C(gP); + + Tensor acc_o = partition_fragment_C(tiled_mma, Shape, Int>{}); // MMA, MMA_M, MMA_K // // Copy Atom retiling @@ -205,51 +165,46 @@ inline __device__ void compute_attn_1rowblock(const Params& params, const int bi auto smem_tiled_copy_Q = make_tiled_copy_A(typename Kernel_traits::SmemCopyAtom{}, tiled_mma); auto smem_thr_copy_Q = smem_tiled_copy_Q.get_thread_slice(tidx); - cute::Tensor tSsQ = smem_thr_copy_Q.partition_S(sQ); + Tensor tSsQ = smem_thr_copy_Q.partition_S(sQ); auto smem_tiled_copy_K = make_tiled_copy_B(typename Kernel_traits::SmemCopyAtom{}, tiled_mma); auto smem_thr_copy_K = smem_tiled_copy_K.get_thread_slice(tidx); - cute::Tensor tSsK = smem_thr_copy_K.partition_S(sK); + Tensor tSsK = smem_thr_copy_K.partition_S(sK); auto smem_tiled_copy_V = make_tiled_copy_B(typename Kernel_traits::SmemCopyAtomTransposed{}, tiled_mma); auto smem_thr_copy_V = smem_tiled_copy_V.get_thread_slice(tidx); - cute::Tensor tOsVt = smem_thr_copy_V.partition_S(sVt); - - // TODO: this might need to change if we change the mma instruction in SM70 - cute::Tensor scores_max = make_tensor(cute::Shape(acc_o)>>{}); - cute::Tensor scores_sum = make_fragment_like(scores_max); + Tensor tOsVt = smem_thr_copy_V.partition_S(sVt); // // PREDICATES // // Construct identity layout for sQ and sK - cute::Tensor cQ = make_identity_tensor(make_shape(cute::size<0>(sQ), cute::size<1>(sQ))); // (BLK_M,BLK_K) -> (blk_m,blk_k) - cute::Tensor cKV = make_identity_tensor(make_shape(cute::size<0>(sK), cute::size<1>(sK))); // (BLK_N,BLK_K) -> (blk_n,blk_k) + Tensor cQ = make_identity_tensor(make_shape(size<0>(sQ), size<1>(sQ))); // (BLK_M,BLK_K) -> (blk_m,blk_k) + Tensor cKV = make_identity_tensor(make_shape(size<0>(sK), size<1>(sK))); // (BLK_N,BLK_K) -> (blk_n,blk_k) // Repeat the partitioning with identity layouts - cute::Tensor tQcQ = gmem_thr_copy_QKV.partition_S(cQ); // (ACPY,ACPY_M,ACPY_K) -> (blk_m,blk_k) - cute::Tensor tKVcKV = gmem_thr_copy_QKV.partition_S(cKV); // (BCPY,BCPY_N,BCPY_K) -> (blk_n,blk_k) + Tensor tQcQ = gmem_thr_copy_QKV.partition_S(cQ); // (ACPY,ACPY_M,ACPY_K) -> (blk_m,blk_k) + Tensor tKVcKV = gmem_thr_copy_QKV.partition_S(cKV); // (BCPY,BCPY_N,BCPY_K) -> (blk_n,blk_k) // Allocate predicate tensors for k - cute::Tensor tQpQ = make_tensor(make_shape(cute::size<2>(tQsQ))); - cute::Tensor tKVpKV = make_tensor(make_shape(cute::size<2>(tKsK))); + Tensor tQpQ = make_tensor(make_shape(size<2>(tQsQ))); + Tensor tKVpKV = make_tensor(make_shape(size<2>(tKsK))); // Set predicates for k bounds if (!Is_even_K) { #pragma unroll - for (int k = 0; k < cute::size(tQpQ); ++k) { + for (int k = 0; k < size(tQpQ); ++k) { tQpQ(k) = get<1>(tQcQ(0, 0, k)) < params.d; } #pragma unroll - for (int k = 0; k < cute::size(tKVpKV); ++k) { + for (int k = 0; k < size(tKVpKV); ++k) { tKVpKV(k) = get<1>(tKVcKV(0, 0, k)) < params.d; } } // Prologue - cute::Tensor tQrQ = make_fragment_like(tQgQ); // We don't need to clear the sQ smem tiles since we'll only write out the valid outputs flash::copy(gmem_tiled_copy_QKV, tQgQ, tQsQ, tQcQ, tQpQ, binfo.actual_seqlen_q - m_block * kBlockM); @@ -260,27 +215,34 @@ inline __device__ void compute_attn_1rowblock(const Params& params, const int bi if (Kernel_traits::Share_Q_K_smem) { flash::cp_async_wait<0>(); __syncthreads(); - cute::Tensor tSrQ_copy_view = smem_thr_copy_Q.retile_D(tSrQ); - CUTE_STATIC_ASSERT_V(cute::size<1>(tSsQ) == cute::size<1>(tSrQ_copy_view)); // M + Tensor tSrQ_copy_view = smem_thr_copy_Q.retile_D(tSrQ); + CUTE_STATIC_ASSERT_V(size<1>(tSsQ) == size<1>(tSrQ_copy_view)); // M cute::copy(smem_tiled_copy_Q, tSsQ, tSrQ_copy_view); __syncthreads(); } int n_block = n_block_max - 1; // We don't need to clear the sK smem tiles since we'll mask out the scores anyway. - flash::copy(gmem_tiled_copy_QKV, tKgK, tKsK, tKVcKV, tKVpKV, + flash::copy(gmem_tiled_copy_QKV, tKgK(_, _, _, n_block), tKsK, tKVcKV, tKVpKV, binfo.actual_seqlen_k - n_block * kBlockN); cute::cp_async_fence(); if (Kernel_traits::Is_Q_in_regs && !Kernel_traits::Share_Q_K_smem) { flash::cp_async_wait<1>(); __syncthreads(); - cute::Tensor tSrQ_copy_view = smem_thr_copy_Q.retile_D(tSrQ); - CUTE_STATIC_ASSERT_V(cute::size<1>(tSsQ) == cute::size<1>(tSrQ_copy_view)); // M + Tensor tSrQ_copy_view = smem_thr_copy_Q.retile_D(tSrQ); + CUTE_STATIC_ASSERT_V(size<1>(tSsQ) == size<1>(tSrQ_copy_view)); // M cute::copy(smem_tiled_copy_Q, tSsQ, tSrQ_copy_view); } clear(acc_o); + flash::Softmax<2 * size<1>(acc_o)> softmax; + + const float alibi_slope = !Has_alibi || params.alibi_slopes_ptr == nullptr + ? 0.0f + : reinterpret_cast(params.alibi_slopes_ptr)[bidb * params.alibi_slopes_batch_stride + bidh] / params.scale_softmax; + flash::Mask mask(binfo.actual_seqlen_k, binfo.actual_seqlen_q, + params.window_size_left, params.window_size_right, alibi_slope); // For performance reason, we separate out two kinds of iterations: // those that need masking on S, and those that don't. @@ -292,22 +254,23 @@ inline __device__ void compute_attn_1rowblock(const Params& params, const int bi // mask 2 blocks (e.g. when kBlockM == kBlockN), not just 1. constexpr int n_masking_steps = (!Is_causal && !Is_local) ? 1 - : ((Is_even_MN && Is_causal) ? cute::ceil_div(kBlockM, kBlockN) : cute::ceil_div(kBlockM, kBlockN) + 1); + : ((Is_even_MN && Is_causal) + ? cute::ceil_div(kBlockM, kBlockN) + : cute::ceil_div(kBlockM, kBlockN) + 1); #pragma unroll for (int masking_step = 0; masking_step < n_masking_steps; ++masking_step, --n_block) { - cute::Tensor acc_s = partition_fragment_C(tiled_mma, cute::Shape, cute::Int>{}); // (MMA=4, MMA_M, MMA_N) + Tensor acc_s = partition_fragment_C(tiled_mma, Shape, Int>{}); // (MMA=4, MMA_M, MMA_N) clear(acc_s); flash::cp_async_wait<0>(); __syncthreads(); // Advance gV if (masking_step > 0) { - tVgV.data() = tVgV.data() + (-int(kBlockN * params.v_row_stride)); - flash::copy(gmem_tiled_copy_QKV, tVgV, tVsV, tKVcKV, tKVpKV); + flash::copy(gmem_tiled_copy_QKV, tVgV(_, _, _, n_block), tVsV, tKVcKV, tKVpKV); } else { // Clear the smem tiles to account for predicated off loads flash::copy( - gmem_tiled_copy_QKV, tVgV, tVsV, tKVcKV, tKVpKV, binfo.actual_seqlen_k - n_block * kBlockN); + gmem_tiled_copy_QKV, tVgV(_, _, _, n_block), tVsV, tKVcKV, tKVpKV, binfo.actual_seqlen_k - n_block * kBlockN); } cute::cp_async_fence(); @@ -316,31 +279,13 @@ inline __device__ void compute_attn_1rowblock(const Params& params, const int bi smem_thr_copy_Q, smem_thr_copy_K); // if (cute::thread0()) { print(acc_s); } - // Reshape acc_s from (MMA=4, MMA_M, MMA_N) to (nrow=(2, MMA_M), ncol=(2, MMA_N)) - cute::Tensor scores = make_tensor(acc_s.data(), flash::convert_layout_acc_rowcol(acc_s.layout())); - - // We don't put the masking before the matmul S = Q K^T because we don't clear sK - // for rows outside actual_seqlen_k. So those rows could have Inf / NaN, and the matmul - // can produce Inf / NaN. - if (!Is_causal && !Is_local) { - if (!Is_even_MN) { - flash::apply_mask(scores, binfo.actual_seqlen_k - n_block * kBlockN); - } - } else { - // I can't get the stride from idx_row - flash::apply_mask_local(scores, n_block * kBlockN, binfo.actual_seqlen_k, - // m_block * kBlockM + get<0>(idx_row(0)), - m_block * kBlockM + (tidx / 32) * 16 + (tidx % 32) / 4, - binfo.actual_seqlen_q, kNWarps * 16, - params.window_size_left, params.window_size_right); - } + mask.template apply_mask( + acc_s, n_block * kBlockN, m_block * kBlockM + (tidx / 32) * 16 + (tidx % 32) / 4, kNWarps * 16); flash::cp_async_wait<0>(); __syncthreads(); if (n_block > n_block_min) { - // Advance gK - tKgK.data() = tKgK.data() + (-int(kBlockN * params.k_row_stride)); - flash::copy(gmem_tiled_copy_QKV, tKgK, tKsK, tKVcKV, tKVpKV); + flash::copy(gmem_tiled_copy_QKV, tKgK(_, _, _, n_block - 1), tKsK, tKVcKV, tKVpKV); // This cp_async_fence needs to be in the if block, otherwise the synchronization // isn't right and we get race conditions. cute::cp_async_fence(); @@ -348,22 +293,15 @@ inline __device__ void compute_attn_1rowblock(const Params& params, const int bi // TODO: when we have key_padding_mask we'll need to Check_inf masking_step == 0 - ? softmax_rescale_o(scores, scores_max, scores_sum, acc_o, params.scale_softmax_log2) - : softmax_rescale_o(scores, scores_max, scores_sum, acc_o, params.scale_softmax_log2); + ? softmax.template softmax_rescale_o(acc_s, acc_o, params.scale_softmax_log2) + : softmax.template softmax_rescale_o(acc_s, acc_o, params.scale_softmax_log2); - // Convert scores from fp32 to fp16/bf16 - cute::Tensor rP = flash::convert_type(scores); - // Reshape rP from (nrow=(2, MMA_M), ncol=(2, MMA_N)) to ((2, 2, 2), MMA_M, MMA_N / 2) + // Convert acc_s from fp32 to fp16/bf16 + Tensor rP = flash::convert_type(acc_s); + // Reshape rP from (MMA=4, MMA_M, MMA_N) to ((4, 2), MMA_M, MMA_N / 2) // if using m16n8k16 or ((2, 2, 1), MMA_M, MMA_N) if using m16n8k8. - cute::Tensor tOrP = make_tensor(rP.data(), flash::convert_layout_rowcol_Aregs(rP.layout())); - // if (Return_softmax) { - // cute::Tensor tOrP_copy = make_fragment_like(tOrP); - // copy(tOrP, tOrP_copy); - // flash::write_softmax_to_gmem(tOrP_copy, tPgP, gmem_thr_copy_P); - // tPgP.data() = tPgP.data() + (-kBlockN); - // } - - flash::gemm_A_in_regs(acc_o, tOrP, tOrVt, tOsVt, tiled_mma, smem_tiled_copy_V, smem_thr_copy_V); + Tensor tOrP = make_tensor(rP.data(), flash::convert_layout_acc_Aregs(rP.layout())); + flash::gemm_rs(acc_o, tOrP, tOrVt, tOsVt, tiled_mma, smem_tiled_copy_V, smem_thr_copy_V); // This check is at the end of the loop since we always have at least 1 iteration if (n_masking_steps > 1 && n_block <= n_block_min) { @@ -374,13 +312,11 @@ inline __device__ void compute_attn_1rowblock(const Params& params, const int bi // These are the iterations where we don't need masking on S for (; n_block >= n_block_min; --n_block) { - cute::Tensor acc_s = partition_fragment_C(tiled_mma, cute::Shape, cute::Int>{}); // (MMA=4, MMA_M, MMA_N) + Tensor acc_s = partition_fragment_C(tiled_mma, Shape, Int>{}); // (MMA=4, MMA_M, MMA_N) clear(acc_s); flash::cp_async_wait<0>(); __syncthreads(); - // Advance gV - tVgV.data() = tVgV.data() + (-int(kBlockN * params.v_row_stride)); - flash::copy(gmem_tiled_copy_QKV, tVgV, tVsV, tKVcKV, tKVpKV); + flash::copy(gmem_tiled_copy_QKV, tVgV(_, _, _, n_block), tVsV, tKVcKV, tKVpKV); cute::cp_async_fence(); flash::gemm( @@ -390,64 +326,36 @@ inline __device__ void compute_attn_1rowblock(const Params& params, const int bi flash::cp_async_wait<0>(); __syncthreads(); if (n_block > n_block_min) { - // Advance gK - tKgK.data() = tKgK.data() + (-int(kBlockN * params.k_row_stride)); - flash::copy(gmem_tiled_copy_QKV, tKgK, tKsK, tKVcKV, tKVpKV); + flash::copy(gmem_tiled_copy_QKV, tKgK(_, _, _, n_block - 1), tKsK, tKVcKV, tKVpKV); // This cp_async_fence needs to be in the if block, otherwise the synchronization // isn't right and we get race conditions. cute::cp_async_fence(); } - // Reshape acc_s from (MMA=4, MMA_M, MMA_N) to (nrow=(2, MMA_M), ncol=(2, MMA_N)) - Tensor scores = make_tensor(acc_s.data(), flash::convert_layout_acc_rowcol(acc_s.layout())); - if (Is_local && n_block * kBlockN < (m_block + 1) * kBlockM + binfo.actual_seqlen_k - binfo.actual_seqlen_q + params.window_size_right) { - flash::apply_mask_local( - scores, n_block * kBlockN, binfo.actual_seqlen_k, - m_block * kBlockM + (tidx / 32) * 16 + (tidx % 32) / 4, - binfo.actual_seqlen_q, kNWarps * 16, - params.window_size_left, params.window_size_right); - } - softmax_rescale_o(scores, scores_max, scores_sum, acc_o, params.scale_softmax_log2); + mask.template apply_mask( + acc_s, n_block * kBlockN, m_block * kBlockM + (tidx / 32) * 16 + (tidx % 32) / 4, kNWarps * 16); - cute::Tensor rP = flash::convert_type(scores); - // Reshape rP from (nrow=(2, MMA_M), ncol=(2, MMA_N)) to ((2, 2, 2), MMA_M, MMA_N / 2) - // if using m16n8k16 or ((2, 2, 1), MMA_M, MMA_N) if using m16n8k8. - cute::Tensor tOrP = make_tensor(rP.data(), flash::convert_layout_rowcol_Aregs(rP.layout())); - // if (Return_softmax) { - // cute::Tensor tOrP_copy = make_fragment_like(tOrP); - // copy(tOrP, tOrP_copy); - // flash::write_softmax_to_gmem(tOrP_copy, tPgP, gmem_thr_copy_P); - // tPgP.data() = tPgP.data() + (-kBlockN); - // } - - flash::gemm_A_in_regs(acc_o, tOrP, tOrVt, tOsVt, tiled_mma, smem_tiled_copy_V, smem_thr_copy_V); + softmax.template softmax_rescale_o(acc_s, acc_o, params.scale_softmax_log2); + + Tensor rP = flash::convert_type(acc_s); + // Reshape rP from (MMA=4, MMA_M, MMA_N) to ((4, 2), MMA_M, MMA_N / 2) + // if using m16n8k16 or (4, MMA_M, MMA_N) if using m16n8k8. + Tensor tOrP = make_tensor(rP.data(), flash::convert_layout_acc_Aregs(rP.layout())); + flash::gemm_rs(acc_o, tOrP, tOrVt, tOsVt, tiled_mma, smem_tiled_copy_V, smem_thr_copy_V); } // Epilogue - // Reshape acc_o from (MMA=4, MMA_M, MMA_K) to (nrow=(2, MMA_M), ncol=(2, MMA_K)) - cute::Tensor acc_o_rowcol = make_tensor(acc_o.data(), flash::convert_layout_acc_rowcol(acc_o.layout())); - cute::Tensor lse = make_fragment_like(scores_sum); -#pragma unroll - for (int mi = 0; mi < cute::size<0>(acc_o_rowcol); ++mi) { - float sum = scores_sum(mi); - float inv_sum = (sum == 0.f || sum != sum) ? 1.f : 1.f / sum; - lse(mi) = (sum == 0.f || sum != sum) ? INFINITY : scores_max(mi) * params.scale_softmax + __logf(sum); - float scale = inv_sum; -#pragma unroll - for (int ni = 0; ni < cute::size<1>(acc_o_rowcol); ++ni) { - acc_o_rowcol(mi, ni) *= scale; - } - } + Tensor lse = softmax.template normalize_softmax_lse<>(acc_o, params.scale_softmax); // Convert acc_o from fp32 to fp16/bf16 - cute::Tensor rO = flash::convert_type(acc_o); - cute::Tensor sO = make_tensor(sQ.data(), typename Kernel_traits::SmemLayoutO{}); // (SMEM_M,SMEM_N) + Tensor rO = flash::convert_type(acc_o); + Tensor sO = make_tensor(sQ.data(), typename Kernel_traits::SmemLayoutO{}); // (SMEM_M,SMEM_N) // Partition sO to match the accumulator partitioning auto smem_tiled_copy_O = make_tiled_copy_C(typename Kernel_traits::SmemCopyAtomO{}, tiled_mma); - auto smem_thr_copy_O = smem_tiled_copy_O.get_thread_slice(tidx); // auto smem_thr_copy_O = make_tiled_copy_C_warpcontiguousM(typename Kernel_traits::SmemCopyAtomO{}, tiled_mma).get_thread_slice(tidx); - cute::Tensor taccOrO = smem_thr_copy_O.retile_S(rO); // ((Atom,AtomNum), MMA_M, MMA_N) - cute::Tensor taccOsO = smem_thr_copy_O.partition_D(sO); // ((Atom,AtomNum),PIPE_M,PIPE_N) + auto smem_thr_copy_O = smem_tiled_copy_O.get_thread_slice(tidx); + Tensor taccOrO = smem_thr_copy_O.retile_S(rO); // ((Atom,AtomNum), MMA_M, MMA_N) + Tensor taccOsO = smem_thr_copy_O.partition_D(sO); // ((Atom,AtomNum),PIPE_M,PIPE_N) // sO has the same size as sQ, so we don't need to sync here. if (Kernel_traits::Share_Q_K_smem) { @@ -456,33 +364,35 @@ inline __device__ void compute_attn_1rowblock(const Params& params, const int bi cute::copy(smem_tiled_copy_O, taccOrO, taccOsO); - const index_t row_offset_o = binfo.q_offset(params.o_batch_stride, params.o_row_stride, bidb) + m_block * kBlockM * params.o_row_stride + bidh * params.o_head_stride; - const index_t row_offset_lse = (bidb * params.h + bidh) * params.seqlen_q + m_block * kBlockM; - cute::Tensor gO = make_tensor(make_gmem_ptr(reinterpret_cast(params.o_ptr) + row_offset_o), - cute::Shape, cute::Int>{}, - make_stride(params.o_row_stride, _1{})); - cute::Tensor gLSE = make_tensor(make_gmem_ptr(reinterpret_cast(params.softmax_lse_ptr) + row_offset_lse), - cute::Shape>{}, cute::Stride<_1>{}); + Tensor mO = make_tensor(make_gmem_ptr(reinterpret_cast(params.o_ptr) + binfo.q_offset(params.o_batch_stride, params.o_row_stride, bidb)), + make_shape(binfo.actual_seqlen_q, params.h, params.d), + make_stride(params.o_row_stride, params.o_head_stride, _1{})); + Tensor gO = local_tile(mO(_, bidh, _), Shape, Int>{}, + make_coord(m_block, 0)); // (kBlockM, kHeadDim) + Tensor mLSE = make_tensor(make_gmem_ptr(reinterpret_cast(params.softmax_lse_ptr)), + make_shape(params.b, params.h, params.seqlen_q), + make_stride(params.h * params.seqlen_q, params.seqlen_q, _1{})); + Tensor gLSE = local_tile(mLSE(bidb, bidh, _), Shape>{}, make_coord(m_block)); typename Kernel_traits::GmemTiledCopyO gmem_tiled_copy_O; auto gmem_thr_copy_O = gmem_tiled_copy_O.get_thread_slice(tidx); - cute::Tensor tOsO = gmem_thr_copy_O.partition_S(sO); // ((Atom,AtomNum),ATOM_M,ATOM_N) - cute::Tensor tOgO = gmem_thr_copy_O.partition_D(gO); + Tensor tOsO = gmem_thr_copy_O.partition_S(sO); // ((Atom,AtomNum),ATOM_M,ATOM_N) + Tensor tOgO = gmem_thr_copy_O.partition_D(gO); __syncthreads(); - cute::Tensor tOrO = make_tensor(cute::shape(tOgO)); + Tensor tOrO = make_tensor(shape(tOgO)); cute::copy(gmem_tiled_copy_O, tOsO, tOrO); - cute::Tensor caccO = make_identity_tensor(cute::Shape, cute::Int>{}); // (BLK_M,BLK_K) -> (blk_m,blk_k) - cute::Tensor taccOcO = thr_mma.partition_C(caccO); // (MMA,MMA_M,MMA_K) - static_assert(decltype(cute::size<0>(taccOcO))::value == 4); + Tensor caccO = make_identity_tensor(Shape, Int>{}); // (BLK_M,BLK_K) -> (blk_m,blk_k) + Tensor taccOcO = thr_mma.partition_C(caccO); // (MMA,MMA_M,MMA_K) + static_assert(decltype(size<0>(taccOcO))::value == 4); // Convert to ((2, 2), MMA_M, MMA_K) then take only the row indices. - cute::Tensor taccOcO_row = logical_divide(taccOcO, cute::Shape<_2>{})(make_coord(0, _), _, 0); - CUTE_STATIC_ASSERT_V(cute::size(lse) == cute::size(taccOcO_row)); // MMA_M + Tensor taccOcO_row = logical_divide(taccOcO, Shape<_2>{})(make_coord(0, _), _, 0); + CUTE_STATIC_ASSERT_V(size(lse) == size(taccOcO_row)); // MMA_M if (get<1>(taccOcO_row(0)) == 0) { #pragma unroll - for (int mi = 0; mi < cute::size(lse); ++mi) { + for (int mi = 0; mi < size(lse); ++mi) { const int row = get<0>(taccOcO_row(mi)); if (row < binfo.actual_seqlen_q - m_block * kBlockM) { gLSE(row) = lse(mi); @@ -491,13 +401,13 @@ inline __device__ void compute_attn_1rowblock(const Params& params, const int bi } // Construct identity layout for sO - cute::Tensor cO = make_identity_tensor(make_shape(cute::size<0>(sO), cute::size<1>(sO))); // (BLK_M,BLK_K) -> (blk_m,blk_k) + Tensor cO = make_identity_tensor(make_shape(size<0>(sO), size<1>(sO))); // (BLK_M,BLK_K) -> (blk_m,blk_k) // Repeat the partitioning with identity layouts - cute::Tensor tOcO = gmem_thr_copy_O.partition_D(cO); // (ACPY,ACPY_M,ACPY_K) -> (blk_m,blk_k) - cute::Tensor tOpO = make_tensor(make_shape(cute::size<2>(tOgO))); + Tensor tOcO = gmem_thr_copy_O.partition_D(cO); // (ACPY,ACPY_M,ACPY_K) -> (blk_m,blk_k) + Tensor tOpO = make_tensor(make_shape(size<2>(tOgO))); if (!Is_even_K) { #pragma unroll - for (int k = 0; k < cute::size(tOpO); ++k) { + for (int k = 0; k < size(tOpO); ++k) { tOpO(k) = get<1>(tOcO(0, 0, k)) < params.d; } } @@ -508,8 +418,10 @@ inline __device__ void compute_attn_1rowblock(const Params& params, const int bi //////////////////////////////////////////////////////////////////////////////////////////////////// -template -inline __device__ void compute_attn_1rowblock_splitkv(const Params& params, const int bidb, const int bidh, const int m_block, const int n_split_idx, const int num_n_splits) { +template +inline __device__ void compute_attn_1rowblock_splitkv(const Params& params, const int bidb, const int bidh, + const int m_block, const int n_split_idx, + const int num_n_splits) { using Element = typename Kernel_traits::Element; using ElementAccum = typename Kernel_traits::ElementAccum; using index_t = typename Kernel_traits::index_t; @@ -527,8 +439,8 @@ inline __device__ void compute_attn_1rowblock_splitkv(const Params& params, cons using GmemTiledCopyO = std::conditional_t< !Split, - typename Kernel_traits::GmemTiledCopyOaccum, - typename Kernel_traits::GmemTiledCopyO>; + typename Kernel_traits::GmemTiledCopyO, + typename Kernel_traits::GmemTiledCopyOaccum>; using ElementO = std::conditional_t; const BlockInfo binfo(params, bidb); @@ -591,15 +503,29 @@ inline __device__ void compute_attn_1rowblock_splitkv(const Params& params, cons // that needs masking when we read K and V from global memory. Moreover, iterating in reverse // might save us 1 register (we just need n_block instead of both n_block and n_block_max). - const index_t row_offset_q = binfo.q_offset(params.q_batch_stride, params.q_row_stride, bidb) + m_block * kBlockM * params.q_row_stride + bidh * params.q_head_stride; // We move K and V to the last block. const int bidb_cache = params.cache_batch_idx == nullptr ? bidb : params.cache_batch_idx[bidb]; - const index_t row_offset_k = binfo.k_offset(params.k_batch_stride, params.k_row_stride, bidb_cache) + (n_block_max - 1) * kBlockN * params.k_row_stride + (bidh / params.h_h_k_ratio) * params.k_head_stride; - const index_t row_offset_v = binfo.k_offset(params.v_batch_stride, params.v_row_stride, bidb_cache) + (n_block_max - 1) * kBlockN * params.v_row_stride + (bidh / params.h_h_k_ratio) * params.v_head_stride; - - Tensor gQ = make_tensor(make_gmem_ptr(reinterpret_cast(params.q_ptr) + row_offset_q), - Shape, Int>{}, - make_stride(params.q_row_stride, _1{})); + const int* block_table = params.block_table == nullptr + ? nullptr + : params.block_table + bidb * params.block_table_batch_stride; + const int block_table_idx = block_table == nullptr + ? 0 + : (n_block_max - 1) * kBlockN / params.page_block_size; + const int block_table_offset = block_table == nullptr + ? 0 + : (n_block_max - 1) * kBlockN - block_table_idx * params.page_block_size; + const index_t row_offset_k = block_table == nullptr + ? binfo.k_offset(params.k_batch_stride, params.k_row_stride, bidb_cache) + (n_block_max - 1) * kBlockN * params.k_row_stride + (bidh / params.h_h_k_ratio) * params.k_head_stride + : block_table[block_table_idx] * params.k_batch_stride + block_table_offset * params.k_row_stride + (bidh / params.h_h_k_ratio) * params.k_head_stride; + const index_t row_offset_v = block_table == nullptr + ? binfo.k_offset(params.v_batch_stride, params.v_row_stride, bidb_cache) + (n_block_max - 1) * kBlockN * params.v_row_stride + (bidh / params.h_h_k_ratio) * params.v_head_stride + : block_table[block_table_idx] * params.v_batch_stride + block_table_offset * params.v_row_stride + (bidh / params.h_h_k_ratio) * params.v_head_stride; + + Tensor mQ = make_tensor(make_gmem_ptr(reinterpret_cast(params.q_ptr) + binfo.q_offset(params.q_batch_stride, params.q_row_stride, bidb)), + make_shape(binfo.actual_seqlen_q, params.h, params.d), + make_stride(params.q_row_stride, params.q_head_stride, _1{})); + Tensor gQ = local_tile(mQ(_, bidh, _), Shape, Int>{}, + make_coord(m_block, 0)); // (kBlockM, kHeadDim) Tensor gK = make_tensor(make_gmem_ptr(reinterpret_cast(params.k_ptr) + row_offset_k), Shape, Int>{}, make_stride(params.k_row_stride, _1{})); @@ -649,10 +575,6 @@ inline __device__ void compute_attn_1rowblock_splitkv(const Params& params, cons auto smem_thr_copy_V = smem_tiled_copy_V.get_thread_slice(tidx); Tensor tOsVt = smem_thr_copy_V.partition_S(sVt); - // TODO: this might need to change if we change the mma instruction in SM70 - Tensor scores_max = make_tensor(Shape(acc_o)>>{}); - Tensor scores_sum = make_fragment_like(scores_max); - // // PREDICATES // @@ -732,10 +654,11 @@ inline __device__ void compute_attn_1rowblock_splitkv(const Params& params, cons Tensor tVgVnew = gmem_thr_copy_QKV.partition_S(gVnew); // (VCPY, VCPY_N, VCPY_K) const int n_block_copy_min = std::max(n_block_min, binfo.seqlen_k_cache / kBlockN); + auto tKgK_data = tKgK.data(); + auto tVgV_data = tVgV.data(); for (int n_block = n_block_max - 1; n_block >= n_block_copy_min; n_block--) { flash::copy_w_min_idx( tVgVnew, tVgV, tKVcKV, tKVpKV, binfo.actual_seqlen_k - n_block * kBlockN, binfo.seqlen_k_cache - n_block * kBlockN); - tVgV.data() = tVgV.data() + (-int(kBlockN * params.v_row_stride)); tVgVnew.data() = tVgVnew.data() + (-int(kBlockN * params.vnew_row_stride)); if (params.rotary_dim == 0) { flash::copy_w_min_idx( @@ -757,19 +680,30 @@ inline __device__ void compute_attn_1rowblock_splitkv(const Params& params, cons tRgSinCont.data() = tRgSinCont.data() + (-int(kBlockN * params.rotary_dim / 2)); } } - tKgK.data() = tKgK.data() + (-int(kBlockN * params.k_row_stride)); tKgKnew.data() = tKgKnew.data() + (-int(kBlockN * params.knew_row_stride)); + if (block_table == nullptr) { + tVgV.data() = tVgV.data() + (-int(kBlockN * params.v_row_stride)); + tKgK.data() = tKgK.data() + (-int(kBlockN * params.k_row_stride)); + } else { + if (n_block > n_block_copy_min) { + const int block_table_idx_cur = n_block * kBlockN / params.page_block_size; + const int block_table_offset_cur = n_block * kBlockN - block_table_idx_cur * params.page_block_size; + const int block_table_idx_next = (n_block - 1) * kBlockN / params.page_block_size; + const int block_table_offset_next = (n_block - 1) * kBlockN - block_table_idx_next * params.page_block_size; + const int table_diff = block_table[block_table_idx_next] - block_table[block_table_idx_cur]; + const int offset_diff = block_table_offset_next - block_table_offset_cur; + tVgV.data() = tVgV.data() + table_diff * params.v_batch_stride + offset_diff * params.v_row_stride; + tKgK.data() = tKgK.data() + table_diff * params.k_batch_stride + offset_diff * params.k_row_stride; + } + } } // Need this before we can read in K again, so that we'll see the updated K values. __syncthreads(); - if (n_block_max > n_block_copy_min) { - tKgK.data() = tKgK.data() + (n_block_max - n_block_copy_min) * kBlockN * params.k_row_stride; - tVgV.data() = tVgV.data() + (n_block_max - n_block_copy_min) * kBlockN * params.v_row_stride; - } + tKgK.data() = tKgK_data; + tVgV.data() = tVgV_data; } // Read Q from gmem to smem, optionally apply rotary embedding. - Tensor tQrQ = make_fragment_like(tQgQ); if (!Append_KV || params.rotary_dim == 0) { // We don't need to clear the sQ smem tiles since we'll only write out the valid outputs flash::copy(gmem_tiled_copy_QKV, tQgQ, tQsQ, tQcQ, tQpQ, @@ -818,6 +752,12 @@ inline __device__ void compute_attn_1rowblock_splitkv(const Params& params, cons clear(acc_o); + flash::Softmax<2 * size<1>(acc_o)> softmax; + + const float alibi_slope = !Has_alibi ? 0.0f + : reinterpret_cast(params.alibi_slopes_ptr)[bidb * params.alibi_slopes_batch_stride + bidh] / params.scale_softmax; + flash::Mask mask(binfo.actual_seqlen_k, binfo.actual_seqlen_q, params.window_size_left, params.window_size_right, alibi_slope); + // For performance reason, we separate out two kinds of iterations: // those that need masking on S, and those that don't. // We need masking on S for the very last block when K and V has length not multiple of kBlockN. @@ -838,7 +778,15 @@ inline __device__ void compute_attn_1rowblock_splitkv(const Params& params, cons // Advance gV if (masking_step > 0) { - tVgV.data() = tVgV.data() + (-int(kBlockN * params.v_row_stride)); + if (block_table == nullptr) { + tVgV.data() = tVgV.data() + (-int(kBlockN * params.v_row_stride)); + } else { + const int block_table_idx_cur = (n_block + 1) * kBlockN / params.page_block_size; + const int block_table_offset_cur = (n_block + 1) * kBlockN - block_table_idx_cur * params.page_block_size; + const int block_table_idx_next = n_block * kBlockN / params.page_block_size; + const int block_table_offset_next = n_block * kBlockN - block_table_idx_next * params.page_block_size; + tVgV.data() = tVgV.data() + (block_table[block_table_idx_next] - block_table[block_table_idx_cur]) * params.v_batch_stride + (block_table_offset_next - block_table_offset_cur) * params.v_row_stride; + } flash::copy(gmem_tiled_copy_QKV, tVgV, tVsV, tKVcKV, tKVpKV); } else { // Clear the smem tiles to account for predicated off loads @@ -852,22 +800,8 @@ inline __device__ void compute_attn_1rowblock_splitkv(const Params& params, cons smem_thr_copy_Q, smem_thr_copy_K); // if (cute::thread0()) { print(acc_s); } - // Reshape acc_s from (MMA=4, MMA_M, MMA_N) to (nrow=(2, MMA_M), ncol=(2, MMA_N)) - Tensor scores = make_tensor(acc_s.data(), flash::convert_layout_acc_rowcol(acc_s.layout())); - // if (cute::thread0()) { print(scores); } - // We don't put the masking before the matmul S = Q K^T because we don't clear sK - // for rows outside actual_seqlen_k. So those rows could have Inf / NaN, and the matmul - // can produce Inf / NaN. - if (!Is_causal && !Is_local) { - if (!Is_even_MN) { - flash::apply_mask(scores, binfo.actual_seqlen_k - n_block * kBlockN); - } - } else { - flash::apply_mask_local(scores, n_block * kBlockN, binfo.actual_seqlen_k, - m_block * kBlockM + (tidx / 32) * 16 + (tidx % 32) / 4, - binfo.actual_seqlen_q, kNWarps * 16, - params.window_size_left, params.window_size_right); - } + mask.template apply_mask( + acc_s, n_block * kBlockN, m_block * kBlockM + (tidx / 32) * 16 + (tidx % 32) / 4, kNWarps * 16); flash::cp_async_wait<0>(); __syncthreads(); @@ -876,7 +810,15 @@ inline __device__ void compute_attn_1rowblock_splitkv(const Params& params, cons if (n_block > n_block_min) { // Advance gK - tKgK.data() = tKgK.data() + (-int(kBlockN * params.k_row_stride)); + if (block_table == nullptr) { + tKgK.data() = tKgK.data() + (-int(kBlockN * params.k_row_stride)); + } else { + const int block_table_idx_cur = n_block * kBlockN / params.page_block_size; + const int block_table_offset_cur = n_block * kBlockN - block_table_idx_cur * params.page_block_size; + const int block_table_idx_next = (n_block - 1) * kBlockN / params.page_block_size; + const int block_table_offset_next = (n_block - 1) * kBlockN - block_table_idx_next * params.page_block_size; + tKgK.data() = tKgK.data() + (block_table[block_table_idx_next] - block_table[block_table_idx_cur]) * params.k_batch_stride + (block_table_offset_next - block_table_offset_cur) * params.k_row_stride; + } flash::copy(gmem_tiled_copy_QKV, tKgK, tKsK, tKVcKV, tKVpKV); // This cp_async_fence needs to be in the if block, otherwise the synchronization // isn't right and we get race conditions. @@ -885,18 +827,17 @@ inline __device__ void compute_attn_1rowblock_splitkv(const Params& params, cons // We have key_padding_mask so we'll need to Check_inf masking_step == 0 - ? softmax_rescale_o(scores, scores_max, scores_sum, acc_o, params.scale_softmax_log2) - : softmax_rescale_o(scores, scores_max, scores_sum, acc_o, params.scale_softmax_log2); + ? softmax.template softmax_rescale_o(acc_s, acc_o, params.scale_softmax_log2) + : softmax.template softmax_rescale_o(acc_s, acc_o, params.scale_softmax_log2); // if (cute::thread0()) { print(scores_max); print(scores_sum); print(scores); } - // Convert scores from fp32 to fp16/bf16 - Tensor rP = flash::convert_type(scores); - // Reshape rP from (nrow=(2, MMA_M), ncol=(2, MMA_N)) to ((2, 2, 2), MMA_M, MMA_N / 2) - // if using m16n8k16 or ((2, 2, 1), MMA_M, MMA_N) if using m16n8k8. - Tensor tOrP = make_tensor(rP.data(), flash::convert_layout_rowcol_Aregs(rP.layout())); + // Convert acc_s from fp32 to fp16/bf16 + Tensor rP = flash::convert_type(acc_s); + // Reshape rP from (MMA=4, MMA_M, MMA_N) to ((4, 2), MMA_M, MMA_N / 2) + // if using m16n8k16 or (4, MMA_M, MMA_N) if using m16n8k8. + Tensor tOrP = make_tensor(rP.data(), flash::convert_layout_acc_Aregs(rP.layout())); - flash::gemm_A_in_regs(acc_o, tOrP, tOrVt, tOsVt, tiled_mma, smem_tiled_copy_V, smem_thr_copy_V); - // if (cute::thread0()) { print(scores); } + flash::gemm_rs(acc_o, tOrP, tOrVt, tOsVt, tiled_mma, smem_tiled_copy_V, smem_thr_copy_V); // This check is at the end of the loop since we always have at least 1 iteration if (n_masking_steps > 1 && n_block <= n_block_min) { @@ -912,7 +853,15 @@ inline __device__ void compute_attn_1rowblock_splitkv(const Params& params, cons flash::cp_async_wait<0>(); __syncthreads(); // Advance gV - tVgV.data() = tVgV.data() + (-int(kBlockN * params.v_row_stride)); + if (block_table == nullptr) { + tVgV.data() = tVgV.data() + (-int(kBlockN * params.v_row_stride)); + } else { + const int block_table_idx_cur = (n_block + 1) * kBlockN / params.page_block_size; + const int block_table_offset_cur = (n_block + 1) * kBlockN - block_table_idx_cur * params.page_block_size; + const int block_table_idx_next = n_block * kBlockN / params.page_block_size; + const int block_table_offset_next = n_block * kBlockN - block_table_idx_next * params.page_block_size; + tVgV.data() = tVgV.data() + (block_table[block_table_idx_next] - block_table[block_table_idx_cur]) * params.v_batch_stride + (block_table_offset_next - block_table_offset_cur) * params.v_row_stride; + } flash::copy(gmem_tiled_copy_QKV, tVgV, tVsV, tKVcKV, tKVpKV); cute::cp_async_fence(); @@ -924,51 +873,36 @@ inline __device__ void compute_attn_1rowblock_splitkv(const Params& params, cons __syncthreads(); if (n_block > n_block_min) { // Advance gK - tKgK.data() = tKgK.data() + (-int(kBlockN * params.k_row_stride)); + if (block_table == nullptr) { + tKgK.data() = tKgK.data() + (-int(kBlockN * params.k_row_stride)); + } else { + const int block_table_idx_cur = n_block * kBlockN / params.page_block_size; + const int block_table_offset_cur = n_block * kBlockN - block_table_idx_cur * params.page_block_size; + const int block_table_idx_next = (n_block - 1) * kBlockN / params.page_block_size; + const int block_table_offset_next = (n_block - 1) * kBlockN - block_table_idx_next * params.page_block_size; + tKgK.data() = tKgK.data() + (block_table[block_table_idx_next] - block_table[block_table_idx_cur]) * params.k_batch_stride + (block_table_offset_next - block_table_offset_cur) * params.k_row_stride; + } flash::copy(gmem_tiled_copy_QKV, tKgK, tKsK, tKVcKV, tKVpKV); // This cp_async_fence needs to be in the if block, otherwise the synchronization // isn't right and we get race conditions. cute::cp_async_fence(); } - // Reshape acc_s from (MMA=4, MMA_M, MMA_N) to (nrow=(2, MMA_M), ncol=(2, MMA_N)) - Tensor scores = make_tensor(acc_s.data(), flash::convert_layout_acc_rowcol(acc_s.layout())); - if (Is_local && n_block * kBlockN < (m_block + 1) * kBlockM + binfo.actual_seqlen_k - binfo.actual_seqlen_q + params.window_size_right) { - flash::apply_mask_local( - scores, n_block * kBlockN, binfo.actual_seqlen_k, - m_block * kBlockM + (tidx / 32) * 16 + (tidx % 32) / 4, - binfo.actual_seqlen_q, kNWarps * 16, - params.window_size_left, params.window_size_right); - } - softmax_rescale_o(scores, scores_max, scores_sum, acc_o, params.scale_softmax_log2); + mask.template apply_mask( + acc_s, n_block * kBlockN, m_block * kBlockM + (tidx / 32) * 16 + (tidx % 32) / 4, kNWarps * 16); + softmax.template softmax_rescale_o(acc_s, acc_o, params.scale_softmax_log2); - Tensor rP = flash::convert_type(scores); - // Reshape rP from (nrow=(2, MMA_M), ncol=(2, MMA_N)) to ((2, 2, 2), MMA_M, MMA_N / 2) - // if using m16n8k16 or ((2, 2, 1), MMA_M, MMA_N) if using m16n8k8. - Tensor tOrP = make_tensor(rP.data(), flash::convert_layout_rowcol_Aregs(rP.layout())); + Tensor rP = flash::convert_type(acc_s); + // Reshape rP from (MMA=4, MMA_M, MMA_N) to ((4, 2), MMA_M, MMA_N / 2) + // if using m16n8k16 or (4, MMA_M, MMA_N) if using m16n8k8. + Tensor tOrP = make_tensor(rP.data(), flash::convert_layout_acc_Aregs(rP.layout())); - flash::gemm_A_in_regs(acc_o, tOrP, tOrVt, tOsVt, tiled_mma, smem_tiled_copy_V, smem_thr_copy_V); + flash::gemm_rs(acc_o, tOrP, tOrVt, tOsVt, tiled_mma, smem_tiled_copy_V, smem_thr_copy_V); } // Epilogue - // Reshape acc_o from (MMA=4, MMA_M, MMA_K) to (nrow=(2, MMA_M), ncol=(2, MMA_K)) - Tensor acc_o_rowcol = make_tensor(acc_o.data(), flash::convert_layout_acc_rowcol(acc_o.layout())); - // if (cute::thread0()) { print(acc_o_rowcol); } - Tensor lse = make_fragment_like(scores_sum); -#pragma unroll - for (int mi = 0; mi < size<0>(acc_o_rowcol); ++mi) { - float sum = scores_sum(mi); - float inv_sum = (sum == 0.f || sum != sum) ? 1.f : 1.f / sum; - lse(mi) = (sum == 0.f || sum != sum) ? (Split ? -INFINITY : INFINITY) : scores_max(mi) * params.scale_softmax + __logf(sum); - float scale = inv_sum; -#pragma unroll - for (int ni = 0; ni < size<1>(acc_o_rowcol); ++ni) { - acc_o_rowcol(mi, ni) *= scale; - } - } - // if (cute::thread0()) { print(lse); } - // if (cute::thread0()) { print(acc_o_rowcol); } + Tensor lse = softmax.template normalize_softmax_lse(acc_o, params.scale_softmax); Tensor sOaccum = make_tensor(make_smem_ptr(reinterpret_cast(smem_)), typename Kernel_traits::SmemLayoutO{}); // (SMEM_M,SMEM_N) // Partition sO to match the accumulator partitioning @@ -1041,13 +975,11 @@ inline __device__ void compute_attn_1rowblock_splitkv(const Params& params, cons // Clear_OOB_K must be false since we don't want to write zeros to gmem flash::copy( gmem_tiled_copy_Oaccum, tOrOaccum, tOgOaccum, tOcO, tOpO, binfo.actual_seqlen_q - m_block * kBlockM); - // __syncthreads(); - // if (cute::thread0()) { print(tOgOaccum); } } //////////////////////////////////////////////////////////////////////////////////////////////////// -template +template inline __device__ void compute_attn(const Params& params) { const int m_block = blockIdx.x; // The block index for the batch. @@ -1063,12 +995,12 @@ inline __device__ void compute_attn(const Params& params) { // the attention matrix. This way, as long as we have the batch, head, and the location of // the 16 x 32 block within the attention matrix, we can generate the exact same dropout pattern. - flash::compute_attn_1rowblock(params, bidb, bidh, m_block); + flash::compute_attn_1rowblock(params, bidb, bidh, m_block); } //////////////////////////////////////////////////////////////////////////////////////////////////// -template +template inline __device__ void compute_attn_splitkv(const Params& params) { const int m_block = blockIdx.x; // The block index for the batch. @@ -1077,7 +1009,7 @@ inline __device__ void compute_attn_splitkv(const Params& params) { const int bidh = Split ? blockIdx.z - bidb * params.h : blockIdx.z; const int n_split_idx = Split ? blockIdx.y : 0; const int num_n_splits = Split ? gridDim.y : 1; - flash::compute_attn_1rowblock_splitkv(params, bidb, bidh, m_block, n_split_idx, num_n_splits); + flash::compute_attn_1rowblock_splitkv(params, bidb, bidh, m_block, n_split_idx, num_n_splits); } //////////////////////////////////////////////////////////////////////////////////////////////////// @@ -1121,7 +1053,7 @@ inline __device__ void combine_attn_seqk_parallel(const Params& params) { if (row < kMaxSplits) { sLSE[row][col] = lse; } - // if (bidx == 0 && tidx < 32) { printf("tidx = %d, row = %d, col = %d, lse = %f\n", tidx, row, col, lse_accum(l)); } + // if (bidx == 0 && tidx < 32) { printf("tidx = %d, row = %d, col = %d, lse = %f\n", tidx, row, col, lse); } } // if (bidx == 1 && tidx < 32) { printf("tidx = %d, row_offset_lse = %d, lse = %f\n", tidx, row_offset_lse, lse_accum(0)); } __syncthreads(); @@ -1129,7 +1061,7 @@ inline __device__ void combine_attn_seqk_parallel(const Params& params) { constexpr int kRowsPerLoadTranspose = std::min(kRowsPerLoadLSE, kMaxSplits); // To make sure that kMaxSplits is within 1 warp: we decide how many elements within kMaxSplits // each thread should hold. If kMaxSplits = 16, then each thread holds 2 elements (128 threads, - // 16 rows, so each time we load we can load 8 rows). + // kBlockM rows, so each time we load we can load 128 / kBlockM rows). // constexpr int kThreadsPerSplit = kMaxSplits / kRowsPerLoadTranspose; // static_assert(kThreadsPerSplit <= 32); static_assert(kRowsPerLoadTranspose <= 32); @@ -1218,11 +1150,11 @@ inline __device__ void combine_attn_seqk_parallel(const Params& params) { tOrO(i, m, k) += lse_scale * tOrOaccum(i, m, k); } } - // if (cute::thread0()) { printf("lse_scale = %f, %f\n", sLSE[split][0], sLSE[split][1]); print(tOrOaccum); print(tOrO); } + // if (cute::thread0()) { printf("lse_scale = %f, %f\n", sLSE[split][0], sLSE[split][1]); print(tOrOaccum); } } tOgOaccum.data() = tOgOaccum.data() + params.b * params.h * params.seqlen_q * params.d_rounded; } - // if (cute::thread0()) { print(tOrO); } + // if (cute::thread0()) { print_tensor(tOrO); } Tensor rO = flash::convert_type(tOrO); // Write to gO @@ -1256,4 +1188,6 @@ inline __device__ void combine_attn_seqk_parallel(const Params& params) { #if defined(__GNUC__) #pragma GCC diagnostic pop +#elif defined(_MSC_VER) +#pragma warning(pop) #endif diff --git a/onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_fwd_launch_template.h b/onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_fwd_launch_template.h index 87d189a803f8a..b1941df75be2c 100644 --- a/onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_fwd_launch_template.h +++ b/onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_fwd_launch_template.h @@ -10,33 +10,42 @@ namespace onnxruntime { namespace flash { -template -__global__ void flash_fwd_kernel(Flash_fwd_params params) { - static_assert(!(Is_causal && Is_local)); // If Is_local is true, Is_causal should be false +// Determine if the architecture supports FLASH and define a macro to handle parameter modifiers #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 - flash::compute_attn(params); +#define ARCH_SUPPORTS_FLASH +#define KERNEL_PARAM_MODIFIER __grid_constant__ #else - (void)params; +#define KERNEL_PARAM_MODIFIER +#endif + +// Define a macro for unsupported architecture handling to centralize the error message +#define FLASH_UNSUPPORTED_ARCH printf("FATAL: FlashAttention requires building with sm version sm80-sm90, but was built for < 8.0!"); + +// Use a macro to clean up kernel definitions +#define DEFINE_FLASH_FORWARD_KERNEL(kernelName, ...) \ + template \ + __global__ void kernelName(KERNEL_PARAM_MODIFIER const Flash_fwd_params params) + +DEFINE_FLASH_FORWARD_KERNEL(flash_fwd_kernel, bool Is_causal, bool Is_local, bool Has_alibi, bool Is_even_MN, bool Is_even_K, bool Return_softmax) { +#if defined(ARCH_SUPPORTS_FLASH) + static_assert(!(Is_causal && Is_local)); // Enforce constraints + flash::compute_attn(params); +#else + FLASH_UNSUPPORTED_ARCH #endif } -template -__global__ void flash_fwd_splitkv_kernel(Flash_fwd_params params) { -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 - flash::compute_attn_splitkv(params); +DEFINE_FLASH_FORWARD_KERNEL(flash_fwd_splitkv_kernel, bool Is_causal, bool Is_local, bool Has_alibi, bool Is_even_MN, bool Is_even_K, bool Split, bool Append_KV) { +#if defined(ARCH_SUPPORTS_FLASH) + flash::compute_attn_splitkv(params); #else - (void)params; + FLASH_UNSUPPORTED_ARCH #endif } -template -__global__ void flash_fwd_splitkv_combine_kernel(Flash_fwd_params params) { -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 +DEFINE_FLASH_FORWARD_KERNEL(flash_fwd_splitkv_combine_kernel, int kBlockM, int Log_max_splits, bool Is_even_K) { static_assert(Log_max_splits >= 1); flash::combine_attn_seqk_parallel(params); -#else - (void)params; -#endif } template @@ -52,25 +61,27 @@ void run_flash_fwd(Flash_fwd_params& params, cudaStream_t stream) { const bool is_even_MN = params.cu_seqlens_q == nullptr && params.cu_seqlens_k == nullptr && params.seqlen_k % Kernel_traits::kBlockN == 0 && params.seqlen_q % Kernel_traits::kBlockM == 0; const bool is_even_K = params.d == Kernel_traits::kHeadDim; BOOL_SWITCH(is_even_MN, IsEvenMNConst, [&] { - BOOL_SWITCH(is_even_K, IsEvenKConst, [&] { - BOOL_SWITCH(params.window_size_left >= 0 || params.window_size_right >= 0, Is_local, [&] { - // Will only return softmax if dropout, to reduce compilation time. - // If not IsEvenKConst, we also set IsEvenMNConst to false to reduce number of templates. - // If head dim > 128, set IsEvenMNConst to false to reduce number of templates - // If Is_local, set Is_causal to false - auto kernel = &flash_fwd_kernel < Kernel_traits, Is_causal && !Is_local, Is_local, IsEvenMNConst && IsEvenKConst && !Is_local && Kernel_traits::kHeadDim <= 128, IsEvenKConst, false > ; - // auto kernel = &flash_fwd_kernel; - if (smem_size >= 48 * 1024) { - cudaFuncSetAttribute( - kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size); - // ORT_ENFORCE(cudaFuncSetAttribute( - // kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); - } - // int ctas_per_sm; - // cudaError status_ = cudaOccupancyMaxActiveBlocksPerMultiprocessor( - // &ctas_per_sm, kernel, Kernel_traits::kNThreads, smem_size); - // printf("smem_size = %d, CTAs per SM = %d\n", int(smem_size), ctas_per_sm); - kernel<<>>(params); + EVENK_SWITCH(is_even_K, IsEvenKConst, [&] { + LOCAL_SWITCH((params.window_size_left >= 0 || params.window_size_right >= 0) && !Is_causal, Is_local, [&] { + ALIBI_SWITCH(params.alibi_slopes_ptr != nullptr, Has_alibi, [&] { + // Will only return softmax if dropout, to reduce compilation time. + // If not IsEvenKConst, we also set IsEvenMNConst to false to reduce number of templates. + // If head dim > 128, set IsEvenMNConst to false to reduce number of templates + // If Is_local, set Is_causal to false + auto kernel = &flash_fwd_kernel < Kernel_traits, Is_causal, Is_local && !Is_causal, Has_alibi, IsEvenMNConst && IsEvenKConst && !Is_local && Kernel_traits::kHeadDim <= 128, IsEvenKConst, false > ; + // auto kernel = &flash_fwd_kernel; + if (smem_size >= 48 * 1024) { + cudaFuncSetAttribute( + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast(smem_size)); + // ORT_ENFORCE(cudaFuncSetAttribute( + // kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); + } + // int ctas_per_sm; + // cudaError status_ = cudaOccupancyMaxActiveBlocksPerMultiprocessor( + // &ctas_per_sm, kernel, Kernel_traits::kNThreads, smem_size); + // printf("smem_size = %d, CTAs per SM = %d\n", int(smem_size), ctas_per_sm); + kernel<<(smem_size), stream>>>(params); + }); }); }); }); @@ -87,20 +98,25 @@ void run_flash_splitkv_fwd(Flash_fwd_params& params, cudaStream_t stream) { const bool is_even_K = params.d == Kernel_traits::kHeadDim; BOOL_SWITCH(params.is_causal, Is_causal, [&] { BOOL_SWITCH(is_even_MN, IsEvenMNConst, [&] { - BOOL_SWITCH(is_even_K, IsEvenKConst, [&] { - BOOL_SWITCH(params.window_size_left >= 0 || params.window_size_right >= 0, Is_local, [&] { - BOOL_SWITCH(params.num_splits > 1, Split, [&] { - BOOL_SWITCH(params.knew_ptr != nullptr, Append_KV, [&] { - // If Append_KV, then we must have seqlen_offsets, which means cu_seqlens_k != nullptr. - // printf("About to launch, Split = %d, Append_KV = %d, knew_ptr = %p\n", Split, Append_KV, params.knew_ptr); - auto kernel = &flash_fwd_splitkv_kernel < Kernel_traits, Is_causal && !Is_local, Is_local, IsEvenMNConst && !Append_KV && IsEvenKConst && !Is_local && Kernel_traits::kHeadDim <= 128, IsEvenKConst, Split, Append_KV > ; - // auto kernel = &flash_fwd_splitkv_kernel; - // auto kernel = &flash_fwd_splitkv_kernel; - if (smem_size >= 48 * 1024) { - cudaFuncSetAttribute( - kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size); - } - kernel<<>>(params); + EVENK_SWITCH(is_even_K, IsEvenKConst, [&] { + LOCAL_SWITCH((params.window_size_left >= 0 || params.window_size_right >= 0) && !Is_causal, Is_Local_Const, [&] { + BOOL_SWITCH(params.num_splits > 1, SplitConst, [&] { + BOOL_SWITCH(params.knew_ptr != nullptr, Append_KV_Const, [&] { + ALIBI_SWITCH(params.alibi_slopes_ptr != nullptr, Has_alibi, [&] { + // If Append_KV_Const, then we must have seqlen_offsets, which means cu_seqlens_k != nullptr. + // If not IsEvenKConst, we also set IsEvenMNConst to false to reduce number of templates. + // If Is_Local_Const, set Is_causal to false + auto kernel = &flash_fwd_splitkv_kernel < Kernel_traits, Is_causal, Is_Local_Const && !Is_causal, Has_alibi, + IsEvenMNConst && !Append_KV_Const && IsEvenKConst && !Is_Local_Const && Kernel_traits::kHeadDim <= 128, + IsEvenKConst, SplitConst, Append_KV_Const >; + // auto kernel = &flash_fwd_splitkv_kernel; + // auto kernel = &flash_fwd_splitkv_kernel; + if (smem_size >= 48 * 1024) { + cudaFuncSetAttribute( + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast(smem_size)); + } + kernel<<(smem_size), stream>>>(params); + }); }); }); }); @@ -113,7 +129,7 @@ void run_flash_splitkv_fwd(Flash_fwd_params& params, cudaStream_t stream) { // If headdim is divisible by 64, then we set kBlockM = 8, etc. constexpr static int kBlockM = Kernel_traits::kHeadDim % 128 == 0 ? 4 : (Kernel_traits::kHeadDim % 64 == 0 ? 8 : 16); dim3 grid_combine((params.b * params.h * params.seqlen_q + kBlockM - 1) / kBlockM); - BOOL_SWITCH(is_even_K, IsEvenKConst, [&] { + EVENK_SWITCH(is_even_K, IsEvenKConst, [&] { if (params.num_splits <= 2) { flash_fwd_splitkv_combine_kernel<<>>(params); } else if (params.num_splits <= 4) { @@ -135,8 +151,11 @@ void run_flash_splitkv_fwd(Flash_fwd_params& params, cudaStream_t stream) { template void run_mha_fwd_splitkv_dispatch(Flash_fwd_params& params, cudaStream_t stream) { - constexpr int kBlockM = 64; // Fixed for all head dimensions - constexpr int kBlockN = Headdim <= 64 ? 256 : (Headdim <= 128 ? 128 : 64); + constexpr static int kBlockM = 64; // Fixed for all head dimensions + // TD [2023-08-28]: nvcc segfaults for headdim 96 with block size 64 x 256, + // and for headdim 192 with block size 64 x 128. + // Also for headdim 160 with block size 64 x 128 after the rotary addition. + constexpr static int kBlockN = Headdim <= 64 ? 256 : (Headdim <= 128 ? 128 : 64); run_flash_splitkv_fwd>(params, stream); } @@ -163,8 +182,8 @@ void run_mha_fwd_hdim64(Flash_fwd_params& params, cudaStream_t stream) { template void run_mha_fwd_hdim96(Flash_fwd_params& params, cudaStream_t stream) { - constexpr int Headdim = 96; - const bool is_sm8x = params.dprops->major == 8 && params.dprops->minor > 0; + constexpr static int Headdim = 96; + bool is_sm8x = params.dprops->major == 8 && params.dprops->minor > 0; BOOL_SWITCH(params.is_causal, Is_causal, [&] { // For sm86 or sm89, 64 x 64 is the fastest for causal (because it's square), if (is_sm8x) { @@ -240,7 +259,7 @@ void run_mha_fwd_hdim160(Flash_fwd_params& params, cudaStream_t stream) { template void run_mha_fwd_hdim192(Flash_fwd_params& params, cudaStream_t stream) { - constexpr int Headdim = 192; + constexpr static int Headdim = 192; BOOL_SWITCH(params.is_causal, Is_causal, [&] { run_flash_fwd, Is_causal>(params, stream); // run_flash_fwd, Is_causal>(params, stream); @@ -254,7 +273,7 @@ void run_mha_fwd_hdim192(Flash_fwd_params& params, cudaStream_t stream) { template void run_mha_fwd_hdim224(Flash_fwd_params& params, cudaStream_t stream) { constexpr static int Headdim = 224; - int max_smem_per_block = params.dprops->sharedMemPerBlockOptin; + size_t max_smem_per_block = params.dprops->sharedMemPerBlockOptin; // printf("max_smem_per_block = %d\n", max_smem_per_block); BOOL_SWITCH(params.is_causal, Is_causal, [&] { if (max_smem_per_block >= 2 * Headdim * (128 + 2 * 64)) { // 112 KB diff --git a/onnxruntime/contrib_ops/cuda/bert/flash_attention/kernel_traits.h b/onnxruntime/contrib_ops/cuda/bert/flash_attention/kernel_traits.h index 52a4e56491c5e..cb08dbc853a91 100644 --- a/onnxruntime/contrib_ops/cuda/bert/flash_attention/kernel_traits.h +++ b/onnxruntime/contrib_ops/cuda/bert/flash_attention/kernel_traits.h @@ -25,7 +25,7 @@ struct Flash_kernel_traits { #endif using ElementAccum = float; - using index_t = uint32_t; + using index_t = int64_t; #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 using MMA_Atom_Arch = std::conditional_t< @@ -76,7 +76,6 @@ struct Flash_fwd_kernel_traits : public Base { typename Base::MMA_Atom_Arch, Layout, _1, _1>>, // 4x1x1 or 8x1x1 thread group Tile, _16, _16>>; - using SmemLayoutAtomQ = decltype(composition(Swizzle{}, // This has to be kBlockKSmem, using kHeadDim gives wrong results for d=128 Layout>, @@ -89,19 +88,9 @@ struct Flash_fwd_kernel_traits : public Base { SmemLayoutAtomQ{}, Shape, Int>{})); - // This has to be kBlockN and not 8, otherwise we get wrong results for d=128 - using SmemLayoutAtomVtransposedNoSwizzle = Layout, Int>, - Stride<_1, Int>>; - using SmemLayoutAtomVtransposed = decltype(composition(Swizzle{}, SmemLayoutAtomVtransposedNoSwizzle{})); - using SmemLayoutVtransposed = decltype(tile_to_shape( - SmemLayoutAtomVtransposed{}, - Shape, Int>{})); - // Maybe the VtransposeNoSwizzle just needs to have the right shape - // And the strides don't matter? - using SmemLayoutVtransposedNoSwizzle = decltype(tile_to_shape( - SmemLayoutAtomVtransposedNoSwizzle{}, - Shape, Int>{})); - // using SmemLayoutVtransposedNoSwizzle = decltype(SmemLayoutVtransposed{}.layout_fn()); + // https://github.com/ColfaxResearch/cutlass-kernels/blob/a222587e6d59b93ba704853d3946fb686d8b8892/src/fmha/fmha_forward.cu#L434 + using SmemLayoutVtransposed = decltype(composition(SmemLayoutKV{}, make_layout(Shape, Int>{}, GenRowMajor{}))); + using SmemLayoutVtransposedNoSwizzle = decltype(get_nonswizzle_portion(SmemLayoutVtransposed{})); using SmemLayoutAtomO = decltype(composition(Swizzle{}, Layout, Int>, @@ -112,10 +101,8 @@ struct Flash_fwd_kernel_traits : public Base { using SmemCopyAtomO = Copy_Atom; using SmemCopyAtomOaccum = Copy_Atom; - static constexpr int kSmemQCount = cute::size(SmemLayoutQ{}); - static constexpr int kSmemKVCount = cute::size(SmemLayoutKV{}) * 2; - static constexpr int kSmemQSize = kSmemQCount * sizeof(Element); - static constexpr int kSmemKVSize = kSmemKVCount * sizeof(Element); + static constexpr int kSmemQSize = size(SmemLayoutQ{}) * sizeof(Element); + static constexpr int kSmemKVSize = size(SmemLayoutKV{}) * 2 * sizeof(Element); static constexpr int kSmemSize = Share_Q_K_smem ? std::max(kSmemQSize, kSmemKVSize) : kSmemQSize + kSmemKVSize; static constexpr int kGmemElemsPerLoad = sizeof(cute::uint128_t) / sizeof(Element); @@ -127,8 +114,8 @@ struct Flash_fwd_kernel_traits : public Base { // to the same banks. static constexpr int kGmemThreadsPerRow = kBlockKSmem / kGmemElemsPerLoad; static_assert(kNThreads % kGmemThreadsPerRow == 0, "kNThreads must be a multiple of kGmemThreadsPerRow"); - using GmemLayoutAtom = cute::Layout, cute::Int>, - cute::Stride, _1>>; + using GmemLayoutAtom = Layout, Int>, + Stride, _1>>; // We use CACHEGLOBAL instead of CACHEALWAYS for both Q and K/V, since we won't be reading // from the same address by the same threadblock. This is slightly faster. @@ -136,27 +123,19 @@ struct Flash_fwd_kernel_traits : public Base { Has_cp_async, SM80_CP_ASYNC_CACHEGLOBAL, DefaultCopy>; - using GmemTiledCopyQKV = decltype(make_tiled_copy(Copy_Atom{}, + using GmemTiledCopyQKV = decltype(make_tiled_copy(Copy_Atom{}, GmemLayoutAtom{}, - cute::Layout>{})); // Val layout, 8 vals per read + Layout>{})); // Val layout, 8 vals per read using GmemTiledCopyO = decltype(make_tiled_copy(Copy_Atom{}, GmemLayoutAtom{}, - cute::Layout>{})); // Val layout, 8 vals per store - static constexpr int kGmemThreadsPerRowP = kBlockN / kGmemElemsPerLoad; - static_assert(kNThreads % kGmemThreadsPerRowP == 0, "kNThreads must be a multiple of kGmemThreadsPerRowP"); - using GmemLayoutAtomP = cute::Layout, cute::Int>, - cute::Stride, _1>>; - - using GmemTiledCopyP = decltype(make_tiled_copy(Copy_Atom{}, - GmemLayoutAtomP{}, - cute::Layout>{})); // Val layout, 8 vals per store + Layout>{})); // Val layout, 8 vals per store using GmemLayoutAtomOaccum = std::conditional_t< kBlockKSmem == 32, - cute::Layout, // Thread layout, 8 threads per row - cute::Stride<_8, _1>>, - cute::Layout, // Thread layout, 16 threads per row - cute::Stride<_16, _1>>>; + Layout, // Thread layout, 8 threads per row + Stride<_8, _1>>, + Layout, // Thread layout, 16 threads per row + Stride<_16, _1>>>; using GmemTiledCopyOaccum = decltype(make_tiled_copy(Copy_Atom{}, GmemLayoutAtomOaccum{}, Layout>{})); // Val layout, 4 vals per store @@ -205,84 +184,61 @@ struct Flash_bwd_kernel_traits : public Base { using TiledMmaSdP = TiledMMA< typename Base::MMA_Atom_Arch, - cute::Layout, cute::Int, _1>>, + Layout, Int, _1>>, Tile, Int<16 * kNWarps / AtomLayoutMSdP>, _16>>; using TiledMmadKV = TiledMMA< typename Base::MMA_Atom_Arch, - cute::Layout, cute::Int, _1>>, + Layout, Int, _1>>, Tile, Int<16 * kNWarps / AtomLayoutNdKV>, _16>>; using TiledMmadQ = TiledMMA< typename Base::MMA_Atom_Arch, - cute::Layout, cute::Int, _1>>, // 2x4x1 or 4x2x1 thread group + Layout, Int, _1>>, // 2x4x1 or 4x2x1 thread group Tile, Int<16 * kNWarps / AtomLayoutMdQ>, _16>>; using SmemLayoutAtomQdO = decltype(composition(Swizzle{}, - cute::Layout>, - cute::Stride, _1>>{})); + Layout>, + Stride, _1>>{})); using SmemLayoutQdO = decltype(tile_to_shape( SmemLayoutAtomQdO{}, - cute::make_shape(cute::Int{}, cute::Int{}))); + make_shape(Int{}, Int{}))); using SmemLayoutAtomKV = decltype(composition(Swizzle{}, - cute::Layout, cute::Int>, - cute::Stride, _1>>{})); + Layout, Int>, + Stride, _1>>{})); using SmemLayoutKV = decltype(tile_to_shape( // SmemLayoutAtomQdO{}, SmemLayoutAtomKV{}, - cute::make_shape(cute::Int{}, cute::Int{}))); - - using SmemLayoutAtomKtransposedNoSwizzle = Layout, Int>, - Stride<_1, Int>>; - using SmemLayoutAtomKtransposed = decltype(composition(Swizzle{}, SmemLayoutAtomKtransposedNoSwizzle{})); - using SmemLayoutKtransposed = decltype(tile_to_shape( - SmemLayoutAtomKtransposed{}, - make_shape(Int{}, Int{}))); - // Maybe the KtransposeNoSwizzle just needs to have the right shape - // And the strides don't matter? - using SmemLayoutKtransposedNoSwizzle = decltype(tile_to_shape( - SmemLayoutAtomKtransposedNoSwizzle{}, - make_shape(Int{}, Int{}))); - // using SmemLayoutKtransposedNoSwizzle = decltype(SmemLayoutKtransposed{}.layout_fn()); + make_shape(Int{}, Int{}))); + + using SmemLayoutKtransposed = decltype(composition(SmemLayoutKV{}, make_layout(Shape, Int>{}, GenRowMajor{}))); + using SmemLayoutKtransposedNoSwizzle = decltype(get_nonswizzle_portion(SmemLayoutKtransposed{})); // TODO: generalize to other values of kBlockN // TODO: what should be the Swizzle here? 3 is faster than 1, and 1 is faster than 2 // static constexpr int kPBlockN = kBlockN; - static_assert(kBlockN >= 64); + // Temporarily disabling this for hdim 256 on sm86 and sm89 + // static_assert(kBlockN >= 64); + static_assert(kBlockN >= 32); // TD [2023-03-19]: Idk why kPBlockN = 16 and kSwizzlePdS=3 is the fastest. - static constexpr int kPBlockN = 64; + static constexpr int kPBlockN = kBlockN >= 64 ? 64 : 32; static_assert(kPBlockN == 16 || kPBlockN == 32 || kPBlockN == 64); // static constexpr int kSwizzlePdS = kPBlockN == 16 ? 1 : (kPBlockN == 32 ? 2 : 3); static constexpr int kSwizzlePdS = 3; using SmemLayoutAtomPdS = decltype(composition(Swizzle{}, - cute::Layout, cute::Int>, - cute::Stride, _1>>{})); + Layout, Int>, + Stride, _1>>{})); using SmemLayoutPdS = decltype(tile_to_shape( SmemLayoutAtomPdS{}, - cute::make_shape(cute::Int{}, cute::Int{}))); - using SmemLayoutAtomPdStransposedNoSwizzle = Layout, Int>, - Stride<_1, Int>>; - using SmemLayoutAtomPdStransposed = decltype(composition(Swizzle{}, SmemLayoutAtomPdStransposedNoSwizzle{})); - using SmemLayoutPdStransposed = decltype(tile_to_shape( - SmemLayoutAtomPdStransposed{}, - make_shape(Int{}, Int{}))); - using SmemLayoutPdStransposedNoSwizzle = decltype(tile_to_shape( - SmemLayoutAtomPdStransposedNoSwizzle{}, - make_shape(Int{}, Int{}))); - // using SmemLayoutPdStransposedNoSwizzle = decltype(SmemLayoutPdStransposed{}.layout_fn()); + make_shape(Int{}, Int{}))); + using SmemLayoutPdStransposed = decltype(composition(SmemLayoutPdS{}, make_layout(Shape, Int>{}, GenRowMajor{}))); + using SmemLayoutPdStransposedNoSwizzle = decltype(get_nonswizzle_portion(SmemLayoutPdStransposed{})); + using SmemCopyAtomPdS = Copy_Atom; - using SmemLayoutAtomQdOtransposedNoSwizzle = Layout, Int>, - Stride<_1, Int>>; - using SmemLayoutAtomQdOtransposed = decltype(composition(Swizzle{}, SmemLayoutAtomQdOtransposedNoSwizzle{})); - using SmemLayoutQdOtransposed = decltype(tile_to_shape( - SmemLayoutAtomQdOtransposed{}, - make_shape(Int{}, Int{}))); - using SmemLayoutQdOtransposedNoSwizzle = decltype(tile_to_shape( - SmemLayoutAtomQdOtransposedNoSwizzle{}, - make_shape(Int{}, Int{}))); - // using SmemLayoutQdOtransposedNoSwizzle = decltype(SmemLayoutQdOtransposed{}.layout_fn()); + using SmemLayoutQdOtransposed = decltype(composition(SmemLayoutQdO{}, make_layout(Shape, Int>{}, GenRowMajor{}))); + using SmemLayoutQdOtransposedNoSwizzle = decltype(get_nonswizzle_portion(SmemLayoutQdOtransposed{})); using SmemLayoutAtomdKV = decltype(composition(Swizzle{}, Layout>, @@ -300,25 +256,18 @@ struct Flash_bwd_kernel_traits : public Base { make_shape(Int{}, Int{}))); using SmemCopyAtomdQ = Copy_Atom; - static constexpr int kSmemQdOCount = cute::size(SmemLayoutQdO{}) * (No_double_buffer ? 2 : 3); // Double buffer for sQ - static constexpr int kSmemKVCount = cute::size(SmemLayoutKV{}) * 2; - static constexpr int kSmemdSCount = cute::size(SmemLayoutPdS{}); - static constexpr int kSmemPCount = cute::size(SmemLayoutPdS{}); - static constexpr int kSmemdQCount = cute::size(SmemLayoutdQ{}); - // static constexpr int kSmemdPsumCount = kBlockM; - static constexpr int kSmemQdOSize = kSmemQdOCount * sizeof(Element); - static constexpr int kSmemKVSize = kSmemKVCount * sizeof(Element); - static constexpr int kSmemdSSize = kSmemdSCount * sizeof(Element); - static constexpr int kSmemPSize = kSmemPCount * sizeof(Element); - static constexpr int kSmemdQSize = kSmemdQCount * sizeof(Element); - // static constexpr int kSmemdPsumSize = kSmemdPsumCount * sizeof(ElementAccum); + // Double buffer for sQ + static constexpr int kSmemQdOSize = size(SmemLayoutQdO{}) * (No_double_buffer ? 2 : 3) * sizeof(Element); + static constexpr int kSmemKVSize = size(SmemLayoutKV{}) * 2 * sizeof(Element); + static constexpr int kSmemdSSize = size(SmemLayoutPdS{}) * sizeof(Element); + static constexpr int kSmemPSize = size(SmemLayoutPdS{}) * sizeof(Element); + static constexpr int kSmemdQSize = size(SmemLayoutdQ{}) * sizeof(Element); static constexpr int kSmemSize = kSmemQdOSize + (!Is_V_in_regs ? kSmemKVSize + kSmemdSSize + std::max(kSmemPSize, kSmemdQSize) : std::max(kSmemKVSize, kSmemKVSize / 2 + kSmemdSSize + std::max(kSmemPSize, kSmemdQSize))); static constexpr int kSmemSize1colblock = kSmemQdOSize + (!Is_V_in_regs ? kSmemKVSize + kSmemdSSize + kSmemPSize : std::max(kSmemKVSize, kSmemKVSize / 2 + kSmemdSSize + kSmemPSize)); - static constexpr int kSmemSize1rowblock = kSmemQdOSize / 3 * 2 + kSmemKVSize / 2 * 3 + kSmemdSSize + kSmemPSize; static constexpr int kGmemElemsPerLoad = sizeof(cute::uint128_t) / sizeof(Element); static_assert(kHeadDim % kGmemElemsPerLoad == 0, "kHeadDim must be a multiple of kGmemElemsPerLoad"); @@ -326,8 +275,8 @@ struct Flash_bwd_kernel_traits : public Base { // to affect speed in practice. static constexpr int kGmemThreadsPerRow = kBlockKSmem / kGmemElemsPerLoad; static_assert(kNThreads % kGmemThreadsPerRow == 0, "kNThreads must be a multiple of kGmemThreadsPerRow"); - using GmemLayoutAtom = cute::Layout, cute::Int>, - cute::Stride, _1>>; + using GmemLayoutAtom = Layout, Int>, + Stride, _1>>; // We use CACHEGLOBAL instead of CACHEALWAYS for both Q and K/V, since we won't be reading // from the same address by the same threadblock. This is slightly faster. @@ -337,30 +286,30 @@ struct Flash_bwd_kernel_traits : public Base { DefaultCopy>; using GmemTiledCopyQKV = decltype(make_tiled_copy(Copy_Atom{}, GmemLayoutAtom{}, - cute::Layout>{})); // Val layout, 8 vals per read + Layout>{})); // Val layout, 8 vals per read using GmemTiledCopydO = decltype(make_tiled_copy(Copy_Atom{}, GmemLayoutAtom{}, - cute::Layout>{})); // Val layout, 8 vals per store + Layout>{})); // Val layout, 8 vals per store using GmemTiledCopydKV = decltype(make_tiled_copy(Copy_Atom{}, GmemLayoutAtom{}, - cute::Layout>{})); // Val layout, 8 vals per store + Layout>{})); // Val layout, 8 vals per store using GmemTiledCopydQ = decltype(make_tiled_copy(Copy_Atom{}, GmemLayoutAtom{}, - cute::Layout>{})); // Val layout, 8 vals per store + Layout>{})); // Val layout, 8 vals per store using GmemLayoutAtomdQaccum = std::conditional_t< kBlockKSmem == 32, - cute::Layout, // Thread layout, 8 threads per row - cute::Stride<_8, _1>>, - cute::Layout, // Thread layout, 16 threads per row - cute::Stride<_16, _1>>>; + Layout, // Thread layout, 8 threads per row + Stride<_8, _1>>, + Layout, // Thread layout, 16 threads per row + Stride<_16, _1>>>; using GmemTiledCopydQaccum = decltype(make_tiled_copy(Copy_Atom{}, GmemLayoutAtomdQaccum{}, - cute::Layout>{})); // Val layout, 4 vals per store + Layout>{})); // Val layout, 4 vals per store using GmemTiledCopydQaccumAtomicAdd = decltype(make_tiled_copy(Copy_Atom{}, - cute::Layout, // Thread layout, 8 threads per row - cute::Stride<_32, _1>>{}, - cute::Layout>{})); // Val layout, 1 val per store + Layout, // Thread layout, 8 threads per row + Stride<_32, _1>>{}, + Layout>{})); // Val layout, 1 val per store }; } // namespace flash diff --git a/onnxruntime/contrib_ops/cuda/bert/flash_attention/mask.h b/onnxruntime/contrib_ops/cuda/bert/flash_attention/mask.h new file mode 100644 index 0000000000000..b225e5e3be559 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/bert/flash_attention/mask.h @@ -0,0 +1,208 @@ +/****************************************************************************** + * Copyright (c) 2024, Tri Dao. + ******************************************************************************/ + +#pragma once + +#include + +namespace onnxruntime { +namespace flash { + +using namespace cute; + +template +__forceinline__ __device__ void apply_mask(Tensor& tensor, const int max_seqlen_k, + const int col_idx_offset_ = 0) { + // tensor has shape (ncol=(2, MMA_M), nrow=(2, MMA_N)) + static_assert(Layout::rank == 2, "Only support 2D Tensor"); + const int lane_id = threadIdx.x % 32; + const int col_idx_offset = col_idx_offset_ + (lane_id % 4) * 2; +#pragma unroll + for (int nj = 0; nj < size<1, 1>(tensor); ++nj) { + const int col_idx_base = col_idx_offset + nj * 8; +#pragma unroll + for (int j = 0; j < size<1, 0>(tensor); ++j) { + const int col_idx = col_idx_base + j; + if (col_idx >= max_seqlen_k) { +// Without the "make_coord" we get wrong results +#pragma unroll + for (int mi = 0; mi < size<0>(tensor); ++mi) { + tensor(mi, make_coord(j, nj)) = -INFINITY; + } + } + } + } +} + +template +__forceinline__ __device__ void apply_mask_local(Tensor& tensor, const int col_idx_offset_, + const int max_seqlen_k, const int row_idx_offset, + const int max_seqlen_q, const int warp_row_stride, + const int window_size_left, const int window_size_right) { + // tensor has shape (ncol=(2, MMA_M), nrow=(2, MMA_N)) + static_assert(Layout::rank == 2, "Only support 2D Tensor"); + const int lane_id = threadIdx.x % 32; + const int col_idx_offset = col_idx_offset_ + (lane_id % 4) * 2; +#pragma unroll + for (int mi = 0; mi < size<0, 1>(tensor); ++mi) { + const int row_idx_base = row_idx_offset + mi * warp_row_stride; +#pragma unroll + for (int i = 0; i < size<0, 0>(tensor); ++i) { + const int row_idx = row_idx_base + i * 8; + const int col_idx_limit_left = std::max(0, row_idx + max_seqlen_k - max_seqlen_q - window_size_left); + const int col_idx_limit_right = std::min(max_seqlen_k, row_idx + 1 + max_seqlen_k - max_seqlen_q + window_size_right); +#pragma unroll + for (int nj = 0; nj < size<1, 1>(tensor); ++nj) { + const int col_idx_base = col_idx_offset + nj * 8; +#pragma unroll + for (int j = 0; j < size<1, 0>(tensor); ++j) { + const int col_idx = col_idx_base + j; + if (col_idx >= col_idx_limit_right || (HasWSLeft && col_idx < col_idx_limit_left)) { + tensor(make_coord(i, mi), make_coord(j, nj)) = -INFINITY; + } + } + } + // if (cute::thread0()) { + // printf("mi = %d, i = %d, row_idx = %d, max_seqlen_k = %d\n", mi, i, row_idx, max_seqlen_k); + // print(tensor(make_coord(i, mi), _)); + // // print(tensor(_, j + nj * size<1, 0>(tensor))); + // } + } + } +} + +template +__forceinline__ __device__ void apply_mask_causal(Tensor& tensor, const int col_idx_offset_, + const int max_seqlen_k, const int row_idx_offset, + const int max_seqlen_q, const int warp_row_stride) { + // Causal masking is equivalent to local masking with window_size_left = infinity and window_size_right = 0 + apply_mask_local(tensor, col_idx_offset_, max_seqlen_k, row_idx_offset, + max_seqlen_q, warp_row_stride, -1, 0); +} + +template +__forceinline__ __device__ void apply_mask_causal_w_idx( + Tensor& tensor, Tensor const& idx_rowcol, + const int col_idx_offset_, const int max_seqlen_k, const int row_idx_offset) { + // tensor has shape (ncol=(2, MMA_M), nrow=(2, MMA_N)) + static_assert(Layout0::rank == 2, "Only support 2D Tensor"); + static_assert(Layout1::rank == 2, "Only support 2D Tensor"); + CUTE_STATIC_ASSERT_V(size<0>(tensor) == size<0>(idx_rowcol)); + CUTE_STATIC_ASSERT_V(size<1>(tensor) == size<1>(idx_rowcol)); +#pragma unroll + for (int mi = 0; mi < size<0>(tensor); ++mi) { + const int col_idx_limit = std::min(max_seqlen_k, 1 + row_idx_offset + get<0>(idx_rowcol(mi, 0))); +#pragma unroll + for (int ni = 0; ni < size<1, 1>(tensor); ++ni) { + if (col_idx_offset_ + get<1>(idx_rowcol(0, ni)) >= col_idx_limit) { + tensor(mi, ni) = -INFINITY; + } + } + // if (cute::thread0()) { + // printf("ni = %d, j = %d, col_idx = %d, max_seqlen_k = %d\n", ni, j, col_idx, max_seqlen_k); + // print(tensor(_, make_coord(j, ni))); + // // print(tensor(_, j + ni * size<1, 0>(tensor))); + // } + } +} + +template +struct Mask { + const int max_seqlen_k, max_seqlen_q; + const int window_size_left, window_size_right; + const float alibi_slope; + + __forceinline__ __device__ Mask(const int max_seqlen_k, const int max_seqlen_q, + const int window_size_left, const int window_size_right, + const float alibi_slope = 0.f) + : max_seqlen_k(max_seqlen_k), max_seqlen_q(max_seqlen_q), window_size_left(window_size_left), window_size_right(window_size_right), alibi_slope(!Has_alibi ? 0.0 : alibi_slope){}; + + // Causal_mask: whether this particular iteration needs causal masking + template + __forceinline__ __device__ void apply_mask(Tensor& tensor_, + const int col_idx_offset_, + const int row_idx_offset, + const int warp_row_stride) { + static_assert(!(Causal_mask && Is_local), "Cannot be both causal and local"); + static_assert(Layout::rank == 3, "Only support 3D Tensor"); + static_assert(decltype(size<0>(tensor_))::value == 4, "First dimension must be 4"); + static constexpr bool Need_masking = Has_alibi || Causal_mask || Is_local || !Is_even_MN; + // if (cute::thread0()) { printf("Has_alibi = %d, Causal_mask=%d, Is_local=%d, Is_even_MN = %d, Need_masking = %d\n", Has_alibi, Causal_mask, Is_local, Is_even_MN, Need_masking); } + if constexpr (Need_masking) { + // Reshape tensor_ from (MMA=4, MMA_M, MMA_N) to (nrow=(2, MMA_M), ncol=(2, MMA_N)) + Tensor tensor = make_tensor(tensor_.data(), flash::convert_layout_acc_rowcol(tensor_.layout())); + // Do we need both row and column indices, or just column incides? + static constexpr bool Col_idx_only = !(Has_alibi && !Is_causal) && !Is_local && !Causal_mask; + const int lane_id = threadIdx.x % 32; + const int col_idx_offset = col_idx_offset_ + (lane_id % 4) * 2; + if constexpr (Col_idx_only) { +#pragma unroll + for (int nj = 0; nj < size<1, 1>(tensor); ++nj) { + const int col_idx_base = col_idx_offset + nj * 8; +#pragma unroll + for (int j = 0; j < size<1, 0>(tensor); ++j) { + const int col_idx = col_idx_base + j; +#pragma unroll + for (int mi = 0; mi < size<0>(tensor); ++mi) { + // No causal, no local + if constexpr (Has_alibi) { + tensor(mi, make_coord(j, nj)) += alibi_slope * col_idx; + } + if constexpr (!Is_even_MN) { + if (col_idx >= max_seqlen_k) { + tensor(mi, make_coord(j, nj)) = -INFINITY; + } + } + } + } + } + } else { +#pragma unroll + for (int mi = 0; mi < size<0, 1>(tensor); ++mi) { + const int row_idx_base = row_idx_offset + mi * warp_row_stride; +#pragma unroll + for (int i = 0; i < size<0, 0>(tensor); ++i) { + const int row_idx = row_idx_base + i * 8; + const int col_idx_limit_left = std::max(0, row_idx + max_seqlen_k - max_seqlen_q - window_size_left); + const int col_idx_limit_right = std::min(max_seqlen_k, row_idx + 1 + max_seqlen_k - max_seqlen_q + window_size_right); +#pragma unroll + for (int nj = 0; nj < size<1, 1>(tensor); ++nj) { + const int col_idx_base = col_idx_offset + nj * 8; +#pragma unroll + for (int j = 0; j < size<1, 0>(tensor); ++j) { + const int col_idx = col_idx_base + j; + if constexpr (Has_alibi) { + if constexpr (Is_causal) { + tensor(make_coord(i, mi), make_coord(j, nj)) += alibi_slope * col_idx; + } else { + tensor(make_coord(i, mi), make_coord(j, nj)) -= alibi_slope * abs(row_idx + max_seqlen_k - max_seqlen_q - col_idx); + } + } + if constexpr (Causal_mask) { + if (col_idx >= col_idx_limit_right) { + tensor(make_coord(i, mi), make_coord(j, nj)) = -INFINITY; + } + } + if constexpr (Is_local) { + if (col_idx >= col_idx_limit_right || col_idx < col_idx_limit_left) { + tensor(make_coord(i, mi), make_coord(j, nj)) = -INFINITY; + } + } + if constexpr (!Causal_mask && !Is_local && !Is_even_MN) { + // Causal and Local already handles MN masking + if (col_idx >= max_seqlen_k) { + tensor(make_coord(i, mi), make_coord(j, nj)) = -INFINITY; + } + } + } + } + } + } + } + } + }; +}; + +} // namespace flash +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/bert/flash_attention/rotary.h b/onnxruntime/contrib_ops/cuda/bert/flash_attention/rotary.h new file mode 100644 index 0000000000000..dfc14ab4b4406 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/bert/flash_attention/rotary.h @@ -0,0 +1,154 @@ +/****************************************************************************** + * Copyright (c) 2024, Tri Dao. + ******************************************************************************/ + +#pragma once + +#include + +#include "contrib_ops/cuda/bert/flash_attention/utils.h" + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +namespace onnxruntime { +namespace flash { + +using namespace cute; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +__forceinline__ __device__ void copy_rotary_interleaved(Tensor const& S, + Tensor& D, + Tensor const& Cos, + Tensor const& Sin, + Tensor const& identity_MN, + const int max_MN, const int min_MN, + const int dim, const int rotary_dim) { + CUTE_STATIC_ASSERT_V(rank(S) == Int<3>{}); + CUTE_STATIC_ASSERT_V(rank(D) == Int<3>{}); + CUTE_STATIC_ASSERT_V(size<0>(S) == size<0>(D)); // MMA + CUTE_STATIC_ASSERT_V(size<1>(S) == size<1>(D)); // MMA_M + CUTE_STATIC_ASSERT_V(size<2>(S) == size<2>(D)); // MMA_K + CUTE_STATIC_ASSERT_V(size<1>(S) == size<1>(Cos)); // MMA_M + CUTE_STATIC_ASSERT_V(size<2>(S) == size<2>(Cos)); // MMA_K + CUTE_STATIC_ASSERT_V(size<1>(S) == size<1>(Sin)); // MMA_M + CUTE_STATIC_ASSERT_V(size<2>(S) == size<2>(Sin)); // MMA_K + CUTE_STATIC_ASSERT_V(size<0>(Cos) == size<0>(Sin)); // MMA_K + static_assert(decltype(size<0>(S))::value == decltype(size<0>(Cos))::value * 2); + static_assert(decltype(size<0>(Cos))::value % 2 == 0); // Since we do fast conversion from fp16/bf16 to fp32 + Tensor rCos = make_fragment_like(Cos); + Tensor rSin = make_fragment_like(Sin); + Tensor rS = make_fragment_like(S); +#pragma unroll + for (int m = 0; m < size<1>(S); ++m) { + if (get<0>(identity_MN(0, m, 0)) >= min_MN && get<0>(identity_MN(0, m, 0)) < max_MN) { +#pragma unroll + for (int k = 0; k < size<2>(S); ++k) { + if (Is_even_K || get<1>(identity_MN(0, 0, k)) < dim) { + cute::copy(S(_, m, k), rS(_, m, k)); + if (get<1>(identity_MN(0, 0, k)) < rotary_dim) { + cute::copy(Cos(_, m, k), rCos(_, m, k)); + cute::copy(Sin(_, m, k), rSin(_, m, k)); + Tensor S_fp32 = convert_type(rS(_, m, k)); + Tensor cos_fp32 = convert_type(rCos(_, m, k)); + Tensor sin_fp32 = convert_type(rSin(_, m, k)); +#pragma unroll + for (int i = 0; i < size<0>(rS) / 2; ++i) { + float real = S_fp32(2 * i) * cos_fp32(i) - S_fp32(2 * i + 1) * sin_fp32(i); + float imag = S_fp32(2 * i) * sin_fp32(i) + S_fp32(2 * i + 1) * cos_fp32(i); + S_fp32(2 * i) = real; + S_fp32(2 * i + 1) = imag; + } + // Idk but I need to copy for the convert_type to work + Tensor S_fp32_copy = make_fragment_like(S_fp32); + cute::copy(S_fp32, S_fp32_copy); + using T = typename Engine0::value_type; + Tensor S_og_type = convert_type(S_fp32_copy); + cute::copy(S_og_type, rS(_, m, k)); + } + cute::copy(rS(_, m, k), D(_, m, k)); + } else if (Clear_OOB_K) { + cute::clear(D(_, m, k)); + } + } + } + } +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +__forceinline__ __device__ void copy_rotary_contiguous(Tensor const& S, + Tensor& D, + Tensor const& Cos, + Tensor const& Sin, + Tensor const& identity_MN, + const int max_MN, const int min_MN, + const int dim, const int rotary_dim) { + CUTE_STATIC_ASSERT_V(rank(S) == Int<3>{}); + CUTE_STATIC_ASSERT_V(rank(D) == Int<3>{}); + CUTE_STATIC_ASSERT_V(size<0>(S) == size<0>(D)); // MMA + CUTE_STATIC_ASSERT_V(size<1>(S) == size<1>(D)); // MMA_M + CUTE_STATIC_ASSERT_V(size<2>(S) == size<2>(D)); // MMA_K + CUTE_STATIC_ASSERT_V(size<1>(S) == size<1>(Cos)); // MMA_M + CUTE_STATIC_ASSERT_V(size<2>(S) == size<2>(Cos)); // MMA_K + CUTE_STATIC_ASSERT_V(size<1>(S) == size<1>(Sin)); // MMA_M + CUTE_STATIC_ASSERT_V(size<2>(S) == size<2>(Sin)); // MMA_K + CUTE_STATIC_ASSERT_V(size<0>(S) == size<0>(Cos)); // MMA + CUTE_STATIC_ASSERT_V(size<0>(Cos) == size<0>(Sin)); + static_assert(decltype(size<0>(Cos))::value % 2 == 0); // Since we do fast conversion from fp16/bf16 to fp32 + Tensor rCos = make_fragment_like(Cos); + Tensor rSin = make_fragment_like(Sin); + Tensor rS = make_fragment_like(S); + Tensor rS_other = make_fragment_like(rS(_, 0, 0)); +#pragma unroll + for (int m = 0; m < size<1>(S); ++m) { + if (get<0>(identity_MN(0, m, 0)) >= min_MN && get<0>(identity_MN(0, m, 0)) < max_MN) { +#pragma unroll + for (int k = 0; k < size<2>(S); ++k) { + if (Is_even_K || get<1>(identity_MN(0, 0, k)) < dim) { + cute::copy(S(_, m, k), rS(_, m, k)); + if (get<1>(identity_MN(0, 0, k)) < rotary_dim) { + const bool is_left = get<1>(identity_MN(0, 0, k)) < rotary_dim / 2; + Tensor gS_other = make_tensor(S(_, m, k).data() + (is_left ? rotary_dim / 2 : -rotary_dim / 2), S(_, m, k).layout()); + cute::copy(gS_other, rS_other); + // if (cute::thread0()) { print_tensor(rS(_, m, k)); print_tensor(rS_other); } + Tensor gCos = make_tensor(Cos(_, m, k).data() + (is_left ? 0 : -rotary_dim / 2), Cos(_, m, k).layout()); + Tensor gSin = make_tensor(Sin(_, m, k).data() + (is_left ? 0 : -rotary_dim / 2), Sin(_, m, k).layout()); + cute::copy(gCos, rCos(_, m, k)); + cute::copy(gSin, rSin(_, m, k)); + // if (cute::thread0()) { print_tensor(rCos(_, m, k)); print_tensor(rSin(_, m, k)); } + Tensor S_fp32 = convert_type(rS(_, m, k)); + Tensor S_other_fp32 = convert_type(rS_other); + Tensor cos_fp32 = convert_type(rCos(_, m, k)); + Tensor sin_fp32 = convert_type(rSin(_, m, k)); +#pragma unroll + for (int i = 0; i < size<0>(rS); ++i) { + S_fp32(i) = S_fp32(i) * cos_fp32(i) + S_other_fp32(i) * (is_left ? -sin_fp32(i) : sin_fp32(i)); + } + // Idk but I need to copy for the convert_type to work + Tensor S_fp32_copy = make_fragment_like(S_fp32); + cute::copy(S_fp32, S_fp32_copy); + using T = typename Engine0::value_type; + Tensor S_og_type = convert_type(S_fp32_copy); + cute::copy(S_og_type, rS(_, m, k)); + // if (cute::thread0()) { print_tensor(rS(_, m, k)); } + } + cute::copy(rS(_, m, k), D(_, m, k)); + } else if (Clear_OOB_K) { + cute::clear(D(_, m, k)); + } + } + } + } +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace flash +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/bert/flash_attention/softmax.h b/onnxruntime/contrib_ops/cuda/bert/flash_attention/softmax.h index 8017f83bbb01d..3c205378f0177 100644 --- a/onnxruntime/contrib_ops/cuda/bert/flash_attention/softmax.h +++ b/onnxruntime/contrib_ops/cuda/bert/flash_attention/softmax.h @@ -7,8 +7,7 @@ #include -#include -#include +#include #include "contrib_ops/cuda/bert/flash_attention/utils.h" @@ -20,7 +19,7 @@ using namespace cute; //////////////////////////////////////////////////////////////////////////////////////////////////// template -__device__ inline void thread_reduce_(Tensor const& tensor, Tensor& summary, Operator& op) { +__device__ __forceinline__ void thread_reduce_(Tensor const& tensor, Tensor& summary, Operator& op) { static_assert(Layout0::rank == 2, "Only support 2D Tensor"); static_assert(Layout1::rank == 1, "Only support 1D Tensor"); CUTE_STATIC_ASSERT_V(size<0>(summary) == size<0>(tensor)); @@ -35,7 +34,7 @@ __device__ inline void thread_reduce_(Tensor const& tensor, Te } template -__device__ inline void quad_allreduce_(Tensor& dst, Tensor& src, Operator& op) { +__device__ __forceinline__ void quad_allreduce_(Tensor& dst, Tensor& src, Operator& op) { CUTE_STATIC_ASSERT_V(size(dst) == size(src)); #pragma unroll for (int i = 0; i < size(dst); i++) { @@ -44,26 +43,26 @@ __device__ inline void quad_allreduce_(Tensor& dst, Tensor -__device__ inline void reduce_(Tensor const& tensor, Tensor& summary, Operator& op) { +__device__ __forceinline__ void reduce_(Tensor const& tensor, Tensor& summary, Operator& op) { thread_reduce_(tensor, summary, op); quad_allreduce_(summary, summary, op); } template -__device__ inline void reduce_max(Tensor const& tensor, Tensor& max) { +__device__ __forceinline__ void reduce_max(Tensor const& tensor, Tensor& max) { MaxOp max_op; reduce_(tensor, max, max_op); } -template -__device__ inline void reduce_sum(Tensor const& tensor, Tensor& sum) { +template +__device__ __forceinline__ void reduce_sum(Tensor const& tensor, Tensor& sum) { SumOp sum_op; - reduce_(tensor, sum, sum_op); + thread_reduce_(tensor, sum, sum_op); } // Apply the exp to all the elements. template -inline __device__ void scale_apply_exp2(Tensor& tensor, Tensor const& max, const float scale) { +__forceinline__ __device__ void scale_apply_exp2(Tensor& tensor, Tensor const& max, const float scale) { static_assert(Layout0::rank == 2, "Only support 2D Tensor"); static_assert(Layout1::rank == 1, "Only support 1D Tensor"); CUTE_STATIC_ASSERT_V(size<0>(max) == size<0>(tensor)); @@ -85,7 +84,7 @@ inline __device__ void scale_apply_exp2(Tensor& tensor, Tensor // Apply the exp to all the elements. template -inline __device__ void max_scale_exp2_sum(Tensor& tensor, Tensor& max, Tensor& sum, const float scale) { +__forceinline__ __device__ void max_scale_exp2_sum(Tensor& tensor, Tensor& max, Tensor& sum, const float scale) { static_assert(Layout0::rank == 2, "Only support 2D Tensor"); static_assert(Layout1::rank == 1, "Only support 1D Tensor"); CUTE_STATIC_ASSERT_V(size<0>(max) == size<0>(tensor)); @@ -115,103 +114,71 @@ inline __device__ void max_scale_exp2_sum(Tensor& tensor, Tens } } -template -inline __device__ void apply_mask(Tensor& tensor, const int max_seqlen_k, - const int col_idx_offset_ = 0) { - // tensor has shape (ncol=(2, MMA_M), nrow=(2, MMA_N)) - static_assert(Layout::rank == 2, "Only support 2D Tensor"); - const int lane_id = threadIdx.x % 32; - const int col_idx_offset = col_idx_offset_ + (lane_id % 4) * 2; -#pragma unroll - for (int nj = 0; nj < size<1, 1>(tensor); ++nj) { - const int col_idx_base = col_idx_offset + nj * 8; -#pragma unroll - for (int j = 0; j < size<1, 0>(tensor); ++j) { - const int col_idx = col_idx_base + j; - if (col_idx >= max_seqlen_k) { -// Without the "make_coord" we get wrong results -#pragma unroll - for (int mi = 0; mi < size<0>(tensor); ++mi) { - tensor(mi, make_coord(j, nj)) = -INFINITY; - } - } - } - } -} +//////////////////////////////////////////////////////////////////////////////////////////////////// -template -inline __device__ void apply_mask_local(Tensor& tensor, const int col_idx_offset_, - const int max_seqlen_k, const int row_idx_offset_, - const int max_seqlen_q, const int warp_row_stride, - const int window_size_left, const int window_size_right) { - // tensor has shape (ncol=(2, MMA_M), nrow=(2, MMA_N)) - static_assert(Layout::rank == 2, "Only support 2D Tensor"); - const int lane_id = threadIdx.x % 32; - // const int row_idx_offset = row_idx_offset_ + lane_id / 4; - const int row_idx_offset = row_idx_offset_; - const int col_idx_offset = col_idx_offset_ + (lane_id % 4) * 2; -#pragma unroll - for (int mi = 0; mi < size<0, 1>(tensor); ++mi) { - const int row_idx_base = row_idx_offset + mi * warp_row_stride; -#pragma unroll - for (int i = 0; i < size<0, 0>(tensor); ++i) { - const int row_idx = row_idx_base + i * 8; - const int col_idx_limit_left = std::max(0, row_idx + max_seqlen_k - max_seqlen_q - window_size_left); - const int col_idx_limit_right = std::min(max_seqlen_k, row_idx + 1 + max_seqlen_k - max_seqlen_q + window_size_right); -#pragma unroll - for (int nj = 0; nj < size<1, 1>(tensor); ++nj) { - const int col_idx_base = col_idx_offset + nj * 8; -#pragma unroll - for (int j = 0; j < size<1, 0>(tensor); ++j) { - const int col_idx = col_idx_base + j; - if (col_idx >= col_idx_limit_right || (HasWSLeft && col_idx < col_idx_limit_left)) { - tensor(make_coord(i, mi), make_coord(j, nj)) = -INFINITY; - } +template +struct Softmax { + using TensorT = decltype(make_tensor(Shape>{})); + TensorT row_max, row_sum; + + __forceinline__ __device__ Softmax(){}; + + template + __forceinline__ __device__ void softmax_rescale_o(Tensor0& acc_s, Tensor1& acc_o, float softmax_scale_log2) { + // Reshape acc_s from (MMA=4, MMA_M, MMA_N) to (nrow=(2, MMA_M), ncol=(2, MMA_N)) + Tensor scores = make_tensor(acc_s.data(), flash::convert_layout_acc_rowcol(acc_s.layout())); + static_assert(decltype(size<0>(scores))::value == kNRows); + if (Is_first) { + flash::template reduce_max(scores, row_max); + flash::scale_apply_exp2(scores, row_max, softmax_scale_log2); + flash::reduce_sum(scores, row_sum); + } else { + Tensor scores_max_prev = make_fragment_like(row_max); + cute::copy(row_max, scores_max_prev); + flash::template reduce_max(scores, row_max); + // Reshape acc_o from (MMA=4, MMA_M, MMA_K) to (nrow=(2, MMA_M), ncol=(2, MMA_K)) + Tensor acc_o_rowcol = make_tensor(acc_o.data(), flash::convert_layout_acc_rowcol(acc_o.layout())); + static_assert(decltype(size<0>(acc_o_rowcol))::value == kNRows); +#pragma unroll + for (int mi = 0; mi < size(row_max); ++mi) { + float scores_max_cur = !Check_inf + ? row_max(mi) + : (row_max(mi) == -INFINITY ? 0.0f : row_max(mi)); + float scores_scale = exp2f((scores_max_prev(mi) - scores_max_cur) * softmax_scale_log2); + row_sum(mi) *= scores_scale; +#pragma unroll + for (int ni = 0; ni < size<1>(acc_o_rowcol); ++ni) { + acc_o_rowcol(mi, ni) *= scores_scale; } } - // if (cute::thread0()) { - // printf("mi = %d, i = %d, row_idx = %d, max_seqlen_k = %d\n", mi, i, row_idx, max_seqlen_k); - // print(tensor(make_coord(i, mi), _)); - // // print(tensor(_, j + nj * size<1, 0>(tensor))); - // } + flash::scale_apply_exp2(scores, row_max, softmax_scale_log2); + // We don't do the reduce across threads here since we don't need to use the row_sum. + // We do that reduce at the end when we need to normalize the softmax. + flash::reduce_sum(scores, row_sum); } - } -} + }; -template -inline __device__ void apply_mask_causal(Tensor& tensor, const int col_idx_offset_, - const int max_seqlen_k, const int row_idx_offset_, - const int max_seqlen_q, const int warp_row_stride) { - // Causal masking is equivalent to local masking with window_size_left = infinity and window_size_right = 0 - apply_mask_local(tensor, col_idx_offset_, max_seqlen_k, row_idx_offset_, - max_seqlen_q, warp_row_stride, -1, 0); -} - -template -inline __device__ void apply_mask_causal_w_idx( - Tensor& tensor, Tensor const& idx_rowcol, - const int col_idx_offset_, const int max_seqlen_k, const int row_idx_offset_) { - // tensor has shape (ncol=(2, MMA_M), nrow=(2, MMA_N)) - static_assert(Layout0::rank == 2, "Only support 2D Tensor"); - static_assert(Layout1::rank == 2, "Only support 2D Tensor"); - CUTE_STATIC_ASSERT_V(size<0>(tensor) == size<0>(idx_rowcol)); - CUTE_STATIC_ASSERT_V(size<1>(tensor) == size<1>(idx_rowcol)); -#pragma unroll - for (int mi = 0; mi < size<0>(tensor); ++mi) { - const int col_idx_limit = std::min(max_seqlen_k, 1 + row_idx_offset_ + get<0>(idx_rowcol(mi, 0))); -#pragma unroll - for (int ni = 0; ni < size<1, 1>(tensor); ++ni) { - if (col_idx_offset_ + get<1>(idx_rowcol(0, ni)) >= col_idx_limit) { - tensor(mi, ni) = -INFINITY; + template + __forceinline__ __device__ TensorT normalize_softmax_lse(Tensor0& acc_o, float softmax_scale) { + SumOp sum_op; + quad_allreduce_(row_sum, row_sum, sum_op); + TensorT lse = make_fragment_like(row_sum); + Tensor acc_o_rowcol = make_tensor(acc_o.data(), flash::convert_layout_acc_rowcol(acc_o.layout())); + static_assert(decltype(size<0>(acc_o_rowcol))::value == kNRows); +#pragma unroll + for (int mi = 0; mi < size<0>(acc_o_rowcol); ++mi) { + float sum = row_sum(mi); + float inv_sum = (sum == 0.f || sum != sum) ? 1.f : 1.f / sum; + lse(mi) = (sum == 0.f || sum != sum) ? (Split ? -INFINITY : INFINITY) : row_max(mi) * softmax_scale + __logf(sum); + float scale = inv_sum; +#pragma unroll + for (int ni = 0; ni < size<1>(acc_o_rowcol); ++ni) { + acc_o_rowcol(mi, ni) *= scale; } } - // if (cute::thread0()) { - // printf("ni = %d, j = %d, col_idx = %d, max_seqlen_k = %d\n", ni, j, col_idx, max_seqlen_k); - // print(tensor(_, make_coord(j, ni))); - // // print(tensor(_, j + ni * size<1, 0>(tensor))); - // } - } -} + return lse; + }; +}; } // namespace flash } // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/bert/flash_attention/static_switch.h b/onnxruntime/contrib_ops/cuda/bert/flash_attention/static_switch.h index 5b70988949bbd..02bd7effd7da6 100644 --- a/onnxruntime/contrib_ops/cuda/bert/flash_attention/static_switch.h +++ b/onnxruntime/contrib_ops/cuda/bert/flash_attention/static_switch.h @@ -23,6 +23,37 @@ } \ }() +#define FLASHATTENTION_DISABLE_ALIBI // TEMP: Remove if we enable alibi +#ifdef FLASHATTENTION_DISABLE_ALIBI +#define ALIBI_SWITCH(COND, CONST_NAME, ...) \ + [&] { \ + constexpr static bool CONST_NAME = false; \ + return __VA_ARGS__(); \ + }() +#else +#define ALIBI_SWITCH BOOL_SWITCH +#endif + +#ifdef FLASHATTENTION_DISABLE_UNEVEN_K +#define EVENK_SWITCH(COND, CONST_NAME, ...) \ + [&] { \ + constexpr static bool CONST_NAME = true; \ + return __VA_ARGS__(); \ + }() +#else +#define EVENK_SWITCH BOOL_SWITCH +#endif + +#ifdef FLASHATTENTION_DISABLE_LOCAL +#define LOCAL_SWITCH(COND, CONST_NAME, ...) \ + [&] { \ + constexpr static bool CONST_NAME = false; \ + return __VA_ARGS__(); \ + }() +#else +#define LOCAL_SWITCH BOOL_SWITCH +#endif + #define FP16_SWITCH(COND, ...) \ [&] { \ if (COND) { \ @@ -34,7 +65,7 @@ } \ }() -#define FWD_HEADDIM_SWITCH(HEADDIM, ...) \ +#define HEADDIM_SWITCH(HEADDIM, ...) \ [&] { \ if (HEADDIM <= 32) { \ constexpr static int kHeadDim = 32; \ diff --git a/onnxruntime/contrib_ops/cuda/bert/flash_attention/utils.h b/onnxruntime/contrib_ops/cuda/bert/flash_attention/utils.h index 7aefd4799bc4d..9ef75120881e4 100644 --- a/onnxruntime/contrib_ops/cuda/bert/flash_attention/utils.h +++ b/onnxruntime/contrib_ops/cuda/bert/flash_attention/utils.h @@ -27,10 +27,10 @@ namespace flash { //////////////////////////////////////////////////////////////////////////////////////////////////// template -inline __device__ uint32_t relu2(const uint32_t x); +__forceinline__ __device__ uint32_t relu2(const uint32_t x); template <> -inline __device__ uint32_t relu2(const uint32_t x) { +__forceinline__ __device__ uint32_t relu2(const uint32_t x) { uint32_t res; const uint32_t zero = 0u; #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 @@ -52,7 +52,7 @@ inline __device__ uint32_t relu2(const uint32_t x) { #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 template <> -inline __device__ uint32_t relu2(const uint32_t x) { +__forceinline__ __device__ uint32_t relu2(const uint32_t x) { uint32_t res; const uint32_t zero = 0u; asm volatile("max.bf16x2 %0, %1, %2;\n" @@ -67,10 +67,10 @@ inline __device__ uint32_t relu2(const uint32_t x) { #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 template -inline __device__ uint32_t convert_relu2(const float2 x); +__forceinline__ __device__ uint32_t convert_relu2(const float2 x); template <> -inline __device__ uint32_t convert_relu2(const float2 x) { +__forceinline__ __device__ uint32_t convert_relu2(const float2 x) { uint32_t res; const uint32_t a = reinterpret_cast(x.x); const uint32_t b = reinterpret_cast(x.y); @@ -81,7 +81,7 @@ inline __device__ uint32_t convert_relu2(const float2 x) { } template <> -inline __device__ uint32_t convert_relu2(const float2 x) { +__forceinline__ __device__ uint32_t convert_relu2(const float2 x) { uint32_t res; const uint32_t a = reinterpret_cast(x.x); const uint32_t b = reinterpret_cast(x.y); @@ -97,20 +97,20 @@ inline __device__ uint32_t convert_relu2(const float2 x) { template struct MaxOp { - __device__ inline T operator()(T const& x, T const& y) { return x > y ? x : y; } + __device__ __forceinline__ T operator()(T const& x, T const& y) { return x > y ? x : y; } }; template <> struct MaxOp { // This is slightly faster - __device__ inline float operator()(float const& x, float const& y) { return max(x, y); } + __device__ __forceinline__ float operator()(float const& x, float const& y) { return max(x, y); } }; //////////////////////////////////////////////////////////////////////////////////////////////////// template struct SumOp { - __device__ inline T operator()(T const& x, T const& y) { return x + y; } + __device__ __forceinline__ T operator()(T const& x, T const& y) { return x + y; } }; //////////////////////////////////////////////////////////////////////////////////////////////////// @@ -119,7 +119,7 @@ template struct Allreduce { static_assert(THREADS == 32 || THREADS == 16 || THREADS == 8 || THREADS == 4); template - static __device__ inline T run(T x, Operator& op) { + static __device__ __forceinline__ T run(T x, Operator& op) { constexpr int OFFSET = THREADS / 2; x = op(x, __shfl_xor_sync(uint32_t(-1), x, OFFSET)); return Allreduce::run(x, op); @@ -131,7 +131,7 @@ struct Allreduce { template <> struct Allreduce<2> { template - static __device__ inline T run(T x, Operator& op) { + static __device__ __forceinline__ T run(T x, Operator& op) { x = op(x, __shfl_xor_sync(uint32_t(-1), x, 1)); return x; } @@ -143,10 +143,10 @@ template -inline __device__ void gemm(Tensor0& acc, Tensor1& tCrA, Tensor2& tCrB, Tensor3 const& tCsA, - Tensor4 const& tCsB, TiledMma tiled_mma, - TiledCopyA smem_tiled_copy_A, TiledCopyB smem_tiled_copy_B, - ThrCopyA smem_thr_copy_A, ThrCopyB smem_thr_copy_B) { +__forceinline__ __device__ void gemm(Tensor0& acc, Tensor1& tCrA, Tensor2& tCrB, Tensor3 const& tCsA, + Tensor4 const& tCsB, TiledMma tiled_mma, + TiledCopyA smem_tiled_copy_A, TiledCopyB smem_tiled_copy_B, + ThrCopyA smem_thr_copy_A, ThrCopyB smem_thr_copy_B) { CUTE_STATIC_ASSERT_V(size<1>(tCrA) == size<1>(acc)); // MMA_M CUTE_STATIC_ASSERT_V(size<1>(tCrB) == size<2>(acc)); // MMA_N CUTE_STATIC_ASSERT_V(size<2>(tCrA) == size<2>(tCrB)); // MMA_K @@ -178,9 +178,9 @@ inline __device__ void gemm(Tensor0& acc, Tensor1& tCrA, Tensor2& tCrB, Tensor3 template -inline __device__ void gemm_A_in_regs(Tensor0& acc, Tensor1& tCrA, Tensor2& tCrB, Tensor3 const& tCsB, - TiledMma tiled_mma, TiledCopy smem_tiled_copy_B, - ThrCopy smem_thr_copy_B) { +__forceinline__ __device__ void gemm_rs(Tensor0& acc, Tensor1& tCrA, Tensor2& tCrB, Tensor3 const& tCsB, + TiledMma tiled_mma, TiledCopy smem_tiled_copy_B, + ThrCopy smem_thr_copy_B) { CUTE_STATIC_ASSERT_V(size<1>(tCrA) == size<1>(acc)); // MMA_M CUTE_STATIC_ASSERT_V(size<1>(tCrB) == size<2>(acc)); // MMA_N CUTE_STATIC_ASSERT_V(size<2>(tCrA) == size<2>(tCrB)); // MMA_K @@ -200,42 +200,48 @@ inline __device__ void gemm_A_in_regs(Tensor0& acc, Tensor1& tCrA, Tensor2& tCrB // Convert acc_layout from (MMA=4, MMA_M, MMA_N) to (nrow=(2, MMA_M), ncol=(2, MMA_N)) template -inline __device__ auto convert_layout_acc_rowcol(Layout acc_layout) { +__forceinline__ __device__ auto convert_layout_acc_rowcol(Layout acc_layout) { static_assert(decltype(size<0>(acc_layout))::value == 4); static_assert(decltype(rank(acc_layout))::value == 3); auto l = logical_divide(acc_layout, Shape<_2>{}); // ((2, 2), MMA_M, MMA_N) - // TD [2023-08-13]: Idk why but get<0, 1>(l) doesn't work for Cutlass 3.2, I'm getting - // "int_tuple.hpp(74): error: conversion to inaccessible base class" - // return make_layout(make_layout(get<0, 1>(l), get<1>(l)), make_layout(get<0, 0>(l), get<2>(l))); - return make_layout(make_layout(get<1>(get<0>(l)), get<1>(l)), make_layout(get<0>(get<0>(l)), get<2>(l))); + return make_layout(make_layout(get<0, 1>(l), get<1>(l)), make_layout(get<0, 0>(l), get<2>(l))); }; //////////////////////////////////////////////////////////////////////////////////////////////////// -// Convert rowcol_layout from (nrow=(2, MMA_M), ncol=(2, MMA_N)) to ((2, 2, 2), MMA_M, MMA_N / 2) -// if using m16n8k16, or to ((2, 2, 1), MMA_M, MMA_N) if using m16n8k8. +// Convert acc_layout from (MMA=4, MMA_M, MMA_N) to ((4, 2), MMA_M, MMA_N / 2) +// if using m16n8k16, or to (4, MMA_M, MMA_N) if using m16n8k8. template -inline __device__ auto convert_layout_rowcol_Aregs(Layout rowcol_layout) { +__forceinline__ __device__ auto convert_layout_acc_Aregs(Layout acc_layout) { using X = Underscore; - static_assert(decltype(size<0, 0>(rowcol_layout))::value == 2); - static_assert(decltype(size<1, 0>(rowcol_layout))::value == 2); + static_assert(decltype(size<0>(acc_layout))::value == 4); + static_assert(decltype(rank(acc_layout))::value == 3); constexpr int mma_shape_K = get<2>(typename MMA_traits::Shape_MNK{}); static_assert(mma_shape_K == 8 || mma_shape_K == 16); - constexpr int MMA_N_divisor = mma_shape_K == 8 ? 1 : 2; - auto l = logical_divide(rowcol_layout, Shape>>{}); // ((2, MMA_M), (2, (2, MMA_N / 2))) - // TD [2023-08-13]: Same error as above on Cutlass 3.2 - // return make_layout(make_layout(get<1, 0>(l), get<0, 0>(l), get<1, 1, 0>(l)), - // get<0, 1>(l), - // get<1, 1, 1>(l)); - return make_layout(make_layout(get<0>(get<1>(l)), get<0>(get<0>(l)), get<0>(get<1>(get<1>(l)))), - get<1>(get<0>(l)), - get<1>(get<1>(get<1>(l)))); + if constexpr (mma_shape_K == 8) { + return acc_layout; + } else { + auto l = logical_divide(acc_layout, Shape{}); // (4, MMA_M, (2, MMA_N / 2))) + return make_layout(make_layout(get<0>(l), get<2, 0>(l)), get<1>(l), get<2, 1>(l)); + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// Convert acc_layout from (MMA=4, MMA_M, MMA_N) to ((4, 2), MMA_M, MMA_N / 2) +template +__forceinline__ __device__ auto convert_layout_acc_dropout(Layout acc_layout) { + using X = Underscore; + static_assert(decltype(size<0>(acc_layout))::value == 4); + static_assert(decltype(rank(acc_layout))::value == 3); + auto l = logical_divide(acc_layout, Shape{}); // (4, MMA_M, (2, MMA_N / 2))) + return make_layout(make_layout(get<0>(l), get<2, 0>(l)), get<1>(l), get<2, 1>(l)); }; //////////////////////////////////////////////////////////////////////////////////////////////////// template -inline __device__ auto convert_type(Tensor const& tensor) { +__forceinline__ __device__ auto convert_type(Tensor const& tensor) { using From_type = typename Engine::value_type; constexpr int numel = decltype(size(tensor))::value; cutlass::NumericArrayConverter convert_op; @@ -247,7 +253,7 @@ inline __device__ auto convert_type(Tensor const& tensor) { //////////////////////////////////////////////////////////////////////////////////////////////////// template -inline __device__ void relu_(Tensor& tensor) { +__forceinline__ __device__ void relu_(Tensor& tensor) { constexpr int numel = decltype(size(tensor))::value; static_assert(numel % 2 == 0); using value_t = typename Engine::value_type; @@ -263,7 +269,7 @@ inline __device__ void relu_(Tensor& tensor) { // On SM80 and above, we can fuse fp32 -> fp16/bf16 conversion and relu into 1 instruction template -inline __device__ auto convert_type_relu(Tensor const& tensor) { +__forceinline__ __device__ auto convert_type_relu(Tensor const& tensor) { using From_type = typename Engine::value_type; static_assert(std::is_same_v || std::is_same_v); static_assert(std::is_same_v); @@ -304,9 +310,9 @@ CUTE_HOST_DEVICE void cp_async_wait() { template -inline __device__ void copy(TiledCopy tiled_copy, Tensor const& S, - Tensor& D, Tensor const& identity_MN, - Tensor const& predicate_K, const int max_MN = 0) { +__forceinline__ __device__ void copy(TiledCopy tiled_copy, Tensor const& S, + Tensor& D, Tensor const& identity_MN, + Tensor const& predicate_K, const int max_MN = 0) { CUTE_STATIC_ASSERT_V(rank(S) == Int<3>{}); CUTE_STATIC_ASSERT_V(rank(D) == Int<3>{}); CUTE_STATIC_ASSERT_V(size<0>(S) == size<0>(D)); // MMA @@ -363,138 +369,5 @@ inline __device__ void copy_w_min_idx(Tensor const& S, //////////////////////////////////////////////////////////////////////////////////////////////////// -template -inline __device__ void copy_rotary_interleaved(Tensor const& S, - Tensor& D, - Tensor const& Cos, - Tensor const& Sin, - Tensor const& identity_MN, - const int max_MN, const int min_MN, - const int dim, const int rotary_dim) { - CUTE_STATIC_ASSERT_V(rank(S) == Int<3>{}); - CUTE_STATIC_ASSERT_V(rank(D) == Int<3>{}); - CUTE_STATIC_ASSERT_V(size<0>(S) == size<0>(D)); // MMA - CUTE_STATIC_ASSERT_V(size<1>(S) == size<1>(D)); // MMA_M - CUTE_STATIC_ASSERT_V(size<2>(S) == size<2>(D)); // MMA_K - CUTE_STATIC_ASSERT_V(size<1>(S) == size<1>(Cos)); // MMA_M - CUTE_STATIC_ASSERT_V(size<2>(S) == size<2>(Cos)); // MMA_K - CUTE_STATIC_ASSERT_V(size<1>(S) == size<1>(Sin)); // MMA_M - CUTE_STATIC_ASSERT_V(size<2>(S) == size<2>(Sin)); // MMA_K - CUTE_STATIC_ASSERT_V(size<0>(Cos) == size<0>(Sin)); // MMA_K - static_assert(decltype(size<0>(S))::value == decltype(size<0>(Cos))::value * 2); - static_assert(decltype(size<0>(Cos))::value % 2 == 0); // Since we do fast conversion from fp16/bf16 to fp32 - Tensor rCos = make_fragment_like(Cos); - Tensor rSin = make_fragment_like(Sin); - Tensor rS = make_fragment_like(S); -#pragma unroll - for (int m = 0; m < size<1>(S); ++m) { - if (get<0>(identity_MN(0, m, 0)) >= min_MN && get<0>(identity_MN(0, m, 0)) < max_MN) { -#pragma unroll - for (int k = 0; k < size<2>(S); ++k) { - if (Is_even_K || get<1>(identity_MN(0, 0, k)) < dim) { - cute::copy(S(_, m, k), rS(_, m, k)); - if (get<1>(identity_MN(0, 0, k)) < rotary_dim) { - cute::copy(Cos(_, m, k), rCos(_, m, k)); - cute::copy(Sin(_, m, k), rSin(_, m, k)); - Tensor S_fp32 = convert_type(rS(_, m, k)); - Tensor cos_fp32 = convert_type(rCos(_, m, k)); - Tensor sin_fp32 = convert_type(rSin(_, m, k)); -#pragma unroll - for (int i = 0; i < size<0>(rS) / 2; ++i) { - float real = S_fp32(2 * i) * cos_fp32(i) - S_fp32(2 * i + 1) * sin_fp32(i); - float imag = S_fp32(2 * i) * sin_fp32(i) + S_fp32(2 * i + 1) * cos_fp32(i); - S_fp32(2 * i) = real; - S_fp32(2 * i + 1) = imag; - } - // Idk but I need to copy for the convert_type to work - Tensor S_fp32_copy = make_fragment_like(S_fp32); - cute::copy(S_fp32, S_fp32_copy); - using T = typename Engine0::value_type; - Tensor S_og_type = convert_type(S_fp32_copy); - cute::copy(S_og_type, rS(_, m, k)); - } - cute::copy(rS(_, m, k), D(_, m, k)); - } else if (Clear_OOB_K) { - cute::clear(D(_, m, k)); - } - } - } - } -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -template -inline __device__ void copy_rotary_contiguous(Tensor const& S, - Tensor& D, - Tensor const& Cos, - Tensor const& Sin, - Tensor const& identity_MN, - const int max_MN, const int min_MN, - const int dim, const int rotary_dim) { - CUTE_STATIC_ASSERT_V(rank(S) == Int<3>{}); - CUTE_STATIC_ASSERT_V(rank(D) == Int<3>{}); - CUTE_STATIC_ASSERT_V(size<0>(S) == size<0>(D)); // MMA - CUTE_STATIC_ASSERT_V(size<1>(S) == size<1>(D)); // MMA_M - CUTE_STATIC_ASSERT_V(size<2>(S) == size<2>(D)); // MMA_K - CUTE_STATIC_ASSERT_V(size<1>(S) == size<1>(Cos)); // MMA_M - CUTE_STATIC_ASSERT_V(size<2>(S) == size<2>(Cos)); // MMA_K - CUTE_STATIC_ASSERT_V(size<1>(S) == size<1>(Sin)); // MMA_M - CUTE_STATIC_ASSERT_V(size<2>(S) == size<2>(Sin)); // MMA_K - CUTE_STATIC_ASSERT_V(size<0>(S) == size<0>(Cos)); // MMA - CUTE_STATIC_ASSERT_V(size<0>(Cos) == size<0>(Sin)); - static_assert(decltype(size<0>(Cos))::value % 2 == 0); // Since we do fast conversion from fp16/bf16 to fp32 - Tensor rCos = make_fragment_like(Cos); - Tensor rSin = make_fragment_like(Sin); - Tensor rS = make_fragment_like(S); - Tensor rS_other = make_fragment_like(rS(_, 0, 0)); -#pragma unroll - for (int m = 0; m < size<1>(S); ++m) { - if (get<0>(identity_MN(0, m, 0)) >= min_MN && get<0>(identity_MN(0, m, 0)) < max_MN) { -#pragma unroll - for (int k = 0; k < size<2>(S); ++k) { - if (Is_even_K || get<1>(identity_MN(0, 0, k)) < dim) { - cute::copy(S(_, m, k), rS(_, m, k)); - if (get<1>(identity_MN(0, 0, k)) < rotary_dim) { - const bool is_left = get<1>(identity_MN(0, 0, k)) < rotary_dim / 2; - Tensor gS_other = make_tensor(S(_, m, k).data() + (is_left ? rotary_dim / 2 : -rotary_dim / 2), S(_, m, k).layout()); - cute::copy(gS_other, rS_other); - // if (cute::thread0()) { print_tensor(rS(_, m, k)); print_tensor(rS_other); } - Tensor gCos = make_tensor(Cos(_, m, k).data() + (is_left ? 0 : -rotary_dim / 2), Cos(_, m, k).layout()); - Tensor gSin = make_tensor(Sin(_, m, k).data() + (is_left ? 0 : -rotary_dim / 2), Sin(_, m, k).layout()); - cute::copy(gCos, rCos(_, m, k)); - cute::copy(gSin, rSin(_, m, k)); - // if (cute::thread0()) { print_tensor(rCos(_, m, k)); print_tensor(rSin(_, m, k)); } - Tensor S_fp32 = convert_type(rS(_, m, k)); - Tensor S_other_fp32 = convert_type(rS_other); - Tensor cos_fp32 = convert_type(rCos(_, m, k)); - Tensor sin_fp32 = convert_type(rSin(_, m, k)); -#pragma unroll - for (int i = 0; i < size<0>(rS); ++i) { - S_fp32(i) = S_fp32(i) * cos_fp32(i) + S_other_fp32(i) * (is_left ? -sin_fp32(i) : sin_fp32(i)); - } - // Idk but I need to copy for the convert_type to work - Tensor S_fp32_copy = make_fragment_like(S_fp32); - cute::copy(S_fp32, S_fp32_copy); - using T = typename Engine0::value_type; - Tensor S_og_type = convert_type(S_fp32_copy); - cute::copy(S_og_type, rS(_, m, k)); - // if (cute::thread0()) { print_tensor(rS(_, m, k)); } - } - cute::copy(rS(_, m, k), D(_, m, k)); - } else if (Clear_OOB_K) { - cute::clear(D(_, m, k)); - } - } - } - } -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - } // namespace flash } // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc b/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc index 0c26f04edef99..3b6ad238cc826 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc @@ -141,7 +141,7 @@ Status GroupQueryAttention::ComputeInternal(OpKernelContext* context) const { auto [num_splits, slse_accum_bytes, o_accum_bytes] = onnxruntime::flash::get_num_splits_and_buffer_sizes( parameters.batch_size, parameters.sequence_length, parameters.sequence_length, parameters.num_heads, parameters.head_size, device_prop.multiProcessorCount); - parameters.num_splits = num_splits; + parameters.num_splits = static_cast(num_splits); softmax_lse_accum_bytes = slse_accum_bytes; out_accum_bytes = o_accum_bytes; } diff --git a/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu index 62974d12003fe..c38929697f3cb 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu @@ -675,7 +675,7 @@ Status FlashAttention( bool past_bsnh = past_kv_format == AttentionQkvFormat::Q_K_V_BSNH; ORT_RETURN_IF_ERROR(onnxruntime::flash::mha_fwd_kvcache( device_prop, stream, query, present_key, present_value, key, value, data.output, - reinterpret_cast(data.softmax_lse), seqlens_k, cos_cache, sin_cache, + reinterpret_cast(data.softmax_lse), seqlens_k, cos_cache, sin_cache, /*block_table*/ nullptr, batch_size, num_heads, kv_num_heads, head_size, sequence_length, parameters.seqlen_present_kv_cache, kv_sequence_length, parameters.rotary_dim, scale, is_causal, is_bf16, past_bsnh, parameters.num_splits, reinterpret_cast(data.softmax_lse_accum), diff --git a/onnxruntime/contrib_ops/cuda/bert/multihead_attention.cc b/onnxruntime/contrib_ops/cuda/bert/multihead_attention.cc index 5ae7c149fa05c..ba8b00df07e06 100644 --- a/onnxruntime/contrib_ops/cuda/bert/multihead_attention.cc +++ b/onnxruntime/contrib_ops/cuda/bert/multihead_attention.cc @@ -166,7 +166,7 @@ Status MultiHeadAttention::ComputeInternal(OpKernelContext* context) const { auto [num_splits, slse_accum_bytes, o_accum_bytes] = onnxruntime::flash::get_num_splits_and_buffer_sizes( parameters.batch_size, parameters.sequence_length, parameters.kv_sequence_length, parameters.num_heads, parameters.head_size, device_prop.multiProcessorCount); - parameters.num_splits = num_splits; + parameters.num_splits = static_cast(num_splits); softmax_lse_accum_bytes = slse_accum_bytes; out_accum_bytes = o_accum_bytes; } diff --git a/onnxruntime/contrib_ops/cuda/bert/packed_multihead_attention_impl.cu b/onnxruntime/contrib_ops/cuda/bert/packed_multihead_attention_impl.cu index a18744d29b1db..3e168189be3d5 100644 --- a/onnxruntime/contrib_ops/cuda/bert/packed_multihead_attention_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/packed_multihead_attention_impl.cu @@ -631,6 +631,8 @@ Status FlashAttention( data.output, cu_seqlens_q, cu_seqlens_k, + nullptr, // seqused_k + nullptr, // block_table softmax_lse_buffer, batch_size, num_heads, @@ -640,7 +642,7 @@ Status FlashAttention( sequence_length, scale, false, // is causal - false // is bf16 + false // is bf16 )); DUMP_TENSOR_INIT(); diff --git a/onnxruntime/test/python/transformers/benchmark_gqa_windows.py b/onnxruntime/test/python/transformers/benchmark_gqa_windows.py new file mode 100644 index 0000000000000..b781ccf03f138 --- /dev/null +++ b/onnxruntime/test/python/transformers/benchmark_gqa_windows.py @@ -0,0 +1,221 @@ +import argparse +import os +import time +from typing import Optional + +import torch +from test_sparse_attention import GroupQueryAttentionConfig, OrtGroupQueryAttention + + +def save_results(results, filename): + import pandas as pd + + df = pd.DataFrame( + results, + columns=[ + "Inference Interval (ms)", + "Throughput (samples/second)", + "Batch Size", + "Max Sequence Length", + "Sequence Length", + "Past Sequence Length", + "Model Name", + ], + ) + # df = df.transpose() # This line swaps the rows and columns + df.to_csv(filename, header=True, index=False) + print(f"Results saved in {filename}!") + + +def benchmark( + batch_size: int, + num_heads: int, + kv_num_heads: int, + head_size: int, + max_seq_len: int, + sequence_length: int = 1, + past_sequence_length: int = 0, + local_window_size: Optional[int] = None, + model_name: str = "Llama3-8B", +): + warmup = 15 + repeat = 100 + + config: GroupQueryAttentionConfig = GroupQueryAttentionConfig( + batch_size=batch_size, + sequence_length=sequence_length, + max_sequence_length=max_seq_len, + past_sequence_length=past_sequence_length, + num_heads=num_heads, + kv_num_heads=kv_num_heads, + head_size=head_size, + local_window_size=local_window_size if local_window_size else -1, + do_rotary=True, # Most models use rotary positional embeddings + is_packed_qkv=model_name in ["Phi-3-mini-128k", "Phi-3-small-128k"], + device="cuda", + ) + + obj = OrtGroupQueryAttention(config) + + for _ in range(warmup): + obj.infer() + + intervals = [] + for _ in range(repeat): + infer_start = time.perf_counter() + obj.infer() + infer_interval = time.perf_counter() - infer_start + intervals.append(infer_interval) + avg_infer_interval = sum(intervals) / len(intervals) + avg_infer_interval_ms = avg_infer_interval * 1000 + print(f"Average inference interval: {avg_infer_interval_ms:.6f} milliseconds") + avg_throughput = batch_size / avg_infer_interval + print(f"Average throughput: {avg_throughput:.6f} samples/second") + + return [avg_infer_interval_ms, avg_throughput] + + +def run_performance_tests(args): + device_id = torch.cuda.current_device() + memory_in_gb = torch.cuda.get_device_properties(device_id).total_memory / (1024 * 1024 * 1024) + + configures = [ + (32, 128, 8, 8192, None, "Llama3-8B"), + (64, 128, 8, 8192, None, "Llama3-70B"), + (48, 128, 8, 32768, None, "Mixtral-8x22B-v0.1"), + (32, 96, 32, 131072, None, "Phi-3-mini-128k"), + (32, 128, 8, 65536, None, "Phi-3-small-128k"), # Sparsity is not used in this test + (40, 128, 10, 32768, None, "Phi-3-medium-128K"), + ] + if args.kernel == "flash_attention": + configures.append((32, 128, 8, 32768, 4096, "Mistral-7B-v0.1")) + + # Reduce max sequence length when GPU memory is not enough. + threshold = 131072 if memory_in_gb > 24 else 65536 if memory_in_gb > 12 else 32768 + + all_metrics = [] + for num_heads, head_size, kv_num_heads, max_seq_len, local_window_size, model_name in configures: + prompt_metrics_model = [] + token_metrics_model = [] + for batch_size in [1, 4]: + # Benchmark prompt + for sequence_length in [ + 1, + 4, + 8, + 16, + 32, + 64, + 128, + 256, + 512, + 1024, + 2048, + 4096, + 8192, + 16384, + 32768, + 65536, + 131072, + ]: + if sequence_length >= min(threshold, max_seq_len): + continue + print( + f"Prompt: batch_size={batch_size}, num_heads={num_heads}, kv_num_heads={kv_num_heads}, head_size={head_size}, sequence_length={sequence_length}, max_seq_len={max_seq_len}, local_window_size={local_window_size}, model_name={model_name}" + ) + metrics = benchmark( + batch_size=batch_size, + num_heads=num_heads, + kv_num_heads=kv_num_heads, + head_size=head_size, + sequence_length=sequence_length, + max_seq_len=min(threshold, max_seq_len), + local_window_size=local_window_size, + model_name=model_name, + ) + metrics = [*metrics, batch_size, max_seq_len, sequence_length, 0, model_name] + prompt_metrics_model.append(metrics) + all_metrics.append(metrics) + # Benchmark token + for past_sequence_length in [ + 0, + 3, + 7, + 15, + 31, + 63, + 127, + 255, + 511, + 1023, + 2047, + 4095, + 8191, + 16383, + 32767, + 65535, + 131071, + ]: + if past_sequence_length >= min(threshold, max_seq_len): + continue + print( + f"Token: batch_size={batch_size}, num_heads={num_heads}, kv_num_heads={kv_num_heads}, head_size={head_size}, past_sequence_length={past_sequence_length}, max_seq_len={max_seq_len}, local_window_size={local_window_size}, model_name={model_name}" + ) + metrics = benchmark( + batch_size=batch_size, + num_heads=num_heads, + kv_num_heads=kv_num_heads, + head_size=head_size, + past_sequence_length=past_sequence_length, + max_seq_len=min(threshold, max_seq_len), + local_window_size=local_window_size, + model_name=model_name, + ) + metrics = [*metrics, batch_size, max_seq_len, 1, past_sequence_length, model_name] + token_metrics_model.append(metrics) + all_metrics.append(metrics) + # Calculate average inference interval and throughput for each model + avg_prompt_infer_interval = sum([metrics[0] for metrics in prompt_metrics_model]) / len(prompt_metrics_model) + avg_prompt_throughput = sum([metrics[1] for metrics in prompt_metrics_model]) / len(prompt_metrics_model) + avg_token_infer_interval = sum([metrics[0] for metrics in token_metrics_model]) / len(token_metrics_model) + avg_token_throughput = sum([metrics[1] for metrics in token_metrics_model]) / len(token_metrics_model) + print(f"Average {model_name} prompt inference interval: {avg_prompt_infer_interval:.6f} milliseconds") + print(f"Average {model_name} prompt throughput: {avg_prompt_throughput:.6f} samples/second") + print(f"Average {model_name} token inference interval: {avg_token_infer_interval:.6f} milliseconds") + print(f"Average {model_name} token throughput: {avg_token_throughput:.6f} samples/second") + all_metrics.append( + [avg_prompt_infer_interval, avg_prompt_throughput, 0, max_seq_len, 0, 0, model_name + " (Average Prompt)"] + ) + all_metrics.append( + [avg_token_infer_interval, avg_token_throughput, 0, max_seq_len, 0, 0, model_name + " (Average Token)"] + ) + + save_results(all_metrics, args.output) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="End-to-end benchmarking for gen-ai") + parser.add_argument( + "-o", + "--output", + type=str, + default="benchmark_results.csv", + help="Output CSV file name or path (with .csv extension)", + ) + parser.add_argument( + "-k", + "--kernel", + type=str, + default="flash_attention", + help="GQA Kernel to use for benchmarking. Options: flash_attention, memory_efficient", + ) + args = parser.parse_args() + + if args.kernel == "memory_efficient": + os.environ["ORT_DISABLE_FLASH_ATTENTION"] = "1" + else: + os.environ["ORT_DISABLE_FLASH_ATTENTION"] = "0" + + s = torch.cuda.Stream() + with torch.cuda.stream(s), torch.no_grad(): + run_performance_tests(args) diff --git a/onnxruntime/test/python/transformers/test_flash_attn_cuda.py b/onnxruntime/test/python/transformers/test_flash_attn_cuda.py index 1d49583b3f20c..84bf30b65a742 100644 --- a/onnxruntime/test/python/transformers/test_flash_attn_cuda.py +++ b/onnxruntime/test/python/transformers/test_flash_attn_cuda.py @@ -20,6 +20,7 @@ from bert_padding import pad_input, unpad_input from einops import rearrange, repeat from onnx import TensorProto, helper +from packaging import version from parameterized import parameterized from onnxruntime import InferenceSession, OrtValue, SessionOptions @@ -1878,11 +1879,31 @@ def parity_check_gqa_past_no_buff( numpy.testing.assert_allclose(out, out_ref, rtol=rtol, atol=atol, equal_nan=True, err_msg=err_msg) +def has_flash_attention(): + if not torch.cuda.is_available(): + return False + major, _ = torch.cuda.get_device_capability() + return major >= 8 and ( + platform.system() == "Linux" + or (platform.system() == "Windows" and version.parse(torch.version.cuda) >= version.parse("12.0")) + ) + + +def has_memory_efficient(): + if not torch.cuda.is_available(): + return False + major, minor = torch.cuda.get_device_capability() + if major < 5 or (major == 5 and minor < 3): + return False + return True + + def packed_mha_test_cases(): batches = [2] if pipeline_mode else [1, 5] - seqs = [8, 97, 256, 1024] if pipeline_mode else [97, 128, 200, 256, 257, 384, 512, 768, 1024, 1025, 2048] + seqs = [1024, 1025] if pipeline_mode else [1024, 1025, 2048] num_h = [1, 3] if pipeline_mode else [1, 6, 16] h_sizes = [16, 256] if pipeline_mode else [32, 40, 64, 80, 96, 128, 160, 192, 224, 256] + for b in batches: for s in seqs: for n in num_h: @@ -1911,6 +1932,7 @@ def mha_test_cases(): ) num_h = [1, 3] if pipeline_mode else [1, 6, 16] h_sizes = [16, 256] if pipeline_mode else [32, 40, 64, 80, 96, 128, 160, 192, 224, 256] + for b in batches: for s, s2 in seqs: for n in num_h: @@ -1922,21 +1944,17 @@ def mha_test_cases(): class TestMHA(unittest.TestCase): @parameterized.expand(packed_mha_test_cases()) def test_packed_mha(self, _, config): - if not torch.cuda.is_available() or platform.system() != "Linux": - return - major, _ = torch.cuda.get_device_capability() - if major < 8: + if not has_flash_attention(): return + os.environ["ORT_DISABLE_FLASH_ATTENTION"] = "0" print("-------- TEST PACKED MHA ---------") parity_check_mha(config, True) @parameterized.expand(mha_test_cases()) def test_mha(self, _, config): - if not torch.cuda.is_available() or platform.system() != "Linux": - return - major, _ = torch.cuda.get_device_capability() - if major < 8: + if not has_flash_attention(): return + os.environ["ORT_DISABLE_FLASH_ATTENTION"] = "0" print("-------- TEST MHA ---------") parity_check_mha(config, False) @@ -2106,10 +2124,7 @@ def gqa_past_flash_attention_test_cases(): class TestGQA(unittest.TestCase): @parameterized.expand(gqa_no_past_memory_efficient_test_cases()) def test_gqa_no_past_memory_efficient(self, _, config, rotary, rotary_interleaved, packed): - if not torch.cuda.is_available(): - return - major, minor = torch.cuda.get_device_capability() - if major < 5 or (major == 5 and minor < 3): + if not has_memory_efficient(): return os.environ["ORT_DISABLE_FLASH_ATTENTION"] = "1" print("------- MEMORY EFFICIENT ATTENTION (PROMPT CASE) ---------") @@ -2135,10 +2150,7 @@ def test_gqa_no_past_memory_efficient(self, _, config, rotary, rotary_interleave @parameterized.expand(gqa_no_past_flash_attention_test_cases()) def test_gqa_no_past_flash_attention(self, _, config, local, rotary, rotary_interleaved, packed): - if not torch.cuda.is_available(): - return - major, _ = torch.cuda.get_device_capability() - if major < 8 or platform.system() != "Linux": + if not has_flash_attention(): return print("------- FLASH ATTENTION (PROMPT CASE) --------") os.environ["ORT_DISABLE_FLASH_ATTENTION"] = "0" @@ -2162,10 +2174,7 @@ def test_gqa_no_past_flash_attention(self, _, config, local, rotary, rotary_inte @parameterized.expand(gqa_past_memory_efficient_test_cases()) def test_gqa_past_memory_efficient(self, _, config, rotary, rotary_interleaved, packed): - if not torch.cuda.is_available(): - return - major, minor = torch.cuda.get_device_capability() - if major < 5 or (major == 5 and minor < 3): + if not has_memory_efficient(): return os.environ["ORT_DISABLE_FLASH_ATTENTION"] = "1" print("-------- MEMORY EFFICIENT (TOKEN GEN) --------") @@ -2191,10 +2200,7 @@ def test_gqa_past_memory_efficient(self, _, config, rotary, rotary_interleaved, @parameterized.expand(gqa_past_flash_attention_test_cases()) def test_gqa_past_flash_attention(self, _, config, local, rotary, rotary_interleaved, packed): - if not torch.cuda.is_available(): - return - major, _ = torch.cuda.get_device_capability() - if major < 8 or platform.system() != "Linux": + if not has_flash_attention(): return print("------- FLASH ATTENTION (TOKEN GEN) -------") os.environ["ORT_DISABLE_FLASH_ATTENTION"] = "0" diff --git a/onnxruntime/test/python/transformers/test_parity_moe.py b/onnxruntime/test/python/transformers/test_parity_moe.py index aa480a1af4587..be288d8b6e360 100644 --- a/onnxruntime/test/python/transformers/test_parity_moe.py +++ b/onnxruntime/test/python/transformers/test_parity_moe.py @@ -10,6 +10,7 @@ # license information. # ------------------------------------------------------------------------- +import platform import time import unittest @@ -375,6 +376,8 @@ def benchmark(self): class TestMoE(unittest.TestCase): def test_moe_small(self): + if platform.system() == "Windows": + pytest.skip("Skip on Windows") rt = MoE( batch_size=2, num_rows=8, diff --git a/tools/ci_build/github/azure-pipelines/win-gpu-ci-pipeline.yml b/tools/ci_build/github/azure-pipelines/win-gpu-ci-pipeline.yml index 291e2f4e19401..438e51175c5b4 100644 --- a/tools/ci_build/github/azure-pipelines/win-gpu-ci-pipeline.yml +++ b/tools/ci_build/github/azure-pipelines/win-gpu-ci-pipeline.yml @@ -44,7 +44,7 @@ stages: buildArch: x64 additionalBuildFlags: >- --enable_pybind --build_java --build_nodejs --use_cuda --cuda_home="$(Agent.TempDirectory)\v11.8" - --enable_cuda_profiling + --enable_cuda_profiling --enable_transformers_tool_test --cmake_extra_defines CMAKE_CUDA_ARCHITECTURES=86 --cmake_extra_defines onnxruntime_BUILD_UNIT_TESTS=ON --cmake_extra_defines onnxruntime_ENABLE_CUDA_EP_INTERNAL_TESTS=ON From 3a917e49fb09e0e6542370bfbcf13d64c4115d1e Mon Sep 17 00:00:00 2001 From: Wanming Lin Date: Tue, 25 Jun 2024 00:52:12 +0800 Subject: [PATCH 24/52] [WebNN EP] Support 4 more ops for TFLite backend (#21134) Recently WebNN TFLite backend supports gelu, expand, softsign, reciprocal. --- js/web/docs/webnn-operators.md | 8 ++++---- onnxruntime/core/providers/webnn/builders/helper.h | 8 ++++---- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/js/web/docs/webnn-operators.md b/js/web/docs/webnn-operators.md index 508f85377a3a2..a49759b9a93c5 100644 --- a/js/web/docs/webnn-operators.md +++ b/js/web/docs/webnn-operators.md @@ -29,11 +29,11 @@ operators and the supported opset domain/versions in **WebNN EP** by ONNX Runtim | Equal | ai.onnx(7-10, 11-12, 13-18, 19+) | equal | ✓ | ✓ | | | Erf | ai.onnx(7-9, 10-12, 13+) | erf | ✗ | ✓ | | | Exp | ai.onnx(7-12, 13+) | exp | ✓ | ✓ | | -| Expand | ai.onnx(8-12, 13+) | expand | ✗ | ✓ | 'shape' input should be a constant | +| Expand | ai.onnx(8-12, 13+) | expand | ✓ | ✓ | 'shape' input should be a constant | | Flatten | ai.onnx(7-8, 9-10, 11-12, 13-20, 21+) | reshape | ✓ | ✓ | | | Floor | ai.onnx(7-12, 13+) | floor | ✓ | ✓ | | | Gather | ai.onnx(7-10, 11-12, 13+) | gather | ✓ | ✓ | | -| Gelu | ai.onnx(20+) | gelu | ✗ | ✓ | | +| Gelu | ai.onnx(20+) | gelu | ✓ | ✓ | | | Gemm | ai.onnx(7-8, 9-10, 11-12, 13+) | gemm | ✓ | ✓ | Only supports 1-D 'C' input | | GlobalAveragePool | ai.onnx(7+) | averagePool2d | ✓ | ✓ | Only supports 4-D input | | GlobalMaxPool | ai.onnx(7+) | maxPool2d | ✓ | ✓ | Only supports 4-D input | @@ -60,7 +60,7 @@ operators and the supported opset domain/versions in **WebNN EP** by ONNX Runtim | Pad | ai.onnx(7-10, 11-12, 13-17, 18, 19-20, 21+) | pad | ✓ | ✓ | modes == 'wrap' is not supported | | Pow | ai.onnx(7-11, 12, 13-14, 15+) | pow | ✓ | ✓ | | | PRelu | ai.onnx(7-8, 9-15, 16+) | prelu | ✓ | ✓ | WebNN CPU backend restricts the last dimension of input and slope to be same (Chromium issue: https://issues.chromium.org/issues/335517470) | -| Reciprocal | ai.onnx(7-12, 13+) | reciprocal | ✗ | ✓ | | +| Reciprocal | ai.onnx(7-12, 13+) | reciprocal | ✓ | ✓ | | | ReduceL1 | ai.onnx(7-10, 11-12, 13-17, 18+) | reduceL1 | ✗ | ✓ | Input 'axes' if present should be a constant | | ReduceL2 | ai.onnx(7-10, 11-12, 13-17, 18+) | reduceL2 | ✗ | ✓ | Input 'axes' if present should be a constant | | ReduceLogSum| ai.onnx(7-10, 11-12, 13-17, 18+) | reduceLogSum| ✗ | ✓ | Input 'axes' if present should be a constant | @@ -77,7 +77,7 @@ operators and the supported opset domain/versions in **WebNN EP** by ONNX Runtim | Shape | ai.onnx(7-12, 13-14, 15-18, 19-20, 21+) | slice | ✓ | ✓ | | | Sigmoid | ai.onnx(7-12, 13+) | sigmoid | ✓ | ✓ | | | Softplus | ai.onnx(7+) | softplus | ✓ | ✓ | | -| Softsign | ai.onnx(7+) | softsign | ✗ | ✓ | | +| Softsign | ai.onnx(7+) | softsign | ✓ | ✓ | | | Sin | ai.onnx(7+) | sin | ✓ | ✓ | | | Slice | ai.onnx(7-9, 10, 11-12, 13+) | slice | ✓ | ✓ | Input 'starts', 'ends', 'axes', and 'steps' if present must be a constant, only supports 'steps' value 1 | | Softmax | ai.onnx(7-10, 11-12, 13+) | softmax | ✓ | ✓ | | diff --git a/onnxruntime/core/providers/webnn/builders/helper.h b/onnxruntime/core/providers/webnn/builders/helper.h index 7240fa37d9cc9..401d2eaa09129 100644 --- a/onnxruntime/core/providers/webnn/builders/helper.h +++ b/onnxruntime/core/providers/webnn/builders/helper.h @@ -176,11 +176,11 @@ static const InlinedHashMap op_map = { {"Equal", {"equal", true}}, {"Erf", {"erf", false}}, {"Exp", {"exp", true}}, - {"Expand", {"expand", false}}, + {"Expand", {"expand", true}}, {"Flatten", {"reshape", true}}, {"Floor", {"floor", true}}, {"Gather", {"gather", true}}, - {"Gelu", {"gelu", false}}, + {"Gelu", {"gelu", true}}, {"Gemm", {"gemm", true}}, {"GlobalAveragePool", {"averagePool2d", true}}, {"GlobalMaxPool", {"maxPool2d", true}}, @@ -208,7 +208,7 @@ static const InlinedHashMap op_map = { {"Pad", {"pad", true}}, {"Pow", {"pow", true}}, {"PRelu", {"prelu", true}}, - {"Reciprocal", {"reciprocal", false}}, + {"Reciprocal", {"reciprocal", true}}, {"ReduceL1", {"reduceL1", false}}, {"ReduceL2", {"reduceL2", false}}, {"ReduceLogSum", {"reduceLogSum", false}}, @@ -225,7 +225,7 @@ static const InlinedHashMap op_map = { {"Shape", {"slice", true}}, {"Sigmoid", {"sigmoid", true}}, {"Softplus", {"softplus", true}}, - {"Softsign", {"softsign", false}}, + {"Softsign", {"softsign", true}}, {"Sin", {"sin", true}}, {"Slice", {"slice", true}}, {"Softmax", {"softmax", true}}, From adaf0e81168a32040839e3bfea09ac69edcb011e Mon Sep 17 00:00:00 2001 From: mindest <30493312+mindest@users.noreply.github.com> Date: Tue, 25 Jun 2024 02:33:17 +0800 Subject: [PATCH 25/52] [Fix] USE_NCCL -> ORT_USE_NCCL (#21136) ### Description Correct the macro used when NCCL enabled. --- onnxruntime/contrib_ops/cuda/collective/custom_reduce_impl.cu | 2 +- onnxruntime/contrib_ops/cuda/collective/custom_reduce_impl.h | 2 +- onnxruntime/contrib_ops/cuda/collective/ipc_utils.cc | 2 +- onnxruntime/contrib_ops/cuda/collective/ipc_utils.h | 4 ++-- onnxruntime/contrib_ops/cuda/collective/nccl_kernels.h | 2 +- 5 files changed, 6 insertions(+), 6 deletions(-) diff --git a/onnxruntime/contrib_ops/cuda/collective/custom_reduce_impl.cu b/onnxruntime/contrib_ops/cuda/collective/custom_reduce_impl.cu index 666ec3a993235..9de3d48417b34 100644 --- a/onnxruntime/contrib_ops/cuda/collective/custom_reduce_impl.cu +++ b/onnxruntime/contrib_ops/cuda/collective/custom_reduce_impl.cu @@ -29,7 +29,7 @@ namespace onnxruntime { namespace cuda { namespace collective { -#if defined(USE_MPI) || defined(USE_NCCL) +#if defined(USE_MPI) || defined(ORT_USE_NCCL) using namespace onnxruntime; using namespace onnxruntime::cuda; diff --git a/onnxruntime/contrib_ops/cuda/collective/custom_reduce_impl.h b/onnxruntime/contrib_ops/cuda/collective/custom_reduce_impl.h index 3ca3c1dd166af..4721fb11ec86d 100644 --- a/onnxruntime/contrib_ops/cuda/collective/custom_reduce_impl.h +++ b/onnxruntime/contrib_ops/cuda/collective/custom_reduce_impl.h @@ -25,7 +25,7 @@ namespace onnxruntime { namespace cuda { namespace collective { -#if defined(USE_MPI) || defined(USE_NCCL) +#if defined(USE_MPI) || defined(ORT_USE_NCCL) constexpr size_t WARP_SIZE = 32; constexpr size_t MAX_ALL_REDUCE_BLOCKS = 24; diff --git a/onnxruntime/contrib_ops/cuda/collective/ipc_utils.cc b/onnxruntime/contrib_ops/cuda/collective/ipc_utils.cc index b4e602228ee63..fefad449bb5b5 100644 --- a/onnxruntime/contrib_ops/cuda/collective/ipc_utils.cc +++ b/onnxruntime/contrib_ops/cuda/collective/ipc_utils.cc @@ -23,7 +23,7 @@ namespace onnxruntime { namespace cuda { namespace collective { -#if defined(USE_MPI) || defined(USE_NCCL) +#if defined(USE_MPI) || defined(ORT_USE_NCCL) using namespace onnxruntime; diff --git a/onnxruntime/contrib_ops/cuda/collective/ipc_utils.h b/onnxruntime/contrib_ops/cuda/collective/ipc_utils.h index cda0f3437b25f..44f352168b411 100644 --- a/onnxruntime/contrib_ops/cuda/collective/ipc_utils.h +++ b/onnxruntime/contrib_ops/cuda/collective/ipc_utils.h @@ -24,7 +24,7 @@ namespace onnxruntime { namespace cuda { namespace collective { -#if defined(USE_MPI) || defined(USE_NCCL) +#if defined(USE_MPI) || defined(ORT_USE_NCCL) struct CudaDeleter { void operator()(void* ptr) const noexcept { @@ -86,4 +86,4 @@ GetCustomAllReduceWorkspace(int rank, int world_size, size_t input_size, IPCMemo } // namespace collective } // namespace cuda -} // namespace onnxruntime \ No newline at end of file +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/collective/nccl_kernels.h b/onnxruntime/contrib_ops/cuda/collective/nccl_kernels.h index 10b6f7dd56ef8..49646637b635e 100644 --- a/onnxruntime/contrib_ops/cuda/collective/nccl_kernels.h +++ b/onnxruntime/contrib_ops/cuda/collective/nccl_kernels.h @@ -5,7 +5,7 @@ #include "core/providers/cuda/cuda_kernel.h" -#if defined(ORT_USE_NCCL) || defined(ORT_USE_MPI) +#if defined(ORT_USE_NCCL) || defined(USE_MPI) #ifndef USE_ROCM #include "custom_reduce_impl.h" #include "ipc_utils.h" From f81c0ec32a5e239dae7904c02ff64660fbc3216d Mon Sep 17 00:00:00 2001 From: Jian Chen Date: Mon, 24 Jun 2024 16:46:21 -0700 Subject: [PATCH 26/52] Remove warning suppression from Java Packaging pipeline. (#21010) ### Description Remove warning suppression from Java Packaging pipeline. ### Motivation and Context We want the CI step not to produce warning. --- java/build.gradle | 3 +++ .../templates/make_java_win_binaries.yml | 14 +++++++++----- .../github/azure-pipelines/templates/win-ci.yml | 6 +++--- 3 files changed, 15 insertions(+), 8 deletions(-) diff --git a/java/build.gradle b/java/build.gradle index cebf67e085446..3219b082994ff 100644 --- a/java/build.gradle +++ b/java/build.gradle @@ -166,11 +166,14 @@ if (cmakeBuildDir != null) { } tasks.register('cmakeCheck', Copy) { + group = 'verification' from layout.buildDirectory.get() include 'reports/**' into cmakeBuildOutputDir dependsOn(check) } +} else { + println "cmakeBuildDir is not set. Skipping cmake tasks." } dependencies { diff --git a/tools/ci_build/github/azure-pipelines/templates/make_java_win_binaries.yml b/tools/ci_build/github/azure-pipelines/templates/make_java_win_binaries.yml index 9a666155028cc..0d62ed7907a67 100644 --- a/tools/ci_build/github/azure-pipelines/templates/make_java_win_binaries.yml +++ b/tools/ci_build/github/azure-pipelines/templates/make_java_win_binaries.yml @@ -9,12 +9,16 @@ parameters: steps: - task: CmdLine@2 displayName: 'Gradle cmakeCheck' - continueOnError: ${{ parameters.buildOnly }} inputs: - script: | - @echo on - call gradlew.bat cmakeCheck -DcmakeBuildDir=$(Build.BinariesDirectory)\RelWithDebInfo --warning-mode all - workingDirectory: $(Build.SourcesDirectory)\java + ${{ if eq(parameters.buildOnly, true) }}: + script: | + call gradlew.bat testClasses -DcmakeBuildDir=$(Build.BinariesDirectory)\RelWithDebInfo + call gradlew.bat cmakeCheck -x test -DcmakeBuildDir=$(Build.BinariesDirectory)\RelWithDebInfo --warning-mode all + workingDirectory: $(Build.SourcesDirectory)\java + ${{ else }}: + script: | + call gradlew.bat cmakeCheck -DcmakeBuildDir=$(Build.BinariesDirectory)\RelWithDebInfo --warning-mode all + workingDirectory: $(Build.SourcesDirectory)\java - task: CmdLine@2 displayName: 'Add symbols and notices to Java' diff --git a/tools/ci_build/github/azure-pipelines/templates/win-ci.yml b/tools/ci_build/github/azure-pipelines/templates/win-ci.yml index c726054d8eb10..52547fd9a796b 100644 --- a/tools/ci_build/github/azure-pipelines/templates/win-ci.yml +++ b/tools/ci_build/github/azure-pipelines/templates/win-ci.yml @@ -196,11 +196,11 @@ stages: parameters: msbuildPlatform: ${{ parameters.msbuildPlatform }} java_artifact_id: ${{ parameters.java_artifact_id }} - ${{ if contains(parameters.ort_build_pool_name, 'CPU') }}: - buildOnly: false + ${{ if or(contains(parameters.buildparameter, 'use_cuda'), contains(parameters.buildparameter, 'use_tensorrt')) }}: # When it is a GPU build, we only assemble the java binaries, testing will be done in the later stage with GPU machine - ${{ else }}: buildOnly: true + ${{ else }}: + buildOnly: false - task: PublishBuildArtifacts@1 displayName: 'Publish Java temp binaries' From 4743803944d6f6333bf93b89a62d8083d4466710 Mon Sep 17 00:00:00 2001 From: Wanming Lin Date: Tue, 25 Jun 2024 10:04:23 +0800 Subject: [PATCH 27/52] [WebNN EP] Support more Normalization ops for TFLite backend (#21151) Following Normalization ops have been supported in Chromium for TFLite backend: - batchNormalization: https://chromium-review.googlesource.com/c/chromium/src/+/5532745 - layerNormalization: https://chromium-review.googlesource.com/c/chromium/src/+/5573326 - instanceNormalization: https://chromium-review.googlesource.com/c/chromium/src/+/5532750 --- js/web/docs/webnn-operators.md | 6 +++--- onnxruntime/core/providers/webnn/builders/helper.h | 6 +++--- .../webnn/builders/impl/normalization_op_builder.cc | 7 ++----- 3 files changed, 8 insertions(+), 11 deletions(-) diff --git a/js/web/docs/webnn-operators.md b/js/web/docs/webnn-operators.md index a49759b9a93c5..987e063485846 100644 --- a/js/web/docs/webnn-operators.md +++ b/js/web/docs/webnn-operators.md @@ -16,7 +16,7 @@ operators and the supported opset domain/versions in **WebNN EP** by ONNX Runtim | ArgMax | ai.onnx(7-10, 11, 12, 13+) | argMax | ✓ | ✓ | WebNN CPU backend only supports 'select_last_index' value is 0 | | ArgMin | ai.onnx(7-10, 11, 12, 13+) | argMin | ✓ | ✓ | WebNN CPU backend only supports 'select_last_index' value is 0 | | AveragePool | ai.onnx(7-9, 10, 11, 12-18, 19+) | averagePool2d | ✓ | ✓ | Only supports 4-D input, 2-D 'kernel_shape', 'count_include_pad' value is 0 | -| BatchNormalization | ai.onnx(7-8, 9-13, 14, 15+) | batchNormalization | ✗ | ✓ | Only supports 'training_mode' value is 0, one output | +| BatchNormalization | ai.onnx(7-8, 9-13, 14, 15+) | batchNormalization | ✓ | ✓ | Only supports 'training_mode' value is 0, one output | | Cast | ai.onnx(7-8, 9-12, 13-18, 19-20, 21+) | cast | ✓ | ✓ | WebNN CPU backend doesn't support casting to uint64 data type | | Ceil | ai.onnx(7-12, 13+) | ceil | ✓ | ✓ | | | Clip | ai.onnx(7-10, 11, 12, 13+) | clamp | ✓ | ✓ | WebNN CPU backend only supports 3 specific ranges: [0.0, infinity], [-1.0, 1.0], [0.0, 6.0] (Chromium issue: https://issues.chromium.org/issues/326156496) | @@ -43,8 +43,8 @@ operators and the supported opset domain/versions in **WebNN EP** by ONNX Runtim | HardSigmoid | ai.onnx(7+) | hardSigmoid | ✓ | ✓ | | | HardSwish | ai.onnx(14+) | hardSwish | ✓ | ✓ | | | Identity | ai.onnx(7-13, 14-15, 16-18, 19-20, 21+) | identity | ✓ | ✓ | | -| InstanceNormalization | ai.onnx(7+) | instanceNormalization | ✗ | ✓ | | -| LayerNormalization | ai.onnx(7-16, 17+) | layerNormalization | ✗ | ✓ | | +| InstanceNormalization | ai.onnx(7+) | instanceNormalization | ✓ | ✓ | | +| LayerNormalization | ai.onnx(7-16, 17+) | layerNormalization | ✓ | ✓ | | | LeakyRelu | ai.onnx(7-15, 16+) | leakyRelu | ✓ | ✓ | | | Less | ai.onnx(7-8, 9-12, 13+) | lesser | ✓ | ✓ | | | LessOrEqual | ai.onnx(12-15, 16+) | lesserOrEqual | ✓ | ✓ | | diff --git a/onnxruntime/core/providers/webnn/builders/helper.h b/onnxruntime/core/providers/webnn/builders/helper.h index 401d2eaa09129..395a0b40e5bbb 100644 --- a/onnxruntime/core/providers/webnn/builders/helper.h +++ b/onnxruntime/core/providers/webnn/builders/helper.h @@ -160,7 +160,7 @@ static const InlinedHashMap op_map = { {"ArgMax", {"argMax", true}}, {"ArgMin", {"argMin", true}}, {"AveragePool", {"averagePool2d", true}}, - {"BatchNormalization", {"batchNormalization", false}}, + {"BatchNormalization", {"batchNormalization", true}}, {"Cast", {"cast", true}}, {"Ceil", {"ceil", true}}, {"Clip", {"clamp", true}}, @@ -190,8 +190,8 @@ static const InlinedHashMap op_map = { {"HardSigmoid", {"hardSigmoid", true}}, {"HardSwish", {"hardSwish", true}}, {"Identity", {"identity", true}}, - {"InstanceNormalization", {"instanceNormalization", false}}, - {"LayerNormalization", {"layerNormalization", false}}, + {"InstanceNormalization", {"instanceNormalization", true}}, + {"LayerNormalization", {"layerNormalization", true}}, {"LeakyRelu", {"leakyRelu", true}}, {"Less", {"lesser", true}}, {"LessOrEqual", {"lesserOrEqual", true}}, diff --git a/onnxruntime/core/providers/webnn/builders/impl/normalization_op_builder.cc b/onnxruntime/core/providers/webnn/builders/impl/normalization_op_builder.cc index 90ad9b48d5866..a2aa0df5586e3 100644 --- a/onnxruntime/core/providers/webnn/builders/impl/normalization_op_builder.cc +++ b/onnxruntime/core/providers/webnn/builders/impl/normalization_op_builder.cc @@ -87,11 +87,8 @@ Status NormalizationOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder int64_t axis = helper.Get("axis", -1); axis = HandleNegativeAxis(axis, rank); std::vector axes(rank - SafeInt(axis)); - if (model_builder.GetPreferredLayout() == DataLayout::NHWC && axis > 1) { - std::iota(axes.begin(), axes.end(), axis - 1); - } else { - std::iota(axes.begin(), axes.end(), axis); - } + std::iota(axes.begin(), axes.end(), axis); + options.set("axes", emscripten::val::array(axes)); output = model_builder.GetBuilder().call("layerNormalization", input, options); } else if (op_type == "InstanceNormalization") { From 41ad83fb004dd67b6927e7b2c3fffa81a90edc37 Mon Sep 17 00:00:00 2001 From: Wanming Lin Date: Wed, 26 Jun 2024 09:30:55 +0800 Subject: [PATCH 28/52] [WebNN EP] Support rest Reduction ops for TFLite backend (#21135) - reduceLogSum, reduceLogSumExp and reduceSumSquare have been landed in https://chromium-review.googlesource.com/c/chromium/src/+/5575815 - reduceL1 and reduceL2 have been landed in https://chromium-review.googlesource.com/c/chromium/src/+/5606091 --- js/web/docs/webnn-operators.md | 10 +++++----- onnxruntime/core/providers/webnn/builders/helper.h | 10 +++++----- 2 files changed, 10 insertions(+), 10 deletions(-) diff --git a/js/web/docs/webnn-operators.md b/js/web/docs/webnn-operators.md index 987e063485846..725d11b9d54c5 100644 --- a/js/web/docs/webnn-operators.md +++ b/js/web/docs/webnn-operators.md @@ -61,16 +61,16 @@ operators and the supported opset domain/versions in **WebNN EP** by ONNX Runtim | Pow | ai.onnx(7-11, 12, 13-14, 15+) | pow | ✓ | ✓ | | | PRelu | ai.onnx(7-8, 9-15, 16+) | prelu | ✓ | ✓ | WebNN CPU backend restricts the last dimension of input and slope to be same (Chromium issue: https://issues.chromium.org/issues/335517470) | | Reciprocal | ai.onnx(7-12, 13+) | reciprocal | ✓ | ✓ | | -| ReduceL1 | ai.onnx(7-10, 11-12, 13-17, 18+) | reduceL1 | ✗ | ✓ | Input 'axes' if present should be a constant | -| ReduceL2 | ai.onnx(7-10, 11-12, 13-17, 18+) | reduceL2 | ✗ | ✓ | Input 'axes' if present should be a constant | -| ReduceLogSum| ai.onnx(7-10, 11-12, 13-17, 18+) | reduceLogSum| ✗ | ✓ | Input 'axes' if present should be a constant | -| ReduceLogSumExp | ai.onnx(7-10, 11-12, 13-17, 18+) | reduceLogSumExp | ✗ | ✓ | Input 'axes' if present should be a constant | +| ReduceL1 | ai.onnx(7-10, 11-12, 13-17, 18+) | reduceL1 | ✓ | ✓ | Input 'axes' if present should be a constant | +| ReduceL2 | ai.onnx(7-10, 11-12, 13-17, 18+) | reduceL2 | ✓ | ✓ | Input 'axes' if present should be a constant | +| ReduceLogSum| ai.onnx(7-10, 11-12, 13-17, 18+) | reduceLogSum| ✓ | ✓ | Input 'axes' if present should be a constant | +| ReduceLogSumExp | ai.onnx(7-10, 11-12, 13-17, 18+) | reduceLogSumExp | ✓ | ✓ | Input 'axes' if present should be a constant | | ReduceMax | ai.onnx(7-10, 11, 12, 13-17, 18-19, 20+) | reduceMax | ✓ | ✓ | Input 'axes' if present should be a constant | | ReduceMean | ai.onnx(7-10, 11-12, 13-17, 18+) | reduceMean | ✓ | ✓ | Input 'axes' if present should be a constant | | ReduceMin | ai.onnx(7-10, 11, 12, 13-17, 18-19, 20+) | reduceMin | ✓ | ✓ | Input 'axes' if present should be a constant | | ReduceProd | ai.onnx(7-10, 11-12, 13-17, 18+) | reduceProduct | ✓ | ✓ | Input 'axes' if present should be a constant | | ReduceSum | ai.onnx(7-10, 11-12, 13+) | reduceSum | ✓ | ✓ | Input 'axes' if present should be a constant | -| ReduceSumSquare | ai.onnx(7-10, 11-12, 13-17, 18+) | reduceSumSquare | ✗ | ✓ | Input 'axes' if present should be a constant | +| ReduceSumSquare | ai.onnx(7-10, 11-12, 13-17, 18+) | reduceSumSquare | ✓ | ✓ | Input 'axes' if present should be a constant | | Relu | ai.onnx(7-12, 13, 14+) | relu | ✓ | ✓ | | | Reshape | ai.onnx(7-12, 13, 14-18, 19-20, 21+) | reshape | ✓ | ✓ | Input 'shape' should be a constant, 0 dimension value in 'shape' is not supported | | Resize | ai.onnx(11-12, 13-17, 18, 19+) | resample2d | ✓ | ✓ | Only supports 4-D input, exclude_outside != 0, input 'scales' and 'sizes' if present must be a constant, 'linear' and 'nearest' modes | diff --git a/onnxruntime/core/providers/webnn/builders/helper.h b/onnxruntime/core/providers/webnn/builders/helper.h index 395a0b40e5bbb..4ee3f891f92ca 100644 --- a/onnxruntime/core/providers/webnn/builders/helper.h +++ b/onnxruntime/core/providers/webnn/builders/helper.h @@ -209,16 +209,16 @@ static const InlinedHashMap op_map = { {"Pow", {"pow", true}}, {"PRelu", {"prelu", true}}, {"Reciprocal", {"reciprocal", true}}, - {"ReduceL1", {"reduceL1", false}}, - {"ReduceL2", {"reduceL2", false}}, - {"ReduceLogSum", {"reduceLogSum", false}}, - {"ReduceLogSumExp", {"reduceLogSumExp", false}}, + {"ReduceL1", {"reduceL1", true}}, + {"ReduceL2", {"reduceL2", true}}, + {"ReduceLogSum", {"reduceLogSum", true}}, + {"ReduceLogSumExp", {"reduceLogSumExp", true}}, {"ReduceMax", {"reduceMax", true}}, {"ReduceMean", {"reduceMean", true}}, {"ReduceMin", {"reduceMin", true}}, {"ReduceProd", {"reduceProduct", true}}, {"ReduceSum", {"reduceSum", true}}, - {"ReduceSumSquare", {"reduceSumSquare", false}}, + {"ReduceSumSquare", {"reduceSumSquare", true}}, {"Relu", {"relu", true}}, {"Reshape", {"reshape", true}}, {"Resize", {"resample2d", true}}, From e2abba18ea9370329ce6894a4eb3e98ad8f11cb6 Mon Sep 17 00:00:00 2001 From: mindest <30493312+mindest@users.noreply.github.com> Date: Wed, 26 Jun 2024 11:15:50 +0800 Subject: [PATCH 29/52] Skip softmax BF16 test for ROCm (#21162) ### Description Skip softmax BF16 test for ROCm, because BFloat16 is unsupported by MIOpen, and `torch.cuda.is_available()` also returns `True` for ROCm. --- .../test/python/orttraining_test_ortmodule_onnx_ops.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/orttraining/orttraining/test/python/orttraining_test_ortmodule_onnx_ops.py b/orttraining/orttraining/test/python/orttraining_test_ortmodule_onnx_ops.py index 88735ff18515e..35c5b736bd962 100644 --- a/orttraining/orttraining/test/python/orttraining_test_ortmodule_onnx_ops.py +++ b/orttraining/orttraining/test/python/orttraining_test_ortmodule_onnx_ops.py @@ -148,8 +148,8 @@ def test_onnx_ops(self): @unittest.skipIf(not torch.cuda.is_bf16_supported(), "Test requires CUDA and BF16 support") def test_softmax_bf16_large(self): - if not torch.cuda.is_available(): - # only test bf16 on cuda + if torch.version.cuda is None: + # Only run this test when CUDA is available, as on ROCm BF16 is not supported by MIOpen. return class Model(torch.nn.Module): @@ -175,7 +175,7 @@ def forward(self, input): data_ort.requires_grad = True ort_res = ort_model(input=data_ort) ort_res.backward(gradient=init_grad) - # compara result + # compare result torch.testing.assert_close(data_torch.grad, data_ort.grad, rtol=1e-5, atol=1e-4) From 337cc56d6f205ee5bb73c4aadad94f2e82fdadd4 Mon Sep 17 00:00:00 2001 From: Scott McKay Date: Thu, 27 Jun 2024 02:54:36 +1000 Subject: [PATCH 30/52] Convert scalars to 1D to satisfy ML Program requirements. (#21159) ### Description Convert scalars to 1D to satisfy ML Program requirements. https://dev.azure.com/onnxruntime/onnxruntime/_build/results?buildId=1418617&view=logs&j=f7cc61a9-cc70-56e7-b06c-4668ca17e426&t=16d281b5-1bfd-5309-f274-36d0dffd9cb1&l=27167 ### Motivation and Context Fixes test failure in #17361 --- .../coreml/builders/impl/builder_utils.cc | 48 ++++++++++++------- .../coreml/builders/impl/builder_utils.h | 4 +- .../coreml/builders/model_builder.cc | 7 +-- 3 files changed, 35 insertions(+), 24 deletions(-) diff --git a/onnxruntime/core/providers/coreml/builders/impl/builder_utils.cc b/onnxruntime/core/providers/coreml/builders/impl/builder_utils.cc index cbea969904ed5..2fcf9a1d7d9ba 100644 --- a/onnxruntime/core/providers/coreml/builders/impl/builder_utils.cc +++ b/onnxruntime/core/providers/coreml/builders/impl/builder_utils.cc @@ -140,30 +140,44 @@ void CreateCoreMLWeight(CoreML::Specification::WeightParams& weight, gsl::span> shape) { + std::optional> shape, bool convert_scalar = false) { tensor_type.set_datatype(data_type); if (shape) { - tensor_type.set_rank(shape->size()); - for (const auto& dim : *shape) { - if (dim >= 0) { - tensor_type.add_dimensions()->mutable_constant()->set_size(narrow(dim)); - } else { - tensor_type.add_dimensions()->mutable_unknown()->set_variadic(false); + auto rank = shape->size(); + if (convert_scalar && rank == 0) { + // CoreML scalar has shape {1} + tensor_type.set_rank(1); + tensor_type.add_dimensions()->mutable_constant()->set_size(1); + } else { + tensor_type.set_rank(rank); + for (const auto& dim : *shape) { + if (dim >= 0) { + tensor_type.add_dimensions()->mutable_constant()->set_size(narrow(dim)); + } else { + tensor_type.add_dimensions()->mutable_unknown()->set_variadic(false); + } } } } } void SetTensorTypeInfo(MILSpec::TensorType& tensor_type, MILSpec::DataType data_type, - const ONNX_NAMESPACE::TensorShapeProto* shape) { + const ONNX_NAMESPACE::TensorShapeProto* shape, bool convert_scalar = false) { tensor_type.set_datatype(data_type); if (shape) { - tensor_type.set_rank(shape->dim_size()); - for (const auto& dim : shape->dim()) { - if (dim.has_dim_value()) { - tensor_type.add_dimensions()->mutable_constant()->set_size(narrow(dim.dim_value())); - } else { - tensor_type.add_dimensions()->mutable_unknown()->set_variadic(false); + auto rank = shape->dim_size(); + if (convert_scalar && rank == 0) { + // CoreML scalar has shape {1} + tensor_type.set_rank(1); + tensor_type.add_dimensions()->mutable_constant()->set_size(1); + } else { + tensor_type.set_rank(rank); + for (const auto& dim : shape->dim()) { + if (dim.has_dim_value()) { + tensor_type.add_dimensions()->mutable_constant()->set_size(narrow(dim.dim_value())); + } else { + tensor_type.add_dimensions()->mutable_unknown()->set_variadic(false); + } } } } @@ -281,13 +295,13 @@ template MILSpec::Value CreateScalarTensorValue(const int32_t& data); template MILSpec::Value CreateScalarTensorValue(const std::string& data); template MILSpec::Value CreateScalarTensorValue(const bool& data); -COREML_SPEC::MILSpec::NamedValueType CreateNamedTensorValueType(const NodeArg& node_arg) { +COREML_SPEC::MILSpec::NamedValueType CreateNamedTensorValueType(const NodeArg& node_arg, bool convert_scalar) { MILSpec::NamedValueType nvt; nvt.set_name(node_arg.Name()); MILSpec::TensorType& tensor_type = *nvt.mutable_type()->mutable_tensortype(); SetTensorTypeInfo(tensor_type, OnnxDataTypeToMILSpec(node_arg.TypeAsProto()->tensor_type().elem_type()), - node_arg.Shape()); + node_arg.Shape(), convert_scalar); return nvt; } @@ -308,7 +322,7 @@ void AddOperationOutput(COREML_SPEC::MILSpec::Operation& op, const NodeArg& outp MILSpec::TensorType& tensor_type = *value.mutable_tensortype(); SetTensorTypeInfo(tensor_type, OnnxDataTypeToMILSpec(output.TypeAsProto()->tensor_type().elem_type()), - output.Shape()); + output.Shape(), /*convert_scalar*/ true); } void AddPadTypeAndPads(COREML_SPEC::MILSpec::Operation& op, ModelBuilder& model_builder, std::string_view op_type, diff --git a/onnxruntime/core/providers/coreml/builders/impl/builder_utils.h b/onnxruntime/core/providers/coreml/builders/impl/builder_utils.h index 2804589065631..3e6c43ab07867 100644 --- a/onnxruntime/core/providers/coreml/builders/impl/builder_utils.h +++ b/onnxruntime/core/providers/coreml/builders/impl/builder_utils.h @@ -114,8 +114,10 @@ template COREML_SPEC::MILSpec::Value CreateScalarTensorValue(const T& data); /// Create a NamedValueType from an ONNX tensor NodeArg. +/// NodeArg to create NamedValueType from. +/// If true, scalar shapes are converted to 1D. /// Used to create inputs for the 'main' function in an ML Program. -COREML_SPEC::MILSpec::NamedValueType CreateNamedTensorValueType(const NodeArg& node_arg); +COREML_SPEC::MILSpec::NamedValueType CreateNamedTensorValueType(const NodeArg& node_arg, bool convert_scalar = false); /// /// Add an input argument to a MILSpec::Operation diff --git a/onnxruntime/core/providers/coreml/builders/model_builder.cc b/onnxruntime/core/providers/coreml/builders/model_builder.cc index eb4723a3b9746..88b518ab2289c 100644 --- a/onnxruntime/core/providers/coreml/builders/model_builder.cc +++ b/onnxruntime/core/providers/coreml/builders/model_builder.cc @@ -838,13 +838,8 @@ Status ModelBuilder::RegisterModelInputOutput(const NodeArg& node_arg, bool is_i if (create_ml_program_) { if (is_input) { // the model inputs need to be wired up as args to the 'main' function. - auto tensor_value_type = CreateNamedTensorValueType(node_arg); + auto tensor_value_type = CreateNamedTensorValueType(node_arg, /*convert_scalar*/ true); tensor_value_type.set_name(name); - if (node_arg.Shape()->dim_size() == 0) { - // update shape from {} to {1} (same change we made at the model input level above). - tensor_value_type.mutable_type()->mutable_tensortype()->set_rank(1); - tensor_value_type.mutable_type()->mutable_tensortype()->add_dimensions()->mutable_constant()->set_size(1); - } mlprogram_main_fn_->mutable_inputs()->Add(std::move(tensor_value_type)); } else { From 3c0b407709fd3c71755ed046edd688b30a786d94 Mon Sep 17 00:00:00 2001 From: Vincent Wang Date: Thu, 27 Jun 2024 01:00:45 +0800 Subject: [PATCH 31/52] Rollback 19832, Remove shape_input_merge Fusion (#21179) The PR caused Big Models pipeline failure for running Llama2. After the rollback, the pipeline is back to normal. --- .../core/optimizer/graph_transformer_utils.cc | 8 +- .../core/optimizer/shape_input_merge.cc | 78 ------------------- .../core/optimizer/shape_input_merge.h | 23 ------ .../test/optimizer/graph_transform_test.cc | 75 ------------------ .../core/optimizer/graph_transformer_utils.cc | 8 +- 5 files changed, 6 insertions(+), 186 deletions(-) delete mode 100644 onnxruntime/core/optimizer/shape_input_merge.cc delete mode 100644 onnxruntime/core/optimizer/shape_input_merge.h diff --git a/onnxruntime/core/optimizer/graph_transformer_utils.cc b/onnxruntime/core/optimizer/graph_transformer_utils.cc index 53c7f39bdd7f1..4298551aec412 100644 --- a/onnxruntime/core/optimizer/graph_transformer_utils.cc +++ b/onnxruntime/core/optimizer/graph_transformer_utils.cc @@ -71,7 +71,6 @@ #include "core/optimizer/reshape_fusion.h" #include "core/optimizer/rocm_blas_alt_impl.h" #include "core/optimizer/rule_based_graph_transformer.h" -#include "core/optimizer/shape_input_merge.h" #include "core/optimizer/skip_layer_norm_fusion.h" #include "core/optimizer/slice_elimination.h" #include "core/optimizer/transpose_optimizer.h" @@ -215,9 +214,9 @@ InlinedVector> GenerateTransformers( transformers.emplace_back(std::make_unique()); } - // Put ConstantSharing and ShapeInputMerge before CommonSubexpressionElimination by intention as it can create - // more opportunities for CSE. For example, if A and B nodes consume same different args but produce same output - // or consume different initializers with same value, by default, CSE will not merge them. + // Put ConstantSharing before CommonSubexpressionElimination by intention as it can create more opportunities for + // CSE. For example, if A and B nodes consume different initializers with same value, by default, + // CSE will not merge them. InlinedHashSet excluded_initializers; excluded_initializers.reserve(session_options.initializers_to_share_map.size()); for (const auto& p : session_options.initializers_to_share_map) { @@ -225,7 +224,6 @@ InlinedVector> GenerateTransformers( } const InlinedHashSet no_limit_empty_ep_list = {}; transformers.emplace_back(std::make_unique(no_limit_empty_ep_list, excluded_initializers)); - transformers.emplace_back(std::make_unique()); transformers.emplace_back(std::make_unique()); transformers.emplace_back(std::make_unique(cpu_execution_provider, !disable_quant_qdq, session_options.config_options)); diff --git a/onnxruntime/core/optimizer/shape_input_merge.cc b/onnxruntime/core/optimizer/shape_input_merge.cc deleted file mode 100644 index dec1382319f16..0000000000000 --- a/onnxruntime/core/optimizer/shape_input_merge.cc +++ /dev/null @@ -1,78 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT License. - -#include "core/optimizer/shape_input_merge.h" - -#include "core/graph/graph_utils.h" - -namespace onnxruntime { - -namespace { -std::string GetShapeString(const NodeArg* input_arg) { - auto shape = input_arg->Shape(); - if (!shape) return ""; - std::stringstream ss; - ss << "["; - for (int i = 0; i < shape->dim_size(); ++i) { - if (i != 0) ss << ","; - auto dim = shape->dim(i); - if (dim.has_dim_value()) { - ss << std::to_string(dim.dim_value()); - } else if (dim.has_dim_param()) { - ss << "'" << dim.dim_param() << "'"; - } else { - return ""; - } - } - ss << "]"; - return ss.str(); -} - -} // namespace - -Status ShapeInputMerge::ApplyImpl(Graph& graph, bool& modified, int graph_level, const logging::Logger& logger) const { - GraphViewer graph_viewer(graph); - const auto& node_topology_list = graph_viewer.GetNodesInTopologicalOrder(); - InlinedHashMap> input_hash_to_nodes; - for (auto node_index : node_topology_list) { - auto* p_node = graph.GetNode(node_index); - if (!p_node) continue; // we removed the node as part of an earlier fusion - ORT_RETURN_IF_ERROR(Recurse(*p_node, modified, graph_level, logger)); - if (!graph_utils::IsSupportedOptypeVersionAndDomain(*p_node, "Shape", {1, 13, 15, 19, 21}) || - !graph_utils::IsSupportedProvider(*p_node, GetCompatibleExecutionProviders())) { - continue; - } - std::string shape_str = GetShapeString(p_node->InputDefs()[0]); - if (shape_str.empty()) continue; - if (input_hash_to_nodes.find(shape_str) == input_hash_to_nodes.end()) { - input_hash_to_nodes[shape_str] = InlinedVector(); - } - input_hash_to_nodes[shape_str].emplace_back(p_node); - } - - // All Shape nodes are processed in topological order, so we can safely merge the inputs to the first node's input. - for (auto& kv : input_hash_to_nodes) { - if (kv.second.size() < 2) continue; - NodeArg* first_input_arg = kv.second[0]->MutableInputDefs()[0]; - bool is_first_input_arg_graph_input = graph.IsInputsIncludingInitializers(first_input_arg); - for (size_t i = 1; i < kv.second.size(); ++i) { - Node* p_node = kv.second[i]; - const NodeArg* input_arg = p_node->InputDefs()[0]; - if (input_arg->Name() == first_input_arg->Name()) continue; - if (!graph.IsInputsIncludingInitializers(input_arg) && p_node->GetInputEdgesCount()) { - const Node::EdgeEnd& input_edge = *p_node->InputEdgesBegin(); - graph.RemoveEdge(input_edge.GetNode().Index(), p_node->Index(), input_edge.GetSrcArgIndex(), 0); - } - graph_utils::ReplaceNodeInput(*p_node, 0, *first_input_arg); - if (!is_first_input_arg_graph_input && kv.second[0]->GetInputEdgesCount()) { - const Node::EdgeEnd& first_input_edge = *kv.second[0]->InputEdgesBegin(); - graph.AddEdge(first_input_edge.GetNode().Index(), p_node->Index(), first_input_edge.GetSrcArgIndex(), 0); - } - modified = true; - } - } - - return Status::OK(); -} - -} // namespace onnxruntime diff --git a/onnxruntime/core/optimizer/shape_input_merge.h b/onnxruntime/core/optimizer/shape_input_merge.h deleted file mode 100644 index 5cb943998487b..0000000000000 --- a/onnxruntime/core/optimizer/shape_input_merge.h +++ /dev/null @@ -1,23 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT License. - -#pragma once - -#include "core/optimizer/graph_transformer.h" - -namespace onnxruntime { - -/** -@Class ShapeInputMerge -Merge all shape inputs having same shape value to a single shape input. -This change will not affect the performance, but it open chances for CSE fusion to merge nodes. -*/ -class ShapeInputMerge : public GraphTransformer { - public: - ShapeInputMerge(const InlinedHashSet& compatible_execution_providers = {}) noexcept - : GraphTransformer("ShapeInputMerge", compatible_execution_providers) {} - - Status ApplyImpl(Graph& graph, bool& modified, int graph_level, const logging::Logger& logger) const override; -}; - -} // namespace onnxruntime diff --git a/onnxruntime/test/optimizer/graph_transform_test.cc b/onnxruntime/test/optimizer/graph_transform_test.cc index 9609ec57c8b26..f83fb8238ff61 100755 --- a/onnxruntime/test/optimizer/graph_transform_test.cc +++ b/onnxruntime/test/optimizer/graph_transform_test.cc @@ -70,7 +70,6 @@ #include "core/optimizer/relu_clip_fusion.h" #include "core/optimizer/reshape_fusion.h" #include "core/optimizer/rule_based_graph_transformer.h" -#include "core/optimizer/shape_input_merge.h" #include "core/optimizer/slice_elimination.h" #include "core/optimizer/unsqueeze_elimination.h" #include "core/optimizer/utils.h" @@ -7691,80 +7690,6 @@ TEST_F(GraphTransformationTests, GatherToSliceFusion) { } } -TEST_F(GraphTransformationTests, ShapeInputMerge) { - auto build_test_case = [&](ModelTestBuilder& builder) { - std::vector> input_shape; - input_shape.reserve(5); - input_shape.emplace_back("dim0"); - input_shape.emplace_back(512); - input_shape.emplace_back(1); - input_shape.emplace_back(1536); - input_shape.emplace_back("dim4"); - auto* input_arg = builder.MakeSymbolicInput(input_shape); - auto* neg_out = builder.MakeIntermediate(); - auto* axes_initializer = builder.MakeInitializer({1}, {static_cast(2)}); - auto* squeeze_out = builder.MakeIntermediate(); - auto* cast_out = builder.MakeIntermediate(); - auto* unsqueeze_out = builder.MakeOutput(); - auto* shape_1_out = builder.MakeOutput(); - auto* shape_2_out = builder.MakeOutput(); - auto* shape_3_out = builder.MakeOutput(); - auto* shape_4_out = builder.MakeOutput(); - auto* shape_5_out = builder.MakeOutput(); - builder.AddNode("Neg", {input_arg}, {neg_out}); - builder.AddNode("Squeeze", {neg_out, axes_initializer}, {squeeze_out}); - builder.AddNode("Cast", {squeeze_out}, {cast_out}).AddAttribute("to", static_cast(10)); - builder.AddNode("Unsqueeze", {cast_out, axes_initializer}, {unsqueeze_out}); - builder.AddNode("Shape", {input_arg}, {shape_1_out}); - builder.AddNode("Shape", {neg_out}, {shape_2_out}); - builder.AddNode("Shape", {squeeze_out}, {shape_3_out}); - builder.AddNode("Shape", {cast_out}, {shape_4_out}); - builder.AddNode("Shape", {unsqueeze_out}, {shape_5_out}); - }; - - auto pre_graph_checker = [&](Graph& graph) { - InlinedHashMap ref_count; - for (auto& node : graph.Nodes()) { - if (node.OpType() == "Shape") { - std::string name = node.InputDefs()[0]->Name(); - if (ref_count.find(name) == ref_count.end()) { - ref_count[name] = 1; - } else { - ref_count[name]++; - } - } - } - TEST_RETURN_IF_NOT(ref_count.size() == 5); - return Status::OK(); - }; - - auto post_graph_checker = [&](Graph& graph) { - InlinedHashMap ref_count; - for (auto& node : graph.Nodes()) { - if (node.OpType() == "Shape") { - std::string name = node.InputDefs()[0]->Name(); - if (ref_count.find(name) == ref_count.end()) { - ref_count[name] = 1; - } else { - ref_count[name]++; - } - } - } - TEST_RETURN_IF_NOT(ref_count.size() == 2); - int sum = 0, mul = 1; - for (auto& entry : ref_count) { - sum += entry.second; - mul *= entry.second; - } - TEST_RETURN_IF_NOT(sum == 5 && mul == 6); - return Status::OK(); - }; - - std::unique_ptr transformer = std::make_unique(); - ASSERT_STATUS_OK(TestGraphTransformer(build_test_case, 14, *logger_, std::move(transformer), TransformerLevel::Level1, - 1, pre_graph_checker, post_graph_checker)); -} - #if !defined(DISABLE_CONTRIB_OPS) TEST_F(GraphTransformationTests, MatMulNBitsBiasFusion) { diff --git a/orttraining/orttraining/core/optimizer/graph_transformer_utils.cc b/orttraining/orttraining/core/optimizer/graph_transformer_utils.cc index 436a24c34ddfd..589e7be455dbc 100644 --- a/orttraining/orttraining/core/optimizer/graph_transformer_utils.cc +++ b/orttraining/orttraining/core/optimizer/graph_transformer_utils.cc @@ -44,7 +44,6 @@ #include "core/optimizer/relu_clip_fusion.h" #include "core/optimizer/reshape_fusion.h" #include "core/optimizer/rule_based_graph_transformer.h" -#include "core/optimizer/shape_input_merge.h" #include "core/optimizer/skip_layer_norm_fusion.h" #include "core/optimizer/slice_elimination.h" #include "core/optimizer/unsqueeze_elimination.h" @@ -117,11 +116,10 @@ std::vector> GeneratePreTrainingTransformers( ORT_THROW_IF_ERROR(rule_transformer->Register(std::make_unique())); #endif - // Put ConstantSharing and ShapeInputMerge before CommonSubexpressionElimination by intention as it can create - // more opportunities for CSE. For example, if A and B nodes consume same different args but produce same output - // or consume different initializers with same value, by default, CSE will not merge them. + // Put ConstantSharing before CommonSubexpressionElimination by intention as it can create more opportunities for + // CSE. For example, if A and B nodes consume different initializers with same value, by default, + // CSE will not merge them. transformers.emplace_back(std::make_unique(compatible_eps)); - transformers.emplace_back(std::make_unique(compatible_eps)); // LayerNormFusion must be applied before CommonSubexpressionElimination as the latter will break the pattern when 2 LayerNormFusion share the same input. transformers.emplace_back(std::make_unique(compatible_eps)); // Remove duplicate nodes. Must be applied before any recompute transformations. From 887a818aa7d79ac0e08f3531e493dac1f980674d Mon Sep 17 00:00:00 2001 From: Scott McKay Date: Thu, 27 Jun 2024 12:51:13 +1000 Subject: [PATCH 32/52] Check for unit test log severity override earlier (#21177) ### Description Setting the log level after environment creation is too late in some cases. If the DML EP is enabled, it will create a composite sink with the original logger using the creation time log severity, as well as additional ETW sink. As it saves the current severity levels for each sink inside the composite sink that prevents being able to get verbose log output to stdout even if you set that at the session level. I don't know enough about the setup that combines ETW with the original sink to say whether we should also be updating the severity of individual sinks in the combined sink, so this change is limited to making the unit tests behave in the expected manner when the default log severity is set in the background and not directly controlled. ### Motivation and Context Make it possible to get verbose output to stdout when the DML EP is enabled. --- onnxruntime/test/unittest_main/test_main.cc | 28 ++++++++++++--------- 1 file changed, 16 insertions(+), 12 deletions(-) diff --git a/onnxruntime/test/unittest_main/test_main.cc b/onnxruntime/test/unittest_main/test_main.cc index b7c3b38538421..1d89272680e47 100644 --- a/onnxruntime/test/unittest_main/test_main.cc +++ b/onnxruntime/test/unittest_main/test_main.cc @@ -26,9 +26,24 @@ #include "test/test_environment.h" std::unique_ptr ort_env; + +// ortenv_setup is used by /onnxruntime/test/xctest/xcgtest.mm so can't be file local void ortenv_setup() { OrtThreadingOptions tpo; - ort_env.reset(new Ort::Env(&tpo, ORT_LOGGING_LEVEL_WARNING, "Default")); + + // allow verbose logging to be enabled by setting this environment variable to a numeric log level + constexpr auto kLogLevelEnvironmentVariableName = "ORT_UNIT_TEST_MAIN_LOG_LEVEL"; + OrtLoggingLevel log_level = ORT_LOGGING_LEVEL_WARNING; + if (auto log_level_override = onnxruntime::ParseEnvironmentVariable(kLogLevelEnvironmentVariableName); + log_level_override.has_value()) { + *log_level_override = std::clamp(*log_level_override, + static_cast(ORT_LOGGING_LEVEL_VERBOSE), + static_cast(ORT_LOGGING_LEVEL_FATAL)); + std::cout << "Setting log level to " << *log_level_override << "\n"; + log_level = static_cast(*log_level_override); + } + + ort_env.reset(new Ort::Env(&tpo, log_level, "Default")); } #ifdef USE_TENSORRT @@ -76,17 +91,6 @@ int TEST_MAIN(int argc, char** argv) { ortenv_setup(); ::testing::InitGoogleTest(&argc, argv); - // allow verbose logging to be enabled by setting this environment variable to a numeric log level - constexpr auto kLogLevelEnvironmentVariableName = "ORT_UNIT_TEST_MAIN_LOG_LEVEL"; - if (auto log_level = onnxruntime::ParseEnvironmentVariable(kLogLevelEnvironmentVariableName); - log_level.has_value()) { - *log_level = std::clamp(*log_level, - static_cast(ORT_LOGGING_LEVEL_VERBOSE), - static_cast(ORT_LOGGING_LEVEL_FATAL)); - std::cout << "Setting log level to " << *log_level << "\n"; - ort_env->UpdateEnvWithCustomLogLevel(static_cast(*log_level)); - } - status = RUN_ALL_TESTS(); } ORT_CATCH(const std::exception& ex) { From eecc11afc7e4dfd130ecebd1b693ec828afdebd4 Mon Sep 17 00:00:00 2001 From: mindest <30493312+mindest@users.noreply.github.com> Date: Thu, 27 Jun 2024 12:04:17 +0800 Subject: [PATCH 33/52] [ROCm] Disable ck_tile in Debug build (#21178) ### Description tmp fix: disable ck_tile for Debug build. ### Motivation and Context Release build works fine for ck_tile, while Debug build fails.
Typical error log to revisit ``` [880/1797] Building HIP object CMakeFiles/onnxruntime_composable_kernel_fmha.dir/_deps/composable_kernel-build/fmha_fwd_d32_fp16_batch_b128x64x16x32x32x32_r2x1x1_w32x32x16_qr_async_vc_psddv.cpp.o FAILED: CMakeFiles/onnxruntime_composable_kernel_fmha.dir/_deps/composable_kernel-build/fmha_fwd_d32_fp16_batch_b128x64x16x32x32x32_r2x1x1_w32x32x16_qr_async_vc_psddv.cpp.o /opt/rocm/llvm/bin/clang++ -DEIGEN_MPL2_ONLY -DENABLE_ROCM_PROFILING -DENABLE_STRIDED_TENSORS -DENABLE_TRAINING -DENABLE_TRAINING_APIS -DENABLE_TRAINING_CORE -DENABLE_TRAINING_OPS -DENABLE_TRAINING_TORCH_INTEROP -DMIOPEN_VERSION=30100 -DORT_ENABLE_STREAM -DROCM_VERSION=60100 -DUSE_ROCM=1 -D_GNU_SOURCE -D__HIP_ROCclr__=1 -D__bf16__ -D__fp16__ -D__fp32__ -I/build/Debug/_deps/utf8_range-src -I/ws/onnxruntime/include/onnxruntime -I/ws/onnxruntime/include/onnxruntime/core/session -I/ws/onnxruntime/orttraining/orttraining/training_api/include -I/build/Debug/_deps/composable_kernel-src/example/ck_tile/01_fmha -I/build/Debug/_deps/composable_kernel-src/include -I/build/Debug/_deps/composable_kernel-build/include -I/build/Debug/_deps/composable_kernel-src/library/include -isystem /opt/rocm-6.1.0/include -g -O -std=gnu++17 --offload-arch=gfx90a -fPIC -x hip -mllvm=-amdgpu-early-inline-all=true -mllvm=-amdgpu-function-calls=false -MD -MT CMakeFiles/onnxruntime_composable_kernel_fmha.dir/_deps/composable_kernel-build/fmha_fwd_d32_fp16_batch_b128x64x16x32x32x32_r2x1x1_w32x32x16_qr_async_vc_psddv.cpp.o -MF CMakeFiles/onnxruntime_composable_kernel_fmha.dir/_deps/composable_kernel-build/fmha_fwd_d32_fp16_batch_b128x64x16x32x32x32_r2x1x1_w32x32x16_qr_async_vc_psddv.cpp.o.d -o CMakeFiles/onnxruntime_composable_kernel_fmha.dir/_deps/composable_kernel-build/fmha_fwd_d32_fp16_batch_b128x64x16x32x32x32_r2x1x1_w32x32x16_qr_async_vc_psddv.cpp.o -x hip -c /build/Debug/_deps/composable_kernel-build/fmha_fwd_d32_fp16_batch_b128x64x16x32x32x32_r2x1x1_w32x32x16_qr_async_vc_psddv.cpp In file included from /build/Debug/_deps/composable_kernel-build/fmha_fwd_d32_fp16_batch_b128x64x16x32x32x32_r2x1x1_w32x32x16_qr_async_vc_psddv.cpp:5: In file included from /build/Debug/_deps/composable_kernel-src/example/ck_tile/01_fmha/fmha_fwd.hpp:6: In file included from /build/Debug/_deps/composable_kernel-src/include/ck_tile/core.hpp:11: /build/Debug/_deps/composable_kernel-src/include/ck_tile/core/arch/utility.hpp:27:18: error: constraint 'n' expects an integer constant expression 27 | asm volatile("s_add_u32 m0, %0, m0" : : "n"(v) : "memory"); | ^ /build/Debug/_deps/composable_kernel-src/include/ck_tile/core/arch/utility.hpp:27:18: error: constraint 'n' expects an integer constant expression /build/Debug/_deps/composable_kernel-src/include/ck_tile/core/arch/utility.hpp:27:18: error: constraint 'n' expects an integer constant expression /build/Debug/_deps/composable_kernel-src/include/ck_tile/core/arch/utility.hpp:27:18: error: constraint 'n' expects an integer constant expression /build/Debug/_deps/composable_kernel-src/include/ck_tile/core/arch/utility.hpp:27:18: error: constraint 'n' expects an integer constant expression /build/Debug/_deps/composable_kernel-src/include/ck_tile/core/arch/utility.hpp:27:18: error: constraint 'n' expects an integer constant expression /build/Debug/_deps/composable_kernel-src/include/ck_tile/core/arch/utility.hpp:27:18: error: constraint 'n' expects an integer constant expression /build/Debug/_deps/composable_kernel-src/include/ck_tile/core/arch/utility.hpp:27:18: error: constraint 'n' expects an integer constant expression /build/Debug/_deps/composable_kernel-src/include/ck_tile/core/arch/utility.hpp:27:18: error: constraint 'n' expects an integer constant expression /build/Debug/_deps/composable_kernel-src/include/ck_tile/core/arch/utility.hpp:27:18: error: constraint 'n' expects an integer constant expression /build/Debug/_deps/composable_kernel-src/include/ck_tile/core/arch/utility.hpp:27:18: error: constraint 'n' expects an integer constant expression /build/Debug/_deps/composable_kernel-src/include/ck_tile/core/arch/utility.hpp:27:18: error: constraint 'n' expects an integer constant expression /build/Debug/_deps/composable_kernel-src/include/ck_tile/core/arch/utility.hpp:27:18: error: constraint 'n' expects an integer constant expression /build/Debug/_deps/composable_kernel-src/include/ck_tile/core/arch/utility.hpp:27:18: error: constraint 'n' expects an integer constant expression /build/Debug/_deps/composable_kernel-src/include/ck_tile/core/arch/utility.hpp:27:18: error: constraint 'n' expects an integer constant expression /build/Debug/_deps/composable_kernel-src/include/ck_tile/core/arch/utility.hpp:27:18: error: constraint 'n' expects an integer constant expression /build/Debug/_deps/composable_kernel-src/include/ck_tile/core/arch/utility.hpp:27:18: error: constraint 'n' expects an integer constant expression /build/Debug/_deps/composable_kernel-src/include/ck_tile/core/arch/utility.hpp:27:18: error: constraint 'n' expects an integer constant expression /build/Debug/_deps/composable_kernel-src/include/ck_tile/core/arch/utility.hpp:27:18: error: constraint 'n' expects an integer constant expression fatal error: too many errors emitted, stopping now [-ferror-limit=] 20 errors generated when compiling for gfx90a. ... ```
--- cmake/CMakeLists.txt | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/cmake/CMakeLists.txt b/cmake/CMakeLists.txt index b2122bf56abd8..9670dcb97abb2 100644 --- a/cmake/CMakeLists.txt +++ b/cmake/CMakeLists.txt @@ -373,6 +373,10 @@ if (onnxruntime_USE_ROCM) message(WARNING "ck_tile can only be enabled on ROCm >= 6.0 due to compatibility and compilation speed, disable automatically") set(onnxruntime_USE_COMPOSABLE_KERNEL_CK_TILE OFF) endif() + if (onnxruntime_USE_COMPOSABLE_KERNEL_CK_TILE AND CMAKE_BUILD_TYPE STREQUAL "Debug") + message(WARNING "ck_tile hits compiler error in Debug build, disable automatically") + set(onnxruntime_USE_COMPOSABLE_KERNEL_CK_TILE OFF) + endif() endif() From 446aa986a18615ed6675678cf2d62127f9a5b255 Mon Sep 17 00:00:00 2001 From: PeixuanZuo <94887879+PeixuanZuo@users.noreply.github.com> Date: Thu, 27 Jun 2024 15:36:04 +0800 Subject: [PATCH 34/52] [ROCm] Extend the Pipeline restriction time (#21158) ROCm EP builds are taking longer. --- tools/ci_build/github/azure-pipelines/templates/rocm.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tools/ci_build/github/azure-pipelines/templates/rocm.yml b/tools/ci_build/github/azure-pipelines/templates/rocm.yml index 231e375bedd9e..47c0ba3eb2d1e 100644 --- a/tools/ci_build/github/azure-pipelines/templates/rocm.yml +++ b/tools/ci_build/github/azure-pipelines/templates/rocm.yml @@ -17,7 +17,7 @@ jobs: - job: wheels_python_${{ replace(parameters.PythonVersion,'.','_') }}_rocm_${{ replace(parameters.RocmVersion,'.','_') }}_${{ parameters.BuildConfig }} workspace: clean: all - timeoutInMinutes: 180 + timeoutInMinutes: 360 pool: Ubuntu-2204-rocm-aiinfra variables: - name: PythonVersion From a1bbfeb3061f291431234c61cf4ed8051e4c00c8 Mon Sep 17 00:00:00 2001 From: kailums <109063327+kailums@users.noreply.github.com> Date: Thu, 27 Jun 2024 18:53:12 +0800 Subject: [PATCH 35/52] add split3inner (#19886) ### Description The split op is using pin_memory when split on different sizes. But pin_memory is not capable for using cudagraph. Add a new implementation for only transformer scenarios, it split the qkv_proj into q, k, v, not using pin_memory. ### Motivation and Context --- .../core/providers/cuda/tensor/split.cc | 25 ++++ .../core/providers/cuda/tensor/split_impl.cu | 107 ++++++++++++++++++ .../core/providers/cuda/tensor/split_impl.h | 4 + .../providers/cpu/tensor/split_op_test.cc | 57 ++++++++++ 4 files changed, 193 insertions(+) diff --git a/onnxruntime/core/providers/cuda/tensor/split.cc b/onnxruntime/core/providers/cuda/tensor/split.cc index 5f73512ab8696..ca82387600085 100644 --- a/onnxruntime/core/providers/cuda/tensor/split.cc +++ b/onnxruntime/core/providers/cuda/tensor/split.cc @@ -76,6 +76,31 @@ Status SplitKernel::ComputeInternal(OpKernelContext* ctx) const { auto input_dims = input_shape.GetDims(); auto output_dimensions{input_shape.AsShapeVector()}; + if (split_sizes.size() == 3 && ((axis + 1) == gsl::narrow_cast(input_shape.NumDimensions()))) { + // we use (axis + 1) == num_dimensions to check if we are splitting on inner most axis. + // only when split on inner axis and output size is 3, we can use Split3Inner. + // this kernel is not using pin_memory, so it is ok for using cuda graph. + output_dimensions[axis] = split_sizes[0]; + Tensor* output0 = ctx->Output(0, TensorShape{output_dimensions}); + output_dimensions[axis] = split_sizes[1]; + Tensor* output1 = ctx->Output(1, TensorShape{output_dimensions}); + output_dimensions[axis] = split_sizes[2]; + Tensor* output2 = ctx->Output(2, TensorShape{output_dimensions}); + + // if input tensor is empty, we don't need to launch kernel, but still need to set output tensor. + if (input_tensor->Shape().Size() <= 0) return Status::OK(); + + return Split3Inner(Stream(ctx), + input_tensor->DataType()->Size(), + split_sizes[0], split_sizes[1], + split_sizes[2], + input_tensor->DataRaw(), + output0->MutableDataRaw(), + output1->MutableDataRaw(), + output2->MutableDataRaw(), + input_dims); + } + CudaAsyncBuffer output_ptr(this, num_outputs); gsl::span output_ptr_span = output_ptr.CpuSpan(); TensorShapeVector axis_dimension_input_output_mapping(input_dims[axis]); diff --git a/onnxruntime/core/providers/cuda/tensor/split_impl.cu b/onnxruntime/core/providers/cuda/tensor/split_impl.cu index b0ff856a43970..00f94694f83c0 100644 --- a/onnxruntime/core/providers/cuda/tensor/split_impl.cu +++ b/onnxruntime/core/providers/cuda/tensor/split_impl.cu @@ -157,5 +157,112 @@ Status SplitImpl(cudaStream_t stream, const size_t element_size, const int block return Status::OK(); } +template +__global__ void _Split3InnerKernel(const int64_t size0_in_byte, + const int64_t size1_in_byte, + const int64_t size2_in_byte, + const void* input_data, + void* output_data0, + void* output_data1, + void* output_data2, + const int64_t inner_size_in_byte) { + // each block copy one row of input data + auto size0 = size0_in_byte / sizeof(T); + auto size1 = size1_in_byte / sizeof(T); + auto size2 = size2_in_byte / sizeof(T); + auto inner_size = inner_size_in_byte / sizeof(T); + auto output0_vec = reinterpret_cast(output_data0) + blockIdx.x * size0; + auto output1_vec = reinterpret_cast(output_data1) + blockIdx.x * size1; + auto output2_vec = reinterpret_cast(output_data2) + blockIdx.x * size2; + auto input_vec = reinterpret_cast(input_data) + blockIdx.x * inner_size; + // all size and pointer are aligned to sizeof(T) + // so here use all threads in the block to do vectorized copy + + for (auto tid = threadIdx.x; tid < inner_size; tid += blockDim.x) { + auto data = input_vec[tid]; + if (tid < size0) { + output0_vec[tid] = data; + } else if (tid < (size0 + size1)) { + output1_vec[tid - size0] = data; + } else { + output2_vec[tid - size0 - size1] = data; + } + } +} + +Status Split3Inner(cudaStream_t stream, const size_t element_size, const int64_t size0, const int64_t size1, + const int64_t size2, const void* input_data, void* output_data0, void* output_data1, + void* output_data2, const gsl::span& input_shape) { + CUDA_LONG outer_size = 1; + for (size_t i = 0; i < input_shape.size() - 1; ++i) { + outer_size *= static_cast(input_shape[i]); + } + CUDA_LONG inner_size_in_byte = static_cast(input_shape[input_shape.size() - 1] * element_size); + + auto select = [](size_t value) { + if (value % 16 == 0) { + return 16; + } else if (value % 8 == 0) { + return 8; + } else if (value % 4 == 0) { + return 4; + } else if (value % 2 == 0) { + return 2; + } else { + return 1; + } + }; + + auto input_v = reinterpret_cast(input_data); + auto output_v0 = reinterpret_cast(output_data0); + auto output_v1 = reinterpret_cast(output_data1); + auto output_v2 = reinterpret_cast(output_data2); + auto size0_in_byte = size0 * element_size; + auto size1_in_byte = size1 * element_size; + auto size2_in_byte = size2 * element_size; + + auto VEC_SIZE = std::min(select(size0_in_byte), std::min(select(size1_in_byte), select(size2_in_byte))); + auto min_output_vec_size = std::min(select(output_v0), std::min(select(output_v1), select(output_v2))); + VEC_SIZE = std::min(VEC_SIZE, std::min(select(input_v), min_output_vec_size)); + + // determine threads based on the size of the output + auto threadsPerBlock = kNumThreadsPerBlock; + if ((inner_size_in_byte / VEC_SIZE) <= 128) { + // use less threads when the size is small + threadsPerBlock = 128; + } + + switch (VEC_SIZE) { +#define CASE_ELEMENT_TYPE(type) \ + _Split3InnerKernel<<>>( \ + size0_in_byte, \ + size1_in_byte, \ + size2_in_byte, \ + input_data, \ + output_data0, \ + output_data1, \ + output_data2, \ + inner_size_in_byte) + case 16: + CASE_ELEMENT_TYPE(int4); + break; + case 8: + CASE_ELEMENT_TYPE(int64_t); + break; + case 4: + CASE_ELEMENT_TYPE(int32_t); + break; + case 2: + CASE_ELEMENT_TYPE(int16_t); + break; + default: + CASE_ELEMENT_TYPE(int8_t); + break; +#undef CASE_ELEMENT_TYPE + } + + return Status::OK(); +} + } // namespace cuda } // namespace onnxruntime diff --git a/onnxruntime/core/providers/cuda/tensor/split_impl.h b/onnxruntime/core/providers/cuda/tensor/split_impl.h index 16961cfb7d22d..62e0da7716608 100644 --- a/onnxruntime/core/providers/cuda/tensor/split_impl.h +++ b/onnxruntime/core/providers/cuda/tensor/split_impl.h @@ -19,5 +19,9 @@ Status SplitImpl(cudaStream_t stream, const size_t element_size, const int block const int64_t* axis_dimension_input_output_mapping, const int num_outputs, const void* input_data, void** output_data, const size_t input_size); +Status Split3Inner(cudaStream_t stream, const size_t element_size, const int64_t size0, const int64_t size1, + const int64_t size2, const void* input_data, void* output_data0, void* output_data1, + void* output_data2, const gsl::span& input_shape); + } // namespace cuda } // namespace onnxruntime diff --git a/onnxruntime/test/providers/cpu/tensor/split_op_test.cc b/onnxruntime/test/providers/cpu/tensor/split_op_test.cc index 15a7d7cd9fdbf..066302a4a37d1 100644 --- a/onnxruntime/test/providers/cpu/tensor/split_op_test.cc +++ b/onnxruntime/test/providers/cpu/tensor/split_op_test.cc @@ -815,5 +815,62 @@ TEST(SplitOperatorTest, Split18_NumOutputsUnevenSplitAxis1) { RunTest(axis, {}, input, outputs, {kTensorrtExecutionProvider, kQnnExecutionProvider}, false, true, num_outputs, false); } +TEST(SplitOperatorTest, Split3Inner) { + constexpr int64_t axis = -1; + using ShapeAndDataT = ShapeAndData; + std::vector outputs; + int64_t num_outputs = -1; // when provides split_sizes, then num_outputs should not be provided + const int batch = 16; + const int data_len = 96; // should be multiple of 3 + + // create input with shape {b, l}, and data from 1 ~ b*l + auto input = CreateInput({batch, data_len}); // input is 1.f ~ 48.f + + // slice the input data by start and end in axis of -1 + auto gen_output = [&](int start, int end) { + std::vector data0; + auto input_data = input.second; + for (int b = 0; b < batch; b++) { + for (int i = start; i < end; i++) { + data0.push_back(input_data[b * data_len + i]); + } + } + return ShapeAndDataT{{batch, end - start}, data0}; + }; + + auto do_test = [&](std::vector& splits) { + outputs.clear(); + outputs.push_back(gen_output(0, splits[0])); + outputs.push_back(gen_output(splits[0], splits[1])); + outputs.push_back(gen_output(splits[1], data_len)); + + RunTest(axis, {splits[0], splits[1] - splits[0], data_len - splits[1]}, input, outputs, {kTensorrtExecutionProvider, kQnnExecutionProvider}, false, true, num_outputs); + }; + + // split into 3 same size, and aligned to 16 + std::vector splits{data_len / 3, data_len / 3 * 2}; + do_test(splits); + + // test split with data alignment is 8 + splits[0] = splits[0] + 8; + splits[1] = splits[1] + 8; + do_test(splits); + + // test split with data alignment is 4 + splits[0] = splits[0] + 4; + splits[1] = splits[1] + 4; + do_test(splits); + + // test split with data alignment is 2 + splits[0] = splits[0] + 2; + splits[1] = splits[1] + 2; + do_test(splits); + + // test split with data alignment is 1 + splits[0] = splits[0] + 1; + splits[1] = splits[1] + 1; + do_test(splits); +} + } // namespace test } // namespace onnxruntime From b49788e68bef20f5f7dbf4f0d1ab8691b02f6087 Mon Sep 17 00:00:00 2001 From: Wanming Lin Date: Thu, 27 Jun 2024 23:09:13 +0800 Subject: [PATCH 36/52] [WebNN EP] Fixed bug in Expand implementation (#21163) ONNX's Expand supports bidirectionally broadcast, while WebNN's expand op only supports unidirectionally broadcast. Thus we should calculate the output shape for 'newShape' input of WebNN's expand op. --- .../core/providers/webnn/builders/helper.cc | 26 ++++++++++++++----- .../core/providers/webnn/builders/helper.h | 6 ++--- .../webnn/builders/impl/expand_op_builder.cc | 20 +++++++------- 3 files changed, 31 insertions(+), 21 deletions(-) diff --git a/onnxruntime/core/providers/webnn/builders/helper.cc b/onnxruntime/core/providers/webnn/builders/helper.cc index 5d1794ed0afa8..44e6953db438e 100644 --- a/onnxruntime/core/providers/webnn/builders/helper.cc +++ b/onnxruntime/core/providers/webnn/builders/helper.cc @@ -136,21 +136,33 @@ bool IsSupportedDataType(const int32_t data_type, supported_data_types.end(); } -bool IsValidMultidirectionalBroadcast(std::vector& shape_a, - std::vector& shape_b, - const logging::Logger& logger) { +bool GetBidirectionalBroadcastShape(std::vector& shape_a, + std::vector& shape_b, + std::vector& output_shape) { size_t size_a = shape_a.size(); size_t size_b = shape_b.size(); size_t smaller_size = std::min(size_a, size_b); - for (size_t i = 0; i < smaller_size; i++) { + size_t larger_size = std::max(size_a, size_b); + + output_shape.resize(larger_size); + + for (size_t i = 0; i < larger_size; i++) { // right alignment size_t axis_a = size_a - i - 1; size_t axis_b = size_b - i - 1; - // Broadcastable tensors must either have each dimension the same size or equal to one. - if (shape_a[axis_a] != shape_b[axis_b] && shape_a[axis_a] != 1 && shape_b[axis_b] != 1) { - return false; + + if (i < smaller_size) { + // Broadcastable tensors must either have each dimension the same size or equal to one. + if (shape_a[axis_a] != shape_b[axis_b] && shape_a[axis_a] != 1 && shape_b[axis_b] != 1) { + return false; + } + output_shape[larger_size - i - 1] = std::max(shape_a[axis_a], shape_b[axis_b]); + } else { + // For the remaining dimensions in the larger tensor, copy the dimension size directly to the output shape. + output_shape[larger_size - i - 1] = (size_a > size_b) ? shape_a[axis_a] : shape_b[axis_b]; } } + return true; } diff --git a/onnxruntime/core/providers/webnn/builders/helper.h b/onnxruntime/core/providers/webnn/builders/helper.h index 4ee3f891f92ca..496f886e5a076 100644 --- a/onnxruntime/core/providers/webnn/builders/helper.h +++ b/onnxruntime/core/providers/webnn/builders/helper.h @@ -277,9 +277,9 @@ static const std::unordered_set webnn_supp bool IsSupportedDataType(const int32_t data_type, const std::unordered_set& supported_data_types); -bool IsValidMultidirectionalBroadcast(std::vector& shape_a, - std::vector& shape_b, - const logging::Logger& logger); +bool GetBidirectionalBroadcastShape(std::vector& shape_a, + std::vector& shape_b, + std::vector& output_shape); bool SetWebnnDataType(emscripten::val& desc, const int32_t data_type); diff --git a/onnxruntime/core/providers/webnn/builders/impl/expand_op_builder.cc b/onnxruntime/core/providers/webnn/builders/impl/expand_op_builder.cc index 9d4de45fdafd1..9c75c00fa9273 100644 --- a/onnxruntime/core/providers/webnn/builders/impl/expand_op_builder.cc +++ b/onnxruntime/core/providers/webnn/builders/impl/expand_op_builder.cc @@ -44,18 +44,19 @@ Status ExpandOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder, const auto& input_defs = node.InputDefs(); const auto& initializers(model_builder.GetInitializerTensors()); const auto& shape_tensor = *initializers.at(input_defs[1]->Name()); - std::vector new_shape; + std::vector new_shape; ORT_RETURN_IF_NOT(ReadIntArrayFrom1DTensor(shape_tensor, new_shape, logger), "Cannot get shape."); emscripten::val input = model_builder.GetOperand(input_defs[0]->Name()); std::vector input_shape; ORT_RETURN_IF_NOT(GetShape(*input_defs[0], input_shape, logger), "Cannot get input's shape."); - if (new_shape.size() < input_shape.size()) { - // Enlarge new shape to input.rank, right aligned with leading ones - new_shape.insert(new_shape.begin(), input_shape.size() - new_shape.size(), 1); - } + + std::vector output_shape; + ORT_RETURN_IF_NOT(GetBidirectionalBroadcastShape(input_shape, new_shape, output_shape), "Cannot get output shape."); + emscripten::val output = model_builder.GetBuilder().call("expand", - input, emscripten::val::array(new_shape)); + input, + emscripten::val::array(GetVecUint32FromVecInt64(output_shape))); model_builder.AddOperand(node.OutputDefs()[0]->Name(), std::move(output)); return Status::OK(); } @@ -95,11 +96,8 @@ bool ExpandOpBuilder::IsOpSupportedImpl(const InitializedTensorSet& initializers return false; } - if (new_shape.size() > input_shape.size()) { - LOGS(logger, VERBOSE) << "The size of shape must be less than or equal to the rank of input."; - } - - if (!IsValidMultidirectionalBroadcast(input_shape, new_shape, logger)) { + std::vector output_shape; + if (!GetBidirectionalBroadcastShape(input_shape, new_shape, output_shape)) { LOGS(logger, VERBOSE) << "The input cannot expand to shape " << GetShapeString(new_shape); return false; } From 8f738d8e9f03baa8958b22eb5a50d723f1ddbad4 Mon Sep 17 00:00:00 2001 From: Yi Zhang Date: Thu, 27 Jun 2024 23:49:32 +0800 Subject: [PATCH 37/52] [Fix] Throwes one excepiton while Llama2 parity_check fails (#21160) ### Description ### Motivation and Context The pipeline is green even Llama2 parity_check fails. The PR should be merged after the below exception is solved. ''' 2024-06-25 03:49:43.621298481 [E:onnxruntime:, sequential_executor.cc:514 ExecuteKernel] Non-zero status code returned while running Expand node. Name:'/model/Expand' Status Message: /model/Expand: left operand cannot broadcast on dim 3 LeftShape: {1,1,9,9}, RightShape: {2,1,9,17} An error occurred while verifying parity: Error in execution: Non-zero status code returned while running Expand node. Name:'/model/Expand' Status Message: /model/Expand: left operand cannot broadcast on dim 3 LeftShape: {1,1,9,9}, RightShape: {2,1,9,17} Traceback (most recent call last): File "/workspace/onnxruntime/python/tools/transformers/models/llama/convert_to_onnx.py", line 1043, in main parity_check(parity_cmd) File "/workspace/onnxruntime/python/tools/transformers/models/llama/llama_parity.py", line 298, in main verify_parity(args, location, use_auth_token, kv_cache_ortvalues, pytorch_model=llama, config=config) File "/workspace/onnxruntime/python/tools/transformers/models/llama/llama_parity.py", line 137, in verify_parity ort_model.run_with_iobinding(io_binding) File "/home/onnxruntimedev/.local/lib/python3.8/site-packages/onnxruntime/capi/onnxruntime_inference_collection.py", line 331, in run_with_iobinding self._sess.run_with_iobinding(iobinding._iobinding, run_options) RuntimeError: Error in execution: Non-zero status code returned while running Expand node. Name:'/model/Expand' Status Message: /model/Expand: left operand cannot broadcast on dim 3 LeftShape: {1,1,9,9}, RightShape: {2,1,9,17} ''' The exception looks caused by #19832 --- .../python/tools/transformers/models/llama/convert_to_onnx.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/onnxruntime/python/tools/transformers/models/llama/convert_to_onnx.py b/onnxruntime/python/tools/transformers/models/llama/convert_to_onnx.py index 8a33544654e05..f701e465b9153 100644 --- a/onnxruntime/python/tools/transformers/models/llama/convert_to_onnx.py +++ b/onnxruntime/python/tools/transformers/models/llama/convert_to_onnx.py @@ -1052,7 +1052,8 @@ def main(): logger.info(f"check parity with cmd: {parity_cmd}") parity_check(parity_cmd) except Exception as e: - logger.warning(f"An error occurred while verifying parity: {e}", exc_info=True) + logger.exception(f"An error occurred while verifying parity: {e}") + sys.exit(-1) if __name__ == "__main__": From 9eb1c2a7a3d46add7a40f12e87aec61ea4ebd9ae Mon Sep 17 00:00:00 2001 From: Guenther Schmuelling Date: Thu, 27 Jun 2024 10:20:48 -0700 Subject: [PATCH 38/52] support for layernorm in webgpu pre opset-17 (#21121) handled the same way cpu does --- js/web/docs/webgpu-operators.md | 2 +- onnxruntime/contrib_ops/js/js_contrib_kernels.cc | 5 +++++ onnxruntime/contrib_ops/js/layer_norm.cc | 13 +++++++++++++ 3 files changed, 19 insertions(+), 1 deletion(-) diff --git a/js/web/docs/webgpu-operators.md b/js/web/docs/webgpu-operators.md index 919b005ec4c21..3ee9441eeb981 100644 --- a/js/web/docs/webgpu-operators.md +++ b/js/web/docs/webgpu-operators.md @@ -58,7 +58,7 @@ Do not modify directly.* | HardSigmoid | ai.onnx(6+) | | | If | ai.onnx(1-10,11-12,13-18,19+) | | | InstanceNormalization | ai.onnx(6+); com.ms.internal.nhwc(6+) | | -| LayerNormalization | ai.onnx(17+) | | +| LayerNormalization | ai.onnx(1-16,17+) | | | LeakyRelu | ai.onnx(6-15,16+) | | | Less | ai.onnx(7-8,9-12,13+) | | | LessOrEqual | ai.onnx(12-15,16+) | | diff --git a/onnxruntime/contrib_ops/js/js_contrib_kernels.cc b/onnxruntime/contrib_ops/js/js_contrib_kernels.cc index 7bc3414c89978..11899feb6e1dc 100644 --- a/onnxruntime/contrib_ops/js/js_contrib_kernels.cc +++ b/onnxruntime/contrib_ops/js/js_contrib_kernels.cc @@ -14,6 +14,8 @@ class ONNX_OPERATOR_KERNEL_CLASS_NAME(kJsExecutionProvider, kMSDomain, 1, FastGe class ONNX_OPERATOR_KERNEL_CLASS_NAME(kJsExecutionProvider, kMSDomain, 1, FusedConv); class ONNX_OPERATOR_KERNEL_CLASS_NAME(kJsExecutionProvider, kMSDomain, 1, Gelu); class ONNX_OPERATOR_KERNEL_CLASS_NAME(kJsExecutionProvider, kMSDomain, 1, GroupQueryAttention); +// LayerNormalization used to be a contrib op that (incorrectly) used kOnnxDomain so we need to version it +class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 1, 16, LayerNormalization); class ONNX_OPERATOR_KERNEL_CLASS_NAME(kJsExecutionProvider, kMSDomain, 1, MatMulNBits); class ONNX_OPERATOR_KERNEL_CLASS_NAME(kJsExecutionProvider, kMSDomain, 1, MultiHeadAttention); class ONNX_OPERATOR_KERNEL_CLASS_NAME(kJsExecutionProvider, kMSDomain, 1, QuickGelu); @@ -23,6 +25,7 @@ class ONNX_OPERATOR_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 1, Simp class ONNX_OPERATOR_KERNEL_CLASS_NAME(kJsExecutionProvider, kMSDomain, 1, SkipSimplifiedLayerNormalization); template <> + KernelCreateInfo BuildKernelCreateInfo() { KernelCreateInfo info; return info; @@ -37,6 +40,8 @@ Status RegisterJsContribKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, + // LayerNormalization used to be a contrib op that (incorrectly) used kOnnxDomain so we need to version it + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, diff --git a/onnxruntime/contrib_ops/js/layer_norm.cc b/onnxruntime/contrib_ops/js/layer_norm.cc index 814543a9905e0..ec4603cc69de4 100644 --- a/onnxruntime/contrib_ops/js/layer_norm.cc +++ b/onnxruntime/contrib_ops/js/layer_norm.cc @@ -8,6 +8,19 @@ namespace onnxruntime { namespace contrib { namespace js { +// LayerNormalization used to be a contrib op +// that (incorrectly) used kOnnxDomain so we need to version it +ONNX_OPERATOR_VERSIONED_KERNEL_EX( + LayerNormalization, + kOnnxDomain, + 1, + 16, + kJsExecutionProvider, + (*KernelDefBuilder::Create()) + .TypeConstraint("T", onnxruntime::js::JsepSupportedFloatTypes()) + .TypeConstraint("U", onnxruntime::js::JsepSupportedFloatTypes()), + onnxruntime::js::LayerNorm); + ONNX_OPERATOR_KERNEL_EX( SimplifiedLayerNormalization, kOnnxDomain, From 78316c8cbe6cf741bcce3c5ce25ee52945306998 Mon Sep 17 00:00:00 2001 From: Wanming Lin Date: Fri, 28 Jun 2024 02:56:56 +0800 Subject: [PATCH 39/52] [WebNN EP] Remove useless variable unpacked_tensors_ (#21189) --- onnxruntime/core/providers/webnn/builders/model_builder.cc | 3 +-- onnxruntime/core/providers/webnn/builders/model_builder.h | 1 - 2 files changed, 1 insertion(+), 3 deletions(-) diff --git a/onnxruntime/core/providers/webnn/builders/model_builder.cc b/onnxruntime/core/providers/webnn/builders/model_builder.cc index c46b04a3c29d9..6b0e1495f552d 100644 --- a/onnxruntime/core/providers/webnn/builders/model_builder.cc +++ b/onnxruntime/core/providers/webnn/builders/model_builder.cc @@ -104,8 +104,7 @@ Status ModelBuilder::RegisterInitializers() { if (tensor.has_raw_data()) { tensor_ptr = reinterpret_cast(const_cast(tensor.raw_data().c_str())); } else { - unpacked_tensors_.push_back({}); - std::vector& unpacked_tensor = unpacked_tensors_.back(); + std::vector unpacked_tensor; ORT_RETURN_IF_ERROR(onnxruntime::utils::UnpackInitializerData(tensor, unpacked_tensor)); tensor_ptr = reinterpret_cast(unpacked_tensor.data()); } diff --git a/onnxruntime/core/providers/webnn/builders/model_builder.h b/onnxruntime/core/providers/webnn/builders/model_builder.h index 80077b3abe56d..6a1688f16d2a6 100644 --- a/onnxruntime/core/providers/webnn/builders/model_builder.h +++ b/onnxruntime/core/providers/webnn/builders/model_builder.h @@ -66,7 +66,6 @@ class ModelBuilder { emscripten::val wnn_builder_ = emscripten::val::object(); DataLayout preferred_layout_; WebnnDeviceType wnn_device_type_; - std::vector> unpacked_tensors_; InlinedHashMap wnn_operands_; std::vector input_names_; std::vector output_names_; From d1ab94c2b0697c0268abca1f096f641675dc3b07 Mon Sep 17 00:00:00 2001 From: Changming Sun Date: Thu, 27 Jun 2024 13:50:53 -0700 Subject: [PATCH 40/52] Add compatibility for NumPy 2.0 (#21085) ### Description As suggested by SciPy's doc, we will `Build against NumPy 2.0.0, then it will work for all NumPy versions with the same major version number (NumPy does maintain backwards ABI compatibility), and as far back as NumPy 1.19 series at the time of writing` I think it works because in [numpyconfig.h#L64](https://github.com/numpy/numpy/blob/main/numpy/_core/include/numpy/numpyconfig.h#L64) there is a macro NPY_FEATURE_VERSION. By default it is set to NPY_1_19_API_VERSION. And the NPY_FEATURE_VERSION macro controls ABI. This PR only upgrade the build time dependency; When a user installs ONNX Runtime, they still can use numpy 1.x. ### Motivation and Context Recently numpy published a new version, 2.0.0, which is incompatible with the latest ONNX Runtime release. --- cmake/CMakeLists.txt | 12 +++---- cmake/onnxruntime_python.cmake | 3 ++ onnxruntime/test/onnx/gen_test_models.py | 2 +- .../python/quantization/test_quant_util.py | 32 +++++++++++-------- .../test_quantizeblockwise_bnb4.py | 6 ++-- requirements.txt.in => requirements.txt | 2 +- .../py-packaging-selectable-stage.yml | 22 ------------- .../templates/py-packaging-stage.yml | 11 ------- .../azure-pipelines/templates/py-win-gpu.yml | 14 -------- .../templates/py-win-x64-qnn.yml | 11 ------- .../python/cpu/scripts/requirements.txt | 5 ++- .../python/cpu/scripts/requirements.txt | 5 ++- .../python/cuda/scripts/requirements.txt | 5 ++- .../docker/scripts/lort/requirements.txt | 3 +- .../docker/scripts/manylinux/requirements.txt | 5 ++- .../linux/docker/scripts/requirements.txt | 5 ++- .../stage1/requirements_rocm/requirements.txt | 5 ++- .../stage1/torch_eager_cpu/requirements.txt | 3 +- .../ortmodule/stage2/requirements.txt | 5 ++- .../linux/test_custom_ops_pytorch_export.sh | 2 +- .../github/windows/eager/requirements.txt | 3 +- tools/ci_build/github/windows/helpers.ps1 | 8 +++-- 22 files changed, 56 insertions(+), 113 deletions(-) rename requirements.txt.in => requirements.txt (60%) diff --git a/cmake/CMakeLists.txt b/cmake/CMakeLists.txt index 9670dcb97abb2..6ba0db789965c 100644 --- a/cmake/CMakeLists.txt +++ b/cmake/CMakeLists.txt @@ -579,6 +579,10 @@ if (onnxruntime_BUILD_SHARED_LIB OR onnxruntime_ENABLE_PYTHON) else() find_package(Python 3.8 COMPONENTS Interpreter Development.Module NumPy) endif() + message("Numpy version: ${Python_NumPy_VERSION}") + if(Python_NumPy_VERSION VERSION_LESS "2.0.0") + message(WARNING "The build binary will not be compatible with NumPy 2.0 because the NumPy installed on this machine is too low.") + endif() else() find_package(Python 3.8 COMPONENTS Interpreter) endif() @@ -1406,14 +1410,6 @@ string(APPEND ORT_BUILD_INFO "build type=${CMAKE_BUILD_TYPE}") string(APPEND ORT_BUILD_INFO ", cmake cxx flags: ${CMAKE_CXX_FLAGS}") configure_file(onnxruntime_config.h.in ${CMAKE_CURRENT_BINARY_DIR}/onnxruntime_config.h) get_property(onnxruntime_GENERATOR_IS_MULTI_CONFIG GLOBAL PROPERTY GENERATOR_IS_MULTI_CONFIG) -if (onnxruntime_GENERATOR_IS_MULTI_CONFIG) - configure_file(../requirements.txt.in ${CMAKE_CURRENT_BINARY_DIR}/Debug/requirements.txt) - configure_file(../requirements.txt.in ${CMAKE_CURRENT_BINARY_DIR}/Release/requirements.txt) - configure_file(../requirements.txt.in ${CMAKE_CURRENT_BINARY_DIR}/RelWithDebInfo/requirements.txt) - configure_file(../requirements.txt.in ${CMAKE_CURRENT_BINARY_DIR}/MinSizeRel/requirements.txt) -else() - configure_file(../requirements.txt.in ${CMAKE_CURRENT_BINARY_DIR}/requirements.txt) -endif() if (onnxruntime_USE_CUDA) set(CMAKE_CUDA_RUNTIME_LIBRARY Shared) diff --git a/cmake/onnxruntime_python.cmake b/cmake/onnxruntime_python.cmake index 3c2833d87d652..062cc8f9dbff3 100644 --- a/cmake/onnxruntime_python.cmake +++ b/cmake/onnxruntime_python.cmake @@ -562,6 +562,9 @@ add_custom_command( COMMAND ${CMAKE_COMMAND} -E copy ${ONNXRUNTIME_ROOT}/__init__.py $/onnxruntime/ + COMMAND ${CMAKE_COMMAND} -E copy + ${REPO_ROOT}/requirements.txt + $ COMMAND ${CMAKE_COMMAND} -E copy ${REPO_ROOT}/ThirdPartyNotices.txt $/onnxruntime/ diff --git a/onnxruntime/test/onnx/gen_test_models.py b/onnxruntime/test/onnx/gen_test_models.py index 1a64df2936810..a5224925251cf 100644 --- a/onnxruntime/test/onnx/gen_test_models.py +++ b/onnxruntime/test/onnx/gen_test_models.py @@ -144,7 +144,7 @@ def test_abs(output_dir): ) generate_abs_op_test( TensorProto.UINT16, - np.uint16([-32767, -4, 0, 3, 32767]), + np.uint16([0, 3, 32767, 65535]), os.path.join(output_dir, "test_abs_uint16"), ) generate_abs_op_test( diff --git a/onnxruntime/test/python/quantization/test_quant_util.py b/onnxruntime/test/python/quantization/test_quant_util.py index 7b3fc08982ac1..96d841654adbd 100644 --- a/onnxruntime/test/python/quantization/test_quant_util.py +++ b/onnxruntime/test/python/quantization/test_quant_util.py @@ -37,41 +37,45 @@ def _compute_scale_zp(rmin, rmax, qmin, qmax, qtype, symmetric=False, min_real_r assert isinstance(scale, numpy.ndarray) return [float(zp), float(scale)] - self.assertEqual(_compute_scale_zp(0.0, 0.0, -127, 127, numpy.int8, symmetric=True), [0, 1.0]) - self.assertEqual(_compute_scale_zp(1.0, -1.0, -127, 127, numpy.int8, symmetric=True), [0, 1.0]) - self.assertEqual(_compute_scale_zp(0.0, 0.0, 0, 255, numpy.uint8, symmetric=True), [0, 1.0]) - self.assertEqual(_compute_scale_zp(1.0, -1.0, 0, 255, numpy.uint8, symmetric=True), [0, 1.0]) + numpy.testing.assert_allclose(_compute_scale_zp(0.0, 0.0, -127, 127, numpy.int8, symmetric=True), [0, 1.0]) + numpy.testing.assert_allclose(_compute_scale_zp(1.0, -1.0, -127, 127, numpy.int8, symmetric=True), [0, 1.0]) + numpy.testing.assert_allclose(_compute_scale_zp(0.0, 0.0, 0, 255, numpy.uint8, symmetric=True), [0, 1.0]) + numpy.testing.assert_allclose(_compute_scale_zp(1.0, -1.0, 0, 255, numpy.uint8, symmetric=True), [0, 1.0]) - self.assertEqual( + numpy.testing.assert_allclose( _compute_scale_zp(-1.0, 2.0, -127, 127, numpy.int8, symmetric=True), [0, numpy.float32(2.0 / 127)] ) - self.assertEqual( + numpy.testing.assert_allclose( _compute_scale_zp(-1.0, 2.0, -127, 127, numpy.int8, symmetric=False), [-42, numpy.float32(3.0 / 254)] ) - self.assertEqual( + numpy.testing.assert_allclose( _compute_scale_zp(-1.0, 2.0, 0, 255, numpy.uint8, symmetric=True), [128, numpy.float32(4.0 / 255)] ) - self.assertEqual( + numpy.testing.assert_allclose( _compute_scale_zp(-1.0, 2.0, 0, 255, numpy.uint8, symmetric=False), [85, numpy.float32(3.0 / 255)] ) tiny_float = numpy.float32(numpy.finfo(numpy.float32).tiny * 0.1) - self.assertEqual(_compute_scale_zp(-tiny_float, tiny_float, 0, 255, numpy.uint8, symmetric=True), [0, 1.0]) - self.assertEqual(_compute_scale_zp(-tiny_float, 0.0, 0, 255, numpy.uint8, symmetric=False), [0, 1.0]) + numpy.testing.assert_allclose( + _compute_scale_zp(-tiny_float, tiny_float, 0, 255, numpy.uint8, symmetric=True), [0, 1.0] + ) + numpy.testing.assert_allclose( + _compute_scale_zp(-tiny_float, 0.0, 0, 255, numpy.uint8, symmetric=False), [0, 1.0] + ) # Test enforcing a minimum floatint-point range. - self.assertEqual( + numpy.testing.assert_allclose( _compute_scale_zp(0.0, 0.0, 0, 255, numpy.uint8, symmetric=False, min_real_range=0.0001), [0, 0.0001 / 255] ) - self.assertEqual( + numpy.testing.assert_allclose( _compute_scale_zp(0.0, 0.0, -128, 127, numpy.int8, symmetric=True, min_real_range=0.0001), [0, 0.0002 / 255] ) - self.assertEqual( + numpy.testing.assert_allclose( _compute_scale_zp(0.0, 0.0, 0, 65535, numpy.uint16, symmetric=False, min_real_range=0.0001), [0, 0.0001 / 65535], ) - self.assertEqual( + numpy.testing.assert_allclose( _compute_scale_zp(0.0, 0.0, -32768, 32767, numpy.int16, symmetric=True, min_real_range=0.0001), [0, 0.0002 / 65535], ) diff --git a/onnxruntime/test/python/quantization/test_quantizeblockwise_bnb4.py b/onnxruntime/test/python/quantization/test_quantizeblockwise_bnb4.py index 9e9d05fae027d..eafab0c03a951 100644 --- a/onnxruntime/test/python/quantization/test_quantizeblockwise_bnb4.py +++ b/onnxruntime/test/python/quantization/test_quantizeblockwise_bnb4.py @@ -87,7 +87,7 @@ def quantize_blockwise_bnb4_ref(matrix_float: npt.ArrayLike, block_size: int, qu absmax[block_idx] = block_absmax if block_len % 2 != 0: - block = np.append(block, 0.0) + block = np.append(block, np.float32(0.0)) block_len += 1 block *= reciprocal_absmax @@ -131,8 +131,8 @@ def test_quantize_blockwise_bnb4(self): matrix_float = np.random.uniform(-1, 1, (k, n)).astype(type) quant_value_ref, absmax_ref = quantize_blockwise_bnb4_ref(matrix_float, block_size, quant_type) quant_value, absmax = quantize_blockwise_bnb4_target(matrix_float, block_size, quant_type) - assert np.allclose(quant_value_ref, quant_value) - assert np.allclose(absmax_ref, absmax) + np.testing.assert_allclose(quant_value_ref, quant_value) + np.testing.assert_allclose(absmax_ref, absmax) if __name__ == "__main__": diff --git a/requirements.txt.in b/requirements.txt similarity index 60% rename from requirements.txt.in rename to requirements.txt index 89242061fb119..2fd9362c949dd 100644 --- a/requirements.txt.in +++ b/requirements.txt @@ -1,6 +1,6 @@ coloredlogs flatbuffers -numpy >= @Python_NumPy_VERSION@ +numpy >= 1.21.6 packaging protobuf sympy diff --git a/tools/ci_build/github/azure-pipelines/templates/py-packaging-selectable-stage.yml b/tools/ci_build/github/azure-pipelines/templates/py-packaging-selectable-stage.yml index cc07df59da619..3f1c4ef0f8d61 100644 --- a/tools/ci_build/github/azure-pipelines/templates/py-packaging-selectable-stage.yml +++ b/tools/ci_build/github/azure-pipelines/templates/py-packaging-selectable-stage.yml @@ -152,17 +152,6 @@ stages: filename: 'C:\Program Files\Intel\openvino_2021.4.752\bin\setupvars.bat' modifyEnvironment: true - - task: PythonScript@0 - inputs: - scriptSource: inline - script: | - import sys - np_version = 'numpy==1.21.6' if sys.version_info < (3, 11) else 'numpy==1.24.2' - import subprocess - subprocess.call(['pip', 'install', '-q', 'setuptools', 'wheel', np_version]) - workingDirectory: '$(Build.BinariesDirectory)' - displayName: 'Install python modules' - - task: PowerShell@2 displayName: 'Install ONNX' inputs: @@ -419,17 +408,6 @@ stages: modifyEnvironment: true workingFolder: '$(Build.BinariesDirectory)' - - task: PythonScript@0 - inputs: - scriptSource: inline - script: | - import sys - np_version = 'numpy==1.21.6' if sys.version_info < (3, 11) else 'numpy==1.24.2' - import subprocess - subprocess.call(['pip', 'install', '-q', 'setuptools', 'wheel', np_version]) - workingDirectory: '$(Build.BinariesDirectory)' - displayName: 'Install python modules' - - task: PowerShell@2 displayName: 'Install ONNX' inputs: diff --git a/tools/ci_build/github/azure-pipelines/templates/py-packaging-stage.yml b/tools/ci_build/github/azure-pipelines/templates/py-packaging-stage.yml index 588ffca30c262..9e14789f3b234 100644 --- a/tools/ci_build/github/azure-pipelines/templates/py-packaging-stage.yml +++ b/tools/ci_build/github/azure-pipelines/templates/py-packaging-stage.yml @@ -153,17 +153,6 @@ stages: modifyEnvironment: true workingFolder: '$(Build.BinariesDirectory)' - - task: PythonScript@0 - inputs: - scriptSource: inline - script: | - import sys - np_version = 'numpy==1.21.6' if sys.version_info < (3, 11) else 'numpy==1.24.2' - import subprocess - subprocess.call(['pip', 'install', '-q', 'setuptools', 'wheel', np_version]) - workingDirectory: '$(Build.BinariesDirectory)' - displayName: 'Install python modules' - - template: download-deps.yml - task: PythonScript@0 diff --git a/tools/ci_build/github/azure-pipelines/templates/py-win-gpu.yml b/tools/ci_build/github/azure-pipelines/templates/py-win-gpu.yml index c7a74a7f0e9c7..e89227d51de32 100644 --- a/tools/ci_build/github/azure-pipelines/templates/py-win-gpu.yml +++ b/tools/ci_build/github/azure-pipelines/templates/py-win-gpu.yml @@ -89,20 +89,6 @@ stages: tsaConfigFilePath: '$(Build.SourcesDirectory)\.config\tsaoptions.json' appendSourceBranchName: false - - task: PythonScript@0 - inputs: - scriptSource: inline - script: | - import sys - np_version = 'numpy==1.21.6' if sys.version_info < (3, 11) else 'numpy==1.26' - import subprocess - try: - subprocess.check_call(['pip', 'install', '-q', 'setuptools', 'wheel', np_version]) - except subprocess.CalledProcessError: - sys.exit(1) - workingDirectory: '$(Build.BinariesDirectory)' - displayName: 'Install python modules' - - template: download-deps.yml - ${{ if ne(parameters.ENV_SETUP_SCRIPT, '') }}: diff --git a/tools/ci_build/github/azure-pipelines/templates/py-win-x64-qnn.yml b/tools/ci_build/github/azure-pipelines/templates/py-win-x64-qnn.yml index 0d1d3c5ced400..884e6eafee965 100644 --- a/tools/ci_build/github/azure-pipelines/templates/py-win-x64-qnn.yml +++ b/tools/ci_build/github/azure-pipelines/templates/py-win-x64-qnn.yml @@ -60,17 +60,6 @@ jobs: tsaConfigFilePath: '$(Build.SourcesDirectory)\.config\tsaoptions.json' appendSourceBranchName: false - - task: PythonScript@0 - inputs: - scriptSource: inline - script: | - import sys - np_version = 'numpy==1.21.6' if sys.version_info < (3, 11) else 'numpy==1.24.2' - import subprocess - subprocess.call(['pip', 'install', '-q', 'setuptools', 'wheel', np_version]) - workingDirectory: '$(Build.BinariesDirectory)' - displayName: 'Install python modules' - - template: download-deps.yml - task: PythonScript@0 diff --git a/tools/ci_build/github/linux/docker/inference/aarch64/python/cpu/scripts/requirements.txt b/tools/ci_build/github/linux/docker/inference/aarch64/python/cpu/scripts/requirements.txt index cc47718f78a46..a977ccae1922f 100644 --- a/tools/ci_build/github/linux/docker/inference/aarch64/python/cpu/scripts/requirements.txt +++ b/tools/ci_build/github/linux/docker/inference/aarch64/python/cpu/scripts/requirements.txt @@ -1,6 +1,5 @@ -numpy==1.21.6 ; python_version < '3.11' -numpy==1.24.2 ; python_version == '3.11' -numpy==1.26.0 ; python_version >= '3.12' +numpy==1.21.6 ; python_version < '3.9' +numpy==2.0.0 ; python_version >= '3.9' mypy pytest setuptools>=68.2.2 diff --git a/tools/ci_build/github/linux/docker/inference/x86_64/python/cpu/scripts/requirements.txt b/tools/ci_build/github/linux/docker/inference/x86_64/python/cpu/scripts/requirements.txt index cc47718f78a46..a977ccae1922f 100644 --- a/tools/ci_build/github/linux/docker/inference/x86_64/python/cpu/scripts/requirements.txt +++ b/tools/ci_build/github/linux/docker/inference/x86_64/python/cpu/scripts/requirements.txt @@ -1,6 +1,5 @@ -numpy==1.21.6 ; python_version < '3.11' -numpy==1.24.2 ; python_version == '3.11' -numpy==1.26.0 ; python_version >= '3.12' +numpy==1.21.6 ; python_version < '3.9' +numpy==2.0.0 ; python_version >= '3.9' mypy pytest setuptools>=68.2.2 diff --git a/tools/ci_build/github/linux/docker/inference/x86_64/python/cuda/scripts/requirements.txt b/tools/ci_build/github/linux/docker/inference/x86_64/python/cuda/scripts/requirements.txt index cc47718f78a46..a977ccae1922f 100644 --- a/tools/ci_build/github/linux/docker/inference/x86_64/python/cuda/scripts/requirements.txt +++ b/tools/ci_build/github/linux/docker/inference/x86_64/python/cuda/scripts/requirements.txt @@ -1,6 +1,5 @@ -numpy==1.21.6 ; python_version < '3.11' -numpy==1.24.2 ; python_version == '3.11' -numpy==1.26.0 ; python_version >= '3.12' +numpy==1.21.6 ; python_version < '3.9' +numpy==2.0.0 ; python_version >= '3.9' mypy pytest setuptools>=68.2.2 diff --git a/tools/ci_build/github/linux/docker/scripts/lort/requirements.txt b/tools/ci_build/github/linux/docker/scripts/lort/requirements.txt index e9b222fe09711..d76a4337e7487 100644 --- a/tools/ci_build/github/linux/docker/scripts/lort/requirements.txt +++ b/tools/ci_build/github/linux/docker/scripts/lort/requirements.txt @@ -8,7 +8,8 @@ onnx==1.16.1 astunparse expecttest!=0.2.0 hypothesis -numpy +numpy==1.21.6 ; python_version < '3.9' +numpy==2.0.0 ; python_version >= '3.9' psutil pyyaml requests diff --git a/tools/ci_build/github/linux/docker/scripts/manylinux/requirements.txt b/tools/ci_build/github/linux/docker/scripts/manylinux/requirements.txt index bdae9d72a1a63..12db3bd132bb7 100644 --- a/tools/ci_build/github/linux/docker/scripts/manylinux/requirements.txt +++ b/tools/ci_build/github/linux/docker/scripts/manylinux/requirements.txt @@ -1,6 +1,5 @@ -numpy==1.21.6 ; python_version < '3.11' -numpy==1.24.2 ; python_version == '3.11' -numpy==1.26.0 ; python_version >= '3.12' +numpy==1.21.6 ; python_version < '3.9' +numpy==2.0.0 ; python_version >= '3.9' mypy pytest setuptools>=68.2.2 diff --git a/tools/ci_build/github/linux/docker/scripts/requirements.txt b/tools/ci_build/github/linux/docker/scripts/requirements.txt index 3e619ea3dfb56..36af6aa71b075 100644 --- a/tools/ci_build/github/linux/docker/scripts/requirements.txt +++ b/tools/ci_build/github/linux/docker/scripts/requirements.txt @@ -1,7 +1,6 @@ cerberus -numpy==1.21.6 ; python_version < '3.11' -numpy==1.24.2 ; python_version == '3.11' -numpy==1.26.0 ; python_version >= '3.12' +numpy==1.24.4 ; python_version < '3.9' +numpy==2.0.0; python_version >= '3.9' mypy pytest setuptools==69.0.3 diff --git a/tools/ci_build/github/linux/docker/scripts/training/ortmodule/stage1/requirements_rocm/requirements.txt b/tools/ci_build/github/linux/docker/scripts/training/ortmodule/stage1/requirements_rocm/requirements.txt index 57331d6df97d9..89bda11737d10 100644 --- a/tools/ci_build/github/linux/docker/scripts/training/ortmodule/stage1/requirements_rocm/requirements.txt +++ b/tools/ci_build/github/linux/docker/scripts/training/ortmodule/stage1/requirements_rocm/requirements.txt @@ -1,3 +1,2 @@ -numpy==1.21.6 ; python_version < '3.11' -numpy==1.24.2 ; python_version == '3.11' -numpy==1.26.0 ; python_version >= '3.12' \ No newline at end of file +numpy==1.21.6 ; python_version < '3.9' +numpy==2.0.0 ; python_version >= '3.9' diff --git a/tools/ci_build/github/linux/docker/scripts/training/ortmodule/stage1/torch_eager_cpu/requirements.txt b/tools/ci_build/github/linux/docker/scripts/training/ortmodule/stage1/torch_eager_cpu/requirements.txt index 08e251eddbf96..ee4f8bd586804 100644 --- a/tools/ci_build/github/linux/docker/scripts/training/ortmodule/stage1/torch_eager_cpu/requirements.txt +++ b/tools/ci_build/github/linux/docker/scripts/training/ortmodule/stage1/torch_eager_cpu/requirements.txt @@ -5,6 +5,7 @@ setuptools>=68.2.2 cerberus h5py scikit-learn -numpy +numpy==1.21.6 ; python_version < '3.9' +numpy==2.0.0 ; python_version >= '3.9' pandas parameterized diff --git a/tools/ci_build/github/linux/docker/scripts/training/ortmodule/stage2/requirements.txt b/tools/ci_build/github/linux/docker/scripts/training/ortmodule/stage2/requirements.txt index 47f64568f424a..d7fab6a1c8a27 100644 --- a/tools/ci_build/github/linux/docker/scripts/training/ortmodule/stage2/requirements.txt +++ b/tools/ci_build/github/linux/docker/scripts/training/ortmodule/stage2/requirements.txt @@ -1,8 +1,7 @@ pandas scikit-learn -numpy==1.21.6 ; python_version < '3.11' -numpy==1.24.2 ; python_version == '3.11' -numpy==1.26.0 ; python_version >= '3.12' +numpy==1.21.6 ; python_version < '3.9' +numpy==2.0.0 ; python_version >= '3.9' transformers==v4.36.0 accelerate==0.25.0 rsa==4.9 diff --git a/tools/ci_build/github/linux/test_custom_ops_pytorch_export.sh b/tools/ci_build/github/linux/test_custom_ops_pytorch_export.sh index 56f5ff9f9eac0..9cd1222cabfa6 100755 --- a/tools/ci_build/github/linux/test_custom_ops_pytorch_export.sh +++ b/tools/ci_build/github/linux/test_custom_ops_pytorch_export.sh @@ -2,7 +2,7 @@ pip3 install --user --upgrade pip -pip3 install --user numpy==1.19.0 torch pytest +pip3 install --user numpy torch pytest pip3 install --user /build/Release/dist/*.whl export PYTHONPATH=/onnxruntime_src/tools:/usr/local/lib/python3.8/site-packages:$PYTHONPATH diff --git a/tools/ci_build/github/windows/eager/requirements.txt b/tools/ci_build/github/windows/eager/requirements.txt index a820174957185..08e7baa76471b 100644 --- a/tools/ci_build/github/windows/eager/requirements.txt +++ b/tools/ci_build/github/windows/eager/requirements.txt @@ -1,6 +1,7 @@ setuptools wheel -numpy +numpy==1.21.6 ; python_version < '3.9' +numpy==2.0.0 ; python_version >= '3.9' typing_extensions torch==1.13.1 parameterized diff --git a/tools/ci_build/github/windows/helpers.ps1 b/tools/ci_build/github/windows/helpers.ps1 index 0e7d279c9fa49..95a36aa24e904 100644 --- a/tools/ci_build/github/windows/helpers.ps1 +++ b/tools/ci_build/github/windows/helpers.ps1 @@ -635,16 +635,18 @@ function Install-ONNX { if ($lastExitCode -ne 0) { exit $lastExitCode } - + $temp_dir = Get-TempDirectory + $new_requirements_text_file = Join-Path $temp_dir "new_requirements.txt" Write-Host "Installing python packages..." - [string[]]$pip_args = "-m", "pip", "install", "-qq", "--disable-pip-version-check", "setuptools>=68.2.2", "wheel", "numpy", "protobuf==$protobuf_version" + Get-Content "$src_root\tools\ci_build\github\linux\docker\inference\x86_64\python\cpu\scripts\requirements.txt" | Select-String -pattern 'onnx' -notmatch | Out-File $new_requirements_text_file + + [string[]]$pip_args = "-m", "pip", "install", "-qq", "--disable-pip-version-check", "-r", $new_requirements_text_file &"python.exe" $pip_args if ($lastExitCode -ne 0) { exit $lastExitCode } $url=Get-DownloadURL -name onnx -src_root $src_root - $temp_dir = Get-TempDirectory $onnx_src_dir = Join-Path $temp_dir "onnx" $download_finished = DownloadAndExtract -Uri $url -InstallDirectory $onnx_src_dir -Force if(-Not $download_finished){ From 587e92c2791b8af512c15688c0d8469217f92122 Mon Sep 17 00:00:00 2001 From: Yi Zhang Date: Fri, 28 Jun 2024 06:18:26 +0800 Subject: [PATCH 41/52] Add FP32 and INT4 test in Llama2 (#21187) ### Description ### Motivation and Context --- .../azure-pipelines/bigmodels-ci-pipeline.yml | 67 ++++++++++++++----- ...rfile.package_ubi8_cuda_tensorrt10_0_torch | 57 ++++++++++++++++ 2 files changed, 108 insertions(+), 16 deletions(-) create mode 100644 tools/ci_build/github/linux/docker/Dockerfile.package_ubi8_cuda_tensorrt10_0_torch diff --git a/tools/ci_build/github/azure-pipelines/bigmodels-ci-pipeline.yml b/tools/ci_build/github/azure-pipelines/bigmodels-ci-pipeline.yml index 0c0cd8d0a870b..41b3c47ba0396 100644 --- a/tools/ci_build/github/azure-pipelines/bigmodels-ci-pipeline.yml +++ b/tools/ci_build/github/azure-pipelines/bigmodels-ci-pipeline.yml @@ -38,14 +38,6 @@ parameters: type: number default: 0 -resources: - repositories: - - repository: LLaMa2Onnx - type: Github - endpoint: Microsoft - name: Microsoft/Llama-2-Onnx - ref: main - variables: - template: templates/common-variables.yml - name: docker_base_image @@ -287,11 +279,12 @@ stages: workingDirectory: $(Build.SourcesDirectory) condition: ne(variables.hitAnother, 'True') -- stage: Llama2_ONNX_FP16 +- stage: Llama2_7B_ONNX dependsOn: - Build_Onnxruntime_Cuda jobs: - - job: Llama2_ONNX_FP16 + - job: Llama2_7B_ONNX + timeoutInMinutes: 120 variables: skipComponentGovernanceDetection: true workspace: @@ -319,7 +312,7 @@ stages: - template: templates/get-docker-image-steps.yml parameters: - Dockerfile: tools/ci_build/github/linux/docker/Dockerfile.package_ubi8_cuda_tensorrt10_0 + Dockerfile: tools/ci_build/github/linux/docker/Dockerfile.package_ubi8_cuda_tensorrt10_0_torch Context: tools/ci_build/github/linux/docker/ ScriptName: tools/ci_build/get_docker_image.py DockerBuildArgs: " @@ -327,7 +320,7 @@ stages: --build-arg BASEIMAGE=${{ variables.docker_base_image }} --build-arg TRT_VERSION=${{ variables.linux_trt_version }} " - Repository: onnxruntimeubi8packagestest + Repository: onnxruntimeubi8packagestest_torch UpdateDepsTxt: false - task: DownloadPackage@1 @@ -343,7 +336,7 @@ stages: docker run --rm --gpus all -v $(Build.SourcesDirectory):/workspace \ -v $(Build.BinariesDirectory)/ort-artifact/:/ort-artifact \ -v $(Agent.TempDirectory)/meta_llama2_7b_hf:/meta-llama2 \ - onnxruntimeubi8packagestest \ + onnxruntimeubi8packagestest_torch \ bash -c " set -ex; \ pushd /workspace/onnxruntime/python/tools/transformers/ ; \ @@ -352,14 +345,56 @@ stages: python3 -m pip install -r requirements.txt ; \ popd ; \ python3 -m pip install /ort-artifact/*.whl ; \ - python3 -m pip uninstall -y torch ; \ - python3 -m pip install torch --index-url https://download.pytorch.org/whl/cu118 ; \ - python3 -m models.llama.convert_to_onnx -m meta-llama/Llama-2-7b-hf --output llama2-7b-fp16 --precision fp16 --execution_provider cuda --input /meta-llama2 --small_gpu ;\ + python3 -m models.llama.convert_to_onnx -m meta-llama/Llama-2-7b-hf --input /meta-llama2 --output llama2-7b-fp16 --precision fp16 --execution_provider cuda --small_gp;\ + ls -l llama2-7b-fp16; \ + du -sh llama2-7b-fp16; \ popd ; \ " displayName: 'Run Llama2 to Onnx F16 and parity Test' workingDirectory: $(Build.SourcesDirectory) + - script: | + docker run --rm --gpus all -v $(Build.SourcesDirectory):/workspace \ + -v $(Build.BinariesDirectory)/ort-artifact/:/ort-artifact \ + -v $(Agent.TempDirectory)/meta_llama2_7b_hf:/meta-llama2 \ + onnxruntimeubi8packagestest_torch \ + bash -c " + set -ex; \ + pushd /workspace/onnxruntime/python/tools/transformers/ ; \ + python3 -m pip install --upgrade pip ; \ + pushd models/llama ; \ + python3 -m pip install -r requirements.txt ; \ + popd ; \ + python3 -m pip install /ort-artifact/*.whl ; \ + python3 -m models.llama.convert_to_onnx -m meta-llama/Llama-2-7b-hf --input /meta-llama2 --output llama2-7b-fp32-gpu --precision fp32 --execution_provider cuda;\ + ls -l llama2-7b-fp32-gpu; \ + du -sh llama2-7b-fp32-gpu; \ + popd ; \ + " + displayName: 'Run Llama2 to Onnx fp32 and parity Test' + workingDirectory: $(Build.SourcesDirectory) + + - script: | + docker run --rm --gpus all -v $(Build.SourcesDirectory):/workspace \ + -v $(Build.BinariesDirectory)/ort-artifact/:/ort-artifact \ + -v $(Agent.TempDirectory)/meta_llama2_7b_hf:/meta-llama2 \ + onnxruntimeubi8packagestest_torch \ + bash -c " + set -ex; \ + pushd /workspace/onnxruntime/python/tools/transformers/ ; \ + python3 -m pip install --upgrade pip ; \ + pushd models/llama ; \ + python3 -m pip install -r requirements.txt ; \ + popd ; \ + python3 -m pip install /ort-artifact/*.whl ; \ + python3 -m models.llama.convert_to_onnx -m meta-llama/Llama-2-7b-hf --input /meta-llama2 --output llama2-7b-int4-gpu --precision int4 --execution_provider cuda --use_gqa;\ + ls -l llama2-7b-int4-gpu; \ + du -sh llama2-7b-int4-gpu; \ + popd ; \ + " + displayName: 'Run Llama2 to Onnx INT4 and parity Test' + workingDirectory: $(Build.SourcesDirectory) + - stage: Whisper_ONNX dependsOn: - Build_Onnxruntime_Cuda diff --git a/tools/ci_build/github/linux/docker/Dockerfile.package_ubi8_cuda_tensorrt10_0_torch b/tools/ci_build/github/linux/docker/Dockerfile.package_ubi8_cuda_tensorrt10_0_torch new file mode 100644 index 0000000000000..4542d3a3f2e4c --- /dev/null +++ b/tools/ci_build/github/linux/docker/Dockerfile.package_ubi8_cuda_tensorrt10_0_torch @@ -0,0 +1,57 @@ +# -------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------- +# Dockerfile to Test ONNX Runtime on UBI8 with TensorRT 10.0 and CUDA 11.8 by default + +# Build base image with required system packages +ARG BASEIMAGE=nvidia/cuda:11.8.0-cudnn8-devel-ubi8 +ARG TRT_VERSION=10.0.1.6-1.cuda11.8 +FROM $BASEIMAGE AS base +ARG TRT_VERSION +ENV PATH /opt/python/cp38-cp38/bin:/usr/local/nvidia/bin:/usr/local/cuda/bin:/usr/src/tensorrt/bin:${PATH} + +RUN dnf install -y bash wget &&\ + dnf clean dbcache + +RUN pip3 install --upgrade pip +RUN pip3 install setuptools>=68.2.2 + +#Install TensorRT only if TRT_VERSION is not empty +RUN if [ -n "$TRT_VERSION" ]; then \ + echo "TRT_VERSION is $TRT_VERSION" && \ + dnf -y install \ + libnvinfer10-${TRT_VERSION} \ + libnvinfer-headers-devel-${TRT_VERSION} \ + libnvinfer-devel-${TRT_VERSION} \ + libnvinfer-lean10-${TRT_VERSION} \ + libnvonnxparsers10-${TRT_VERSION} \ + libnvonnxparsers-devel-${TRT_VERSION} \ + libnvinfer-dispatch10-${TRT_VERSION} \ + libnvinfer-plugin10-${TRT_VERSION} \ + libnvinfer-vc-plugin10-${TRT_VERSION} \ + libnvinfer-bin-${TRT_VERSION} \ + libnvinfer-plugin10-${TRT_VERSION} \ + libnvinfer-plugin-devel-${TRT_VERSION} \ + libnvinfer-vc-plugin-devel-${TRT_VERSION} \ + libnvinfer-lean-devel-${TRT_VERSION} \ + libnvinfer-dispatch-devel-${TRT_VERSION} \ + libnvinfer-headers-plugin-devel-${TRT_VERSION} && \ + dnf clean dbcache ; \ +else \ + echo "TRT_VERSION is none skipping Tensor RT Installation" ; \ +fi + +ADD scripts /tmp/scripts +RUN cd /tmp/scripts && /tmp/scripts/install_dotnet.sh && /tmp/scripts/install_java.sh && rm -rf /tmp/scripts + +RUN python3 -m pip uninstall -y torch +RUN python3 -m pip install torch --index-url https://download.pytorch.org/whl/cu118 + +# Build final image from base. +FROM base as final +ARG BUILD_USER=onnxruntimedev +ARG BUILD_UID=1000 +RUN adduser --uid $BUILD_UID $BUILD_USER +WORKDIR /home/$BUILD_USER +USER $BUILD_USER From 21ad0042379f9e5f226f7d54b841274268d46289 Mon Sep 17 00:00:00 2001 From: Hector Li Date: Thu, 27 Jun 2024 22:09:13 -0700 Subject: [PATCH 42/52] Add QNN UTs for QNN Pad Op with FP16 data on HTP backend (#21142) ### Description 1. Add QNN UTs for QNN Pad Op with FP16 data on HTP backend 2. Improve Pad op builder to handle invalid optional input 3. Add UT for ReduceSum for FP16 precision with 5D for issue reproduce --- .../qnn/builder/opbuilder/pad_op_builder.cc | 12 +- .../test/providers/qnn/pad_op_test.cpp | 77 ++++++- .../test/providers/qnn/reduce_op_test.cc | 209 ++++++++++-------- 3 files changed, 198 insertions(+), 100 deletions(-) diff --git a/onnxruntime/core/providers/qnn/builder/opbuilder/pad_op_builder.cc b/onnxruntime/core/providers/qnn/builder/opbuilder/pad_op_builder.cc index 3f73ef76e9def..b7455314578de 100644 --- a/onnxruntime/core/providers/qnn/builder/opbuilder/pad_op_builder.cc +++ b/onnxruntime/core/providers/qnn/builder/opbuilder/pad_op_builder.cc @@ -51,7 +51,7 @@ Status PadOpBuilder::ProcessInputs(QnnModelWrapper& qnn_model_wrapper, auto& pads_input_name = inputs[1].node_arg.Name(); ORT_RETURN_IF_NOT(qnn_model_wrapper.IsInitializerInput(pads_input_name), "Qnn doesn't support dynamic pad input"); - if (node_unit.Inputs().size() > 2) { + if (inputs.size() > 2 && inputs[2].node_arg.Exists()) { auto& constant_value_input_name = inputs[2].node_arg.Name(); ORT_RETURN_IF_NOT(qnn_model_wrapper.IsInitializerInput(constant_value_input_name), "Qnn doesn't support dynamic constant_value input"); @@ -227,13 +227,13 @@ Status PadOpBuilder::ProcessAttributesAndOutputs(QnnModelWrapper& qnn_model_wrap param_tensor_names.push_back(mode_param.GetParamTensorName()); qnn_model_wrapper.AddParamWrapper(std::move(mode_param)); - QnnParamWrapper multiples_param(node_unit.Index(), node_unit.Name(), QNN_OP_PAD_PARAM_PAD_AMOUNT, - std::move(pad_amount_dim), std::move(pad_amount)); - param_tensor_names.push_back(multiples_param.GetParamTensorName()); - qnn_model_wrapper.AddParamWrapper(std::move(multiples_param)); + QnnParamWrapper pad_amount_param(node_unit.Index(), node_unit.Name(), QNN_OP_PAD_PARAM_PAD_AMOUNT, + std::move(pad_amount_dim), std::move(pad_amount)); + param_tensor_names.push_back(pad_amount_param.GetParamTensorName()); + qnn_model_wrapper.AddParamWrapper(std::move(pad_amount_param)); // Process optional input constant_value - if (node_unit.Inputs().size() > 2) { + if (inputs.size() > 2 && inputs[2].node_arg.Exists()) { ORT_RETURN_IF_ERROR(ProcessConstantValue(qnn_model_wrapper, param_tensor_names, node_unit, inputs[2])); } // constant_value diff --git a/onnxruntime/test/providers/qnn/pad_op_test.cpp b/onnxruntime/test/providers/qnn/pad_op_test.cpp index 4ef71457d5bfe..a6b8664c6c0c9 100644 --- a/onnxruntime/test/providers/qnn/pad_op_test.cpp +++ b/onnxruntime/test/providers/qnn/pad_op_test.cpp @@ -98,18 +98,33 @@ static void RunPadOpTest(const TestInputDef& data_def, const std::vector& attrs, ExpectedEPNodeAssignment expected_ep_assignment, bool has_constant_value = true, - int opset = 18) { + int opset = 18, + bool use_htp = false, + bool enable_fp16_precision = false, + float f32_abs_err = 1e-5f) { ProviderOptions provider_options; + if (use_htp) { #if defined(_WIN32) - provider_options["backend_path"] = "QnnCpu.dll"; + provider_options["backend_path"] = "QnnHtp.dll"; #else - provider_options["backend_path"] = "libQnnCpu.so"; + provider_options["backend_path"] = "libQnnHtp.so"; #endif + } else { +#if defined(_WIN32) + provider_options["backend_path"] = "QnnCpu.dll"; +#else + provider_options["backend_path"] = "libQnnCpu.so"; +#endif + } + + if (enable_fp16_precision) { + provider_options["enable_htp_fp16_precision"] = "1"; + } RunQnnModelTest(BuildPadTestCase(data_def, pads_def, constant_value_def, attrs, has_constant_value), provider_options, opset, - expected_ep_assignment); + expected_ep_assignment, f32_abs_err); } // Runs a QDQ Pad model on the QNN HTP backend. Checks the graph node assignment, and that inference @@ -229,6 +244,60 @@ TEST_F(QnnCPUBackendTests, Pad6d) { #if defined(__aarch64__) || defined(_M_ARM64) || defined(__linux__) // // HTP tests: +TEST_F(QnnHTPBackendTests, PadNoConstantValue_fp16_test) { + bool has_constant_value_input = false; + bool use_htp = true; + bool enable_fp16_precision = true; + RunPadOpTest(TestInputDef({3, 2}, false, {1.0f, 1.2f, 2.3f, 3.4f, 4.5f, 5.6f}), + TestInputDef({4}, true, {0, 2, 0, 0}), + TestInputDef({1}, true, {0.0f}), + {utils::MakeAttribute("mode", "constant")}, + ExpectedEPNodeAssignment::All, + has_constant_value_input, + 18, // opset + use_htp, + enable_fp16_precision, + 2e-3f); +} + +TEST_F(QnnHTPBackendTests, PadReflectMode_fp16) { + bool has_constant_value_input = false; + bool use_htp = true; + bool enable_fp16_precision = true; + RunPadOpTest(TestInputDef({3, 2}, false, {1.0f, 1.2f, 2.3f, 3.4f, 4.5f, 5.6f}), + TestInputDef({4}, true, {0, 1, 0, 0}), + TestInputDef({1}, true, {0.0f}), + {utils::MakeAttribute("mode", "reflect")}, + ExpectedEPNodeAssignment::All, + has_constant_value_input, + 18, // opset + use_htp, + enable_fp16_precision, + 2e-3f); +} + +// HTP\HTP\src\hexagon\prepare\graph_prepare.cc:203:ERROR:could not create op: q::flat_from_vtcm +// HTP\HTP\src\hexagon\prepare\graph_prepare.cc:1238:ERROR:Op 0x104100000011 preparation failed with err:-1 +// Completed stage: Graph Transformations and Optimizations (13372 us) +// QnnDsp "node" generated: could not create op +// QnnDsp RouterWindows graph prepare failed 12 +// QnnDsp Failed to finalize graph (id: 1) with err 1002 +TEST_F(QnnHTPBackendTests, DISABLED_PadReflectMode_FP16_big_data) { + bool has_constant_value_input = false; + bool use_htp = true; + bool enable_fp16_precision = true; + RunPadOpTest(TestInputDef({1, 4, 512, 512}, false, GetFloatDataInRange(1.0f, 10.0f, 4 * 512 * 512)), + TestInputDef({8}, true, {0, 0, 3, 3, 0, 0, 3, 3}), + TestInputDef({1}, true, {0.0f}), + {utils::MakeAttribute("mode", "reflect")}, + ExpectedEPNodeAssignment::All, + has_constant_value_input, + 18, // opset + use_htp, + enable_fp16_precision, + 2e-3f); +} + // // QDQ Pad TEST_F(QnnHTPBackendTests, PadNoConstantValue) { diff --git a/onnxruntime/test/providers/qnn/reduce_op_test.cc b/onnxruntime/test/providers/qnn/reduce_op_test.cc index e39ba5fb40cf7..13173d9a87f55 100644 --- a/onnxruntime/test/providers/qnn/reduce_op_test.cc +++ b/onnxruntime/test/providers/qnn/reduce_op_test.cc @@ -60,30 +60,42 @@ static GetTestModelFn BuildReduceOpTestCase(const std::string& reduce_op_type, } /** - * Runs a ReduceOp model on the QNN CPU backend. Checks the graph node assignment, and that inference + * Runs a ReduceOp model on the QNN CPU/NPU backend. Checks the graph node assignment, and that inference * outputs for QNN and CPU match. * * \param op_type The ReduceOp type (e.g., ReduceSum). * \param input_def The input definition (shape, data, etc.) * \param axes The axes of reduction. + * \param keepdims Common attribute for all reduce operations. * \param opset The opset version. Some opset versions have "axes" as an attribute or input. * \param expected_ep_assignment How many nodes are expected to be assigned to QNN (All, Some, or None) - * \param keepdims Common attribute for all reduce operations. + * \param fp32_abs_err Error tolerance. + * \param enable_fp16 Enable fp32 model with FP16 precision on NPU. */ template -static void RunReduceOpCpuTest(const std::string& op_type, - const TestInputDef& input_def, - const std::vector& axes, - bool keepdims, - int opset, - ExpectedEPNodeAssignment expected_ep_assignment, - float fp32_abs_err = 1e-5f) { +static void RunReduceTest(const std::string& op_type, + const TestInputDef& input_def, + const std::vector& axes, + bool keepdims, + int opset, + ExpectedEPNodeAssignment expected_ep_assignment, + float fp32_abs_err = 1e-5f, + bool enable_fp16 = false) { ProviderOptions provider_options; + if (enable_fp16) { #if defined(_WIN32) - provider_options["backend_path"] = "QnnCpu.dll"; + provider_options["backend_path"] = "QnnHtp.dll"; #else - provider_options["backend_path"] = "libQnnCpu.so"; + provider_options["backend_path"] = "libQnnHtp.so"; #endif + provider_options["enable_htp_fp16_precision"] = "1"; + } else { +#if defined(_WIN32) + provider_options["backend_path"] = "QnnCpu.dll"; +#else + provider_options["backend_path"] = "libQnnCpu.so"; +#endif + } RunQnnModelTest(BuildReduceOpTestCase(op_type, input_def, //{2, 2}, // input shape @@ -107,12 +119,12 @@ static void RunReduceOpCpuTest(const std::string& op_type, // - The input and output data type is int32. // - Uses opset 13, which has "axes" as an input. TEST_F(QnnCPUBackendTests, ReduceSumOpset13_Int32) { - RunReduceOpCpuTest("ReduceSum", - TestInputDef({2, 2}, false, -10.0f, 10.0f), - std::vector{0, 1}, - true, // keepdims - 13, - ExpectedEPNodeAssignment::All); + RunReduceTest("ReduceSum", + TestInputDef({2, 2}, false, -10.0f, 10.0f), + std::vector{0, 1}, + true, // keepdims + 13, + ExpectedEPNodeAssignment::All); } // Test creates a graph with a ReduceSum node, and checks that all @@ -121,12 +133,12 @@ TEST_F(QnnCPUBackendTests, ReduceSumOpset13_Int32) { // - The input and output data type is int32. // - Uses opset 11, which has "axes" as an attribute. TEST_F(QnnCPUBackendTests, ReduceSumOpset11_Int32) { - RunReduceOpCpuTest("ReduceSum", - TestInputDef({2, 2}, false, -10.0f, 10.0f), - std::vector{0, 1}, - true, // keepdims - 11, - ExpectedEPNodeAssignment::All); + RunReduceTest("ReduceSum", + TestInputDef({2, 2}, false, -10.0f, 10.0f), + std::vector{0, 1}, + true, // keepdims + 11, + ExpectedEPNodeAssignment::All); } // Test creates a graph with a ReduceSum node, and checks that all @@ -135,12 +147,12 @@ TEST_F(QnnCPUBackendTests, ReduceSumOpset11_Int32) { // - The input and output data type is float. // - Uses opset 13, which has "axes" as an input. TEST_F(QnnCPUBackendTests, ReduceSumOpset13_Float) { - RunReduceOpCpuTest("ReduceSum", - TestInputDef({2, 2}, false, -10.0f, 10.0f), - std::vector{0, 1}, - true, // keepdims - 13, - ExpectedEPNodeAssignment::All); + RunReduceTest("ReduceSum", + TestInputDef({2, 2}, false, -10.0f, 10.0f), + std::vector{0, 1}, + true, // keepdims + 13, + ExpectedEPNodeAssignment::All); } // Test creates a graph with a ReduceSum node, and checks that all @@ -149,12 +161,12 @@ TEST_F(QnnCPUBackendTests, ReduceSumOpset13_Float) { // - The input and output data type is float. // - Uses opset 11, which has "axes" as an attribute. TEST_F(QnnCPUBackendTests, ReduceSumOpset11_Float) { - RunReduceOpCpuTest("ReduceSum", - TestInputDef({2, 2}, false, -10.0f, 10.0f), - std::vector{0, 1}, - true, // keepdims - 11, - ExpectedEPNodeAssignment::All); + RunReduceTest("ReduceSum", + TestInputDef({2, 2}, false, -10.0f, 10.0f), + std::vector{0, 1}, + true, // keepdims + 11, + ExpectedEPNodeAssignment::All); } // @@ -167,24 +179,24 @@ TEST_F(QnnCPUBackendTests, ReduceSumOpset11_Float) { // - The input and output data type is float. // - Uses opset 18, which has "axes" as an input. TEST_F(QnnCPUBackendTests, ReduceProdOpset18) { - RunReduceOpCpuTest("ReduceProd", - TestInputDef({2, 2}, false, {-10.0f, -8.2f, 0.0f, 10.0f}), - std::vector{0, 1}, - true, // keepdims - 18, - ExpectedEPNodeAssignment::All); + RunReduceTest("ReduceProd", + TestInputDef({2, 2}, false, {-10.0f, -8.2f, 0.0f, 10.0f}), + std::vector{0, 1}, + true, // keepdims + 18, + ExpectedEPNodeAssignment::All); } // TODO: Investigate slight inaccuracy. x64 Windows/Linux require a slightly larger error tolerance greater than 1.5e-5f. // LOG: ... the value pair (208.881729, 208.881744) at index #0 don't match, which is 1.52588e-05 from 208.882 TEST_F(QnnCPUBackendTests, ReduceProdOpset18_SlightlyInaccurate_WindowsLinuxX64) { - RunReduceOpCpuTest("ReduceProd", - TestInputDef({2, 2}, false, {3.21289f, -5.9981f, -1.72799f, 6.27263f}), - std::vector{0, 1}, - true, // keepdims - 18, - ExpectedEPNodeAssignment::All, - 2e-5f); // x64 Linux & Windows require larger tolerance. + RunReduceTest("ReduceProd", + TestInputDef({2, 2}, false, {3.21289f, -5.9981f, -1.72799f, 6.27263f}), + std::vector{0, 1}, + true, // keepdims + 18, + ExpectedEPNodeAssignment::All, + 2e-5f); // x64 Linux & Windows require larger tolerance. } // Test creates a graph with a ReduceProd node, and checks that all @@ -193,12 +205,12 @@ TEST_F(QnnCPUBackendTests, ReduceProdOpset18_SlightlyInaccurate_WindowsLinuxX64) // - The input and output data type is float. // - Uses opset 13, which has "axes" as an attribute. TEST_F(QnnCPUBackendTests, ReduceProdOpset13) { - RunReduceOpCpuTest("ReduceProd", - TestInputDef({2, 2}, false, {-10.0f, -8.2f, 0.0f, 10.0f}), - std::vector{0, 1}, - true, // keepdims - 13, - ExpectedEPNodeAssignment::All); + RunReduceTest("ReduceProd", + TestInputDef({2, 2}, false, {-10.0f, -8.2f, 0.0f, 10.0f}), + std::vector{0, 1}, + true, // keepdims + 13, + ExpectedEPNodeAssignment::All); } // @@ -211,12 +223,12 @@ TEST_F(QnnCPUBackendTests, ReduceProdOpset13) { // - The input and output data type is float. // - Uses opset 18, which has "axes" as an input. TEST_F(QnnCPUBackendTests, ReduceMaxOpset18) { - RunReduceOpCpuTest("ReduceMax", - TestInputDef({2, 2}, false, -10.0f, 10.0f), - std::vector{0, 1}, - true, // keepdims - 18, - ExpectedEPNodeAssignment::All); + RunReduceTest("ReduceMax", + TestInputDef({2, 2}, false, -10.0f, 10.0f), + std::vector{0, 1}, + true, // keepdims + 18, + ExpectedEPNodeAssignment::All); } // Test creates a graph with a ReduceMax node, and checks that all @@ -225,12 +237,12 @@ TEST_F(QnnCPUBackendTests, ReduceMaxOpset18) { // - The input and output data type is float. // - Uses opset 13, which has "axes" as an attribute. TEST_F(QnnCPUBackendTests, ReduceMaxOpset13) { - RunReduceOpCpuTest("ReduceMax", - TestInputDef({2, 2}, false, -10.0f, 10.0f), - std::vector{0, 1}, - true, // keepdims - 13, - ExpectedEPNodeAssignment::All); + RunReduceTest("ReduceMax", + TestInputDef({2, 2}, false, -10.0f, 10.0f), + std::vector{0, 1}, + true, // keepdims + 13, + ExpectedEPNodeAssignment::All); } // @@ -243,12 +255,12 @@ TEST_F(QnnCPUBackendTests, ReduceMaxOpset13) { // - The input and output data type is float. // - Uses opset 18, which has "axes" as an input. TEST_F(QnnCPUBackendTests, ReduceMinOpset18) { - RunReduceOpCpuTest("ReduceMin", - TestInputDef({2, 2}, false, -10.0f, 10.0f), - std::vector{0, 1}, - true, // keepdims - 18, - ExpectedEPNodeAssignment::All); + RunReduceTest("ReduceMin", + TestInputDef({2, 2}, false, -10.0f, 10.0f), + std::vector{0, 1}, + true, // keepdims + 18, + ExpectedEPNodeAssignment::All); } // Test creates a graph with a ReduceMin node, and checks that all @@ -257,12 +269,12 @@ TEST_F(QnnCPUBackendTests, ReduceMinOpset18) { // - The input and output data type is float. // - Uses opset 13, which has "axes" as an attribute. TEST_F(QnnCPUBackendTests, ReduceMinOpset13) { - RunReduceOpCpuTest("ReduceMin", - TestInputDef({2, 2}, false, -10.0f, 10.0f), - std::vector{0, 1}, - true, // keepdims - 13, - ExpectedEPNodeAssignment::All); + RunReduceTest("ReduceMin", + TestInputDef({2, 2}, false, -10.0f, 10.0f), + std::vector{0, 1}, + true, // keepdims + 13, + ExpectedEPNodeAssignment::All); } // @@ -275,12 +287,12 @@ TEST_F(QnnCPUBackendTests, ReduceMinOpset13) { // - The input and output data type is float. // - Uses opset 18, which has "axes" as an input. TEST_F(QnnCPUBackendTests, ReduceMeanOpset18) { - RunReduceOpCpuTest("ReduceMean", - TestInputDef({2, 2}, false, -10.0f, 10.0f), - std::vector{0, 1}, - true, // keepdims - 18, - ExpectedEPNodeAssignment::All); + RunReduceTest("ReduceMean", + TestInputDef({2, 2}, false, -10.0f, 10.0f), + std::vector{0, 1}, + true, // keepdims + 18, + ExpectedEPNodeAssignment::All); } // Test creates a graph with a ReduceMean node, and checks that all @@ -289,16 +301,33 @@ TEST_F(QnnCPUBackendTests, ReduceMeanOpset18) { // - The input and output data type is float. // - Uses opset 13, which has "axes" as an attribute. TEST_F(QnnCPUBackendTests, ReduceMeanOpset13) { - RunReduceOpCpuTest("ReduceMean", - TestInputDef({2, 2}, false, -10.0f, 10.0f), - std::vector{0, 1}, - true, // keepdims - 13, - ExpectedEPNodeAssignment::All); + RunReduceTest("ReduceMean", + TestInputDef({2, 2}, false, -10.0f, 10.0f), + std::vector{0, 1}, + true, // keepdims + 13, + ExpectedEPNodeAssignment::All); } #if defined(__aarch64__) || defined(_M_ARM64) || defined(__linux__) +// Test creates a graph with a ReduceSum node, and checks that all nodes are supported by the QNN EP +// HTP backend with FP16 precision, and that the inference results match the CPU EP results. +// +// Failed QNN Opvalidation because of 5D input. It runs OK if bypass the op validation +TEST_F(QnnHTPBackendTests, DISABLED_ReduceSumOpset11_5D_FP16) { + float fp32_abs_err = 3e-2f; + bool enable_fp16 = true; + RunReduceTest("ReduceSum", + TestInputDef({1, 12, 249, 2, 4}, false, -10.0f, 10.0f), + std::vector{-1}, + false, // keepdims + 13, + ExpectedEPNodeAssignment::All, + fp32_abs_err, + enable_fp16); +} + // Creates the following graph if axes is an input (newer opsets): // _______________________ // input (f32) -> Q -> DQ -> | | -> Q -> DQ -> output (f32) From 6baaaf516538f9059da3558b2cd22128a9e42c07 Mon Sep 17 00:00:00 2001 From: Preetha Veeramalai Date: Fri, 28 Jun 2024 08:31:02 -0700 Subject: [PATCH 43/52] OVEP options to disable CPU fallback at compile time (#21166) ### Description Provide user level options to control the fallback on CPU for models not supported on Intel's NPU hardware. ### Motivation and Context - Current workflow of OVEP allows safe fallback from OV NPU to OV CPU on compilation failures. Also supports MLAS CPU fallback in presence of unsupported custom ops. - The PR provides a build-time option to disable fallback from OV NPU to OV CPU. - The session Option "kOrtSessionOptionsDisableCPUEPFallback" disables OV CPU and MLAS CPU fallback. - Also has bug fix for proto creation. --------- Co-authored-by: jatinwadhwa921 Co-authored-by: ankitm3k --- cmake/CMakeLists.txt | 4 + onnxruntime/core/framework/config_options.cc | 4 +- .../core/graph/graph_proto_serializer.cc | 16 +- .../providers/openvino/backend_manager.cc | 18 +- .../openvino/backends/basic_backend.cc | 7 +- .../core/providers/openvino/contexts.h | 1 + .../openvino/openvino_execution_provider.cc | 1 + .../openvino/openvino_execution_provider.h | 6 +- .../openvino/openvino_provider_factory.cc | 24 +- .../openvino_provider_factory_creator.h | 4 +- .../core/session/provider_bridge_ort.cc | 16 +- .../core/session/provider_registration.cc | 3 +- .../python/onnxruntime_pybind_schema.cc | 3 +- .../python/onnxruntime_pybind_state.cc | 2 +- .../python/onnxruntime_pybind_state_common.h | 4 + onnxruntime/test/perftest/ort_test_session.cc | 294 +++++++++--------- onnxruntime/test/util/default_providers.cc | 3 +- tools/ci_build/build.py | 2 + 18 files changed, 247 insertions(+), 165 deletions(-) diff --git a/cmake/CMakeLists.txt b/cmake/CMakeLists.txt index 6ba0db789965c..4483e4d5cb17f 100644 --- a/cmake/CMakeLists.txt +++ b/cmake/CMakeLists.txt @@ -1341,6 +1341,10 @@ if (onnxruntime_USE_OPENVINO) add_definitions(-DUSE_OPENVINO=1) + if(onnxruntime_NPU_NO_FALLBACK) + add_definitions(-DOPENVINO_DISABLE_NPU_FALLBACK=1) + endif() + if (onnxruntime_USE_OPENVINO_GPU) add_definitions(-DOPENVINO_CONFIG_GPU=1) endif() diff --git a/onnxruntime/core/framework/config_options.cc b/onnxruntime/core/framework/config_options.cc index 1a4acb6dabf71..9fe5beafd6e7e 100644 --- a/onnxruntime/core/framework/config_options.cc +++ b/onnxruntime/core/framework/config_options.cc @@ -30,11 +30,11 @@ std::string ConfigOptions::GetConfigOrDefault(const std::string& config_key, } Status ConfigOptions::AddConfigEntry(const char* config_key, const char* config_value) noexcept { - std::string key(config_key); + std::string key = config_key; if (key.empty() || key.length() > 128) return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Config key is empty or longer than maximum length 128"); - std::string val(config_value); + std::string val = config_value; if (val.length() > onnxruntime::kMaxStrLen) return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Config value is longer than maximum length: ", diff --git a/onnxruntime/core/graph/graph_proto_serializer.cc b/onnxruntime/core/graph/graph_proto_serializer.cc index aefad28eb37e8..eb0fb22346f37 100644 --- a/onnxruntime/core/graph/graph_proto_serializer.cc +++ b/onnxruntime/core/graph/graph_proto_serializer.cc @@ -21,7 +21,21 @@ void GraphViewerToProto(const GraphViewer& graph_view, *(graph_proto.mutable_output()->Add()) = output_arg->ToProto(); } - for (const auto* value_info : graph_view.GetValueInfo()) { + std::unordered_set value_info_ = graph_view.GetValueInfo(); + + // Reserve memory for the vector to avoid reallocations + std::vector value_info_sorted; + value_info_sorted.reserve(value_info_.size()); + + value_info_sorted.assign(value_info_.begin(), value_info_.end()); + auto sort_predicate = [](const NodeArg* v1, const NodeArg* v2) { + return v1->Name() < v2->Name(); + }; + + // This ensures consistent ordering of value_info entries in the output graph + std::sort(value_info_sorted.begin(), value_info_sorted.end(), sort_predicate); + + for (const auto* value_info : value_info_sorted) { *(graph_proto.mutable_value_info()->Add()) = value_info->ToProto(); } diff --git a/onnxruntime/core/providers/openvino/backend_manager.cc b/onnxruntime/core/providers/openvino/backend_manager.cc index d0ef447a46d21..1c027e39fa5f5 100644 --- a/onnxruntime/core/providers/openvino/backend_manager.cc +++ b/onnxruntime/core/providers/openvino/backend_manager.cc @@ -105,7 +105,11 @@ BackendManager::BackendManager(const GlobalContext& global_context, subgraph_context_, ep_ctx_handle_); } catch (const OnnxRuntimeException& ex) { - if (device_type.find("NPU") != std::string::npos) { +#if defined(OPENVINO_DISABLE_NPU_FALLBACK) + ORT_THROW(ex.what()); +#else + if (device_type.find("NPU") != std::string::npos && + !GetGlobalContext().disable_cpu_fallback) { LOGS_DEFAULT(WARNING) << ex.what(); LOGS_DEFAULT(WARNING) << "Model compilation failed at OV NPU." << "Falling back to OV CPU for execution"; @@ -122,6 +126,7 @@ BackendManager::BackendManager(const GlobalContext& global_context, } else { ORT_THROW(ex.what()); } +#endif } } } @@ -419,7 +424,13 @@ void BackendManager::Compute(OrtKernelContext* context) { subgraph_context_, ep_ctx_handle_); } catch (const OnnxRuntimeException& ex) { - if (GetGlobalContext().device_type.find("NPU") != std::string::npos) { + // Build option disables fallback to CPU on compilation failures with NPU. +#if defined(OPENVINO_DISABLE_NPU_FALLBACK) + LOGS_DEFAULT(WARNING) << "Model compilation failed at OV NPU."; + ORT_THROW(ex.what()); +#else + if (GetGlobalContext().device_type.find("NPU") != std::string::npos && + !GetGlobalContext().disable_cpu_fallback) { LOGS_DEFAULT(WARNING) << ex.what(); LOGS_DEFAULT(WARNING) << "Model compilation failed at OV NPU." << "Falling back to OV CPU for execution"; @@ -434,7 +445,10 @@ void BackendManager::Compute(OrtKernelContext* context) { } catch (std::string const& msg) { ORT_THROW(msg); } + } else { + ORT_THROW(ex.what()); } +#endif } backend_map_.insert({key, dynamic_backend}); } else { diff --git a/onnxruntime/core/providers/openvino/backends/basic_backend.cc b/onnxruntime/core/providers/openvino/backends/basic_backend.cc index 9da6e5945ab83..f8046bcb3a06f 100644 --- a/onnxruntime/core/providers/openvino/backends/basic_backend.cc +++ b/onnxruntime/core/providers/openvino/backends/basic_backend.cc @@ -545,6 +545,11 @@ void BasicBackend::Infer(OrtKernelContext* ctx) { std::cout << "Inference successful" << std::endl; } + // Create a duplicate infer_request_ shared ptr on the stack in the current local scope, + // as the infer_request gets freed in the next stage the reference count for the infer_request decrements & + // thus we dont have any dangling ptr leading to seg faults in the debug mode subsequent execution call + OVInferRequestPtr infer_request_ = infer_request; + // Once the inference is completed, the infer_request becomes free and is placed back into pool of infer_requests_ inferRequestsQueue_->putIdleRequest(std::move(infer_request)); #ifndef NDEBUG @@ -552,7 +557,7 @@ void BasicBackend::Infer(OrtKernelContext* ctx) { if (openvino_ep::backend_utils::IsDebugEnabled()) { inferRequestsQueue_->printstatus(); // Printing the elements of infer_requests_ vector pool only in debug mode std::string& hw_target = global_context_.device_type; - printPerformanceCounts(infer_request, std::cout, hw_target); + printPerformanceCounts(std::move(infer_request_), std::cout, hw_target); } #endif #endif diff --git a/onnxruntime/core/providers/openvino/contexts.h b/onnxruntime/core/providers/openvino/contexts.h index 6e11cbf4a699f..598e985676f8d 100644 --- a/onnxruntime/core/providers/openvino/contexts.h +++ b/onnxruntime/core/providers/openvino/contexts.h @@ -21,6 +21,7 @@ struct GlobalContext { bool ep_context_embed_mode = true; bool export_ep_ctx_blob = false; bool enable_qdq_optimizer = false; + bool disable_cpu_fallback = false; size_t num_of_threads; std::string device_type; std::string precision_str; diff --git a/onnxruntime/core/providers/openvino/openvino_execution_provider.cc b/onnxruntime/core/providers/openvino/openvino_execution_provider.cc index 6f7e1fb607864..040c56926a803 100644 --- a/onnxruntime/core/providers/openvino/openvino_execution_provider.cc +++ b/onnxruntime/core/providers/openvino/openvino_execution_provider.cc @@ -33,6 +33,7 @@ OpenVINOExecutionProvider::OpenVINOExecutionProvider(const OpenVINOExecutionProv global_context_->OpenVINO_Version = {OPENVINO_VERSION_MAJOR, OPENVINO_VERSION_MINOR}; global_context_->export_ep_ctx_blob = info.export_ep_ctx_blob_; global_context_->enable_qdq_optimizer = info.enable_qdq_optimizer_; + global_context_->disable_cpu_fallback = info.disable_cpu_fallback_; // to check if target device is available // using ie_core capability GetAvailableDevices to fetch list of devices plugged in diff --git a/onnxruntime/core/providers/openvino/openvino_execution_provider.h b/onnxruntime/core/providers/openvino/openvino_execution_provider.h index d950255c7727b..050fb91c51771 100644 --- a/onnxruntime/core/providers/openvino/openvino_execution_provider.h +++ b/onnxruntime/core/providers/openvino/openvino_execution_provider.h @@ -74,6 +74,7 @@ struct OpenVINOExecutionProviderInfo { bool disable_dynamic_shapes_{false}; bool export_ep_ctx_blob_{false}; bool enable_qdq_optimizer_{false}; + bool disable_cpu_fallback_{false}; OpenVINOExecutionProviderInfo() = delete; @@ -81,7 +82,7 @@ struct OpenVINOExecutionProviderInfo { size_t num_of_threads, std::string cache_dir, std::string model_priority, int num_streams, void* context, bool enable_opencl_throttling, bool disable_dynamic_shapes, bool export_ep_ctx_blob, - bool enable_qdq_optimizer) + bool enable_qdq_optimizer, bool disable_cpu_fallback) : precision_(precision), enable_npu_fast_compile_(enable_npu_fast_compile), num_of_threads_(num_of_threads), @@ -92,7 +93,8 @@ struct OpenVINOExecutionProviderInfo { enable_opencl_throttling_(enable_opencl_throttling), disable_dynamic_shapes_(disable_dynamic_shapes), export_ep_ctx_blob_(export_ep_ctx_blob), - enable_qdq_optimizer_(enable_qdq_optimizer) { + enable_qdq_optimizer_(enable_qdq_optimizer), + disable_cpu_fallback_(disable_cpu_fallback) { std::set ov_supported_device_types = {"CPU", "GPU", "GPU.0", "GPU.1", "NPU"}; if (dev_type == "") { diff --git a/onnxruntime/core/providers/openvino/openvino_provider_factory.cc b/onnxruntime/core/providers/openvino/openvino_provider_factory.cc index a45c1fd236af1..45bba431741c5 100644 --- a/onnxruntime/core/providers/openvino/openvino_provider_factory.cc +++ b/onnxruntime/core/providers/openvino/openvino_provider_factory.cc @@ -13,7 +13,8 @@ struct OpenVINOProviderFactory : IExecutionProviderFactory { const char* cache_dir, const char* model_priority, int num_streams, void* context, bool enable_opencl_throttling, bool disable_dynamic_shapes, - bool export_ep_ctx_blob, bool enable_qdq_optimizer) + bool export_ep_ctx_blob, bool enable_qdq_optimizer, + bool disable_cpu_fallback) : precision_(precision), enable_npu_fast_compile_(enable_npu_fast_compile), num_of_threads_(num_of_threads), @@ -23,7 +24,8 @@ struct OpenVINOProviderFactory : IExecutionProviderFactory { enable_opencl_throttling_(enable_opencl_throttling), disable_dynamic_shapes_(disable_dynamic_shapes), export_ep_ctx_blob_(export_ep_ctx_blob), - enable_qdq_optimizer_(enable_qdq_optimizer) { + enable_qdq_optimizer_(enable_qdq_optimizer), + disable_cpu_fallback_(disable_cpu_fallback) { device_type_ = (device_type == nullptr) ? "" : device_type; cache_dir_ = (cache_dir == nullptr) ? "" : cache_dir; } @@ -45,12 +47,14 @@ struct OpenVINOProviderFactory : IExecutionProviderFactory { bool disable_dynamic_shapes_; bool export_ep_ctx_blob_; bool enable_qdq_optimizer_; + bool disable_cpu_fallback_; }; std::unique_ptr OpenVINOProviderFactory::CreateProvider() { OpenVINOExecutionProviderInfo info(device_type_, precision_, enable_npu_fast_compile_, num_of_threads_, cache_dir_, model_priority_, num_streams_, context_, enable_opencl_throttling_, - disable_dynamic_shapes_, export_ep_ctx_blob_, enable_qdq_optimizer_); + disable_dynamic_shapes_, export_ep_ctx_blob_, enable_qdq_optimizer_, + disable_cpu_fallback_); return std::make_unique(info); } @@ -99,6 +103,8 @@ struct OpenVINO_Provider : Provider { bool enable_qdq_optimizer = false; + bool disable_cpu_fallback = false; + if (provider_options_map.find("device_type") != provider_options_map.end()) { device_type = provider_options_map.at("device_type").c_str(); @@ -256,6 +262,15 @@ struct OpenVINO_Provider : Provider { export_ep_ctx_blob = false; bool_flag = ""; } + + if (provider_options_map.find("disable_cpu_fallback") != provider_options_map.end()) { + bool_flag = provider_options_map.at("disable_cpu_fallback"); + if (bool_flag == "true" || bool_flag == "True") + disable_cpu_fallback = true; + else if (bool_flag == "false" || bool_flag == "False") + disable_cpu_fallback = false; + bool_flag = ""; + } return std::make_shared(const_cast(device_type.c_str()), const_cast(precision.c_str()), enable_npu_fast_compile, @@ -267,7 +282,8 @@ struct OpenVINO_Provider : Provider { enable_opencl_throttling, disable_dynamic_shapes, export_ep_ctx_blob, - enable_qdq_optimizer); + enable_qdq_optimizer, + disable_cpu_fallback); } void Initialize() override { diff --git a/onnxruntime/core/providers/openvino/openvino_provider_factory_creator.h b/onnxruntime/core/providers/openvino/openvino_provider_factory_creator.h index 4df653b022a66..bff70a90b6a70 100644 --- a/onnxruntime/core/providers/openvino/openvino_provider_factory_creator.h +++ b/onnxruntime/core/providers/openvino/openvino_provider_factory_creator.h @@ -11,9 +11,11 @@ struct OrtOpenVINOProviderOptions; namespace onnxruntime { +struct SessionOptions; // defined in provider_bridge_ort.cc struct OpenVINOProviderFactoryCreator { - static std::shared_ptr Create(const ProviderOptions* provider_options_map); + static std::shared_ptr Create(ProviderOptions* provider_options_map, + const SessionOptions* session_options); static std::shared_ptr Create(const OrtOpenVINOProviderOptions* provider_options); }; } // namespace onnxruntime diff --git a/onnxruntime/core/session/provider_bridge_ort.cc b/onnxruntime/core/session/provider_bridge_ort.cc index 7f7ed5e436afe..d4c6e3d506f18 100644 --- a/onnxruntime/core/session/provider_bridge_ort.cc +++ b/onnxruntime/core/session/provider_bridge_ort.cc @@ -35,6 +35,7 @@ #include "core/framework/model_metadef_id_generator.h" #include "core/optimizer/qdq_transformer/selectors_actions/qdq_selectors.h" #include "core/optimizer/qdq_transformer/selectors_actions/shared/utils.h" +#include "core/session/onnxruntime_session_options_config_keys.h" #include "core/session/onnxruntime_c_api.h" #include "core/common/string_helper.h" @@ -1800,7 +1801,18 @@ std::shared_ptr OpenVINOProviderFactoryCreator::Creat return s_library_openvino.Get().CreateExecutionProviderFactory(&ov_options_converted_map); } -std::shared_ptr OpenVINOProviderFactoryCreator::Create(const ProviderOptions* provider_options_map) { +void ORTSessionOptionsToOrtOpenVINOProviderOptions(ProviderOptions& ov_options, + const SessionOptions* session_options) { + bool disable_cpu_fallback = session_options->config_options.GetConfigOrDefault( + kOrtSessionOptionsDisableCPUEPFallback, "0") == "1"; + if (disable_cpu_fallback) + ov_options["disable_cpu_fallback"] = "true"; +} + +std::shared_ptr OpenVINOProviderFactoryCreator::Create(ProviderOptions* provider_options_map, + const SessionOptions* session_options) { + if (session_options) + onnxruntime::ORTSessionOptionsToOrtOpenVINOProviderOptions(*provider_options_map, session_options); return s_library_openvino.Get().CreateExecutionProviderFactory(provider_options_map); } @@ -2075,7 +2087,7 @@ ORT_API_STATUS_IMPL(OrtApis::SessionOptionsAppendExecutionProvider_OpenVINO_V2, provider_options[provider_options_keys[i]] = provider_options_values[i]; } - auto factory = onnxruntime::OpenVINOProviderFactoryCreator::Create(&provider_options); + auto factory = onnxruntime::OpenVINOProviderFactoryCreator::Create(&provider_options, &(options->value)); if (!factory) { return OrtApis::CreateStatus(ORT_FAIL, "SessionOptionsAppendExecutionProvider_OpenVINO_V2: Failed to load shared library"); } diff --git a/onnxruntime/core/session/provider_registration.cc b/onnxruntime/core/session/provider_registration.cc index 05408db9884cd..688ee76c591f6 100644 --- a/onnxruntime/core/session/provider_registration.cc +++ b/onnxruntime/core/session/provider_registration.cc @@ -108,11 +108,10 @@ ORT_API_STATUS_IMPL(OrtApis::SessionOptionsAppendExecutionProvider, #endif } else if (strcmp(provider_name, "OpenVINO") == 0) { #if defined(USE_OPENVINO) - options->provider_factories.push_back(OpenVINOProviderFactoryCreator::Create(&provider_options)); + options->provider_factories.push_back(OpenVINOProviderFactoryCreator::Create(&provider_options, &(options->value))); #else status = create_not_supported_status(); #endif - } else if (strcmp(provider_name, "SNPE") == 0) { #if defined(USE_SNPE) options->provider_factories.push_back(SNPEProviderFactoryCreator::Create(provider_options)); diff --git a/onnxruntime/python/onnxruntime_pybind_schema.cc b/onnxruntime/python/onnxruntime_pybind_schema.cc index 4da25eac32040..218b59688b01c 100644 --- a/onnxruntime/python/onnxruntime_pybind_schema.cc +++ b/onnxruntime/python/onnxruntime_pybind_schema.cc @@ -40,7 +40,8 @@ void addGlobalSchemaFunctions(pybind11::module& m) { #ifdef USE_OPENVINO []() { ProviderOptions provider_options_map; - return onnxruntime::OpenVINOProviderFactoryCreator::Create(&provider_options_map); + SessionOptions session_options; + return onnxruntime::OpenVINOProviderFactoryCreator::Create(&provider_options_map, &session_options); }(), #endif #ifdef USE_TENSORRT diff --git a/onnxruntime/python/onnxruntime_pybind_state.cc b/onnxruntime/python/onnxruntime_pybind_state.cc index fa4c906dd054c..e539614fd6d1d 100644 --- a/onnxruntime/python/onnxruntime_pybind_state.cc +++ b/onnxruntime/python/onnxruntime_pybind_state.cc @@ -1084,7 +1084,7 @@ std::unique_ptr CreateExecutionProviderInstance( } } if (std::shared_ptr openvino_provider_factory = onnxruntime::OpenVINOProviderFactoryCreator::Create( - &OV_provider_options_map)) { + &OV_provider_options_map, &session_options)) { auto p = openvino_provider_factory->CreateProvider(); // Reset global variables config to avoid it being accidentally passed on to the next session openvino_device_type.clear(); diff --git a/onnxruntime/python/onnxruntime_pybind_state_common.h b/onnxruntime/python/onnxruntime_pybind_state_common.h index dc9394a83a4ea..4d6e411defae3 100644 --- a/onnxruntime/python/onnxruntime_pybind_state_common.h +++ b/onnxruntime/python/onnxruntime_pybind_state_common.h @@ -65,7 +65,11 @@ struct OrtStatus { #elif OPENVINO_CONFIG_HETERO #define BACKEND_OPENVINO "-OPENVINO_HETERO" + +#elif OPENVINO_DISABLE_NPU_FALLBACK +#define BACKEND_OPENVINO "-OPENVINO_DISABLE_NPU_FALLBACK" #endif + #else #define BACKEND_OPENVINO "" #endif diff --git a/onnxruntime/test/perftest/ort_test_session.cc b/onnxruntime/test/perftest/ort_test_session.cc index 81053bf400a63..1485a4456d326 100644 --- a/onnxruntime/test/perftest/ort_test_session.cc +++ b/onnxruntime/test/perftest/ort_test_session.cc @@ -47,6 +47,7 @@ OnnxRuntimeTestSession::OnnxRuntimeTestSession(Ort::Env& env, std::random_device const TestModelInfo& m) : rand_engine_(rd()), input_names_(m.GetInputCount()), input_names_str_(m.GetInputCount()), input_length_(m.GetInputCount()) { Ort::SessionOptions session_options; + provider_name_ = performance_test_config.machine_config.provider_type_name; if (provider_name_ == onnxruntime::kDnnlExecutionProvider) { #ifdef USE_DNNL @@ -221,150 +222,6 @@ OnnxRuntimeTestSession::OnnxRuntimeTestSession(Ort::Env& env, std::random_device session_options.AppendExecutionProvider_CUDA(cuda_options); #else ORT_THROW("TensorRT is not supported in this build\n"); -#endif - } else if (provider_name_ == onnxruntime::kOpenVINOExecutionProvider) { -#ifdef USE_OPENVINO -#ifdef _MSC_VER - std::string ov_string = ToUTF8String(performance_test_config.run_config.ep_runtime_config_string); -#else - std::string ov_string = performance_test_config.run_config.ep_runtime_config_string; -#endif - std::unordered_map ov_options; - std::istringstream ss(ov_string); - std::string token; - while (ss >> token) { - if (token == "") { - continue; - } - auto pos = token.find("|"); - if (pos == std::string::npos || pos == 0 || pos == token.length()) { - ORT_THROW("[ERROR] [OpenVINO] Use a '|' to separate the key and value for the run-time option you are trying to use.\n"); - } - - auto key = token.substr(0, pos); - auto value = token.substr(pos + 1); - - if (key == "device_type") { - std::set ov_supported_device_types = {"CPU", "GPU", - "GPU.0", "GPU.1", "NPU"}; - std::set deprecated_device_types = {"CPU_FP32", "GPU_FP32", - "GPU.0_FP32", "GPU.1_FP32", "GPU_FP16", - "GPU.0_FP16", "GPU.1_FP16"}; - if (ov_supported_device_types.find(value) != ov_supported_device_types.end()) { - ov_options[key] = value; - } else if (deprecated_device_types.find(value) != deprecated_device_types.end()) { - ov_options[key] = value; - } else if (value.find("HETERO:") == 0) { - ov_options[key] = value; - } else if (value.find("MULTI:") == 0) { - ov_options[key] = value; - } else if (value.find("AUTO:") == 0) { - ov_options[key] = value; - } else { - ORT_THROW( - "[ERROR] [OpenVINO] You have selcted wrong configuration value for the key 'device_type'. " - "Select from 'CPU', 'GPU', 'GPU.0', 'GPU.1', 'NPU' or from" - " HETERO/MULTI/AUTO options available. \n"); - } - } else if (key == "device_id") { - if (value == "CPU" || value == "GPU" || value == "NPU") { - ov_options[key] = value; - } else { - ORT_THROW("[ERROR] [OpenVINO] Unsupported device_id is selected. Select from available options."); - } - } else if (key == "precision") { - auto device_type = ov_options["device_type"]; - if (device_type.find("GPU") != std::string::npos) { - if (value == "") { - ov_options[key] = "FP16"; - continue; - } else if (value == "ACCURACY" || value == "FP16" || value == "FP32") { - ov_options[key] = value; - continue; - } else { - ORT_THROW( - "[ERROR] [OpenVINO] Unsupported inference precision is selected. " - "GPU only supported FP32 / FP16. \n"); - } - } else if (device_type.find("NPU") != std::string::npos) { - if (value == "" || value == "ACCURACY" || value == "FP16") { - ov_options[key] = "FP16"; - continue; - } else { - ORT_THROW("[ERROR] [OpenVINO] Unsupported inference precision is selected. NPU only supported FP16. \n"); - } - } else if (device_type.find("CPU") != std::string::npos) { - if (value == "" || value == "ACCURACY" || value == "FP32") { - ov_options[key] = "FP32"; - continue; - } else { - ORT_THROW("[ERROR] [OpenVINO] Unsupported inference precision is selected. CPU only supports FP32 . \n"); - } - } - } else if (key == "enable_npu_fast_compile") { - if (value == "true" || value == "True" || - value == "false" || value == "False") { - ov_options[key] = value; - } else { - ORT_THROW("[ERROR] [OpenVINO] The value for the key 'enable_npu_fast_compile' should be a boolean i.e. true or false. Default value is false.\n"); - } - } else if (key == "enable_opencl_throttling") { - if (value == "true" || value == "True" || - value == "false" || value == "False") { - ov_options[key] = value; - } else { - ORT_THROW("[ERROR] [OpenVINO] The value for the key 'enable_opencl_throttling' should be a boolean i.e. true or false. Default value is false.\n"); - } - } else if (key == "enable_qdq_optimizer") { - if (value == "true" || value == "True" || - value == "false" || value == "False") { - ov_options[key] = value; - } else { - ORT_THROW("[ERROR] [OpenVINO] The value for the key 'enable_qdq_optimizer' should be a boolean i.e. true or false. Default value is false.\n"); - } - } else if (key == "disable_dynamic_shapes") { - if (value == "true" || value == "True" || - value == "false" || value == "False") { - ov_options[key] = value; - } else { - ORT_THROW( - "[ERROR] [OpenVINO] The value for the key 'enable_dynamic_shapes' " - "should be a boolean i.e. true or false. Default value is false.\n"); - } - } else if (key == "num_of_threads") { - if (std::stoi(value) <= 0) { - ORT_THROW("[ERROR] [OpenVINO] The value for the key 'num_of_threads' should be greater than 0\n"); - } else { - ov_options[key] = value; - } - } else if (key == "model_priority") { - ov_options[key] = value; - } else if (key == "cache_dir") { - ov_options[key] = value; - } else if (key == "context") { - ov_options[key] = value; - } else if (key == "num_streams") { - if (std::stoi(value) <= 0 && std::stoi(value) > 8) { - ORT_THROW("[ERROR] [OpenVINO] The value for the key 'num_streams' should be in the range of 1-8 \n"); - } else { - ov_options[key] = value; - } - } else if (key == "export_ep_ctx_blob") { - if (value == "true" || value == "True" || - value == "false" || value == "False") { - ov_options[key] = value; - } else { - ORT_THROW( - "[ERROR] [OpenVINO] The value for the key 'export_ep_ctx_blob' " - "should be a boolean i.e. true or false. Default value is false.\n"); - } - } else { - ORT_THROW("[ERROR] [OpenVINO] wrong key type entered. Choose from the following runtime key options that are available for OpenVINO. ['device_type', 'device_id', 'enable_npu_fast_compile', 'num_of_threads', 'cache_dir', 'num_streams', 'enable_opencl_throttling', 'disable_dynamic_shapes'] \n"); - } - } - session_options.AppendExecutionProvider_OpenVINO_V2(ov_options); -#else - ORT_THROW("OpenVINO is not supported in this build\n"); #endif } else if (provider_name_ == onnxruntime::kQnnExecutionProvider) { #ifdef USE_QNN @@ -716,7 +573,9 @@ select from 'TF8', 'TF16', 'UINT8', 'FLOAT', 'ITENSOR'. \n)"); #else ORT_THROW("VitisAI is not supported in this build\n"); #endif - } else if (!provider_name_.empty() && provider_name_ != onnxruntime::kCpuExecutionProvider) { + } else if (!provider_name_.empty() && + provider_name_ != onnxruntime::kCpuExecutionProvider && + provider_name_ != onnxruntime::kOpenVINOExecutionProvider) { ORT_THROW("This backend is not included in perf test runner.\n"); } @@ -805,6 +664,151 @@ select from 'TF8', 'TF16', 'UINT8', 'FLOAT', 'ITENSOR'. \n)"); } } } + if (provider_name_ == onnxruntime::kOpenVINOExecutionProvider) { +#ifdef USE_OPENVINO +#ifdef _MSC_VER + std::string ov_string = ToUTF8String(performance_test_config.run_config.ep_runtime_config_string); +#else + std::string ov_string = performance_test_config.run_config.ep_runtime_config_string; +#endif + std::unordered_map ov_options; + std::istringstream ss(ov_string); + std::string token; + while (ss >> token) { + if (token == "") { + continue; + } + auto pos = token.find("|"); + if (pos == std::string::npos || pos == 0 || pos == token.length()) { + ORT_THROW("[ERROR] [OpenVINO] Use a '|' to separate the key and value for the run-time option you are trying to use.\n"); + } + + auto key = token.substr(0, pos); + auto value = token.substr(pos + 1); + + if (key == "device_type") { + std::set ov_supported_device_types = {"CPU", "GPU", + "GPU.0", "GPU.1", "NPU"}; + std::set deprecated_device_types = {"CPU_FP32", "GPU_FP32", + "GPU.0_FP32", "GPU.1_FP32", "GPU_FP16", + "GPU.0_FP16", "GPU.1_FP16"}; + if (ov_supported_device_types.find(value) != ov_supported_device_types.end()) { + ov_options[key] = value; + } else if (deprecated_device_types.find(value) != deprecated_device_types.end()) { + ov_options[key] = value; + } else if (value.find("HETERO:") == 0) { + ov_options[key] = value; + } else if (value.find("MULTI:") == 0) { + ov_options[key] = value; + } else if (value.find("AUTO:") == 0) { + ov_options[key] = value; + } else { + ORT_THROW( + "[ERROR] [OpenVINO] You have selcted wrong configuration value for the key 'device_type'. " + "Select from 'CPU', 'GPU', 'GPU.0', 'GPU.1', 'NPU' or from" + " HETERO/MULTI/AUTO options available. \n"); + } + } else if (key == "device_id") { + if (value == "CPU" || value == "GPU" || value == "NPU") { + ov_options[key] = value; + } else { + ORT_THROW("[ERROR] [OpenVINO] Unsupported device_id is selected. Select from available options."); + } + } else if (key == "precision") { + auto device_type = ov_options["device_type"]; + if (device_type.find("GPU") != std::string::npos) { + if (value == "") { + ov_options[key] = "FP16"; + continue; + } else if (value == "ACCURACY" || value == "FP16" || value == "FP32") { + ov_options[key] = value; + continue; + } else { + ORT_THROW( + "[ERROR] [OpenVINO] Unsupported inference precision is selected. " + "GPU only supported FP32 / FP16. \n"); + } + } else if (device_type.find("NPU") != std::string::npos) { + if (value == "" || value == "ACCURACY" || value == "FP16") { + ov_options[key] = "FP16"; + continue; + } else { + ORT_THROW("[ERROR] [OpenVINO] Unsupported inference precision is selected. NPU only supported FP16. \n"); + } + } else if (device_type.find("CPU") != std::string::npos) { + if (value == "" || value == "ACCURACY" || value == "FP32") { + ov_options[key] = "FP32"; + continue; + } else { + ORT_THROW("[ERROR] [OpenVINO] Unsupported inference precision is selected. CPU only supports FP32 . \n"); + } + } + } else if (key == "enable_npu_fast_compile") { + if (value == "true" || value == "True" || + value == "false" || value == "False") { + ov_options[key] = value; + } else { + ORT_THROW("[ERROR] [OpenVINO] The value for the key 'enable_npu_fast_compile' should be a boolean i.e. true or false. Default value is false.\n"); + } + } else if (key == "enable_opencl_throttling") { + if (value == "true" || value == "True" || + value == "false" || value == "False") { + ov_options[key] = value; + } else { + ORT_THROW("[ERROR] [OpenVINO] The value for the key 'enable_opencl_throttling' should be a boolean i.e. true or false. Default value is false.\n"); + } + } else if (key == "enable_qdq_optimizer") { + if (value == "true" || value == "True" || + value == "false" || value == "False") { + ov_options[key] = value; + } else { + ORT_THROW("[ERROR] [OpenVINO] The value for the key 'enable_qdq_optimizer' should be a boolean i.e. true or false. Default value is false.\n"); + } + } else if (key == "disable_dynamic_shapes") { + if (value == "true" || value == "True" || + value == "false" || value == "False") { + ov_options[key] = value; + } else { + ORT_THROW( + "[ERROR] [OpenVINO] The value for the key 'enable_dynamic_shapes' " + "should be a boolean i.e. true or false. Default value is false.\n"); + } + } else if (key == "num_of_threads") { + if (std::stoi(value) <= 0) { + ORT_THROW("[ERROR] [OpenVINO] The value for the key 'num_of_threads' should be greater than 0\n"); + } else { + ov_options[key] = value; + } + } else if (key == "model_priority") { + ov_options[key] = value; + } else if (key == "cache_dir") { + ov_options[key] = value; + } else if (key == "context") { + ov_options[key] = value; + } else if (key == "num_streams") { + if (std::stoi(value) <= 0 && std::stoi(value) > 8) { + ORT_THROW("[ERROR] [OpenVINO] The value for the key 'num_streams' should be in the range of 1-8 \n"); + } else { + ov_options[key] = value; + } + } else if (key == "export_ep_ctx_blob") { + if (value == "true" || value == "True" || + value == "false" || value == "False") { + ov_options[key] = value; + } else { + ORT_THROW( + "[ERROR] [OpenVINO] The value for the key 'export_ep_ctx_blob' " + "should be a boolean i.e. true or false. Default value is false.\n"); + } + } else { + ORT_THROW("[ERROR] [OpenVINO] wrong key type entered. Choose from the following runtime key options that are available for OpenVINO. ['device_type', 'device_id', 'enable_npu_fast_compile', 'num_of_threads', 'cache_dir', 'num_streams', 'enable_opencl_throttling', 'disable_dynamic_shapes'] \n"); + } + } + session_options.AppendExecutionProvider_OpenVINO_V2(ov_options); +#else + ORT_THROW("OpenVINO is not supported in this build\n"); +#endif + } session_ = Ort::Session(env, performance_test_config.model_info.model_file_path.c_str(), session_options); diff --git a/onnxruntime/test/util/default_providers.cc b/onnxruntime/test/util/default_providers.cc index 6f07385729555..f15ac100f4e3f 100644 --- a/onnxruntime/test/util/default_providers.cc +++ b/onnxruntime/test/util/default_providers.cc @@ -109,7 +109,8 @@ std::unique_ptr OpenVINOExecutionProviderWithOptions(const O std::unique_ptr DefaultOpenVINOExecutionProvider() { #ifdef USE_OPENVINO ProviderOptions provider_options_map; - return OpenVINOProviderFactoryCreator::Create(&provider_options_map)->CreateProvider(); + SessionOptions session_options; + return OpenVINOProviderFactoryCreator::Create(&provider_options_map, &session_options)->CreateProvider(); #else return nullptr; #endif diff --git a/tools/ci_build/build.py b/tools/ci_build/build.py index 3e587e9b56e2e..f431f471c4082 100644 --- a/tools/ci_build/build.py +++ b/tools/ci_build/build.py @@ -79,6 +79,7 @@ def _openvino_verify_device_type(device_read): "CPU_NO_PARTITION", "GPU_NO_PARTITION", "NPU_NO_PARTITION", + "NPU_NO_CPU_FALLBACK", ] status_hetero = True res = False @@ -1227,6 +1228,7 @@ def generate_build_tree( if args.use_openvino: cmake_args += [ "-Donnxruntime_USE_OPENVINO=ON", + "-Donnxruntime_NPU_NO_FALLBACK=" + ("ON" if args.use_openvino == "NPU_NO_CPU_FALLBACK" else "OFF"), "-Donnxruntime_USE_OPENVINO_GPU=" + ("ON" if args.use_openvino == "GPU" else "OFF"), "-Donnxruntime_USE_OPENVINO_CPU=" + ("ON" if args.use_openvino == "CPU" else "OFF"), "-Donnxruntime_USE_OPENVINO_NPU=" + ("ON" if args.use_openvino == "NPU" else "OFF"), From 7e93cd7f8b21d00b8d406cb043ff5f0e61217d85 Mon Sep 17 00:00:00 2001 From: mingyueliuh <131847423+mingyueliuh@users.noreply.github.com> Date: Fri, 28 Jun 2024 20:19:20 -0400 Subject: [PATCH 44/52] [VitisAI] Align TensorProto_DataType with onnx1.16 (#21067) ### Description Vitis AI EP synchronously supports the TensorProto data types supported by ONNX 1.16. Add error message show when graph resolve fail for troubleshooting. ### Motivation and Context ONNX 1.15 & 1.16 add support some new TensorProto DataType , such as - FLOAT8E4M3FN - FLOAT8E4M3FNUZ - FLOAT8E5M2 - FLOAT8E5M2FNUZ - UINT4 - INT4 --------- Co-authored-by: liumingyue --- .../core/providers/vitisai/imp/global_api.cc | 3 +++ .../providers/vitisai/include/vaip/my_ort.h | 8 ++++++- .../vitisai/include/vaip/vaip_ort_api.h | 2 +- .../core/session/provider_bridge_ort.cc | 24 +++++++++++++++++++ 4 files changed, 35 insertions(+), 2 deletions(-) diff --git a/onnxruntime/core/providers/vitisai/imp/global_api.cc b/onnxruntime/core/providers/vitisai/imp/global_api.cc index 29a1231fdce18..1133751d82d65 100644 --- a/onnxruntime/core/providers/vitisai/imp/global_api.cc +++ b/onnxruntime/core/providers/vitisai/imp/global_api.cc @@ -270,6 +270,9 @@ vaip_core::OrtApiForVaip* create_org_api_hook() { graph.SetGraphResolveNeeded(); } auto status = graph.Resolve(); + if (!status.IsOK()) { + std::cerr << "graph resolve error:" << status.ErrorMessage() << std::endl; + } return status.Code(); }; the_global_api.graph_get_consumer_nodes_unsafe = [](const Graph& graph, const std::string& node_arg_name) -> auto { diff --git a/onnxruntime/core/providers/vitisai/include/vaip/my_ort.h b/onnxruntime/core/providers/vitisai/include/vaip/my_ort.h index 46fc4ac9b2a5d..74482d8e9ee0e 100644 --- a/onnxruntime/core/providers/vitisai/include/vaip/my_ort.h +++ b/onnxruntime/core/providers/vitisai/include/vaip/my_ort.h @@ -38,7 +38,13 @@ enum TensorProto_DataType : int { TensorProto_DataType_UINT64 = 13, TensorProto_DataType_COMPLEX64 = 14, TensorProto_DataType_COMPLEX128 = 15, - TensorProto_DataType_BFLOAT16 = 16 + TensorProto_DataType_BFLOAT16 = 16, + TensorProto_DataType_FLOAT8E4M3FN = 17, + TensorProto_DataType_FLOAT8E4M3FNUZ = 18, + TensorProto_DataType_FLOAT8E5M2 = 19, + TensorProto_DataType_FLOAT8E5M2FNUZ = 20, + TensorProto_DataType_UINT4 = 21, + TensorProto_DataType_INT4 = 22 }; enum AttributeProto_AttributeType : int { AttributeProto_AttributeType_UNDEFINED = 0, diff --git a/onnxruntime/core/providers/vitisai/include/vaip/vaip_ort_api.h b/onnxruntime/core/providers/vitisai/include/vaip/vaip_ort_api.h index 62a7bb602e7e8..3346739890484 100644 --- a/onnxruntime/core/providers/vitisai/include/vaip/vaip_ort_api.h +++ b/onnxruntime/core/providers/vitisai/include/vaip/vaip_ort_api.h @@ -13,7 +13,7 @@ struct OrtApi; namespace vaip_core { #define VAIP_ORT_API_MAJOR (3u) -#define VAIP_ORT_API_MINOR (0u) +#define VAIP_ORT_API_MINOR (1u) #define VAIP_ORT_API_PATCH (0u) struct OrtApiForVaip { uint32_t magic; // 'VAIP' or something else to make sure the following field diff --git a/onnxruntime/core/session/provider_bridge_ort.cc b/onnxruntime/core/session/provider_bridge_ort.cc index d4c6e3d506f18..408ad7815835f 100644 --- a/onnxruntime/core/session/provider_bridge_ort.cc +++ b/onnxruntime/core/session/provider_bridge_ort.cc @@ -613,8 +613,12 @@ struct ProviderHostImpl : ProviderHost { elemType = ONNX_NAMESPACE::TensorProto_DataType_UINT8; } else if (data_type->s() == "int32") { elemType = ONNX_NAMESPACE::TensorProto_DataType_INT32; + } else if (data_type->s() == "uint32") { + elemType = ONNX_NAMESPACE::TensorProto_DataType_UINT32; } else if (data_type->s() == "int64") { elemType = ONNX_NAMESPACE::TensorProto_DataType_INT64; + } else if (data_type->s() == "uint64") { + elemType = ONNX_NAMESPACE::TensorProto_DataType_UINT64; } else if (data_type->s() == "int1") { elemType = ONNX_NAMESPACE::TensorProto_DataType_BOOL; } else if (data_type->s() == "bfloat16") { @@ -625,6 +629,26 @@ struct ProviderHostImpl : ProviderHost { elemType = ONNX_NAMESPACE::TensorProto_DataType_UINT16; } else if (data_type->s() == "int16") { elemType = ONNX_NAMESPACE::TensorProto_DataType_INT16; + } else if (data_type->s() == "double") { + elemType = ONNX_NAMESPACE::TensorProto_DataType_DOUBLE; + } else if (data_type->s() == "string") { + elemType = ONNX_NAMESPACE::TensorProto_DataType_STRING; + } else if (data_type->s() == "complex64") { + elemType = ONNX_NAMESPACE::TensorProto_DataType_COMPLEX64; + } else if (data_type->s() == "complex128") { + elemType = ONNX_NAMESPACE::TensorProto_DataType_COMPLEX128; + } else if (data_type->s() == "float8e4m3fn") { + elemType = ONNX_NAMESPACE::TensorProto_DataType_FLOAT8E4M3FN; + } else if (data_type->s() == "float8e4m3fnuz") { + elemType = ONNX_NAMESPACE::TensorProto_DataType_FLOAT8E4M3FNUZ; + } else if (data_type->s() == "float8e5m2") { + elemType = ONNX_NAMESPACE::TensorProto_DataType_FLOAT8E5M2; + } else if (data_type->s() == "float8e5m2funz") { + elemType = ONNX_NAMESPACE::TensorProto_DataType_FLOAT8E5M2FNUZ; + } else if (data_type->s() == "uint4") { + elemType = ONNX_NAMESPACE::TensorProto_DataType_UINT4; + } else if (data_type->s() == "int4") { + elemType = ONNX_NAMESPACE::TensorProto_DataType_INT4; } else { return; } From 0cbe7eec5e4c8cdd7146a44a4a49236c99b54597 Mon Sep 17 00:00:00 2001 From: Jian Chen Date: Fri, 28 Jun 2024 19:49:54 -0700 Subject: [PATCH 45/52] Uppdate nuget to Use Nuget 6.10.x (#21209) ### Description Uppdate nuget to Use Nuget 6.10.x --- .../azure-pipelines/c-api-noopenmp-packaging-pipelines.yml | 4 ++-- .../github/azure-pipelines/nuget/templates/dml-vs-2022.yml | 4 ++-- .../github/azure-pipelines/nuget/templates/test_win.yml | 4 ++-- .../azure-pipelines/stages/nuget-cuda-packaging-stage.yml | 4 ++-- tools/ci_build/github/azure-pipelines/templates/c-api-cpu.yml | 4 ++-- .../azure-pipelines/templates/jobs/win-ci-vs-2022-job.yml | 4 ++-- .../templates/ondevice-training-cpu-packaging-pipeline.yml | 4 ++-- tools/ci_build/github/azure-pipelines/win-ci-fuzz-testing.yml | 4 ++-- 8 files changed, 16 insertions(+), 16 deletions(-) diff --git a/tools/ci_build/github/azure-pipelines/c-api-noopenmp-packaging-pipelines.yml b/tools/ci_build/github/azure-pipelines/c-api-noopenmp-packaging-pipelines.yml index 990d0c253c789..3aadefecaab87 100644 --- a/tools/ci_build/github/azure-pipelines/c-api-noopenmp-packaging-pipelines.yml +++ b/tools/ci_build/github/azure-pipelines/c-api-noopenmp-packaging-pipelines.yml @@ -307,9 +307,9 @@ stages: displayName: 'Create models link' - task: NuGetToolInstaller@0 - displayName: Use Nuget 6.2.1 + displayName: Use Nuget 6.10.x inputs: - versionSpec: 6.2.1 + versionSpec: 6.10.x - task: MSBuild@1 displayName: 'Restore NuGet Packages and create project.assets.json' diff --git a/tools/ci_build/github/azure-pipelines/nuget/templates/dml-vs-2022.yml b/tools/ci_build/github/azure-pipelines/nuget/templates/dml-vs-2022.yml index 344c9a8f14022..f20a1ae3e1cd9 100644 --- a/tools/ci_build/github/azure-pipelines/nuget/templates/dml-vs-2022.yml +++ b/tools/ci_build/github/azure-pipelines/nuget/templates/dml-vs-2022.yml @@ -135,9 +135,9 @@ stages: - task: NuGetToolInstaller@0 - displayName: Use Nuget 5.7.0 + displayName: Use Nuget 6.10.x inputs: - versionSpec: 5.7.0 + versionSpec: 6.10.x - task: MSBuild@1 displayName: 'Restore NuGet Packages' diff --git a/tools/ci_build/github/azure-pipelines/nuget/templates/test_win.yml b/tools/ci_build/github/azure-pipelines/nuget/templates/test_win.yml index c582a836c7dbd..869374aa5a6e7 100644 --- a/tools/ci_build/github/azure-pipelines/nuget/templates/test_win.yml +++ b/tools/ci_build/github/azure-pipelines/nuget/templates/test_win.yml @@ -45,9 +45,9 @@ stages: architecture: x64 - task: NuGetToolInstaller@0 - displayName: Use Nuget 5.7.0 + displayName: Use Nuget 6.10.x inputs: - versionSpec: 5.7.0 + versionSpec: 6.10.x - ${{ if ne( parameters.CudaVersion, '') }}: - template: ../../templates/jobs/download_win_gpu_library.yml parameters: diff --git a/tools/ci_build/github/azure-pipelines/stages/nuget-cuda-packaging-stage.yml b/tools/ci_build/github/azure-pipelines/stages/nuget-cuda-packaging-stage.yml index 18615b6ca18b1..424ed6237260b 100644 --- a/tools/ci_build/github/azure-pipelines/stages/nuget-cuda-packaging-stage.yml +++ b/tools/ci_build/github/azure-pipelines/stages/nuget-cuda-packaging-stage.yml @@ -76,9 +76,9 @@ stages: displayName: 'Create models link' - task: NuGetToolInstaller@0 - displayName: Use Nuget 6.2.1 + displayName: Use Nuget 6.10.x inputs: - versionSpec: 6.2.1 + versionSpec: 6.10.x - task: PowerShell@2 displayName: Install MAUI workloads diff --git a/tools/ci_build/github/azure-pipelines/templates/c-api-cpu.yml b/tools/ci_build/github/azure-pipelines/templates/c-api-cpu.yml index c41f9589d8469..7ba1179e7ad4d 100644 --- a/tools/ci_build/github/azure-pipelines/templates/c-api-cpu.yml +++ b/tools/ci_build/github/azure-pipelines/templates/c-api-cpu.yml @@ -376,9 +376,9 @@ stages: workingDirectory: '$(Build.BinariesDirectory)' displayName: 'Create models link' - task: NuGetToolInstaller@0 - displayName: Use Nuget 6.2.1 + displayName: Use Nuget 6.10.x inputs: - versionSpec: 6.2.1 + versionSpec: 6.10.x - task: PowerShell@2 displayName: Install mobile workloads diff --git a/tools/ci_build/github/azure-pipelines/templates/jobs/win-ci-vs-2022-job.yml b/tools/ci_build/github/azure-pipelines/templates/jobs/win-ci-vs-2022-job.yml index fee4a82ceac78..9a35d7b75c708 100644 --- a/tools/ci_build/github/azure-pipelines/templates/jobs/win-ci-vs-2022-job.yml +++ b/tools/ci_build/github/azure-pipelines/templates/jobs/win-ci-vs-2022-job.yml @@ -123,9 +123,9 @@ jobs: displayName: 'API Documentation Check and generate' - task: NuGetToolInstaller@0 - displayName: Use Nuget 5.7.0 + displayName: Use Nuget 6.10.x inputs: - versionSpec: 5.7.0 + versionSpec: 6.10.x - task: NuGetCommand@2 displayName: 'NuGet restore' diff --git a/tools/ci_build/github/azure-pipelines/templates/ondevice-training-cpu-packaging-pipeline.yml b/tools/ci_build/github/azure-pipelines/templates/ondevice-training-cpu-packaging-pipeline.yml index bc75a115326f6..fb9ff65fe8534 100644 --- a/tools/ci_build/github/azure-pipelines/templates/ondevice-training-cpu-packaging-pipeline.yml +++ b/tools/ci_build/github/azure-pipelines/templates/ondevice-training-cpu-packaging-pipeline.yml @@ -191,9 +191,9 @@ stages: displayName: 'Create models link' - task: NuGetToolInstaller@0 - displayName: Use Nuget 6.2.1 + displayName: Use Nuget 6.10.x inputs: - versionSpec: 6.2.1 + versionSpec: 6.10.x - task: PowerShell@2 displayName: Install mobile workloads diff --git a/tools/ci_build/github/azure-pipelines/win-ci-fuzz-testing.yml b/tools/ci_build/github/azure-pipelines/win-ci-fuzz-testing.yml index d6fa5184f6882..b5120f01bff3e 100644 --- a/tools/ci_build/github/azure-pipelines/win-ci-fuzz-testing.yml +++ b/tools/ci_build/github/azure-pipelines/win-ci-fuzz-testing.yml @@ -39,9 +39,9 @@ jobs: versionSpec: '18.x' - task: NuGetToolInstaller@0 - displayName: Use Nuget 5.7.0 + displayName: Use Nuget 6.10.x inputs: - versionSpec: 5.7.0 + versionSpec: 6.10.x - task: PythonScript@0 displayName: 'Generate cmake config' From 3a83f8b317ba1e52cb8ee824ad2c5d02f677c82f Mon Sep 17 00:00:00 2001 From: Changming Sun Date: Fri, 28 Jun 2024 20:03:57 -0700 Subject: [PATCH 46/52] Update the functions in tensorprotoutils.h to use std::filesystem::path instead (#20920) ### Description 1. Update the functions in tensorprotoutils.h to use std::filesystem::path instead of onnxruntime::Path. Eventually we can remove the whole onnxruntime::Path class, but to this PR small I am not doing that. 2. Remove the _SILENCE_EXPERIMENTAL_FILESYSTEM_DEPRECATION_WARNING macro def when TensorRT EP is enabled. --- cmake/onnxruntime_providers_tensorrt.cmake | 1 - include/onnxruntime/core/graph/graph.h | 22 +- include/onnxruntime/core/graph/graph_viewer.h | 5 +- .../core/framework/graph_partitioner.cc | 24 +- .../framework/model_metadef_id_generator.cc | 2 +- onnxruntime/core/framework/node_unit.cc | 2 +- onnxruntime/core/framework/node_unit.h | 3 +- onnxruntime/core/framework/session_options.h | 3 +- .../core/framework/tensorprotoutils.cc | 412 ++++++++---------- onnxruntime/core/framework/tensorprotoutils.h | 25 +- onnxruntime/core/graph/graph.cc | 27 +- .../core/graph/graph_flatbuffers_utils.cc | 6 +- .../core/graph/graph_flatbuffers_utils.h | 7 +- onnxruntime/core/graph/model.cc | 26 +- onnxruntime/core/graph/model.h | 22 +- onnxruntime/core/optimizer/initializer.cc | 6 +- onnxruntime/core/optimizer/initializer.h | 4 +- .../core/optimizer/matmul_scale_fusion.cc | 3 +- .../optimizer/optimizer_execution_frame.cc | 9 +- .../optimizer/optimizer_execution_frame.h | 5 +- .../optimizer/qdq_transformer/qdq_util.cc | 4 +- .../core/optimizer/qdq_transformer/qdq_util.h | 5 +- .../ort_optimizer_api_impl.cc | 11 +- .../core/providers/cpu/ml/label_encoder.h | 6 +- .../src/DmlRuntimeFusedGraphKernel.cpp | 6 +- .../src/DmlRuntimeFusedGraphKernel.h | 4 +- .../src/GraphDescBuilder.cpp | 2 +- .../src/GraphDescBuilder.h | 3 +- .../src/MLOperatorAuthorImpl.cpp | 8 +- .../src/MLOperatorAuthorImpl.h | 6 +- .../dml/DmlExecutionProvider/src/Utility.h | 17 +- .../nnapi/nnapi_builtin/builders/helper.cc | 3 +- .../nnapi/nnapi_builtin/builders/helper.h | 3 +- .../builders/op_builder_helpers.cc | 4 +- .../builders/op_builder_helpers.h | 5 +- .../openvino/openvino_execution_provider.cc | 10 +- .../qnn/builder/opbuilder/base_op_builder.cc | 5 +- .../qnn/builder/opbuilder/slice_op_builder.cc | 6 +- .../providers/qnn/qnn_execution_provider.cc | 2 +- .../provider_bridge_provider.cc | 2 +- .../shared_library/provider_interfaces.h | 11 +- .../shared_library/provider_wrappedtypes.h | 8 +- .../tensorrt/onnx_ctx_model_helper.cc | 2 +- .../tensorrt/onnx_ctx_model_helper.h | 2 +- .../tensorrt/tensorrt_execution_provider.cc | 7 +- .../tensorrt_execution_provider_utils.h | 6 +- .../core/providers/vitisai/imp/graph.cc | 2 +- onnxruntime/core/session/custom_ops.cc | 4 +- onnxruntime/core/session/inference_session.cc | 9 +- onnxruntime/core/session/inference_session.h | 3 +- .../core/session/provider_bridge_ort.cc | 10 +- .../test/flatbuffers/flatbuffer_utils_test.cc | 4 +- .../test/framework/allocation_planner_test.cc | 2 +- .../test/framework/inference_session_test.cc | 12 +- .../save_model_with_external_initializers.cc | 62 +-- .../test/framework/sparse_kernels_test.cc | 92 ++-- .../test/framework/tensorutils_test.cc | 16 +- .../test/framework/test_tensor_loader.cc | 10 +- onnxruntime/test/ir/graph_test.cc | 6 +- .../test/optimizer/initializer_test.cc | 60 +-- .../test/optimizer/resnet50_fusion_test.cc | 12 +- .../providers/tensorrt/tensorrt_basic_test.cc | 46 +- onnxruntime/test/util/test_utils.cc | 2 +- .../core/framework/checkpoint_common.cc | 4 +- .../core/framework/ortmodule_graph_builder.cc | 2 +- .../core/optimizer/graph_transformer_config.h | 1 + .../models/runner/training_util.cc | 2 +- .../python/orttraining_pybind_state.cc | 2 +- .../test/gradient/allreduce_op_test.cc | 12 +- .../test/optimizer/graph_transform_test.cc | 28 +- .../orttraining/training_api/checkpoint.cc | 2 +- .../orttraining/training_api/module.cc | 2 +- 72 files changed, 566 insertions(+), 603 deletions(-) diff --git a/cmake/onnxruntime_providers_tensorrt.cmake b/cmake/onnxruntime_providers_tensorrt.cmake index e56de0c7124dc..90203216600fa 100644 --- a/cmake/onnxruntime_providers_tensorrt.cmake +++ b/cmake/onnxruntime_providers_tensorrt.cmake @@ -13,7 +13,6 @@ set(OLD_CMAKE_CXX_FLAGS ${CMAKE_CXX_FLAGS}) set(PROTOBUF_LIBRARY ${PROTOBUF_LIB}) if (WIN32) - add_definitions(-D_SILENCE_EXPERIMENTAL_FILESYSTEM_DEPRECATION_WARNING=1) set(OLD_CMAKE_CUDA_FLAGS ${CMAKE_CUDA_FLAGS}) set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} /wd4099 /wd4551 /wd4505 /wd4515 /wd4706 /wd4456 /wd4324 /wd4701 /wd4804 /wd4702 /wd4458 /wd4703") if (CMAKE_BUILD_TYPE STREQUAL "Debug") diff --git a/include/onnxruntime/core/graph/graph.h b/include/onnxruntime/core/graph/graph.h index 4f3377f0aa0c0..538cbfdcefc47 100644 --- a/include/onnxruntime/core/graph/graph.h +++ b/include/onnxruntime/core/graph/graph.h @@ -10,16 +10,7 @@ #include #include #include - -#ifdef _WIN32 -#pragma warning(push) -// disable some warnings from protobuf to pass Windows build -#pragma warning(disable : 4244) -#endif - -#ifdef _WIN32 -#pragma warning(pop) -#endif +#include #include "core/common/flatbuffers.h" @@ -147,7 +138,7 @@ class Node { const std::string& Domain() const noexcept { return domain_; } /** Gets the path of the owning model if any. */ - const Path& ModelPath() const noexcept; + const std::filesystem::path& ModelPath() const noexcept; /** Gets the Node's execution priority. @remarks Lower value means higher priority */ @@ -693,7 +684,7 @@ class Graph { // NOLINT(clang-analyzer-optin.performance.Padding): preserve exi const std::string& Description() const noexcept; /** Gets the path of the owning model, if any. */ - const Path& ModelPath() const; + const std::filesystem::path& ModelPath() const; /** Returns true if this is a subgraph or false if it is a high-level graph. */ bool IsSubgraph() const { return parent_graph_ != nullptr; } @@ -1149,13 +1140,14 @@ class Graph { // NOLINT(clang-analyzer-optin.performance.Padding): preserve exi ONNX_NAMESPACE::GraphProto ToGraphProto() const; /** Gets the GraphProto representation of this Graph - @params external_file_name name of the binary file to use for initializers + @param external_file_path File path of the binary file to use for initializers. + @param model_file_path path of the model file. @param initializer_size_threshold initializers larger or equal to this threshold (in bytes) are saved in the external file. Initializer smaller than this threshold are included in the onnx file. @returns GraphProto serialization of the graph. */ - ONNX_NAMESPACE::GraphProto ToGraphProtoWithExternalInitializers(const std::string& external_file_name, - const PathString& file_path, + ONNX_NAMESPACE::GraphProto ToGraphProtoWithExternalInitializers(const std::filesystem::path& external_file_path, + const std::filesystem::path& model_file_path, size_t initializer_size_threshold) const; /** Gets the ISchemaRegistry instances being used with this Graph. */ diff --git a/include/onnxruntime/core/graph/graph_viewer.h b/include/onnxruntime/core/graph/graph_viewer.h index 1816099d3210f..9385e2f092e58 100644 --- a/include/onnxruntime/core/graph/graph_viewer.h +++ b/include/onnxruntime/core/graph/graph_viewer.h @@ -2,10 +2,11 @@ // Licensed under the MIT License. #pragma once +#include +#include #include "core/graph/graph.h" #include "core/framework/session_options.h" -#include namespace onnxruntime { class Function; @@ -43,7 +44,7 @@ class GraphViewer { const std::string& Description() const noexcept; /** Gets the path of the owning model if any **/ - const Path& ModelPath() const noexcept { return graph_->ModelPath(); } + const std::filesystem::path& ModelPath() const noexcept { return graph_->ModelPath(); } /** Gets a tensor created from an initializer. diff --git a/onnxruntime/core/framework/graph_partitioner.cc b/onnxruntime/core/framework/graph_partitioner.cc index 90ee8a46f66a9..4f745b74abce7 100644 --- a/onnxruntime/core/framework/graph_partitioner.cc +++ b/onnxruntime/core/framework/graph_partitioner.cc @@ -637,7 +637,7 @@ static Status InlineFunctionsAOTImpl(const ExecutionProviders& execution_provide static Status CreateEpContextModel(const ExecutionProviders& execution_providers, const Graph& graph, - const std::string& ep_context_path, + const std::filesystem::path& ep_context_path, const logging::Logger& logger) { InlinedVector all_ep_context_nodes; for (const auto& ep : execution_providers) { @@ -658,22 +658,20 @@ static Status CreateEpContextModel(const ExecutionProviders& execution_providers return std::make_pair(false, static_cast(nullptr)); }; - onnxruntime::PathString context_cache_path; - PathString model_pathstring = graph.ModelPath().ToPathString(); + std::filesystem::path context_cache_path; + const std::filesystem::path& model_path = graph.ModelPath(); if (!ep_context_path.empty()) { - context_cache_path = ToPathString(ep_context_path); - } else if (!model_pathstring.empty()) { - context_cache_path = model_pathstring + ToPathString("_ctx.onnx"); + context_cache_path = ep_context_path; + } else if (!model_path.empty()) { + context_cache_path = model_path.native() + ORT_TSTR("_ctx.onnx"); + } else { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Both ep_context_path and model_path are empty"); } - { -#ifdef _WIN32 - std::wifstream fs(context_cache_path); -#else - std::ifstream fs(context_cache_path); -#endif - ORT_RETURN_IF(fs.good(), "Failed to generate EP context model since the file exist already."); + if (std::filesystem::exists(context_cache_path)) { + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Failed to generate EP context model since the file '", + context_cache_path, "' exist already."); } Model ep_context_model(graph.Name(), false, ModelMetaData(), PathString(), IOnnxRuntimeOpSchemaRegistryList(), diff --git a/onnxruntime/core/framework/model_metadef_id_generator.cc b/onnxruntime/core/framework/model_metadef_id_generator.cc index e51c6ebc29975..8b1d1f4f304c9 100644 --- a/onnxruntime/core/framework/model_metadef_id_generator.cc +++ b/onnxruntime/core/framework/model_metadef_id_generator.cc @@ -40,7 +40,7 @@ int ModelMetadefIdGenerator::GenerateId(const onnxruntime::GraphViewer& graph_vi // prefer path the model was loaded from // this may not be available if the model was loaded from a stream or in-memory bytes - const auto& model_path_str = main_graph.ModelPath().ToPathString(); + const auto model_path_str = main_graph.ModelPath().string(); if (!model_path_str.empty()) { MurmurHash3::x86_128(model_path_str.data(), gsl::narrow_cast(model_path_str.size()), hash[0], &hash); } else { diff --git a/onnxruntime/core/framework/node_unit.cc b/onnxruntime/core/framework/node_unit.cc index 4e2f22dea164d..ac4301641105a 100644 --- a/onnxruntime/core/framework/node_unit.cc +++ b/onnxruntime/core/framework/node_unit.cc @@ -277,7 +277,7 @@ const std::string& NodeUnit::OpType() const noexcept { return target_node_.OpTyp const std::string& NodeUnit::Name() const noexcept { return target_node_.Name(); } int NodeUnit::SinceVersion() const noexcept { return target_node_.SinceVersion(); } NodeIndex NodeUnit::Index() const noexcept { return target_node_.Index(); } -const Path& NodeUnit::ModelPath() const noexcept { return target_node_.ModelPath(); } +const std::filesystem::path& NodeUnit::ModelPath() const noexcept { return target_node_.ModelPath(); } ProviderType NodeUnit::GetExecutionProviderType() const noexcept { return target_node_.GetExecutionProviderType(); } void NodeUnit::InitForSingleNode() { diff --git a/onnxruntime/core/framework/node_unit.h b/onnxruntime/core/framework/node_unit.h index a168495f12ebf..e84e62479162f 100644 --- a/onnxruntime/core/framework/node_unit.h +++ b/onnxruntime/core/framework/node_unit.h @@ -9,6 +9,7 @@ #include #include #include +#include #include "core/graph/basic_types.h" #include "core/graph/graph.h" @@ -78,7 +79,7 @@ class NodeUnit { const std::string& Name() const noexcept; int SinceVersion() const noexcept; NodeIndex Index() const noexcept; - const Path& ModelPath() const noexcept; + const std::filesystem::path& ModelPath() const noexcept; ProviderType GetExecutionProviderType() const noexcept; const Node& GetNode() const noexcept { return target_node_; } diff --git a/onnxruntime/core/framework/session_options.h b/onnxruntime/core/framework/session_options.h index 0453a7ecac81f..13da26d5e6053 100644 --- a/onnxruntime/core/framework/session_options.h +++ b/onnxruntime/core/framework/session_options.h @@ -7,6 +7,7 @@ #include #include #include +#include #include "core/common/gsl.h" #include "core/common/inlined_containers.h" #include "core/framework/config_options.h" @@ -89,7 +90,7 @@ struct SessionOptions { // // If session config value is not set, it will be assumed to be ONNX // unless the filepath ends in '.ort' (case insensitive). - std::basic_string optimized_model_filepath; + std::filesystem::path optimized_model_filepath; // enable the memory pattern optimization. // The idea is if the input shapes are the same, we could trace the internal memory allocation diff --git a/onnxruntime/core/framework/tensorprotoutils.cc b/onnxruntime/core/framework/tensorprotoutils.cc index 6af78f18fb82f..77323f268a27d 100644 --- a/onnxruntime/core/framework/tensorprotoutils.cc +++ b/onnxruntime/core/framework/tensorprotoutils.cc @@ -6,7 +6,7 @@ #include #include #include - +#include #if defined(__wasm__) #include #endif @@ -111,8 +111,8 @@ namespace onnxruntime { namespace { // This function doesn't support string tensors -static Status UnpackTensorWithRawDataImpl(const void* raw_data, size_t raw_data_len, - size_t expected_num_elements, size_t element_size, +static Status UnpackTensorWithRawDataImpl(const void* raw_data, size_t raw_data_len, size_t expected_num_elements, + size_t element_size, /*out*/ unsigned char* p_data) { auto src = gsl::make_span(static_cast(raw_data), raw_data_len); auto dst = gsl::make_span(p_data, expected_num_elements * element_size); @@ -152,8 +152,8 @@ Status UnpackTensorWithRawData(const void* raw_data, size_t raw_data_len, size_t size_t num_packed_pairs = INT4_TYPE::CalcNumInt4Pairs(expected_num_elements); \ ORT_RETURN_IF_NOT(num_packed_pairs == raw_data_len, "Unexpected number of packed int4 pairs"); \ \ - gsl::span src_span = gsl::make_span(reinterpret_cast(raw_data), \ - num_packed_pairs); \ + gsl::span src_span = \ + gsl::make_span(reinterpret_cast(raw_data), num_packed_pairs); \ gsl::span dst_span = gsl::make_span(p_data, num_packed_pairs); \ \ std::memcpy(dst_span.data(), src_span.data(), num_packed_pairs); \ @@ -165,7 +165,7 @@ DEFINE_INT4_UNPACK_TENSOR_WITH_RAW_DATA_IMPL(Int4x2) DEFINE_INT4_UNPACK_TENSOR_WITH_RAW_DATA_IMPL(UInt4x2) static Status GetExternalDataInfo(const ONNX_NAMESPACE::TensorProto& tensor_proto, - const ORTCHAR_T* tensor_proto_dir, + const std::filesystem::path& tensor_proto_dir, std::basic_string& external_file_path, onnxruntime::FileOffsetType& file_offset, SafeInt& tensor_byte_size) { @@ -180,22 +180,15 @@ static Status GetExternalDataInfo(const ONNX_NAMESPACE::TensorProto& tensor_prot const auto& location = external_data_info->GetRelPath(); - if (location == onnxruntime::utils::kTensorProtoMemoryAddressTag) { - external_file_path = location; - } else { - if (tensor_proto_dir != nullptr) { - external_file_path = onnxruntime::ConcatPathComponent(tensor_proto_dir, - external_data_info->GetRelPath()); - } else { - external_file_path = external_data_info->GetRelPath(); - } - } + external_file_path = location == onnxruntime::utils::kTensorProtoMemoryAddressTag ? std::filesystem::path(location) + : (tensor_proto_dir / location); ORT_RETURN_IF_ERROR(onnxruntime::utils::GetSizeInBytesFromTensorProto<0>(tensor_proto, &tensor_byte_size)); const size_t external_data_length = external_data_info->GetLength(); ORT_RETURN_IF_NOT(external_data_length == 0 || external_data_length == tensor_byte_size, - "TensorProto: ", tensor_proto.name(), " external data size mismatch. Computed size: ", - *&tensor_byte_size, ", external_data.length: ", external_data_length); + "TensorProto: ", tensor_proto.name(), + " external data size mismatch. Computed size: ", *&tensor_byte_size, + ", external_data.length: ", external_data_length); file_offset = external_data_info->GetOffset(); @@ -207,17 +200,13 @@ static Status GetExternalDataInfo(const ONNX_NAMESPACE::TensorProto& tensor_prot // then uses the current directory instead. // This function does not unpack string_data of an initializer tensor Status ReadExternalDataForTensor(const ONNX_NAMESPACE::TensorProto& tensor_proto, - const ORTCHAR_T* tensor_proto_dir, + const std::filesystem::path& tensor_proto_dir, std::vector& unpacked_tensor) { std::basic_string external_file_path; onnxruntime::FileOffsetType file_offset; SafeInt tensor_byte_size; - ORT_RETURN_IF_ERROR(GetExternalDataInfo( - tensor_proto, - tensor_proto_dir, - external_file_path, - file_offset, - tensor_byte_size)); + ORT_RETURN_IF_ERROR( + GetExternalDataInfo(tensor_proto, tensor_proto_dir, external_file_path, file_offset, tensor_byte_size)); unpacked_tensor.resize(tensor_byte_size); ORT_RETURN_IF_ERROR(onnxruntime::Env::Default().ReadFileIntoBuffer( @@ -229,12 +218,9 @@ Status ReadExternalDataForTensor(const ONNX_NAMESPACE::TensorProto& tensor_proto return Status::OK(); } -// TODO(unknown): Change the current interface to take Path object for model path -// so that validating and manipulating path for reading external data becomes easy -Status TensorProtoToOrtValueImpl(const Env& env, const ORTCHAR_T* model_path, - const ONNX_NAMESPACE::TensorProto& tensor_proto, - const MemBuffer* m, AllocatorPtr alloc, - OrtValue& value) { +Status TensorProtoToOrtValueImpl(const Env& env, const std::filesystem::path& model_path, + const ONNX_NAMESPACE::TensorProto& tensor_proto, const MemBuffer* m, + AllocatorPtr alloc, OrtValue& value) { if (m && m->GetBuffer() == nullptr) { return Status(common::ONNXRUNTIME, common::INVALID_ARGUMENT, "MemBuffer has not been allocated."); } @@ -276,7 +262,7 @@ namespace utils { #if !defined(ORT_MINIMAL_BUILD) static Status UnpackTensorWithExternalDataImpl(const ONNX_NAMESPACE::TensorProto& tensor, - const ORTCHAR_T* tensor_proto_dir, + const std::filesystem::path& tensor_proto_dir, size_t expected_num_elements, size_t element_size, /*out*/ unsigned char* p_data) { ORT_RETURN_IF(nullptr == p_data, "nullptr == p_data"); @@ -292,7 +278,7 @@ static Status UnpackTensorWithExternalDataImpl(const ONNX_NAMESPACE::TensorProto template Status UnpackTensorWithExternalData(const ONNX_NAMESPACE::TensorProto& tensor, - const ORTCHAR_T* tensor_proto_dir, size_t expected_num_elements, + const std::filesystem::path& tensor_proto_dir, size_t expected_num_elements, /*out*/ T* p_data) { static_assert(std::is_trivially_copyable::value, "T must be trivially copyable"); @@ -300,34 +286,35 @@ Status UnpackTensorWithExternalData(const ONNX_NAMESPACE::TensorProto& tensor, reinterpret_cast(p_data)); } -#define DEFINE_INT4_UNPACK_TENSOR_WITH_EXT_DATA_IMPL(INT4_TYPE) \ - template <> \ - Status UnpackTensorWithExternalData(const ONNX_NAMESPACE::TensorProto& tensor, \ - const ORTCHAR_T* tensor_proto_dir, size_t expected_num_elements, \ - /*out*/ INT4_TYPE* p_data) { \ - static_assert(std::is_trivially_copyable::value, "T must be trivially copyable"); \ - \ - ORT_RETURN_IF(nullptr == p_data, "nullptr == p_data"); \ - std::vector unpacked_tensor; \ - ORT_RETURN_IF_ERROR(ReadExternalDataForTensor(tensor, tensor_proto_dir, unpacked_tensor)); \ - \ - size_t num_packed_pairs = INT4_TYPE::CalcNumInt4Pairs(expected_num_elements); \ - ORT_RETURN_IF_NOT(num_packed_pairs == unpacked_tensor.size(), "Unexpected number of packed int4 pairs"); \ - \ - gsl::span src_span = gsl::make_span(reinterpret_cast(unpacked_tensor.data()), \ - num_packed_pairs); \ - gsl::span dst_span = gsl::make_span(p_data, expected_num_elements); \ - \ - std::memcpy(dst_span.data(), src_span.data(), num_packed_pairs); \ - \ - return Status::OK(); \ +#define DEFINE_INT4_UNPACK_TENSOR_WITH_EXT_DATA_IMPL(INT4_TYPE) \ + template <> \ + Status UnpackTensorWithExternalData(const ONNX_NAMESPACE::TensorProto& tensor, \ + const std::filesystem::path& tensor_proto_dir, \ + size_t expected_num_elements, /*out*/ INT4_TYPE* p_data) { \ + static_assert(std::is_trivially_copyable::value, "T must be trivially copyable"); \ + \ + ORT_RETURN_IF(nullptr == p_data, "nullptr == p_data"); \ + std::vector unpacked_tensor; \ + ORT_RETURN_IF_ERROR(ReadExternalDataForTensor(tensor, tensor_proto_dir, unpacked_tensor)); \ + \ + size_t num_packed_pairs = INT4_TYPE::CalcNumInt4Pairs(expected_num_elements); \ + ORT_RETURN_IF_NOT(num_packed_pairs == unpacked_tensor.size(), "Unexpected number of packed int4 pairs"); \ + \ + gsl::span src_span = \ + gsl::make_span(reinterpret_cast(unpacked_tensor.data()), num_packed_pairs); \ + gsl::span dst_span = gsl::make_span(p_data, expected_num_elements); \ + \ + std::memcpy(dst_span.data(), src_span.data(), num_packed_pairs); \ + \ + return Status::OK(); \ } DEFINE_INT4_UNPACK_TENSOR_WITH_EXT_DATA_IMPL(Int4x2) DEFINE_INT4_UNPACK_TENSOR_WITH_EXT_DATA_IMPL(UInt4x2) -#define INSTANTIATE_UNPACK_EXTERNAL_TENSOR(type) \ - template Status UnpackTensorWithExternalData(const ONNX_NAMESPACE::TensorProto&, const ORTCHAR_T*, size_t, type*); +#define INSTANTIATE_UNPACK_EXTERNAL_TENSOR(type) \ + template Status UnpackTensorWithExternalData(const ONNX_NAMESPACE::TensorProto&, const std::filesystem::path&, \ + size_t, type*); INSTANTIATE_UNPACK_EXTERNAL_TENSOR(float) INSTANTIATE_UNPACK_EXTERNAL_TENSOR(double) @@ -352,7 +339,7 @@ INSTANTIATE_UNPACK_EXTERNAL_TENSOR(Float8E5M2FNUZ) template <> Status UnpackTensorWithExternalData(const ONNX_NAMESPACE::TensorProto& /*tensor*/, - const ORTCHAR_T* /*tensor_proto_dir*/, size_t /*expected_num_elements*/, + const std::filesystem::path& /*tensor_proto_dir*/, size_t /*expected_num_elements*/, /*out*/ std::string* /*p_data*/) { return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "External data type cannot be STRING."); } @@ -369,7 +356,8 @@ Status UnpackTensor(const ONNX_NAMESPACE::TensorProto& tensor, const void* raw_d /*out*/ T* p_data, size_t expected_num_elements) { \ if (nullptr == p_data) { \ const size_t size = raw_data != nullptr ? raw_data_len : tensor.field_size(); \ - if (size == 0) return Status::OK(); \ + if (size == 0) \ + return Status::OK(); \ return Status(common::ONNXRUNTIME, common::INVALID_ARGUMENT); \ } \ if (nullptr == p_data || Type != tensor.data_type()) { \ @@ -379,9 +367,9 @@ Status UnpackTensor(const ONNX_NAMESPACE::TensorProto& tensor, const void* raw_d return UnpackTensorWithRawData(raw_data, raw_data_len, expected_num_elements, p_data); \ } \ if (static_cast(tensor.field_size()) != expected_num_elements) \ - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, \ - "corrupted protobuf data: tensor shape size(", expected_num_elements, \ - ") does not match the data size(", tensor.field_size(), ") in proto"); \ + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "corrupted protobuf data: tensor shape size(", \ + expected_num_elements, ") does not match the data size(", tensor.field_size(), \ + ") in proto"); \ auto& data = tensor.field_name(); \ for (auto data_iter = data.cbegin(); data_iter != data.cend(); ++data_iter) \ *p_data++ = static_cast(*data_iter); \ @@ -409,7 +397,8 @@ template <> Status UnpackTensor(const ONNX_NAMESPACE::TensorProto& tensor, const void* /*raw_data*/, size_t /*raw_data_len*/, /*out*/ std::string* p_data, size_t expected_size) { if (nullptr == p_data) { - if (tensor.string_data_size() == 0) return Status::OK(); + if (tensor.string_data_size() == 0) + return Status::OK(); return Status(common::ONNXRUNTIME, common::INVALID_ARGUMENT); } if (ONNX_NAMESPACE::TensorProto_DataType_STRING != tensor.data_type()) { @@ -434,7 +423,8 @@ Status UnpackTensor(const ONNX_NAMESPACE::TensorProto& tensor, const void* raw_d /*out*/ bool* p_data, size_t expected_size) { if (nullptr == p_data) { const size_t size = raw_data != nullptr ? raw_data_len : tensor.int32_data_size(); - if (size == 0) return Status::OK(); + if (size == 0) + return Status::OK(); return Status(common::ONNXRUNTIME, common::INVALID_ARGUMENT); } if (ONNX_NAMESPACE::TensorProto_DataType_BOOL != tensor.data_type()) { @@ -461,7 +451,8 @@ Status UnpackTensor(const ONNX_NAMESPACE::TensorProto& tensor, const void* raw_d /*out*/ MLFloat16* p_data, size_t expected_size) { if (nullptr == p_data) { const size_t size = raw_data != nullptr ? raw_data_len : tensor.int32_data_size(); - if (size == 0) return Status::OK(); + if (size == 0) + return Status::OK(); return Status(common::ONNXRUNTIME, common::INVALID_ARGUMENT); } if (ONNX_NAMESPACE::TensorProto_DataType_FLOAT16 != tensor.data_type()) { @@ -705,15 +696,12 @@ DEFINE_INT4_UNPACK_TENSOR_IMPL(UInt4x2, TensorProto_DataType_UINT4) // Uses the model path to construct the full path for loading external data. In case when model_path is empty // it uses current directory. template -Status UnpackTensor(const ONNX_NAMESPACE::TensorProto& tensor, const Path& model_path, +Status UnpackTensor(const ONNX_NAMESPACE::TensorProto& tensor, const std::filesystem::path& model_path, /*out*/ T* p_data, size_t expected_num_elements) { #if !defined(ORT_MINIMAL_BUILD) if (HasExternalData(tensor)) { - return UnpackTensorWithExternalData( - tensor, - model_path.IsEmpty() ? nullptr : model_path.ParentPath().ToPathString().c_str(), - expected_num_elements, - p_data); + return UnpackTensorWithExternalData(tensor, model_path.parent_path(), + expected_num_elements, p_data); } #else ORT_UNUSED_PARAMETER(model_path); @@ -727,7 +715,7 @@ Status UnpackTensor(const ONNX_NAMESPACE::TensorProto& tensor, const Path& model // instantiate the UnpackTensor variant that supports external data #define INSTANTIATE_UNPACK_TENSOR(type) \ - template Status UnpackTensor(const ONNX_NAMESPACE::TensorProto&, const Path&, type* p_data, size_t); + template Status UnpackTensor(const ONNX_NAMESPACE::TensorProto&, const std::filesystem::path&, type* p_data, size_t); INSTANTIATE_UNPACK_TENSOR(float) INSTANTIATE_UNPACK_TENSOR(double) @@ -812,8 +800,8 @@ TensorShape GetTensorShapeFromTensorShapeProto(const ONNX_NAMESPACE::TensorShape const auto& dims = tensor_shape_proto.dim(); std::vector tensor_shape_vec(static_cast(dims.size())); for (int i = 0; i < dims.size(); ++i) { - tensor_shape_vec[i] = HasDimValue(dims[i]) ? dims[i].dim_value() - : -1; /* symbolic dimensions are represented as -1 in onnxruntime*/ + tensor_shape_vec[i] = + HasDimValue(dims[i]) ? dims[i].dim_value() : -1; /* symbolic dimensions are represented as -1 in onnxruntime*/ } return TensorShape(std::move(tensor_shape_vec)); } @@ -838,7 +826,8 @@ ORT_API_STATUS_IMPL(OrtInitializeBufferForTensor, _In_opt_ void* input, size_t i enum ONNXTensorElementDataType type) { OrtStatus* status = nullptr; ORT_TRY { - if (type != ONNX_TENSOR_ELEMENT_DATA_TYPE_STRING || input == nullptr) return nullptr; + if (type != ONNX_TENSOR_ELEMENT_DATA_TYPE_STRING || input == nullptr) + return nullptr; size_t tensor_size = input_len / sizeof(std::string); std::string* ptr = reinterpret_cast(input); for (size_t i = 0, n = tensor_size; i < n; ++i) { @@ -846,16 +835,15 @@ ORT_API_STATUS_IMPL(OrtInitializeBufferForTensor, _In_opt_ void* input, size_t i } } ORT_CATCH(const std::exception& ex) { - ORT_HANDLE_EXCEPTION([&]() { - status = OrtApis::CreateStatus(ORT_RUNTIME_EXCEPTION, ex.what()); - }); + ORT_HANDLE_EXCEPTION([&]() { status = OrtApis::CreateStatus(ORT_RUNTIME_EXCEPTION, ex.what()); }); } return status; } ORT_API(void, OrtUninitializeBuffer, _In_opt_ void* input, size_t input_len, enum ONNXTensorElementDataType type) { - if (type != ONNX_TENSOR_ELEMENT_DATA_TYPE_STRING || input == nullptr) return; + if (type != ONNX_TENSOR_ELEMENT_DATA_TYPE_STRING || input == nullptr) + return; size_t tensor_size = input_len / sizeof(std::string); std::string* ptr = reinterpret_cast(input); using std::string; @@ -884,18 +872,18 @@ static void DeleteCharArray(void* param) noexcept { } #if !defined(__wasm__) -static Status GetFileContent( - const Env& env, const ORTCHAR_T* file_path, FileOffsetType offset, size_t length, - void*& raw_buffer, OrtCallback& deleter) { +static Status GetFileContent(const Env& env, const std::filesystem::path& file_path, FileOffsetType offset, + size_t length, void*& raw_buffer, OrtCallback& deleter) { // query length if it is 0 if (length == 0) { - ORT_RETURN_IF_ERROR(env.GetFileLength(file_path, length)); + // The return type of std::filesystem::file_size is uintmax_t which could be bigger than size_t + length = narrow(std::filesystem::file_size(file_path)); } // first, try to map into memory { Env::MappedMemoryPtr mapped_memory{}; - auto status = env.MapFileIntoMemory(file_path, offset, length, mapped_memory); + auto status = env.MapFileIntoMemory(file_path.native().c_str(), offset, length, mapped_memory); if (status.IsOK()) { deleter = mapped_memory.get_deleter().callback; raw_buffer = mapped_memory.release(); @@ -905,8 +893,8 @@ static Status GetFileContent( // if that fails, try to copy auto buffer = std::make_unique(length); - ORT_RETURN_IF_ERROR(env.ReadFileIntoBuffer( - file_path, offset, length, gsl::make_span(buffer.get(), length))); + ORT_RETURN_IF_ERROR( + env.ReadFileIntoBuffer(file_path.native().c_str(), offset, length, gsl::make_span(buffer.get(), length))); deleter = OrtCallback{DeleteCharArray, buffer.get()}; raw_buffer = buffer.release(); @@ -914,20 +902,19 @@ static Status GetFileContent( } #endif -Status GetExtDataFromTensorProto(const Env& env, const ORTCHAR_T* model_path, - const ONNX_NAMESPACE::TensorProto& tensor_proto, - void*& ext_data_buf, SafeInt& ext_data_len, OrtCallback& ext_data_deleter) { +Status GetExtDataFromTensorProto(const Env& env, const std::filesystem::path& model_path, + const ONNX_NAMESPACE::TensorProto& tensor_proto, void*& ext_data_buf, + SafeInt& ext_data_len, OrtCallback& ext_data_deleter) { ORT_ENFORCE(utils::HasExternalData(tensor_proto)); std::basic_string tensor_proto_dir; - if (model_path != nullptr) { + if (!model_path.empty()) { ORT_RETURN_IF_ERROR(GetDirNameFromFilePath(model_path, tensor_proto_dir)); } - const ORTCHAR_T* t_prot_dir_s = tensor_proto_dir.size() == 0 ? nullptr : tensor_proto_dir.c_str(); std::basic_string external_data_file_path; FileOffsetType file_offset; SafeInt raw_data_safe_len = 0; - ORT_RETURN_IF_ERROR(GetExternalDataInfo(tensor_proto, t_prot_dir_s, external_data_file_path, file_offset, - raw_data_safe_len)); + ORT_RETURN_IF_ERROR( + GetExternalDataInfo(tensor_proto, tensor_proto_dir, external_data_file_path, file_offset, raw_data_safe_len)); if (external_data_file_path == onnxruntime::utils::kTensorProtoMemoryAddressTag) { // the value in location is the memory address of the data @@ -937,8 +924,8 @@ Status GetExtDataFromTensorProto(const Env& env, const ORTCHAR_T* model_path, } else { #if defined(__wasm__) ORT_RETURN_IF(file_offset < 0 || file_offset + raw_data_safe_len >= 4294967296, - "External initializer: ", tensor_proto.name(), - " offset: ", file_offset, " size to read: ", static_cast(raw_data_safe_len), + "External initializer: ", tensor_proto.name(), " offset: ", file_offset, + " size to read: ", static_cast(raw_data_safe_len), " are out of bounds or can not be read in full (>4GB)."); auto buffer = std::make_unique(raw_data_safe_len); @@ -969,7 +956,8 @@ Status GetExtDataFromTensorProto(const Env& env, const ORTCHAR_T* model_path, } try { - // Copy the file data (fileData,offset,length) into WebAssembly memory (HEAPU8,buffer,length). + // Copy the file data (fileData,offset,length) into WebAssembly memory + // (HEAPU8,buffer,length). HEAPU8.set(fileData.subarray(offset, offset + length), buffer); return 0; } catch { @@ -996,22 +984,19 @@ Status GetExtDataFromTensorProto(const Env& env, const ORTCHAR_T* model_path, default: err_msg = "Unknown error occurred in memory copy."; } - return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Failed to load external data file \"", external_data_file_path, "\", error: ", err_msg); + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Failed to load external data file \"", external_data_file_path, + "\", error: ", err_msg); #else - size_t file_length; - // error reporting is inconsistent across platforms. Make sure the full path we attempted to open is included. - auto status = env.GetFileLength(external_data_file_path.c_str(), file_length); - if (!status.IsOK()) { - return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "GetFileLength for ", ToUTF8String(external_data_file_path), - " failed:", status.ErrorMessage()); - } + // The GetFileContent function doesn't report error if the requested data range is invalid. Therefore we need to + // manually check file size first. + std::uintmax_t file_length = std::filesystem::file_size(external_data_file_path); SafeInt end_of_read(file_offset); end_of_read += raw_data_safe_len; - ORT_RETURN_IF(file_offset < 0 || end_of_read > narrow(file_length), - "External initializer: ", tensor_proto.name(), - " offset: ", file_offset, " size to read: ", static_cast(raw_data_safe_len), - " given file_length: ", file_length, " are out of bounds or can not be read in full."); + ORT_RETURN_IF(file_offset < 0 || static_cast(end_of_read) > file_length, + "External initializer: ", tensor_proto.name(), " offset: ", file_offset, + " size to read: ", static_cast(raw_data_safe_len), " given file_length: ", file_length, + " are out of bounds or can not be read in full."); ORT_RETURN_IF_ERROR(GetFileContent(env, external_data_file_path.c_str(), file_offset, raw_data_safe_len, ext_data_buf, ext_data_deleter)); ext_data_len = raw_data_safe_len; @@ -1021,11 +1006,10 @@ Status GetExtDataFromTensorProto(const Env& env, const ORTCHAR_T* model_path, return Status::OK(); } -#define CASE_PROTO(X, Y) \ - case ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_##X: \ - ORT_RETURN_IF_ERROR( \ - UnpackTensor(tensor_proto, raw_data, raw_data_len, \ - (Y*)preallocated, static_cast(tensor_size))); \ +#define CASE_PROTO(X, Y) \ + case ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_##X: \ + ORT_RETURN_IF_ERROR( \ + UnpackTensor(tensor_proto, raw_data, raw_data_len, (Y*)preallocated, static_cast(tensor_size))); \ break; /** @@ -1036,15 +1020,15 @@ Status GetExtDataFromTensorProto(const Env& env, const ORTCHAR_T* model_path, * @param tensor pre-allocated tensor object, where we store the data * @return */ -Status TensorProtoToTensor(const Env& env, const ORTCHAR_T* model_path, - const ONNX_NAMESPACE::TensorProto& tensor_proto, - Tensor& tensor) { +Status TensorProtoToTensor(const Env& env, const std::filesystem::path& model_path, + const ONNX_NAMESPACE::TensorProto& tensor_proto, Tensor& tensor) { // Validate tensor compatibility TensorShape tensor_shape = GetTensorShapeFromTensorProto(tensor_proto); if (tensor_shape != tensor.Shape()) { return Status(common::ONNXRUNTIME, common::INVALID_ARGUMENT, "TensorProtoToTensor() tensor shape mismatch!"); } - const DataTypeImpl* const source_type = DataTypeImpl::TensorTypeFromONNXEnum(tensor_proto.data_type())->GetElementType(); + const DataTypeImpl* const source_type = + DataTypeImpl::TensorTypeFromONNXEnum(tensor_proto.data_type())->GetElementType(); if (source_type->Size() > tensor.DataType()->Size()) { return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "TensorProto type ", DataTypeImpl::ToString(source_type), " can not be written into Tensor type ", DataTypeImpl::ToString(tensor.DataType())); @@ -1125,15 +1109,13 @@ Status TensorProtoToTensor(const Env& env, const ORTCHAR_T* model_path, return Status::OK(); } -Status TensorProtoToOrtValue(const Env& env, const ORTCHAR_T* model_path, - const ONNX_NAMESPACE::TensorProto& tensor_proto, - const MemBuffer& m, OrtValue& value) { +Status TensorProtoToOrtValue(const Env& env, const std::filesystem::path& model_path, + const ONNX_NAMESPACE::TensorProto& tensor_proto, const MemBuffer& m, OrtValue& value) { return TensorProtoToOrtValueImpl(env, model_path, tensor_proto, &m, nullptr, value); } -Status TensorProtoToOrtValue(const Env& env, const ORTCHAR_T* model_path, - const ONNX_NAMESPACE::TensorProto& tensor_proto, - AllocatorPtr alloc, OrtValue& value) { +Status TensorProtoToOrtValue(const Env& env, const std::filesystem::path& model_path, + const ONNX_NAMESPACE::TensorProto& tensor_proto, AllocatorPtr alloc, OrtValue& value) { return TensorProtoToOrtValueImpl(env, model_path, tensor_proto, nullptr, alloc, value); } @@ -1207,7 +1189,7 @@ ONNX_NAMESPACE::TensorProto TensorToTensorProto(const Tensor& tensor, const std: } common::Status ConstantNodeProtoToTensorProto(const ONNX_NAMESPACE::NodeProto& node, - const Path& model_path, + const std::filesystem::path& model_path, ONNX_NAMESPACE::TensorProto& tensor, const std::string& tensor_name) { ORT_RETURN_IF_NOT(node.attribute_size() > 0, "Constant node: ", node.name(), " has no data attributes"); @@ -1255,8 +1237,8 @@ common::Status ConstantNodeProtoToTensorProto(const ONNX_NAMESPACE::NodeProto& n ORT_UNUSED_PARAMETER(model_path); #endif default: - ORT_THROW("Unsupported attribute value type of ", constant_attribute.type(), - " in 'Constant' node '", node.name(), "'"); + ORT_THROW("Unsupported attribute value type of ", constant_attribute.type(), " in 'Constant' node '", node.name(), + "'"); } // set name last in case attribute type was tensor (would copy over name) @@ -1266,7 +1248,7 @@ common::Status ConstantNodeProtoToTensorProto(const ONNX_NAMESPACE::NodeProto& n } common::Status ConstantNodeProtoToTensorProto(const ONNX_NAMESPACE::NodeProto& node, - const Path& model_path, + const std::filesystem::path& model_path, ONNX_NAMESPACE::TensorProto& tensor) { return ConstantNodeProtoToTensorProto(node, model_path, tensor, node.output(0)); } @@ -1274,9 +1256,11 @@ common::Status ConstantNodeProtoToTensorProto(const ONNX_NAMESPACE::NodeProto& n #if !defined(DISABLE_SPARSE_TENSORS) static Status CopySparseData(size_t n_sparse_elements, const ONNX_NAMESPACE::TensorProto& indices, - const Path& model_path, - gsl::span dims, - std::function copier) { + const std::filesystem::path& model_path, + gsl::span + dims, + std::function + copier) { Status status = Status::OK(); TensorShape indices_shape(indices.dims().data(), indices.dims().size()); const auto elements = narrow(indices_shape.Size()); @@ -1293,7 +1277,8 @@ static Status CopySparseData(size_t n_sparse_elements, ORT_RETURN_IF_ERROR(UnpackInitializerData(indices, model_path, unpack_buffer)); indices_data = ReinterpretAsSpan(gsl::make_span(unpack_buffer)); } else { - ORT_RETURN_IF_NOT(indices.int64_data_size() == static_cast(elements), "Sparse indices int64 data size does not match expected"); + ORT_RETURN_IF_NOT(indices.int64_data_size() == static_cast(elements), + "Sparse indices int64 data size does not match expected"); indices_data = gsl::make_span(indices.int64_data().data(), elements); } break; @@ -1307,7 +1292,8 @@ static Status CopySparseData(size_t n_sparse_elements, unpack_buffer.clear(); unpack_buffer.shrink_to_fit(); } else { - ORT_RETURN_IF_NOT(indices.int32_data_size() == static_cast(elements), "Sparse indices int32 data size does not match expected"); + ORT_RETURN_IF_NOT(indices.int32_data_size() == static_cast(elements), + "Sparse indices int32 data size does not match expected"); indices_values.insert(indices_values.cend(), indices.int32_data().cbegin(), indices.int32_data().cend()); } indices_data = gsl::make_span(indices_values); @@ -1346,8 +1332,9 @@ static Status CopySparseData(size_t n_sparse_elements, break; } default: - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_GRAPH, - "Invalid SparseTensor indices. Should one of the following types: int8, int16, int32 or int64"); + return ORT_MAKE_STATUS( + ONNXRUNTIME, INVALID_GRAPH, + "Invalid SparseTensor indices. Should one of the following types: int8, int16, int32 or int64"); } if (indices_shape.NumDimensions() == 1) { @@ -1385,15 +1372,15 @@ static Status CopySparseData(size_t n_sparse_elements, ORT_ENFORCE(cur_index == indices_data.end()); } else { - status = ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_GRAPH, "Invalid SparseTensor indices. Should be rank 0 or 1. Got:", - indices_shape); + status = ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_GRAPH, + "Invalid SparseTensor indices. Should be rank 0 or 1. Got:", indices_shape); } return status; } common::Status SparseTensorProtoToDenseTensorProto(const ONNX_NAMESPACE::SparseTensorProto& sparse, - const Path& model_path, + const std::filesystem::path& model_path, ONNX_NAMESPACE::TensorProto& dense) { Status status = Status::OK(); @@ -1434,53 +1421,45 @@ common::Status SparseTensorProtoToDenseTensorProto(const ONNX_NAMESPACE::SparseT switch (element_size) { case 1: { status = CopySparseData( - n_sparse_elements, - indices, model_path, dims, - [sparse_data, dense_data](size_t from_idx, size_t to_idx) { + n_sparse_elements, indices, model_path, dims, [sparse_data, dense_data](size_t from_idx, size_t to_idx) { static_cast(dense_data)[to_idx] = static_cast(sparse_data)[from_idx]; }); break; } case 2: { - status = CopySparseData( - n_sparse_elements, - indices, model_path, dims, - [sparse_data, dense_data](size_t from_idx, size_t to_idx) { - const auto* src = static_cast(sparse_data) + from_idx; - auto* dst = static_cast(dense_data) + to_idx; - memcpy(dst, src, sizeof(uint16_t)); - }); + status = CopySparseData(n_sparse_elements, indices, model_path, dims, + [sparse_data, dense_data](size_t from_idx, size_t to_idx) { + const auto* src = static_cast(sparse_data) + from_idx; + auto* dst = static_cast(dense_data) + to_idx; + memcpy(dst, src, sizeof(uint16_t)); + }); break; } case 4: { - status = CopySparseData( - n_sparse_elements, - indices, model_path, dims, - [sparse_data, dense_data](size_t from_idx, size_t to_idx) { - const auto* src = static_cast(sparse_data) + from_idx; - auto* dst = static_cast(dense_data) + to_idx; - memcpy(dst, src, sizeof(uint32_t)); - }); + status = CopySparseData(n_sparse_elements, indices, model_path, dims, + [sparse_data, dense_data](size_t from_idx, size_t to_idx) { + const auto* src = static_cast(sparse_data) + from_idx; + auto* dst = static_cast(dense_data) + to_idx; + memcpy(dst, src, sizeof(uint32_t)); + }); break; } case 8: { - status = CopySparseData( - n_sparse_elements, - indices, model_path, dims, - [sparse_data, dense_data](size_t from_idx, size_t to_idx) { - const auto* src = static_cast(sparse_data) + from_idx; - auto* dst = static_cast(dense_data) + to_idx; - memcpy(dst, src, sizeof(uint64_t)); - }); + status = CopySparseData(n_sparse_elements, indices, model_path, dims, + [sparse_data, dense_data](size_t from_idx, size_t to_idx) { + const auto* src = static_cast(sparse_data) + from_idx; + auto* dst = static_cast(dense_data) + to_idx; + memcpy(dst, src, sizeof(uint64_t)); + }); break; } default: - return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, - "Element_size of: ", element_size, " is not supported.", " type: ", type); + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Element_size of: ", element_size, " is not supported.", + " type: ", type); } ORT_RETURN_IF_ERROR(status); @@ -1521,9 +1500,7 @@ inline void CopyElement(void* dst, const void* src, int64_t dst_index, } template -static void SetIndices(gsl::span gathered_indices, - std::string& raw_indices, - TensorProto& indices) { +static void SetIndices(gsl::span gathered_indices, std::string& raw_indices, TensorProto& indices) { raw_indices.resize(gathered_indices.size() * sizeof(T)); auto* ind_dest = reinterpret_cast(raw_indices.data()); size_t dest_index = 0; @@ -1541,8 +1518,7 @@ static void SetIndices(gsl::span gathered_indices, } static void SparsifyGeneric(const void* dense_raw_data, size_t n_dense_elements, size_t element_size, - IsZeroFunc is_zero, CopyElementFunc copy, - TensorProto& values, TensorProto& indices, + IsZeroFunc is_zero, CopyElementFunc copy, TensorProto& values, TensorProto& indices, size_t& nnz) { auto advance = [element_size](const void* start, size_t elements) -> const void* { return (reinterpret_cast(start) + elements * element_size); @@ -1591,7 +1567,7 @@ static void SparsifyGeneric(const void* dense_raw_data, size_t n_dense_elements, } common::Status DenseTensorToSparseTensorProto(const ONNX_NAMESPACE::TensorProto& dense_proto, - const Path& model_path, + const std::filesystem::path& model_path, ONNX_NAMESPACE::SparseTensorProto& result) { ORT_ENFORCE(HasDataType(dense_proto), "Must have a valid data type"); @@ -1623,28 +1599,28 @@ common::Status DenseTensorToSparseTensorProto(const ONNX_NAMESPACE::TensorProto& void* dense_data = dense_raw_data.data(); switch (element_size) { case 1: { - SparsifyGeneric(dense_data, n_dense_elements, element_size, - IsZero, CopyElement, values, indices, nnz); + SparsifyGeneric(dense_data, n_dense_elements, element_size, IsZero, CopyElement, values, + indices, nnz); break; } case 2: { - SparsifyGeneric(dense_data, n_dense_elements, element_size, - IsZero, CopyElement, values, indices, nnz); + SparsifyGeneric(dense_data, n_dense_elements, element_size, IsZero, CopyElement, values, + indices, nnz); break; } case 4: { - SparsifyGeneric(dense_data, n_dense_elements, element_size, - IsZero, CopyElement, values, indices, nnz); + SparsifyGeneric(dense_data, n_dense_elements, element_size, IsZero, CopyElement, values, + indices, nnz); break; } case 8: { - SparsifyGeneric(dense_data, n_dense_elements, element_size, - IsZero, CopyElement, values, indices, nnz); + SparsifyGeneric(dense_data, n_dense_elements, element_size, IsZero, CopyElement, values, + indices, nnz); break; } default: - return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, - "Element_size of: ", element_size, " is not supported.", " data_type: ", data_type); + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Element_size of: ", element_size, " is not supported.", + " data_type: ", data_type); } // Fix up shapes @@ -1664,42 +1640,40 @@ template common::Status GetSizeInBytesFromTensorProto(const ONN size_t* out); template common::Status GetSizeInBytesFromTensorProto<0>(const ONNX_NAMESPACE::TensorProto& tensor_proto, size_t* out); -#define CASE_UNPACK(TYPE, ELEMENT_TYPE, DATA_SIZE) \ - case ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_##TYPE: { \ - SafeInt tensor_byte_size; \ - size_t element_count = 0; \ - if (initializer.has_raw_data()) { \ - tensor_byte_size = initializer.raw_data().size(); \ - element_count = tensor_byte_size / sizeof(ELEMENT_TYPE); \ - } else { \ - element_count = initializer.DATA_SIZE(); \ - tensor_byte_size = element_count * sizeof(ELEMENT_TYPE); \ - } \ - unpacked_tensor.resize(tensor_byte_size); \ - return onnxruntime::utils::UnpackTensor( \ - initializer, \ - initializer.has_raw_data() ? initializer.raw_data().data() : nullptr, \ - initializer.has_raw_data() ? initializer.raw_data().size() : 0, \ - reinterpret_cast(unpacked_tensor.data()), element_count); \ - break; \ - } - -#define CASE_UNPACK_INT4(TYPE, ELEMENT_TYPE, DATA_SIZE) \ - case ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_##TYPE: { \ - TensorShape tensor_shape = GetTensorShapeFromTensorProto(initializer); \ - size_t element_count = static_cast(tensor_shape.Size()); \ - size_t packed_element_count = ELEMENT_TYPE::CalcNumInt4Pairs(element_count); \ - unpacked_tensor.resize(packed_element_count * sizeof(ELEMENT_TYPE)); \ - return onnxruntime::utils::UnpackTensor( \ - initializer, \ - initializer.has_raw_data() ? initializer.raw_data().data() : nullptr, \ - initializer.has_raw_data() ? initializer.raw_data().size() : 0, \ - reinterpret_cast(unpacked_tensor.data()), element_count); \ - break; \ +#define CASE_UNPACK(TYPE, ELEMENT_TYPE, DATA_SIZE) \ + case ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_##TYPE: { \ + SafeInt tensor_byte_size; \ + size_t element_count = 0; \ + if (initializer.has_raw_data()) { \ + tensor_byte_size = initializer.raw_data().size(); \ + element_count = tensor_byte_size / sizeof(ELEMENT_TYPE); \ + } else { \ + element_count = initializer.DATA_SIZE(); \ + tensor_byte_size = element_count * sizeof(ELEMENT_TYPE); \ + } \ + unpacked_tensor.resize(tensor_byte_size); \ + return onnxruntime::utils::UnpackTensor(initializer, \ + initializer.has_raw_data() ? initializer.raw_data().data() : nullptr, \ + initializer.has_raw_data() ? initializer.raw_data().size() : 0, \ + reinterpret_cast(unpacked_tensor.data()), element_count); \ + break; \ + } + +#define CASE_UNPACK_INT4(TYPE, ELEMENT_TYPE, DATA_SIZE) \ + case ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_##TYPE: { \ + TensorShape tensor_shape = GetTensorShapeFromTensorProto(initializer); \ + size_t element_count = static_cast(tensor_shape.Size()); \ + size_t packed_element_count = ELEMENT_TYPE::CalcNumInt4Pairs(element_count); \ + unpacked_tensor.resize(packed_element_count * sizeof(ELEMENT_TYPE)); \ + return onnxruntime::utils::UnpackTensor(initializer, \ + initializer.has_raw_data() ? initializer.raw_data().data() : nullptr, \ + initializer.has_raw_data() ? initializer.raw_data().size() : 0, \ + reinterpret_cast(unpacked_tensor.data()), element_count); \ + break; \ } Status UnpackInitializerData(const onnx::TensorProto& initializer, - const Path& model_path, + const std::filesystem::path& model_path, std::vector& unpacked_tensor) { // TODO, if std::vector does not use a custom allocator, the default std::allocator will // allocation the memory aligned to std::max_align_t, need look into allocating @@ -1707,7 +1681,7 @@ Status UnpackInitializerData(const onnx::TensorProto& initializer, if (initializer.data_location() == TensorProto_DataLocation_EXTERNAL) { ORT_RETURN_IF_ERROR(ReadExternalDataForTensor( initializer, - (model_path.IsEmpty() || model_path.ParentPath().IsEmpty()) ? nullptr : model_path.ParentPath().ToPathString().c_str(), + model_path.parent_path(), unpacked_tensor)); return Status::OK(); } @@ -1737,16 +1711,14 @@ Status UnpackInitializerData(const onnx::TensorProto& initializer, default: break; } - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "Unsupported type: ", initializer.data_type()); + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Unsupported type: ", initializer.data_type()); } #undef CASE_UNPACK -Status UnpackInitializerData(const ONNX_NAMESPACE::TensorProto& initializer, - std::vector& unpacked_tensor) { +Status UnpackInitializerData(const ONNX_NAMESPACE::TensorProto& initializer, std::vector& unpacked_tensor) { ORT_RETURN_IF(initializer.data_location() == TensorProto_DataLocation_EXTERNAL, "The given initializer contains external data"); - return UnpackInitializerData(initializer, Path(), unpacked_tensor); + return UnpackInitializerData(initializer, std::filesystem::path(), unpacked_tensor); } } // namespace utils diff --git a/onnxruntime/core/framework/tensorprotoutils.h b/onnxruntime/core/framework/tensorprotoutils.h index 000502ba47594..2f3f942e75578 100644 --- a/onnxruntime/core/framework/tensorprotoutils.h +++ b/onnxruntime/core/framework/tensorprotoutils.h @@ -5,6 +5,7 @@ #include #include +#include #ifndef SHARED_PROVIDER #include "core/common/common.h" @@ -41,23 +42,23 @@ TensorShape GetTensorShapeFromTensorProto(const ONNX_NAMESPACE::TensorProto& ten /** * deserialize a TensorProto into a preallocated memory buffer on CPU. * \param tensor_proto_path A local file path of where the 'input' was loaded from. - * Can be NULL if the tensor proto doesn't have external data or it was loaded from + * Can be empty if the tensor proto doesn't have external data or it was loaded from * the current working dir. This path could be either a relative path or an absolute path. * \return Status::OK on success with 'value' containing the Tensor in CPU based memory. */ -common::Status TensorProtoToOrtValue(const Env& env, const ORTCHAR_T* tensor_proto_path, +common::Status TensorProtoToOrtValue(const Env& env, const std::filesystem::path& tensor_proto_path, const ONNX_NAMESPACE::TensorProto& input, const MemBuffer& m, OrtValue& value); /** * deserialize a TensorProto into a buffer on CPU allocated using 'alloc'. * \param tensor_proto_path A local file path of where the 'input' was loaded from. - * Can be NULL if the tensor proto doesn't have external data or it was loaded from + * Can be empty if the tensor proto doesn't have external data or it was loaded from * the current working dir. This path could be either a relative path or an absolute path. * \param alloc Allocator to use for allocating the buffer. Must allocate CPU based memory. * \return Status::OK on success with 'value' containing the Tensor in CPU based memory. */ -common::Status TensorProtoToOrtValue(const Env& env, const ORTCHAR_T* tensor_proto_path, +common::Status TensorProtoToOrtValue(const Env& env, const std::filesystem::path& tensor_proto_path, const ONNX_NAMESPACE::TensorProto& input, AllocatorPtr alloc, OrtValue& value); @@ -69,7 +70,7 @@ common::Status TensorProtoToOrtValue(const Env& env, const ORTCHAR_T* tensor_pro * @param tensorp destination empty tensor * @return */ -common::Status TensorProtoToTensor(const Env& env, const ORTCHAR_T* model_path, +common::Status TensorProtoToTensor(const Env& env, const std::filesystem::path& model_path, const ONNX_NAMESPACE::TensorProto& tensor_proto, Tensor& tensor); @@ -100,7 +101,7 @@ constexpr const ORTCHAR_T* kTensorProtoMemoryAddressTag = ORT_TSTR("*/_ORT_MEM_A // Given a tensor proto with external data obtain a pointer to the data and its length. // The ext_data_deleter argument is updated with a callback that owns/releases the data. -common::Status GetExtDataFromTensorProto(const Env& env, const ORTCHAR_T* model_path, +common::Status GetExtDataFromTensorProto(const Env& env, const std::filesystem::path& model_path, const ONNX_NAMESPACE::TensorProto& tensor_proto, void*& ext_data_buf, SafeInt& ext_data_len, OrtCallback& ext_data_deleter); @@ -113,11 +114,11 @@ common::Status GetExtDataFromTensorProto(const Env& env, const ORTCHAR_T* model_ // model_path is used for contructing full path for external_data // tensor_name specifies the name for the new TensorProto TensorProto common::Status ConstantNodeProtoToTensorProto(const ONNX_NAMESPACE::NodeProto& node, - const Path& model_path, + const std::filesystem::path& model_path, ONNX_NAMESPACE::TensorProto& tensor, const std::string& tensor_name); common::Status ConstantNodeProtoToTensorProto(const ONNX_NAMESPACE::NodeProto& node, - const Path& model_path, + const std::filesystem::path& model_path, ONNX_NAMESPACE::TensorProto& tensor); #if !defined(DISABLE_SPARSE_TENSORS) @@ -126,7 +127,7 @@ common::Status ConstantNodeProtoToTensorProto(const ONNX_NAMESPACE::NodeProto& n // The resulting TensorProto will contain the data as raw data. // model_path is used for contructing full path for external_data common::Status SparseTensorProtoToDenseTensorProto(const ONNX_NAMESPACE::SparseTensorProto& sparse, - const Path& model_path, + const std::filesystem::path& model_path, ONNX_NAMESPACE::TensorProto& dense); #if !defined(ORT_MINIMAL_BUILD) @@ -135,7 +136,7 @@ common::Status SparseTensorProtoToDenseTensorProto(const ONNX_NAMESPACE::SparseT // The resulting SparseTensorProto will contain the data as raw data // model_path is used for contructing full path for external_data common::Status DenseTensorToSparseTensorProto(const ONNX_NAMESPACE::TensorProto& dense, - const Path& model_path, + const std::filesystem::path& model_path, ONNX_NAMESPACE::SparseTensorProto& sparse); #endif // !ORT_MINIMAL_BUILD #endif // !defined(DISABLE_SPARSE_TENSORS) @@ -446,7 +447,7 @@ Status UnpackTensor(const ONNX_NAMESPACE::TensorProto& tensor, const void* raw_d // Uses the model path to construct the full path for loading external data. In case when model_path is empty // it uses current directory. template -Status UnpackTensor(const ONNX_NAMESPACE::TensorProto& tensor, const Path& model_path, +Status UnpackTensor(const ONNX_NAMESPACE::TensorProto& tensor, const std::filesystem::path& model_path, /*out*/ T* p_data, size_t expected_size); /** @@ -458,7 +459,7 @@ Status UnpackTensor(const ONNX_NAMESPACE::TensorProto& tensor, const Path& model * @returns Status::OK() if data is unpacked successfully */ common::Status UnpackInitializerData(const ONNX_NAMESPACE::TensorProto& initializer, - const Path& model_path, + const std::filesystem::path& model_path, std::vector& unpacked_tensor); /** diff --git a/onnxruntime/core/graph/graph.cc b/onnxruntime/core/graph/graph.cc index 0c1d79532f120..67451301023e5 100644 --- a/onnxruntime/core/graph/graph.cc +++ b/onnxruntime/core/graph/graph.cc @@ -560,7 +560,7 @@ void Node::SetPriority(int priority) noexcept { priority_ = priority; } -const Path& Node::ModelPath() const noexcept { +const std::filesystem::path& Node::ModelPath() const noexcept { return graph_->ModelPath(); } @@ -3025,7 +3025,8 @@ Status Graph::VerifyNodeAndOpMatch(const ResolveOptions& options) { ctx.set_opset_imports(DomainToVersionMap()); ctx.set_schema_registry(schema_registry_.get()); // Set the parent directory of model path to load external tensors if exist - ctx.set_model_dir(ToUTF8String(ModelPath().ParentPath().ToPathString())); + // ONNX expects a UTF-8 string here. + ctx.set_model_dir(ToUTF8String(ModelPath().parent_path().native())); LexicalScopeContext parent; if (parent_node_) { @@ -3370,7 +3371,7 @@ const std::string& Graph::Description() const noexcept { return graph_proto_->doc_string(); } -const Path& Graph::ModelPath() const { +const std::filesystem::path& Graph::ModelPath() const { return owning_model_.ModelPath(); } @@ -3971,21 +3972,17 @@ ONNX_NAMESPACE::GraphProto Graph::ToGraphProto() const { return result; } -ONNX_NAMESPACE::GraphProto Graph::ToGraphProtoWithExternalInitializers(const std::string& external_file_name, - const PathString& destination_file_path, +ONNX_NAMESPACE::GraphProto Graph::ToGraphProtoWithExternalInitializers(const std::filesystem::path& external_file_path, + const std::filesystem::path& model_file_path, size_t initializer_size_threshold) const { GraphProto result; ToGraphProtoInternal(result); + ORT_ENFORCE(external_file_path.is_relative()); + // If model_file_path is just a file name without a path separator, for example: "model.onnx". Its parent path could + // be empty. Else, save external data file in same directory as the model. + const std::filesystem::path modified_external_file_path = model_file_path.parent_path() / external_file_path; - Path parent_path = Path::Parse(destination_file_path).ParentPath(); - Path external_file_path = Path::Parse(ToPathString(external_file_name)); - // Check if parent_path is relative path (length = 0) - if (parent_path.ToPathString().length()) { - // Save external data file in same directory as model - external_file_path = parent_path.Append(external_file_path); - } - - std::ofstream external_stream(external_file_path.ToPathString(), std::ofstream::out | std::ofstream::binary); + std::ofstream external_stream(modified_external_file_path, std::ofstream::out | std::ofstream::binary); ORT_ENFORCE(external_stream.is_open()); int64_t external_offset = 0; @@ -4022,7 +4019,7 @@ ONNX_NAMESPACE::GraphProto Graph::ToGraphProtoWithExternalInitializers(const std output_proto->set_data_location(ONNX_NAMESPACE::TensorProto_DataLocation::TensorProto_DataLocation_EXTERNAL); ONNX_NAMESPACE::StringStringEntryProto* location = output_proto->add_external_data(); location->set_key("location"); - location->set_value(external_file_name); + location->set_value(ToUTF8String(external_file_path.native())); ONNX_NAMESPACE::StringStringEntryProto* offset = output_proto->add_external_data(); offset->set_key("offset"); offset->set_value(std::to_string(external_offset)); diff --git a/onnxruntime/core/graph/graph_flatbuffers_utils.cc b/onnxruntime/core/graph/graph_flatbuffers_utils.cc index 7dfdba687517f..922759b02e75f 100644 --- a/onnxruntime/core/graph/graph_flatbuffers_utils.cc +++ b/onnxruntime/core/graph/graph_flatbuffers_utils.cc @@ -28,7 +28,7 @@ SaveDims(flatbuffers::FlatBufferBuilder& builder, const DimsFieldType& dims) { Status SaveInitializerOrtFormat(flatbuffers::FlatBufferBuilder& builder, const TensorProto& initializer, - const Path& model_path, + const std::filesystem::path& model_path, flatbuffers::Offset& fbs_tensor, const ExternalDataWriter& external_writer) { auto name = SaveStringToOrtFormat(builder, initializer.has_name(), initializer.name()); @@ -85,7 +85,7 @@ Status SaveInitializerOrtFormat(flatbuffers::FlatBufferBuilder& builder, #if !defined(DISABLE_SPARSE_TENSORS) Status SaveSparseInitializerOrtFormat(flatbuffers::FlatBufferBuilder& builder, const ONNX_NAMESPACE::SparseTensorProto& initializer, - const Path& model_path, + const std::filesystem::path& model_path, flatbuffers::Offset& fbs_sparse_tensor) { // values const auto& values = initializer.values(); @@ -126,7 +126,7 @@ Status SaveSparseInitializerOrtFormat(flatbuffers::FlatBufferBuilder& builder, Status SaveAttributeOrtFormat(flatbuffers::FlatBufferBuilder& builder, const AttributeProto& attr_proto, flatbuffers::Offset& fbs_attr, - const Path& model_path, + const std::filesystem::path& model_path, const onnxruntime::Graph* subgraph) { auto name = SaveStringToOrtFormat(builder, attr_proto.has_name(), attr_proto.name()); auto doc_string = SaveStringToOrtFormat(builder, attr_proto.has_doc_string(), attr_proto.doc_string()); diff --git a/onnxruntime/core/graph/graph_flatbuffers_utils.h b/onnxruntime/core/graph/graph_flatbuffers_utils.h index 33eba34fbaff0..224d966500e18 100644 --- a/onnxruntime/core/graph/graph_flatbuffers_utils.h +++ b/onnxruntime/core/graph/graph_flatbuffers_utils.h @@ -4,6 +4,7 @@ #pragma once #include +#include #include "core/common/flatbuffers.h" @@ -71,13 +72,13 @@ constexpr uint32_t kMinimumSizeForExternalData = 64; /// if the initializer contains kMinimumSizeForExternalData bytes or more, and not string data. Status SaveInitializerOrtFormat( flatbuffers::FlatBufferBuilder& builder, const ONNX_NAMESPACE::TensorProto& initializer, - const Path& model_path, flatbuffers::Offset& fbs_tensor, + const std::filesystem::path& model_path, flatbuffers::Offset& fbs_tensor, const ExternalDataWriter& external_writer = nullptr); #if !defined(DISABLE_SPARSE_TENSORS) Status SaveSparseInitializerOrtFormat( flatbuffers::FlatBufferBuilder& builder, const ONNX_NAMESPACE::SparseTensorProto& initializer, - const Path& model_path, flatbuffers::Offset& fbs_sparse_tensor); + const std::filesystem::path& model_path, flatbuffers::Offset& fbs_sparse_tensor); #endif // !defined(DISABLE_SPARSE_TENSORS) // Convert a given AttributeProto into fbs::Attribute @@ -86,7 +87,7 @@ Status SaveSparseInitializerOrtFormat( // instead of the GraphProto in attr_proto Status SaveAttributeOrtFormat( flatbuffers::FlatBufferBuilder& builder, const ONNX_NAMESPACE::AttributeProto& attr_proto, - flatbuffers::Offset& fbs_attr, const Path& model_path, + flatbuffers::Offset& fbs_attr, const std::filesystem::path& model_path, const onnxruntime::Graph* subgraph); /// diff --git a/onnxruntime/core/graph/model.cc b/onnxruntime/core/graph/model.cc index b3935e69ad7b1..e9d1b4e944edd 100644 --- a/onnxruntime/core/graph/model.cc +++ b/onnxruntime/core/graph/model.cc @@ -81,7 +81,7 @@ Model::Model(const std::string& graph_name, const std::vector& model_local_functions, const logging::Logger& logger, const ModelOptions& options) - : model_path_(Path::Parse(model_path)) { + : model_path_(model_path) { model_proto_.set_ir_version(ONNX_NAMESPACE::Version::IR_VERSION); model_proto_.mutable_graph()->set_name(graph_name); model_metadata_ = model_metadata; @@ -159,7 +159,7 @@ Model::Model(const ModelProto& model_proto, const PathString& model_path, Model::Model(ModelProto&& model_proto, const PathString& model_path, const IOnnxRuntimeOpSchemaRegistryList* local_registries, const logging::Logger& logger, const ModelOptions& options) - : model_path_(Path::Parse(model_path)) { + : model_path_(model_path) { if (!utils::HasGraph(model_proto)) { ORT_THROW("ModelProto does not have a graph."); } @@ -378,8 +378,8 @@ ModelProto Model::ToProto() const { return result; } -ModelProto Model::ToGraphProtoWithExternalInitializers(const std::string& external_file_name, - const PathString& file_path, +ModelProto Model::ToGraphProtoWithExternalInitializers(const std::filesystem::path& external_file_name, + const std::filesystem::path& file_path, size_t initializer_size_threshold) const { ModelProto result(model_proto_); const auto& graph = *graph_; @@ -593,16 +593,14 @@ static Status SaveModel(Model& model, const T& file_path) { #endif } -#ifdef _WIN32 -Status Model::Save(Model& model, const std::wstring& file_path) { +Status Model::Save(Model& model, const PathString& file_path) { return SaveModel(model, file_path); } -#endif template static Status SaveModelWithExternalInitializers(Model& model, const T& file_path, - const std::string& external_file_name, + const std::filesystem::path& external_file_name, size_t initializer_size_threshold) { int fd = 0; Status status = Env::Default().FileOpenWr(file_path, fd); @@ -638,12 +636,8 @@ Status Model::Load(const PathString& file_path, std::shared_ptr& p_model, return LoadModel(file_path, p_model, local_registries, logger, options); } -Status Model::Save(Model& model, const std::string& file_path) { - return SaveModel(model, file_path); -} - -Status Model::SaveWithExternalInitializers(Model& model, const PathString& file_path, - const std::string& external_file_name, +Status Model::SaveWithExternalInitializers(Model& model, const std::filesystem::path& file_path, + const std::filesystem::path& external_file_name, size_t initializer_size_threshold) { return SaveModelWithExternalInitializers(model, file_path, external_file_name, initializer_size_threshold); } @@ -759,8 +753,8 @@ Status Model::Save(Model& model, int p_fd) { Status Model::SaveWithExternalInitializers(Model& model, int fd, - const PathString& file_path, - const std::string& external_file_name, + const std::filesystem::path& file_path, + const std::filesystem::path& external_file_name, size_t initializer_size_threshold) { if (fd < 0) { return Status(ONNXRUNTIME, INVALID_ARGUMENT, " is less than 0."); diff --git a/onnxruntime/core/graph/model.h b/onnxruntime/core/graph/model.h index 6f4b7f4f9f00b..9c73ee16963bd 100644 --- a/onnxruntime/core/graph/model.h +++ b/onnxruntime/core/graph/model.h @@ -7,6 +7,7 @@ #include #include #include +#include #include "core/common/flatbuffers.h" @@ -174,7 +175,7 @@ class Model { const ModelMetaData& MetaData() const noexcept; // Gets the path from which the model was loaded, if any. - const Path& ModelPath() const noexcept { return model_path_; } + const std::filesystem::path& ModelPath() const noexcept { return model_path_; } // Get model's main graph. Graph& MainGraph() noexcept; @@ -187,27 +188,24 @@ class Model { // Get model's serialization proto data. // Save initializer larger than the given threshold (in bytes) into an external binary file // with the given name. This function is useful to avoid hitting the size limit of protobuf files. - ONNX_NAMESPACE::ModelProto ToGraphProtoWithExternalInitializers(const std::string& external_file_name, - const PathString& file_path, + ONNX_NAMESPACE::ModelProto ToGraphProtoWithExternalInitializers(const std::filesystem::path& external_file_name, + const std::filesystem::path& file_path, size_t initializer_size_threshold) const; -#ifdef _WIN32 - static common::Status Save(Model& model, const std::wstring& file_path); -#endif - static common::Status Save(Model& model, const std::string& file_path); + static common::Status Save(Model& model, const PathString& file_path); static common::Status Save(Model& model, int fd); // Save the model to file using an external file for initializers larger than the given threshold (in bytes). static common::Status SaveWithExternalInitializers(Model& model, - const PathString& file_path, - const std::string& external_file_name, + const std::filesystem::path& file_path, + const std::filesystem::path& external_file_path, size_t initializer_size_threshold); static common::Status SaveWithExternalInitializers(Model& model, int fd, - const PathString& file_path, - const std::string& external_file_name, + const std::filesystem::path& file_path, + const std::filesystem::path& external_file_path, size_t initializer_size_threshold); static common::Status Load(std::istream& model_istream, ONNX_NAMESPACE::ModelProto* p_model_proto); @@ -332,7 +330,7 @@ class Model { ModelMetaData model_metadata_; // Path to model file. May be empty. - const Path model_path_; + const std::filesystem::path model_path_; // Main graph of the model. std::unique_ptr graph_; diff --git a/onnxruntime/core/optimizer/initializer.cc b/onnxruntime/core/optimizer/initializer.cc index 3679a40d32eee..7d80e6e5d3a76 100644 --- a/onnxruntime/core/optimizer/initializer.cc +++ b/onnxruntime/core/optimizer/initializer.cc @@ -25,13 +25,13 @@ Initializer::Initializer(ONNX_NAMESPACE::TensorProto_DataType data_type, } } -Initializer::Initializer(const ONNX_NAMESPACE::TensorProto& tensor_proto, const Path& model_path) { +Initializer::Initializer(const ONNX_NAMESPACE::TensorProto& tensor_proto, const std::filesystem::path& model_path) { ORT_ENFORCE(utils::HasDataType(tensor_proto), "Initializer must have a datatype"); #if !defined(__wasm__) // using full filepath is required by utils::TensorProtoToTensor(). One exception is WebAssembly platform, where // external data is not loaded from real file system. if (utils::HasExternalData(tensor_proto)) { - ORT_ENFORCE(!model_path.IsEmpty(), + ORT_ENFORCE(!model_path.empty(), "model_path must not be empty. Ensure that a path is provided when the model is created or loaded."); } #endif @@ -46,7 +46,7 @@ Initializer::Initializer(const ONNX_NAMESPACE::TensorProto& tensor_proto, const // This must be pre-allocated Tensor w(DataTypeImpl::TensorTypeFromONNXEnum(proto_data_type)->GetElementType(), proto_shape, std::make_shared()); - ORT_THROW_IF_ERROR(utils::TensorProtoToTensor(Env::Default(), model_path.ToPathString().c_str(), tensor_proto, w)); + ORT_THROW_IF_ERROR(utils::TensorProtoToTensor(Env::Default(), model_path, tensor_proto, w)); data_ = std::move(w); } diff --git a/onnxruntime/core/optimizer/initializer.h b/onnxruntime/core/optimizer/initializer.h index 78e3fd6a3d24e..b8ae2188beb5d 100644 --- a/onnxruntime/core/optimizer/initializer.h +++ b/onnxruntime/core/optimizer/initializer.h @@ -7,7 +7,7 @@ #include #include #include - +#include #include "core/common/common.h" #include "core/common/narrow.h" #include "core/common/path.h" @@ -28,7 +28,7 @@ class Initializer final { gsl::span dims); Initializer(const ONNX_NAMESPACE::TensorProto& tensor_proto, - const Path& model_path = {}); + const std::filesystem::path& model_path = {}); ~Initializer() = default; diff --git a/onnxruntime/core/optimizer/matmul_scale_fusion.cc b/onnxruntime/core/optimizer/matmul_scale_fusion.cc index e4cdeadbf54d7..338722fb00782 100644 --- a/onnxruntime/core/optimizer/matmul_scale_fusion.cc +++ b/onnxruntime/core/optimizer/matmul_scale_fusion.cc @@ -17,7 +17,8 @@ namespace onnxruntime { namespace { template struct ExtractScalarAsFloatDispatchTarget { - Status operator()(const ONNX_NAMESPACE::TensorProto& tensor_proto, const Path& model_path, float& scalar_float) { + Status operator()(const ONNX_NAMESPACE::TensorProto& tensor_proto, const std::filesystem::path& model_path, + float& scalar_float) { T scalar; ORT_RETURN_IF_ERROR(utils::UnpackTensor(tensor_proto, model_path, &scalar, 1)); scalar_float = static_cast(scalar); diff --git a/onnxruntime/core/optimizer/optimizer_execution_frame.cc b/onnxruntime/core/optimizer/optimizer_execution_frame.cc index 1eabc079f3a20..ed7d5feb2beb3 100644 --- a/onnxruntime/core/optimizer/optimizer_execution_frame.cc +++ b/onnxruntime/core/optimizer/optimizer_execution_frame.cc @@ -30,7 +30,7 @@ static size_t EstimateInputsOutputs(gsl::span nodes) { OptimizerExecutionFrame::Info::Info(const std::vector& nodes, const InitializedTensorSet& initialized_tensor_set, - const Path& model_path, + const std::filesystem::path& model_path, const IExecutionProvider& execution_provider, const std::function& is_sparse_initializer_func) : execution_provider_(execution_provider), @@ -52,7 +52,7 @@ OptimizerExecutionFrame::Info::Info(const std::vector& nodes, OrtValue ort_value; ORT_RETURN_IF_ERROR( utils::TensorProtoToOrtValue(Env::Default(), - model_path.IsEmpty() ? nullptr : model_path.ToPathString().c_str(), + model_path, tensor_proto, allocator_ptr_, ort_value)); initializers_[idx] = std::move(ort_value); @@ -77,7 +77,7 @@ OptimizerExecutionFrame::Info::Info(const std::vector& nodes, OptimizerExecutionFrame::Info::Info(const std::vector& nodes, const std::unordered_map& initialized_tensor_set, - const Path& model_path, + const std::filesystem::path& /* model_path */, const IExecutionProvider& execution_provider, const std::function& is_sparse_initializer_func) : execution_provider_(execution_provider), @@ -88,8 +88,7 @@ OptimizerExecutionFrame::Info::Info(const std::vector& nodes, ORT_THROW_IF_ERROR(data_transfer_mgr_.RegisterDataTransfer(std::make_unique())); // Create MLValues related maps - auto initialize_maps = [this, &initialized_tensor_set, &model_path](const NodeArg& arg, size_t /*index*/) -> Status { - (void)model_path; + auto initialize_maps = [this, &initialized_tensor_set](const NodeArg& arg, size_t /*index*/) -> Status { int idx = ort_value_name_idx_map_.Add(arg.Name()); ort_value_idx_nodearg_map_.insert_or_assign(idx, &arg); diff --git a/onnxruntime/core/optimizer/optimizer_execution_frame.h b/onnxruntime/core/optimizer/optimizer_execution_frame.h index 3dbf6c1d97aa6..b0f7f461661b5 100644 --- a/onnxruntime/core/optimizer/optimizer_execution_frame.h +++ b/onnxruntime/core/optimizer/optimizer_execution_frame.h @@ -4,6 +4,7 @@ #pragma once #include +#include #include "core/common/inlined_containers.h" #include "core/graph/graph.h" @@ -24,13 +25,13 @@ class OptimizerExecutionFrame final : public IExecutionFrame { public: Info(const std::vector& nodes, const InitializedTensorSet& initialized_tensor_set, - const Path& model_path, + const std::filesystem::path& model_path, const IExecutionProvider& execution_provider, const std::function& is_sparse_initializer_func); Info(const std::vector& nodes, const std::unordered_map& initialized_tensor_set, - const Path& model_path, + const std::filesystem::path& model_path, const IExecutionProvider& execution_provider, const std::function& is_sparse_initializer_func); diff --git a/onnxruntime/core/optimizer/qdq_transformer/qdq_util.cc b/onnxruntime/core/optimizer/qdq_transformer/qdq_util.cc index 5a6c47a8d8454..a4d1ea3c7cf56 100644 --- a/onnxruntime/core/optimizer/qdq_transformer/qdq_util.cc +++ b/onnxruntime/core/optimizer/qdq_transformer/qdq_util.cc @@ -17,7 +17,7 @@ namespace onnxruntime::QDQ { bool IsQDQPairSupported( const Node& q_node, const Node& dq_node, const GetConstantInitializerFn& get_const_initializer, - const Path& model_path, + const std::filesystem::path& model_path, bool check_op_type) { if (check_op_type) { if (!MatchQNode(q_node) || !MatchDQNode(dq_node)) { @@ -86,7 +86,7 @@ bool IsQDQPairSupported( bool IsDQQConversion( const Node& dq_node, const Node& q_node, const GetConstantInitializerFn& get_const_initializer, - const Path& model_path) { + const std::filesystem::path& model_path) { ConstPointerContainer> dq_input_defs = dq_node.InputDefs(); ConstPointerContainer> q_input_defs = q_node.InputDefs(); diff --git a/onnxruntime/core/optimizer/qdq_transformer/qdq_util.h b/onnxruntime/core/optimizer/qdq_transformer/qdq_util.h index c5f7cd601a2f0..5d11b8bfd5558 100644 --- a/onnxruntime/core/optimizer/qdq_transformer/qdq_util.h +++ b/onnxruntime/core/optimizer/qdq_transformer/qdq_util.h @@ -5,6 +5,7 @@ #include #include +#include namespace ONNX_NAMESPACE { class TensorProto; @@ -36,7 +37,7 @@ using GetConstantInitializerFn = std::function Q sequence represents a conversion in quantization data type. @@ -49,7 +50,7 @@ bool IsQDQPairSupported( bool IsDQQConversion( const Node& dq_node, const Node& q_node, const GetConstantInitializerFn& get_const_initializer, - const Path& model_path); + const std::filesystem::path& model_path); // Check if DQ is supported in extended level QDQ transformers. It requires: // 1. DQ doesn't have optional input. diff --git a/onnxruntime/core/optimizer/transpose_optimization/ort_optimizer_api_impl.cc b/onnxruntime/core/optimizer/transpose_optimization/ort_optimizer_api_impl.cc index c532f56b3d3d9..1f7e54cb807ea 100644 --- a/onnxruntime/core/optimizer/transpose_optimization/ort_optimizer_api_impl.cc +++ b/onnxruntime/core/optimizer/transpose_optimization/ort_optimizer_api_impl.cc @@ -46,11 +46,12 @@ class ApiValueInfo final : public api::ValueInfoRef { class ApiTensor final : public api::TensorRef { private: const onnx::TensorProto& tensor_proto_; - const Path& model_path_; + const std::filesystem::path& model_path_; AllocatorPtr cpu_allocator_; public: - explicit ApiTensor(const onnx::TensorProto& tensor_proto, const Path& model_path, AllocatorPtr cpu_allocator) + explicit ApiTensor(const onnx::TensorProto& tensor_proto, const std::filesystem::path& model_path, + AllocatorPtr cpu_allocator) : tensor_proto_(tensor_proto), model_path_(model_path), cpu_allocator_(std::move(cpu_allocator)) {} const onnx::TensorProto& TensorProto() { @@ -289,10 +290,12 @@ std::vector ApiTensor::Data() const { auto tensor_shape_dims = utils::GetTensorShapeFromTensorProto(tensor_proto_); TensorShape tensor_shape{std::move(tensor_shape_dims)}; onnxruntime::Tensor tensor(tensor_dtype, tensor_shape, cpu_allocator_); - ORT_THROW_IF_ERROR(utils::TensorProtoToTensor(Env::Default(), model_path_.ToPathString().c_str(), + ORT_THROW_IF_ERROR(utils::TensorProtoToTensor(Env::Default(), model_path_, tensor_proto_, tensor)); size_t num_bytes = gsl::narrow_cast(tensor.SizeInBytes()); const uint8_t* data = static_cast(tensor.DataRaw()); + // TODO: the returned data is unaligned, which does not meet the alignment requirement that mlas requires. Because + // the returned type is a vector, not a Tensor or tensor buffer that is allocated from a CPU allocator. return std::vector(data, data + num_bytes); } // @@ -554,7 +557,7 @@ void ApiGraph::TransposeInitializer(std::string_view name, const std::vector #include "core/common/common.h" #include "core/framework/op_kernel.h" #include "core/providers/cpu/ml/ml_common.h" @@ -123,7 +123,7 @@ std::vector GetAttribute(const OpKernelInfo& info, const std::string& name, c } const SafeInt tensor_size(element_count); std::vector out(tensor_size); - result = utils::UnpackTensor(attr_tensor_proto, Path(), out.data(), tensor_size); + result = utils::UnpackTensor(attr_tensor_proto, std::filesystem::path(), out.data(), tensor_size); ORT_ENFORCE(result.IsOK(), "LabelEncoder could not unpack tensor attribute ", name); return out; } @@ -134,7 +134,7 @@ T GetDefault(const OpKernelInfo& info, const std::string& attr_name, const T& ba auto result = info.GetAttr("default_tensor", &attr_tensor_proto); if (result.IsOK() && utils::HasDataType(attr_tensor_proto)) { T default_value; - result = utils::UnpackTensor(attr_tensor_proto, Path(), &default_value, 1); + result = utils::UnpackTensor(attr_tensor_proto, std::filesystem::path(), &default_value, 1); ORT_ENFORCE(result.IsOK(), "LabelEncoder could not unpack default tensor ", attr_name); return default_value; } else if constexpr (std::is_same_v || std::is_same_v || std::is_same_v) { diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/DmlRuntimeFusedGraphKernel.cpp b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/DmlRuntimeFusedGraphKernel.cpp index 10b8b7fe42f86..2f110ba339beb 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/DmlRuntimeFusedGraphKernel.cpp +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/DmlRuntimeFusedGraphKernel.cpp @@ -20,7 +20,7 @@ namespace Dml DmlRuntimeFusedGraphKernel( const onnxruntime::OpKernelInfo& kernelInfo, std::shared_ptr indexedSubGraph, - const onnxruntime::Path& modelPath, + const std::filesystem::path& modelPath, std::vector>&& subgraphNodes, std::vector&& subgraphInputs, std::vector&& subgraphOutputs, @@ -314,7 +314,7 @@ namespace Dml mutable std::optional m_persistentResourceBinding; std::shared_ptr m_indexedSubGraph; - const onnxruntime::Path& m_modelPath; + const std::filesystem::path& m_modelPath; std::vector> m_subgraphNodes; std::vector m_subgraphInputs; @@ -341,7 +341,7 @@ namespace Dml onnxruntime::OpKernel* CreateRuntimeFusedGraphKernel( const onnxruntime::OpKernelInfo& info, std::shared_ptr indexedSubGraph, - const onnxruntime::Path& modelPath, + const std::filesystem::path& modelPath, std::vector>&& subgraphNodes, std::vector&& subgraphInputs, std::vector&& subgraphOutputs, diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/DmlRuntimeFusedGraphKernel.h b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/DmlRuntimeFusedGraphKernel.h index d679c5aa5667c..e800175268557 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/DmlRuntimeFusedGraphKernel.h +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/DmlRuntimeFusedGraphKernel.h @@ -1,6 +1,6 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. - +#include #include "core/framework/op_kernel.h" #include "GraphDescBuilder.h" #include "DmlRuntimeGraphFusionTransformer.h" @@ -10,7 +10,7 @@ namespace Dml onnxruntime::OpKernel* CreateRuntimeFusedGraphKernel( const onnxruntime::OpKernelInfo& info, std::shared_ptr indexedSubGraph, - const onnxruntime::Path& modelPath, + const std::filesystem::path& modelPath, std::vector>&& subgraphNodes, std::vector&& subgraphInputs, std::vector&& subgraphOutputs, diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/GraphDescBuilder.cpp b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/GraphDescBuilder.cpp index 2bd9377e4c2fa..387767f821b3e 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/GraphDescBuilder.cpp +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/GraphDescBuilder.cpp @@ -232,7 +232,7 @@ namespace Dml::GraphDescBuilder const std::unordered_map>& isInitializerTransferable, const std::unordered_map& graphNodePropertyMap, const ExecutionProviderImpl* executionHandle, - const onnxruntime::Path& modelPath, + const std::filesystem::path& modelPath, gsl::span subgraphNodes, gsl::span subgraphInputs, gsl::span subgraphOutputs, diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/GraphDescBuilder.h b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/GraphDescBuilder.h index 4055984b40405..3f778b3a7feba 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/GraphDescBuilder.h +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/GraphDescBuilder.h @@ -2,6 +2,7 @@ // Licensed under the MIT License. #pragma once +#include #include "MLOperatorAuthorImpl.h" #include "ExecutionProvider.h" @@ -41,7 +42,7 @@ namespace Dml const std::unordered_map>& isInitializerTransferable, const std::unordered_map& graphNodePropertyMap, const ExecutionProviderImpl* executionHandle, - const onnxruntime::Path& modelPath, + const std::filesystem::path& modelPath, gsl::span subgraphNodes, gsl::span subgraphInputs, gsl::span subgraphOutputs, diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/MLOperatorAuthorImpl.cpp b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/MLOperatorAuthorImpl.cpp index f29fbc7a1a65b..0a2a5bbcbedaf 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/MLOperatorAuthorImpl.cpp +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/MLOperatorAuthorImpl.cpp @@ -842,7 +842,7 @@ namespace Windows::AI::MachineLearning::Adapter const onnx::TensorProto* tensorProto = &attributeProto->t(); // An empty path is used as external weights are not currently supported in this case - Microsoft::WRL::ComPtr tensorWrapper = wil::MakeOrThrow(const_cast(tensorProto), onnxruntime::Path()); + Microsoft::WRL::ComPtr tensorWrapper = wil::MakeOrThrow(const_cast(tensorProto), std::filesystem::path()); *tensor = tensorWrapper.Detach(); return S_OK; } @@ -1545,7 +1545,7 @@ namespace Windows::AI::MachineLearning::Adapter ORT_CATCH_RETURN } - OnnxTensorWrapper::OnnxTensorWrapper(onnx::TensorProto* impl, const onnxruntime::Path& modelPath) : m_impl(impl) + OnnxTensorWrapper::OnnxTensorWrapper(onnx::TensorProto* impl, const std::filesystem::path& modelPath) : m_impl(impl) { // The tensor may be stored as raw data or in typed fields. if (impl->data_location() == onnx::TensorProto_DataLocation_EXTERNAL) @@ -2826,7 +2826,7 @@ namespace Windows::AI::MachineLearning::Adapter { // An empty path is used as external weights are not currently supported in this case Microsoft::WRL::ComPtr tensorWrapper = wil::MakeOrThrow( - const_cast(ctx->getInputData(index)), onnxruntime::Path()); + const_cast(ctx->getInputData(index)), std::filesystem::path()); return tensorWrapper; } ); @@ -3018,7 +3018,7 @@ namespace Windows::AI::MachineLearning::Adapter std::tuple, size_t> UnpackTensor( const onnx::TensorProto& initializer, - const onnxruntime::Path& modelPath) + const std::filesystem::path& modelPath) { std::unique_ptr unpackedTensor; size_t tensorByteSize = 0; diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/MLOperatorAuthorImpl.h b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/MLOperatorAuthorImpl.h index 59e253e88457a..7e51ce026d365 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/MLOperatorAuthorImpl.h +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/MLOperatorAuthorImpl.h @@ -2,6 +2,8 @@ // Licensed under the MIT License. #pragma once +#include + #include "core/providers/dml/DmlExecutionProvider/inc/IWinmlExecutionProvider.h" #include "core/providers/dml/OperatorAuthorHelper/MLOperatorAuthorHelper.h" #include "core/providers/dml/DmlExecutionProvider/src/DmlEdgeShapes.h" @@ -283,7 +285,7 @@ class OnnxTensorWrapper : public WRL::Base, public Closable public: OnnxTensorWrapper() = default; - OnnxTensorWrapper(onnx::TensorProto* impl, const onnxruntime::Path& modelPath); + OnnxTensorWrapper(onnx::TensorProto* impl, const std::filesystem::path& modelPath); uint32_t STDMETHODCALLTYPE GetDimensionCount() const noexcept override; @@ -681,5 +683,5 @@ bool TryGetStaticInputShapes(const onnxruntime::Node& node, EdgeShapes& inputSha bool TryGetStaticOutputShapes(const onnxruntime::Node& node, EdgeShapes& outputShapes); bool ContainsEmptyDimensions(const EdgeShapes& shapes, gsl::span ignoredShapeIndices = gsl::span()); -std::tuple, size_t> UnpackTensor(const onnx::TensorProto& initializer, const onnxruntime::Path& modelPath); +std::tuple, size_t> UnpackTensor(const onnx::TensorProto& initializer, const std::filesystem::path& modelPath); } // namespace Windows::AI::MachineLearning::Adapter diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Utility.h b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Utility.h index 02166f992449e..a3f2777a0c805 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Utility.h +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Utility.h @@ -6,7 +6,7 @@ #include #include #include - +#include namespace Dml { @@ -16,21 +16,14 @@ namespace Dml return g_converterToUtf16.from_bytes(str.data()); } - static inline std::wstring GetModelName(const onnxruntime::Path& modelPath) + static inline std::wstring GetModelName(const std::filesystem::path& modelPath) { - if (modelPath.GetComponents().empty()) + if (modelPath.empty() || !modelPath.has_filename() || !modelPath.has_extension()) { return L""; } - const onnxruntime::PathString& pathString = modelPath.GetComponents().back(); - size_t dotPosition = pathString.find_last_of('.'); - if (dotPosition == std::string::npos) - { - return L""; - } - - return pathString.substr(0, dotPosition); + return modelPath.stem().native(); } static inline std::wstring GetSanitizedFileName(std::wstring_view name) @@ -138,4 +131,4 @@ namespace StringUtil return {}; } -} \ No newline at end of file +} diff --git a/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/helper.cc b/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/helper.cc index 745504ca04941..5108f90fc763a 100644 --- a/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/helper.cc +++ b/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/helper.cc @@ -185,7 +185,8 @@ bool HasValidBinaryOpQuantizedInputTypes(const NodeUnit& node_unit) { } common::Status GetQuantizationScaleAndZeroPoint(const GraphViewer& graph_viewer, const NodeUnitIODef& io_def, - const Path& model_path, float& scale, int32_t& zero_point) { + const std::filesystem::path& model_path, float& scale, + int32_t& zero_point) { scale = 0.0f; zero_point = 0; diff --git a/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/helper.h b/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/helper.h index a606b8aceb63d..d4967b6251824 100644 --- a/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/helper.h +++ b/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/helper.h @@ -5,6 +5,7 @@ #include #include +#include #include "core/common/inlined_containers.h" #include "core/graph/basic_types.h" #include "core/providers/nnapi/nnapi_builtin/nnapi_lib/NeuralNetworksTypes.h" @@ -132,7 +133,7 @@ bool IsQuantizedBinaryOp(QuantizedOpType quant_op_type); bool HasValidBinaryOpQuantizedInputTypes(const NodeUnit& node_unit); common::Status GetQuantizationScaleAndZeroPoint( - const GraphViewer& graph_viewer, const NodeUnitIODef& io_def, const Path& model_path, + const GraphViewer& graph_viewer, const NodeUnitIODef& io_def, const std::filesystem::path& model_path, float& scale, int32_t& zero_point); common::Status GetQuantizationScaleAndZeroPoint( diff --git a/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/op_builder_helpers.cc b/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/op_builder_helpers.cc index dab7bccf43396..c1770e0119b25 100644 --- a/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/op_builder_helpers.cc +++ b/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/op_builder_helpers.cc @@ -1142,7 +1142,7 @@ bool IsQuantizationScaleSupported(const GraphViewer& graph_viewer, bool IsQuantizationZeroPointSupported(const GraphViewer& graph_viewer, const NodeUnitIODef& io_def, const std::string& op_type, - const Path& model_path, + const std::filesystem::path& model_path, bool is_quant_matmul, bool is_conv_matmul_u8s8_weight) { // zero point is optional here @@ -1282,7 +1282,7 @@ bool IsQuantizedIOSupported(const GraphViewer& graph_viewer, const NodeUnit& nod bool HasRequiredScaleAndZeroPoint(const GraphViewer& graph_viewer, const std::string& op_desc, const NodeUnitIODef& io_def, - const Path& path, + const std::filesystem::path& path, float required_scale, int32_t required_zp) { float scale = 0.0f; int32_t zp = 0; diff --git a/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/op_builder_helpers.h b/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/op_builder_helpers.h index 0844857a06d61..94e511e04dff3 100644 --- a/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/op_builder_helpers.h +++ b/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/op_builder_helpers.h @@ -5,6 +5,7 @@ #include #include +#include #include "core/common/common.h" #include "core/framework/node_unit.h" @@ -200,7 +201,7 @@ bool IsQuantizationScaleSupported(const GraphViewer& graph_viewer, bool IsQuantizationZeroPointSupported(const GraphViewer& graph_viewer, const NodeUnitIODef& io_def, const std::string& op_type, - const Path& model_path, + const std::filesystem::path& model_path, bool is_quant_matmul, bool is_conv_matmul_u8s8_weight); @@ -214,7 +215,7 @@ bool IsQuantizedIOSupported(const GraphViewer& graph_viewer, const NodeUnit& nod bool HasRequiredScaleAndZeroPoint(const GraphViewer& graph_viewer, const std::string& op_desc, const NodeUnitIODef& io_def, - const Path& path, + const std::filesystem::path& path, float required_scale, int32_t required_zp); // performs broadcasting operation on two shapes to make them compatible diff --git a/onnxruntime/core/providers/openvino/openvino_execution_provider.cc b/onnxruntime/core/providers/openvino/openvino_execution_provider.cc index 040c56926a803..655e1b180388b 100644 --- a/onnxruntime/core/providers/openvino/openvino_execution_provider.cc +++ b/onnxruntime/core/providers/openvino/openvino_execution_provider.cc @@ -90,14 +90,8 @@ OpenVINOExecutionProvider::GetCapability(const GraphViewer& graph_viewer, if (!(GetEnvironmentVar("ORT_OPENVINO_ENABLE_CI_LOG").empty())) { std::cout << "In the OpenVINO EP" << std::endl; } -#ifdef _WIN32 - std::wstring onnx_path = graph_viewer.ModelPath().ToPathString(); - global_context_->onnx_model_path_name = - std::string(onnx_path.begin(), onnx_path.end()); -#else - global_context_->onnx_model_path_name = - graph_viewer.ModelPath().ToPathString(); -#endif + global_context_->onnx_model_path_name = graph_viewer.ModelPath().string(); + global_context_->onnx_opset_version = graph_viewer.DomainToVersionMap().at(kOnnxDomain); diff --git a/onnxruntime/core/providers/qnn/builder/opbuilder/base_op_builder.cc b/onnxruntime/core/providers/qnn/builder/opbuilder/base_op_builder.cc index e1156288d2f8f..2fbe59bf0d578 100644 --- a/onnxruntime/core/providers/qnn/builder/opbuilder/base_op_builder.cc +++ b/onnxruntime/core/providers/qnn/builder/opbuilder/base_op_builder.cc @@ -235,9 +235,8 @@ Status BaseOpBuilder::TransposeInitializer(const QnnModelWrapper& qnn_model_wrap TensorShape new_tensor_shape(new_tensor_shape_dims); Tensor out_tensor = Tensor(tensor_dtype, new_tensor_shape, cpu_allocator); - onnxruntime::PathString model_path = qnn_model_wrapper.GetGraphViewer().ModelPath().ToPathString(); - const ORTCHAR_T* model_path_str = model_path.empty() ? nullptr : model_path.c_str(); - ORT_RETURN_IF_ERROR(onnxruntime::utils::TensorProtoToTensor(Env::Default(), model_path_str, initializer, in_tensor)); + ORT_RETURN_IF_ERROR(onnxruntime::utils::TensorProtoToTensor( + Env::Default(), qnn_model_wrapper.GetGraphViewer().ModelPath(), initializer, in_tensor)); ORT_RETURN_IF_ERROR(Transpose::DoTranspose(permutations, in_tensor, out_tensor)); onnx::TensorProto new_tensor_proto = onnxruntime::utils::TensorToTensorProto(out_tensor, "test"); ORT_RETURN_IF_ERROR(qnn_model_wrapper.UnpackInitializerData(new_tensor_proto, transposed_data)); diff --git a/onnxruntime/core/providers/qnn/builder/opbuilder/slice_op_builder.cc b/onnxruntime/core/providers/qnn/builder/opbuilder/slice_op_builder.cc index 88c94581a8887..b033c8723ea86 100644 --- a/onnxruntime/core/providers/qnn/builder/opbuilder/slice_op_builder.cc +++ b/onnxruntime/core/providers/qnn/builder/opbuilder/slice_op_builder.cc @@ -92,10 +92,8 @@ static Status GetInitializerInputData(const NodeUnitIODef& input, const QnnModel Tensor tensor(dtype, shape, std::make_shared()); // Deserialize initializer into Tensor. - onnxruntime::PathString model_path = qnn_model_wrapper.GetGraphViewer().ModelPath().ToPathString(); - const ORTCHAR_T* model_path_str = model_path.empty() ? nullptr : model_path.c_str(); - ORT_RETURN_IF_ERROR(onnxruntime::utils::TensorProtoToTensor(onnxruntime::Env::Default(), model_path_str, - *initializer_proto, tensor)); + ORT_RETURN_IF_ERROR(onnxruntime::utils::TensorProtoToTensor( + onnxruntime::Env::Default(), qnn_model_wrapper.GetGraphViewer().ModelPath(), *initializer_proto, tensor)); Status status; diff --git a/onnxruntime/core/providers/qnn/qnn_execution_provider.cc b/onnxruntime/core/providers/qnn/qnn_execution_provider.cc index c159730d46cf1..0ddaa97694217 100644 --- a/onnxruntime/core/providers/qnn/qnn_execution_provider.cc +++ b/onnxruntime/core/providers/qnn/qnn_execution_provider.cc @@ -811,7 +811,7 @@ Status QNNExecutionProvider::Compile(const std::vector& fused const onnxruntime::GraphViewer& graph_viewer_0(fused_nodes_and_graphs[0].filtered_graph); is_ctx_file_exist = qnn::ValidateContextCacheFilePath(is_qnn_ctx_model, context_cache_path_cfg_, - graph_viewer_0.ModelPath().ToPathString(), + graph_viewer_0.ModelPath().native(), context_cache_path); } diff --git a/onnxruntime/core/providers/shared_library/provider_bridge_provider.cc b/onnxruntime/core/providers/shared_library/provider_bridge_provider.cc index 27d8a0f06f565..6e6a80f097c12 100644 --- a/onnxruntime/core/providers/shared_library/provider_bridge_provider.cc +++ b/onnxruntime/core/providers/shared_library/provider_bridge_provider.cc @@ -503,7 +503,7 @@ template <> Status UnpackTensor(const ONNX_NAMESPACE::TensorProto& tensor, const void* raw_data, size_t raw_data_len, /*out*/ int64_t* p_data, size_t expected_size) { return g_host->UnpackTensor(tensor, raw_data, raw_data_len, p_data, expected_size); } template <> Status UnpackTensor(const ONNX_NAMESPACE::TensorProto& tensor, const void* raw_data, size_t raw_data_len, /*out*/ uint64_t* p_data, size_t expected_size) { return g_host->UnpackTensor(tensor, raw_data, raw_data_len, p_data, expected_size); } -Status UnpackInitializerData(const ONNX_NAMESPACE::TensorProto& tensor, const Path& model_path, +Status UnpackInitializerData(const ONNX_NAMESPACE::TensorProto& tensor, const std::filesystem::path& model_path, /*out*/ std::vector& unpacked_tensor) { return g_host->UnpackInitializerData(tensor, model_path, unpacked_tensor); } diff --git a/onnxruntime/core/providers/shared_library/provider_interfaces.h b/onnxruntime/core/providers/shared_library/provider_interfaces.h index cc3b13f696a96..7454b322a310c 100644 --- a/onnxruntime/core/providers/shared_library/provider_interfaces.h +++ b/onnxruntime/core/providers/shared_library/provider_interfaces.h @@ -5,6 +5,7 @@ #include #include #include +#include // Public wrappers around internal ort interfaces (currently) #include "core/providers/shared_library/provider_host_api.h" @@ -209,7 +210,7 @@ struct ProviderHost { virtual Status UnpackTensor(const ONNX_NAMESPACE::TensorProto& tensor, const void* raw_data, size_t raw_data_len, /*out*/ uint32_t* p_data, size_t expected_size) = 0; virtual Status UnpackTensor(const ONNX_NAMESPACE::TensorProto& tensor, const void* raw_data, size_t raw_data_len, /*out*/ int64_t* p_data, size_t expected_size) = 0; virtual Status UnpackTensor(const ONNX_NAMESPACE::TensorProto& tensor, const void* raw_data, size_t raw_data_len, /*out*/ uint64_t* p_data, size_t expected_size) = 0; - virtual Status UnpackInitializerData(const ONNX_NAMESPACE::TensorProto& tensor, const Path& model_path, + virtual Status UnpackInitializerData(const ONNX_NAMESPACE::TensorProto& tensor, const std::filesystem::path& model_path, /*out*/ std::vector& unpacked_tensor) = 0; virtual uint16_t math__floatToHalf(float f) = 0; @@ -784,7 +785,7 @@ struct ProviderHost { virtual const std::string& NodeUnit__Name(const NodeUnit* p) noexcept = 0; virtual int NodeUnit__SinceVersion(const NodeUnit* p) noexcept = 0; virtual NodeIndex NodeUnit__Index(const NodeUnit* p) noexcept = 0; - virtual const Path& NodeUnit__ModelPath(const NodeUnit* p) noexcept = 0; + virtual const std::filesystem::path& NodeUnit__ModelPath(const NodeUnit* p) noexcept = 0; virtual ProviderType NodeUnit__GetExecutionProviderType(const NodeUnit* p) noexcept = 0; virtual const Node& NodeUnit__GetNode(const NodeUnit* p) noexcept = 0; @@ -806,7 +807,7 @@ struct ProviderHost { virtual void Model__operator_delete(Model* p) = 0; virtual Graph& Model__MainGraph(Model* p) = 0; virtual std::unique_ptr Model__ToProto(Model* p) = 0; - virtual std::unique_ptr Model__ToGraphProtoWithExternalInitializers(Model* p, const std::string& external_file_name, const PathString& file_path, size_t initializer_size_threshold) = 0; + virtual std::unique_ptr Model__ToGraphProtoWithExternalInitializers(Model* p, const std::filesystem::path& external_file_name, const std::filesystem::path& file_path, size_t initializer_size_threshold) = 0; virtual const ModelMetaData& Model__MetaData(const Model* p) const noexcept = 0; virtual Status Model__Load(const PathString& file_path, /*out*/ ONNX_NAMESPACE::ModelProto& model_proto) = 0; @@ -834,7 +835,7 @@ struct ProviderHost { virtual const Graph* Graph__ParentGraph(const Graph* p) const = 0; virtual Graph* Graph__MutableParentGraph(Graph* p) = 0; virtual const std::string& Graph__Name(const Graph* p) const noexcept = 0; - virtual const Path& Graph__ModelPath(const Graph* p) const = 0; + virtual const std::filesystem::path& Graph__ModelPath(const Graph* p) const = 0; virtual const std::vector& Graph__GetInputsIncludingInitializers(const Graph* p) const noexcept = 0; virtual bool Graph__IsSubgraph(const Graph* p) = 0; virtual const Node* Graph__GetProducerNode(const Graph* p, const std::string& node_arg_name) const = 0; @@ -868,7 +869,7 @@ struct ProviderHost { virtual std::unique_ptr GraphViewer__CreateModel(const GraphViewer* p, const logging::Logger& logger) = 0; virtual const std::string& GraphViewer__Name(const GraphViewer* p) noexcept = 0; - virtual const Path& GraphViewer__ModelPath(const GraphViewer* p) noexcept = 0; + virtual const std::filesystem::path& GraphViewer__ModelPath(const GraphViewer* p) noexcept = 0; virtual const Node* GraphViewer__GetNode(const GraphViewer* p, NodeIndex node_index) = 0; virtual const NodeArg* GraphViewer__GetNodeArg(const GraphViewer* p, const std::string& name) = 0; diff --git a/onnxruntime/core/providers/shared_library/provider_wrappedtypes.h b/onnxruntime/core/providers/shared_library/provider_wrappedtypes.h index fd2540b42a3db..2ccd05fe9df60 100644 --- a/onnxruntime/core/providers/shared_library/provider_wrappedtypes.h +++ b/onnxruntime/core/providers/shared_library/provider_wrappedtypes.h @@ -811,7 +811,7 @@ struct NodeUnit final { const std::string& Name() const noexcept { return g_host->NodeUnit__Name(this); } int SinceVersion() const noexcept { return g_host->NodeUnit__SinceVersion(this); } NodeIndex Index() const noexcept { return g_host->NodeUnit__Index(this); } - const Path& ModelPath() const noexcept { return g_host->NodeUnit__ModelPath(this); } + const std::filesystem::path& ModelPath() const noexcept { return g_host->NodeUnit__ModelPath(this); } ProviderType GetExecutionProviderType() const noexcept { return g_host->NodeUnit__GetExecutionProviderType(this); } const Node& GetNode() const noexcept { return g_host->NodeUnit__GetNode(this); } @@ -840,7 +840,7 @@ struct Model final { Graph& MainGraph() { return g_host->Model__MainGraph(this); } std::unique_ptr ToProto() { return g_host->Model__ToProto(this); } - std::unique_ptr ToGraphProtoWithExternalInitializers(const std::string& external_file_name, const PathString& file_path, size_t initializer_size_threshold) { return g_host->Model__ToGraphProtoWithExternalInitializers(this, external_file_name, file_path, initializer_size_threshold); } + std::unique_ptr ToGraphProtoWithExternalInitializers(const std::filesystem::path& external_file_name, const std::filesystem::path& file_path, size_t initializer_size_threshold) { return g_host->Model__ToGraphProtoWithExternalInitializers(this, external_file_name, file_path, initializer_size_threshold); } const ModelMetaData& MetaData() const noexcept { return g_host->Model__MetaData(this); } Model() = delete; @@ -873,7 +873,7 @@ struct Graph final { const Graph* ParentGraph() const { return g_host->Graph__ParentGraph(this); } Graph* MutableParentGraph() { return g_host->Graph__MutableParentGraph(this); } const std::string& Name() const noexcept { return g_host->Graph__Name(this); } - const Path& ModelPath() const { return g_host->Graph__ModelPath(this); } + const std::filesystem::path& ModelPath() const { return g_host->Graph__ModelPath(this); } const std::vector& GetInputsIncludingInitializers() const noexcept { return g_host->Graph__GetInputsIncludingInitializers(this); } bool IsSubgraph() const { return g_host->Graph__IsSubgraph(this); } const Node* GetProducerNode(const std::string& node_arg_name) const { return g_host->Graph__GetProducerNode(this, node_arg_name); } @@ -923,7 +923,7 @@ class GraphViewer final { std::unique_ptr CreateModel(const logging::Logger& logger) const { return g_host->GraphViewer__CreateModel(this, logger); } const std::string& Name() const noexcept { return g_host->GraphViewer__Name(this); } - const Path& ModelPath() const noexcept { return g_host->GraphViewer__ModelPath(this); } + const std::filesystem::path& ModelPath() const noexcept { return g_host->GraphViewer__ModelPath(this); } const Node* GetNode(NodeIndex node_index) const { return g_host->GraphViewer__GetNode(this, node_index); } const NodeArg* GetNodeArg(const std::string& name) const { return g_host->GraphViewer__GetNodeArg(this, name); } diff --git a/onnxruntime/core/providers/tensorrt/onnx_ctx_model_helper.cc b/onnxruntime/core/providers/tensorrt/onnx_ctx_model_helper.cc index 2171ce056e029..42788f2960197 100644 --- a/onnxruntime/core/providers/tensorrt/onnx_ctx_model_helper.cc +++ b/onnxruntime/core/providers/tensorrt/onnx_ctx_model_helper.cc @@ -29,7 +29,7 @@ bool GraphHasCtxNode(const GraphViewer& graph_viewer) { return false; } -const onnxruntime::Path& GetModelPath(const GraphViewer& graph_viewer) { +const std::filesystem::path& GetModelPath(const GraphViewer& graph_viewer) { // find the top level graph const Graph* cur_graph = &graph_viewer.GetGraph(); while (cur_graph->IsSubgraph()) { diff --git a/onnxruntime/core/providers/tensorrt/onnx_ctx_model_helper.h b/onnxruntime/core/providers/tensorrt/onnx_ctx_model_helper.h index f8fefc12c3453..3be08d043da48 100644 --- a/onnxruntime/core/providers/tensorrt/onnx_ctx_model_helper.h +++ b/onnxruntime/core/providers/tensorrt/onnx_ctx_model_helper.h @@ -24,7 +24,7 @@ static const std::string EPCONTEXT_WARNING = for the best model loading time"; bool GraphHasCtxNode(const GraphViewer& graph_viewer); -const onnxruntime::Path& GetModelPath(const GraphViewer& graph_viewer); +const std::filesystem::path& GetModelPath(const GraphViewer& graph_viewer); std::filesystem::path GetPathOrParentPathOfCtxModel(const std::string& ep_context_file_path); ONNX_NAMESPACE::ModelProto* CreateCtxModel(const GraphViewer& graph_viewer, const std::string engine_cache_path, diff --git a/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc b/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc index 13316d6cbc749..3ca0935b9e46c 100644 --- a/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc +++ b/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc @@ -2327,12 +2327,13 @@ TensorrtExecutionProvider::GetCapability(const GraphViewer& graph, // Construct subgraph capability from node list std::vector> result; // Get ModelPath - const auto& path_string = graph.ModelPath().ToPathString(); + const auto& path_string = graph.ModelPath().string(); #ifdef _WIN32 - wcstombs_s(nullptr, model_path_, sizeof(model_path_), path_string.c_str(), sizeof(model_path_)); + strncpy_s(model_path_, path_string.c_str(), sizeof(model_path_) - 1); #else - strcpy(model_path_, path_string.c_str()); + strncpy(model_path_, path_string.c_str(), sizeof(model_path_) - 1); #endif + model_path_[sizeof(model_path_) - 1] = '\0'; // If the model consists of only a single "EPContext" contrib op, it means TRT EP can fetch the precompiled engine info from the node and // load the engine directly without having to go through the processes of graph proto reconstruction, calling TRT parser and engine compilation. diff --git a/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider_utils.h b/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider_utils.h index df12d90338782..95abcd1bad2b8 100644 --- a/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider_utils.h +++ b/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider_utils.h @@ -537,10 +537,8 @@ HashValue TRTGenerateId(const GraphViewer& graph_viewer) { }; // Use the model's file name instead of the entire path to avoid cache regeneration if path changes - const auto& model_path_components = main_graph.ModelPath().GetComponents(); - - if (!model_path_components.empty()) { - std::string model_name = PathToUTF8String(model_path_components.back()); + if (main_graph.ModelPath().has_filename()) { + std::string model_name = PathToUTF8String(main_graph.ModelPath().filename()); LOGS_DEFAULT(INFO) << "[TensorRT EP] Model name is " << model_name; // Ensure enough characters are hashed in case model names are too short diff --git a/onnxruntime/core/providers/vitisai/imp/graph.cc b/onnxruntime/core/providers/vitisai/imp/graph.cc index 7cd6da206a6cd..40b396fda6135 100644 --- a/onnxruntime/core/providers/vitisai/imp/graph.cc +++ b/onnxruntime/core/providers/vitisai/imp/graph.cc @@ -112,7 +112,7 @@ void graph_save(const Graph& graph, const std::string& filename, const std::stri if (initializer_size_threshold == std::numeric_limits::max()) { model_proto = model->ToProto(); } else { - model_proto = model->ToGraphProtoWithExternalInitializers(filename_dat, ToPathString(filename), initializer_size_threshold); + model_proto = model->ToGraphProtoWithExternalInitializers(filename_dat, graph.ModelPath(), initializer_size_threshold); } auto& metadata = model->MetaData(); if (!metadata.empty()) { diff --git a/onnxruntime/core/session/custom_ops.cc b/onnxruntime/core/session/custom_ops.cc index d0c46142ac060..7102dbfc750ed 100644 --- a/onnxruntime/core/session/custom_ops.cc +++ b/onnxruntime/core/session/custom_ops.cc @@ -584,7 +584,9 @@ ORT_API_STATUS_IMPL(OrtApis::KernelInfoGetAttribute_tensor, _In_ const OrtKernel auto tensorp = std::make_unique(type, tensor_shape, std::move(alloc_ptr)); // Deserialize TensorProto into pre-allocated, empty Tensor. - status = onnxruntime::utils::TensorProtoToTensor(onnxruntime::Env::Default(), nullptr, tensor_proto, *tensorp); + // TODO: here the TensorProto loses model path information, so it cannot be an external tensor. + status = onnxruntime::utils::TensorProtoToTensor(onnxruntime::Env::Default(), std::filesystem::path(), + tensor_proto, *tensorp); if (!status.IsOK()) { return onnxruntime::ToOrtStatus(status); } diff --git a/onnxruntime/core/session/inference_session.cc b/onnxruntime/core/session/inference_session.cc index 03049c4b51c9c..3ef6490a56ded 100644 --- a/onnxruntime/core/session/inference_session.cc +++ b/onnxruntime/core/session/inference_session.cc @@ -880,7 +880,7 @@ common::Status InferenceSession::RegisterGraphTransformer( return graph_transformer_mgr_.Register(std::move(p_graph_transformer), level); } -common::Status InferenceSession::SaveToOrtFormat(const PathString& filepath) const { +common::Status InferenceSession::SaveToOrtFormat(const std::filesystem::path& filepath) const { ORT_RETURN_IF_NOT(FLATBUFFERS_LITTLEENDIAN, "ort format only supports little-endian machines"); // Get the byte size of the ModelProto and round it to the next MB and use it as flatbuffers' init_size @@ -920,7 +920,7 @@ common::Status InferenceSession::SaveToOrtFormat(const PathString& filepath) con uint8_t* buf = builder.GetBufferPointer(); int size = builder.GetSize(); file.write(reinterpret_cast(buf), size); - ORT_RETURN_IF_NOT(file, "Failed to save ORT format model to file: ", ToUTF8String(filepath)); + ORT_RETURN_IF_NOT(file, "Failed to save ORT format model to file: ", ToUTF8String(filepath.native())); } return Status::OK(); @@ -1272,8 +1272,9 @@ common::Status InferenceSession::TransformGraph(onnxruntime::Graph& graph, bool // for the result of the first step in layout transformation debug_graph_fn = [counter = 1, this](const Graph& graph) mutable { if (graph.GraphProtoSyncNeeded()) { - ORT_THROW_IF_ERROR( - Model::Save(*model_, "post_layout_transform_step_" + std::to_string(counter) + ".onnx")); + std::basic_ostringstream modelpath; + modelpath << ORT_TSTR("post_layout_transform_step_") << counter << ORT_TSTR(".onnx"); + ORT_THROW_IF_ERROR(Model::Save(*model_, modelpath.str())); } // counter is used to denote the step, so increment regardless of whether we wrote out the model in this step. diff --git a/onnxruntime/core/session/inference_session.h b/onnxruntime/core/session/inference_session.h index 77fba90b56b1e..e1cd085d2c271 100644 --- a/onnxruntime/core/session/inference_session.h +++ b/onnxruntime/core/session/inference_session.h @@ -7,6 +7,7 @@ #include #include #include +#include #include "core/common/common.h" #include "core/common/inlined_containers.h" @@ -621,7 +622,7 @@ class InferenceSession { return !custom_schema_registries_.empty(); } - common::Status SaveToOrtFormat(const PathString& filepath) const; + common::Status SaveToOrtFormat(const std::filesystem::path& filepath) const; #endif /** diff --git a/onnxruntime/core/session/provider_bridge_ort.cc b/onnxruntime/core/session/provider_bridge_ort.cc index 408ad7815835f..0494616a9ca0c 100644 --- a/onnxruntime/core/session/provider_bridge_ort.cc +++ b/onnxruntime/core/session/provider_bridge_ort.cc @@ -284,7 +284,7 @@ struct ProviderHostImpl : ProviderHost { Status UnpackTensor(const ONNX_NAMESPACE::TensorProto& tensor, const void* raw_data, size_t raw_data_len, /*out*/ uint32_t* p_data, size_t expected_size) override { return utils::UnpackTensor(tensor, raw_data, raw_data_len, p_data, expected_size); } Status UnpackTensor(const ONNX_NAMESPACE::TensorProto& tensor, const void* raw_data, size_t raw_data_len, /*out*/ int64_t* p_data, size_t expected_size) override { return utils::UnpackTensor(tensor, raw_data, raw_data_len, p_data, expected_size); } Status UnpackTensor(const ONNX_NAMESPACE::TensorProto& tensor, const void* raw_data, size_t raw_data_len, /*out*/ uint64_t* p_data, size_t expected_size) override { return utils::UnpackTensor(tensor, raw_data, raw_data_len, p_data, expected_size); } - Status UnpackInitializerData(const ONNX_NAMESPACE::TensorProto& tensor, const Path& model_path, + Status UnpackInitializerData(const ONNX_NAMESPACE::TensorProto& tensor, const std::filesystem::path& model_path, /*out*/ std::vector& unpacked_tensor) override { return utils::UnpackInitializerData(tensor, model_path, unpacked_tensor); } @@ -1024,7 +1024,7 @@ struct ProviderHostImpl : ProviderHost { const std::string& NodeUnit__Name(const NodeUnit* p) noexcept override { return p->Name(); } int NodeUnit__SinceVersion(const NodeUnit* p) noexcept override { return p->SinceVersion(); } NodeIndex NodeUnit__Index(const NodeUnit* p) noexcept override { return p->Index(); } - const Path& NodeUnit__ModelPath(const NodeUnit* p) noexcept override { return p->ModelPath(); } + const std::filesystem::path& NodeUnit__ModelPath(const NodeUnit* p) noexcept override { return p->ModelPath(); } ProviderType NodeUnit__GetExecutionProviderType(const NodeUnit* p) noexcept override { return p->GetExecutionProviderType(); } @@ -1064,7 +1064,7 @@ struct ProviderHostImpl : ProviderHost { void Model__operator_delete(Model* p) override { delete p; } Graph& Model__MainGraph(Model* p) override { return p->MainGraph(); } std::unique_ptr Model__ToProto(Model* p) override { return std::make_unique(p->ToProto()); } - std::unique_ptr Model__ToGraphProtoWithExternalInitializers(Model* p, const std::string& external_file_name, const PathString& file_path, size_t initializer_size_threshold) override { return std::make_unique(p->ToGraphProtoWithExternalInitializers(external_file_name, file_path, initializer_size_threshold)); }; + std::unique_ptr Model__ToGraphProtoWithExternalInitializers(Model* p, const std::filesystem::path& external_file_name, const std::filesystem::path& file_path, size_t initializer_size_threshold) override { return std::make_unique(p->ToGraphProtoWithExternalInitializers(external_file_name, file_path, initializer_size_threshold)); }; const ModelMetaData& Model__MetaData(const Model* p) const noexcept override { return p->MetaData(); }; Status Model__Load(const PathString& file_path, /*out*/ ONNX_NAMESPACE::ModelProto& model_proto) override { return Model::Load(file_path, model_proto); } @@ -1101,7 +1101,7 @@ struct ProviderHostImpl : ProviderHost { const Graph* Graph__ParentGraph(const Graph* p) const override { return p->ParentGraph(); } Graph* Graph__MutableParentGraph(Graph* p) override { return p->MutableParentGraph(); } const std::string& Graph__Name(const Graph* p) const noexcept override { return p->Name(); } - const Path& Graph__ModelPath(const Graph* p) const override { return p->ModelPath(); } + const std::filesystem::path& Graph__ModelPath(const Graph* p) const override { return p->ModelPath(); } const std::vector& Graph__GetInputsIncludingInitializers(const Graph* p) const noexcept override { return p->GetInputsIncludingInitializers(); } bool Graph__IsSubgraph(const Graph* p) override { return p->IsSubgraph(); } const Node* Graph__GetProducerNode(const Graph* p, const std::string& node_arg_name) const override { return p->GetProducerNode(node_arg_name); } @@ -1157,7 +1157,7 @@ struct ProviderHostImpl : ProviderHost { } const std::string& GraphViewer__Name(const GraphViewer* p) noexcept override { return p->Name(); } - const Path& GraphViewer__ModelPath(const GraphViewer* p) noexcept override { return p->ModelPath(); } + const std::filesystem::path& GraphViewer__ModelPath(const GraphViewer* p) noexcept override { return p->ModelPath(); } const Node* GraphViewer__GetNode(const GraphViewer* p, NodeIndex node_index) override { return p->GetNode(node_index); } const NodeArg* GraphViewer__GetNodeArg(const GraphViewer* p, const std::string& name) override { return p->GetNodeArg(name); } diff --git a/onnxruntime/test/flatbuffers/flatbuffer_utils_test.cc b/onnxruntime/test/flatbuffers/flatbuffer_utils_test.cc index f36dbaf3d1aca..7289f92c65663 100644 --- a/onnxruntime/test/flatbuffers/flatbuffer_utils_test.cc +++ b/onnxruntime/test/flatbuffers/flatbuffer_utils_test.cc @@ -230,7 +230,7 @@ TEST(FlatbufferUtilsTest, ExternalWriteReadWithLoadInitializers) { std::vector> fbs_tensors; for (const auto& initializer : initializers) { flatbuffers::Offset fbs_tensor; - ASSERT_STATUS_OK(SaveInitializerOrtFormat(builder, initializer, Path(), fbs_tensor, writer)); + ASSERT_STATUS_OK(SaveInitializerOrtFormat(builder, initializer, std::filesystem::path(), fbs_tensor, writer)); fbs_tensors.push_back(fbs_tensor); } @@ -313,7 +313,7 @@ TEST(FlatbufferUtilsTest, ExternalWriteReadWithLoadOrtTensor) { std::vector> fbs_tensors; for (const auto& initializer : initializers) { flatbuffers::Offset fbs_tensor; - ASSERT_STATUS_OK(SaveInitializerOrtFormat(builder, initializer, Path(), fbs_tensor, writer)); + ASSERT_STATUS_OK(SaveInitializerOrtFormat(builder, initializer, std::filesystem::path(), fbs_tensor, writer)); fbs_tensors.push_back(fbs_tensor); } diff --git a/onnxruntime/test/framework/allocation_planner_test.cc b/onnxruntime/test/framework/allocation_planner_test.cc index 72eee5aca2638..bf15a9d35b56a 100644 --- a/onnxruntime/test/framework/allocation_planner_test.cc +++ b/onnxruntime/test/framework/allocation_planner_test.cc @@ -1855,7 +1855,7 @@ TEST_F(PlannerTest, ParaPlanCreation) { status = sess.RegisterExecutionProvider(DefaultCpuExecutionProvider()); ASSERT_TRUE(status.IsOK()); - ASSERT_TRUE(model.Save(model, "./simplified_ssd.onnx").IsOK()); + ASSERT_TRUE(model.Save(model, ORT_TSTR("./simplified_ssd.onnx")).IsOK()); std::string s1; const bool rc = model.ToProto().SerializeToString(&s1); diff --git a/onnxruntime/test/framework/inference_session_test.cc b/onnxruntime/test/framework/inference_session_test.cc index d0520ebbcba5a..84389c1d9711c 100644 --- a/onnxruntime/test/framework/inference_session_test.cc +++ b/onnxruntime/test/framework/inference_session_test.cc @@ -1278,7 +1278,7 @@ TEST(InferenceSessionTests, TestOptionalInputs) { } } -static void CreateFuseOpModel(const std::string& model_file_name) { +static void CreateFuseOpModel(const PathString& model_file_name) { onnxruntime::Model model("graph_1", false, ModelMetaData(), PathString(), IOnnxRuntimeOpSchemaRegistryList(), {{kOnnxDomain, 12}}, {}, DefaultLoggingManager().DefaultLogger()); auto& graph = model.MainGraph(); @@ -1312,7 +1312,7 @@ static void CreateFuseOpModel(const std::string& model_file_name) { } TEST(ExecutionProviderTest, FunctionTest) { - std::string model_file_name = "execution_provider_test_graph.onnx"; + PathString model_file_name = ORT_TSTR("execution_provider_test_graph.onnx"); CreateFuseOpModel(model_file_name); SessionOptions so; @@ -1365,7 +1365,7 @@ TEST(ExecutionProviderTest, FunctionTest) { } TEST(ExecutionProviderTest, ShapeInferenceForFusedFunctionTest) { - std::string model_file_name = "fused_node_shape_inference_test_graph.onnx"; + PathString model_file_name = ORT_TSTR("fused_node_shape_inference_test_graph.onnx"); CreateFuseOpModel(model_file_name); @@ -1393,7 +1393,7 @@ TEST(ExecutionProviderTest, ShapeInferenceForFusedFunctionTest) { } TEST(ExecutionProviderTest, OpKernelInfoCanReadConfigOptions) { - std::string model_file_name = "OpKernelInfoCanReadConfigOptions.onnx"; + PathString model_file_name = ORT_TSTR("OpKernelInfoCanReadConfigOptions.onnx"); CreateFuseOpModel(model_file_name); SessionOptions so; @@ -1580,7 +1580,7 @@ TEST(InferenceSessionTests, Test3LayerNestedSubgraph) { auto status = graph.Resolve(); ASSERT_TRUE(status.IsOK()); - std::string model_file_name = "3-layer-nested-subgraph-test.onnx"; + PathString model_file_name = ORT_TSTR("3-layer-nested-subgraph-test.onnx"); status = onnxruntime::Model::Save(model, model_file_name); ASSERT_TRUE(status.IsOK()); @@ -1732,7 +1732,7 @@ TEST(InferenceSessionTests, Test2LayerNestedSubgraph) { auto status = graph.Resolve(); ASSERT_TRUE(status.IsOK()); - std::string model_file_name = "2-layer-nested-subgraph-test.onnx"; + PathString model_file_name = ORT_TSTR("2-layer-nested-subgraph-test.onnx"); status = onnxruntime::Model::Save(model, model_file_name); ASSERT_TRUE(status.IsOK()); diff --git a/onnxruntime/test/framework/save_model_with_external_initializers.cc b/onnxruntime/test/framework/save_model_with_external_initializers.cc index 19c7bf476e6e1..447b0edef879b 100644 --- a/onnxruntime/test/framework/save_model_with_external_initializers.cc +++ b/onnxruntime/test/framework/save_model_with_external_initializers.cc @@ -1,6 +1,8 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. +#include "core/common/common.h" +#include "core/common/status.h" #include "core/common/path_string.h" #include "core/framework/data_types.h" #include "core/graph/model.h" @@ -17,75 +19,77 @@ using namespace onnxruntime; namespace onnxruntime { namespace test { -void LoadSaveAndCompareModel(const std::string& input_onnx, - const std::string& input_external_init_file, - const std::string& output_onnx, - const std::string& output_external_init_file, - size_t initializer_size_threshold) { +Status LoadSaveAndCompareModel(const std::filesystem::path& input_onnx, + const std::filesystem::path& input_external_init_file, + const std::filesystem::path& output_onnx, + const std::filesystem::path& output_external_init_file, + size_t initializer_size_threshold) { + auto logger = DefaultLoggingManager().CreateLogger("LoadSaveAndCompareModel"); std::shared_ptr model; - ASSERT_STATUS_OK(Model::Load(ToPathString(input_onnx), model, nullptr, DefaultLoggingManager().DefaultLogger())); - std::remove(output_onnx.c_str()); - std::remove(output_external_init_file.c_str()); - ASSERT_STATUS_OK(Model::SaveWithExternalInitializers(*model, ToPathString(output_onnx), output_external_init_file, initializer_size_threshold)); + ORT_RETURN_IF_ERROR(Model::Load(input_onnx, model, nullptr, *logger)); + std::filesystem::remove(output_onnx); + std::filesystem::remove(output_external_init_file); + ORT_RETURN_IF_ERROR(Model::SaveWithExternalInitializers(*model, output_onnx, output_external_init_file, initializer_size_threshold)); std::shared_ptr model_from_external; - ASSERT_STATUS_OK(Model::Load(ToPathString(output_onnx), model_from_external, nullptr, DefaultLoggingManager().DefaultLogger())); + ORT_RETURN_IF_ERROR(Model::Load(output_onnx.native(), model_from_external, nullptr, *logger)); Graph& graph = model->MainGraph(); // Perform shape inference on the graph, if this succeeds then it means that we could correctly read the // integer initializers used by reshape and transpose. - ASSERT_STATUS_OK(graph.Resolve()); + ORT_RETURN_IF_ERROR(graph.Resolve()); Graph& graph_from_external = model_from_external->MainGraph(); InitializedTensorSet initializers = graph.GetAllInitializedTensors(); InitializedTensorSet initializers_from_external = graph_from_external.GetAllInitializedTensors(); - ASSERT_EQ(initializers.size(), initializers_from_external.size()); + ORT_RETURN_IF_NOT(initializers.size() == initializers_from_external.size(), "size mismatch"); // Compare the initializers of the two versions. - Path model_path{}; - Path external_data_path{}; - for (auto i : initializers) { + std::filesystem::path model_path{}; + std::filesystem::path external_data_path{}; + for (const auto& i : initializers) { const std::string kInitName = i.first; const ONNX_NAMESPACE::TensorProto* tensor_proto = i.second; const ONNX_NAMESPACE::TensorProto* from_external_tensor_proto = initializers_from_external[kInitName]; std::vector tensor_proto_data; - model_path = Path::Parse(ToPathString(input_onnx)); - external_data_path = (input_external_init_file.size()) ? model_path.ParentPath().Append(Path::Parse(ToPathString(input_external_init_file))) : Path(); - ORT_THROW_IF_ERROR(utils::UnpackInitializerData(*tensor_proto, external_data_path, tensor_proto_data)); + model_path = input_onnx; + external_data_path = (!input_external_init_file.empty()) ? (model_path.parent_path() / input_external_init_file) : std::filesystem::path(); + ORT_RETURN_IF_ERROR(utils::UnpackInitializerData(*tensor_proto, external_data_path, tensor_proto_data)); size_t tensor_proto_size = tensor_proto_data.size(); std::vector from_external_tensor_proto_data; - model_path = Path::Parse(ToPathString(output_onnx)); - external_data_path = model_path.ParentPath().Append(Path::Parse(ToPathString(output_external_init_file))); - ORT_THROW_IF_ERROR(utils::UnpackInitializerData(*from_external_tensor_proto, model_path, from_external_tensor_proto_data)); + model_path = output_onnx; + external_data_path = model_path.parent_path() / output_external_init_file; + ORT_RETURN_IF_ERROR(utils::UnpackInitializerData(*from_external_tensor_proto, model_path, from_external_tensor_proto_data)); size_t from_external_tensor_proto_size = from_external_tensor_proto_data.size(); if (from_external_tensor_proto_size < initializer_size_threshold) { // 'Small' tensors should be embedded in the onnx file. - EXPECT_EQ(from_external_tensor_proto->data_location(), ONNX_NAMESPACE::TensorProto_DataLocation::TensorProto_DataLocation_DEFAULT); + ORT_RETURN_IF_NOT(from_external_tensor_proto->data_location() == ONNX_NAMESPACE::TensorProto_DataLocation::TensorProto_DataLocation_DEFAULT, "location mismatch"); } else { // 'Large' tensors should be added to the external binary file. - EXPECT_EQ(from_external_tensor_proto->data_location(), ONNX_NAMESPACE::TensorProto_DataLocation::TensorProto_DataLocation_EXTERNAL); + ORT_RETURN_IF_NOT(from_external_tensor_proto->data_location() == ONNX_NAMESPACE::TensorProto_DataLocation::TensorProto_DataLocation_EXTERNAL, "location mismatch"); } - ASSERT_EQ(tensor_proto_size, from_external_tensor_proto_size); - EXPECT_EQ(memcmp(tensor_proto_data.data(), from_external_tensor_proto_data.data(), tensor_proto_size), 0); + ORT_RETURN_IF_NOT(tensor_proto_size == from_external_tensor_proto_size, "size mismatch"); + ORT_RETURN_IF_NOT(memcmp(tensor_proto_data.data(), from_external_tensor_proto_data.data(), tensor_proto_size) == 0, "data mismatch"); } // Cleanup. - ASSERT_EQ(std::remove(output_onnx.c_str()), 0); - ASSERT_EQ(std::remove(PathToUTF8String(external_data_path.ToPathString()).c_str()), 0); + ORT_RETURN_IF_NOT(std::filesystem::remove(output_onnx), "delete file failed"); + ORT_RETURN_IF_NOT(std::filesystem::remove(external_data_path), "delete file failed"); + return Status::OK(); } // Original model does not have external initializers TEST(SaveWithExternalInitializers, Mnist) { - LoadSaveAndCompareModel("testdata/mnist.onnx", "", "testdata/mnist_with_external_initializers.onnx", "mnist_external_initializers.bin", 100); + ASSERT_STATUS_OK(LoadSaveAndCompareModel(ORT_TSTR("testdata/mnist.onnx"), ORT_TSTR(""), ORT_TSTR("testdata/mnist_with_external_initializers.onnx"), ORT_TSTR("mnist_external_initializers.bin"), 100)); } // Original model has external initializers TEST(SaveWithExternalInitializers, ModelWithOriginalExternalData) { - LoadSaveAndCompareModel("testdata/model_with_orig_ext_data.onnx", "model_with_orig_ext_data.onnx.data", "testdata/model_with_new_external_initializers.onnx", "model_with_new_external_initializers.bin", 0); + ASSERT_STATUS_OK(LoadSaveAndCompareModel(ORT_TSTR("testdata/model_with_orig_ext_data.onnx"), ORT_TSTR("model_with_orig_ext_data.onnx.data"), ORT_TSTR("testdata/model_with_new_external_initializers.onnx"), ORT_TSTR("model_with_new_external_initializers.bin"), 0)); } } // namespace test diff --git a/onnxruntime/test/framework/sparse_kernels_test.cc b/onnxruntime/test/framework/sparse_kernels_test.cc index 80f23b054a4ad..fa42bb6e96cd5 100644 --- a/onnxruntime/test/framework/sparse_kernels_test.cc +++ b/onnxruntime/test/framework/sparse_kernels_test.cc @@ -795,7 +795,7 @@ static void TestConversion(bool use_1D_indices, int32_t indices_type, TensorProto dense; // Path is required for loading external data (if any) // When path is empty it will look for the data in current dir - ASSERT_STATUS_OK(utils::ConstantNodeProtoToTensorProto(node, Path(), dense)); + ASSERT_STATUS_OK(utils::ConstantNodeProtoToTensorProto(node, std::filesystem::path(), dense)); gsl::span expected_span = gsl::make_span(expected.data(), expected.size()); checker(expected_span, dense); @@ -810,7 +810,7 @@ static void TestConversionAllZeros(bool use_1D_indices, TensorProto dense; // Path is required for loading external data (if any) // When path is empty it will look for the data in current dir - ASSERT_STATUS_OK(utils::ConstantNodeProtoToTensorProto(node, Path(), dense)); + ASSERT_STATUS_OK(utils::ConstantNodeProtoToTensorProto(node, std::filesystem::path(), dense)); gsl::span expected_span = gsl::make_span(expected.data(), expected.size()); checker(expected_span, dense); @@ -1109,30 +1109,31 @@ void RawSparseDataChecker(gsl::span expected_bfloat, } template -static void TestDenseToSparseConversionValues(size_t indices_start, - std::function& values, TensorProto& tp)> inserter, - std::function expected, - gsl::span expected_indicies, - const SparseTensorProto& actual)> - checker) { +static Status TestDenseToSparseConversionValues(size_t indices_start, + std::function& values, TensorProto& tp)> inserter, + std::function expected, + gsl::span expected_indicies, + const SparseTensorProto& actual)> + checker) { std::vector expected_values; std::vector expected_indicies; // Path is required for loading external data // Using empty path here since the data is not external - Path model_path; + std::filesystem::path model_path; TensorProto dense_tensor = CreateDenseTensor(indices_start, inserter, expected_values, expected_indicies); SparseTensorProto sparse_tensor; - ASSERT_STATUS_OK(utils::DenseTensorToSparseTensorProto(dense_tensor, model_path, sparse_tensor)); + ORT_RETURN_IF_ERROR(utils::DenseTensorToSparseTensorProto(dense_tensor, model_path, sparse_tensor)); gsl::span expected_values_span = gsl::make_span(expected_values.data(), expected_values.size()); gsl::span expected_ind_span = gsl::make_span(expected_indicies.data(), expected_indicies.size()); checker(expected_values_span, expected_ind_span, sparse_tensor); + return Status::OK(); } template -static void TestDenseAllZerosToSparseConversion( +static Status TestDenseAllZerosToSparseConversion( std::function& values, TensorProto& tp)> inserter, std::function expected, gsl::span expected_indicies, @@ -1142,55 +1143,56 @@ static void TestDenseAllZerosToSparseConversion( std::vector expected_indicies; // Path is required for loading external data // Using empty path here since the data is not external - Path model_path; + std::filesystem::path model_path; TensorProto dense_tensor = CreateDenseTensorAllZeros(inserter); SparseTensorProto sparse_tensor; - ASSERT_STATUS_OK(utils::DenseTensorToSparseTensorProto(dense_tensor, model_path, sparse_tensor)); + ORT_RETURN_IF_ERROR(utils::DenseTensorToSparseTensorProto(dense_tensor, model_path, sparse_tensor)); gsl::span expected_values_span = gsl::make_span(expected_values.data(), expected_values.size()); gsl::span expected_ind_span = gsl::make_span(expected_indicies.data(), expected_indicies.size()); checker(expected_values_span, expected_ind_span, sparse_tensor); + return Status::OK(); } template -static void TestDenseToSparseConversion(size_t indices_start, - std::function& values, TensorProto& tp)> inserter, - std::function expected, - gsl::span expected_indicies, - const SparseTensorProto& actual)> - checker) { - TestDenseToSparseConversionValues(indices_start, inserter, checker); - TestDenseAllZerosToSparseConversion(inserter, checker); +static Status TestDenseToSparseConversion(size_t indices_start, + std::function& values, TensorProto& tp)> inserter, + std::function expected, + gsl::span expected_indicies, + const SparseTensorProto& actual)> + checker) { + ORT_RETURN_IF_ERROR(TestDenseToSparseConversionValues(indices_start, inserter, checker)); + return TestDenseAllZerosToSparseConversion(inserter, checker); } TEST(SparseTensorConversionTests, TestDenseToSparseConversion) { // This one will test indices that are less than max int8 value // which should result in int8 indices - TestDenseToSparseConversion( + ASSERT_STATUS_OK(TestDenseToSparseConversion( 20U, [](const std::vector& values, TensorProto& tp) { tp.set_data_type(TensorProto_DataType_FLOAT); tp.set_name("dense_float"); tp.mutable_float_data()->Add(values.cbegin(), values.cend()); }, - RawSparseDataChecker); + RawSparseDataChecker)); // This one will test indices that are max(int8) < ind < max(int16) value // which should result in int16 indices - TestDenseToSparseConversion( + ASSERT_STATUS_OK(TestDenseToSparseConversion( static_cast(std::numeric_limits::max()) + 20U, [](const std::vector& values, TensorProto& tp) { tp.set_data_type(TensorProto_DataType_DOUBLE); tp.set_name("dense_double"); tp.mutable_double_data()->Add(values.cbegin(), values.cend()); }, - RawSparseDataChecker); + RawSparseDataChecker)); // This one will test indices that are max(int16) < ind < max(int32) value // which should result in int32 indices - TestDenseToSparseConversion( + ASSERT_STATUS_OK(TestDenseToSparseConversion( static_cast(std::numeric_limits::max()) + 20U, [](const std::vector& values, TensorProto& tp) { tp.set_data_type(TensorProto_DataType_BFLOAT16); @@ -1199,12 +1201,12 @@ TEST(SparseTensorConversionTests, TestDenseToSparseConversion) { tp.mutable_int32_data()->Add(v.val); } }, - RawSparseDataChecker); + RawSparseDataChecker)); // Protobuf can not hold anything more than 2Gb and it overflows. Can't test 64-bit indices // on conversion unless explicitly created. // which should result in int32 indices - TestDenseToSparseConversion( + ASSERT_STATUS_OK(TestDenseToSparseConversion( 20U, [](const std::vector& values, TensorProto& tp) { tp.set_data_type(TensorProto_DataType_FLOAT16); @@ -1213,78 +1215,78 @@ TEST(SparseTensorConversionTests, TestDenseToSparseConversion) { tp.mutable_int32_data()->Add(v.val); } }, - RawSparseDataChecker); + RawSparseDataChecker)); - TestDenseToSparseConversion( + ASSERT_STATUS_OK(TestDenseToSparseConversion( 20U, [](const std::vector& values, TensorProto& tp) { tp.set_name("dense_int16"); tp.set_data_type(TensorProto_DataType_INT16); tp.mutable_int32_data()->Add(values.cbegin(), values.cend()); }, - RawSparseDataChecker); + RawSparseDataChecker)); - TestDenseToSparseConversion( + ASSERT_STATUS_OK(TestDenseToSparseConversion( 20U, [](const std::vector& values, TensorProto& tp) { tp.set_name("dense_uint16"); tp.set_data_type(TensorProto_DataType_UINT16); tp.mutable_int32_data()->Add(values.cbegin(), values.cend()); }, - RawSparseDataChecker); + RawSparseDataChecker)); - TestDenseToSparseConversion( + ASSERT_STATUS_OK(TestDenseToSparseConversion( 20U, [](const std::vector& values, TensorProto& tp) { tp.set_name("dense_int32"); tp.set_data_type(TensorProto_DataType_INT32); tp.mutable_int32_data()->Add(values.cbegin(), values.cend()); }, - RawSparseDataChecker); + RawSparseDataChecker)); - TestDenseToSparseConversion( + ASSERT_STATUS_OK(TestDenseToSparseConversion( 20U, [](const std::vector& values, TensorProto& tp) { tp.set_name("dense_uint32"); tp.set_data_type(TensorProto_DataType_UINT32); tp.mutable_uint64_data()->Add(values.cbegin(), values.cend()); }, - RawSparseDataChecker); + RawSparseDataChecker)); - TestDenseToSparseConversion( + ASSERT_STATUS_OK(TestDenseToSparseConversion( 20U, [](const std::vector& values, TensorProto& tp) { tp.set_name("dense_int64"); tp.set_data_type(TensorProto_DataType_INT64); tp.mutable_int64_data()->Add(values.cbegin(), values.cend()); }, - RawSparseDataChecker); + RawSparseDataChecker)); - TestDenseToSparseConversion( + ASSERT_STATUS_OK(TestDenseToSparseConversion( 20U, [](const std::vector& values, TensorProto& tp) { tp.set_name("dense_uint64"); tp.set_data_type(TensorProto_DataType_UINT64); tp.mutable_uint64_data()->Add(values.cbegin(), values.cend()); }, - RawSparseDataChecker); + RawSparseDataChecker)); - TestDenseToSparseConversion( + ASSERT_STATUS_OK(TestDenseToSparseConversion( 20U, [](const std::vector& values, TensorProto& tp) { tp.set_name("dense_int8"); tp.set_data_type(TensorProto_DataType_INT8); tp.mutable_int32_data()->Add(values.cbegin(), values.cend()); }, - RawSparseDataChecker); + RawSparseDataChecker)); - TestDenseToSparseConversion( + ASSERT_STATUS_OK(TestDenseToSparseConversion( 20U, [](const std::vector& values, TensorProto& tp) { tp.set_name("dense_int64"); RawDataWriter(values, tp, TensorProto_DataType_UINT8); }, - RawSparseDataChecker); + RawSparseDataChecker)); } TEST(SparseTensorConversionTests, CsrConversion) { diff --git a/onnxruntime/test/framework/tensorutils_test.cc b/onnxruntime/test/framework/tensorutils_test.cc index 42f82c374348e..05bdb3a9a033d 100644 --- a/onnxruntime/test/framework/tensorutils_test.cc +++ b/onnxruntime/test/framework/tensorutils_test.cc @@ -21,7 +21,7 @@ namespace test { // T must be float for double, and it must match with the 'type' argument template -void TestUnpackFloatTensor(TensorProto_DataType type, const Path& model_path) { +void TestUnpackFloatTensor(TensorProto_DataType type, const std::filesystem::path& model_path) { TensorProto float_tensor_proto; float_tensor_proto.set_data_type(type); T f[4] = {1.1f, 2.2f, 3.3f, 4.4f}; @@ -45,7 +45,7 @@ TEST(TensorProtoUtilsTest, UnpackTensor) { // Path is required for loading external data. // Using empty path here since this test does not test // external data utils - Path model_path; + std::filesystem::path model_path; bool_tensor_proto.set_data_type(TensorProto_DataType_BOOL); bool_tensor_proto.add_int32_data(1); @@ -142,7 +142,7 @@ void CreateTensorWithExternalData(TensorProto_DataType type, const std::vector -void UnpackAndValidate(const TensorProto& tensor_proto, const Path& model_path, const std::vector& test_data) { +void UnpackAndValidate(const TensorProto& tensor_proto, const std::filesystem::path& model_path, const std::vector& test_data) { // Unpack tensor with external data std::vector val(test_data.size()); auto st = utils::UnpackTensor(tensor_proto, model_path, val.data(), test_data.size()); @@ -155,7 +155,7 @@ void UnpackAndValidate(const TensorProto& tensor_proto, const Path& model_path, } template <> -void UnpackAndValidate(const TensorProto& tensor_proto, const Path& model_path, +void UnpackAndValidate(const TensorProto& tensor_proto, const std::filesystem::path& model_path, const std::vector& test_data) { // Unpack tensor with external data auto arr = std::make_unique(test_data.size()); @@ -169,7 +169,7 @@ void UnpackAndValidate(const TensorProto& tensor_proto, const Path& model_ } template -void TestUnpackExternalTensor(TensorProto_DataType type, const Path& model_path) { +void TestUnpackExternalTensor(TensorProto_DataType type, const std::filesystem::path& model_path) { // Create external data std::basic_string filename(ORT_TSTR("tensor_XXXXXX")); TensorProto tensor_proto; @@ -181,7 +181,7 @@ void TestUnpackExternalTensor(TensorProto_DataType type, const Path& model_path) } } // namespace TEST(TensorProtoUtilsTest, UnpackTensorWithExternalData) { - Path model_path; + std::filesystem::path model_path; TestUnpackExternalTensor(TensorProto_DataType_FLOAT, model_path); TestUnpackExternalTensor(TensorProto_DataType_DOUBLE, model_path); TestUnpackExternalTensor(TensorProto_DataType_INT32, model_path); @@ -225,7 +225,7 @@ static void TestConstantNodeConversion(const std::string& attrib_name, [&input, &add_data](AttributeProto& attrib) { add_data(attrib, input); }); TensorProto tp; - Path model_path; + std::filesystem::path model_path; EXPECT_STATUS_OK(utils::ConstantNodeProtoToTensorProto(c, model_path, tp)); EXPECT_THAT(get_data(tp), ::testing::ContainerEq(input)); @@ -311,7 +311,7 @@ template static void TestConstantNodeConversionWithExternalData(TensorProto_DataType type) { // Create a constant node with external data auto test_data = CreateValues(); - Path model_path; + std::filesystem::path model_path; PathString tensor_filename(ORT_TSTR("tensor_XXXXXX")); auto c = CreateConstantNodeWithExternalData(type, tensor_filename, test_data); std::unique_ptr file_deleter(const_cast(tensor_filename.c_str()), diff --git a/onnxruntime/test/framework/test_tensor_loader.cc b/onnxruntime/test/framework/test_tensor_loader.cc index 71d70abceb82e..17edad73085c9 100644 --- a/onnxruntime/test/framework/test_tensor_loader.cc +++ b/onnxruntime/test/framework/test_tensor_loader.cc @@ -34,7 +34,7 @@ TEST(CApiTensorTest, load_simple_float_tensor_not_enough_space) { OrtMemoryInfo cpu_memory_info(onnxruntime::CPU, OrtDeviceAllocator, OrtDevice(), 0, OrtMemTypeDefault); ASSERT_STATUS_NOT_OK( - utils::TensorProtoToOrtValue(Env::Default(), nullptr, p, + utils::TensorProtoToOrtValue(Env::Default(), std::filesystem::path(), p, MemBuffer(output.data(), output.size() * sizeof(float), cpu_memory_info), value)); } @@ -55,7 +55,7 @@ TEST(CApiTensorTest, load_simple_float_tensor_membuffer) { OrtValue value; OrtMemoryInfo cpu_memory_info(onnxruntime::CPU, OrtDeviceAllocator, OrtDevice(), 0, OrtMemTypeDefault); ASSERT_STATUS_OK( - utils::TensorProtoToOrtValue(Env::Default(), nullptr, p, + utils::TensorProtoToOrtValue(Env::Default(), std::filesystem::path(), p, MemBuffer(output.data(), output.size() * sizeof(float), cpu_memory_info), value)); float* real_output; @@ -83,7 +83,7 @@ TEST(CApiTensorTest, load_simple_float_tensor_allocator) { AllocatorPtr tmp_allocator = std::make_shared(); OrtValue value; - ASSERT_STATUS_OK(utils::TensorProtoToOrtValue(Env::Default(), nullptr, p, tmp_allocator, value)); + ASSERT_STATUS_OK(utils::TensorProtoToOrtValue(Env::Default(), std::filesystem::path(), p, tmp_allocator, value)); float* real_output; auto ort_st = g_ort->GetTensorMutableData(&value, (void**)&real_output); @@ -139,7 +139,7 @@ static void run_external_data_test() { OrtValue value; OrtMemoryInfo cpu_memory_info(onnxruntime::CPU, OrtDeviceAllocator, OrtDevice(), 0, OrtMemTypeDefault); ASSERT_STATUS_OK(utils::TensorProtoToOrtValue( - Env::Default(), nullptr, p, MemBuffer(output.data(), output.size() * sizeof(float), cpu_memory_info), value)); + Env::Default(), std::filesystem::path(), p, MemBuffer(output.data(), output.size() * sizeof(float), cpu_memory_info), value)); float* real_output; auto ort_st = g_ort->GetTensorMutableData(&value, (void**)&real_output); @@ -190,7 +190,7 @@ TEST(CApiTensorTest, load_huge_tensor_with_external_data) { OrtValue value; OrtMemoryInfo cpu_memory_info(onnxruntime::CPU, OrtDeviceAllocator, OrtDevice(), 0, OrtMemTypeDefault); ASSERT_STATUS_OK( - utils::TensorProtoToOrtValue(Env::Default(), nullptr, p, + utils::TensorProtoToOrtValue(Env::Default(), std::filesystem::path(), p, MemBuffer(output.data(), output.size() * sizeof(int), cpu_memory_info), value)); int* buffer; diff --git a/onnxruntime/test/ir/graph_test.cc b/onnxruntime/test/ir/graph_test.cc index ff10765741bbe..4766ef6fbc621 100644 --- a/onnxruntime/test/ir/graph_test.cc +++ b/onnxruntime/test/ir/graph_test.cc @@ -388,7 +388,7 @@ TEST_F(GraphTest, UnusedValueInfoSerializes) { std::shared_ptr model; ASSERT_STATUS_OK(Model::Load(std::move(m), model, nullptr, *logger_)); model->MainGraph().SetGraphProtoSyncNeeded(); - EXPECT_TRUE(Model::Save(*model, "graph_with_unused_value_info.onnx").IsOK()); + EXPECT_TRUE(Model::Save(*model, ORT_TSTR("graph_with_unused_value_info.onnx")).IsOK()); } TEST_F(GraphTest, WrongOpset) { @@ -762,7 +762,7 @@ TEST_F(GraphTest, GraphConstruction_CheckIsAcyclic) { auto status = graph.Resolve(); EXPECT_TRUE(status.IsOK()) << status.ErrorMessage(); - EXPECT_TRUE(Model::Save(model, "graph_1.onnx").IsOK()); + EXPECT_TRUE(Model::Save(model, ORT_TSTR("graph_1.onnx")).IsOK()); std::shared_ptr model2; EXPECT_TRUE(Model::Load(ORT_TSTR("graph_1.onnx"), model2, nullptr, *logger_).IsOK()); @@ -1476,7 +1476,7 @@ TEST_F(GraphTest, GraphConstruction_TypeInference) { EXPECT_EQ("node_4_out_1", graph.GetOutputs()[0]->Name()); EXPECT_EQ(2u, graph.GetInputs().size()); - EXPECT_TRUE(Model::Save(model, "model_x.onnx").IsOK()); + EXPECT_TRUE(Model::Save(model, ORT_TSTR("model_x.onnx")).IsOK()); std::shared_ptr loaded_model; EXPECT_TRUE(Model::Load(ORT_TSTR("model_x.onnx"), loaded_model, nullptr, *logger_).IsOK()); EXPECT_EQ(2u, loaded_model->MainGraph().GetInputs().size()); diff --git a/onnxruntime/test/optimizer/initializer_test.cc b/onnxruntime/test/optimizer/initializer_test.cc index ee93cfaa67e2a..9e55d9b2ef921 100644 --- a/onnxruntime/test/optimizer/initializer_test.cc +++ b/onnxruntime/test/optimizer/initializer_test.cc @@ -51,12 +51,12 @@ TEST(OptimizerInitializerTest, LoadExternalData) { return tensor_data; }(); const gsl::span tensor_data_span = gsl::make_span(tensor_data); - const auto tensor_data_dir_path = Path::Parse(ToPathString(".")); - const auto tensor_data_dir_relative_path = Path::Parse(ToPathString("OptimizerInitializerTest_LoadExternalData.bin")); + const std::filesystem::path tensor_data_dir_path = ORT_TSTR("."); + const std::filesystem::path tensor_data_dir_relative_path = ORT_TSTR("OptimizerInitializerTest_LoadExternalData.bin"); ScopedFileDeleter file_deleter{}; ASSERT_STATUS_OK(WriteExternalDataFile( - tensor_data_span, (tensor_data_dir_path / tensor_data_dir_relative_path).ToPathString(), file_deleter)); + tensor_data_span, tensor_data_dir_path / tensor_data_dir_relative_path, file_deleter)); const auto tensor_proto_base = [&]() { @@ -65,7 +65,7 @@ TEST(OptimizerInitializerTest, LoadExternalData) { tensor_proto.add_dims(tensor_data.size()); tensor_proto.set_data_type(ONNX_NAMESPACE::TensorProto_DataType_INT32); tensor_proto.set_data_location(ONNX_NAMESPACE::TensorProto_DataLocation_EXTERNAL); - SetTensorProtoExternalData("location", ToUTF8String(tensor_data_dir_relative_path.ToPathString()), tensor_proto); + SetTensorProtoExternalData("location", ToUTF8String(tensor_data_dir_relative_path.native()), tensor_proto); SetTensorProtoExternalData("offset", "0", tensor_proto); SetTensorProtoExternalData("length", std::to_string(tensor_data.size() * sizeof(int32_t)), tensor_proto); return tensor_proto; @@ -95,8 +95,8 @@ TEST(OptimizerInitializerTest, LoadExternalData) { check_initializer_load(0, tensor_data.size() + 1); // bad model paths - EXPECT_THROW(Initializer i(tensor_proto_base, Path{}), OnnxRuntimeException); - EXPECT_THROW(Initializer i(tensor_proto_base, Path::Parse(ToPathString("invalid/directory"))), OnnxRuntimeException); + EXPECT_THROW(Initializer i(tensor_proto_base, std::filesystem::path()), OnnxRuntimeException); + EXPECT_THROW(Initializer i(tensor_proto_base, ORT_TSTR("invalid/directory")), std::filesystem::filesystem_error); // bad length { @@ -165,7 +165,7 @@ void TestInitializerRawData() { tensor_proto.add_dims(4); tensor_proto.set_raw_data(data.data(), data.size() * sizeof(T)); - const Initializer init(tensor_proto, Path()); + const Initializer init(tensor_proto, std::filesystem::path()); for (size_t idx = 0; idx < data.size(); idx++) { EXPECT_EQ(data[idx], init.data()[idx]); @@ -220,35 +220,35 @@ void TestInitializerDataField() { AddData(data, idx, tensor_proto); } - const Initializer init(tensor_proto, Path()); + const Initializer init(tensor_proto, std::filesystem::path()); for (size_t idx = 0; idx < data.size(); idx++) { EXPECT_EQ(data[idx], init.data()[idx]); } } -#define TestInitializerDataFieldSpecialized(type) \ - template <> \ - void TestInitializerDataField() { \ - std::vector data{ \ - 0, 1, 2, 3, \ - 4, 5, 6, 7, \ - 8, 9, 10, 11}; \ - \ - ONNX_NAMESPACE::TensorProto tensor_proto; \ - tensor_proto.set_data_type(GetTensorProtoDataType()); \ - tensor_proto.set_name("OptimizerInitializerTest_DataField"); \ - tensor_proto.add_dims(3); \ - tensor_proto.add_dims(4); \ - for (size_t idx = 0; idx < data.size(); idx++) { \ - tensor_proto.add_##type##_data(data[idx]); \ - } \ - \ - const Initializer init(tensor_proto, Path()); \ - \ - for (size_t idx = 0; idx < data.size(); idx++) { \ - EXPECT_EQ(data[idx], init.data()[idx]); \ - } \ +#define TestInitializerDataFieldSpecialized(type) \ + template <> \ + void TestInitializerDataField() { \ + std::vector data{ \ + 0, 1, 2, 3, \ + 4, 5, 6, 7, \ + 8, 9, 10, 11}; \ + \ + ONNX_NAMESPACE::TensorProto tensor_proto; \ + tensor_proto.set_data_type(GetTensorProtoDataType()); \ + tensor_proto.set_name("OptimizerInitializerTest_DataField"); \ + tensor_proto.add_dims(3); \ + tensor_proto.add_dims(4); \ + for (size_t idx = 0; idx < data.size(); idx++) { \ + tensor_proto.add_##type##_data(data[idx]); \ + } \ + \ + const Initializer init(tensor_proto, std::filesystem::path()); \ + \ + for (size_t idx = 0; idx < data.size(); idx++) { \ + EXPECT_EQ(data[idx], init.data()[idx]); \ + } \ } typedef int64_t int64; diff --git a/onnxruntime/test/optimizer/resnet50_fusion_test.cc b/onnxruntime/test/optimizer/resnet50_fusion_test.cc index 04b11b46e5002..5cb0206156a84 100644 --- a/onnxruntime/test/optimizer/resnet50_fusion_test.cc +++ b/onnxruntime/test/optimizer/resnet50_fusion_test.cc @@ -61,13 +61,13 @@ TEST_F(ResNet50FusionTests, FuseConvIntegrationTest) { ASSERT_STATUS_OK(graph_transformation_mgr_32.Register(std::make_unique(), TransformerLevel::Level3)); ASSERT_STATUS_OK(graph_transformation_mgr_32.Register(std::make_unique(), TransformerLevel::Level3)); ASSERT_STATUS_OK(graph_transformation_mgr_32.ApplyTransformers(fp32_graph, TransformerLevel::Level3, *logger)); - ASSERT_STATUS_OK(Model::Save(*fp32_model, "resnet50_fused.onnx")); + ASSERT_STATUS_OK(Model::Save(*fp32_model, ORT_TSTR("resnet50_fused.onnx"))); onnxruntime::GraphTransformerManager graph_transformation_mgr_16{5}; ASSERT_STATUS_OK(graph_transformation_mgr_16.Register(std::make_unique(), TransformerLevel::Level3)); ASSERT_STATUS_OK(graph_transformation_mgr_16.Register(std::make_unique(), TransformerLevel::Level3)); ASSERT_STATUS_OK(graph_transformation_mgr_16.ApplyTransformers(fp16_graph, TransformerLevel::Level3, *logger)); - ASSERT_STATUS_OK(Model::Save(*fp16_model, "resnet50_fp16_fused.onnx")); + ASSERT_STATUS_OK(Model::Save(*fp16_model, ORT_TSTR("resnet50_fp16_fused.onnx"))); // std::cout << "-------Op Counts After Fusion---------" << std::endl; fp32_op_count = CountOpsInGraph(fp32_graph); fp16_op_count = CountOpsInGraph(fp16_graph); @@ -91,7 +91,7 @@ TEST_F(ResNet50FusionTests, FuseConvAddReluUnitTest) { ASSERT_STATUS_OK(graph_transformation_mgr.Register(std::make_unique(), TransformerLevel::Level3)); ASSERT_STATUS_OK(graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level3, *logger)); op_to_count = CountOpsInGraph(graph); - ASSERT_STATUS_OK(Model::Save(*p_model, "conv_add_relu_fp16_fused.onnx")); + ASSERT_STATUS_OK(Model::Save(*p_model, ORT_TSTR("conv_add_relu_fp16_fused.onnx"))); ASSERT_TRUE(op_to_count["Add"] == 0); // Add removed from graph ASSERT_TRUE(op_to_count["Relu"] == 0); // Relu removed from graph } @@ -109,7 +109,7 @@ TEST_F(ResNet50FusionTests, FuseConvAddUnitTest) { ASSERT_STATUS_OK(graph_transformation_mgr.Register(std::make_unique(), TransformerLevel::Level3)); ASSERT_STATUS_OK(graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level3, *logger)); op_to_count = CountOpsInGraph(graph); - ASSERT_STATUS_OK(Model::Save(*p_model, "conv_add_fp16_fused.onnx")); + ASSERT_STATUS_OK(Model::Save(*p_model, ORT_TSTR("conv_add_fp16_fused.onnx"))); ASSERT_TRUE(op_to_count["Add"] == 0); // Add removed from graph } TEST_F(ResNet50FusionTests, FuseConvReluUnitTest) { @@ -126,9 +126,9 @@ TEST_F(ResNet50FusionTests, FuseConvReluUnitTest) { ASSERT_STATUS_OK(graph_transformation_mgr.Register(std::make_unique(), TransformerLevel::Level3)); ASSERT_STATUS_OK(graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level3, *logger)); op_to_count = CountOpsInGraph(graph); - ASSERT_STATUS_OK(Model::Save(*p_model, "conv_relu_fp16_fused.onnx")); + ASSERT_STATUS_OK(Model::Save(*p_model, ORT_TSTR("conv_relu_fp16_fused.onnx"))); ASSERT_TRUE(op_to_count["Relu"] == 0); // Add removed from graph } #endif // defined(MLAS_F16VEC_INTRINSICS_SUPPORTED) && !defined(DISABLE_CONTRIB_OPS) } // namespace test -} // namespace onnxruntime \ No newline at end of file +} // namespace onnxruntime diff --git a/onnxruntime/test/providers/tensorrt/tensorrt_basic_test.cc b/onnxruntime/test/providers/tensorrt/tensorrt_basic_test.cc index 4d2538c947dcc..2b5b82d0fc16a 100644 --- a/onnxruntime/test/providers/tensorrt/tensorrt_basic_test.cc +++ b/onnxruntime/test/providers/tensorrt/tensorrt_basic_test.cc @@ -66,7 +66,7 @@ void VerifyOutputs(const std::vector& fetches, const std::vector dims, bool add_non_zero_node = false) { @@ -165,7 +165,7 @@ void RunSession2(InferenceSession& session_object, VerifyOutputs(fetches, expected_dims, expected_values); } -void RunWithOneSessionSingleThreadInference(std::string model_name, std::string sess_log_id) { +void RunWithOneSessionSingleThreadInference(PathString model_name, std::string sess_log_id) { SessionOptions so; so.session_logid = sess_log_id; RunOptions run_options; @@ -222,7 +222,7 @@ void RunWithOneSessionSingleThreadInference(std::string model_name, std::string ASSERT_TRUE(HasCacheFileWithPrefix(params.trt_ep_context_file_path)); } -void RunWithOneSessionMultiThreadsInference(std::string model_name, std::string sess_log_id, bool has_non_zero_node = false) { +void RunWithOneSessionMultiThreadsInference(PathString model_name, std::string sess_log_id, bool has_non_zero_node = false) { SessionOptions so; so.session_logid = sess_log_id; RunOptions run_options; @@ -289,7 +289,7 @@ void RunWithOneSessionMultiThreadsInference(std::string model_name, std::string TEST(TensorrtExecutionProviderTest, SessionCreationWithMultiThreadsAndInferenceWithMultiThreads) { std::vector threads; - std::string model_name = "trt_execution_provider_multithreading_test.onnx"; + PathString model_name = ORT_TSTR("trt_execution_provider_multithreading_test.onnx"); std::string graph_name = "multithreading_test"; std::string sess_log_id = "TRTEPMultiThreadingTestWithOneSessionSingleThread"; std::vector dims = {1, 3, 2}; @@ -305,7 +305,7 @@ TEST(TensorrtExecutionProviderTest, SessionCreationWithMultiThreadsAndInferenceW } TEST(TensorrtExecutionProviderTest, SessionCreationWithSingleThreadAndInferenceWithMultiThreads) { - std::string model_name = "trt_execution_provider_multithreading_test.onnx"; + PathString model_name = ORT_TSTR("trt_execution_provider_multithreading_test.onnx"); std::string graph_name = "multithreading_test"; std::string sess_log_id = "TRTEPMultiThreadingTestWithOneSessionMultiThreads"; std::vector dims = {1, 3, 2}; @@ -360,7 +360,7 @@ TEST(TensorrtExecutionProviderTest, TRTModelIdGeneratorUsingModelHashing) { } TEST(TensorrtExecutionProviderTest, EPContextNode) { - std::string model_name = "EPContextNode_test.onnx"; + PathString model_name = ORT_TSTR("EPContextNode_test.onnx"); std::string graph_name = "EPContextNode_test"; std::string sess_log_id = "EPContextNode_test"; std::vector dims = {1, 3, 2}; @@ -461,7 +461,7 @@ TEST(TensorrtExecutionProviderTest, EPContextNode) { */ InferenceSession session_object3{so, GetEnvironment()}; OrtTensorRTProviderOptionsV2 params3; - model_name = params.trt_ep_context_file_path; + model_name = ToPathString(params.trt_ep_context_file_path); params3.trt_engine_cache_enable = 1; execution_provider = TensorrtExecutionProviderWithOptions(¶ms3); EXPECT_TRUE(session_object3.RegisterExecutionProvider(std::move(execution_provider)).IsOK()); @@ -490,7 +490,7 @@ TEST(TensorrtExecutionProviderTest, EPContextNode) { */ InferenceSession session_object4{so, GetEnvironment()}; OrtTensorRTProviderOptionsV2 params4; - model_name = "./context_model_folder/EPContextNode_test_ctx.onnx"; + model_name = ORT_TSTR("./context_model_folder/EPContextNode_test_ctx.onnx"); execution_provider = TensorrtExecutionProviderWithOptions(¶ms4); EXPECT_TRUE(session_object4.RegisterExecutionProvider(std::move(execution_provider)).IsOK()); status = session_object4.Load(model_name); @@ -514,7 +514,7 @@ TEST(TensorrtExecutionProviderTest, EPContextNode) { params5.trt_dump_ep_context_model = 1; params5.trt_ep_context_embed_mode = 1; params5.trt_ep_context_file_path = "EP_Context_model_2.onnx"; - model_name = "EPContextNode_test.onnx"; + model_name = ORT_TSTR("EPContextNode_test.onnx"); execution_provider = TensorrtExecutionProviderWithOptions(¶ms5); EXPECT_TRUE(session_object5.RegisterExecutionProvider(std::move(execution_provider)).IsOK()); status = session_object5.Load(model_name); @@ -528,7 +528,7 @@ TEST(TensorrtExecutionProviderTest, EPContextNode) { InferenceSession session_object6{so, GetEnvironment()}; OrtTensorRTProviderOptionsV2 params6; params6.trt_ep_context_embed_mode = 1; - model_name = params5.trt_ep_context_file_path; + model_name = ToPathString(params5.trt_ep_context_file_path); execution_provider = TensorrtExecutionProviderWithOptions(¶ms6); EXPECT_TRUE(session_object6.RegisterExecutionProvider(std::move(execution_provider)).IsOK()); status = session_object6.Load(model_name); @@ -546,7 +546,7 @@ TEST(TensorrtExecutionProviderTest, EPContextNode) { } TEST(TensorrtExecutionProviderTest, TRTPluginsCustomOpTest) { - std::string model_name = "testdata/trt_plugin_custom_op_test.onnx"; + PathString model_name = ORT_TSTR("testdata/trt_plugin_custom_op_test.onnx"); SessionOptions so; so.session_logid = "TensorrtExecutionProviderTRTPluginsTest"; RunOptions run_options; @@ -575,7 +575,6 @@ TEST(TensorrtExecutionProviderTest, TRTPluginsCustomOpTest) { OrtTensorRTProviderOptionsV2 params; std::unique_ptr execution_provider = TensorrtExecutionProviderWithOptions(¶ms); EXPECT_TRUE(session_object.RegisterExecutionProvider(std::move(execution_provider)).IsOK()); - std::cout << model_name << std::endl; auto status = session_object.Load(model_name); ASSERT_TRUE(status.IsOK()); status = session_object.Initialize(); @@ -591,9 +590,12 @@ TEST_P(TensorrtExecutionProviderCacheTest, Run) { size_t pos = param.find("_"); std::string input_type = param.substr(pos + 1); ASSERT_NE(pos, std::string::npos); - std::string cache_type = ToUTF8String(param.substr(0, pos)); - - std::string model_name = "trt_execution_provider_" + cache_type + "caching_test_" + input_type + ".onnx"; + std::string cache_type_mbs = param.substr(0, pos); + PathString cache_type = ToPathString(cache_type_mbs); + std::basic_ostringstream oss; + oss << ORT_TSTR("trt_execution_provider_") << cache_type << ORT_TSTR("_caching_test_") << ToPathString(input_type) + << ORT_TSTR(".onnx"); + PathString model_name = oss.str(); std::vector dims; if (input_type.compare("dynamic") == 0) { dims = {1, -1, -1}; // dynamic shape input @@ -601,10 +603,10 @@ TEST_P(TensorrtExecutionProviderCacheTest, Run) { dims = {1, 3, 2}; } - CreateBaseModel(model_name, cache_type + "cachingtest", dims); + CreateBaseModel(model_name, cache_type_mbs + "cachingtest", dims); SessionOptions so; - so.session_logid = "TensorrtExecutionProvider" + cache_type + "cacheTest"; + so.session_logid = "TensorrtExecutionProvider" + cache_type_mbs + "cacheTest"; RunOptions run_options; run_options.run_tag = so.session_logid; InferenceSession session_object{so, GetEnvironment()}; @@ -633,7 +635,7 @@ TEST_P(TensorrtExecutionProviderCacheTest, Run) { std::vector expected_values_mul_m = {3.0f, 6.0f, 9.0f, 12.0f, 15.0f, 18.0f}; OrtTensorRTProviderOptionsV2 params; - if (cache_type.compare("engine") == 0) { + if (cache_type_mbs.compare("engine") == 0) { /* Following code block tests the functionality of engine and optimization profile of ORT TRT, including: * - engine cache serialization/de-serialization * - profile cache serialization/de-serialization @@ -807,7 +809,7 @@ TEST_P(TensorrtExecutionProviderCacheTest, Run) { } } - } else if (cache_type.compare("timing") == 0) { + } else if (cache_type_mbs.compare("timing") == 0) { /* Following code block tests the functionality of timing cache, including: * - timing cache cache serialization/de-serialization * - TODO: benefir of usign a timing cache no matter if dynamic / static input @@ -917,7 +919,7 @@ TEST(TensorrtExecutionProviderTest, FunctionTest) { auto status = graph.Resolve(); ASSERT_TRUE(status.IsOK()); - std::string model_file_name = "trt_execution_provider_function_test.onnx"; + PathString model_file_name = ORT_TSTR("trt_execution_provider_function_test.onnx"); status = onnxruntime::Model::Save(model, model_file_name); SessionOptions so; @@ -1019,7 +1021,7 @@ TEST(TensorrtExecutionProviderTest, DISABLED_NodeIndexMappingTest) { // [W:onn auto status = graph.Resolve(); ASSERT_TRUE(status.IsOK()); - std::string model_file_name = "trt_execution_provider_nodeindexmapping_test.onnx"; + PathString model_file_name = ORT_TSTR("trt_execution_provider_nodeindexmapping_test.onnx"); status = onnxruntime::Model::Save(model, model_file_name); SessionOptions so; @@ -1131,7 +1133,7 @@ TEST(TensorrtExecutionProviderTest, RemoveCycleTest) { auto status = graph.Resolve(); ASSERT_TRUE(status.IsOK()); - std::string model_file_name = "trt_execution_provider_removecycle_test.onnx"; + PathString model_file_name = ORT_TSTR("trt_execution_provider_removecycle_test.onnx"); status = onnxruntime::Model::Save(model, model_file_name); std::vector dims_mul_x = {1, 3, 2}; diff --git a/onnxruntime/test/util/test_utils.cc b/onnxruntime/test/util/test_utils.cc index 598147b81dd89..6bc0f8d105495 100644 --- a/onnxruntime/test/util/test_utils.cc +++ b/onnxruntime/test/util/test_utils.cc @@ -233,7 +233,7 @@ void CheckShapeEquality(const ONNX_NAMESPACE::TensorShapeProto* shape1, #if !defined(DISABLE_SPARSE_TENSORS) void SparseIndicesChecker(const ONNX_NAMESPACE::TensorProto& indices_proto, gsl::span expected_indicies) { using namespace ONNX_NAMESPACE; - Path model_path; + std::filesystem::path model_path; std::vector unpack_buffer; gsl::span ind_span; std::vector converted_indices; diff --git a/orttraining/orttraining/core/framework/checkpoint_common.cc b/orttraining/orttraining/core/framework/checkpoint_common.cc index 295f17b894095..2c36895de2ac5 100644 --- a/orttraining/orttraining/core/framework/checkpoint_common.cc +++ b/orttraining/orttraining/core/framework/checkpoint_common.cc @@ -16,7 +16,7 @@ namespace onnxruntime { namespace training { /** - * @brief Create OrtValues From TensorProto objects + * @brief Create OrtValues From TensorProto objects. Doesn't support external tensor. * * @param tensor_protos vector of TensorProto * @param name_to_ort_value saved results. @@ -42,7 +42,7 @@ Status CreateOrtValuesFromTensorProtos( tensor_proto.data_type()) ->GetElementType(); auto p_tensor = std::make_unique(tensor_dtype, tensor_shape, cpu_allocator); - ORT_RETURN_IF_ERROR(utils::TensorProtoToTensor(Env::Default(), nullptr, tensor_proto, *p_tensor)); + ORT_RETURN_IF_ERROR(utils::TensorProtoToTensor(Env::Default(), std::filesystem::path(), tensor_proto, *p_tensor)); OrtValue ort_value; ort_value.Init(p_tensor.release(), diff --git a/orttraining/orttraining/core/framework/ortmodule_graph_builder.cc b/orttraining/orttraining/core/framework/ortmodule_graph_builder.cc index c5948e563fcf8..e01456ee3d769 100644 --- a/orttraining/orttraining/core/framework/ortmodule_graph_builder.cc +++ b/orttraining/orttraining/core/framework/ortmodule_graph_builder.cc @@ -183,7 +183,7 @@ Status OrtModuleGraphBuilder::OptimizeForwardGraph(const TrainingGraphTransforme } if (!config.optimized_pre_grad_filepath.empty()) { - ORT_RETURN_IF_ERROR(Model::Save(*forward_model_, config.optimized_pre_grad_filepath)); + ORT_RETURN_IF_ERROR(Model::Save(*forward_model_, ToPathString(config.optimized_pre_grad_filepath))); } return Status::OK(); diff --git a/orttraining/orttraining/core/optimizer/graph_transformer_config.h b/orttraining/orttraining/core/optimizer/graph_transformer_config.h index f72dbfa3fdfc3..c496e36689de1 100644 --- a/orttraining/orttraining/core/optimizer/graph_transformer_config.h +++ b/orttraining/orttraining/core/optimizer/graph_transformer_config.h @@ -28,6 +28,7 @@ struct TrainingGraphTransformerConfiguration : public GraphTransformerConfigurat bool print_input_density{false}; // Path for serialization of the transformed optimized model. If empty, serialization is disabled. + // A UTF-8 string. std::string optimized_pre_grad_filepath; }; diff --git a/orttraining/orttraining/models/runner/training_util.cc b/orttraining/orttraining/models/runner/training_util.cc index 7764508d9a091..6af3bf4410065 100644 --- a/orttraining/orttraining/models/runner/training_util.cc +++ b/orttraining/orttraining/models/runner/training_util.cc @@ -53,7 +53,7 @@ common::Status DataSet::AddData(const vector& featu OrtMemoryInfo info("Cpu", OrtDeviceAllocator, OrtDevice{}, 0, OrtMemTypeDefault); std::unique_ptr buffer = std::make_unique(cpu_tensor_length); ORT_RETURN_IF_ERROR(utils::TensorProtoToOrtValue( - Env::Default(), nullptr, tensor_proto, MemBuffer(buffer.get(), cpu_tensor_length, info), ort_value)); + Env::Default(), std::filesystem::path(), tensor_proto, MemBuffer(buffer.get(), cpu_tensor_length, info), ort_value)); sample->push_back(ort_value); ortvalue_buffers_.emplace_back(std::move(buffer)); diff --git a/orttraining/orttraining/python/orttraining_pybind_state.cc b/orttraining/orttraining/python/orttraining_pybind_state.cc index e2616e1b441f7..a81ea76e807ca 100644 --- a/orttraining/orttraining/python/orttraining_pybind_state.cc +++ b/orttraining/orttraining/python/orttraining_pybind_state.cc @@ -621,7 +621,7 @@ void addObjectMethodsForTraining(py::module& m) { ORT_THROW_IF_ERROR(gradient_graph_builder->builder_->Build()); }) .def("save", [](PyGradientGraphBuilderContext* gradient_graph_builder, const std::string& path) { - ORT_THROW_IF_ERROR(Model::Save(*(gradient_graph_builder->model_), path)); + ORT_THROW_IF_ERROR(Model::Save(*(gradient_graph_builder->model_), ToPathString(path))); }) .def("get_model", [](PyGradientGraphBuilderContext* gradient_graph_builder) { std::string model_str; diff --git a/orttraining/orttraining/test/gradient/allreduce_op_test.cc b/orttraining/orttraining/test/gradient/allreduce_op_test.cc index 82f01a3c43681..1b1bd680a1191 100644 --- a/orttraining/orttraining/test/gradient/allreduce_op_test.cc +++ b/orttraining/orttraining/test/gradient/allreduce_op_test.cc @@ -472,7 +472,7 @@ TEST(AllreduceTest, GPUHierarchicalAdasumAllreduceOptimizerTest) { build_allreduce_graph(graph, adasum_graph_configs, training::AdasumReductionType::GpuHierarchicalReduction, true /*build_optimizer*/, false /*half_precision*/); - std::string model_file_name = "GPUHierarchicalAdasumAllreduceOptimizerTest.onnx"; + PathString model_file_name = ORT_TSTR("GPUHierarchicalAdasumAllreduceOptimizerTest.onnx"); auto status = onnxruntime::Model::Save(model, model_file_name); SessionOptions so; @@ -649,7 +649,7 @@ TEST(AllreduceTest, GPUHierarchicalAdasumAllreduceOptimizerFP16Test) { build_allreduce_graph(graph, adasum_graph_configs, training::AdasumReductionType::GpuHierarchicalReduction, true /*build_optimizer*/, true /*half_precision*/); - std::string model_file_name = "GPUHierarchicalAdasumAllreduceOptimizerFP16Test.onnx"; + PathString model_file_name = ORT_TSTR("GPUHierarchicalAdasumAllreduceOptimizerFP16Test.onnx"); auto status = onnxruntime::Model::Save(model, model_file_name); SessionOptions so; @@ -791,7 +791,7 @@ TEST(AllreduceTest, GPUHierarchicalAdasumAllreduceTest) { adasum_graph_configs.push_back(adasum_graph_config); build_allreduce_graph(graph, adasum_graph_configs, training::AdasumReductionType::GpuHierarchicalReduction); - std::string model_file_name = "GPUHierarchicalAdasumAllreduceTest.onnx"; + PathString model_file_name = ORT_TSTR("GPUHierarchicalAdasumAllreduceTest.onnx"); auto status = onnxruntime::Model::Save(model, model_file_name); SessionOptions so; @@ -896,7 +896,7 @@ TEST(AllreduceTest, GPUHierarchicalAdasumFP16AllreduceTest) { false /*build_optimizer*/, true /*half_precision*/); - std::string model_file_name = "GPUHierarchicalAdasumFP16AllreduceTest.onnx"; + PathString model_file_name = ORT_TSTR("GPUHierarchicalAdasumFP16AllreduceTest.onnx"); auto status = onnxruntime::Model::Save(model, model_file_name); SessionOptions so; @@ -1003,7 +1003,7 @@ TEST(AllreduceTest, GPUAdasumAllreduceTest) { build_allreduce_graph(graph, adasum_graph_configs, training::AdasumReductionType::CpuReduction); - std::string model_file_name = "GPUAdasumAllreduceTest.onnx"; + PathString model_file_name = ORT_TSTR("GPUAdasumAllreduceTest.onnx"); auto status = onnxruntime::Model::Save(model, model_file_name); SessionOptions so; @@ -1110,7 +1110,7 @@ TEST(AllreduceTest, GPUAdasumFP16AllreduceTest) { build_allreduce_graph(graph, adasum_graph_configs, training::AdasumReductionType::CpuReduction, true /*half_precision*/); - std::string model_file_name = "GPUAdasumFP16AllreduceTest.onnx"; + PathString model_file_name = ORT_TSTR("GPUAdasumFP16AllreduceTest.onnx"); auto status = onnxruntime::Model::Save(model, model_file_name); SessionOptions so; diff --git a/orttraining/orttraining/test/optimizer/graph_transform_test.cc b/orttraining/orttraining/test/optimizer/graph_transform_test.cc index 4ab035a171430..b2ab4891f2e1e 100644 --- a/orttraining/orttraining/test/optimizer/graph_transform_test.cc +++ b/orttraining/orttraining/test/optimizer/graph_transform_test.cc @@ -627,7 +627,7 @@ TEST_F(GraphTransformationTests, MegatronMLPPartitionRank0) { TransformerLevel::Level1)); ASSERT_STATUS_OK(graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level1, *logger_)); - auto model_uri2 = "mlp_megatron_basic_test_partition_rank0.onnx"; + PathString model_uri2 = ORT_TSTR("mlp_megatron_basic_test_partition_rank0.onnx"); ASSERT_STATUS_OK(Model::Save(*p_model, model_uri2)); { @@ -705,7 +705,7 @@ TEST_F(GraphTransformationTests, MegatronMLPPartitionRank1) { TransformerLevel::Level1)); ASSERT_STATUS_OK(graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level1, *logger_)); - auto model_uri2 = "mlp_megatron_basic_test_partition_rank1.onnx"; + PathString model_uri2 = ORT_TSTR("mlp_megatron_basic_test_partition_rank1.onnx"); ASSERT_STATUS_OK(Model::Save(*p_model, model_uri2)); { @@ -765,7 +765,7 @@ TEST_F(GraphTransformationTests, MegatronMLPPartitionRank1) { } TEST_F(GraphTransformationTests, MegatronSelfAttentionPartitionRank0) { - auto model_uri = MODEL_FOLDER "model_parallel/self_attention_megatron_basic_test.onnx"; + PathString model_uri = MODEL_FOLDER "model_parallel/self_attention_megatron_basic_test.onnx"; std::shared_ptr p_model; ASSERT_STATUS_OK(Model::Load(model_uri, p_model, nullptr, *logger_)); Graph& graph = p_model->MainGraph(); @@ -781,7 +781,7 @@ TEST_F(GraphTransformationTests, MegatronSelfAttentionPartitionRank0) { TransformerLevel::Level1)); ASSERT_STATUS_OK(graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level1, *logger_)); - auto model_uri2 = "self_attention_megatron_basic_test_partition_rank0.onnx"; + PathString model_uri2 = ORT_TSTR("self_attention_megatron_basic_test_partition_rank0.onnx"); ASSERT_STATUS_OK(Model::Save(*p_model, model_uri2)); { @@ -838,7 +838,7 @@ TEST_F(GraphTransformationTests, MegatronSelfAttentionPartitionRank0) { } TEST_F(GraphTransformationTests, MegatronSelfAttentionPartitionRank1) { - auto model_uri = MODEL_FOLDER "model_parallel/self_attention_megatron_basic_test.onnx"; + PathString model_uri = MODEL_FOLDER "model_parallel/self_attention_megatron_basic_test.onnx"; std::shared_ptr p_model; ASSERT_STATUS_OK(Model::Load(model_uri, p_model, nullptr, *logger_)); Graph& graph = p_model->MainGraph(); @@ -856,7 +856,7 @@ TEST_F(GraphTransformationTests, MegatronSelfAttentionPartitionRank1) { TransformerLevel::Level1)); ASSERT_STATUS_OK(graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level1, *logger_)); - auto model_uri2 = "self_attention_megatron_basic_test_partition_rank1.onnx"; + PathString model_uri2 = ORT_TSTR("self_attention_megatron_basic_test_partition_rank1.onnx"); ASSERT_STATUS_OK(Model::Save(*p_model, model_uri2)); { @@ -913,7 +913,7 @@ TEST_F(GraphTransformationTests, MegatronSelfAttentionPartitionRank1) { } TEST_F(GraphTransformationTests, BiasGeluRecomputeTest) { - auto model_uri = MODEL_FOLDER "fusion/bias_gelu_fusion_recompute.onnx"; + PathString model_uri = MODEL_FOLDER "fusion/bias_gelu_fusion_recompute.onnx"; std::shared_ptr p_model; ASSERT_STATUS_OK(Model::Load(model_uri, p_model, nullptr, *logger_)); Graph& graph = p_model->MainGraph(); @@ -1397,7 +1397,7 @@ static void RunPartitionCorrectnessTest(std::string model_path, TransformerLevel::Level1)); ASSERT_STATUS_OK(graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level1, logger)); graphs.push_back(&graph); - auto model_uri2 = ToPathString(model_path) + ORT_TSTR("_partition_rank_") + ToPathString(std::to_string(i)) + ORT_TSTR(".onnx"); + PathString model_uri2 = ToPathString(model_path) + ORT_TSTR("_partition_rank_") + ToPathString(std::to_string(i)) + ORT_TSTR(".onnx"); ASSERT_STATUS_OK(Model::Save(*p_models[i], model_uri2)); } @@ -1405,7 +1405,7 @@ static void RunPartitionCorrectnessTest(std::string model_path, auto& combine_graph = combine_model.MainGraph(); auto ret = horizontal_parallel_test_utils::MergeGraphsOnAllWorkers(graphs, combine_graph); ORT_ENFORCE(ret.IsOK()); - auto model_uri2 = ToPathString(model_path) + ORT_TSTR("_partition_combine.onnx"); + PathString model_uri2 = ToPathString(model_path) + ORT_TSTR("_partition_combine.onnx"); ASSERT_STATUS_OK(Model::Save(combine_model, model_uri2)); float scale = 1.f; @@ -1790,7 +1790,7 @@ TEST_F(GraphTransformationTests, ScaledSumFusionTwoInputs) { #ifdef ENABLE_TRITON TEST_F(GraphTransformationTests, TritonFusion) { - auto model_uri = MODEL_FOLDER "bert_toy_opset14.onnx"; + PathString model_uri = MODEL_FOLDER "bert_toy_opset14.onnx"; std::shared_ptr model; ASSERT_STATUS_OK(Model::Load(model_uri, model, nullptr, *logger_)); Graph& graph = model->MainGraph(); @@ -1805,7 +1805,7 @@ TEST_F(GraphTransformationTests, TritonFusion) { ASSERT_TRUE(op_to_count["LayerNormalization"] == 4); { - auto model_uri = MODEL_FOLDER "bert_toy_opset14.onnx"; + PathString model_uri = MODEL_FOLDER "bert_toy_opset14.onnx"; std::shared_ptr model; ASSERT_STATUS_OK(Model::Load(model_uri, model, nullptr, *logger_)); Graph& graph = model->MainGraph(); @@ -1845,7 +1845,7 @@ TEST_F(GraphTransformationTests, TritonFusion) { // No Dropout. { - auto model_uri = MODEL_FOLDER "bert_toy_opset14.onnx"; + PathString model_uri = MODEL_FOLDER "bert_toy_opset14.onnx"; std::shared_ptr model; ASSERT_STATUS_OK(Model::Load(model_uri, model, nullptr, *logger_)); Graph& graph = model->MainGraph(); @@ -1884,7 +1884,7 @@ TEST_F(GraphTransformationTests, TritonFusion) { // Ignore min nodes. { - auto model_uri = MODEL_FOLDER "bert_toy_opset14.onnx"; + PathString model_uri = MODEL_FOLDER "bert_toy_opset14.onnx"; std::shared_ptr model; ASSERT_STATUS_OK(Model::Load(model_uri, model, nullptr, *logger_)); Graph& graph = model->MainGraph(); @@ -1924,7 +1924,7 @@ TEST_F(GraphTransformationTests, TritonFusion) { // Exclude Softmax using axis attribute. { - auto model_uri = MODEL_FOLDER "bert_toy_opset14.onnx"; + PathString model_uri = MODEL_FOLDER "bert_toy_opset14.onnx"; std::shared_ptr model; ASSERT_STATUS_OK(Model::Load(model_uri, model, nullptr, *logger_)); Graph& graph = model->MainGraph(); diff --git a/orttraining/orttraining/training_api/checkpoint.cc b/orttraining/orttraining/training_api/checkpoint.cc index cb355ed04e907..56029b34c24d7 100644 --- a/orttraining/orttraining/training_api/checkpoint.cc +++ b/orttraining/orttraining/training_api/checkpoint.cc @@ -330,7 +330,7 @@ Status FromTensorProtos(gsl::span trainable_t for (const auto& tensor_proto : tensor_protos) { flatbuffers::Offset fbs_tensor; ORT_RETURN_IF_ERROR( - fbs::utils::SaveInitializerOrtFormat(builder, tensor_proto, Path(), fbs_tensor, external_data_writer)); + fbs::utils::SaveInitializerOrtFormat(builder, tensor_proto, std::filesystem::path(), fbs_tensor, external_data_writer)); fbs_tensors.push_back(fbs_tensor); } diff --git a/orttraining/orttraining/training_api/module.cc b/orttraining/orttraining/training_api/module.cc index 347673628e106..dc724fbae48eb 100644 --- a/orttraining/orttraining/training_api/module.cc +++ b/orttraining/orttraining/training_api/module.cc @@ -685,7 +685,7 @@ Status Module::ExportModelForInferencing(const std::string& inference_model_path ORT_THROW_IF_ERROR( Model::SaveWithExternalInitializers(*inference_model, inference_model_pathstring, external_data_name, 64)); } else { - ORT_THROW_IF_ERROR(Model::Save(*inference_model, inference_model_path)); + ORT_THROW_IF_ERROR(Model::Save(*inference_model, ToPathString(inference_model_path))); } // Save the model at the desired location. return Status::OK(); From 9007ede102162883c71457cc2872244c374009c8 Mon Sep 17 00:00:00 2001 From: Jian Chen Date: Fri, 28 Jun 2024 21:40:09 -0700 Subject: [PATCH 47/52] Update upstream packaging pipeline name to make it more meaningful. (#21154) ### Description Update upstream packaging pipeline name to make it more meaningful. ### Motivation and Context The upstream pipeline used to only building Nuget packages, but now it also builds Zip and Java. So change the name will make it more meaningful. --- .../github/azure-pipelines/nuget-cuda-publishing-pipeline.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tools/ci_build/github/azure-pipelines/nuget-cuda-publishing-pipeline.yml b/tools/ci_build/github/azure-pipelines/nuget-cuda-publishing-pipeline.yml index b0cc253ae0973..4bfd726f5c58c 100644 --- a/tools/ci_build/github/azure-pipelines/nuget-cuda-publishing-pipeline.yml +++ b/tools/ci_build/github/azure-pipelines/nuget-cuda-publishing-pipeline.yml @@ -1,7 +1,7 @@ resources: pipelines: - pipeline: build - source: 'Nuget-CUDA-Packaging-Pipeline' + source: 'CUDA-Zip-Nuget-Java-Packaging-Pipeline' trigger: branches: include: From 56b36a58baa5738af7e1beb7dd4aa4acf7e54d9e Mon Sep 17 00:00:00 2001 From: Chen Feiyue <69809761+chenfeiyue-cfy@users.noreply.github.com> Date: Sat, 29 Jun 2024 12:48:34 +0800 Subject: [PATCH 48/52] Initial PR for VSINPU execution provider (#20903) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### Description -It is an initial PR for VSINPU execution provider ### Motivation and Context - For support VeriSilicon hardware - TIM-VX(Tensor Interface Module) (https://github.com/VeriSilicon/TIM-VX) is an integrated software solution by Verisilicon for our hardware(A311D/i.MX 8M Plus etc.) design, it is easy to use Verisilicon’s hardware by simply connecting onnxruntime with the TIM-VX API by this VSINPU execution provider. --- cmake/CMakeLists.txt | 6 + cmake/onnxruntime.cmake | 1 + cmake/onnxruntime_providers.cmake | 32 ++ cmake/onnxruntime_unittests.cmake | 5 + include/onnxruntime/core/graph/constants.h | 1 + .../vsinpu/vsinpu_provider_factory.h | 34 ++ onnxruntime/core/framework/node_unit.cc | 9 +- onnxruntime/core/framework/utils.cc | 1 + .../core/providers/get_execution_providers.cc | 8 + .../providers/provider_factory_creators.h | 4 + .../builders/impl/activation_op_builder.h | 130 +++++ .../vsinpu/builders/impl/base_op_builder.cc | 205 +++++++ .../vsinpu/builders/impl/base_op_builder.h | 75 +++ .../vsinpu/builders/impl/cast_op_builder.h | 47 ++ .../vsinpu/builders/impl/clip_op_builder.cc | 115 ++++ .../vsinpu/builders/impl/clip_op_builder.h | 57 ++ .../vsinpu/builders/impl/concat_op_builder.h | 65 +++ .../vsinpu/builders/impl/conv_op_builder.h | 162 ++++++ .../builders/impl/dequantize_op_builder.h | 83 +++ .../builders/impl/elementwise_op_builder.h | 98 ++++ .../vsinpu/builders/impl/flatten_op_builder.h | 65 +++ .../vsinpu/builders/impl/gather_op_builder.h | 86 +++ .../vsinpu/builders/impl/gemm_op_builder.h | 148 ++++++ .../vsinpu/builders/impl/matmul_op_builder.h | 56 ++ .../vsinpu/builders/impl/norm_op_builder.h | 86 +++ .../vsinpu/builders/impl/pool_op_builder.h | 152 ++++++ .../builders/impl/qlinear_binary_op_builder.h | 85 +++ .../builders/impl/qlinearconcat_op_builder.h | 48 ++ .../builders/impl/qlinearconv_op_builder.h | 151 ++++++ .../builders/impl/qlinearmatmul_op_builder.h | 83 +++ .../builders/impl/quantize_op_builder.h | 79 +++ .../vsinpu/builders/impl/reduce_op_builder.h | 82 +++ .../vsinpu/builders/impl/resize_op_builder.h | 153 ++++++ .../vsinpu/builders/impl/softmax_op_builder.h | 101 ++++ .../vsinpu/builders/impl/squeeze_op_builder.h | 88 +++ .../vsinpu/builders/impl/tensor_op_builder.h | 142 +++++ .../vsinpu/builders/impl/tile_op_builder.h | 71 +++ .../builders/impl/unsqueeze_op_builder.h | 89 ++++ .../providers/vsinpu/builders/op_builder.h | 48 ++ .../vsinpu/builders/op_builder_factory.h | 133 +++++ .../vsinpu/patches/AccuracyCorrection.patch | 26 + .../patches/local_testing_record_res.patch | 343 ++++++++++++ .../vsinpu/patches/mlas_crosscompiling.patch | 34 ++ .../test_scripts/compare_cosine_sim.py | 29 + .../patches/test_scripts/compare_topn.py | 34 ++ .../patches/test_scripts/result_compare.sh | 23 + onnxruntime/core/providers/vsinpu/symbols.txt | 1 + .../core/providers/vsinpu/vsinpu_ep_graph.cc | 296 +++++++++++ .../core/providers/vsinpu/vsinpu_ep_graph.h | 116 ++++ .../vsinpu/vsinpu_execution_provider.cc | 277 ++++++++++ .../vsinpu/vsinpu_execution_provider.h | 53 ++ .../vsinpu/vsinpu_provider_factory.cc | 59 ++ .../vsinpu/vsinpu_provider_factory_creator.h | 34 ++ .../core/providers/vsinpu/vsinpu_util.cc | 502 ++++++++++++++++++ .../core/providers/vsinpu/vsinpu_util.h | 131 +++++ onnxruntime/test/onnx/TestCase.cc | 2 +- onnxruntime/test/onnx/main.cc | 13 +- .../test/perftest/command_args_parser.cc | 2 + onnxruntime/test/perftest/ort_test_session.cc | 6 + onnxruntime/test/providers/base_tester.cc | 4 + onnxruntime/test/providers/cpu/model_tests.cc | 15 + onnxruntime/test/util/default_providers.cc | 8 + .../test/util/include/default_providers.h | 2 + onnxruntime/test/util/include/providers.h | 3 + tools/ci_build/build.py | 2 + 65 files changed, 5096 insertions(+), 3 deletions(-) create mode 100644 include/onnxruntime/core/providers/vsinpu/vsinpu_provider_factory.h create mode 100644 onnxruntime/core/providers/vsinpu/builders/impl/activation_op_builder.h create mode 100644 onnxruntime/core/providers/vsinpu/builders/impl/base_op_builder.cc create mode 100644 onnxruntime/core/providers/vsinpu/builders/impl/base_op_builder.h create mode 100644 onnxruntime/core/providers/vsinpu/builders/impl/cast_op_builder.h create mode 100644 onnxruntime/core/providers/vsinpu/builders/impl/clip_op_builder.cc create mode 100644 onnxruntime/core/providers/vsinpu/builders/impl/clip_op_builder.h create mode 100644 onnxruntime/core/providers/vsinpu/builders/impl/concat_op_builder.h create mode 100644 onnxruntime/core/providers/vsinpu/builders/impl/conv_op_builder.h create mode 100644 onnxruntime/core/providers/vsinpu/builders/impl/dequantize_op_builder.h create mode 100644 onnxruntime/core/providers/vsinpu/builders/impl/elementwise_op_builder.h create mode 100644 onnxruntime/core/providers/vsinpu/builders/impl/flatten_op_builder.h create mode 100644 onnxruntime/core/providers/vsinpu/builders/impl/gather_op_builder.h create mode 100644 onnxruntime/core/providers/vsinpu/builders/impl/gemm_op_builder.h create mode 100644 onnxruntime/core/providers/vsinpu/builders/impl/matmul_op_builder.h create mode 100644 onnxruntime/core/providers/vsinpu/builders/impl/norm_op_builder.h create mode 100644 onnxruntime/core/providers/vsinpu/builders/impl/pool_op_builder.h create mode 100644 onnxruntime/core/providers/vsinpu/builders/impl/qlinear_binary_op_builder.h create mode 100644 onnxruntime/core/providers/vsinpu/builders/impl/qlinearconcat_op_builder.h create mode 100644 onnxruntime/core/providers/vsinpu/builders/impl/qlinearconv_op_builder.h create mode 100644 onnxruntime/core/providers/vsinpu/builders/impl/qlinearmatmul_op_builder.h create mode 100644 onnxruntime/core/providers/vsinpu/builders/impl/quantize_op_builder.h create mode 100644 onnxruntime/core/providers/vsinpu/builders/impl/reduce_op_builder.h create mode 100644 onnxruntime/core/providers/vsinpu/builders/impl/resize_op_builder.h create mode 100644 onnxruntime/core/providers/vsinpu/builders/impl/softmax_op_builder.h create mode 100644 onnxruntime/core/providers/vsinpu/builders/impl/squeeze_op_builder.h create mode 100644 onnxruntime/core/providers/vsinpu/builders/impl/tensor_op_builder.h create mode 100644 onnxruntime/core/providers/vsinpu/builders/impl/tile_op_builder.h create mode 100644 onnxruntime/core/providers/vsinpu/builders/impl/unsqueeze_op_builder.h create mode 100644 onnxruntime/core/providers/vsinpu/builders/op_builder.h create mode 100644 onnxruntime/core/providers/vsinpu/builders/op_builder_factory.h create mode 100644 onnxruntime/core/providers/vsinpu/patches/AccuracyCorrection.patch create mode 100644 onnxruntime/core/providers/vsinpu/patches/local_testing_record_res.patch create mode 100644 onnxruntime/core/providers/vsinpu/patches/mlas_crosscompiling.patch create mode 100644 onnxruntime/core/providers/vsinpu/patches/test_scripts/compare_cosine_sim.py create mode 100644 onnxruntime/core/providers/vsinpu/patches/test_scripts/compare_topn.py create mode 100644 onnxruntime/core/providers/vsinpu/patches/test_scripts/result_compare.sh create mode 100644 onnxruntime/core/providers/vsinpu/symbols.txt create mode 100644 onnxruntime/core/providers/vsinpu/vsinpu_ep_graph.cc create mode 100644 onnxruntime/core/providers/vsinpu/vsinpu_ep_graph.h create mode 100644 onnxruntime/core/providers/vsinpu/vsinpu_execution_provider.cc create mode 100644 onnxruntime/core/providers/vsinpu/vsinpu_execution_provider.h create mode 100644 onnxruntime/core/providers/vsinpu/vsinpu_provider_factory.cc create mode 100644 onnxruntime/core/providers/vsinpu/vsinpu_provider_factory_creator.h create mode 100644 onnxruntime/core/providers/vsinpu/vsinpu_util.cc create mode 100644 onnxruntime/core/providers/vsinpu/vsinpu_util.h diff --git a/cmake/CMakeLists.txt b/cmake/CMakeLists.txt index 4483e4d5cb17f..c4412e0934f17 100644 --- a/cmake/CMakeLists.txt +++ b/cmake/CMakeLists.txt @@ -101,6 +101,7 @@ option(onnxruntime_BUILD_OBJC "Build Objective-C library" OFF) option(onnxruntime_USE_PREINSTALLED_EIGEN "Use pre-installed EIGEN. Need to provide eigen_SOURCE_PATH if turn this on." OFF) option(onnxruntime_BUILD_BENCHMARKS "Build ONNXRuntime micro-benchmarks" OFF) option(onnxruntime_USE_LLVM "Build TVM with LLVM" OFF) +option(onnxruntime_USE_VSINPU "Build with VSINPU support" OFF) cmake_dependent_option(onnxruntime_USE_FLASH_ATTENTION "Build flash attention kernel for scaled dot product attention" ON "onnxruntime_USE_CUDA" OFF) option(onnxruntime_USE_MEMORY_EFFICIENT_ATTENTION "Build memory efficient attention kernel for scaled dot product attention" ON) @@ -797,6 +798,11 @@ if (onnxruntime_USE_RKNPU) list(APPEND ORT_PROVIDER_CMAKE_FLAGS -Donnxruntime_USE_RKNPU=1) list(APPEND ONNXRUNTIME_PROVIDER_NAMES rknpu) endif() +if (onnxruntime_USE_VSINPU) + list(APPEND ORT_PROVIDER_FLAGS -DUSE_VSINPU=1) + list(APPEND ORT_PROVIDER_CMAKE_FLAGS -Donnxruntime_USE_VSINPU=1) + list(APPEND ONNXRUNTIME_PROVIDER_NAMES vsinpu) +endif() if (onnxruntime_USE_NNAPI_BUILTIN) list(APPEND ORT_PROVIDER_FLAGS -DUSE_NNAPI=1) list(APPEND ORT_PROVIDER_CMAKE_FLAGS -Donnxruntime_USE_NNAPI_BUILTIN=1) diff --git a/cmake/onnxruntime.cmake b/cmake/onnxruntime.cmake index 977aa44b0e8d7..ec98047750a91 100644 --- a/cmake/onnxruntime.cmake +++ b/cmake/onnxruntime.cmake @@ -189,6 +189,7 @@ set(onnxruntime_INTERNAL_LIBRARIES ${PROVIDERS_SNPE} ${PROVIDERS_TVM} ${PROVIDERS_RKNPU} + ${PROVIDERS_VSINPU} ${PROVIDERS_XNNPACK} ${PROVIDERS_WEBNN} ${PROVIDERS_AZURE} diff --git a/cmake/onnxruntime_providers.cmake b/cmake/onnxruntime_providers.cmake index 7e7819ac31a19..402135adbdd89 100644 --- a/cmake/onnxruntime_providers.cmake +++ b/cmake/onnxruntime_providers.cmake @@ -80,6 +80,9 @@ endif() if(onnxruntime_USE_RKNPU) set(PROVIDERS_RKNPU onnxruntime_providers_rknpu) endif() +if(onnxruntime_USE_VSINPU) + set(PROVIDERS_VSINPU onnxruntime_providers_vsinpu) +endif() if(onnxruntime_USE_DML) set(PROVIDERS_DML onnxruntime_providers_dml) endif() @@ -188,6 +191,35 @@ if (onnxruntime_USE_TVM) include(onnxruntime_providers_tvm.cmake) endif() +if (onnxruntime_USE_VSINPU) + add_definitions(-DUSE_VSINPU=1) + set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-unused-parameter") + file(GLOB_RECURSE onnxruntime_providers_vsinpu_srcs + "${ONNXRUNTIME_ROOT}/core/providers/vsinpu/builders/*.h" + "${ONNXRUNTIME_ROOT}/core/providers/vsinpu/builders/*.cc" + "${ONNXRUNTIME_ROOT}/core/providers/vsinpu/*.h" + "${ONNXRUNTIME_ROOT}/core/providers/vsinpu/*.cc" + "${ONNXRUNTIME_ROOT}/core/providers/shared/utils/utils.h" + "${ONNXRUNTIME_ROOT}/core/providers/shared/utils/utils.cc" + ) + source_group(TREE ${ONNXRUNTIME_ROOT}/core FILES ${onnxruntime_providers_vsinpu_srcs}) + add_library(onnxruntime_providers_vsinpu ${onnxruntime_providers_vsinpu_srcs}) + onnxruntime_add_include_to_target(onnxruntime_providers_vsinpu + onnxruntime_common onnxruntime_framework onnx onnx_proto protobuf::libprotobuf-lite flatbuffers Boost::mp11 + safeint_interface nsync::nsync_cpp) + add_dependencies(onnxruntime_providers_vsinpu ${onnxruntime_EXTERNAL_DEPENDENCIES}) + set_target_properties(onnxruntime_providers_vsinpu PROPERTIES FOLDER "ONNXRuntime" LINKER_LANGUAGE CXX) + target_include_directories(onnxruntime_providers_vsinpu PRIVATE ${ONNXRUNTIME_ROOT} $ENV{TIM_VX_INSTALL}/include) + + find_library(TIMVX_LIBRARY NAMES tim-vx PATHS $ENV{TIM_VX_INSTALL}/lib NO_DEFAULT_PATH) + if(TIMVX_LIBRARY) + target_link_libraries(onnxruntime_providers_vsinpu PRIVATE ${TIMVX_LIBRARY}) + else() + message(FATAL_ERROR "Cannot find TIM-VX library!") + endif() + +endif() + if (onnxruntime_USE_XNNPACK) include(onnxruntime_providers_xnnpack.cmake) endif() diff --git a/cmake/onnxruntime_unittests.cmake b/cmake/onnxruntime_unittests.cmake index ed71e7a57a500..711a9f77f9094 100644 --- a/cmake/onnxruntime_unittests.cmake +++ b/cmake/onnxruntime_unittests.cmake @@ -546,6 +546,10 @@ if(onnxruntime_USE_NNAPI_BUILTIN) list(APPEND onnxruntime_test_providers_dependencies onnxruntime_providers_nnapi) endif() +if(onnxruntime_USE_VSINPU) + list(APPEND onnxruntime_test_providers_dependencies onnxruntime_providers_vsinpu) +endif() + if(onnxruntime_USE_JSEP) list(APPEND onnxruntime_test_providers_dependencies onnxruntime_providers_js) endif() @@ -589,6 +593,7 @@ set(ONNXRUNTIME_TEST_LIBS ${onnxruntime_libs} # CUDA, ROCM, TENSORRT, MIGRAPHX, DNNL, and OpenVINO are dynamically loaded at runtime ${PROVIDERS_NNAPI} + ${PROVIDERS_VSINPU} ${PROVIDERS_JS} ${PROVIDERS_QNN} ${PROVIDERS_SNPE} diff --git a/include/onnxruntime/core/graph/constants.h b/include/onnxruntime/core/graph/constants.h index c4a46cd422219..39acb6b4f2aa4 100644 --- a/include/onnxruntime/core/graph/constants.h +++ b/include/onnxruntime/core/graph/constants.h @@ -52,6 +52,7 @@ constexpr const char* kXnnpackExecutionProvider = "XnnpackExecutionProvider"; constexpr const char* kWebNNExecutionProvider = "WebNNExecutionProvider"; constexpr const char* kCannExecutionProvider = "CANNExecutionProvider"; constexpr const char* kAzureExecutionProvider = "AzureExecutionProvider"; +constexpr const char* kVSINPUExecutionProvider = "VSINPUExecutionProvider"; constexpr const char* kExecutionProviderSharedLibraryPath = "shared_lib_path"; constexpr const char* kExecutionProviderSharedLibraryEntry = "provider_factory_entry_point"; diff --git a/include/onnxruntime/core/providers/vsinpu/vsinpu_provider_factory.h b/include/onnxruntime/core/providers/vsinpu/vsinpu_provider_factory.h new file mode 100644 index 0000000000000..a84067a19aa8a --- /dev/null +++ b/include/onnxruntime/core/providers/vsinpu/vsinpu_provider_factory.h @@ -0,0 +1,34 @@ +/**************************************************************************** + * + * Copyright (c) 2023 Vivante Corporation + * + * Permission is hereby granted, free of charge, to any person obtaining a + * copy of this software and associated documentation files (the "Software"), + * to deal in the Software without restriction, including without limitation + * the rights to use, copy, modify, merge, publish, distribute, sublicense, + * and/or sell copies of the Software, and to permit persons to whom the + * Software is furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING + * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER + * DEALINGS IN THE SOFTWARE. + * + *****************************************************************************/ +#include "onnxruntime_c_api.h" + +#ifdef __cplusplus +extern "C" { +#endif + +ORT_API_STATUS(OrtSessionOptionsAppendExecutionProvider_VSINPU, _In_ OrtSessionOptions* options); + +#ifdef __cplusplus +} +#endif diff --git a/onnxruntime/core/framework/node_unit.cc b/onnxruntime/core/framework/node_unit.cc index ac4301641105a..e2c06fbdfa621 100644 --- a/onnxruntime/core/framework/node_unit.cc +++ b/onnxruntime/core/framework/node_unit.cc @@ -285,7 +285,7 @@ void NodeUnit::InitForSingleNode() { const auto& output_defs = target_node_.OutputDefs(); const auto& node_attrs = target_node_.GetAttributes(); auto qlinear_type = GetQLinearOpType(target_node_); - if (qlinear_type == QLinearOpType::Unknown || IsVariadicQLinearOp(qlinear_type)) { // TODO, add variadic support + if (qlinear_type == QLinearOpType::Unknown) { // Not a Qlinear op, add all inputs / outputs auto add_all_io = [](std::vector& defs, const ConstPointerContainer>& node_defs) { @@ -351,6 +351,13 @@ void NodeUnit::InitForSingleNode() { NodeUnitIODef::QuantParam{*input_defs[1], input_defs.size() == 3 ? input_defs[2] : nullptr, axis}}); + } else if (IsVariadicQLinearOp(qlinear_type)) { + size_t input_num = (input_defs.size() - 2) / 3; + for (size_t i = 0; i < input_num; i++) { + inputs_.push_back(NodeUnitIODef{*input_defs[3 * i + 2], NodeUnitIODef::QuantParam{*input_defs[3 * i + 3], + input_defs[3 * i + 4]}}); + } + outputs_.push_back(NodeUnitIODef{*output_defs[0], NodeUnitIODef::QuantParam{*input_defs[0], input_defs[1]}}); } else { ORT_THROW("The QLinear op [", static_cast(qlinear_type), "] is not supported"); } diff --git a/onnxruntime/core/framework/utils.cc b/onnxruntime/core/framework/utils.cc index 9c282210d2169..9eed0249711f9 100644 --- a/onnxruntime/core/framework/utils.cc +++ b/onnxruntime/core/framework/utils.cc @@ -61,6 +61,7 @@ bool ProviderIsCpuBased(const std::string& provider_type) { provider_type == onnxruntime::kVitisAIExecutionProvider || provider_type == onnxruntime::kOpenVINOExecutionProvider || provider_type == onnxruntime::kNnapiExecutionProvider || + provider_type == onnxruntime::kVSINPUExecutionProvider || provider_type == onnxruntime::kAclExecutionProvider || provider_type == onnxruntime::kArmNNExecutionProvider || provider_type == onnxruntime::kRknpuExecutionProvider || diff --git a/onnxruntime/core/providers/get_execution_providers.cc b/onnxruntime/core/providers/get_execution_providers.cc index b0f510f054a03..61c035bc29ed5 100644 --- a/onnxruntime/core/providers/get_execution_providers.cc +++ b/onnxruntime/core/providers/get_execution_providers.cc @@ -98,6 +98,14 @@ constexpr ProviderInfo kProvidersInPriorityOrder[] = true, #else false, +#endif + }, + { + kVSINPUExecutionProvider, +#ifdef USE_VSINPU + true, +#else + false, #endif }, { diff --git a/onnxruntime/core/providers/provider_factory_creators.h b/onnxruntime/core/providers/provider_factory_creators.h index 42a58097e1635..47d3f2f793d7c 100644 --- a/onnxruntime/core/providers/provider_factory_creators.h +++ b/onnxruntime/core/providers/provider_factory_creators.h @@ -46,6 +46,10 @@ #include "core/providers/nnapi/nnapi_provider_factory_creator.h" #endif +#if defined(USE_VSINPU) +#include "core/providers/vsinpu/vsinpu_provider_factory_creator.h" +#endif + #if defined(USE_JSEP) #include "core/providers/js/js_provider_factory_creator.h" #endif diff --git a/onnxruntime/core/providers/vsinpu/builders/impl/activation_op_builder.h b/onnxruntime/core/providers/vsinpu/builders/impl/activation_op_builder.h new file mode 100644 index 0000000000000..9a59d90365f64 --- /dev/null +++ b/onnxruntime/core/providers/vsinpu/builders/impl/activation_op_builder.h @@ -0,0 +1,130 @@ +/**************************************************************************** + * + * Copyright (c) 2023 Vivante Corporation + * + * Permission is hereby granted, free of charge, to any person obtaining a + * copy of this software and associated documentation files (the "Software"), + * to deal in the Software without restriction, including without limitation + * the rights to use, copy, modify, merge, publish, distribute, sublicense, + * and/or sell copies of the Software, and to permit persons to whom the + * Software is furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING + * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER + * DEALINGS IN THE SOFTWARE. + * + *****************************************************************************/ +#include +#include +#include +#include "core/providers/shared/utils/utils.h" +#include "core/providers/vsinpu/builders/impl/base_op_builder.h" + +namespace onnxruntime { +namespace vsi { +namespace npu { +class ReluOpBuilder : public BaseOpBuilder { + public: + bool HandleBuildOp(vsi::npu::GraphEP* graph_ep, + std::vector>& inputs, + std::vector>& outputs, + const NodeUnit& node_unit) override { + LOGS_DEFAULT(VERBOSE) << "Creating Relu Activation."; + auto op = graph_ep->GetGraph()->CreateOperation(); + (*op).BindInputs(inputs).BindOutputs(outputs); + graph_ep->GetOps().push_back(std::move(op)); + return true; + } +}; +class SigmoidOpBuilder : public BaseOpBuilder { + public: + bool HandleBuildOp(vsi::npu::GraphEP* graph_ep, + std::vector>& inputs, + std::vector>& outputs, + const NodeUnit& node_unit) override { + LOGS_DEFAULT(VERBOSE) << "Creating Sigmoid Activation."; + auto op = graph_ep->GetGraph()->CreateOperation(); + (*op).BindInputs(inputs).BindOutputs(outputs); + graph_ep->GetOps().push_back(std::move(op)); + return true; + } +}; +class TanhOpBuilder : public BaseOpBuilder { + public: + bool HandleBuildOp(vsi::npu::GraphEP* graph_ep, + std::vector>& inputs, + std::vector>& outputs, + const NodeUnit& node_unit) override { + LOGS_DEFAULT(VERBOSE) << "Creating Tanh activation."; + auto op = graph_ep->GetGraph()->CreateOperation(); + (*op).BindInputs(inputs).BindOutputs(outputs); + graph_ep->GetOps().push_back(std::move(op)); + return true; + } +}; + +class LeakyReluOpBuilder : public BaseOpBuilder { + public: + bool HandleBuildOp(vsi::npu::GraphEP* graph_ep, + std::vector>& inputs, + std::vector>& outputs, + const NodeUnit& node_unit) override { + LOGS_DEFAULT(VERBOSE) << "Creating LeakyRelu activation."; + const auto& node = node_unit.GetNode(); + NodeAttrHelper helper(node); + auto alpha = helper.Get("alpha", 1.0f); + auto op = + graph_ep->GetGraph()->CreateOperation(alpha); + (*op).BindInputs(inputs).BindOutputs(outputs); + graph_ep->GetOps().push_back(std::move(op)); + return true; + } +}; + +class EluOpBuilder : public BaseOpBuilder { + public: + bool HandleBuildOp(vsi::npu::GraphEP* graph_ep, + std::vector>& inputs, + std::vector>& outputs, + const NodeUnit& node_unit) override { + LOGS_DEFAULT(VERBOSE) << "Creating Elu activation."; + const auto& node = node_unit.GetNode(); + NodeAttrHelper helper(node); + auto alpha = helper.Get("alpha", 1.0f); + auto op = + graph_ep->GetGraph()->CreateOperation(alpha); + (*op).BindInputs(inputs).BindOutputs(outputs); + graph_ep->GetOps().push_back(std::move(op)); + return true; + } +}; + +class HardSigmoidOpBuilder : public BaseOpBuilder { + public: + bool HandleBuildOp(vsi::npu::GraphEP* graph_ep, + std::vector>& inputs, + std::vector>& outputs, + const NodeUnit& node_unit) override { + LOGS_DEFAULT(VERBOSE) << "Creating HardSigmoid activation."; + const auto& node = node_unit.GetNode(); + NodeAttrHelper helper(node); + auto alpha = helper.Get("alpha", 1.0f); + auto beta = helper.Get("beta", 1.0f); + auto op = graph_ep->GetGraph()->CreateOperation( + alpha, beta); + (*op).BindInputs(inputs).BindOutputs(outputs); + graph_ep->GetOps().push_back(std::move(op)); + return true; + } +}; +} // namespace npu + +} // namespace vsi +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/vsinpu/builders/impl/base_op_builder.cc b/onnxruntime/core/providers/vsinpu/builders/impl/base_op_builder.cc new file mode 100644 index 0000000000000..894bf8e4444f8 --- /dev/null +++ b/onnxruntime/core/providers/vsinpu/builders/impl/base_op_builder.cc @@ -0,0 +1,205 @@ +/**************************************************************************** + * + * Copyright (c) 2023 Vivante Corporation + * + * Permission is hereby granted, free of charge, to any person obtaining a + * copy of this software and associated documentation files (the "Software"), + * to deal in the Software without restriction, including without limitation + * the rights to use, copy, modify, merge, publish, distribute, sublicense, + * and/or sell copies of the Software, and to permit persons to whom the + * Software is furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING + * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER + * DEALINGS IN THE SOFTWARE. + * + *****************************************************************************/ +#include +#include "core/providers/vsinpu/builders/impl/base_op_builder.h" + +namespace onnxruntime { +namespace vsi { +namespace npu { +bool BaseOpBuilder::IsSupported(const onnxruntime::GraphViewer& graph_viewer, + const NodeUnit& node_unit) const { + auto initializers = graph_viewer.GetAllInitializedTensors(); + if (!HasSupportedOpSet(node_unit)) { + return false; + } + if (!HasSupportedInputOutputs(initializers, node_unit)) { + return false; + } + return IsOpSupported(graph_viewer, &node_unit.GetNode()); +} + +bool BaseOpBuilder::HasSupportedInputOutputs(const InitializedTensorSet& initializers, + const NodeUnit& node_unit) const { + // We do not support unknown(null) input shape + auto has_supported_shape = [](const NodeArg& node_arg, const std::string& name, const std::string& op_type) { + const auto* shape_proto = node_arg.Shape(); + if (!shape_proto) { + LOGS_DEFAULT(WARNING) << "Node [" << name << "] type [" << op_type + << "] Input [" << node_arg.Name() << "] has no shape"; + return false; + } + + // We do not support dynamic shape input yet, but resize op's second input can be empty + for (const auto& dim : shape_proto->dim()) { + if (!dim.has_dim_value()) { + LOGS_DEFAULT(WARNING) << "Dynamic shape is not supported for now, for input:" << node_arg.Name(); + return false; + } + if (dim.dim_value() == 0 && op_type != "Resize") { + LOGS_DEFAULT(WARNING) << "Zero in shape is not supported for now, for input:" << node_arg.Name(); + return false; + } + } + return true; + }; + + auto has_initialized_quant_param = [](const NodeArg& arg, const InitializedTensorSet& initializers) { + auto it = initializers.find(arg.Name()); + if (it == initializers.end()) { + LOGS_DEFAULT(WARNING) << "The quantization param must be an initializer tensor"; + return false; + } + return true; + }; + + for (const auto& input : node_unit.Inputs()) { + if (!input.node_arg.Exists()) { + continue; + } + if (!has_supported_shape(input.node_arg, node_unit.Name(), node_unit.OpType())) + return false; + + if (input.quant_param.has_value()) { + if (!has_supported_shape(input.quant_param->scale, node_unit.Name(), node_unit.OpType())) + return false; + + if (!has_initialized_quant_param(input.quant_param->scale, initializers)) + return false; + // zero point is optional + if (input.quant_param->zero_point) { + if (!has_supported_shape(*input.quant_param->zero_point, node_unit.Name(), node_unit.OpType())) + return false; + if (!has_initialized_quant_param(*input.quant_param->zero_point, initializers)) + return false; + if (input.quant_param->zero_point->Type() != input.node_arg.Type()) { + LOGS_DEFAULT(ERROR) << "Invalid input type because the data type mismatch with its' quant param type."; + return false; + } + } + } + } + for (const auto& output : node_unit.Outputs()) { + if (output.quant_param.has_value()) { + if (!has_supported_shape(output.quant_param->scale, node_unit.Name(), node_unit.OpType())) + return false; + + if (!has_initialized_quant_param(output.quant_param->scale, initializers)) + return false; + // zero point is optional + if (output.quant_param->zero_point) { + if (!has_supported_shape(*output.quant_param->zero_point, node_unit.Name(), node_unit.OpType())) + return false; + if (!has_initialized_quant_param(*output.quant_param->zero_point, initializers)) + return false; + } + } + } + return HasSupportedInputOutputsImpl(initializers, node_unit); +} + +bool BaseOpBuilder::HasSupportedInputOutputsImpl( + const InitializedTensorSet& /* initializers */, const NodeUnit& node_unit) const { + // Check input/output data type, int64 is generally unsupported + // specific op builder can override this if the int64 input corresponds to VSINPU param + for (const auto& input : node_unit.Inputs()) { + auto input_type = input.node_arg.Type(); + if (*input_type == "tensor(int64)" || !util::IsTypeSupported(&input.node_arg)) { + LOGS_DEFAULT(WARNING) << node_unit.OpType() << " has unsupported input type : " + << *input_type; + return false; + } + } + for (const auto& output : node_unit.Outputs()) { + auto output_type = output.node_arg.Type(); + if (*output_type == "tensor(int64)" || !util::IsTypeSupported(&output.node_arg)) { + LOGS_DEFAULT(WARNING) << node_unit.OpType() << " has unsupported output type : " + << *output_type; + return false; + } + } + return true; +} + +bool BaseOpBuilder::HasSupportedOpSet(const NodeUnit& node_unit) const { + auto since_version = node_unit.SinceVersion(); + if (since_version < GetMinSupportedOpSet(node_unit) || since_version > GetMaxSupportedOpSet(node_unit)) { + LOGS_DEFAULT(VERBOSE) << node_unit.OpType() << " opset [" << since_version + << "] is only supported for opset [" + << GetMinSupportedOpSet(node_unit) << ", " + << GetMaxSupportedOpSet(node_unit) << "]"; + return false; + } + + return true; +} + +bool BaseOpBuilder::BuildOp(vsi::npu::GraphEP* graph_ep, + const onnxruntime::GraphViewer& graph_viewer, + const NodeUnit& node_unit) { + std::vector> inputs; + std::vector input_defs = node_unit.Inputs(); + std::vector output_defs = node_unit.Outputs(); + + for (const auto input_def : input_defs) { + auto it = std::find_if( + graph_ep->GetGraphInputs().begin(), graph_ep->GetGraphInputs().end(), + [input_def](const std::shared_ptr& info) { + return info->name == input_def.node_arg.Name(); + }); + tim::vx::TensorAttribute attr; + if (graph_viewer.IsConstantInitializer(input_def.node_arg.Name(), true)) { + attr = tim::vx::TensorAttribute::CONSTANT; + } else if (it == graph_ep->GetGraphInputs().end()) { + attr = tim::vx::TensorAttribute::TRANSIENT; + } else { + attr = tim::vx::TensorAttribute::INPUT; + } + + auto tensor = graph_ep->MapTIMVXTensor(graph_ep->GetGraph(), input_def, node_unit, + &graph_viewer, attr); + inputs.push_back(tensor); + } + + std::vector> outputs; + + for (auto output_def : output_defs) { + auto it = std::find_if( + graph_ep->GetGraphOutputs().begin(), graph_ep->GetGraphOutputs().end(), + [output_def](const std::shared_ptr& info) { + return info->name == output_def.node_arg.Name(); + }); + tim::vx::TensorAttribute attribute = + it == graph_ep->GetGraphOutputs().end() + ? tim::vx::TensorAttribute::TRANSIENT + : tim::vx::TensorAttribute::OUTPUT; + auto tensor = graph_ep->MapTIMVXTensor(graph_ep->GetGraph(), output_def, node_unit, + &graph_viewer, attribute); + outputs.push_back(tensor); + } + return HandleBuildOp(graph_ep, inputs, outputs, node_unit); +} +} // namespace npu + +} // namespace vsi +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/vsinpu/builders/impl/base_op_builder.h b/onnxruntime/core/providers/vsinpu/builders/impl/base_op_builder.h new file mode 100644 index 0000000000000..c0cf3365f46e3 --- /dev/null +++ b/onnxruntime/core/providers/vsinpu/builders/impl/base_op_builder.h @@ -0,0 +1,75 @@ +/**************************************************************************** + * + * Copyright (c) 2023 Vivante Corporation + * + * Permission is hereby granted, free of charge, to any person obtaining a + * copy of this software and associated documentation files (the "Software"), + * to deal in the Software without restriction, including without limitation + * the rights to use, copy, modify, merge, publish, distribute, sublicense, + * and/or sell copies of the Software, and to permit persons to whom the + * Software is furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING + * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER + * DEALINGS IN THE SOFTWARE. + * + *****************************************************************************/ +#pragma once +#include +#include +#include "core/providers/vsinpu/builders/op_builder.h" +#include "core/providers/vsinpu/vsinpu_ep_graph.h" +#include "core/providers/vsinpu/vsinpu_util.h" +#include "tim/vx/operation.h" +#include "tim/vx/ops.h" + +namespace onnxruntime { +namespace vsi { +namespace npu { +class BaseOpBuilder : public IOpBuilder { + public: + virtual ~BaseOpBuilder() = default; + + bool IsSupported(const onnxruntime::GraphViewer& graph_viewer, + const NodeUnit& node_unit) const override; + bool BuildOp(vsi::npu::GraphEP* graph_ep, + const onnxruntime::GraphViewer& graph_viewer, const NodeUnit& node_unit); + virtual bool IsOpSupported(const onnxruntime::GraphViewer& graph_viewer, + const Node* node) const { + return true; + } + + virtual bool IsQuantizedOp(const NodeUnit& /* node_unit */) const { return false; } + + virtual int GetMinSupportedOpSet(const NodeUnit& /* node_unit */) const { return 1; } + virtual int GetMaxSupportedOpSet(const NodeUnit& /* node_unit */) const { return 22; } + + virtual bool HasSupportedInputOutputsImpl( + const InitializedTensorSet& initializers, const NodeUnit& node_unit) const; + + // TODO(cfy): Check if this node_unit's type is supported + virtual bool IsNodeUnitTypeSupported(const NodeUnit& node_unit) const { return true; } + + virtual bool HandleBuildOp( + vsi::npu::GraphEP* graph_ep, + std::vector>& inputs, + std::vector>& outputs, + const NodeUnit& node_unit) { + return true; + } + + private: + bool HasSupportedOpSet(const NodeUnit& node_unit) const; + bool HasSupportedInputOutputs(const InitializedTensorSet& initializers, + const NodeUnit& node_unit) const; +}; +} // namespace npu +} // namespace vsi +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/vsinpu/builders/impl/cast_op_builder.h b/onnxruntime/core/providers/vsinpu/builders/impl/cast_op_builder.h new file mode 100644 index 0000000000000..6579f0ca9045f --- /dev/null +++ b/onnxruntime/core/providers/vsinpu/builders/impl/cast_op_builder.h @@ -0,0 +1,47 @@ +/**************************************************************************** + * + * Copyright (c) 2024 Vivante Corporation + * + * Permission is hereby granted, free of charge, to any person obtaining a + * copy of this software and associated documentation files (the "Software"), + * to deal in the Software without restriction, including without limitation + * the rights to use, copy, modify, merge, publish, distribute, sublicense, + * and/or sell copies of the Software, and to permit persons to whom the + * Software is furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING + * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER + * DEALINGS IN THE SOFTWARE. + * + *****************************************************************************/ +#include +#include +#include +#include "core/providers/shared/utils/utils.h" +#include "core/providers/vsinpu/builders/impl/base_op_builder.h" +namespace onnxruntime { +namespace vsi { +namespace npu { +class CastOpBuilder : public BaseOpBuilder { + protected: + bool HandleBuildOp(vsi::npu::GraphEP* graph_ep, std::vector>& inputs, + std::vector>& outputs, const NodeUnit& node_unit) override { + LOGS_DEFAULT(VERBOSE) << "Creating Cast Op."; + NodeAttrHelper helper(node_unit.GetNode()); + auto op = graph_ep->GetGraph()->CreateOperation(); + (*op).BindInput(inputs[0]).BindOutputs(outputs); + graph_ep->GetOps().push_back(std::move(op)); + return true; + } +}; + +} // namespace npu +} // namespace vsi +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/vsinpu/builders/impl/clip_op_builder.cc b/onnxruntime/core/providers/vsinpu/builders/impl/clip_op_builder.cc new file mode 100644 index 0000000000000..85096d0e262d7 --- /dev/null +++ b/onnxruntime/core/providers/vsinpu/builders/impl/clip_op_builder.cc @@ -0,0 +1,115 @@ +/**************************************************************************** + * + * Copyright (c) 2024 Vivante Corporation + * + * Permission is hereby granted, free of charge, to any person obtaining a + * copy of this software and associated documentation files (the "Software"), + * to deal in the Software without restriction, including without limitation + * the rights to use, copy, modify, merge, publish, distribute, sublicense, + * and/or sell copies of the Software, and to permit persons to whom the + * Software is furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING + * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER + * DEALINGS IN THE SOFTWARE. + * + *****************************************************************************/ +#include +#include +#include "core/providers/vsinpu/builders/impl/clip_op_builder.h" + +namespace onnxruntime { +namespace vsi { +namespace npu { + +namespace clip_internal { +template +struct LowMax { + constexpr static T low() { + return std::numeric_limits::lowest(); + } + constexpr static T max() { + return std::numeric_limits::max(); + } +}; +} // namespace clip_internal + +template +struct ClipOpBuilder::ClipImpl { + ClipImpl(vsi::npu::GraphEP* graph_ep, std::vector>& inputs, + std::vector>& outputs) { + T min_default = clip_internal::LowMax::low(); + T max_default = clip_internal::LowMax::max(); + + T* min_data = &min_default; + T* max_data = &max_default; + std::shared_ptr min_tensor = nullptr; + std::shared_ptr max_tensor = nullptr; + if (inputs.size() > 1) { + min_tensor = inputs[1]; + if (inputs.size() > 2) { + max_tensor = inputs[2]; + } + } + if (min_tensor) { + min_tensor->CopyDataFromTensor(min_data); + } + if (max_tensor) { + max_tensor->CopyDataFromTensor(max_data); + } + auto op = graph_ep->GetGraph()->CreateOperation( + static_cast(*min_data), static_cast(*max_data)); + (*op).BindInputs(inputs).BindOutputs(outputs); + graph_ep->GetOps().push_back(std::move(op)); + } +}; + +bool ClipOpBuilder::HandleBuildOp(vsi::npu::GraphEP* graph_ep, + std::vector>& inputs, + std::vector>& outputs, + const NodeUnit& node_unit) { + LOGS_DEFAULT(INFO) << "Creating Clip Op."; + if (node_unit.SinceVersion() <= 6) { + NodeAttrHelper helper(node_unit.GetNode()); + auto min = helper.Get("min", -3.402e+38f); + auto max = helper.Get("max", 3.402e+38f); + auto op = graph_ep->GetGraph()->CreateOperation(min, max); + (*op).BindInputs(inputs).BindOutputs(outputs); + graph_ep->GetOps().push_back(std::move(op)); + } else { + switch (inputs[0]->GetDataType()) { + case tim::vx::DataType::INT8: + ClipImpl(graph_ep, inputs, outputs); + break; + case tim::vx::DataType::UINT8: + ClipImpl(graph_ep, inputs, outputs); + break; + case tim::vx::DataType::INT16: + ClipImpl(graph_ep, inputs, outputs); + break; + case tim::vx::DataType::INT32: + ClipImpl(graph_ep, inputs, outputs); + break; + case tim::vx::DataType::FLOAT16: + ClipImpl(graph_ep, inputs, outputs); + break; + case tim::vx::DataType::FLOAT32: + default: + ClipImpl(graph_ep, inputs, outputs); + break; + } + } + return true; +} + +} // namespace npu + +} // namespace vsi +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/vsinpu/builders/impl/clip_op_builder.h b/onnxruntime/core/providers/vsinpu/builders/impl/clip_op_builder.h new file mode 100644 index 0000000000000..368cb092657c8 --- /dev/null +++ b/onnxruntime/core/providers/vsinpu/builders/impl/clip_op_builder.h @@ -0,0 +1,57 @@ +/**************************************************************************** + * + * Copyright (c) 2024 Vivante Corporation + * + * Permission is hereby granted, free of charge, to any person obtaining a + * copy of this software and associated documentation files (the "Software"), + * to deal in the Software without restriction, including without limitation + * the rights to use, copy, modify, merge, publish, distribute, sublicense, + * and/or sell copies of the Software, and to permit persons to whom the + * Software is furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING + * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER + * DEALINGS IN THE SOFTWARE. + * + *****************************************************************************/ +#include +#include +#include "core/providers/vsinpu/builders/impl/base_op_builder.h" +#include "core/providers/shared/utils/utils.h" + +namespace onnxruntime { +namespace vsi { +namespace npu { +class ClipOpBuilder final : public BaseOpBuilder { + bool IsOpSupported(const onnxruntime::GraphViewer& graph_viewer, + const Node* node) const override { + if (node->SinceVersion() > 6) { + if (node->InputDefs().size() > 1 && + !Contains(graph_viewer.GetAllInitializedTensors(), node->InputDefs()[1]->Name())) { + LOGS_DEFAULT(WARNING) << "Min/Max value must be const input or attribute."; + return false; + } + } + return true; + } + + bool HandleBuildOp(vsi::npu::GraphEP* graph_ep, + std::vector>& inputs, + std::vector>& outputs, + const NodeUnit& node_unit) override; + + private: + template + struct ClipImpl; +}; +} // namespace npu + +} // namespace vsi +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/vsinpu/builders/impl/concat_op_builder.h b/onnxruntime/core/providers/vsinpu/builders/impl/concat_op_builder.h new file mode 100644 index 0000000000000..4d3fc658b7bef --- /dev/null +++ b/onnxruntime/core/providers/vsinpu/builders/impl/concat_op_builder.h @@ -0,0 +1,65 @@ +/**************************************************************************** + * + * Copyright (c) 2024 Vivante Corporation + * + * Permission is hereby granted, free of charge, to any person obtaining a + * copy of this software and associated documentation files (the "Software"), + * to deal in the Software without restriction, including without limitation + * the rights to use, copy, modify, merge, publish, distribute, sublicense, + * and/or sell copies of the Software, and to permit persons to whom the + * Software is furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING + * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER + * DEALINGS IN THE SOFTWARE. + * + *****************************************************************************/ +#include +#include +#include +#include "core/providers/vsinpu/builders/impl/base_op_builder.h" +#include "core/providers/common.h" +#include "core/providers/shared/utils/utils.h" + +namespace onnxruntime { +namespace vsi { +namespace npu { +class ConcatOpBuilder : public BaseOpBuilder { + bool IsOpSupported(const onnxruntime::GraphViewer& graph_viewer, + const Node* node) const override { + NodeAttrHelper helper(*node); + auto axis = helper.Get("axis", 0); + auto input_defs = node->InputDefs(); + auto input_shape = vsi::npu::util::GetTensorShape(*input_defs[0]); + int32_t rank = input_shape.NumDimensions(); + if (axis >= rank || axis < -rank) { + LOGS_DEFAULT(ERROR) << "Axis is invalid in Concat."; + return false; + } + return true; + } + bool HandleBuildOp(vsi::npu::GraphEP* graph_ep, + std::vector>& inputs, + std::vector>& outputs, + const NodeUnit& node_unit) override { + LOGS_DEFAULT(VERBOSE) << "Creating Concat Op."; + NodeAttrHelper helper(node_unit.GetNode()); + auto axis = helper.Get("axis", 0); + axis = util::ReverseAxis(axis, inputs[0]->GetShape().size()); + auto op = graph_ep->GetGraph()->CreateOperation(static_cast(axis), inputs.size()); + (*op).BindInputs(inputs).BindOutputs(outputs); + graph_ep->GetOps().push_back(std::move(op)); + return true; + } +}; +} // namespace npu + +} // namespace vsi +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/vsinpu/builders/impl/conv_op_builder.h b/onnxruntime/core/providers/vsinpu/builders/impl/conv_op_builder.h new file mode 100644 index 0000000000000..d44e1ce1799c1 --- /dev/null +++ b/onnxruntime/core/providers/vsinpu/builders/impl/conv_op_builder.h @@ -0,0 +1,162 @@ +/**************************************************************************** + * + * Copyright (c) 2023 Vivante Corporation + * + * Permission is hereby granted, free of charge, to any person obtaining a + * copy of this software and associated documentation files (the "Software"), + * to deal in the Software without restriction, including without limitation + * the rights to use, copy, modify, merge, publish, distribute, sublicense, + * and/or sell copies of the Software, and to permit persons to whom the + * Software is furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING + * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER + * DEALINGS IN THE SOFTWARE. + * + *****************************************************************************/ +#include +#include +#include +#include +#include "core/providers/shared/utils/utils.h" +#include "core/providers/vsinpu/builders/impl/base_op_builder.h" +namespace onnxruntime { +namespace vsi { +namespace npu { +class ConvOpBuilder : public BaseOpBuilder { + bool IsOpSupported(const onnxruntime::GraphViewer& graph_viewer, + const Node* node) const override { + auto input_defs = node->InputDefs(); + auto shape = vsi::npu::util::GetTensorShape(*input_defs[0]); + if (shape.NumDimensions() == 5) { + LOGS_DEFAULT(WARNING) << "Not support conv3d yet."; + return false; + } + return true; + } + + bool HandleBuildOp(vsi::npu::GraphEP* graph_ep, + std::vector>& inputs, + std::vector>& outputs, + const NodeUnit& node_unit) override { + auto input_tensor = inputs[0]; + auto weight_tensor = inputs[1]; + auto OChannel_idx = weight_tensor->GetShape().size() - 1; + const bool is_1d_conv = + weight_tensor->GetShape().size() == 3 ? true : false; + NodeAttrHelper helper(node_unit.GetNode()); + auto padtype = helper.Get("auto_pad", std::string("")); + auto group = helper.Get("group", static_cast(1)); + + std::string op_type = (group != 1 && group == weight_tensor->GetShape()[OChannel_idx]) + ? "DepthwiseConv" + : (group != 1) ? "GroupConv" + : "Conv"; + op_type += is_1d_conv ? "1D" : "2D"; + std::string op_name = std::string("Creating ") + op_type + " Op"; + LOGS_DEFAULT(INFO) << op_name; + + uint32_t default_uint = 1; + std::vector default_vec = {1, 1}; + + auto stride = + helper.Get("strides", is_1d_conv ? std::vector{default_uint} + : default_vec); + auto dilation = + helper.Get("dilations", is_1d_conv ? std::vector{default_uint} + : default_vec); + + std::shared_ptr op; + if (padtype != "NOTSET") { // array "pads" is not set + if (group != 1 && group != weight_tensor->GetShape()[OChannel_idx]) { + if (is_1d_conv) { + op = graph_ep->GetGraph() + ->CreateOperation( + vsi::npu::util::GetPadType(padtype), stride[0], + dilation[0], group, tim::vx::DataLayout::WCN, + tim::vx::DataLayout::WIcOc); + } else { + op = graph_ep->GetGraph() + ->CreateOperation( + vsi::npu::util::GetPadType(padtype), + /* W_stride, H_stride*/ + std::array{stride[1], stride[0]}, + /* W_dilation, H_dilation*/ + std::array{dilation[1], dilation[0]}, group, + tim::vx::DataLayout::WHCN, tim::vx::DataLayout::WHIcOc); + } + } else { + int32_t multiplier = group == 1 + ? 0 + : weight_tensor->GetShape()[OChannel_idx] / input_tensor->GetShape()[OChannel_idx - 1]; + if (is_1d_conv) { + op = graph_ep->GetGraph()->CreateOperation( + vsi::npu::util::GetPadType(padtype), stride[0], dilation[0], multiplier, + tim::vx::DataLayout::WCN, tim::vx::DataLayout::WIcOc); + } else { + op = graph_ep->GetGraph()->CreateOperation( + vsi::npu::util::GetPadType(padtype), + /* W_stride, H_stride*/ + std::array{stride[1], stride[0]}, + /* W_dilation, H_dilation*/ + std::array{dilation[1], dilation[0]}, multiplier, + tim::vx::DataLayout::WHCN, tim::vx::DataLayout::WHIcOc); + } + } + } else { + auto pads = helper.Get("pads", std::vector{0U, 0U}); + if (group != 1 && group != weight_tensor->GetShape()[OChannel_idx]) { + if (is_1d_conv) { + op = graph_ep->GetGraph() + ->CreateOperation( + vsi::npu::util::GetPadType(padtype), + std::array{pads[0], pads[1]}, stride[0], + dilation[0], group, tim::vx::DataLayout::WCN, + tim::vx::DataLayout::WIcOc); + } else { + op = graph_ep->GetGraph() + ->CreateOperation( + /* W_begin,W_end, H_begin,H_end*/ std::array< + uint32_t, 4>{pads[1], pads[3], pads[0], pads[2]}, + /* W_stride, H_stide*/ + std::array{stride[1], stride[0]}, + /* W_dilation, H_dilation*/ + std::array{dilation[1], dilation[0]}, group, + tim::vx::DataLayout::WHCN, tim::vx::DataLayout::WHIcOc); + } + } else { + int32_t multiplier = group == 1 + ? 0 + : weight_tensor->GetShape()[OChannel_idx] / input_tensor->GetShape()[OChannel_idx - 1]; + if (is_1d_conv) { + op = graph_ep->GetGraph()->CreateOperation( + std::array{pads[0], pads[1]}, stride[0], dilation[0], + multiplier, tim::vx::DataLayout::WCN, tim::vx::DataLayout::WIcOc); + } else { + op = graph_ep->GetGraph()->CreateOperation( + /* W_begin,W_end, H_begin,H_end*/ std::array{pads[1], pads[3], + pads[0], pads[2]}, + /* W_stride, H_stride*/ + std::array{stride[1], stride[0]}, + /* W_dilation, H_dilation*/ + std::array{dilation[1], dilation[0]}, multiplier, + tim::vx::DataLayout::WHCN, tim::vx::DataLayout::WHIcOc); + } + } + } + (*op).BindInputs(inputs).BindOutputs(outputs); + graph_ep->GetOps().push_back(std::move(op)); + return true; + } +}; +} // namespace npu + +} // namespace vsi +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/vsinpu/builders/impl/dequantize_op_builder.h b/onnxruntime/core/providers/vsinpu/builders/impl/dequantize_op_builder.h new file mode 100644 index 0000000000000..50b295f2fb539 --- /dev/null +++ b/onnxruntime/core/providers/vsinpu/builders/impl/dequantize_op_builder.h @@ -0,0 +1,83 @@ +/**************************************************************************** + * + * Copyright (c) 2024 Vivante Corporation + * + * Permission is hereby granted, free of charge, to any person obtaining a + * copy of this software and associated documentation files (the "Software"), + * to deal in the Software without restriction, including without limitation + * the rights to use, copy, modify, merge, publish, distribute, sublicense, + * and/or sell copies of the Software, and to permit persons to whom the + * Software is furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING + * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER + * DEALINGS IN THE SOFTWARE. + * + *****************************************************************************/ +#include +#include +#include +#include "core/providers/vsinpu/builders/impl/base_op_builder.h" +#include "core/providers/common.h" +#include "core/providers/shared/utils/utils.h" + +namespace onnxruntime { +namespace vsi { +namespace npu { +class DequantizeLinearOpBuilder : public BaseOpBuilder { + enum DequantizeINPUTS { + input_tensor = 0, + scale_tensor = 1, + zero_point_tensor = 2 + }; + bool HasSupportedInputOutputsImpl(const InitializedTensorSet& initializers, + const NodeUnit& node_unit) const override { + auto input_type = node_unit.Inputs()[0].node_arg.Type(); + if (*input_type == "tensor(int64)" || !util::IsTypeSupported(&node_unit.Inputs()[0].node_arg)) { + LOGS_DEFAULT(WARNING) << node_unit.OpType() << " has unsupported input type : " + << *input_type; + return false; + } + if (!node_unit.Inputs()[0].quant_param.has_value()) { + LOGS_DEFAULT(WARNING) << "The quantization params must be known."; + return false; + } + if (node_unit.Inputs()[0].quant_param->scale.Shape()->dim_size() != 0 && + node_unit.Inputs()[0].quant_param->scale.Shape()->dim(0).dim_value() != 1) { + LOGS_DEFAULT(WARNING) << "Per channel quantized input is not support in DequantizeLinear op."; + return false; + } + return true; + } + bool IsOpSupported(const onnxruntime::GraphViewer& graph_viewer, + const Node* node) const override { + NodeAttrHelper helper(*node); + if (helper.HasAttr("block_size") && helper.Get("block_size", 0) != 0) { + LOGS_DEFAULT(WARNING) << "Not support block quantization yet."; + return false; + } + return true; + } + + bool HandleBuildOp(vsi::npu::GraphEP* graph_ep, + std::vector>& inputs, + std::vector>& outputs, + const NodeUnit& node_unit) override { + LOGS_DEFAULT(INFO) << "Creating Dequantize Op."; + auto op = graph_ep->GetGraph()->CreateOperation(); + (*op).BindInputs(inputs).BindOutputs(outputs); + graph_ep->GetOps().push_back(std::move(op)); + return true; + } +}; +} // namespace npu + +} // namespace vsi +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/vsinpu/builders/impl/elementwise_op_builder.h b/onnxruntime/core/providers/vsinpu/builders/impl/elementwise_op_builder.h new file mode 100644 index 0000000000000..89809a4513340 --- /dev/null +++ b/onnxruntime/core/providers/vsinpu/builders/impl/elementwise_op_builder.h @@ -0,0 +1,98 @@ + +/**************************************************************************** + * + * Copyright (c) 2023 Vivante Corporation + * + * Permission is hereby granted, free of charge, to any person obtaining a + * copy of this software and associated documentation files (the "Software"), + * to deal in the Software without restriction, including without limitation + * the rights to use, copy, modify, merge, publish, distribute, sublicense, + * and/or sell copies of the Software, and to permit persons to whom the + * Software is furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING + * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER + * DEALINGS IN THE SOFTWARE. + * + *****************************************************************************/ +#include +#include +#include +#include "core/providers/vsinpu/builders/impl/base_op_builder.h" +namespace onnxruntime { +namespace vsi { +namespace npu { +#define ELEMENTWISE_OP_BUILDER(onnx_op_type, vsinpu_op_kind) \ + class onnx_op_type##OpBuilder : public BaseOpBuilder { \ + bool IsOpSupported(const onnxruntime::GraphViewer& graph_viewer, \ + const Node* node) const override { \ + for (auto input : node->InputDefs()) { \ + if (*input->Type() == "tensor(int64)") { \ + LOGS_DEFAULT(WARNING) << "Int64 type is not suppoted as elementwise operation input."; \ + return false; \ + } \ + } \ + return true; \ + } \ + bool HandleBuildOp(vsi::npu::GraphEP* graph_ep, \ + std::vector>& inputs, \ + std::vector>& outputs, \ + const NodeUnit& node_unit) override { \ + LOGS_DEFAULT(INFO) << "Creating " << #onnx_op_type << " Op"; \ + auto op = graph_ep->GetGraph() -> CreateOperation(); \ + (*op).BindInputs(inputs).BindOutputs(outputs); \ + return true; \ + ; \ + } \ + }; + +ELEMENTWISE_OP_BUILDER(Add, Add); +ELEMENTWISE_OP_BUILDER(Sub, Sub); +ELEMENTWISE_OP_BUILDER(Mul, Multiply); +ELEMENTWISE_OP_BUILDER(Div, Div); // not consider zero +ELEMENTWISE_OP_BUILDER(Abs, Abs); +ELEMENTWISE_OP_BUILDER(Sqrt, Sqrt); +ELEMENTWISE_OP_BUILDER(Exp, Exp); +ELEMENTWISE_OP_BUILDER(Floor, Floor); +ELEMENTWISE_OP_BUILDER(Log, Log); +ELEMENTWISE_OP_BUILDER(Sin, Sin); +ELEMENTWISE_OP_BUILDER(HardSwish, HardSwish); + +class PowOpBuilder : public BaseOpBuilder { + bool IsOpSupported(const onnxruntime::GraphViewer& graph_viewer, + const Node* node) const override { + auto input0_type = *node->InputDefs()[0]->Type(); + auto input1_type = *node->InputDefs()[1]->Type(); + if (input0_type != input1_type) { + if ((input0_type == "tensor(float)" && input1_type == "tensor(int32)") || + (input0_type == "tensor(int32)" && input1_type == "tensor(float)")) { + LOGS_DEFAULT(WARNING) << "Pow op does not support one of input is float32 while the other one is int32 type."; + return false; + } + } + return true; + } + bool HandleBuildOp(vsi::npu::GraphEP* graph_ep, + std::vector>& inputs, + std::vector>& outputs, + const NodeUnit& node_unit) override { + LOGS_DEFAULT(INFO) << "Creating Pow Op"; + auto op = graph_ep->GetGraph() + ->CreateOperation(); + (*op).BindInputs(inputs).BindOutputs(outputs); + graph_ep->GetOps().push_back(std::move(op)); + return true; + } +}; + +} // namespace npu + +} // namespace vsi +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/vsinpu/builders/impl/flatten_op_builder.h b/onnxruntime/core/providers/vsinpu/builders/impl/flatten_op_builder.h new file mode 100644 index 0000000000000..dfb0bb9c1b99f --- /dev/null +++ b/onnxruntime/core/providers/vsinpu/builders/impl/flatten_op_builder.h @@ -0,0 +1,65 @@ +/**************************************************************************** + * + * Copyright (c) 2024 Vivante Corporation + * + * Permission is hereby granted, free of charge, to any person obtaining a + * copy of this software and associated documentation files (the "Software"), + * to deal in the Software without restriction, including without limitation + * the rights to use, copy, modify, merge, publish, distribute, sublicense, + * and/or sell copies of the Software, and to permit persons to whom the + * Software is furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING + * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER + * DEALINGS IN THE SOFTWARE. + * + *****************************************************************************/ +#include +#include +#include +#include "core/providers/vsinpu/builders/impl/base_op_builder.h" +#include "core/providers/common.h" +#include "core/providers/shared/utils/utils.h" + +namespace onnxruntime { +namespace vsi { +namespace npu { +class FlattenOpBuilder : public BaseOpBuilder { + bool HandleBuildOp(vsi::npu::GraphEP* graph_ep, + std::vector>& inputs, + std::vector>& outputs, + const NodeUnit& node_unit) override { + LOGS_DEFAULT(VERBOSE) << "Creating Flatten Op."; + std::vector reshape_param; + if (outputs[0]->GetShape().size() == 2) { + reshape_param = outputs[0]->GetShape(); + } else { + auto input_shape = inputs[0]->GetShape(); + NodeAttrHelper helper(node_unit.GetNode()); + int64_t axis = helper.Get("axis", 1); + axis = util::ReverseAxis(static_cast(axis), input_shape.size()); + uint32_t first_dim = 1; + for (int64_t i = 0; i < axis; i++) { + first_dim *= inputs[0]->GetShape()[i]; + } + uint32_t second_dim = inputs[0]->GetSpec().GetElementNum() / first_dim; + reshape_param.push_back(first_dim); + reshape_param.push_back(second_dim); + } + auto op = graph_ep->GetGraph()->CreateOperation(reshape_param); + (*op).BindInputs(inputs).BindOutputs(outputs); + graph_ep->GetOps().push_back(std::move(op)); + return true; + } +}; +} // namespace npu + +} // namespace vsi +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/vsinpu/builders/impl/gather_op_builder.h b/onnxruntime/core/providers/vsinpu/builders/impl/gather_op_builder.h new file mode 100644 index 0000000000000..0325b68ae0ad7 --- /dev/null +++ b/onnxruntime/core/providers/vsinpu/builders/impl/gather_op_builder.h @@ -0,0 +1,86 @@ +/**************************************************************************** + * + * Copyright (c) 2024 Vivante Corporation + * + * Permission is hereby granted, free of charge, to any person obtaining a + * copy of this software and associated documentation files (the "Software"), + * to deal in the Software without restriction, including without limitation + * the rights to use, copy, modify, merge, publish, distribute, sublicense, + * and/or sell copies of the Software, and to permit persons to whom the + * Software is furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING + * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER + * DEALINGS IN THE SOFTWARE. + * + *****************************************************************************/ +#include +#include +#include +#include "core/providers/vsinpu/builders/impl/base_op_builder.h" +#include "core/providers/shared/utils/utils.h" + +namespace onnxruntime { +namespace vsi { +namespace npu { +class GatherOpBuilder : public BaseOpBuilder { + bool HasSupportedInputOutputsImpl(const InitializedTensorSet& initializers, + const NodeUnit& node_unit) const override { + auto input = node_unit.Inputs()[0]; + auto indices = node_unit.Inputs()[1]; + if (util::IsTypeSupported(&input.node_arg) && util::IsTypeSupported(&indices.node_arg)) { + if (*input.node_arg.Type() == "tensor(int64)") { + LOGS_DEFAULT(WARNING) << "Only support indices tensor to be int64 type in gather op."; + return false; + } + if (*indices.node_arg.Type() != "tensor(int64)" && *indices.node_arg.Type() != "tensor(int32)") { + LOGS_DEFAULT(WARNING) << "Unsupported indices tensor type in gather op."; + return false; + } + if (*indices.node_arg.Type() == "tensor(int64)" && !Contains(initializers, indices.node_arg.Name())) { + LOGS_DEFAULT(WARNING) << "Only support const attribute if indice tensor is in int64 type."; + return false; + } + return true; + } + return false; + } + + bool HandleBuildOp(vsi::npu::GraphEP* graph_ep, + std::vector>& inputs, + std::vector>& outputs, + const NodeUnit& node_unit) override { + LOGS_DEFAULT(VERBOSE) << "Creating Gather Op."; + NodeAttrHelper helper(node_unit.GetNode()); + auto axis = helper.Get("axis", 0); + axis = util::ReverseAxis(axis, inputs[0]->GetShape().size()); + auto op = graph_ep->GetGraph()->CreateOperation(axis, 0); + + bool is_i64_indices = inputs[1]->GetDataType() == tim::vx::DataType::INT64; + if (!is_i64_indices) { + (*op).BindInputs(inputs).BindOutputs(outputs); + } else { + std::vector origin_data(inputs[1]->GetSpec().GetElementNum()); + inputs[1]->CopyDataFromTensor(origin_data.data()); + std::vector transformed_data(origin_data.begin(), origin_data.end()); + tim::vx::TensorSpec ts = inputs[1]->GetSpec().SetAttribute(tim::vx::TensorAttribute::INPUT); + ts.SetDataType(tim::vx::DataType::INT32); + auto transformed_indices = graph_ep->GetGraph()->CreateTensor(ts, transformed_data.data()); + (*op).BindInput(inputs[0]).BindInput(transformed_indices).BindOutput(outputs[0]); + } + graph_ep->GetOps().push_back(std::move(op)); + return true; + } +}; + +} // namespace npu + +} // namespace vsi +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/vsinpu/builders/impl/gemm_op_builder.h b/onnxruntime/core/providers/vsinpu/builders/impl/gemm_op_builder.h new file mode 100644 index 0000000000000..6f2c590b862b6 --- /dev/null +++ b/onnxruntime/core/providers/vsinpu/builders/impl/gemm_op_builder.h @@ -0,0 +1,148 @@ +/**************************************************************************** + * + * Copyright (c) 2023 Vivante Corporation + * + * Permission is hereby granted, free of charge, to any person obtaining a + * copy of this software and associated documentation files (the "Software"), + * to deal in the Software without restriction, including without limitation + * the rights to use, copy, modify, merge, publish, distribute, sublicense, + * and/or sell copies of the Software, and to permit persons to whom the + * Software is furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING + * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER + * DEALINGS IN THE SOFTWARE. + * + *****************************************************************************/ +#include +#include +#include +#include "core/providers/shared/utils/utils.h" +#include "core/providers/vsinpu/builders/impl/base_op_builder.h" +namespace onnxruntime { +namespace vsi { +namespace npu { +class GemmOpBuilder : public BaseOpBuilder { + bool IsOpSupported(const onnxruntime::GraphViewer& graph_viewer, + const Node* node) const override { + auto input_defs = node->InputDefs(); + NodeAttrHelper helper(*node); + auto weight_units = helper.Get("transB", 0) == 1 + ? vsi::npu::util::GetTensorShape(*input_defs[1]).GetDims()[0] + : vsi::npu::util::GetTensorShape(*input_defs[1]).GetDims()[1]; + if (input_defs.size() > 2) { + auto bias_shape = vsi::npu::util::GetTensorShape(*input_defs[2]); + if (bias_shape.NumDimensions() == 1 && bias_shape.GetDims()[0] != weight_units) { + LOGS_DEFAULT(WARNING) << "Not support to broadcast bias shape."; + return false; + } else if (bias_shape.NumDimensions() == 2 && + (bias_shape.Size() != weight_units || + (bias_shape.GetDims()[0] != 1 && bias_shape.GetDims()[1] != 1))) { + LOGS_DEFAULT(WARNING) << "Not support 2-dims bias shape."; + return false; + } + + if (*input_defs[2]->Type() == "tensor(float16)" && + !graph_viewer.IsConstantInitializer(input_defs[2]->Name(), true)) { + LOGS_DEFAULT(WARNING) << "Not support f16 bias with input attr."; + return false; + } + } + return true; + } + bool HandleBuildOp(vsi::npu::GraphEP* graph_ep, + std::vector>& inputs, + std::vector>& outputs, + const NodeUnit& node_unit) override { + LOGS_DEFAULT(VERBOSE) << "Creating Gemm Op."; + auto input_A = inputs[0]; + auto input_B = inputs[1]; + NodeAttrHelper helper(node_unit.GetNode()); + + auto trans_A = helper.Get("transA", 0); + auto trans_B = helper.Get("transB", 0); + const bool has_alpha = (helper.Get("alpha", 1.0f) != 1.0); + const bool has_beta = (helper.Get("beta", 1.0f) != 1.0); + const bool has_C = (inputs.size() == 3); + auto weight_units = helper.Get("transB", 0) == 1 ? inputs[1]->GetShape()[1] : inputs[1]->GetShape()[0]; + + tim::vx::TensorSpec coef_spec(tim::vx::DataType::FLOAT32, {1}, + tim::vx::TensorAttribute::CONSTANT); + + auto multiply_impl = [&](std::shared_ptr input, + std::shared_ptr coef, + std::shared_ptr output) { + auto multiply_op = graph_ep->GetGraph()->CreateOperation(); + (*multiply_op).BindInput(input).BindInput(coef).BindOutput(output); + graph_ep->GetOps().push_back(multiply_op); + }; + + auto transpose_impl = [&](std::shared_ptr input, + std::shared_ptr output) { + std::vector perm = {1U, 0U}; + auto transpose_op = graph_ep->GetGraph()->CreateOperation(perm); + (*transpose_op).BindInput(input).BindOutput(output); + graph_ep->GetOps().push_back(std::move(transpose_op)); + }; + + auto fc_impl = [&](std::vector> inputs, + std::shared_ptr output) { + auto fc_op = graph_ep->GetGraph()->CreateOperation(0, weight_units); + (*fc_op).BindInputs(inputs).BindOutput(output); + graph_ep->GetOps().push_back(std::move(fc_op)); + }; + + auto alpha_A = input_A; + std::shared_ptr beta_C; + auto final_A = input_A; + auto final_B = input_B; + + if (has_alpha) { + auto alpha_tensor = graph_ep->GetGraph()->CreateTensor(coef_spec); + auto alpha = helper.Get("alpha", 1.0f); + alpha_tensor->CopyDataToTensor(&alpha); + alpha_A = graph_ep->GetGraph()->CreateTensor( + input_A->GetSpec().AsTransientSpec()); + multiply_impl(input_A, alpha_tensor, alpha_A); + final_A = alpha_A; + } + if (has_beta) { + auto beta_tensor = graph_ep->GetGraph()->CreateTensor(coef_spec); + auto beta = helper.Get("beta", 1.0f); + beta_tensor->CopyDataToTensor(&beta); + beta_C = graph_ep->GetGraph()->CreateTensor( + inputs[2]->GetSpec().AsTransientSpec()); + multiply_impl(inputs[2], beta_tensor, beta_C); + } else if (has_C) { + beta_C = inputs[2]; + } + + if (trans_A) { + final_A = graph_ep->GetGraph()->CreateTensor( + input_A->GetSpec().AsTransientSpec()); + transpose_impl(alpha_A, final_A); + } + if (!trans_B) { + final_B = graph_ep->GetGraph()->CreateTensor( + input_B->GetSpec().AsTransientSpec()); + transpose_impl(input_B, final_B); + } + std::vector> fc_inputs = {final_A, final_B}; + + if (has_C) fc_inputs.push_back(beta_C); + fc_impl(fc_inputs, outputs[0]); + + return true; + } +}; +} // namespace npu + +} // namespace vsi +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/vsinpu/builders/impl/matmul_op_builder.h b/onnxruntime/core/providers/vsinpu/builders/impl/matmul_op_builder.h new file mode 100644 index 0000000000000..8cdf72906b644 --- /dev/null +++ b/onnxruntime/core/providers/vsinpu/builders/impl/matmul_op_builder.h @@ -0,0 +1,56 @@ +/**************************************************************************** + * + * Copyright (c) 2024 Vivante Corporation + * + * Permission is hereby granted, free of charge, to any person obtaining a + * copy of this software and associated documentation files (the "Software"), + * to deal in the Software without restriction, including without limitation + * the rights to use, copy, modify, merge, publish, distribute, sublicense, + * and/or sell copies of the Software, and to permit persons to whom the + * Software is furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING + * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER + * DEALINGS IN THE SOFTWARE. + * + *****************************************************************************/ +#include +#include +#include +#include "core/providers/vsinpu/builders/impl/base_op_builder.h" + +namespace onnxruntime { +namespace vsi { +namespace npu { +class MatMulOpBuilder : public BaseOpBuilder { + bool IsOpSupported(const onnxruntime::GraphViewer& graph_viewer, + const Node* node) const override { + auto output_defs = node->OutputDefs(); + if (output_defs[0]->Shape()->dim_size() == 0) { + LOGS_DEFAULT(WARNING) << "Inner product of 1-D tensor is not supported in MatMul op."; + return false; + } + return true; + } + bool HandleBuildOp(vsi::npu::GraphEP* graph_ep, + std::vector>& inputs, + std::vector>& outputs, + const NodeUnit& node_unit) override { + LOGS_DEFAULT(VERBOSE) << "Creating Matmul Op."; + auto op = graph_ep->GetGraph()->CreateOperation(); + (*op).BindInputs(inputs).BindOutputs(outputs); + graph_ep->GetOps().push_back(std::move(op)); + return true; + } +}; +} // namespace npu + +} // namespace vsi +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/vsinpu/builders/impl/norm_op_builder.h b/onnxruntime/core/providers/vsinpu/builders/impl/norm_op_builder.h new file mode 100644 index 0000000000000..997163c6b1a6d --- /dev/null +++ b/onnxruntime/core/providers/vsinpu/builders/impl/norm_op_builder.h @@ -0,0 +1,86 @@ +/**************************************************************************** + * + * Copyright (c) 2024 Vivante Corporation + * + * Permission is hereby granted, free of charge, to any person obtaining a + * copy of this software and associated documentation files (the "Software"), + * to deal in the Software without restriction, including without limitation + * the rights to use, copy, modify, merge, publish, distribute, sublicense, + * and/or sell copies of the Software, and to permit persons to whom the + * Software is furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING + * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER + * DEALINGS IN THE SOFTWARE. + * + *****************************************************************************/ +#include +#include +#include +#include "core/providers/vsinpu/builders/impl/base_op_builder.h" +#include "core/providers/shared/utils/utils.h" + +namespace onnxruntime { +namespace vsi { +namespace npu { + +class BatchNormOpBuilder : public BaseOpBuilder { + enum NormINPUTS { + input_tensor = 0, + scale_tensor = 1, + Bias_tensor = 2, + mean_tensor = 3, + var_tensor = 4 + }; + int GetMinSupportedOpSet(const NodeUnit& /* node_unit */) const override { return 9; } + + bool IsOpSupported(const onnxruntime::GraphViewer& graph_viewer, + const Node* node) const override { + auto input_defs = node->InputDefs(); + NodeAttrHelper helper(*node); + auto training_mode = helper.Get("training_mode", 0); + if (training_mode) { + LOGS_DEFAULT(WARNING) << "Training is not supported in batch_norm op."; + return false; + } + if (helper.HasAttr("spatial")) { + LOGS_DEFAULT(WARNING) << "VSINPU does not support 'spatial' parameter."; + return false; + } + if (!graph_viewer.IsConstantInitializer(input_defs[NormINPUTS::scale_tensor]->Name(), true)) { + LOGS_DEFAULT(WARNING) << "Not support mean/var/gamma/beta set as dynamic input yet."; + return false; + } + + return true; + } + bool HandleBuildOp(vsi::npu::GraphEP* graph_ep, + std::vector>& inputs, + std::vector>& outputs, + const NodeUnit& node_unit) override { + LOGS_DEFAULT(INFO) << "Creating BatchNorm Op."; + NodeAttrHelper helper(node_unit.GetNode()); + auto epsilon = helper.Get("epsilon", 1e-5f); + auto op = graph_ep->GetGraph()->CreateOperation(epsilon); + std::vector> reordered_inputs; + int indices[] = {NormINPUTS::input_tensor, NormINPUTS::mean_tensor, NormINPUTS::var_tensor, + NormINPUTS::scale_tensor, NormINPUTS::Bias_tensor}; + for (int i : indices) { + reordered_inputs.push_back(inputs[i]); + } + (*op).BindInputs(reordered_inputs).BindOutputs(outputs); + graph_ep->GetOps().push_back(std::move(op)); + return true; + } +}; +} // namespace npu + +} // namespace vsi +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/vsinpu/builders/impl/pool_op_builder.h b/onnxruntime/core/providers/vsinpu/builders/impl/pool_op_builder.h new file mode 100644 index 0000000000000..7cfa9faf68480 --- /dev/null +++ b/onnxruntime/core/providers/vsinpu/builders/impl/pool_op_builder.h @@ -0,0 +1,152 @@ +/**************************************************************************** + * + * Copyright (c) 2023 Vivante Corporation + * + * Permission is hereby granted, free of charge, to any person obtaining a + * copy of this software and associated documentation files (the "Software"), + * to deal in the Software without restriction, including without limitation + * the rights to use, copy, modify, merge, publish, distribute, sublicense, + * and/or sell copies of the Software, and to permit persons to whom the + * Software is furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING + * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER + * DEALINGS IN THE SOFTWARE. + * + *****************************************************************************/ +#include +#include +#include +#include "core/providers/vsinpu/builders/impl/base_op_builder.h" +#include "core/providers/shared/utils/utils.h" + +namespace onnxruntime { +namespace vsi { +namespace npu { +class BasePoolOpBuilder : public BaseOpBuilder { + public: + explicit BasePoolOpBuilder(tim::vx::PoolType pool_type) : pool_type_(pool_type) {} + + protected: + bool IsOpSupported(const onnxruntime::GraphViewer& graph_viewer, const Node* node) const override { + auto shape = vsi::npu::util::GetTensorShape(*node->InputDefs()[0]); + if (shape.NumDimensions() == 5) { + LOGS_DEFAULT(WARNING) << "3DPool is not supported yet."; + return false; + } + + NodeAttrHelper helper(*node); + if (helper.HasAttr("dilations")) { + LOGS_DEFAULT(WARNING) << "NonMaxPool with Dilation parameter is not supported."; + return false; + } + return true; + } + bool CreatePoolingOp(vsi::npu::GraphEP* graph_ep, + std::vector>& inputs, + std::vector>& outputs, + const std::array& kernel_size, + const std::array& strides, + const std::array& pads, + bool is_global, + const tim::vx::RoundType ceil_mode) { + const bool is_1d_pool = inputs[0]->GetShape().size() == 3; + std::shared_ptr op; + + // Create the appropriate pooling operation + if (is_global) { + if (is_1d_pool) { + op = graph_ep->GetGraph()->CreateOperation(pool_type_, inputs[0]->GetShape()[0], + ceil_mode); + } else { + std::array input_size = {inputs[0]->GetShape()[0], inputs[0]->GetShape()[1]}; + op = graph_ep->GetGraph()->CreateOperation(pool_type_, input_size, ceil_mode); + } + + } else { + if (is_1d_pool) { + std::array arr = {pads[2], pads[0]}; + op = graph_ep->GetGraph()->CreateOperation(pool_type_, arr, + kernel_size[1], strides[1], ceil_mode); + } else { + op = graph_ep->GetGraph()->CreateOperation(pool_type_, pads, kernel_size, + strides, ceil_mode); + } + } + + (*op).BindInputs(inputs).BindOutputs(outputs); + graph_ep->GetOps().push_back(std::move(op)); + return true; + } + tim::vx::PoolType pool_type_; +}; + +class TraditionalPoolOpBuilder : public BasePoolOpBuilder { + public: + TraditionalPoolOpBuilder() : BasePoolOpBuilder(tim::vx::PoolType::MAX) {} + + protected: + bool HandleBuildOp(vsi::npu::GraphEP* graph_ep, + std::vector>& inputs, + std::vector>& outputs, + const NodeUnit& node_unit) override { + NodeAttrHelper helper(node_unit.GetNode()); + auto ksize = helper.Get("kernel_shape", std::vector{1U, 1U}); + auto strides = helper.Get("strides", std::vector{1U, 1U}); + auto pads = helper.Get("pads", std::vector{0U, 0U, 0U, 0U}); + tim::vx::RoundType ceil_mode = helper.Get("ceil_mode", 0U) == 0 + ? tim::vx::RoundType::FLOOR + : tim::vx::RoundType::CEILING; + return CreatePoolingOp(graph_ep, inputs, outputs, {ksize[1], ksize[0]}, {strides[1], strides[0]}, + {pads[1], pads[3], pads[0], pads[2]}, false, ceil_mode); + } +}; + +class GlobalPoolOpBuilder : public BasePoolOpBuilder { + public: + GlobalPoolOpBuilder() : BasePoolOpBuilder(tim::vx::PoolType::MAX) {} + + protected: + bool HandleBuildOp(vsi::npu::GraphEP* graph_ep, + std::vector>& inputs, + std::vector>& outputs, + const NodeUnit& node_unit) override { + NodeAttrHelper helper(node_unit.GetNode()); + tim::vx::RoundType ceil_mode = helper.Get("ceil_mode", 0U) == 0 + ? tim::vx::RoundType::FLOOR + : tim::vx::RoundType::CEILING; + return CreatePoolingOp(graph_ep, inputs, outputs, {}, {}, {}, true, ceil_mode); + } +}; + +class GlobalAveragePoolOpBuilder : public GlobalPoolOpBuilder { + public: + GlobalAveragePoolOpBuilder() { pool_type_ = tim::vx::PoolType::AVG; } +}; + +class GlobalMaxPoolOpBuilder : public GlobalPoolOpBuilder { + public: + GlobalMaxPoolOpBuilder() { pool_type_ = tim::vx::PoolType::MAX; } +}; + +class AveragePoolOpBuilder : public TraditionalPoolOpBuilder { + public: + AveragePoolOpBuilder() { pool_type_ = tim::vx::PoolType::AVG; } +}; + +class MaxPoolOpBuilder : public TraditionalPoolOpBuilder { + public: + MaxPoolOpBuilder() { pool_type_ = tim::vx::PoolType::MAX; } +}; + +} // namespace npu + +} // namespace vsi +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/vsinpu/builders/impl/qlinear_binary_op_builder.h b/onnxruntime/core/providers/vsinpu/builders/impl/qlinear_binary_op_builder.h new file mode 100644 index 0000000000000..def37b1ec1019 --- /dev/null +++ b/onnxruntime/core/providers/vsinpu/builders/impl/qlinear_binary_op_builder.h @@ -0,0 +1,85 @@ +/**************************************************************************** + * + * Copyright (c) 2024 Vivante Corporation + * + * Permission is hereby granted, free of charge, to any person obtaining a + * copy of this software and associated documentation files (the "Software"), + * to deal in the Software without restriction, including without limitation + * the rights to use, copy, modify, merge, publish, distribute, sublicense, + * and/or sell copies of the Software, and to permit persons to whom the + * Software is furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING + * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER + * DEALINGS IN THE SOFTWARE. + * + *****************************************************************************/ +#include +#include +#include +#include "core/providers/vsinpu/builders/impl/base_op_builder.h" + +namespace onnxruntime { +namespace vsi { +namespace npu { +class BaseQLinearOpBuilder : public BaseOpBuilder { + enum { + INPUT_A = 0, + INPUT_A_SCALE = 1, + INPUT_A_ZP = 2, + INPUT_B = 3, + INPUT_B_SCALE = 4, + INPUT_B_ZP = 5, + OUTPUT_SCALE = 6, + OUTPUT_ZP = 7, + }; + + protected: + bool IsOpSupported(const onnxruntime::GraphViewer& graph_viewer, const Node* node) const override { + for (int i = 0; i < node->InputDefs().size(); i++) { + if (i == INPUT_A || i == INPUT_B) continue; + if (!graph_viewer.IsConstantInitializer(node->InputDefs()[i]->Name(), true)) { + LOGS_DEFAULT(WARNING) << "Only support const scale / zero point."; + return false; + } + } + return true; + } +}; + +class QLinearAddOpBuilder : public BaseQLinearOpBuilder { + bool HandleBuildOp(vsi::npu::GraphEP* graph_ep, + std::vector>& inputs, + std::vector>& outputs, + const NodeUnit& node_unit) override { + LOGS_DEFAULT(VERBOSE) << "Creating QLinearAdd Op."; + auto op = graph_ep->GetGraph()->CreateOperation(); + (*op).BindInputs(inputs).BindOutputs(outputs); + graph_ep->GetOps().push_back(std::move(op)); + return true; + } +}; + +class QLinearMulOpBuilder : public BaseQLinearOpBuilder { + bool HandleBuildOp(vsi::npu::GraphEP* graph_ep, + std::vector>& inputs, + std::vector>& outputs, + const NodeUnit& node_unit) override { + LOGS_DEFAULT(VERBOSE) << "Creating QLinearMul Op."; + auto op = graph_ep->GetGraph()->CreateOperation(); + (*op).BindInputs(inputs).BindOutputs(outputs); + graph_ep->GetOps().push_back(std::move(op)); + return true; + } +}; + +} // namespace npu +} // namespace vsi +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/vsinpu/builders/impl/qlinearconcat_op_builder.h b/onnxruntime/core/providers/vsinpu/builders/impl/qlinearconcat_op_builder.h new file mode 100644 index 0000000000000..dc51e99730c15 --- /dev/null +++ b/onnxruntime/core/providers/vsinpu/builders/impl/qlinearconcat_op_builder.h @@ -0,0 +1,48 @@ +/**************************************************************************** + * + * Copyright (c) 2024 Vivante Corporation + * + * Permission is hereby granted, free of charge, to any person obtaining a + * copy of this software and associated documentation files (the "Software"), + * to deal in the Software without restriction, including without limitation + * the rights to use, copy, modify, merge, publish, distribute, sublicense, + * and/or sell copies of the Software, and to permit persons to whom the + * Software is furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING + * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER + * DEALINGS IN THE SOFTWARE. + * + *****************************************************************************/ +#include +#include +#include +#include "core/providers/shared/utils/utils.h" +#include "core/providers/vsinpu/builders/impl/base_op_builder.h" +namespace onnxruntime { +namespace vsi { +namespace npu { +class QLinearConcatOpBuilder : public BaseOpBuilder { + bool HandleBuildOp(vsi::npu::GraphEP* graph_ep, std::vector>& inputs, + std::vector>& outputs, const NodeUnit& node_unit) override { + LOGS_DEFAULT(VERBOSE) << "Creating QLinearConcat Op."; + NodeAttrHelper helper(node_unit.GetNode()); + int axis = helper.Get("axis", 0); + axis = util::ReverseAxis(axis, inputs[0]->GetShape().size()); + auto op = graph_ep->GetGraph()->CreateOperation(axis, inputs.size()); + (*op).BindInputs(inputs).BindOutputs(outputs); + graph_ep->GetOps().push_back(std::move(op)); + return true; + } +}; + +} // namespace npu +} // namespace vsi +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/vsinpu/builders/impl/qlinearconv_op_builder.h b/onnxruntime/core/providers/vsinpu/builders/impl/qlinearconv_op_builder.h new file mode 100644 index 0000000000000..8b63a07e17f1d --- /dev/null +++ b/onnxruntime/core/providers/vsinpu/builders/impl/qlinearconv_op_builder.h @@ -0,0 +1,151 @@ +/**************************************************************************** + * + * Copyright (c) 2024 Vivante Corporation + * + * Permission is hereby granted, free of charge, to any person obtaining a + * copy of this software and associated documentation files (the "Software"), + * to deal in the Software without restriction, including without limitation + * the rights to use, copy, modify, merge, publish, distribute, sublicense, + * and/or sell copies of the Software, and to permit persons to whom the + * Software is furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING + * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER + * DEALINGS IN THE SOFTWARE. + * + *****************************************************************************/ +#include +#include +#include +#include +#include "core/providers/shared/utils/utils.h" +#include "core/providers/vsinpu/builders/impl/base_op_builder.h" +#include "core/framework/tensorprotoutils.h" +namespace onnxruntime { +namespace vsi { +namespace npu { +class QLinearConvOpBuilder : public BaseOpBuilder { + enum QLinearConvINPUTS { + INPUT_TENSOR = 0, + INPUT_TENSOR_SCALE = 1, + INPUT_TENSOR_ZP = 2, + WEIGHT_TENSOR = 3, + WEIGHT_TENSOR_SCALE = 4, + WEIGHT_TENSOR_ZP = 5, + OUTPUT_TENSOR_SCALE = 6, + OUTPUT_TENSOR_ZP = 7, + BIAS_TENSOR = 8, + }; + + bool IsOpSupported(const onnxruntime::GraphViewer& graph_viewer, + const Node* node) const override { + auto input_defs = node->InputDefs(); + auto input_shape = vsi::npu::util::GetTensorShape(*input_defs[QLinearConvINPUTS::INPUT_TENSOR]); + auto w_scale_shape = vsi::npu::util::GetTensorShape(*input_defs[QLinearConvINPUTS::WEIGHT_TENSOR_SCALE]); + auto w_shape_dims = vsi::npu::util::GetTensorShape(*input_defs[QLinearConvINPUTS::WEIGHT_TENSOR]).GetDims(); + if (input_shape.NumDimensions() != 4) { + LOGS_DEFAULT(WARNING) << "Not support conv3d&& conv1d yet."; + return false; + } + + if (!graph_viewer.IsConstantInitializer(input_defs[QLinearConvINPUTS::INPUT_TENSOR_SCALE]->Name(), true) || + !graph_viewer.IsConstantInitializer(input_defs[WEIGHT_TENSOR]->Name(), true)) { + LOGS_DEFAULT(WARNING) << "Not support quantization definitions or weights that are not constant yet."; + return false; + } + + if (w_shape_dims[2] > 15) { + LOGS_DEFAULT(WARNING) << "Not support weight kernel with height higher than 15."; + return false; + } + + if (w_scale_shape.Size() != 1 && *input_defs[WEIGHT_TENSOR]->Type() == "tensor(int8)") { + const ONNX_NAMESPACE::TensorProto* tensor_proto = + graph_viewer.GetConstantInitializer(input_defs[QLinearConvINPUTS::WEIGHT_TENSOR_ZP]->Name(), true); + std::vector w_zp(tensor_proto->dims_size() == 0 ? 1 : tensor_proto->dims()[0]); + + auto status = onnxruntime::utils::UnpackTensor( + *tensor_proto, + tensor_proto->has_raw_data() ? tensor_proto->raw_data().data() : nullptr, + tensor_proto->has_raw_data() ? tensor_proto->raw_data().size() : 0, + w_zp.data(), w_zp.size()); + if (!status.IsOK()) { + LOGS_DEFAULT(ERROR) << "Failed to get data from weight zp tensor."; + return false; + } + if (std::any_of(w_zp.begin(), w_zp.end(), [](int i) { return i != 0; })) { + LOGS_DEFAULT(WARNING) << "Asymmetric perchannel quantization only allows uint8 datatype or int8 with all zero."; + return false; + } + } + return true; + } + bool HandleBuildOp(vsi::npu::GraphEP* graph_ep, + std::vector>& inputs, + std::vector>& outputs, + const NodeUnit& node_unit) override { + LOGS_DEFAULT(VERBOSE) << "Creating QLinearConv Op."; + + NodeAttrHelper helper(node_unit.GetNode()); + auto padtype = helper.Get("auto_pad", std::string("")); + auto group = helper.Get("group", static_cast(1)); + std::vector default_vec = {1, 1, 1, 1}; + auto stride = + helper.Get("strides", default_vec); + auto dilation = + helper.Get("dilations", default_vec); + std::shared_ptr op; + if (padtype != "NOTSET") { // array "pads" is not set + if (group != 1 && group != inputs[1]->GetShape()[3]) { + op = graph_ep->GetGraph() + ->CreateOperation( + vsi::npu::util::GetPadType(padtype), + std::array{stride[1], stride[0]}, + std::array{dilation[1], dilation[0]}, group, + tim::vx::DataLayout::WHCN, tim::vx::DataLayout::WHIcOc); + + } else { + int32_t multiplier = group == 1 ? 0 : inputs[1]->GetShape()[3] / inputs[0]->GetShape()[2]; + op = graph_ep->GetGraph()->CreateOperation( + vsi::npu::util::GetPadType(padtype), + std::array{stride[1], stride[0]}, + std::array{dilation[1], dilation[0]}, multiplier, + tim::vx::DataLayout::WHCN, tim::vx::DataLayout::WHIcOc); + } + } else { + std::vector default_pads(4, 0); + auto pads = helper.Get("pads", default_pads); + if (group != 1 && group != inputs[1]->GetShape()[3]) { + op = graph_ep->GetGraph() + ->CreateOperation( + std::array{pads[1], pads[3], pads[0], pads[2]}, + std::array{stride[1], stride[0]}, + std::array{dilation[1], dilation[0]}, group, + tim::vx::DataLayout::WHCN, tim::vx::DataLayout::WHIcOc); + + } else { + int32_t multiplier = group == 1 ? 0 : inputs[1]->GetShape()[3] / inputs[0]->GetShape()[2]; + op = graph_ep->GetGraph()->CreateOperation( + std::array{pads[1], pads[3], + pads[0], pads[2]}, + std::array{stride[1], stride[0]}, + std::array{dilation[1], dilation[0]}, multiplier, + tim::vx::DataLayout::WHCN, tim::vx::DataLayout::WHIcOc); + } + } + (*op).BindInputs(inputs).BindOutputs(outputs); + graph_ep->GetOps().push_back(std::move(op)); + return true; + } +}; +} // namespace npu + +} // namespace vsi +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/vsinpu/builders/impl/qlinearmatmul_op_builder.h b/onnxruntime/core/providers/vsinpu/builders/impl/qlinearmatmul_op_builder.h new file mode 100644 index 0000000000000..7447c8b6b0b91 --- /dev/null +++ b/onnxruntime/core/providers/vsinpu/builders/impl/qlinearmatmul_op_builder.h @@ -0,0 +1,83 @@ +/**************************************************************************** + * + * Copyright (c) 2024 Vivante Corporation + * + * Permission is hereby granted, free of charge, to any person obtaining a + * copy of this software and associated documentation files (the "Software"), + * to deal in the Software without restriction, including without limitation + * the rights to use, copy, modify, merge, publish, distribute, sublicense, + * and/or sell copies of the Software, and to permit persons to whom the + * Software is furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING + * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER + * DEALINGS IN THE SOFTWARE. + * + *****************************************************************************/ +#include +#include +#include +#include "core/providers/vsinpu/builders/impl/base_op_builder.h" + +namespace onnxruntime { +namespace vsi { +namespace npu { + +class QLinearMatMulOpBuilder : public BaseOpBuilder { + enum { + matrixA = 0, + A_scale = 1, + A_zero_point = 2, + matrixB = 3, + B_scale = 4, + B_zero_point = 5, + out_scale = 6, + out_zero_point = 7 + }; + bool IsOpSupported(const onnxruntime::GraphViewer& graph_viewer, + const Node* node) const override { + auto input_defs = node->InputDefs(); + auto A_def = input_defs[matrixA]; + auto B_def = input_defs[matrixB]; + for (auto def : input_defs) { + if (def->Name() == A_def->Name() || def->Name() == B_def->Name()) { + continue; + } else { + if (!graph_viewer.IsConstantInitializer(def->Name(), true)) { + LOGS_DEFAULT(WARNING) << "Scale and zero point must be known before setting graph."; + return false; + } + } + } + int64_t A_elements = util::GetTensorShape(*input_defs[A_scale]).Size(); + int64_t B_elements = util::GetTensorShape(*input_defs[B_scale]).Size(); + int64_t Out_elements = util::GetTensorShape(*input_defs[out_scale]).Size(); + if (A_elements > 1 || B_elements > 1 || Out_elements > 1) { + LOGS_DEFAULT(WARNING) << "Per channel quantized input/output is not supported in QLinearMatmul Op."; + return false; + } + + return true; + } + bool HandleBuildOp(vsi::npu::GraphEP* graph_ep, + std::vector>& inputs, + std::vector>& outputs, + const NodeUnit& node_unit) override { + LOGS_DEFAULT(INFO) << "Creating QLinearMatmul Op."; + auto op = graph_ep->GetGraph()->CreateOperation(); + (*op).BindInputs(inputs).BindOutputs(outputs); + graph_ep->GetOps().push_back(std::move(op)); + return true; + } +}; +} // namespace npu + +} // namespace vsi +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/vsinpu/builders/impl/quantize_op_builder.h b/onnxruntime/core/providers/vsinpu/builders/impl/quantize_op_builder.h new file mode 100644 index 0000000000000..63ae491909bdc --- /dev/null +++ b/onnxruntime/core/providers/vsinpu/builders/impl/quantize_op_builder.h @@ -0,0 +1,79 @@ +/**************************************************************************** + * + * Copyright (c) 2024 Vivante Corporation + * + * Permission is hereby granted, free of charge, to any person obtaining a + * copy of this software and associated documentation files (the "Software"), + * to deal in the Software without restriction, including without limitation + * the rights to use, copy, modify, merge, publish, distribute, sublicense, + * and/or sell copies of the Software, and to permit persons to whom the + * Software is furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING + * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER + * DEALINGS IN THE SOFTWARE. + * + *****************************************************************************/ +#include +#include +#include +#include "core/providers/vsinpu/builders/impl/base_op_builder.h" +#include "core/providers/common.h" +#include "core/providers/shared/utils/utils.h" + +namespace onnxruntime { +namespace vsi { +namespace npu { + +class QuantizeLinearOpBuilder : public BaseOpBuilder { + enum QuantizeINPUTS { + input_tensor = 0, + scale_tensor = 1, + zero_point_tensor = 2 + }; + + bool IsOpSupported(const onnxruntime::GraphViewer& graph_viewer, + const Node* node) const override { + auto input_defs = node->InputDefs(); + auto scale_shape = npu::util::GetTensorShape(*input_defs[QuantizeINPUTS::scale_tensor]); + NodeAttrHelper helper(*node); + if (helper.HasAttr("block_size") && helper.Get("block_size", 0) != 0) { + LOGS_DEFAULT(WARNING) << "Not support block quantization."; + return false; + } + if (!graph_viewer.IsConstantInitializer(input_defs[QuantizeINPUTS::scale_tensor]->Name(), true) || + (input_defs.size() == 3 && !graph_viewer.IsConstantInitializer( + input_defs[QuantizeINPUTS::zero_point_tensor]->Name(), true))) { + LOGS_DEFAULT(WARNING) << "Only support const scale / zero point."; + return false; + } + + if (scale_shape.Size() != 1) { + LOGS_DEFAULT(WARNING) << "Per channel quantized output is not supported in QuantizeLinearOp."; + return false; + } + return true; + } + + bool HandleBuildOp(vsi::npu::GraphEP* graph_ep, + std::vector>& inputs, + std::vector>& outputs, + const NodeUnit& node_unit) override { + LOGS_DEFAULT(INFO) << "Creating Quantize Op."; + auto op = graph_ep->GetGraph()->CreateOperation(); + (*op).BindInputs(inputs).BindOutputs(outputs); + graph_ep->GetOps().push_back(std::move(op)); + return true; + } +}; +} // namespace npu + +} // namespace vsi +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/vsinpu/builders/impl/reduce_op_builder.h b/onnxruntime/core/providers/vsinpu/builders/impl/reduce_op_builder.h new file mode 100644 index 0000000000000..3b0a282c5de89 --- /dev/null +++ b/onnxruntime/core/providers/vsinpu/builders/impl/reduce_op_builder.h @@ -0,0 +1,82 @@ +/**************************************************************************** + * + * Copyright (c) 2024 Vivante Corporation + * + * Permission is hereby granted, free of charge, to any person obtaining a + * copy of this software and associated documentation files (the "Software"), + * to deal in the Software without restriction, including without limitation + * the rights to use, copy, modify, merge, publish, distribute, sublicense, + * and/or sell copies of the Software, and to permit persons to whom the + * Software is furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING + * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER + * DEALINGS IN THE SOFTWARE. + * + *****************************************************************************/ +#include +#include +#include +#include "core/providers/vsinpu/builders/impl/base_op_builder.h" +#include "core/providers/shared/utils/utils.h" + +namespace onnxruntime { +namespace vsi { +namespace npu { +class ReduceMeanOpBuilder : public BaseOpBuilder { + bool IsOpSupported(const onnxruntime::GraphViewer& graph_viewer, + const Node* node) const override { + auto input_defs = node->InputDefs(); + if (*input_defs[0]->Type() == "tensor(int32)") { + LOGS_DEFAULT(WARNING) << "Not support int32 reduce mean yet."; + return false; + } + return true; + } + bool HandleBuildOp(vsi::npu::GraphEP* graph_ep, + std::vector>& inputs, + std::vector>& outputs, + const NodeUnit& node_unit) override { + LOGS_DEFAULT(INFO) << "Creating ReduceMean Op."; + + NodeAttrHelper helper(node_unit.GetNode()); + std::vector def_axes; + auto input_shape_size = inputs[0]->GetShape().size(); + + if (node_unit.SinceVersion() < 18 && helper.HasAttr("axes")) { + def_axes = helper.Get("axes", def_axes); + } else if (inputs.size() > 1) { + def_axes.resize(inputs[1]->GetSpec().GetElementNum()); + inputs[1]->CopyDataFromTensor(def_axes.data()); + } else { + for (int64_t i = 0; i < input_shape_size; ++i) { + def_axes.push_back(i); + } + } + + std::vector axes(def_axes.begin(), def_axes.end()); + axes = util::ReverseAxis(axes, input_shape_size); + + if (helper.HasAttr("noop_with_empty_axes") && inputs.size() == 1 && helper.Get("noop_with_empty_axes", 0) == 1) { + outputs[0] = inputs[0]; + return true; + } + + bool keepdims = helper.Get("keepdims", 1) == 1; + auto op = graph_ep->GetGraph()->CreateOperation(axes, keepdims); + (*op).BindInput(inputs[0]).BindOutputs(outputs); + graph_ep->GetOps().push_back(std::move(op)); + return true; + } +}; +} // namespace npu + +} // namespace vsi +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/vsinpu/builders/impl/resize_op_builder.h b/onnxruntime/core/providers/vsinpu/builders/impl/resize_op_builder.h new file mode 100644 index 0000000000000..8857efe3537ec --- /dev/null +++ b/onnxruntime/core/providers/vsinpu/builders/impl/resize_op_builder.h @@ -0,0 +1,153 @@ +/**************************************************************************** + * + * Copyright (c) 2024 Vivante Corporation + * + * Permission is hereby granted, free of charge, to any person obtaining a + * copy of this software and associated documentation files (the "Software"), + * to deal in the Software without restriction, including without limitation + * the rights to use, copy, modify, merge, publish, distribute, sublicense, + * and/or sell copies of the Software, and to permit persons to whom the + * Software is furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING + * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER + * DEALINGS IN THE SOFTWARE. + * + *****************************************************************************/ +#include +#include +#include +#include "core/providers/vsinpu/builders/impl/base_op_builder.h" +#include "core/providers/shared/utils/utils.h" + +namespace onnxruntime { +namespace vsi { +namespace npu { +class ResizeOpBuilder : public BaseOpBuilder { + bool HasSupportedInputOutputsImpl(const InitializedTensorSet& initializers, + const NodeUnit& node_unit) const override { + auto input_type = node_unit.Inputs()[0].node_arg.Type(); + if (*input_type == "tensor(int64)" || !util::IsTypeSupported(&node_unit.Inputs()[0].node_arg)) { + LOGS_DEFAULT(WARNING) << node_unit.OpType() << " has unsupported input type : " + << *input_type; + return false; + } + if (node_unit.SinceVersion() > 10) { + if (node_unit.Inputs().size() > 2 && !Contains(initializers, node_unit.Inputs()[2].node_arg.Name())) { + LOGS_DEFAULT(WARNING) << "Scale tensor must be constant."; + return false; + } + if (node_unit.Inputs().size() > 3 && !Contains(initializers, node_unit.Inputs()[3].node_arg.Name())) { + LOGS_DEFAULT(WARNING) << "Size tensor must be constant."; + return false; + } + } else { + if (!Contains(initializers, node_unit.Inputs()[1].node_arg.Name())) { + LOGS_DEFAULT(WARNING) << "Scale tensor must be constant."; + return false; + } + } + return true; + } + bool IsOpSupported(const onnxruntime::GraphViewer& graph_viewer, const Node* node) const override { + auto shape = vsi::npu::util::GetTensorShape(*node->InputDefs()[0]); + if (shape.NumDimensions() > 4) { + LOGS_DEFAULT(WARNING) << "3D or more dimesions resize is not supported."; + return false; + } + + NodeAttrHelper helper(*node); + if (helper.Get("antialiax", 0) != 0) { + LOGS_DEFAULT(WARNING) << "Antialias attribute is not supported."; + return false; + } + auto& cooridinate = helper.Get("coordinate_transoformation_mode", "half_pixel"); + if (cooridinate != "align_corners" && cooridinate != "half_pixel") { + LOGS_DEFAULT(WARNING) << "Only support half_pixel and align_corners attributes now."; + return false; + } + if (helper.Get("keep_aspect_ratio_policy", "stretch") != "stretch") { + LOGS_DEFAULT(WARNING) << "Not support to keep aspect ratio."; + return false; + } + if (helper.Get("mode", "nearest") == "cubic") { + LOGS_DEFAULT(WARNING) << "Not support the cubic resize type yet."; + return false; + } + if (helper.HasAttr("axes")) { + LOGS_DEFAULT(WARNING) << "Axes-specifying is not support."; + return false; + } + return true; + } + + bool HandleBuildOp(vsi::npu::GraphEP* graph_ep, + std::vector>& inputs, + std::vector>& outputs, + const NodeUnit& node_unit) override { + LOGS_DEFAULT(VERBOSE) << "Creating Resize Op."; + auto inputs_num = inputs.size(); + bool is_1dresize = inputs[0]->GetShape().size() == 1; + NodeAttrHelper helper(node_unit.GetNode()); + auto onnx_mode = helper.Get("mode", "nearest"); + auto coordinate_transformation = helper.Get("coordinate_transformation_mode", "half_pixel"); + bool is_size_set = helper.HasAttr("size"); + int32_t scale_index = node_unit.SinceVersion() > 10 ? 2 : 1; + + auto resize_type = onnx_mode == "nearest" ? tim::vx::ResizeType::NEAREST_NEIGHBOR : tim::vx::ResizeType::BILINEAR; + bool align_corners = coordinate_transformation == "align_corners"; + bool half_pixel_center = coordinate_transformation == "half_pixel"; + std::shared_ptr op = nullptr; + if (is_1dresize) { + int target_size; + if (is_size_set) { + int64_t onnx_size; + inputs[3]->CopyDataFromTensor(&onnx_size); + target_size = static_cast(onnx_size); + op = graph_ep->GetGraph()->CreateOperation(resize_type, 0.0f, align_corners, + half_pixel_center, target_size); + } else { + float scale; + inputs[scale_index]->CopyDataFromTensor(&scale); + op = graph_ep->GetGraph()->CreateOperation(resize_type, scale, align_corners, + half_pixel_center, 0); + } + } else { + int target_h, target_w; + if (is_size_set) { + std::vector onnx_sizes(inputs[3]->GetShape().size()); + inputs[3]->CopyDataFromTensor(onnx_sizes.data()); + target_h = static_cast(onnx_sizes[1]); + target_w = static_cast(onnx_sizes[0]); + op = graph_ep->GetGraph()->CreateOperation(resize_type, 0.0f, align_corners, + half_pixel_center, target_h, target_w); + } else { + auto input_shape = inputs[0]->GetShape(); + std::vector scales(input_shape.size()); + std::vector out_shape(input_shape.size()); + inputs[scale_index]->CopyDataFromTensor(scales.data()); + for (int i = 0; i < input_shape.size(); i++) { + out_shape[i] = input_shape[i] * scales[input_shape.size() - 1 - i]; + } + op = graph_ep->GetGraph()->CreateOperation(resize_type, 0, align_corners, + half_pixel_center, out_shape[1], out_shape[0]); + } + } + + (*op).BindInput(inputs[0]).BindOutputs(outputs); + graph_ep->GetOps().push_back(std::move(op)); + return true; + } +}; + +} // namespace npu + +} // namespace vsi +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/vsinpu/builders/impl/softmax_op_builder.h b/onnxruntime/core/providers/vsinpu/builders/impl/softmax_op_builder.h new file mode 100644 index 0000000000000..dad10c1a02518 --- /dev/null +++ b/onnxruntime/core/providers/vsinpu/builders/impl/softmax_op_builder.h @@ -0,0 +1,101 @@ +/**************************************************************************** + * + * Copyright (c) 2024 Vivante Corporation + * + * Permission is hereby granted, free of charge, to any person obtaining a + * copy of this software and associated documentation files (the "Software"), + * to deal in the Software without restriction, including without limitation + * the rights to use, copy, modify, merge, publish, distribute, sublicense, + * and/or sell copies of the Software, and to permit persons to whom the + * Software is furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING + * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER + * DEALINGS IN THE SOFTWARE. + * + *****************************************************************************/ +#include +#include +#include +#include +#include "core/providers/vsinpu/builders/impl/base_op_builder.h" +#include "core/providers/common.h" +#include "core/providers/shared/utils/utils.h" + +namespace onnxruntime { +namespace vsi { +namespace npu { +class SoftmaxOpBuilder : public BaseOpBuilder { + bool IsOpSupported(const onnxruntime::GraphViewer& graph_viewer, + const Node* node) const override { + NodeAttrHelper helper(*node); + auto axis = helper.Get("axis", -1); + auto input_defs = node->InputDefs(); + auto input_shape = vsi::npu::util::GetTensorShape(*input_defs[0]); + int32_t rank = input_shape.NumDimensions(); + if (axis >= rank || axis < -rank) { + LOGS_DEFAULT(ERROR) << "Axis is invalid in Softmax."; + return false; + } + return true; + } + bool HandleBuildOp(vsi::npu::GraphEP* graph_ep, + std::vector>& inputs, + std::vector>& outputs, + const NodeUnit& node_unit) override { + LOGS_DEFAULT(VERBOSE) << "Creating Softmax Op."; + NodeAttrHelper helper(node_unit.GetNode()); + int32_t def_val = node_unit.SinceVersion() < 13 ? 1 : -1; + auto axis = helper.Get("axis", def_val); + + if (def_val == 1) { + // In earlier opset version of softmax, input is coerced into 2D shape + // Attribute "axis" is to describe the axis of the inputs coerced to 2D but not take part in softmax computation + const bool is_2d_shape = inputs[0]->GetShape().size() == 2 ? true : false; + if (!is_2d_shape) { + axis = HandleNegativeAxis(axis, inputs[0]->GetShape().size()); + auto it = inputs[0]->GetShape().end(); + uint32_t last_dim = std::accumulate(it - axis, it, 1, std::multiplies()); + uint32_t first_dim = std::accumulate(inputs[0]->GetShape().begin(), it - axis, 1, std::multiplies()); + auto reshaped_spec = inputs[0]->GetSpec().AsTransientSpec().SetShape( + std::vector{first_dim, last_dim}); + auto reshaped_input = graph_ep->GetGraph()->CreateTensor(reshaped_spec); + auto reshaped_output = graph_ep->GetGraph()->CreateTensor(inputs[0]->GetSpec().AsTransientSpec()); + + auto reshape_input_op = graph_ep->GetGraph()->CreateOperation( + std::vector{first_dim, last_dim}); + auto softmax_op = graph_ep->GetGraph()->CreateOperation(1, 0); + auto reshaped_output_op = graph_ep->GetGraph()->CreateOperation(inputs[0]->GetShape()); + + (*reshape_input_op).BindInputs(inputs).BindOutput(reshaped_input); + (*softmax_op).BindInput(reshaped_input).BindOutput(reshaped_output); + (*reshaped_output_op).BindInput(reshaped_output).BindOutputs(outputs); + + graph_ep->GetOps().push_back(std::move(reshape_input_op)); + graph_ep->GetOps().push_back(std::move(softmax_op)); + graph_ep->GetOps().push_back(std::move(reshaped_output_op)); + } else { + auto op = graph_ep->GetGraph()->CreateOperation(1, 0); + (*op).BindInputs(inputs).BindOutputs(outputs); + graph_ep->GetOps().push_back(std::move(op)); + } + } else { + axis = util::ReverseAxis(axis, inputs[0]->GetShape().size()); + auto op = graph_ep->GetGraph()->CreateOperation(1, static_cast(axis)); + (*op).BindInputs(inputs).BindOutputs(outputs); + graph_ep->GetOps().push_back(std::move(op)); + } + return true; + } +}; +} // namespace npu + +} // namespace vsi +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/vsinpu/builders/impl/squeeze_op_builder.h b/onnxruntime/core/providers/vsinpu/builders/impl/squeeze_op_builder.h new file mode 100644 index 0000000000000..2e1837384618d --- /dev/null +++ b/onnxruntime/core/providers/vsinpu/builders/impl/squeeze_op_builder.h @@ -0,0 +1,88 @@ +/**************************************************************************** + * + * Copyright (c) 2024 Vivante Corporation + * + * Permission is hereby granted, free of charge, to any person obtaining a + * copy of this software and associated documentation files (the "Software"), + * to deal in the Software without restriction, including without limitation + * the rights to use, copy, modify, merge, publish, distribute, sublicense, + * and/or sell copies of the Software, and to permit persons to whom the + * Software is furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING + * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER + * DEALINGS IN THE SOFTWARE. + * + *****************************************************************************/ +#include +#include +#include +#include "core/providers/vsinpu/builders/impl/base_op_builder.h" +#include "core/providers/shared/utils/utils.h" + +namespace onnxruntime { +namespace vsi { +namespace npu { +class SqueezeOpBuilder : public BaseOpBuilder { + bool HasSupportedInputOutputsImpl(const InitializedTensorSet& initializers, + const NodeUnit& node_unit) const override { + auto input_type = node_unit.Inputs()[0].node_arg.Type(); + if (*input_type == "tensor(int64)" || !util::IsTypeSupported(&node_unit.Inputs()[0].node_arg)) { + LOGS_DEFAULT(WARNING) << node_unit.OpType() << " has unsupported input type : " + << *input_type; + return false; + } + if (node_unit.SinceVersion() > 11) { + if (node_unit.Inputs().size() > 1 && !Contains(initializers, node_unit.Inputs()[1].node_arg.Name())) { + LOGS_DEFAULT(WARNING) << "Only support const axes in Squeeze op."; + return false; + } + } + return true; + } + + bool HandleBuildOp(vsi::npu::GraphEP* graph_ep, + std::vector>& inputs, + std::vector>& outputs, + const NodeUnit& node_unit) override { + LOGS_DEFAULT(INFO) << "Creating Squeeze Op."; + + NodeAttrHelper helper(node_unit.GetNode()); + std::vector def_axes; + auto input_shape_size = inputs[0]->GetShape().size(); + + if (node_unit.SinceVersion() < 13 && helper.HasAttr("axes")) { + def_axes = helper.Get("axes", def_axes); + } else if (inputs.size() > 1) { + def_axes.resize(inputs[1]->GetSpec().GetElementNum()); + inputs[1]->CopyDataFromTensor(def_axes.data()); + } else { // if axes is empty from onnx, check input shape to determine + for (int64_t i = 0; i < input_shape_size; ++i) { + if (inputs[0]->GetShape()[i] == 1) { + def_axes.push_back(i); + } + } + } + + std::vector axes(def_axes.begin(), def_axes.end()); + axes = util::ReverseAxis(axes, input_shape_size); + + std::vector timvx_axes(axes.begin(), axes.end()); + + auto op = graph_ep->GetGraph()->CreateOperation(timvx_axes); + (*op).BindInput(inputs[0]).BindOutputs(outputs); + graph_ep->GetOps().push_back(std::move(op)); + return true; + } +}; +} // namespace npu + +} // namespace vsi +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/vsinpu/builders/impl/tensor_op_builder.h b/onnxruntime/core/providers/vsinpu/builders/impl/tensor_op_builder.h new file mode 100644 index 0000000000000..427457b521b61 --- /dev/null +++ b/onnxruntime/core/providers/vsinpu/builders/impl/tensor_op_builder.h @@ -0,0 +1,142 @@ +/**************************************************************************** + * + * Copyright (c) 2024 Vivante Corporation + * + * Permission is hereby granted, free of charge, to any person obtaining a + * copy of this software and associated documentation files (the "Software"), + * to deal in the Software without restriction, including without limitation + * the rights to use, copy, modify, merge, publish, distribute, sublicense, + * and/or sell copies of the Software, and to permit persons to whom the + * Software is furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING + * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER + * DEALINGS IN THE SOFTWARE. + * + *****************************************************************************/ +#include +#include +#include +#include "core/providers/vsinpu/builders/impl/base_op_builder.h" +#include "core/providers/shared/utils/utils.h" + +namespace onnxruntime { +namespace vsi { +namespace npu { +class ReshapeOpBuilder : public BaseOpBuilder { + int GetMinSupportedOpSet(const NodeUnit& /* node_unit */) const override { return 5; } + + bool HasSupportedInputOutputsImpl(const InitializedTensorSet& initializers, + const NodeUnit& node_unit) const override { + auto input = node_unit.Inputs()[0]; + auto shape = node_unit.Inputs()[1]; + if (initializers.end() == initializers.find(shape.node_arg.Name())) { + LOGS_DEFAULT(VERBOSE) << "Target shape of reshape op must be known."; + return false; + } + if (util::IsTypeSupported(&input.node_arg) && util::IsTypeSupported(&shape.node_arg)) { + if (*input.node_arg.Type() != "tensor(int64)") { + return true; + } + } + return false; + } + + bool IsOpSupported(const onnxruntime::GraphViewer& graph_viewer, + const Node* node) const override { + auto input_defs = node->InputDefs(); + + NodeAttrHelper helper(*node); + const bool allow_zero = helper.Get("allowzero", 0) == 1; + auto& perm_tensor_proto = *graph_viewer.GetConstantInitializer(input_defs[1]->Name(), true); + std::vector perm(perm_tensor_proto.dims()[0]); + auto status = onnxruntime::utils::UnpackTensor( + perm_tensor_proto, + perm_tensor_proto.has_raw_data() ? perm_tensor_proto.raw_data().data() : nullptr, + perm_tensor_proto.has_raw_data() ? perm_tensor_proto.raw_data().size() : 0, + perm.data(), perm.size()); + + // Check if perm has any 0's when allow zero is enabled. + if (allow_zero && std::find(perm.begin(), perm.end(), 0L) != perm.end()) { + LOGS_DEFAULT(VERBOSE) << "Reshape doesn't support 0 as dimension when allowzero is enabled"; + return false; + } + + return true; + } + bool HandleBuildOp(vsi::npu::GraphEP* graph_ep, + std::vector>& inputs, + std::vector>& outputs, + const NodeUnit& node_unit) override { + LOGS_DEFAULT(VERBOSE) << "Creating Reshape Op."; + std::vector new_shape(inputs[1]->GetShape()[0]); + inputs[1]->CopyDataFromTensor(new_shape.data()); + for (size_t i = 0; i < new_shape.size(); i++) { + if (new_shape[i] == 0) { + new_shape[i] = inputs[0]->GetShape()[inputs[0]->GetShape().size() - i - 1]; + } + } + + int64_t element_count = std::accumulate(new_shape.begin(), new_shape.end(), static_cast(1), + [&](int64_t a, int64_t b) { + return b == -1 ? a : a * b; + }); + auto negative_it = std::find(new_shape.begin(), new_shape.end(), -1); + if (negative_it != new_shape.end()) { + *negative_it = inputs[0]->GetSpec().GetElementNum() / element_count; + } + + std::vector new_shape_uint32(new_shape.begin(), new_shape.end()); + std::reverse(new_shape_uint32.begin(), new_shape_uint32.end()); + auto op = graph_ep->GetGraph()->CreateOperation(new_shape_uint32); + (*op).BindInput(inputs[0]).BindOutputs(outputs); + graph_ep->GetOps().push_back(std::move(op)); + return true; + } +}; + +class TransposeOpBuilder : public BaseOpBuilder { + bool IsOpSupported(const onnxruntime::GraphViewer& graph_viewer, + const Node* node) const override { + auto input_defs = node->InputDefs(); + auto shape_dim = vsi::npu::util::GetTensorShape(*input_defs[0]).NumDimensions(); + NodeAttrHelper helper(*node); + auto perm = helper.Get("perm", std::vector(shape_dim, 1)); + if (perm.size() != shape_dim) { + LOGS_DEFAULT(VERBOSE) << "Size mismatch between perm vector and input shape."; + return false; + } + return true; + } + bool HandleBuildOp(vsi::npu::GraphEP* graph_ep, + std::vector>& inputs, + std::vector>& outputs, + const NodeUnit& node_unit) override { + LOGS_DEFAULT(VERBOSE) << "Creating Transpose Op."; + std::vector def_val(inputs[0]->GetShape().size()); + for (int64_t i = 0; i < def_val.size(); i++) def_val[i] = def_val.size() - i - 1; + + NodeAttrHelper helper(node_unit.GetNode()); + def_val = helper.Get("perm", def_val); + std::vector timvx_perm; + for (uint32_t i = 0; i < def_val.size(); i++) { + timvx_perm.push_back(def_val.size() - 1 - def_val[def_val.size() - i - 1]); + } + auto op = graph_ep->GetGraph()->CreateOperation(timvx_perm); + (*op).BindInputs(inputs).BindOutputs(outputs); + graph_ep->GetOps().push_back(std::move(op)); + return true; + } +}; + +} // namespace npu + +} // namespace vsi +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/vsinpu/builders/impl/tile_op_builder.h b/onnxruntime/core/providers/vsinpu/builders/impl/tile_op_builder.h new file mode 100644 index 0000000000000..d42624c31557c --- /dev/null +++ b/onnxruntime/core/providers/vsinpu/builders/impl/tile_op_builder.h @@ -0,0 +1,71 @@ +/**************************************************************************** + * + * Copyright (c) 2024 Vivante Corporation + * + * Permission is hereby granted, free of charge, to any person obtaining a + * copy of this software and associated documentation files (the "Software"), + * to deal in the Software without restriction, including without limitation + * the rights to use, copy, modify, merge, publish, distribute, sublicense, + * and/or sell copies of the Software, and to permit persons to whom the + * Software is furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING + * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER + * DEALINGS IN THE SOFTWARE. + * + *****************************************************************************/ +#include +#include +#include +#include "core/providers/vsinpu/builders/impl/base_op_builder.h" + +namespace onnxruntime { +namespace vsi { +namespace npu { +class TileOpBuilder : public BaseOpBuilder { + int GetMinSupportedOpSet(const NodeUnit& /* node_unit */) const override { return 6; } + + bool HasSupportedInputOutputsImpl(const InitializedTensorSet& initializers, + const NodeUnit& node_unit) const override { + auto input = node_unit.Inputs()[0]; + auto multipliers = node_unit.Inputs()[1]; + if (initializers.end() == initializers.find(multipliers.node_arg.Name())) { + LOGS_DEFAULT(WARNING) << "Multipliers of tile op must be known."; + return false; + } + if (util::IsTypeSupported(&input.node_arg) && util::IsTypeSupported(&multipliers.node_arg)) { + if (*input.node_arg.Type() != "tensor(int64)") { + return true; + } + } + LOGS_DEFAULT(WARNING) << "Input type not supported."; + return false; + } + + bool HandleBuildOp(vsi::npu::GraphEP* graph_ep, + std::vector>& inputs, + std::vector>& outputs, + const NodeUnit& node_unit) override { + LOGS_DEFAULT(VERBOSE) << "Creating Tile Op."; + std::vector multipliers(inputs[1]->GetShape()[0]); + inputs[1]->CopyDataFromTensor(multipliers.data()); + std::reverse(multipliers.begin(), multipliers.end()); + std::vector timvx_multipliers(multipliers.begin(), multipliers.end()); + auto op = graph_ep->GetGraph()->CreateOperation(timvx_multipliers); + (*op).BindInput(inputs[0]).BindOutputs(outputs); + graph_ep->GetOps().push_back(std::move(op)); + return true; + } +}; + +} // namespace npu + +} // namespace vsi +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/vsinpu/builders/impl/unsqueeze_op_builder.h b/onnxruntime/core/providers/vsinpu/builders/impl/unsqueeze_op_builder.h new file mode 100644 index 0000000000000..c49c93008b25a --- /dev/null +++ b/onnxruntime/core/providers/vsinpu/builders/impl/unsqueeze_op_builder.h @@ -0,0 +1,89 @@ +/**************************************************************************** + * + * Copyright (c) 2024 Vivante Corporation + * + * Permission is hereby granted, free of charge, to any person obtaining a + * copy of this software and associated documentation files (the "Software"), + * to deal in the Software without restriction, including without limitation + * the rights to use, copy, modify, merge, publish, distribute, sublicense, + * and/or sell copies of the Software, and to permit persons to whom the + * Software is furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING + * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER + * DEALINGS IN THE SOFTWARE. + * + *****************************************************************************/ +#include +#include +#include +#include "core/providers/vsinpu/builders/impl/base_op_builder.h" +#include "core/providers/shared/utils/utils.h" + +namespace onnxruntime { +namespace vsi { +namespace npu { +class UnsqueezeOpBuilder : public BaseOpBuilder { + bool HasSupportedInputOutputsImpl(const InitializedTensorSet& initializers, + const NodeUnit& node_unit) const override { + auto input_type = node_unit.Inputs()[0].node_arg.Type(); + if (*input_type == "tensor(int64)" || !util::IsTypeSupported(&node_unit.Inputs()[0].node_arg)) { + LOGS_DEFAULT(WARNING) << node_unit.OpType() << " has unsupported input type : " + << *input_type; + return false; + } + if (node_unit.SinceVersion() > 11 && !Contains(initializers, node_unit.Inputs()[1].node_arg.Name())) { + LOGS_DEFAULT(WARNING) << "Only support const axes in Unsqueeze op."; + return false; + } + return true; + } + + bool HandleBuildOp(vsi::npu::GraphEP* graph_ep, + std::vector>& inputs, + std::vector>& outputs, + const NodeUnit& node_unit) override { + LOGS_DEFAULT(INFO) << "Creating Unsqueeze Op."; + + NodeAttrHelper helper(node_unit.GetNode()); + std::vector def_axes; + auto input_shape_size = inputs[0]->GetShape().size(); + + if (node_unit.SinceVersion() < 13 && helper.HasAttr("axes")) { + def_axes = helper.Get("axes", def_axes); + } else if (inputs.size() > 1) { + def_axes.resize(inputs[1]->GetSpec().GetElementNum()); + inputs[1]->CopyDataFromTensor(def_axes.data()); + } else { // if axes is empty from onnx, check input shape to determine + for (int64_t i = 0; i < input_shape_size; ++i) { + if (inputs[0]->GetShape()[i] == 1) { + def_axes.push_back(i); + } + } + } + + std::vector axes(def_axes.begin(), def_axes.end()); + axes = util::ReverseAxis(axes, input_shape_size + axes.size()); + + std::vector timvx_axes(inputs[0]->GetShape().begin(), inputs[0]->GetShape().end()); + for (int32_t dim : axes) { + timvx_axes.insert(timvx_axes.begin() + dim, 1); + } + + auto op = graph_ep->GetGraph()->CreateOperation(timvx_axes); + (*op).BindInput(inputs[0]).BindOutputs(outputs); + graph_ep->GetOps().push_back(std::move(op)); + return true; + } +}; +} // namespace npu + +} // namespace vsi +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/vsinpu/builders/op_builder.h b/onnxruntime/core/providers/vsinpu/builders/op_builder.h new file mode 100644 index 0000000000000..d81a478149c6b --- /dev/null +++ b/onnxruntime/core/providers/vsinpu/builders/op_builder.h @@ -0,0 +1,48 @@ +/**************************************************************************** + * + * Copyright (c) 2023 Vivante Corporation + * + * Permission is hereby granted, free of charge, to any person obtaining a + * copy of this software and associated documentation files (the "Software"), + * to deal in the Software without restriction, including without limitation + * the rights to use, copy, modify, merge, publish, distribute, sublicense, + * and/or sell copies of the Software, and to permit persons to whom the + * Software is furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING + * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER + * DEALINGS IN THE SOFTWARE. + * + *****************************************************************************/ +#pragma once +#include "core/graph/graph_viewer.h" +#include "core/framework/node_unit.h" + +namespace onnxruntime { +namespace vsi { +namespace npu { +class GraphEP; + +class IOpBuilder { + public: + IOpBuilder() {} + virtual ~IOpBuilder() {} + virtual bool IsSupported(const onnxruntime::GraphViewer& graph_viewer, + const NodeUnit& node_unit) const { + return true; + } + virtual bool BuildOp(GraphEP* graph_ep, + const onnxruntime::GraphViewer& graph_viewer, + const NodeUnit& node_unit) = 0; +}; +} // namespace npu + +} // namespace vsi +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/vsinpu/builders/op_builder_factory.h b/onnxruntime/core/providers/vsinpu/builders/op_builder_factory.h new file mode 100644 index 0000000000000..3a9190d8cb03a --- /dev/null +++ b/onnxruntime/core/providers/vsinpu/builders/op_builder_factory.h @@ -0,0 +1,133 @@ +/**************************************************************************** + * + * Copyright (c) 2023 Vivante Corporation + * + * Permission is hereby granted, free of charge, to any person obtaining a + * copy of this software and associated documentation files (the "Software"), + * to deal in the Software without restriction, including without limitation + * the rights to use, copy, modify, merge, publish, distribute, sublicense, + * and/or sell copies of the Software, and to permit persons to whom the + * Software is furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING + * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER + * DEALINGS IN THE SOFTWARE. + * + *****************************************************************************/ +#pragma once +#include +#include +#include +#include +#include "impl/activation_op_builder.h" +#include "impl/conv_op_builder.h" +#include "impl/elementwise_op_builder.h" +#include "impl/gemm_op_builder.h" +#include "impl/pool_op_builder.h" +#include "impl/qlinearconv_op_builder.h" +#include "impl/flatten_op_builder.h" +#include "impl/matmul_op_builder.h" +#include "impl/tensor_op_builder.h" +#include "impl/concat_op_builder.h" +#include "impl/softmax_op_builder.h" +#include "impl/norm_op_builder.h" +#include "impl/clip_op_builder.h" +#include "impl/reduce_op_builder.h" +#include "impl/quantize_op_builder.h" +#include "impl/dequantize_op_builder.h" +#include "impl/qlinearmatmul_op_builder.h" +#include "impl/qlinear_binary_op_builder.h" +#include "impl/qlinearconcat_op_builder.h" +#include "impl/gather_op_builder.h" +#include "impl/tile_op_builder.h" +#include "impl/squeeze_op_builder.h" +#include "impl/unsqueeze_op_builder.h" +#include "impl/resize_op_builder.h" +#include "impl/cast_op_builder.h" +namespace onnxruntime { +namespace vsi { +namespace npu { +using createIOpBuildItemFunc = std::function()>; +using OpBuildItemType = std::map>; + +static const std::map reg = { +#define REGISTER_OP_BUILDER(ONNX_NODE_TYPE, BUILDER_TYPE) \ + { \ + ONNX_NODE_TYPE, [] { return std::make_unique(); } \ + } + + REGISTER_OP_BUILDER("Add", AddOpBuilder), + REGISTER_OP_BUILDER("Sub", SubOpBuilder), + REGISTER_OP_BUILDER("Mul", MulOpBuilder), + REGISTER_OP_BUILDER("Div", DivOpBuilder), + REGISTER_OP_BUILDER("Abs", AbsOpBuilder), + REGISTER_OP_BUILDER("Pow", PowOpBuilder), + REGISTER_OP_BUILDER("Sqrt", SqrtOpBuilder), + REGISTER_OP_BUILDER("Exp", ExpOpBuilder), + REGISTER_OP_BUILDER("Floor", FloorOpBuilder), + REGISTER_OP_BUILDER("Log", LogOpBuilder), + REGISTER_OP_BUILDER("Sin", SinOpBuilder), + REGISTER_OP_BUILDER("Conv", ConvOpBuilder), + REGISTER_OP_BUILDER("Gemm", GemmOpBuilder), + REGISTER_OP_BUILDER("Relu", ReluOpBuilder), + REGISTER_OP_BUILDER("LeakyRelu", LeakyReluOpBuilder), + REGISTER_OP_BUILDER("Tanh", TanhOpBuilder), + REGISTER_OP_BUILDER("Sigmoid", SigmoidOpBuilder), + REGISTER_OP_BUILDER("HardSigmoid", HardSigmoidOpBuilder), + REGISTER_OP_BUILDER("HardSwish", HardSwishOpBuilder), + REGISTER_OP_BUILDER("GlobalAveragePool", GlobalAveragePoolOpBuilder), + REGISTER_OP_BUILDER("QLinearConv", QLinearConvOpBuilder), + REGISTER_OP_BUILDER("Flatten", FlattenOpBuilder), + REGISTER_OP_BUILDER("MatMul", MatMulOpBuilder), + REGISTER_OP_BUILDER("GlobalMaxPool", GlobalMaxPoolOpBuilder), + REGISTER_OP_BUILDER("AveragePool", AveragePoolOpBuilder), + REGISTER_OP_BUILDER("MaxPool", MaxPoolOpBuilder), + REGISTER_OP_BUILDER("Reshape", ReshapeOpBuilder), + REGISTER_OP_BUILDER("Concat", ConcatOpBuilder), + REGISTER_OP_BUILDER("Softmax", SoftmaxOpBuilder), + REGISTER_OP_BUILDER("Transpose", TransposeOpBuilder), + REGISTER_OP_BUILDER("BatchNormalization", BatchNormOpBuilder), + REGISTER_OP_BUILDER("Clip", ClipOpBuilder), + REGISTER_OP_BUILDER("ReduceMean", ReduceMeanOpBuilder), + REGISTER_OP_BUILDER("QuantizeLinear", QuantizeLinearOpBuilder), + REGISTER_OP_BUILDER("DequantizeLinear", DequantizeLinearOpBuilder), + REGISTER_OP_BUILDER("QLinearMatMul", QLinearMatMulOpBuilder), + REGISTER_OP_BUILDER("QLinearAdd", QLinearAddOpBuilder), + REGISTER_OP_BUILDER("QLinearMul", QLinearMulOpBuilder), + REGISTER_OP_BUILDER("QLinearConcat", QLinearConcatOpBuilder), + REGISTER_OP_BUILDER("Gather", GatherOpBuilder), + REGISTER_OP_BUILDER("Tile", TileOpBuilder), + REGISTER_OP_BUILDER("Squeeze", SqueezeOpBuilder), + REGISTER_OP_BUILDER("Unsqueeze", UnsqueezeOpBuilder), + REGISTER_OP_BUILDER("Resize", ResizeOpBuilder), + REGISTER_OP_BUILDER("Cast", CastOpBuilder), + +#undef REGISTER_OP_BUILDER +}; + +template +struct OpBuildConstructor { + T supported_builtins; + OpBuildConstructor( + const std::map reg) { + LOGS_DEFAULT(INFO) << "Initialize supported ops"; + for (const auto& kv : reg) { + supported_builtins.insert(std::make_pair(kv.first, kv.second())); + } + } +}; + +inline const OpBuildItemType& SupportedBuiltinOps() { + static OpBuildConstructor c(reg); + return c.supported_builtins; +} +} // namespace npu +} // namespace vsi +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/vsinpu/patches/AccuracyCorrection.patch b/onnxruntime/core/providers/vsinpu/patches/AccuracyCorrection.patch new file mode 100644 index 0000000000000..d44190101d9fa --- /dev/null +++ b/onnxruntime/core/providers/vsinpu/patches/AccuracyCorrection.patch @@ -0,0 +1,26 @@ +diff --git a/onnxruntime/test/providers/checkers.cc b/onnxruntime/test/providers/checkers.cc +index 47c18c478d..93b44501cd 100644 +--- a/onnxruntime/test/providers/checkers.cc ++++ b/onnxruntime/test/providers/checkers.cc +@@ -195,7 +195,7 @@ struct TensorCheck { + // For any other EPs, we still expect an exact match for the results + // TODO: Verify if DML can possibly have a ROUNDING_MODE parameter and conform to the other EPs #41968513 + if ((provider_type == kNnapiExecutionProvider || provider_type == kDmlExecutionProvider || +- provider_type == kXnnpackExecutionProvider) && ++ provider_type == kXnnpackExecutionProvider || provider_type == kVSINPUExecutionProvider) && + (has_abs_err || has_rel_err)) { + double threshold = has_abs_err ? *(params.absolute_error) + : 0.0; +diff --git a/onnxruntime/test/providers/cpu/nn/qlinearconv_op_test.cc b/onnxruntime/test/providers/cpu/nn/qlinearconv_op_test.cc +index 2bc0df5e36..7beb78c2ff 100644 +--- a/onnxruntime/test/providers/cpu/nn/qlinearconv_op_test.cc ++++ b/onnxruntime/test/providers/cpu/nn/qlinearconv_op_test.cc +@@ -498,7 +498,7 @@ class QLinearConvOpTester { + // NOTE, for now the tolerance will only apply if the NNAPI is actually used, + // if for any reason the execution falls back to CPU, we still expect an exact match + // See, 'void Check(...' in onnxruntime/test/providers/provider_test_utils.cc +-#if defined(USE_NNAPI) || defined(USE_DML) ++#if defined(USE_NNAPI) || defined(USE_DML) || defined(USE_VSINPU) + // TODO: Verify if DML can possibly have a ROUNDING_MODE parameter and conform to the other EPs #41968513 + abs_error = 1.0f; + #endif diff --git a/onnxruntime/core/providers/vsinpu/patches/local_testing_record_res.patch b/onnxruntime/core/providers/vsinpu/patches/local_testing_record_res.patch new file mode 100644 index 0000000000000..e118ee104912f --- /dev/null +++ b/onnxruntime/core/providers/vsinpu/patches/local_testing_record_res.patch @@ -0,0 +1,343 @@ +diff --git a/onnxruntime/test/onnx/dataitem_request.cc b/onnxruntime/test/onnx/dataitem_request.cc +index 1ee302d5d5..5c2dd5ab00 100644 +--- a/onnxruntime/test/onnx/dataitem_request.cc ++++ b/onnxruntime/test/onnx/dataitem_request.cc +@@ -135,6 +135,7 @@ std::pair DataTaskRequestContext::RunImpl() { + } + + EXECUTE_RESULT res = EXECUTE_RESULT::SUCCESS; ++ int32_t out_idx = 0; + for (auto& output : expected_output_values) { + const std::string& output_name = output.first; + OrtValue* expected_output_value = output.second; // Automatic cast +@@ -170,7 +171,7 @@ std::pair DataTaskRequestContext::RunImpl() { + } else { // Both expect and actual OrtValues are not None, proceed with data checking + ret = + CompareOrtValue(*actual_output_value, *expected_output_value, per_sample_tolerance, +- relative_per_sample_tolerance, post_procesing); ++ relative_per_sample_tolerance, post_procesing, out_idx); + } + } else { // Expected output is None, ensure that the received output OrtValue is None as well + if (actual_output_value->IsAllocated()) { +@@ -223,9 +224,10 @@ std::pair DataTaskRequestContext::RunImpl() { + if (compare_result != COMPARE_RESULT::SUCCESS && !ret.second.empty()) { + LOGS_DEFAULT(ERROR) << test_case_.GetTestCaseName() << ":output=" << output_name << ":" << ret.second; + } +- if (compare_result != COMPARE_RESULT::SUCCESS) { +- break; +- } ++ // if (compare_result != COMPARE_RESULT::SUCCESS) { ++ // break; ++ // } ++ out_idx ++; + } + return std::make_pair(res, spent_time_); + } +diff --git a/onnxruntime/test/providers/checkers.cc b/onnxruntime/test/providers/checkers.cc +index f1a7240ea3..436031dfa8 100644 +--- a/onnxruntime/test/providers/checkers.cc ++++ b/onnxruntime/test/providers/checkers.cc +@@ -154,6 +154,7 @@ struct TensorCheck { + } + + const bool has_abs_err = params.absolute_error.has_value(); ++ const int8_t default_abs_err = 1; + if (has_abs_err) { + double threshold = *(params.absolute_error); + +@@ -162,7 +163,8 @@ struct TensorCheck { + } + } else { + for (int i = 0; i < size; ++i) { +- EXPECT_EQ(cur_expected[i], cur_actual[i]) << "i:" << i; ++ // EXPECT_EQ(cur_expected[i], cur_actual[i]) << "i:" << i; ++ EXPECT_NEAR(cur_expected[i], cur_actual[i], default_abs_err) << "i:" << i; + } + } + } +diff --git a/onnxruntime/test/util/compare_ortvalue.cc b/onnxruntime/test/util/compare_ortvalue.cc +index 3d53d4a3a0..8129af1820 100644 +--- a/onnxruntime/test/util/compare_ortvalue.cc ++++ b/onnxruntime/test/util/compare_ortvalue.cc +@@ -138,11 +138,75 @@ std::pair CompareFloatResult(const Tensor& outvalue + return res; + } + ++template ++std::pair CompareFloatResult(const Tensor& outvalue, const Tensor& expected_value, ++ double per_sample_tolerance, ++ double relative_per_sample_tolerance, bool post_processing, int32_t out_idx) { ++ const size_t size1 = static_cast(expected_value.Shape().Size()); ++ const FLOAT_TYPE* expected_output = expected_value.Data(); ++ const FLOAT_TYPE* real_output = outvalue.Data(); ++ ++ std::string expected_name = "expected_res"+ std::to_string(out_idx)+ ".txt"; ++ std::string npures_name = "npu_res"+ std::to_string(out_idx)+ ".txt"; ++ std::ofstream expected_res(expected_name), npu_res(npures_name); ++ for(size_t i = 0 ; i < size1; i++){ ++ expected_res << expected_output[i] << std::endl; ++ npu_res << real_output[i] << std::endl; ++ } ++ expected_res.close(); ++ npu_res.close(); ++ ++ std::pair res = std::make_pair(COMPARE_RESULT::SUCCESS, ""); ++ double max_diff = 0; ++ size_t diff_count = 0; ++ for (size_t di = 0; di != size1; ++di) { ++ const double real_value = ++ post_processing ? std::max(0.0, std::min(255.0, real_output[di])) : real_output[di]; ++ const double diff = std::fabs(expected_output[di] - real_value); ++ const double tol = per_sample_tolerance + relative_per_sample_tolerance * std::fabs(expected_output[di]); ++ if (!IsResultCloselyMatch(real_value, expected_output[di], diff, tol)) { ++ res.first = COMPARE_RESULT::RESULT_DIFFERS; ++ // update error message if this is a larger diff ++ if (diff > max_diff || (std::isnan(diff) && !std::isnan(max_diff))) { ++ int64_t expected_int = 0; ++ int64_t real_int = 0; ++ memcpy(&expected_int, &expected_output[di], sizeof(FLOAT_TYPE)); ++ memcpy(&real_int, &real_output[di], sizeof(FLOAT_TYPE)); ++ ++ std::ostringstream oss; ++ oss << std::hex << "expected " << expected_output[di] << " (" << expected_int << "), got " << real_value << " (" ++ << real_int << ")" ++ << ", diff: " << diff << ", tol=" << tol << std::dec << " idx=" << di << "."; ++ res.second = oss.str(); ++ max_diff = diff; ++ } ++ ++diff_count; ++ } ++ } ++ ++ if (res.first == COMPARE_RESULT::SUCCESS) return res; ++ ++ std::ostringstream oss; ++ oss << res.second << " " << diff_count << " of " << size1 << " differ"; ++ res.second = oss.str(); ++ return res; ++} ++ ++ + template +-std::pair IsResultExactlyMatch(const Tensor& outvalue, const Tensor& expected_value) { ++std::pair IsResultExactlyMatch(const Tensor& outvalue, const Tensor& expected_value, int32_t out_idx) { + const size_t size1 = static_cast(expected_value.Shape().Size()); + const T* expected_output = expected_value.Data(); + const T* real_output = outvalue.Data(); ++ std::string expected_name = "expected_res"+ std::to_string(out_idx)+ ".txt"; ++ std::string npures_name = "npu_res"+ std::to_string(out_idx)+ ".txt"; ++ std::ofstream expected_res(expected_name), npu_res(npures_name); ++ for(size_t i = 0 ; i < size1; i++){ ++ expected_res << expected_output[i] << std::endl; ++ npu_res << real_output[i] << std::endl; ++ } ++ expected_res.close(); ++ npu_res.close(); + for (size_t di = 0; di != size1; ++di) { + if (expected_output[di] != real_output[di]) { + std::ostringstream oss; +@@ -201,7 +265,7 @@ std::pair CompareBFloat16Result(const Tensor& outva + + std::pair CompareTwoTensors(const Tensor& outvalue, const Tensor& expected_tensor, + double per_sample_tolerance, +- double relative_per_sample_tolerance, bool post_processing) { ++ double relative_per_sample_tolerance, bool post_processing, int32_t out_idx) { + if (expected_tensor.Shape() != outvalue.Shape()) { + std::ostringstream oss; + oss << "shape mismatch, expect " << expected_tensor.Shape().ToString() << " got " << outvalue.Shape().ToString(); +@@ -209,30 +273,30 @@ std::pair CompareTwoTensors(const Tensor& outvalue, + } + if (outvalue.IsDataType()) { + return CompareFloatResult(outvalue, expected_tensor, per_sample_tolerance, relative_per_sample_tolerance, +- post_processing); ++ post_processing, out_idx); + } else if (outvalue.IsDataType()) { + return CompareFloatResult(outvalue, expected_tensor, per_sample_tolerance, relative_per_sample_tolerance, +- post_processing); ++ post_processing, out_idx); + } else if (outvalue.IsDataTypeString()) { +- return IsResultExactlyMatch(outvalue, expected_tensor); ++ return IsResultExactlyMatch(outvalue, expected_tensor, out_idx); + } else if (outvalue.IsDataType()) { +- return IsResultExactlyMatch(outvalue, expected_tensor); ++ return IsResultExactlyMatch(outvalue, expected_tensor, out_idx); + } else if (outvalue.IsDataType()) { +- return IsResultExactlyMatch(outvalue, expected_tensor); ++ return IsResultExactlyMatch(outvalue, expected_tensor, out_idx); + } else if (outvalue.IsDataType()) { +- return IsResultExactlyMatch(outvalue, expected_tensor); ++ return IsResultExactlyMatch(outvalue, expected_tensor, out_idx); + } else if (outvalue.IsDataType()) { +- return IsResultExactlyMatch(outvalue, expected_tensor); ++ return IsResultExactlyMatch(outvalue, expected_tensor, out_idx); + } else if (outvalue.IsDataType()) { +- return IsResultExactlyMatch(outvalue, expected_tensor); ++ return IsResultExactlyMatch(outvalue, expected_tensor, out_idx); + } else if (outvalue.IsDataType()) { +- return IsResultExactlyMatch(outvalue, expected_tensor); ++ return IsResultExactlyMatch(outvalue, expected_tensor, out_idx); + } else if (outvalue.IsDataType()) { +- return IsResultExactlyMatch(outvalue, expected_tensor); ++ return IsResultExactlyMatch(outvalue, expected_tensor, out_idx); + } else if (outvalue.IsDataType()) { +- return IsResultExactlyMatch(outvalue, expected_tensor); ++ return IsResultExactlyMatch(outvalue, expected_tensor, out_idx); + } else if (outvalue.IsDataType()) { +- return IsResultExactlyMatch(outvalue, expected_tensor); ++ return IsResultExactlyMatch(outvalue, expected_tensor, out_idx); + } else if (outvalue.IsDataType()) { + return CompareFloat16Result(outvalue, expected_tensor, per_sample_tolerance, relative_per_sample_tolerance, + post_processing); +@@ -300,7 +364,7 @@ std::pair CompareSparseTensors(const SparseTensor& + " actual: ", actual.Format()); + + TEST_RETURN_IF_ERROR(CompareTwoTensors(actual.Values(), expected.Values(), +- per_sample_tolerance, relative_per_sample_tolerance, post_processing), ++ per_sample_tolerance, relative_per_sample_tolerance, post_processing, 0), + "While comparing sparse values"); + + if (actual.Format() == SparseFormat::kCoo) { +@@ -308,16 +372,16 @@ std::pair CompareSparseTensors(const SparseTensor& + auto expected_view = expected.AsCoo(); + + TEST_RETURN_IF_ERROR(CompareTwoTensors(actual_view.Indices(), expected_view.Indices(), +- per_sample_tolerance, relative_per_sample_tolerance, post_processing), ++ per_sample_tolerance, relative_per_sample_tolerance, post_processing, 0), + "Comparing COO indices"); + } else if (actual.Format() == SparseFormat::kCsrc) { + auto actual_view = actual.AsCsr(); + auto expected_view = expected.AsCsr(); + TEST_RETURN_IF_ERROR(CompareTwoTensors(actual_view.Inner(), expected_view.Inner(), +- per_sample_tolerance, relative_per_sample_tolerance, post_processing), ++ per_sample_tolerance, relative_per_sample_tolerance, post_processing, 0), + "Comparing Csr(c) inner indices"); + TEST_RETURN_IF_ERROR(CompareTwoTensors(actual_view.Outer(), expected_view.Outer(), +- per_sample_tolerance, relative_per_sample_tolerance, post_processing), ++ per_sample_tolerance, relative_per_sample_tolerance, post_processing, 0), + "Comparing Csr(c) outer indices"); + } + +@@ -385,7 +449,83 @@ std::pair CompareOrtValue(const OrtValue& o, const + return std::make_pair(COMPARE_RESULT::TYPE_MISMATCH, oss.str()); + } + return CompareTwoTensors(outvalue, expected_tensor, per_sample_tolerance, relative_per_sample_tolerance, +- post_processing); ++ post_processing, 0); ++ } else if (o.IsSparseTensor()) { ++#if !defined(DISABLE_SPARSE_TENSORS) ++ TEST_RETURN_IF_NOT(expected_mlvalue.IsSparseTensor(), COMPARE_RESULT::TYPE_MISMATCH, ++ "SparseTensor is not expected as output"); ++ TEST_RETURN_IF_ERROR(CompareSparseTensors(o.Get(), expected_mlvalue.Get(), ++ per_sample_tolerance, relative_per_sample_tolerance, ++ post_processing), ++ "while comaring sparse tensors"); ++#endif ++ return std::make_pair(COMPARE_RESULT::SUCCESS, ""); ++ } else if (o.IsTensorSequence()) { ++ auto& expected_tensor_seq = expected_mlvalue.Get(); ++ auto expected_tensor_count = expected_tensor_seq.Size(); ++ ++ auto& actual_tensor_seq = o.Get(); ++ auto actual_tensor_count = actual_tensor_seq.Size(); ++ ++ if (expected_tensor_count != actual_tensor_count) { ++ std::ostringstream oss; ++ oss << "expected tensor count in the sequence: " << expected_tensor_count << " got " ++ << actual_tensor_count; ++ return std::make_pair(COMPARE_RESULT::RESULT_DIFFERS, oss.str()); ++ } ++ ++ if (!expected_tensor_seq.IsSameDataType(actual_tensor_seq)) { ++ std::ostringstream oss; ++ oss << "expected tensor type in the sequence: " << expected_tensor_seq.DataType() << " got " ++ << actual_tensor_seq.DataType(); ++ return std::make_pair(COMPARE_RESULT::TYPE_MISMATCH, oss.str()); ++ } ++ ++ for (size_t i = 0; i < expected_tensor_count; ++i) { ++ auto res = CompareTwoTensors(actual_tensor_seq.Get(i), expected_tensor_seq.Get(i), per_sample_tolerance, relative_per_sample_tolerance, ++ post_processing,0); ++ if (res.first != COMPARE_RESULT::SUCCESS) { ++ return res; ++ } ++ } ++ ++ return std::make_pair(COMPARE_RESULT::SUCCESS, ""); ++ ++ } else { ++ // Maps ++#if !defined(DISABLE_ML_OPS) ++ if (o.Type() == DataTypeImpl::GetType()) { ++ return CompareSeqOfMapToFloat(o.Get(), expected_mlvalue.Get(), ++ per_sample_tolerance, relative_per_sample_tolerance, post_processing); ++ } ++ if (o.Type() == DataTypeImpl::GetType()) { ++ return CompareSeqOfMapToFloat(o.Get(), expected_mlvalue.Get(), ++ per_sample_tolerance, relative_per_sample_tolerance, post_processing); ++ } ++ return std::make_pair(COMPARE_RESULT::NOT_SUPPORT, ""); ++#else ++ return std::make_pair(COMPARE_RESULT::NOT_SUPPORT, "Map type is not supported in this build."); ++#endif ++ } ++} ++ ++std::pair CompareOrtValue(const OrtValue& o, const OrtValue& expected_mlvalue, ++ double per_sample_tolerance, ++ double relative_per_sample_tolerance, bool post_processing, int32_t out_idx) { ++ if (o.Type() != expected_mlvalue.Type()) { ++ return std::make_pair(COMPARE_RESULT::TYPE_MISMATCH, ""); ++ } ++ if (o.IsTensor()) { ++ const Tensor& outvalue = o.Get(); ++ const Tensor& expected_tensor = expected_mlvalue.Get(); ++ if (outvalue.DataType() != expected_tensor.DataType()) { ++ std::ostringstream oss; ++ oss << "expect " << ElementTypeToString(expected_tensor.DataType()) << " got " ++ << ElementTypeToString(outvalue.DataType()); ++ return std::make_pair(COMPARE_RESULT::TYPE_MISMATCH, oss.str()); ++ } ++ return CompareTwoTensors(outvalue, expected_tensor, per_sample_tolerance, relative_per_sample_tolerance, ++ post_processing, out_idx); + } else if (o.IsSparseTensor()) { + #if !defined(DISABLE_SPARSE_TENSORS) + TEST_RETURN_IF_NOT(expected_mlvalue.IsSparseTensor(), COMPARE_RESULT::TYPE_MISMATCH, +@@ -419,7 +559,7 @@ std::pair CompareOrtValue(const OrtValue& o, const + + for (size_t i = 0; i < expected_tensor_count; ++i) { + auto res = CompareTwoTensors(actual_tensor_seq.Get(i), expected_tensor_seq.Get(i), per_sample_tolerance, relative_per_sample_tolerance, +- post_processing); ++ post_processing, out_idx); + if (res.first != COMPARE_RESULT::SUCCESS) { + return res; + } +diff --git a/onnxruntime/test/util/include/compare_ortvalue.h b/onnxruntime/test/util/include/compare_ortvalue.h +index 24b74b9002..8269346528 100644 +--- a/onnxruntime/test/util/include/compare_ortvalue.h ++++ b/onnxruntime/test/util/include/compare_ortvalue.h +@@ -24,7 +24,9 @@ enum class COMPARE_RESULT { SUCCESS, + std::pair CompareOrtValue(const OrtValue& real, const OrtValue& expected, + double per_sample_tolerance, + double relative_per_sample_tolerance, bool post_processing); +- ++std::pair CompareOrtValue(const OrtValue& real, const OrtValue& expected, ++ double per_sample_tolerance, ++ double relative_per_sample_tolerance, bool post_processing, int32_t out_idx); + // verify if the 'value' matches the 'expected' ValueInfoProto. 'value' is a model output + std::pair VerifyValueInfo(const ONNX_NAMESPACE::ValueInfoProto& expected, + const OrtValue* value); +diff --git a/onnxruntime/test/util/include/test/compare_ortvalue.h b/onnxruntime/test/util/include/test/compare_ortvalue.h +index 545df706c9..170eb9dc4c 100644 +--- a/onnxruntime/test/util/include/test/compare_ortvalue.h ++++ b/onnxruntime/test/util/include/test/compare_ortvalue.h +@@ -28,7 +28,9 @@ enum class COMPARE_RESULT { + std::pair CompareOrtValue(const OrtValue& real, const OrtValue& expected, + double per_sample_tolerance, + double relative_per_sample_tolerance, bool post_processing); +- ++std::pair CompareOrtValue(const OrtValue& real, const OrtValue& expected, ++ double per_sample_tolerance, ++ double relative_per_sample_tolerance, bool post_processing, int32_t out_idx); + // Compare two OrtValue numerically equal or not. The difference with CompareOrtValue is that this function + // will only check the numerical values of the OrtValue, and ignore the type, shape, etc. + // diff --git a/onnxruntime/core/providers/vsinpu/patches/mlas_crosscompiling.patch b/onnxruntime/core/providers/vsinpu/patches/mlas_crosscompiling.patch new file mode 100644 index 0000000000000..a9d02765cf34d --- /dev/null +++ b/onnxruntime/core/providers/vsinpu/patches/mlas_crosscompiling.patch @@ -0,0 +1,34 @@ +diff --git a/cmake/onnxruntime_mlas.cmake b/cmake/onnxruntime_mlas.cmake +index e0ccc504d7..6c5aa6ea53 100644 +--- a/cmake/onnxruntime_mlas.cmake ++++ b/cmake/onnxruntime_mlas.cmake +@@ -335,7 +335,7 @@ else() + ${MLAS_SRC_DIR}/qgemm_kernel_udot.cpp + ${MLAS_SRC_DIR}/qgemm_kernel_sdot.cpp + ) +- if (NOT APPLE) ++ if (NOT APPLE AND NOT onnxruntime_USE_VSINPU) + set(mlas_platform_srcs + ${mlas_platform_srcs} + ${MLAS_SRC_DIR}/aarch64/HalfGemmKernelNeon.S +diff --git a/onnxruntime/core/mlas/inc/mlas.h b/onnxruntime/core/mlas/inc/mlas.h +index fd6b3df934..f81f1c42b6 100644 +--- a/onnxruntime/core/mlas/inc/mlas.h ++++ b/onnxruntime/core/mlas/inc/mlas.h +@@ -79,6 +79,7 @@ Abstract: + + #if (!defined(_MSC_VER)) || (_MSC_VER >= 1930) + #if defined(MLAS_TARGET_ARM64) || defined(MLAS_TARGET_ARM64EC) ++#if !defined(USE_VSINPU) + #if !defined(__APPLE__) + // Had to temporary disable fp16 under APPLE ARM64, as compiling + // the source files require a hardware specific compilation flag. +@@ -87,7 +88,8 @@ Abstract: + + #define MLAS_F16VEC_INTRINSICS_SUPPORTED + +-#endif // ++#endif ++#endif // + #endif // ARM64 + #endif // Visual Studio 16 or earlier does not support fp16 intrinsic diff --git a/onnxruntime/core/providers/vsinpu/patches/test_scripts/compare_cosine_sim.py b/onnxruntime/core/providers/vsinpu/patches/test_scripts/compare_cosine_sim.py new file mode 100644 index 0000000000000..e4e9b44fdc252 --- /dev/null +++ b/onnxruntime/core/providers/vsinpu/patches/test_scripts/compare_cosine_sim.py @@ -0,0 +1,29 @@ +import sys + +import numpy as np +from numpy.linalg import norm + + +def read_values(filename): + with open(filename) as file: + values = np.array([float(line.strip()) for line in file]) + return values + + +def cosine_similarity(vec1, vec2): + return np.dot(vec1, vec2) / (norm(vec1) * norm(vec2)) + + +if __name__ == "__main__": + if len(sys.argv) != 3: + print("Usage: python cosine_similarity.py ") + sys.exit(1) + + file1 = sys.argv[1] + file2 = sys.argv[2] + + vec1 = read_values(file1) + vec2 = read_values(file2) + + similarity = cosine_similarity(vec1, vec2) + print(f"Cosine Similarity: {similarity}") diff --git a/onnxruntime/core/providers/vsinpu/patches/test_scripts/compare_topn.py b/onnxruntime/core/providers/vsinpu/patches/test_scripts/compare_topn.py new file mode 100644 index 0000000000000..cde75b7f18c1e --- /dev/null +++ b/onnxruntime/core/providers/vsinpu/patches/test_scripts/compare_topn.py @@ -0,0 +1,34 @@ +import sys + + +def read_values(filename): + with open(filename) as file: + values = [(float(line.strip()), i + 1) for i, line in enumerate(file)] + return values + + +def top_n(values, N): + return sorted(values, key=lambda x: x[0], reverse=True)[:N] + + +def compare_files(cpu_file, npu_file, N): + cpu_values = read_values(cpu_file) + npu_values = read_values(npu_file) + + cpu_topn = top_n(cpu_values, N) + npu_topn = top_n(npu_values, N) + + print(f"Top-{N} values in {cpu_file}: {cpu_topn}") + print(f"Top-{N} values in {npu_file}: {npu_topn}") + + +if __name__ == "__main__": + if len(sys.argv) != 4: + print("Usage: python compare_topn.py ") + sys.exit(1) + + N = int(sys.argv[1]) + cpu_file = sys.argv[2] + npu_file = sys.argv[3] + + compare_files(cpu_file, npu_file, N) diff --git a/onnxruntime/core/providers/vsinpu/patches/test_scripts/result_compare.sh b/onnxruntime/core/providers/vsinpu/patches/test_scripts/result_compare.sh new file mode 100644 index 0000000000000..c27af51c26799 --- /dev/null +++ b/onnxruntime/core/providers/vsinpu/patches/test_scripts/result_compare.sh @@ -0,0 +1,23 @@ +#!/bin/bash +res_file_dir=$1 +output_num=$2 + +# specifying N value +N=5 + +for i in $(seq 0 $((output_num-1))); +do + # 构建文件名 + golden_file="${res_file_dir}/expected_res${i}.txt" + npu_file="${res_file_dir}/npu_res${i}.txt" + + echo "Comparing Top-${N} for the output_${i}" + python3 compare_topn.py $N $golden_file $npu_file + + echo "--------------------------------" + + echo "Comparing Cosine Similarity for output_${i}:" + python3 compare_cosine_sim.py $golden_file $npu_file + + echo "" +done diff --git a/onnxruntime/core/providers/vsinpu/symbols.txt b/onnxruntime/core/providers/vsinpu/symbols.txt new file mode 100644 index 0000000000000..d69c92692f5fe --- /dev/null +++ b/onnxruntime/core/providers/vsinpu/symbols.txt @@ -0,0 +1 @@ +OrtSessionOptionsAppendExecutionProvider_VSINPU diff --git a/onnxruntime/core/providers/vsinpu/vsinpu_ep_graph.cc b/onnxruntime/core/providers/vsinpu/vsinpu_ep_graph.cc new file mode 100644 index 0000000000000..e51b0713ea41d --- /dev/null +++ b/onnxruntime/core/providers/vsinpu/vsinpu_ep_graph.cc @@ -0,0 +1,296 @@ + +/**************************************************************************** + * + * Copyright (c) 2023 Vivante Corporation + * + * Permission is hereby granted, free of charge, to any person obtaining a + * copy of this software and associated documentation files (the "Software"), + * to deal in the Software without restriction, including without limitation + * the rights to use, copy, modify, merge, publish, distribute, sublicense, + * and/or sell copies of the Software, and to permit persons to whom the + * Software is furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING + * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER + * DEALINGS IN THE SOFTWARE. + * + *****************************************************************************/ +#include +#include "core/providers/vsinpu/vsinpu_ep_graph.h" +#include "core/providers/vsinpu/builders/op_builder_factory.h" +#include "core/providers/vsinpu/vsinpu_util.h" +#include "core/framework/node_unit.h" +#include "core/optimizer/qdq_transformer/selectors_actions/qdq_selectors.h" +#include "core/optimizer/qdq_transformer/selectors_actions/shared/utils.h" + +namespace onnxruntime { + +namespace vsi { +namespace npu { +GraphEP::GraphEP(const onnxruntime::GraphViewer& graph_viewer) : graph_viewer_(graph_viewer) { + Prepare(); + context_ = tim::vx::Context::Create(); + graph_ = context_->CreateGraph(); + compiled_ = false; +} + +bool GraphEP::Prepare() { + std::tie(node_unit_holder_, node_unit_map_) = QDQ::GetAllNodeUnits(graph_viewer_); + for (const auto& node_unit : node_unit_holder_) { + auto quant_op_type = util::GetQuantizedOpType(*node_unit); + + // Not a qlinear op or qdq node group + if (quant_op_type == util::QuantizedOpType::Unknown) + continue; + + const auto add_quantized_input = + [&all_quantized_op_inputs = all_quantized_op_inputs_](const NodeUnit& node_unit, size_t input_idx) { + const auto& input_name = node_unit.Inputs()[input_idx].node_arg.Name(); + all_quantized_op_inputs[input_name].push_back(&node_unit); + }; + + // All quantized ops EXCEPT QuantizeLinear has quantized input + if (quant_op_type != util::QuantizedOpType::QuantizeLinear) { + add_quantized_input(*node_unit, 0); + } + + if (util::IsQuantizedBinaryOp(quant_op_type)) { + add_quantized_input(*node_unit, 1); + if (util::IsQuantizedConv(quant_op_type) && node_unit->Inputs().size() == 3) { + add_quantized_input(*node_unit, 2); + } + } + } // All quantized inputs is recorded + return true; +} + +bool GraphEP::SupportedOp(const onnxruntime::GraphViewer& graph_viewer, + const NodeUnit& node_unit) { + const auto& supported_builtins = vsi::npu::SupportedBuiltinOps(); + const auto& target_node = node_unit.GetNode(); + const auto& it = supported_builtins.find(target_node.OpType()); + if (supported_builtins.end() != it) { + return it->second->IsSupported(graph_viewer, node_unit); + } + LOGS_DEFAULT(WARNING) << "Fallback unsupported op (node_unit) " << node_unit.OpType() + << " to cpu."; + return false; +} + +bool GraphEP::IsNodeSupportedInGroup(const NodeUnit& node_unit, const GraphViewer& graph_viewer) { + return SupportedOp(graph_viewer, node_unit); +} + +const NodeUnit& GraphEP::GetNodeUnit(const Node* node) const { + const auto node_unit_it = node_unit_map_.find(node); + ORT_ENFORCE(node_unit_it != node_unit_map_.end(), "Node does not have corresponding NodeUnit."); + return *node_unit_it->second; +} + +void GraphEP::UpdateTensorMap(const std::string& name, const std::shared_ptr& dst_tensor) { + auto it = tensors_.find(name); + if (it != tensors_.end()) { + it->second = dst_tensor; + } + for (auto& IO : graph_inputs_) { + if (IO->name == name) { + IO->tensor = dst_tensor; + break; + } + } + for (auto& IO : graph_outputs_) { + if (IO->name == name) { + IO->tensor = dst_tensor; + break; + } + } +} + +std::shared_ptr GraphEP::ConstructNodeIO(const std::shared_ptr& op, std::vector input_arg, std::vector output_arg) { + auto info = std::make_shared(); + info->op_ = op; + std::vector input_names, output_names; + if (input_arg.empty()) { + info->input_names_ = std::vector(); + } else { + input_names.reserve(input_arg.size()); + std::transform(input_arg.begin(), input_arg.end(), std::back_inserter(input_names), + [](const NodeArg* node) -> std::string { + return node->Name(); + }); + info->input_names_ = input_names; + } + if (output_arg.empty()) { + info->output_names_ = std::vector(); + } else { + output_names.reserve(output_arg.size()); + std::transform(output_arg.begin(), output_arg.end(), std::back_inserter(output_names), + [](const NodeArg* node) -> std::string { + return node->Name(); + }); + info->output_names_ = output_names; + } + + return info; +} + +bool GraphEP::BindTensors(const std::shared_ptr& nodeio_info) { + auto op = nodeio_info->op_; + auto input_names = nodeio_info->input_names_; + auto output_names = nodeio_info->output_names_; + if (!input_names.empty()) { + for (auto& name : input_names) { + if (tensors_.find(name) == tensors_.end() || tensors_[name] == nullptr) { + LOGS_DEFAULT(ERROR) << "Input tensor not defined or not found!"; + return false; + } + (*op).BindInput(tensors_[name]); + } + } + if (!output_names.empty()) { + for (auto& name : output_names) { + if (tensors_.find(name) == tensors_.end() || tensors_[name] == nullptr) { + LOGS_DEFAULT(ERROR) << "Output tensor not defined or not found!"; + return false; + } + (*op).BindOutput(tensors_[name]); + } + } + return true; +} + +std::shared_ptr GraphEP::MapTIMVXTensor( + std::shared_ptr& graph, const NodeUnitIODef nudef, + const NodeUnit& node_unit, + const GraphViewer* graph_viewer, tim::vx::TensorAttribute attribute) { + const auto& arg = nudef.node_arg; + + if (tensors_.end() != tensors_.find(nudef.node_arg.Name())) { + // if (!quant_param.has_value() || quant_param.has_value() && tensors_[arg.Name()]->GetSpec().GetQuantization().Type() != tim::vx::QuantType::NONE) + return tensors_.find(arg.Name())->second; + } + auto shape = vsi::npu::util::OnnxShapeToTIMVXShape(vsi::npu::util::GetTensorShape(arg)); + std::reverse(shape.begin(), shape.end()); + tim::vx::DataType dt = vsi::npu::util::OnnxDtypeToTIMVXDtype(arg.Type()); + tim::vx::TensorSpec spec = tim::vx::TensorSpec(dt, shape, attribute); + + // Tensors have same name may not have same status of quant_param existence, such as QLinearConv->MaxPool->QLinearConv + // Maxpool output tensor is not set quantization at first pass + bool is_qtensor = nudef.quant_param.has_value() || Contains(all_quantized_op_inputs_, arg.Name()); + if (is_qtensor) { + float scale = 0.0f; + int32_t zp = 0; + std::optional> scales; + std::optional> zps; + if (nudef.quant_param.has_value()) { + util::GetQuantizationScaleAndZeroPoint(graph_viewer_.GetAllInitializedTensors(), + nudef, node_unit.ModelPath(), + scale, zp, scales, zps); + } else { + auto target_nodeunit = all_quantized_op_inputs_[arg.Name()][0]; + auto qinput = all_quantized_op_inputs_[arg.Name()][0]->Inputs(); + auto it = std::find_if(qinput.begin(), qinput.end(), [&arg](const NodeUnitIODef& nud) { return nud.node_arg.Name() == arg.Name(); }); + bool is_conv_bias = std::distance(qinput.begin(), it) == 2; + if (!is_conv_bias || it->quant_param.has_value()) { + util::GetQuantizationScaleAndZeroPoint(graph_viewer_.GetAllInitializedTensors(), + *it, target_nodeunit->ModelPath(), + scale, zp, scales, zps); + } else if (!it->quant_param.has_value()) { + float in_scale, w_scale; + int32_t in_zp, w_zp; + std::optional> in_scales, w_scales; + std::optional> in_zps, w_zps; + + // onnx defines conv bias with non quantization, but it must be quantized in VSINPU support + // The bias scale is set as input_scale * weight_scale if per layer quantized, input_scale* weight_scale[i] if per channel quantized + util::GetQuantizationScaleAndZeroPoint(graph_viewer_.GetAllInitializedTensors(), + qinput[0], target_nodeunit->ModelPath(), + in_scale, in_zp, in_scales, in_zps); + util::GetQuantizationScaleAndZeroPoint(graph_viewer_.GetAllInitializedTensors(), + qinput[1], target_nodeunit->ModelPath(), + w_scale, w_zp, w_scales, w_zps); + scale = in_scale * w_scale; + zp = 0; + if (w_scales) { + std::vector temp; + for (size_t i = 0; i < w_scales->size(); i++) { + temp.push_back(w_scales.value()[i] * in_scale); + } + scales = temp; + } + } + } + tim::vx::Quantization quant; + // per tensor quantization + if (!scales.has_value()) { + quant.SetType(tim::vx::QuantType::ASYMMETRIC); + quant.SetScales({scale}); + quant.SetZeroPoints({zp}); + } else { // per channel quantization + if (zps.has_value()) { + bool has_nonzero = std::find_if(zps->begin(), zps->end(), [](int elem) { return elem != 0; }) != zps->end(); + if (has_nonzero && *arg.Type() == "tensor(uint8)") { + quant.SetType(tim::vx::QuantType::ASYMMETRIC_PER_CHANNEL); + } else { + quant.SetType(tim::vx::QuantType::SYMMETRIC_PER_CHANNEL); + } + quant.SetZeroPoints(zps.value()); + } else { + if (*arg.Type() == "tensor(int32)" || zp == 0) { + // set bias quant type + quant.SetType(tim::vx::QuantType::SYMMETRIC_PER_CHANNEL); + } else { + quant.SetType(tim::vx::QuantType::ASYMMETRIC_PER_CHANNEL); + } + quant.SetZeroPoints({zp}); + } + quant.SetScales(scales.value()); + quant.SetChannelDim(shape.size() - 1); + } + spec.SetQuantization(quant); + } + + std::shared_ptr tensor; + if (attribute == + tim::vx::TensorAttribute::CONSTANT) { // create const tensor + const ONNX_NAMESPACE::TensorProto* tensor_proto = + graph_viewer_.GetConstantInitializer(arg.Name(), true); + std::shared_ptr unpackedTensor = + vsi::npu::util::UnpackTensor(&arg, *tensor_proto); + + const void* valueAddr = + reinterpret_cast(unpackedTensor.get()); + tensor = graph->CreateTensor(spec, valueAddr); + + } else { + tensor = graph->CreateTensor(spec); + } + for (auto& input : graph_inputs_) { + if (input->name == arg.Name()) { + input->tensor = tensor; + input->shape = vsi::npu::util::GetTensorShape(arg); + break; + } + } + for (auto& output : graph_outputs_) { + if (output->name == arg.Name()) { + output->tensor = tensor; + output->shape = utils::GetTensorShapeFromTensorShapeProto(*arg.Shape()); + break; + } + } + tensors_.insert({arg.Name(), tensor}); + return tensor; +} + +} // namespace npu + +} // namespace vsi +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/vsinpu/vsinpu_ep_graph.h b/onnxruntime/core/providers/vsinpu/vsinpu_ep_graph.h new file mode 100644 index 0000000000000..bd0f377b820b0 --- /dev/null +++ b/onnxruntime/core/providers/vsinpu/vsinpu_ep_graph.h @@ -0,0 +1,116 @@ +/**************************************************************************** + * + * Copyright (c) 2023 Vivante Corporation + * + * Permission is hereby granted, free of charge, to any person obtaining a + * copy of this software and associated documentation files (the "Software"), + * to deal in the Software without restriction, including without limitation + * the rights to use, copy, modify, merge, publish, distribute, sublicense, + * and/or sell copies of the Software, and to permit persons to whom the + * Software is furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING + * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER + * DEALINGS IN THE SOFTWARE. + * + *****************************************************************************/ + +#pragma once +#include +#include +#include +#include +#include +#include "builders/op_builder.h" +#include "tim/vx/context.h" +#include "tim/vx/graph.h" +#include "tim/vx/tensor.h" + +namespace onnxruntime { +namespace vsi { +namespace npu { +struct GraphIOInfo { + std::string name; + bool is_initializer; + std::shared_ptr tensor; + TensorShape shape; +}; + +struct NodeIOInfo { + std::shared_ptr op_; + std::vector input_names_; + std::vector output_names_; +}; + +class GraphEP { + public: + explicit GraphEP(const GraphViewer& graph_viewer); + ~GraphEP() {} + + bool Prepare(); + + static bool SupportedOp(const onnxruntime::GraphViewer& graph_viewer, + const NodeUnit& node_unit); + + // If a node is supported by VSINPU in a partition node group + // `node_outputs_in_group` is the set of the output names of the nodes added to this group so far + static bool IsNodeSupportedInGroup(const NodeUnit& node_unit, const GraphViewer& graph_viewer); + + const NodeUnit& GetNodeUnit(const Node* node) const; + + bool& GetCompiled() { return compiled_; } + std::shared_ptr& GetGraph() { return graph_; } + std::vector>& GetOps() { return ops_; } + std::map>& GetTensors() { + return tensors_; + } + + std::vector>& GetGraphInputs() { + return graph_inputs_; + } + + std::vector>& GetGraphOutputs() { + return graph_outputs_; + } + + void UpdateTensorMap(const std::string& name, const std::shared_ptr& dst_tensor); + + std::shared_ptr ConstructNodeIO(const std::shared_ptr& op, std::vector input_arg, std::vector output_arg); + + bool BindTensors(const std::shared_ptr& nodeio_info); + + std::shared_ptr MapTIMVXTensor( + std::shared_ptr& graph, const NodeUnitIODef nudef, + const NodeUnit& nodeunit, + const GraphViewer* graph_viewer, tim::vx::TensorAttribute attribute); + + private: + std::shared_ptr context_; + std::shared_ptr graph_; + std::map> tensors_; + std::vector> ops_; + std::vector> graph_inputs_; + std::vector> graph_outputs_; + + // Contains all quantized operators' input and the NodeUnit(s) using the input + // In the form of {input_name, [NodeUnit(s) using the input]} + std::unordered_map> all_quantized_op_inputs_; + const GraphViewer& graph_viewer_; + + // Holder for the NodeUnits in the graph, this will guarantee the NodeUnits is + // valid throughout the lifetime of the ModelBuilder + std::vector> node_unit_holder_; + std::unordered_map node_unit_map_; + bool compiled_; +}; +} // namespace npu + +} // namespace vsi +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/vsinpu/vsinpu_execution_provider.cc b/onnxruntime/core/providers/vsinpu/vsinpu_execution_provider.cc new file mode 100644 index 0000000000000..7444dcfec09a2 --- /dev/null +++ b/onnxruntime/core/providers/vsinpu/vsinpu_execution_provider.cc @@ -0,0 +1,277 @@ +/**************************************************************************** + * + * Copyright (c) 2023 Vivante Corporation + * + * Permission is hereby granted, free of charge, to any person obtaining a + * copy of this software and associated documentation files (the "Software"), + * to deal in the Software without restriction, including without limitation + * the rights to use, copy, modify, merge, publish, distribute, sublicense, + * and/or sell copies of the Software, and to permit persons to whom the + * Software is furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING + * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER + * DEALINGS IN THE SOFTWARE. + * + *****************************************************************************/ +#include +#include +#include +#include "core/framework/compute_capability.h" +#include "core/providers/vsinpu/vsinpu_execution_provider.h" +#include "core/providers/vsinpu/vsinpu_ep_graph.h" +#include "core/providers/vsinpu/builders/op_builder.h" +#include "core/providers/vsinpu/builders/op_builder_factory.h" +#include "core/providers/vsinpu/vsinpu_util.h" +#include "core/framework/kernel_registry.h" +#include "core/framework/node_unit.h" +#include "core/optimizer/qdq_transformer/selectors_actions/qdq_selectors.h" +#include "core/optimizer/qdq_transformer/selectors_actions/shared/utils.h" +#include "core/providers/partitioning_utils.h" + +namespace onnxruntime { +VSINPUExecutionProvider::VSINPUExecutionProvider(const VSINPUExecutionProviderInfo& info) + : IExecutionProvider{onnxruntime::kVSINPUExecutionProvider}, + device_id_(info.device_id) { + AllocatorCreationInfo default_memory_info{ + [](int) { + return std::make_unique( + OrtMemoryInfo("VSINPU", OrtAllocatorType::OrtDeviceAllocator)); + }}; + + CreateAllocator(default_memory_info); + + AllocatorCreationInfo cpu_memory_info{ + [](int) { + return std::make_unique( + OrtMemoryInfo("VSINPU", OrtAllocatorType::OrtDeviceAllocator, OrtDevice(), 0, OrtMemTypeCPUOutput)); + }}; + + CreateAllocator(cpu_memory_info); +} + +VSINPUExecutionProvider::~VSINPUExecutionProvider() {} + +std::vector> +VSINPUExecutionProvider::GetCapability(const onnxruntime::GraphViewer& graph_viewer, + const IKernelLookup& /*kernel_lookup*/) const { + std::vector> result; + + if (graph_viewer.IsSubgraph()) { + return result; + } + + for (const auto& tensor : graph_viewer.GetAllInitializedTensors()) { + if (tensor.second->has_data_location()) { + LOGS_DEFAULT(VERBOSE) << "location:" << tensor.second->data_location(); + if (tensor.second->data_location() == + ONNX_NAMESPACE::TensorProto_DataLocation_EXTERNAL) { + LOGS_DEFAULT(WARNING) << "VSINPU: Initializers with external data location are not " + "currently supported"; + return result; + } + } + } + // Get all the NodeUnits in the graph_viewer + std::vector> node_unit_holder; + std::unordered_map node_unit_map; + std::tie(node_unit_holder, node_unit_map) = QDQ::GetAllNodeUnits(graph_viewer); + + // This holds the result of whether a NodeUnit is supported or not, + // to prevent nodes in a NodeUnit to be checked for multiple times + std::unordered_map node_unit_supported_result; + node_unit_supported_result.reserve(node_unit_holder.size()); + std::unordered_set node_outputs_in_current_group{}; + + const auto is_node_supported = [&](const Node& node) -> bool { + const NodeUnit* node_unit = node_unit_map.at(&node); + bool supported = false; + + // If we have visited one of the nodes in the node_unit, use the result directly + const auto it = node_unit_supported_result.find(node_unit); + if (it != node_unit_supported_result.cend()) { + supported = it->second; + } else { + // We only check the target node of the node unit + supported = vsi::npu::GraphEP::IsNodeSupportedInGroup(*node_unit, graph_viewer); + node_unit_supported_result[node_unit] = supported; + } + + LOGS_DEFAULT(VERBOSE) << "Node supported: [" << supported + << "] Operator type: [" << node.OpType() + << "] index: [" << node.Index() + << "] name: [" << node.Name() + << "] as part of the NodeUnit type: [" << node_unit->OpType() + << "] index: [" << node_unit->Index() + << "] name: [" << node_unit->Name() + << "]"; + + if (supported) { + // We want to save all the output names of nodes in the current group for easy query + for (const auto* output : node.OutputDefs()) { + node_outputs_in_current_group.insert(output->Name()); + } + } + return supported; + }; + + const auto on_group_closed = [&](const std::vector& group) -> bool { + // reset per-partition node group tracking + node_outputs_in_current_group.clear(); + return true; + }; + + const auto gen_metadef_name = [&]() { + static size_t group_counter = 0; + return "VSINPU_" + std::to_string(++group_counter); + }; + result = utils::CreateSupportedPartitions(graph_viewer, is_node_supported, on_group_closed, + gen_metadef_name, "VSINPU", kVSINPUExecutionProvider, &node_unit_map); + std::for_each(result.begin(), result.end(), [&graph_viewer](auto& capability) { + if (capability && capability->sub_graph && capability->sub_graph->GetMetaDef()) { + const auto* meta_def = capability->sub_graph->GetMetaDef(); + bool has_any_non_constant_inputs = std::any_of(meta_def->inputs.begin(), meta_def->inputs.end(), [&graph_viewer](const auto& input) { + return !graph_viewer.IsConstantInitializer(input, true); + }); + + // ALL inputs are constant + if (!has_any_non_constant_inputs) { + capability.reset(); + } + } + }); + + const auto num_of_partitions = result.size(); + const auto num_of_supported_nodes = std::accumulate( + result.begin(), result.end(), size_t{0}, + [](const auto& acc, const auto& partition) -> size_t { + return acc + (partition && partition->sub_graph ? partition->sub_graph->nodes.size() : 0); + }); + + const auto summary_msg = MakeString( + "VSINPUExecutionProvider::GetCapability,", + " number of partitions supported by VSINPU: ", num_of_partitions, + "; number of nodes in the graph: ", graph_viewer.NumberOfNodes(), + "; number of nodes supported by VSINPU: ", num_of_supported_nodes); + + // If the graph is partitioned in multiple subgraphs, and this may impact performance, + // we want to give users a summary message at warning level. + if (num_of_partitions > 1) { + LOGS_DEFAULT(WARNING) << summary_msg; + } else { + LOGS_DEFAULT(INFO) << summary_msg; + } + + return result; +} + +Status ComputeStateFunc(vsi::npu::GraphEP* graph_ep, + OrtKernelContext* context) { + Ort::KernelContext ctx(context); + size_t num_in = ctx.GetInputCount(); + const size_t num_inputs = graph_ep->GetGraphInputs().size(); + + for (size_t i = 0, j = 0; i < num_inputs; i++) { + if (!graph_ep->GetGraphInputs()[i]->is_initializer) { + const auto onnx_input_tensor = ctx.GetInput(i); + const auto tensor_info = onnx_input_tensor.GetTensorTypeAndShapeInfo(); + + auto origin_tensor = graph_ep->GetGraphInputs()[i]->tensor; + origin_tensor->CopyDataToTensor(onnx_input_tensor.GetTensorRawData(), vsi::npu::util::GetTensorBytes(tensor_info)); + j++; + } + } + + if (!graph_ep->GetGraph()->Run()) { + LOGS_DEFAULT(ERROR) << "Failed to run graph."; + } + for (size_t i = 0; i < ctx.GetOutputCount(); i++) { + auto timvx_tensor = graph_ep->GetGraphOutputs()[i]->tensor; + auto out_shape = graph_ep->GetGraphOutputs()[i]->shape.GetDims(); + auto onnx_output_tensor = + ctx.GetOutput(i, out_shape.data(), out_shape.size()); + timvx_tensor->CopyDataFromTensor(const_cast(onnx_output_tensor.GetTensorRawData())); + } + + return Status::OK(); +} + +Status VSINPUExecutionProvider::Compile(const std::vector& fused_nodes_and_graphs, + std::vector& node_compute_funcs) { + for (const auto& fused_node_graph : fused_nodes_and_graphs) { + const GraphViewer& graph_viewer = fused_node_graph.filtered_graph; + std::shared_ptr graph_ep = std::make_shared(graph_viewer); + + for (auto tensor : graph_viewer.GetInputsIncludingInitializers()) { + LOGS_DEFAULT(VERBOSE) << "subgraph input init:" << vsi::npu::util::PrintNode(*tensor) << "#" + << graph_viewer.IsInitializedTensor(tensor->Name()); + auto input = std::make_shared(); + input->name = tensor->Name(); + input->is_initializer = graph_viewer.IsConstantInitializer(tensor->Name(), true); + graph_ep->GetGraphInputs().push_back(input); + } + for (auto tensor : graph_viewer.GetOutputs()) { + LOGS_DEFAULT(VERBOSE) << "subgraph output:" << vsi::npu::util::PrintNode(*tensor); + auto output = std::make_shared(); + output->name = tensor->Name(); + output->is_initializer = false; + graph_ep->GetGraphOutputs().push_back(output); + } + + auto node_indices = graph_viewer.GetNodesInTopologicalOrder(); + for (const auto& node_index : node_indices) { + const auto node = graph_viewer.GetNode(node_index); + const NodeUnit& node_unit = graph_ep->GetNodeUnit(node); + + // Only add op when we hit the target node + if (node != &node_unit.GetNode()) { + continue; + } + LOGS_DEFAULT(VERBOSE) << "Adding node: [" << node->OpType() << "]"; + vsi::npu::SupportedBuiltinOps().at(node->OpType())->BuildOp(graph_ep.get(), graph_viewer, node_unit); + } + + LOGS_DEFAULT(INFO) << "Verifying graph"; + graph_ep->GetCompiled() = graph_ep->GetGraph()->Compile(); + if (!graph_ep->GetCompiled()) { + LOGS_DEFAULT(ERROR) << "Failed to verify graph."; + } else { + LOGS_DEFAULT(INFO) << "Graph has been verified successfully."; + } + + NodeComputeInfo compute_info; + compute_info.create_state_func = [graph_ep](ComputeContext* /*context*/, + FunctionState* state) { + *state = graph_ep.get(); + return 0; + }; + + compute_info.compute_func = + [graph_ep, this](FunctionState /*state*/, const OrtApi* /* api */, + OrtKernelContext* context) { + std::lock_guard lock(this->GetMutex()); + Status res = ComputeStateFunc(graph_ep.get(), context); + return res; + }; + + compute_info.release_state_func = [](FunctionState /*state*/) {}; + + node_compute_funcs.push_back(compute_info); + } + + return Status::OK(); +} + +std::shared_ptr VSINPUExecutionProvider::GetKernelRegistry() const { + static std::shared_ptr kernel_registry = std::make_shared(); + return kernel_registry; +} + +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/vsinpu/vsinpu_execution_provider.h b/onnxruntime/core/providers/vsinpu/vsinpu_execution_provider.h new file mode 100644 index 0000000000000..44318c332fdd0 --- /dev/null +++ b/onnxruntime/core/providers/vsinpu/vsinpu_execution_provider.h @@ -0,0 +1,53 @@ +/**************************************************************************** + * + * Copyright (c) 2023 Vivante Corporation + * + * Permission is hereby granted, free of charge, to any person obtaining a + * copy of this software and associated documentation files (the "Software"), + * to deal in the Software without restriction, including without limitation + * the rights to use, copy, modify, merge, publish, distribute, sublicense, + * and/or sell copies of the Software, and to permit persons to whom the + * Software is furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING + * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER + * DEALINGS IN THE SOFTWARE. + * + *****************************************************************************/ +#pragma once +#include +#include +#include "core/framework/execution_provider.h" +#include "core/session/abi_session_options_impl.h" + +namespace onnxruntime { +struct VSINPUExecutionProviderInfo { + int device_id{0}; +}; + +class VSINPUExecutionProvider : public IExecutionProvider { + public: + explicit VSINPUExecutionProvider(const VSINPUExecutionProviderInfo& info); + virtual ~VSINPUExecutionProvider(); + + std::vector> GetCapability( + const onnxruntime::GraphViewer& graph_viewer, + const IKernelLookup& kernel_lookup) const override; + std::shared_ptr GetKernelRegistry() const override; + Status Compile(const std::vector& fused_nodes_and_graphs, + std::vector& node_compute_funcs) override; + OrtMutex& GetMutex() { return mutex_; } + + private: + int device_id_; + OrtMutex mutex_; +}; + +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/vsinpu/vsinpu_provider_factory.cc b/onnxruntime/core/providers/vsinpu/vsinpu_provider_factory.cc new file mode 100644 index 0000000000000..5f2f961d95c09 --- /dev/null +++ b/onnxruntime/core/providers/vsinpu/vsinpu_provider_factory.cc @@ -0,0 +1,59 @@ +/**************************************************************************** + * + * Copyright (c) 2023 Vivante Corporation + * + * Permission is hereby granted, free of charge, to any person obtaining a + * copy of this software and associated documentation files (the "Software"), + * to deal in the Software without restriction, including without limitation + * the rights to use, copy, modify, merge, publish, distribute, sublicense, + * and/or sell copies of the Software, and to permit persons to whom the + * Software is furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING + * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER + * DEALINGS IN THE SOFTWARE. + * + *****************************************************************************/ +#include "core/framework/compute_capability.h" +#include "core/providers/vsinpu/vsinpu_provider_factory.h" +#include "core/providers/vsinpu/vsinpu_provider_factory_creator.h" +#include "core/providers/vsinpu/vsinpu_execution_provider.h" + +namespace onnxruntime { + +struct VSINPUProviderFactory : IExecutionProviderFactory { + VSINPUProviderFactory() {} + ~VSINPUProviderFactory() override {} + + std::unique_ptr CreateProvider() override; +}; + +std::unique_ptr VSINPUProviderFactory::CreateProvider() { + onnxruntime::VSINPUExecutionProviderInfo info; + return std::make_unique(info); +} + +std::shared_ptr CreateExecutionProviderFactory_VSINPU() { + return std::make_shared(); +} + +std::shared_ptr +VSINPUProviderFactoryCreator::Create() { + return std::make_shared(); +} + +} // namespace onnxruntime + +ORT_API_STATUS_IMPL(OrtSessionOptionsAppendExecutionProvider_VSINPU, + _In_ OrtSessionOptions* options) { + options->provider_factories.push_back( + onnxruntime::VSINPUProviderFactoryCreator::Create()); + return nullptr; +} diff --git a/onnxruntime/core/providers/vsinpu/vsinpu_provider_factory_creator.h b/onnxruntime/core/providers/vsinpu/vsinpu_provider_factory_creator.h new file mode 100644 index 0000000000000..e69185c0df816 --- /dev/null +++ b/onnxruntime/core/providers/vsinpu/vsinpu_provider_factory_creator.h @@ -0,0 +1,34 @@ +/**************************************************************************** + * + * Copyright (c) 2023 Vivante Corporation + * + * Permission is hereby granted, free of charge, to any person obtaining a + * copy of this software and associated documentation files (the "Software"), + * to deal in the Software without restriction, including without limitation + * the rights to use, copy, modify, merge, publish, distribute, sublicense, + * and/or sell copies of the Software, and to permit persons to whom the + * Software is furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING + * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER + * DEALINGS IN THE SOFTWARE. + * + *****************************************************************************/ +#pragma once + +#include + +#include "core/providers/providers.h" + +namespace onnxruntime { +struct VSINPUProviderFactoryCreator { + static std::shared_ptr Create(); +}; +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/vsinpu/vsinpu_util.cc b/onnxruntime/core/providers/vsinpu/vsinpu_util.cc new file mode 100644 index 0000000000000..8008ec1f436a4 --- /dev/null +++ b/onnxruntime/core/providers/vsinpu/vsinpu_util.cc @@ -0,0 +1,502 @@ +/**************************************************************************** + * + * Copyright (c) 2023 Vivante Corporation + * + * Permission is hereby granted, free of charge, to any person obtaining a + * copy of this software and associated documentation files (the "Software"), + * to deal in the Software without restriction, including without limitation + * the rights to use, copy, modify, merge, publish, distribute, sublicense, + * and/or sell copies of the Software, and to permit persons to whom the + * Software is furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING + * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER + * DEALINGS IN THE SOFTWARE. + * + *****************************************************************************/ + +#include +#include +#include +#include +#include "core/providers/vsinpu/vsinpu_util.h" + +#include "core/optimizer/initializer.h" +#include "core/providers/shared/utils/utils.h" +namespace onnxruntime { + +template +struct shared_array_deletor { + void operator()(T const* ptr) { delete[] ptr; } +}; +namespace vsi { +namespace npu { +namespace util { +tim::vx::DataType OnnxDtypeToTIMVXDtype(const int32_t dtype) { + switch (dtype) { + case onnx::TensorProto_DataType_FLOAT: + return tim::vx::DataType::FLOAT32; + case onnx::TensorProto_DataType_FLOAT16: + return tim::vx::DataType::FLOAT16; + case onnx::TensorProto_DataType_INT8: + return tim::vx::DataType::INT8; + case onnx::TensorProto_DataType_UINT8: + return tim::vx::DataType::UINT8; + case onnx::TensorProto_DataType_INT32: + return tim::vx::DataType::INT32; + case onnx::TensorProto_DataType_INT16: + return tim::vx::DataType::INT16; + case onnx::TensorProto_DataType_UINT16: + return tim::vx::DataType::UINT16; + case onnx::TensorProto_DataType_BOOL: + return tim::vx::DataType::BOOL8; + default: + LOGS_DEFAULT(WARNING) << "Unsupported data type: " << dtype; + break; + } + return tim::vx::DataType::FLOAT32; +} + +tim::vx::DataType OnnxDtypeToTIMVXDtype(const ONNX_NAMESPACE::DataType type) { + static const std::map type_table = { + {"tensor(float)", tim::vx::DataType::FLOAT32}, + {"tensor(float16)", tim::vx::DataType::FLOAT16}, + {"tensor(int8)", tim::vx::DataType::INT8}, + {"tensor(uint8)", tim::vx::DataType::UINT8}, + {"tensor(int32)", tim::vx::DataType::INT32}, + {"tensor(int16)", tim::vx::DataType::INT16}, + {"tensor(uint16)", tim::vx::DataType::UINT16}, + {"tensor(int64)", tim::vx::DataType::INT64}, + {"tensor(bool)", tim::vx::DataType::BOOL8}, + }; + auto search = type_table.find(*type); + if (search != type_table.end()) { + return search->second; + } + LOGS_DEFAULT(WARNING) << "Unsupported data type: " << *type; + return tim::vx::DataType::FLOAT32; +} + +tim::vx::ShapeType OnnxShapeToTIMVXShape(const onnxruntime::TensorShape& ts) { + tim::vx::ShapeType timvx_shape(ts.NumDimensions()); + if (ts.NumDimensions() == 0) { + timvx_shape.push_back(1); + } else { + for (size_t i = 0; i < ts.NumDimensions(); i++) { + timvx_shape[i] = ts.GetDims()[i]; + } + } + return timvx_shape; +} + +std::string PrintNode(const onnxruntime::NodeArg& node_arg) { + auto shape = node_arg.Shape(); + if (shape == nullptr) { + return ""; + } + std::string s = node_arg.Name() + ":<"; + if (shape->dim_size() == 0) { + s += "1>, is a scalar"; + return s; + } + for (int i = 0; i < shape->dim_size(); i++) { + auto dim = shape->dim(i); + std::string s1; + std::stringstream ss; + ss << dim.dim_value(); + ss >> s1; + s += s1; + if (i < shape->dim_size() - 1) { + s += ","; + } else { + s += ">"; + } + } + return s; +} + +std::string PrintNode(const std::vector shape) { + if (shape.size() == 0) { + return ""; + } + std::string s = "<"; + for (std::size_t i = 0; i < shape.size(); i++) { + auto dim = shape[i]; + std::string s1; + std::stringstream ss; + ss << dim; + ss >> s1; + s += s1; + if (i < shape.size() - 1) { + s += ","; + } else { + s += ">"; + } + } + return s; +} + +size_t GetTensorElementSize(const ONNXTensorElementDataType type) { + switch (type) { + case onnx::TensorProto_DataType_INT64: + return 8; + case onnx::TensorProto_DataType_FLOAT: + case onnx::TensorProto_DataType_INT32: + return 4; + case onnx::TensorProto_DataType_FLOAT16: + case onnx::TensorProto_DataType_INT16: + case onnx::TensorProto_DataType_UINT16: + return 2; + case onnx::TensorProto_DataType_INT8: + case onnx::TensorProto_DataType_UINT8: + case onnx::TensorProto_DataType_BOOL: + return 1; + default: + break; + } + return 0; +} + +size_t GetTensorBytes(const Ort::TensorTypeAndShapeInfo& info) { + return info.GetElementCount() * GetTensorElementSize(info.GetElementType()); +} + +TensorShape GetTensorShape(const onnxruntime::NodeArg& node_arg) { + auto shape_proto = node_arg.Shape(); + std::vector dims; + if (shape_proto != nullptr) { + for (int i = 0; i < shape_proto->dim_size(); i++) { + auto dim = shape_proto->dim(i); + dims.push_back(dim.dim_value()); + } + } + if (dims.size() == 0) { + dims.push_back(1); + } + TensorShape ts(dims); + return ts; +} + +std::shared_ptr UnpackTensor( + const NodeArg* node_arg, const ONNX_NAMESPACE::TensorProto& initializer) { + std::shared_ptr unpackedTensor; + auto shape = GetTensorShape(*node_arg); + size_t elementCount = shape.Size(); + +#define CASE_PROTO(X, Y) \ + case ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_##X: { \ + size_t tensorByteSize = elementCount * sizeof(Y); \ + unpackedTensor.reset(new uint8_t[tensorByteSize], \ + shared_array_deletor()); \ + auto status = onnxruntime::utils::UnpackTensor( \ + initializer, \ + initializer.has_raw_data() ? initializer.raw_data().data() : nullptr, \ + initializer.has_raw_data() ? initializer.raw_data().size() : 0, \ + reinterpret_cast(unpackedTensor.get()), elementCount); \ + if (!status.IsOK()) { \ + LOGS_DEFAULT(ERROR) << "Unpack tensor data failed."; \ + } \ + break; \ + } + switch (initializer.data_type()) { + CASE_PROTO(FLOAT, float); + CASE_PROTO(DOUBLE, double); + CASE_PROTO(BOOL, bool); + CASE_PROTO(INT8, int8_t); + CASE_PROTO(INT16, int16_t); + CASE_PROTO(INT32, int32_t); + CASE_PROTO(INT64, int64_t); + CASE_PROTO(UINT8, uint8_t); + CASE_PROTO(UINT16, uint16_t); + CASE_PROTO(UINT32, uint32_t); + CASE_PROTO(FLOAT16, onnxruntime::MLFloat16); + default: + return nullptr; + } + + return unpackedTensor; +} + +tim::vx::PadType GetPadType(const std::string type) { + static const std::map type_table = { + {"NOTSET", tim::vx::PadType::AUTO}, + {"SAME_UPPER", tim::vx::PadType::SAME}, + {"SAME_LOWER", tim::vx::PadType::SAME}, + {"VALID", tim::vx::PadType::VALID}, + }; + auto search = type_table.find(type); + if (search != type_table.end()) { + return search->second; + } + return tim::vx::PadType::NONE; +} + +int32_t ReverseAxis(int32_t origin_axis, int32_t length) { + int32_t axis = 0; + if (origin_axis < 0) { + origin_axis += length; + } + axis = length - origin_axis - 1; + return axis; +} + +std::vector ReverseAxis(std::vector origin_axes, int32_t length) { + std::vector axes; + for (int32_t& axis : origin_axes) { + if (axis < 0) { + axis += length; + } + axes.push_back(length - axis - 1); + } + std::sort(axes.begin(), axes.end()); + return axes; +} + +bool IsTypeSupported(const NodeArg* node_arg) { + const auto* type_proto = node_arg->TypeAsProto(); + if (!type_proto) { + return false; + } + + switch (type_proto->tensor_type().elem_type()) { + case ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_BOOL: + case ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_FLOAT: + case ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_FLOAT16: + case ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_INT8: + case ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_UINT8: + case ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_INT32: + case ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_INT64: + return true; + default: + return false; + } +} + +QuantizedOpType GetQuantizedOpType(const NodeUnit& node_unit) { + const auto& op_type = node_unit.OpType(); + if (node_unit.UnitType() == NodeUnit::Type::SingleNode) { + if (op_type == "DequantizeLinear") + return QuantizedOpType::DequantizeLinear; + else if (op_type == "QuantizeLinear") + return QuantizedOpType::QuantizeLinear; + else if (op_type == "QLinearConv") + return QuantizedOpType::QLinearConv; + else if (op_type == "QLinearMatMul") + return QuantizedOpType::QLinearMatMul; + else if (op_type == "QLinearAdd") + return QuantizedOpType::QLinearAdd; + else if (op_type == "QLinearMul") + return QuantizedOpType::QLinearMul; + else if (op_type == "QLinearSigmoid") + return QuantizedOpType::QLinearSigmoid; + else if (op_type == "QLinearAveragePool") + return QuantizedOpType::QLinearAveragePool; + } else if (node_unit.UnitType() == NodeUnit::Type::QDQGroup) { + if (op_type == "Conv") + return QuantizedOpType::QDQConv; + else if (op_type == "Resize") + return QuantizedOpType::QDQResize; + else if (op_type == "AveragePool") + return QuantizedOpType::QDQAveragePool; + else if (op_type == "Add") + return QuantizedOpType::QDQAdd; + else if (op_type == "Mul") + return QuantizedOpType::QDQMul; + else if (op_type == "Transpose") + return QuantizedOpType::QDQTranspose; + else if (op_type == "Reshape") + return QuantizedOpType::QDQReshape; + else if (op_type == "Softmax") + return QuantizedOpType::QDQSoftmax; + else if (op_type == "Concat") + return QuantizedOpType::QDQConcat; + else if (op_type == "Gemm") + return QuantizedOpType::QDQGemm; + else if (op_type == "MatMul") + return QuantizedOpType::QDQMatMul; + } + return QuantizedOpType::Unknown; +} + +ConvType GetConvType(const NodeUnit& node_unit, const InitializedTensorSet& initializers) { + NodeAttrHelper helper(node_unit); + const auto group = helper.Get("group", 1); + + const auto& weight = node_unit.Inputs()[1].node_arg.Name(); + const auto& weight_tensor = *initializers.at(weight); + + // For ONNX we only have 1 conv ops + // For VSINPU we have 3 + // Input is (W, H, C, N) + // group == 1, --> regular conv + // group != 1 && weight is (kW, kH, group, M), --> depthwise conv + // group != 1 && weight is (kW, kH, C/group, M), --> grouped conv + if (group == 1) + return ConvType::Regular; + else if ((weight_tensor.dims()[1] == group)) + return ConvType::Depthwise; + else + return ConvType::Grouped; +} + +bool IsQuantizedConv(QuantizedOpType quant_op_type) { + return (quant_op_type == QuantizedOpType::QLinearConv) || + (quant_op_type == QuantizedOpType::QDQConv); +} + +bool IsQuantizedPool(QuantizedOpType quant_op_type) { + return (quant_op_type == QuantizedOpType::QLinearAveragePool) || + (quant_op_type == QuantizedOpType::QDQAveragePool); +} + +bool IsQuantizedGemm(QuantizedOpType quant_op_type) { + return (quant_op_type == QuantizedOpType::QLinearMatMul) || + (quant_op_type == QuantizedOpType::QDQGemm) || + (quant_op_type == QuantizedOpType::QDQMatMul); +} + +bool IsQuantizedBinaryOp(QuantizedOpType quant_op_type) { + return quant_op_type == QuantizedOpType::QLinearMatMul || + quant_op_type == QuantizedOpType::QLinearAdd || + quant_op_type == QuantizedOpType::QLinearMul || + quant_op_type == QuantizedOpType::QDQAdd || + quant_op_type == QuantizedOpType::QDQMul || + quant_op_type == QuantizedOpType::QDQGemm || + quant_op_type == QuantizedOpType::QDQMatMul || + IsQuantizedConv(quant_op_type); +} + +bool HasValidBinaryOpQuantizedInputTypes(const NodeUnit& node_unit) { + auto quant_op_type = GetQuantizedOpType(node_unit); + int32_t a_input_type, b_input_type; + if (!IsQuantizedBinaryOp(quant_op_type)) { + LOGS_DEFAULT(VERBOSE) << "[" << node_unit.OpType() << "] is not a binary qlinear op"; + return false; + } + + const auto& inputs = node_unit.Inputs(); + if (!GetType(inputs[0].node_arg, a_input_type)) + return false; + if (!GetType(inputs[1].node_arg, b_input_type)) + return false; + + // QlinearConv/MatMul/QDQGemm/QDQMatMul supports u8u8 or u8s8 + // QLinearAdd/QLinearMul only support u8u8 + bool is_quant_conv_or_gemm = IsQuantizedConv(quant_op_type) || IsQuantizedGemm(quant_op_type); + + bool has_valid_qlinear_conv_weight = + (b_input_type == ONNX_NAMESPACE::TensorProto_DataType_UINT8 || + b_input_type == ONNX_NAMESPACE::TensorProto_DataType_INT8); + + bool has_valid_qlinear_conv_input = + (a_input_type == ONNX_NAMESPACE::TensorProto_DataType_UINT8 || + a_input_type == ONNX_NAMESPACE::TensorProto_DataType_INT8); + + if ((is_quant_conv_or_gemm && !has_valid_qlinear_conv_weight) || + (!is_quant_conv_or_gemm && a_input_type != b_input_type)) { + LOGS_DEFAULT(VERBOSE) << "[" << node_unit.OpType() + << "] A Input type: [" << a_input_type + << "] B Input type: [" << b_input_type + << "] is not supported for now"; + return false; + } + + return true; +} + +void GetQuantizationScaleAndZeroPoint( + const InitializedTensorSet& initializers, const NodeUnitIODef& io_def, const Path& model_path, + float& scale, int32_t& zero_point, std::optional>& pcq_scales, + std::optional>& pcq_zps) { + scale = 0.0f; + zero_point = 0; + + const auto& quant_param = *io_def.quant_param; + { // get the scale + const auto& name = quant_param.scale.Name(); + Initializer unpacked_tensor(*initializers.at(name), model_path); + scale = unpacked_tensor.DataAsSpan()[0]; + + // per channel quantized handling + if (!unpacked_tensor.dims().empty() && unpacked_tensor.dims()[0] != 0 && unpacked_tensor.dims()[0] != 1) { + auto scales = unpacked_tensor.DataAsSpan(); + std::vector scales_vec(scales.begin(), scales.end()); + pcq_scales = onnxruntime::make_optional(std::move(scales_vec)); + } + } + + if (quant_param.zero_point) { // get the zero point if it exists + const auto& name = quant_param.zero_point->Name(); + Initializer unpacked_tensor(*initializers.at(name), model_path); + bool is_i8_zp = unpacked_tensor.data_type() == onnx::TensorProto_DataType_INT8; + // some qdq conv bias is int32 quantized + bool is_int32_zp = unpacked_tensor.data_type() == onnx::TensorProto_DataType_INT32; + zero_point = is_i8_zp ? static_cast(unpacked_tensor.DataAsSpan()[0]) : is_int32_zp ? static_cast(unpacked_tensor.DataAsSpan()[0]) + : static_cast(unpacked_tensor.DataAsByteSpan()[0]); + + // per channel quantized handling + if (!unpacked_tensor.dims().empty() && unpacked_tensor.dims()[0] != 0 && unpacked_tensor.dims()[0] != 1) { + auto type = unpacked_tensor.data_type(); + if (is_i8_zp) { + auto zps = unpacked_tensor.DataAsSpan(); + std::vector zps_vec(zps.begin(), zps.end()); + pcq_zps = onnxruntime::make_optional(std::move(zps_vec)); + } else if (is_int32_zp) { + auto zps = unpacked_tensor.DataAsByteSpan(); + std::vector zps_vec(zps.begin(), zps.end()); + pcq_zps = onnxruntime::make_optional(std::move(zps_vec)); + } else { + auto zps = unpacked_tensor.DataAsSpan(); + std::vector zps_vec(zps.begin(), zps.end()); + pcq_zps = onnxruntime::make_optional(std::move(zps_vec)); + } + } + } +} + +static bool IsInternalQuantizedNodeUnit(const NodeUnit& node_unit) { + // First, ignore QDQ NodeUnit which is not internal quantized node + if (node_unit.UnitType() == NodeUnit::Type::QDQGroup) + return false; + + // These operators can use uint8 input without specific QLinear version of it + // However, the mode has to be internal to the graph/partition (they cannot consume graph inputs) + static const std::unordered_set internal_quantized_op_types = { + "Transpose", + "Resize", + "Concat", + "MaxPool", + }; + + const auto& node = node_unit.GetNode(); + if (!Contains(internal_quantized_op_types, node.OpType())) + return false; + + int32_t input_type; + ORT_ENFORCE(GetType(*node.InputDefs()[0], input_type)); + + return input_type == ONNX_NAMESPACE::TensorProto_DataType_UINT8 || input_type == ONNX_NAMESPACE::TensorProto_DataType_INT8; +} + +bool GetType(const NodeArg& node_arg, int32_t& type) { + type = ONNX_NAMESPACE::TensorProto_DataType_UNDEFINED; + const auto* type_proto = node_arg.TypeAsProto(); + if (!type_proto || !type_proto->has_tensor_type() || !type_proto->tensor_type().has_elem_type()) { + LOGS_DEFAULT(WARNING) << "NodeArg [" << node_arg.Name() << "] has no input type"; + return false; + } + + type = type_proto->tensor_type().elem_type(); + return true; +} +} // namespace util +} // namespace npu +} // namespace vsi +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/vsinpu/vsinpu_util.h b/onnxruntime/core/providers/vsinpu/vsinpu_util.h new file mode 100644 index 0000000000000..9ec580bf02e77 --- /dev/null +++ b/onnxruntime/core/providers/vsinpu/vsinpu_util.h @@ -0,0 +1,131 @@ +/**************************************************************************** + * + * Copyright (c) 2023 Vivante Corporation + * + * Permission is hereby granted, free of charge, to any person obtaining a + * copy of this software and associated documentation files (the "Software"), + * to deal in the Software without restriction, including without limitation + * the rights to use, copy, modify, merge, publish, distribute, sublicense, + * and/or sell copies of the Software, and to permit persons to whom the + * Software is furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING + * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER + * DEALINGS IN THE SOFTWARE. + * + *****************************************************************************/ + +#pragma once +#include +#include +#include +#include "core/framework/op_kernel.h" +#include "core/framework/tensor_type_and_shape.h" +#include "core/framework/tensorprotoutils.h" +#include "core/session/onnxruntime_cxx_api.h" +#include "core/framework/node_unit.h" +#include "tim/vx/tensor.h" +#include "tim/vx/types.h" + +namespace onnxruntime { +namespace vsi { +namespace npu { +namespace util { + +tim::vx::DataType OnnxDtypeToTIMVXDtype(const int32_t dtype); + +tim::vx::DataType OnnxDtypeToTIMVXDtype(const ONNX_NAMESPACE::DataType type); + +tim::vx::ShapeType OnnxShapeToTIMVXShape(const onnxruntime::TensorShape& ts); + +std::string PrintNode(const onnxruntime::NodeArg& node_arg); + +std::string PrintNode(const std::vector shape); + +size_t GetTensorElementSize(const ONNXTensorElementDataType type); + +size_t GetTensorBytes(const Ort::TensorTypeAndShapeInfo& info); + +TensorShape GetTensorShape(const onnxruntime::NodeArg& node_arg); + +std::shared_ptr UnpackTensor( + const NodeArg* node, const ONNX_NAMESPACE::TensorProto& initializer); + +tim::vx::PadType GetPadType(const std::string type); + +int32_t ReverseAxis(int32_t origin_axis, int32_t length); + +std::vector ReverseAxis(std::vector origin_axes, int32_t length); + +bool IsTypeSupported(const NodeArg* node_arg); + +enum class QuantizedOpType : uint8_t { + Unknown, // Unknown or not a quantized NodeUnit + DequantizeLinear, + QuantizeLinear, + QLinearConv, + QLinearMatMul, + QLinearAdd, + QLinearSigmoid, + QLinearAveragePool, + QLinearMul, + // Not yet supported + // QLinearReduceMean, + QDQConv, + QDQResize, + QDQAveragePool, + QDQAdd, + QDQMul, + QDQTranspose, + QDQReshape, + QDQSoftmax, + QDQConcat, + QDQGemm, + QDQMatMul, + // TODO(cfy) :Add other QDQ NodeUnit types +}; + +enum class ConvType : uint8_t { + Regular, + Depthwise, + Grouped, +}; +QuantizedOpType GetQuantizedOpType(const NodeUnit& node_unit); + +ConvType GetConvType(const NodeUnit& node_unit, const InitializedTensorSet& initializers); + +// If this is a quantized Conv (QLinearConv or QDQConv) +bool IsQuantizedConv(QuantizedOpType quant_op_type); + +// If this is a quantized Pool (QLinearAveragePool or QDQAveragePool) +bool IsQuantizedPool(QuantizedOpType quant_op_type); + +// If this is a quantized Gemm (QLinearMatMul or QDQMatMul/QDQGemm) +bool IsQuantizedGemm(QuantizedOpType quant_op_type); + +// This quantized op is an operator or qdq node unit takes 2 inputs and produces 1 output +// Such as QLinearConv, QLinearMatMul, QLinearAdd, QDQConv,... +bool IsQuantizedBinaryOp(QuantizedOpType quant_op_type); + +// Check if a qlinear binary op has valid inputs, Qlinear[Conv/MatMul/Add] +bool HasValidBinaryOpQuantizedInputTypes(const NodeUnit& node_unit); + +void GetQuantizationScaleAndZeroPoint( + const InitializedTensorSet& initializers, const NodeUnitIODef& io_def, const Path& model_path, + float& scale, int32_t& zero_point, + std::optional>& pcq_scales, + std::optional>& pcq_zps); + +bool GetType(const NodeArg& node_arg, int32_t& type); + +} // namespace util +} // namespace npu +} // namespace vsi +} // namespace onnxruntime diff --git a/onnxruntime/test/onnx/TestCase.cc b/onnxruntime/test/onnx/TestCase.cc index 6d3e9c2cb7865..3319fdd34646b 100644 --- a/onnxruntime/test/onnx/TestCase.cc +++ b/onnxruntime/test/onnx/TestCase.cc @@ -796,7 +796,7 @@ void LoadTests(const std::vector>& input_paths auto test_case_dir = model_info->GetDir(); auto test_case_name_in_log = test_case_name + ORT_TSTR(" in ") + test_case_dir.native(); -#if !defined(ORT_MINIMAL_BUILD) && !defined(USE_QNN) +#if !defined(ORT_MINIMAL_BUILD) && !defined(USE_QNN) && !defined(USE_VSINPU) // to skip some models like *-int8 or *-qdq if ((reinterpret_cast(model_info.get()))->HasDomain(ONNX_NAMESPACE::AI_ONNX_TRAINING_DOMAIN) || (reinterpret_cast(model_info.get()))->HasDomain(ONNX_NAMESPACE::AI_ONNX_PREVIEW_TRAINING_DOMAIN)) { diff --git a/onnxruntime/test/onnx/main.cc b/onnxruntime/test/onnx/main.cc index 0356bf5218cc2..fc29756a1ff98 100644 --- a/onnxruntime/test/onnx/main.cc +++ b/onnxruntime/test/onnx/main.cc @@ -44,7 +44,7 @@ void usage() { "\t-r [repeat]: Specifies the number of times to repeat\n" "\t-v: verbose\n" "\t-n [test_case_name]: Specifies a single test case to run.\n" - "\t-e [EXECUTION_PROVIDER]: EXECUTION_PROVIDER could be 'cpu', 'cuda', 'dnnl', 'tensorrt', " + "\t-e [EXECUTION_PROVIDER]: EXECUTION_PROVIDER could be 'cpu', 'cuda', 'dnnl', 'tensorrt', 'vsinpu'" "'openvino', 'rocm', 'migraphx', 'acl', 'armnn', 'xnnpack', 'nnapi', 'qnn', 'snpe' or 'coreml'. " "Default: 'cpu'.\n" "\t-p: Pause after launch, can attach debugger and continue\n" @@ -169,6 +169,7 @@ int real_main(int argc, char* argv[], Ort::Env& env) { bool enable_mem_pattern = true; bool enable_qnn = false; bool enable_nnapi = false; + bool enable_vsinpu = false; bool enable_coreml = false; bool enable_snpe = false; bool enable_dml = false; @@ -248,6 +249,8 @@ int real_main(int argc, char* argv[], Ort::Env& env) { enable_qnn = true; } else if (!CompareCString(optarg, ORT_TSTR("nnapi"))) { enable_nnapi = true; + } else if (!CompareCString(optarg, ORT_TSTR("vsinpu"))) { + enable_vsinpu = true; } else if (!CompareCString(optarg, ORT_TSTR("coreml"))) { enable_coreml = true; } else if (!CompareCString(optarg, ORT_TSTR("snpe"))) { @@ -561,6 +564,14 @@ int real_main(int argc, char* argv[], Ort::Env& env) { #else fprintf(stderr, "NNAPI is not supported in this build"); return -1; +#endif + } + if (enable_vsinpu) { +#ifdef USE_VSINPU + Ort::ThrowOnError(OrtSessionOptionsAppendExecutionProvider_VSINPU(sf)); +#else + fprintf(stderr, "VSINPU is not supported in this build"); + return -1; #endif } if (enable_coreml) { diff --git a/onnxruntime/test/perftest/command_args_parser.cc b/onnxruntime/test/perftest/command_args_parser.cc index 175079d8197bf..b7c99fa66a1ea 100644 --- a/onnxruntime/test/perftest/command_args_parser.cc +++ b/onnxruntime/test/perftest/command_args_parser.cc @@ -261,6 +261,8 @@ static bool ParseSessionConfigs(const std::string& configs_string, test_config.machine_config.provider_type_name = onnxruntime::kSnpeExecutionProvider; } else if (!CompareCString(optarg, ORT_TSTR("nnapi"))) { test_config.machine_config.provider_type_name = onnxruntime::kNnapiExecutionProvider; + } else if (!CompareCString(optarg, ORT_TSTR("vsinpu"))) { + test_config.machine_config.provider_type_name = onnxruntime::kVSINPUExecutionProvider; } else if (!CompareCString(optarg, ORT_TSTR("coreml"))) { test_config.machine_config.provider_type_name = onnxruntime::kCoreMLExecutionProvider; } else if (!CompareCString(optarg, ORT_TSTR("dml"))) { diff --git a/onnxruntime/test/perftest/ort_test_session.cc b/onnxruntime/test/perftest/ort_test_session.cc index 1485a4456d326..ff782da35cbe6 100644 --- a/onnxruntime/test/perftest/ort_test_session.cc +++ b/onnxruntime/test/perftest/ort_test_session.cc @@ -397,6 +397,12 @@ select from 'TF8', 'TF16', 'UINT8', 'FLOAT', 'ITENSOR'. \n)"); Ort::ThrowOnError(OrtSessionOptionsAppendExecutionProvider_Nnapi(session_options, nnapi_flags)); #else ORT_THROW("NNAPI is not supported in this build\n"); +#endif + } else if (provider_name_ == onnxruntime::kVSINPUExecutionProvider) { +#ifdef USE_VSINPU + Ort::ThrowOnError(OrtSessionOptionsAppendExecutionProvider_VSINPU(session_options)); +#else + ORT_THROW("VSINPU is not supported in this build\n"); #endif } else if (provider_name_ == onnxruntime::kCoreMLExecutionProvider) { #ifdef __APPLE__ diff --git a/onnxruntime/test/providers/base_tester.cc b/onnxruntime/test/providers/base_tester.cc index 8d84c689cd23e..1db8616c85daa 100644 --- a/onnxruntime/test/providers/base_tester.cc +++ b/onnxruntime/test/providers/base_tester.cc @@ -428,6 +428,7 @@ bool SetEpsForAllNodes(Graph& graph, if (provider_type == onnxruntime::kOpenVINOExecutionProvider || provider_type == onnxruntime::kTensorrtExecutionProvider || provider_type == onnxruntime::kNnapiExecutionProvider || + provider_type == onnxruntime::kVSINPUExecutionProvider || provider_type == onnxruntime::kCoreMLExecutionProvider || provider_type == onnxruntime::kDnnlExecutionProvider || provider_type == onnxruntime::kQnnExecutionProvider || @@ -649,6 +650,7 @@ void BaseTester::RunWithConfig(size_t* number_of_pre_packed_weights_counter, kAclExecutionProvider, kArmNNExecutionProvider, kNnapiExecutionProvider, + kVSINPUExecutionProvider, kRocmExecutionProvider, kCoreMLExecutionProvider, kCoreMLExecutionProviderMLProgram, @@ -688,6 +690,8 @@ void BaseTester::RunWithConfig(size_t* number_of_pre_packed_weights_counter, execution_provider = DefaultTensorrtExecutionProvider(); else if (provider_type == onnxruntime::kNnapiExecutionProvider) execution_provider = DefaultNnapiExecutionProvider(); + else if (provider_type == onnxruntime::kVSINPUExecutionProvider) + execution_provider = DefaultVSINPUExecutionProvider(); else if (provider_type == onnxruntime::kRknpuExecutionProvider) execution_provider = DefaultRknpuExecutionProvider(); else if (provider_type == onnxruntime::kAclExecutionProvider) diff --git a/onnxruntime/test/providers/cpu/model_tests.cc b/onnxruntime/test/providers/cpu/model_tests.cc index dcb592a4a254e..cb9887314eb66 100644 --- a/onnxruntime/test/providers/cpu/model_tests.cc +++ b/onnxruntime/test/providers/cpu/model_tests.cc @@ -30,6 +30,10 @@ #include "core/providers/nnapi/nnapi_provider_factory.h" #endif +#ifdef USE_VSINPU +#include "core/providers/vsinpu/vsinpu_provider_factory.h" +#endif + #ifdef USE_RKNPU #include "core/providers/rknpu/rknpu_provider_factory.h" #endif @@ -238,6 +242,11 @@ TEST_P(ModelTest, Run) { ASSERT_ORT_STATUS_OK(OrtSessionOptionsAppendExecutionProvider_Nnapi(ortso, 0)); } #endif +#ifdef USE_VSINPU + else if (provider_name == "vsinpu") { + ASSERT_ORT_STATUS_OK(OrtSessionOptionsAppendExecutionProvider_VSINPU(ortso)); + } +#endif #ifdef USE_RKNPU else if (provider_name == "rknpu") { ASSERT_ORT_STATUS_OK(OrtSessionOptionsAppendExecutionProvider_Rknpu(ortso)); @@ -406,6 +415,9 @@ static constexpr ORT_STRING_VIEW provider_name_dnnl = ORT_TSTR("dnnl"); #if defined(USE_NNAPI) && defined(__ANDROID__) static constexpr ORT_STRING_VIEW provider_name_nnapi = ORT_TSTR("nnapi"); #endif +#ifdef USE_VSINPU +static ORT_STRING_VIEW provider_name_vsinpu = ORT_TSTR("vsinpu"); +#endif #ifdef USE_RKNPU static constexpr ORT_STRING_VIEW provider_name_rknpu = ORT_TSTR("rknpu"); #endif @@ -447,6 +459,9 @@ ::std::vector<::std::basic_string> GetParameterStrings() { #if defined(USE_NNAPI) && defined(__ANDROID__) provider_names[provider_name_nnapi] = {opset7, opset8, opset9, opset10, opset11, opset12, opset13, opset14, opset15, opset16, opset17, opset18}; #endif +#ifdef USE_VSINPU + provider_names[provider_name_vsinpu] = {}; +#endif #ifdef USE_RKNPU provider_names[provider_name_rknpu] = {}; #endif diff --git a/onnxruntime/test/util/default_providers.cc b/onnxruntime/test/util/default_providers.cc index f15ac100f4e3f..312aa86277994 100644 --- a/onnxruntime/test/util/default_providers.cc +++ b/onnxruntime/test/util/default_providers.cc @@ -190,6 +190,14 @@ std::unique_ptr DefaultNnapiExecutionProvider() { #endif } +std::unique_ptr DefaultVSINPUExecutionProvider() { +#if defined(USE_VSINPU) + return VSINPUProviderFactoryCreator::Create()->CreateProvider(); +#else + return nullptr; +#endif +} + std::unique_ptr DefaultRknpuExecutionProvider() { #ifdef USE_RKNPU return RknpuProviderFactoryCreator::Create()->CreateProvider(); diff --git a/onnxruntime/test/util/include/default_providers.h b/onnxruntime/test/util/include/default_providers.h index ae8e89c386994..606dfc068d399 100644 --- a/onnxruntime/test/util/include/default_providers.h +++ b/onnxruntime/test/util/include/default_providers.h @@ -20,6 +20,7 @@ std::shared_ptr CreateExecutionProviderFactory_MIGrap std::shared_ptr CreateExecutionProviderFactory_Nnapi( uint32_t flags, const optional& partitioning_stop_ops_list); // std::shared_ptr CreateExecutionProviderFactory_Tvm(const char*); +std::shared_ptr CreateExecutionProviderFactory_VSINPU(); std::shared_ptr CreateExecutionProviderFactory_Rknpu(); std::shared_ptr CreateExecutionProviderFactory_Rocm(const OrtROCMProviderOptions* provider_options); std::shared_ptr CreateExecutionProviderFactory_Tensorrt(const OrtTensorRTProviderOptions* params); @@ -50,6 +51,7 @@ std::unique_ptr MIGraphXExecutionProviderWithOptions(const O std::unique_ptr OpenVINOExecutionProviderWithOptions(const OrtOpenVINOProviderOptions* params); std::unique_ptr DefaultOpenVINOExecutionProvider(); std::unique_ptr DefaultNnapiExecutionProvider(); +std::unique_ptr DefaultVSINPUExecutionProvider(); std::unique_ptr DefaultRknpuExecutionProvider(); std::unique_ptr DefaultAclExecutionProvider(bool enable_arena = true); std::unique_ptr DefaultArmNNExecutionProvider(bool enable_arena = true); diff --git a/onnxruntime/test/util/include/providers.h b/onnxruntime/test/util/include/providers.h index aa489e6cd958b..a73b237ae10df 100644 --- a/onnxruntime/test/util/include/providers.h +++ b/onnxruntime/test/util/include/providers.h @@ -16,6 +16,9 @@ #ifdef USE_NNAPI #include "core/providers/nnapi/nnapi_provider_factory.h" #endif +#ifdef USE_VSINPU +#include "core/providers/vsinpu/vsinpu_provider_factory.h" +#endif #ifdef USE_COREML #include "core/providers/coreml/coreml_provider_factory.h" #endif diff --git a/tools/ci_build/build.py b/tools/ci_build/build.py index f431f471c4082..b73a17db3ce13 100644 --- a/tools/ci_build/build.py +++ b/tools/ci_build/build.py @@ -561,6 +561,7 @@ def convert_arg_line_to_args(self, arg_line): parser.add_argument("--use_snpe", action="store_true", help="Build with SNPE support.") parser.add_argument("--snpe_root", help="Path to SNPE SDK root.") parser.add_argument("--use_nnapi", action="store_true", help="Build with NNAPI support.") + parser.add_argument("--use_vsinpu", action="store_true", help="Build with VSINPU support.") parser.add_argument( "--nnapi_min_api", type=int, help="Minimum Android API level to enable NNAPI, should be no less than 27" ) @@ -1020,6 +1021,7 @@ def generate_build_tree( "-Donnxruntime_BUILD_APPLE_FRAMEWORK=" + ("ON" if args.build_apple_framework else "OFF"), "-Donnxruntime_USE_DNNL=" + ("ON" if args.use_dnnl else "OFF"), "-Donnxruntime_USE_NNAPI_BUILTIN=" + ("ON" if args.use_nnapi else "OFF"), + "-Donnxruntime_USE_VSINPU=" + ("ON" if args.use_vsinpu else "OFF"), "-Donnxruntime_USE_RKNPU=" + ("ON" if args.use_rknpu else "OFF"), "-Donnxruntime_USE_LLVM=" + ("ON" if args.use_tvm else "OFF"), "-Donnxruntime_ENABLE_MICROSOFT_INTERNAL=" + ("ON" if args.enable_msinternal else "OFF"), From 8c2689877fb48bfb4a6a133b020cdb5ec7c9b066 Mon Sep 17 00:00:00 2001 From: Scott McKay Date: Sun, 30 Jun 2024 05:19:51 +1000 Subject: [PATCH 49/52] CoreML: Disable 1D ML Program matmul due to bug in coreml (#21186) ### Description Disable using CoreML ML Program for a matmul where one of the inputs is 1D as the CoreML implementation appears to be broken. See https://github.com/apple/coremltools/issues/2263 Add some debugging notes. ### Motivation and Context Fix failing test on macos-14. --- .github/workflows/mac.yml | 3 +- .../core/providers/coreml/DebugMLProgram.md | 85 +++++++++++++++++++ .../coreml/builders/impl/gemm_op_builder.cc | 33 ++++--- .../coreml/builders/model_builder.cc | 1 + .../providers/coreml/dump_mlprogram_model.py | 9 +- 5 files changed, 114 insertions(+), 17 deletions(-) create mode 100644 onnxruntime/core/providers/coreml/DebugMLProgram.md diff --git a/.github/workflows/mac.yml b/.github/workflows/mac.yml index 8aaec8adef979..3d94d30947c76 100644 --- a/.github/workflows/mac.yml +++ b/.github/workflows/mac.yml @@ -54,11 +54,10 @@ jobs: --test \ --build_shared_lib \ --build_objc \ + --use_coreml \ --use_xnnpack \ --use_binskim_compliant_compile_flags - # TODO add --use_coreml once unit test failures are addressed - Objective-C-StaticAnalysis: runs-on: macos-14 diff --git a/onnxruntime/core/providers/coreml/DebugMLProgram.md b/onnxruntime/core/providers/coreml/DebugMLProgram.md new file mode 100644 index 0000000000000..e41a515594303 --- /dev/null +++ b/onnxruntime/core/providers/coreml/DebugMLProgram.md @@ -0,0 +1,85 @@ +# Steps to debug an ML Program operator implementation + +Basic debugging of everything, excluding model execution, (e.g. partitioning, checking if operator is supported, +adding CoreML operator input/outputs) can be done anywhere as the code is setup to build and be able to create the +protobuf based CoreML Model on all platforms. + +To debug model execution issues you will need a macOS machine. + +## Debugging invalid output + +If there is a crash during execution or unexpected output, the best approach is to see what using coremltools directly +produces. + +NOTE: that doesn't guarantee coremltools is correct as there could be a bug in their implementation. It does however +provide a data point on whether we are generating the same CoreML model as the coremltools python. + +### Comparing to coremltools output + +Create a small test script that replicates the inputs/outputs of the operator you are debugging. +This script should use the coremltools library to run the operator and print the output. +This can be used to compare the CoreML EP's output with the coremltools output. + +https://apple.github.io/coremltools/docs-guides/source/model-intermediate-language.html#create-a-mil-program + +Usage is reasonably intuitive. The below example defines a model with 2 inputs and a matmul operator. +The model is printed, and run with randomly generated inputs. The output from doing so is printed. + +```python +import numpy as np +import coremltools as ct +from coremltools.converters.mil import Builder as mb + +target = ct.target.iOS15 + +x_shape = (1, 4) +y_shape = (10, 4, 3) + +@mb.program(input_specs=[mb.TensorSpec(shape=x_shape), mb.TensorSpec(shape=y_shape)], + opset_version=target) +def prog(x, y): + # For reference, a constant can be added using `mb.const` and specifying the data in the `val` parameter. + # c_shape = (3, ) + # c_data = np.random.random_sample(c_shape) + # c = mb.const(val=c_data) + + # call the operator you are debugging with the inputs/constants. + # See the spec for the operator names, input/outputs and supported data types. + # https://apple.github.io/coremltools/source/coremltools.converters.mil.mil.ops.defs.html + z = mb.matmul(x=x, y=y) + + # can have additional function calls here if there are multiple operators involved. + # Contrived example that uses a constant and the output from a previous operator: + # z = mb.add(x=z, y=c) + + return z + +# Prints the MIL program in a reasonably concise manner. +print(prog) + +# Convert to ML Program model +m = ct.convert(prog, minimum_deployment_target=target) + +# If you want to dump the full protobuf of the model uncomment this. +# You can compare the values to what is being set by the ORT CoreML EP code if you suspect any issues there. +# spec = m.get_spec() +# print(spec) + +# run the model to generate output for comparison with the CoreML EP output +x = np.random.rand(*x_shape) +y = np.random.rand(*y_shape) + +print(m.predict({'x': x, 'y': y})) +``` + +## Dumping the ORT generated mlmodel + +You can also dump the mlmodel generated by the ORT CoreML EP. This can be handy with larger models. + +In a debug build, set the ORT_COREML_EP_MODEL_DIR environment variable to a directory where you want the ML Package +containing the mlmodel to be saved. The model will remain after the CoreML EP exits, unlike the default behavior +where we write it to a temporary directory that is automatically removed on application exit. + +Script to dump: [dump_mlprogram_model.py](dump_mlprogram_model.py) + +See [here](https://github.com/microsoft/onnxruntime/blob/3c0b407709fd3c71755ed046edd688b30a786d94/onnxruntime/core/providers/coreml/model/host_utils.h#L70-L75) for environment variable setup and [usage](https://github.com/search?q=repo%3Amicrosoft%2Fonnxruntime%20kOverrideModelOutputDirectoryEnvVar%20&type=code). diff --git a/onnxruntime/core/providers/coreml/builders/impl/gemm_op_builder.cc b/onnxruntime/core/providers/coreml/builders/impl/gemm_op_builder.cc index 8daf64dc4a457..7338fc18fe779 100644 --- a/onnxruntime/core/providers/coreml/builders/impl/gemm_op_builder.cc +++ b/onnxruntime/core/providers/coreml/builders/impl/gemm_op_builder.cc @@ -109,19 +109,11 @@ Status GemmOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder, const N ORT_IGNORE_RETURN_VALUE(GetShape(b, b_shape, logger)); int64_t b0 = -1, b1 = -1; - // ML Program MatMul supports N-D input if (model_builder.CreateMLProgram() && is_matmul) { - if (b_shape.size() == 1) { - // B is treated as {b_shape[0], 1} according to the numpy rules. - b0 = b_shape[0]; - b1 = 1; - } else { - // last 2 dims are used - b0 = b_shape[b_shape.size() - 2]; - b1 = b_shape[b_shape.size() - 1]; - } + // ML Program MatMul supports N-D input, however we don't use the 'K' or 'N' values calculated below for it + // so we don't need to update b0 or b1. } else { - // we only support 2D input + // we only support 2D input for all other combinations b0 = b_shape[0]; b1 = b_shape[1]; } @@ -182,7 +174,6 @@ Status GemmOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder, const N model_builder.AddOperation(std::move(gemm_op)); } else { // CoreML implementation is the same as ONNX MatMul. - // https://apple.github.io/coremltools/source/coremltools.converters.mil.mil.ops.defs.html#coremltools.converters.mil.mil.ops.defs.iOS15.linear.matmul auto matmul_op = model_builder.CreateOperation(node, "matmul"); AddOperationInput(*matmul_op, "x", a.Name()); AddOperationInput(*matmul_op, "y", b.Name()); @@ -268,14 +259,28 @@ bool GemmOpBuilder::IsOpSupportedImpl(const Node& node, const OpBuilderInputPara } if (is_matmul) { + const auto a_rank = a_shape.size(); + const auto b_rank = b_shape.size(); + if (input_params.create_mlprogram) { - // ML Program matmul op has numpy semantics the same as the ONNX spec so we can use directly + // ML Program matmul op has numpy semantics the same as the ONNX spec, so we can use directly. + // https://apple.github.io/coremltools/source/coremltools.converters.mil.mil.ops.defs.html#coremltools.converters.mil.mil.ops.defs.iOS15.linear.matmul + // + // There does appear to be a bug in handling one of the inputs being 1D, so for now skip these. + // See https://github.com/apple/coremltools/issues/2263 + // + // If required for perf we could manually do the shape alterations the spec documents (convert input to 2D, + // and remove extra dimension from output), as the 2D input is correctly handled by CoreML matmul. + if ((a_rank == 1 && b_rank > 1) || (a_rank > 1 && b_rank == 1)) { + LOGS(logger, VERBOSE) << "Skipping due to bug in CoreML ML Program when one of the inputs is 1D."; + return false; + } } else { // we could potentially support 1D and 3D if required. beyond 3D the dims that merge diverge. // https://github.com/apple/coremltools/blob/1931758aae383c83daddfc56f11a24a9d2bf4b87/coremltools/converters/onnx/_operators.py#L1607 // https://github.com/apple/coremltools/blob/1931758aae383c83daddfc56f11a24a9d2bf4b87/coremltools/converters/mil/backend/nn/op_mapping.py#L1374 // https://apple.github.io/coremltools/mlmodel/Format/NeuralNetwork.html#innerproductlayerparams - if (a_shape.size() != 2 || b_shape.size() != 2) { + if (a_rank != 2 || b_rank != 2) { LOGS(logger, VERBOSE) << "a and b inputs must be 2D. "; return false; } diff --git a/onnxruntime/core/providers/coreml/builders/model_builder.cc b/onnxruntime/core/providers/coreml/builders/model_builder.cc index 88b518ab2289c..eec0fcce51dbc 100644 --- a/onnxruntime/core/providers/coreml/builders/model_builder.cc +++ b/onnxruntime/core/providers/coreml/builders/model_builder.cc @@ -906,6 +906,7 @@ Status ModelBuilder::SaveModel() { #if defined(COREML_ENABLE_MLPROGRAM) if (create_ml_program_) { + // we need to jump through some hoops to get the model path the ML Program load wants. std::string tmp_model_path = model_output_path_ + "/tmp/model.mlmodel"; CreateEmptyFile(tmp_model_path); diff --git a/onnxruntime/core/providers/coreml/dump_mlprogram_model.py b/onnxruntime/core/providers/coreml/dump_mlprogram_model.py index a3ceee70684dc..dce98e5138d98 100644 --- a/onnxruntime/core/providers/coreml/dump_mlprogram_model.py +++ b/onnxruntime/core/providers/coreml/dump_mlprogram_model.py @@ -5,6 +5,11 @@ if len(sys.argv) < 2: print(f"Usage: {sys.argv[0]} ") print("If generated by onnxruntime this will be /Data/com.microsoft.onnxruntime/model.mlmodel") + print( + "The ML Package created by the CoreML EP can saved to a specific directory in a debug build of onnxruntime " + "by setting the environment variable ORT_COREML_EP_MODEL_DIR to the desired directory." + ) + sys.exit(-1) model_path = sys.argv[1] @@ -13,7 +18,9 @@ spec = m.get_spec() print(spec) -# Example code if you want to filter output or do more advanced things +# Example code if you want to filter output or do more advanced things. +# In the below example we print out the value of an attribute of one specific node from a larger model. +# # main = spec.mlProgram.functions["main"] # block = main.block_specializations[main.opset] # print(f"{len(block.operations)} operators") From beb2496748b112ba0b2525c14f1093acbd98c7aa Mon Sep 17 00:00:00 2001 From: Yi Zhang Date: Tue, 2 Jul 2024 09:24:19 +0800 Subject: [PATCH 50/52] Templatize publishing nuget package (#21199) ### Description It's the prerequisite step of reducing complexity of current zip-nuget pipeline. Some packaging tasks could be cut from the most complex nuget pipline and easily be published ### Motivation and Context --- .../github/azure-pipelines/publish-nuget.yml | 179 +++--------------- .../templates/publish-nuget-steps.yml | 136 +++++++++++++ 2 files changed, 164 insertions(+), 151 deletions(-) create mode 100644 tools/ci_build/github/azure-pipelines/templates/publish-nuget-steps.yml diff --git a/tools/ci_build/github/azure-pipelines/publish-nuget.yml b/tools/ci_build/github/azure-pipelines/publish-nuget.yml index 367977ff59192..5e827980e039c 100644 --- a/tools/ci_build/github/azure-pipelines/publish-nuget.yml +++ b/tools/ci_build/github/azure-pipelines/publish-nuget.yml @@ -10,154 +10,31 @@ resources: branch: main stages: -- stage: Publish_NuGet_Package_And_Report - jobs: - - job: Publish_NuGet_Package_And_Report - workspace: - clean: all - variables: - - name: GDN_CODESIGN_TARGETDIRECTORY - value: '$(Agent.TempDirectory)\binfiles' - pool: 'onnxruntime-Win-CPU-2022' - - steps: - # https://learn.microsoft.com/en-us/azure/devops/pipelines/yaml-schema/resources-pipelines-pipeline?view=azure-pipelines#pipeline-resource-metadata-as-predefined-variables - - script: | - echo $(resources.pipeline.build.sourceBranch) - echo $(Build.Reason) - displayName: 'Print triggering sourceBranch Name in resources' - - - checkout: self - submodules: false - - - task: UsePythonVersion@0 - inputs: - versionSpec: '3.9' - addToPath: true - - - template: templates/set-version-number-variables-step.yml - - - script: mkdir "$(Build.BinariesDirectory)\nuget-artifact\final-package" - - - download: build - displayName: 'Download Pipeline Artifact - Signed NuGet Package' - artifact: 'drop-signed-nuget-CPU' - - - script: move "$(Pipeline.Workspace)\build\drop-signed-nuget-CPU\*" "$(Build.BinariesDirectory)\nuget-artifact\final-package" - - - template: nuget/templates/get-nuget-package-version-as-variable.yml - parameters: - packageFolder: '$(Build.BinariesDirectory)/nuget-artifact/final-package' - - - task: CmdLine@2 - displayName: 'Post binary sizes to the dashboard database using command line' - inputs: - script: | - echo changing directory to artifact download path - cd $(Build.BinariesDirectory)/nuget-artifact/final-package - echo processing nupkg - SETLOCAL EnableDelayedExpansion - FOR /R %%i IN (*.nupkg) do ( - set filename=%%~ni - IF NOT "!filename:~25,7!"=="Managed" ( - echo processing %%~ni.nupkg - copy %%~ni.nupkg %%~ni.zip - echo copied to zip - echo listing lib files in the zip - REM use a single .csv file to put the data - echo os,arch,build_config,size > $(Build.BinariesDirectory)\binary_size_data.txt - 7z.exe l -slt %%~ni.zip runtimes\linux-arm64\native\libonnxruntime.so | findstr /R /C:"^Size = [0-9]*" | for /F "tokens=3" %%a in ('more') do if not "%%a" == "" echo linux,aarch64,default,%%a >> $(Build.BinariesDirectory)\binary_size_data.txt - 7z.exe l -slt %%~ni.zip runtimes\osx-x64\native\libonnxruntime.dylib | findstr /R /C:"^Size = [0-9]*" | for /F "tokens=3" %%a in ('more') do if not "%%a" == "" echo osx,x64,default,%%a >> $(Build.BinariesDirectory)\binary_size_data.txt - 7z.exe l -slt %%~ni.zip runtimes\win-x64\native\onnxruntime.dll | findstr /R /C:"^Size = [0-9]*" | for /F "tokens=3" %%a in ('more') do if not "%%a" == "" echo win,x64,default,%%a >> $(Build.BinariesDirectory)\binary_size_data.txt - 7z.exe l -slt %%~ni.zip runtimes\win-x86\native\onnxruntime.dll | findstr /R /C:"^Size = [0-9]*" | for /F "tokens=3" %%a in ('more') do if not "%%a" == "" echo win,x86,default,%%a >> $(Build.BinariesDirectory)\binary_size_data.txt - ) - ) - - - task: AzureCLI@2 - displayName: 'Azure CLI' - #Only report binary sizes to database if the build build was auto-triggered from the main branch - condition: and (succeeded(), and(eq(variables['resources.pipeline.build.sourceBranch'], 'refs/heads/main'), eq(variables['Build.Reason'], 'ResourceTrigger'))) - inputs: - azureSubscription: AIInfraBuildOnnxRuntimeOSS - scriptLocation: inlineScript - scriptType: batch - inlineScript: | - python.exe -m pip install -r $(Build.SourcesDirectory)\tools\ci_build\github\windows\post_to_dashboard\requirements.txt && ^ - python.exe $(Build.SourcesDirectory)\tools\ci_build\github\windows\post_binary_sizes_to_dashboard.py --commit_hash=$(Build.SourceVersion) --size_data_file=binary_size_data.txt --build_project=Lotus --build_id=$(Build.BuildId) - workingDirectory: '$(Build.BinariesDirectory)' - - - download: build - displayName: 'Download Pipeline Artifact - Signed NuGet Package' - artifact: 'drop-signed-nuget-dml' - - - script: move "$(Pipeline.Workspace)\build\drop-signed-nuget-dml\*" $(Build.BinariesDirectory)\nuget-artifact\final-package - - - download: build - displayName: 'Download Pipeline Artifact - Signed NuGet Package' - artifact: 'drop-signed-nuget-Training-CPU' - - script: move "$(Pipeline.Workspace)\build\drop-signed-nuget-Training-CPU\*" $(Build.BinariesDirectory)\nuget-artifact\final-package - - - download: build - displayName: 'Download Pipeline Artifact - Signed NuGet Package' - artifact: 'drop-signed-nuget-GPU' - - script: move "$(Pipeline.Workspace)\build\drop-signed-nuget-GPU\*" $(Build.BinariesDirectory)\nuget-artifact\final-package - - - download: build - displayName: 'Download Pipeline Artifact - Signed NuGet ROCm Package' - artifact: 'drop-signed-nuget-ROCm' - - script: move "$(Pipeline.Workspace)\build\drop-signed-nuget-ROCm\*" $(Build.BinariesDirectory)\nuget-artifact\final-package - - - download: build - displayName: 'Download Pipeline Artifact - Signed NuGet Qnn Package' - artifact: 'drop-signed-nuget-qnn' - - script: move "$(Pipeline.Workspace)\build\drop-signed-nuget-qnn\*" $(Build.BinariesDirectory)\nuget-artifact\final-package - - - script: | - dir $(Build.BinariesDirectory)\nuget-artifact\final-package - cd $(Build.BinariesDirectory)\nuget-artifact\final-package - nuget verify -Signatures *.nupkg - displayName: List Downloaded Package - - - powershell: | - New-Item -Path $(Agent.TempDirectory) -Name "binfiles" -ItemType "directory" - $base_path_name = Join-Path -Path $(Agent.TempDirectory) -ChildPath "binfiles" - Get-ChildItem $Env:BUILD_BINARIESDIRECTORY\nuget-artifact\final-package -Filter *.nupkg | - Foreach-Object { - $dir_name = Join-Path -Path $base_path_name -ChildPath $_.Basename - $cmd = "7z.exe x $($_.FullName) -y -o$dir_name" - Write-Output $cmd - Invoke-Expression -Command $cmd - } - dir $(Agent.TempDirectory) - tree $(Agent.TempDirectory) - workingDirectory: '$(Agent.TempDirectory)' - - - task: CodeSign@1 - displayName: 'Run Codesign Validation' - - - - task: PublishSecurityAnalysisLogs@3 - displayName: 'Publish Security Analysis Logs' - continueOnError: true - - - task: PostAnalysis@2 - inputs: - GdnBreakAllTools: true - GdnBreakPolicy: M365 - GdnBreakPolicyMinSev: Error - - #TODO: allow choosing different feeds - - task: NuGetCommand@2 - displayName: 'Copy Signed Native NuGet Package to ORT-NIGHTLY' - inputs: - command: 'push' - packagesToPush: '$(Build.BinariesDirectory)/nuget-artifact/final-package/*.nupkg' - publishVstsFeed: '2692857e-05ef-43b4-ba9c-ccf1c22c437c/7982ae20-ed19-4a35-a362-a96ac99897b7' - allowPackageConflicts: true - - - template: templates/component-governance-component-detection-steps.yml - parameters : - condition : 'succeeded' - - task: mspremier.PostBuildCleanup.PostBuildCleanup-task.PostBuildCleanup@3 - displayName: 'Clean Agent Directories' - condition: always() + - template: templates/publish-nuget-steps.yml + parameters: + include_cpu_ep: true + download_artifacts_steps: + - download: build + displayName: 'Download Pipeline Artifact - Signed NuGet Package' + artifact: 'drop-signed-nuget-dml' + - script: move "$(Pipeline.Workspace)\build\drop-signed-nuget-dml\*" $(Build.BinariesDirectory)\nuget-artifact\final-package + + - download: build + displayName: 'Download Pipeline Artifact - Signed NuGet Package' + artifact: 'drop-signed-nuget-Training-CPU' + - script: move "$(Pipeline.Workspace)\build\drop-signed-nuget-Training-CPU\*" $(Build.BinariesDirectory)\nuget-artifact\final-package + + - download: build + displayName: 'Download Pipeline Artifact - Signed NuGet Package' + artifact: 'drop-signed-nuget-GPU' + - script: move "$(Pipeline.Workspace)\build\drop-signed-nuget-GPU\*" $(Build.BinariesDirectory)\nuget-artifact\final-package + + - download: build + displayName: 'Download Pipeline Artifact - Signed NuGet ROCm Package' + artifact: 'drop-signed-nuget-ROCm' + - script: move "$(Pipeline.Workspace)\build\drop-signed-nuget-ROCm\*" $(Build.BinariesDirectory)\nuget-artifact\final-package + + - download: build + displayName: 'Download Pipeline Artifact - Signed NuGet Qnn Package' + artifact: 'drop-signed-nuget-qnn' + - script: move "$(Pipeline.Workspace)\build\drop-signed-nuget-qnn\*" $(Build.BinariesDirectory)\nuget-artifact\final-package diff --git a/tools/ci_build/github/azure-pipelines/templates/publish-nuget-steps.yml b/tools/ci_build/github/azure-pipelines/templates/publish-nuget-steps.yml new file mode 100644 index 0000000000000..6698501e74bad --- /dev/null +++ b/tools/ci_build/github/azure-pipelines/templates/publish-nuget-steps.yml @@ -0,0 +1,136 @@ +parameters: +- name: include_cpu_ep + type: boolean + default: false +- name: download_artifacts_steps + type: stepList + +stages: +- stage: Publish_NuGet_Package_And_Report + jobs: + - job: Publish_NuGet_Package_And_Report + workspace: + clean: all + variables: + - name: GDN_CODESIGN_TARGETDIRECTORY + value: '$(Agent.TempDirectory)\binfiles' + pool: 'onnxruntime-Win-CPU-2022' + + steps: + - task: mspremier.PostBuildCleanup.PostBuildCleanup-task.PostBuildCleanup@3 + displayName: 'Clean Agent Directories' + condition: always() + # https://learn.microsoft.com/en-us/azure/devops/pipelines/yaml-schema/resources-pipelines-pipeline?view=azure-pipelines#pipeline-resource-metadata-as-predefined-variables + - script: | + echo $(resources.pipeline.build.sourceBranch) + echo $(Build.Reason) + displayName: 'Print triggering sourceBranch Name in resources' + + - checkout: self + submodules: false + + - task: UsePythonVersion@0 + inputs: + versionSpec: '3.9' + addToPath: true + + - template: set-version-number-variables-step.yml + + - script: mkdir "$(Build.BinariesDirectory)\nuget-artifact\final-package" + + - template: ../nuget/templates/get-nuget-package-version-as-variable.yml + parameters: + packageFolder: '$(Build.BinariesDirectory)/nuget-artifact/final-package' + + - ${{if eq(parameters.include_cpu_ep, true)}}: + - download: build + displayName: 'Download Pipeline Artifact - Signed NuGet Package' + artifact: 'drop-signed-nuget-CPU' + + - script: move "$(Pipeline.Workspace)\build\drop-signed-nuget-CPU\*" "$(Build.BinariesDirectory)\nuget-artifact\final-package" + + - task: CmdLine@2 + displayName: 'Post binary sizes to the dashboard database using command line' + inputs: + script: | + echo changing directory to artifact download path + cd $(Build.BinariesDirectory)/nuget-artifact/final-package + echo processing nupkg + SETLOCAL EnableDelayedExpansion + FOR /R %%i IN (*.nupkg) do ( + set filename=%%~ni + IF NOT "!filename:~25,7!"=="Managed" ( + echo processing %%~ni.nupkg + copy %%~ni.nupkg %%~ni.zip + echo copied to zip + echo listing lib files in the zip + REM use a single .csv file to put the data + echo os,arch,build_config,size > $(Build.BinariesDirectory)\binary_size_data.txt + 7z.exe l -slt %%~ni.zip runtimes\linux-arm64\native\libonnxruntime.so | findstr /R /C:"^Size = [0-9]*" | for /F "tokens=3" %%a in ('more') do if not "%%a" == "" echo linux,aarch64,default,%%a >> $(Build.BinariesDirectory)\binary_size_data.txt + 7z.exe l -slt %%~ni.zip runtimes\osx-x64\native\libonnxruntime.dylib | findstr /R /C:"^Size = [0-9]*" | for /F "tokens=3" %%a in ('more') do if not "%%a" == "" echo osx,x64,default,%%a >> $(Build.BinariesDirectory)\binary_size_data.txt + 7z.exe l -slt %%~ni.zip runtimes\win-x64\native\onnxruntime.dll | findstr /R /C:"^Size = [0-9]*" | for /F "tokens=3" %%a in ('more') do if not "%%a" == "" echo win,x64,default,%%a >> $(Build.BinariesDirectory)\binary_size_data.txt + 7z.exe l -slt %%~ni.zip runtimes\win-x86\native\onnxruntime.dll | findstr /R /C:"^Size = [0-9]*" | for /F "tokens=3" %%a in ('more') do if not "%%a" == "" echo win,x86,default,%%a >> $(Build.BinariesDirectory)\binary_size_data.txt + ) + ) + + - task: AzureCLI@2 + displayName: 'Azure CLI' + #Only report binary sizes to database if the build build was auto-triggered from the main branch + condition: and (succeeded(), and(eq(variables['resources.pipeline.build.sourceBranch'], 'refs/heads/main'), eq(variables['Build.Reason'], 'ResourceTrigger'))) + inputs: + azureSubscription: AIInfraBuildOnnxRuntimeOSS + scriptLocation: inlineScript + scriptType: batch + inlineScript: | + python.exe -m pip install -r $(Build.SourcesDirectory)\tools\ci_build\github\windows\post_to_dashboard\requirements.txt && ^ + python.exe $(Build.SourcesDirectory)\tools\ci_build\github\windows\post_binary_sizes_to_dashboard.py --commit_hash=$(Build.SourceVersion) --size_data_file=binary_size_data.txt --build_project=Lotus --build_id=$(Build.BuildId) + workingDirectory: '$(Build.BinariesDirectory)' + + - ${{ parameters.download_artifacts_steps }} + + - script: | + dir $(Build.BinariesDirectory)\nuget-artifact\final-package + cd $(Build.BinariesDirectory)\nuget-artifact\final-package + nuget verify -Signatures *.nupkg + displayName: List Downloaded Package + + - powershell: | + New-Item -Path $(Agent.TempDirectory) -Name "binfiles" -ItemType "directory" + $base_path_name = Join-Path -Path $(Agent.TempDirectory) -ChildPath "binfiles" + Get-ChildItem $Env:BUILD_BINARIESDIRECTORY\nuget-artifact\final-package -Filter *.nupkg | + Foreach-Object { + $dir_name = Join-Path -Path $base_path_name -ChildPath $_.Basename + $cmd = "7z.exe x $($_.FullName) -y -o$dir_name" + Write-Output $cmd + Invoke-Expression -Command $cmd + } + dir $(Agent.TempDirectory) + tree $(Agent.TempDirectory) + workingDirectory: '$(Agent.TempDirectory)' + + - task: CodeSign@1 + displayName: 'Run Codesign Validation' + + + - task: PublishSecurityAnalysisLogs@3 + displayName: 'Publish Security Analysis Logs' + continueOnError: true + + - task: PostAnalysis@2 + inputs: + GdnBreakAllTools: true + GdnBreakPolicy: M365 + GdnBreakPolicyMinSev: Error + + #TODO: allow choosing different feeds + - task: NuGetCommand@2 + displayName: 'Copy Signed Native NuGet Package to ORT-NIGHTLY' + inputs: + command: 'push' + packagesToPush: '$(Build.BinariesDirectory)/nuget-artifact/final-package/*.nupkg' + publishVstsFeed: '2692857e-05ef-43b4-ba9c-ccf1c22c437c/7982ae20-ed19-4a35-a362-a96ac99897b7' + allowPackageConflicts: true + + - template: component-governance-component-detection-steps.yml + parameters : + condition : 'succeeded' From 7be1d4aad3f984ebe2c4fb0f7db0b9ca67cc8964 Mon Sep 17 00:00:00 2001 From: Yifan Li <109183385+yf711@users.noreply.github.com> Date: Mon, 1 Jul 2024 22:55:20 -0700 Subject: [PATCH 51/52] [TensorRT EP] Update TRT10.0 deprecated api (#20989) ### Description Note: * This PR would remove C4996 suppression in tensorrt_execution_provider.cc only (according to Nvidia, places with nvinfer.h included need C4996 suppression, when /Zc:__cplusplus is enabled in ORT win build) * A follow-up PR will be raised to update deprecated TRT Plugin api usage. Here are deprecated apis to be updated in this PR: | deprecated api | Update | | ------------------------------------------------------------ | ------------------------------------------------------------ | | [kCUBLAS](https://docs.nvidia.com/deeplearning/tensorrt/api/c_api/namespacenvinfer1.html#a9e1d81e5a8bfeb38b86e22a66d5f836a) | / | | [kCUBLAS_LT](https://docs.nvidia.com/deeplearning/tensorrt/api/c_api/namespacenvinfer1.html#a9e1d81e5a8bfeb38b86e22a66d5f836a) | / | | [kCUDNN](https://docs.nvidia.com/deeplearning/tensorrt/api/c_api/namespacenvinfer1.html#a9e1d81e5a8bfeb38b86e22a66d5f836a) | / | | [reallocateOutput](https://docs.nvidia.com/deeplearning/tensorrt/api/c_api/classnvinfer1_1_1v__1__0_1_1_i_output_allocator.html#acae6441d4029584cc1c6550917518691) | Superseded by [reallocateOutputAsync](https://docs.nvidia.com/deeplearning/tensorrt/api/c_api/classnvinfer1_1_1v__1__0_1_1_i_output_allocator.html#aa40eeb891c1dfe4c1bbf1eabe8c705ab) with cudaStream_t argument | | [createExecutionContextWithoutDeviceMemory](https://docs.nvidia.com/deeplearning/tensorrt/api/c_api/classnvinfer1_1_1_i_cuda_engine.html#adc86bcc42b098204997396ef2b1093fb) | Superseded by [createExecutionContext()](https://docs.nvidia.com/deeplearning/tensorrt/api/c_api/classnvinfer1_1_1_i_cuda_engine.html#a35de29aa6134165a5b14a537e6d99e82) with parameter.
Check [ExecutionContextAllocationStrategy::kUSER_MANAGED](https://docs.nvidia.com/deeplearning/tensorrt/api/c_api/namespacenvinfer1.html#ac6251a050df629edfc0ce037fa366503) for more detail | ### Motivation and Context TRT deprecated api list: https://docs.nvidia.com/deeplearning/tensorrt/api/c_api/deprecated.html --- .../tensorrt/tensorrt_execution_provider.cc | 58 +++++++++++++------ .../tensorrt/tensorrt_execution_provider.h | 5 +- 2 files changed, 43 insertions(+), 20 deletions(-) diff --git a/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc b/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc index 3ca0935b9e46c..be924d6a68268 100644 --- a/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc +++ b/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc @@ -169,11 +169,20 @@ nvinfer1::TacticSources GetTacticSourceFromString(std::string& tactic_string) { nvinfer1::TacticSource source{}; t = toUpper(t); if (t == "CUBLAS") { + LOGS_DEFAULT(WARNING) << "[TensorRT EP] Tactic kCUBLAS is deprecated in TensorRT 10.0"; +#if NV_TENSORRT_MAJOR < 10 source = nvinfer1::TacticSource::kCUBLAS; +#endif } else if (t == "CUBLASLT" || t == "CUBLAS_LT") { + LOGS_DEFAULT(WARNING) << "[TensorRT EP] Tactic kCUBLAS_LT is deprecated in TensorRT 9.0"; +#if NV_TENSORRT_MAJOR < 9 source = nvinfer1::TacticSource::kCUBLAS_LT; +#endif } else if (t == "CUDNN") { + LOGS_DEFAULT(WARNING) << "[TensorRT EP] Tactic kCUDNN is deprecated in TensorRT 10.0"; +#if NV_TENSORRT_MAJOR < 10 source = nvinfer1::TacticSource::kCUDNN; +#endif } else if (t == "EDGE_MASK_CONVOLUTIONS") { source = nvinfer1::TacticSource::kEDGE_MASK_CONVOLUTIONS; } else if (t == "JIT_CONVOLUTIONS") { @@ -298,6 +307,25 @@ void CudaCall(cudnnStatus_t retCode, const char* exprString return g_host->CudaCall_true(retCode, exprString, libName, successCode, msg, file, line); } +#if NV_TENSORRT_MAJOR >= 10 +void* OutputAllocator::reallocateOutputAsync(char const* /*tensorName*/, void* /*currentMemory*/, uint64_t size, + uint64_t /*alignment*/, cudaStream_t /*stream*/) noexcept { + // Some memory allocators return nullptr when allocating zero bytes, but TensorRT requires a non-null ptr + // even for empty tensors, so allocate a dummy byte. + size = std::max(size, static_cast(1)); + if (size > allocated_size) { + cudaFree(outputPtr); + outputPtr = nullptr; + allocated_size = 0; + if (cudaMalloc(&outputPtr, size) == cudaSuccess) { + allocated_size = size; + } + } + // if cudaMalloc fails, returns nullptr. + return outputPtr; +} +#else +// Only override this method when TensorRT <= 8.6 void* OutputAllocator::reallocateOutput(char const* /*tensorName*/, void* /*currentMemory*/, uint64_t size, uint64_t /*alignment*/) noexcept { // Some memory allocators return nullptr when allocating zero bytes, but TensorRT requires a non-null ptr @@ -314,6 +342,7 @@ void* OutputAllocator::reallocateOutput(char const* /*tensorName*/, void* /*curr // if cudaMalloc fails, returns nullptr. return outputPtr; } +#endif void OutputAllocator::notifyShape(char const* /*tensorName*/, nvinfer1::Dims const& dims) noexcept { output_shapes.clear(); @@ -3152,14 +3181,10 @@ Status TensorrtExecutionProvider::CreateNodeComputeInfoFromGraph(const GraphView if (mem_size > max_ctx_mem_size_) { max_ctx_mem_size_ = mem_size; } - -#if defined(_MSC_VER) -#pragma warning(push) -#pragma warning(disable : 4996) // nvinfer1::ICudaEngine::createExecutionContextWithoutDeviceMemory was deprecated -#endif +#if NV_TENSORRT_MAJOR < 10 trt_context = std::unique_ptr(trt_engine->createExecutionContextWithoutDeviceMemory()); -#if defined(_MSC_VER) -#pragma warning(pop) +#else + trt_context = std::unique_ptr(trt_engine->createExecutionContext(nvinfer1::ExecutionContextAllocationStrategy::kUSER_MANAGED)); #endif } else { trt_context = std::unique_ptr(trt_engine->createExecutionContext()); @@ -3606,14 +3631,12 @@ Status TensorrtExecutionProvider::CreateNodeComputeInfoFromGraph(const GraphView if (context_update) { if (trt_state->context_memory_sharing_enable) { -#if defined(_MSC_VER) -#pragma warning(push) -#pragma warning(disable : 4996) // nvinfer1::ICudaEngine::createExecutionContextWithoutDeviceMemory was deprecated -#endif +#if NV_TENSORRT_MAJOR < 10 *(trt_state->context) = std::unique_ptr( trt_state->engine->get()->createExecutionContextWithoutDeviceMemory()); -#if defined(_MSC_VER) -#pragma warning(pop) +#else + *(trt_state->context) = std::unique_ptr( + trt_state->engine->get()->createExecutionContext(nvinfer1::ExecutionContextAllocationStrategy::kUSER_MANAGED)); #endif } else { *(trt_state->context) = std::unique_ptr( @@ -3830,13 +3853,10 @@ Status TensorrtExecutionProvider::CreateNodeComputeInfoFromPrecompiledEngine(con if (mem_size > max_ctx_mem_size_) { max_ctx_mem_size_ = mem_size; } -#if defined(_MSC_VER) -#pragma warning(push) -#pragma warning(disable : 4996) // nvinfer1::ICudaEngine::createExecutionContextWithoutDeviceMemory was deprecated -#endif +#if NV_TENSORRT_MAJOR < 10 trt_context = std::unique_ptr(trt_engine->createExecutionContextWithoutDeviceMemory()); -#if defined(_MSC_VER) -#pragma warning(pop) +#else + trt_context = std::unique_ptr(trt_engine->createExecutionContext(nvinfer1::ExecutionContextAllocationStrategy::kUSER_MANAGED)); #endif } else { trt_context = std::unique_ptr(trt_engine->createExecutionContext()); diff --git a/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.h b/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.h index f4dae57487f51..ec140579569b9 100644 --- a/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.h +++ b/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.h @@ -116,8 +116,11 @@ using unique_pointer = std::unique_ptr; // class OutputAllocator : public nvinfer1::IOutputAllocator { public: +#if NV_TENSORRT_MAJOR >= 10 + void* reallocateOutputAsync(char const* tensorName, void* currentMemory, uint64_t size, uint64_t alignment, cudaStream_t stream) noexcept override; +#else void* reallocateOutput(char const* tensorName, void* currentMemory, uint64_t size, uint64_t alignment) noexcept override; - +#endif void notifyShape(char const* tensorName, nvinfer1::Dims const& dims) noexcept override; void* getBuffer() { From 7df97f1987dcdb798e0c22b3d3ae8f27dfa6a82e Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Tue, 2 Jul 2024 11:24:04 -0700 Subject: [PATCH 52/52] Add debugging helper to dump string, vector and thread id (#21224) ### Description Add some macro to help print data to console for debugging purpose. Example usage: ``` int input_id; vector some_vector; DUMP_CPU_TENSOR_INIT() DUMP_CPU_TENSOR("some vector", some_vector); DUMP_STRING("input_id=", input_id); ``` - To enable dump thread id, set environment variable `ORT_DUMP_THREAD_ID=0`. - User can disable dumping by environment variable `ORT_ENABLE_CPU_DUMP=0`. ### Motivation and Context --- .../contrib_ops/cpu/utils/console_dumper.h | 2 + .../contrib_ops/cpu/utils/debug_macros.h | 3 ++ .../contrib_ops/cpu/utils/dump_tensor.cc | 52 ++++++++++++++++++- .../contrib_ops/cpu/utils/dump_tensor.h | 11 +++- .../cuda/utils/dump_cuda_tensor.cc | 8 +++ .../contrib_ops/cuda/utils/dump_cuda_tensor.h | 2 + 6 files changed, 75 insertions(+), 3 deletions(-) diff --git a/onnxruntime/contrib_ops/cpu/utils/console_dumper.h b/onnxruntime/contrib_ops/cpu/utils/console_dumper.h index 3c255879df199..2782a59d4326d 100644 --- a/onnxruntime/contrib_ops/cpu/utils/console_dumper.h +++ b/onnxruntime/contrib_ops/cpu/utils/console_dumper.h @@ -37,6 +37,8 @@ class IConsoleDumper { virtual void Print(const char* name, int index, bool end_line) const = 0; virtual void Print(const char* name, const std::string& value, bool end_line) const = 0; + virtual void Print(const std::string& value) const = 0; + protected: bool is_enabled_; }; diff --git a/onnxruntime/contrib_ops/cpu/utils/debug_macros.h b/onnxruntime/contrib_ops/cpu/utils/debug_macros.h index 37a9b0160ade9..d5cbaa0a3e6b7 100644 --- a/onnxruntime/contrib_ops/cpu/utils/debug_macros.h +++ b/onnxruntime/contrib_ops/cpu/utils/debug_macros.h @@ -1,4 +1,5 @@ #pragma once +#include "core/common/make_string.h" // #define DEBUG_GENERATION 1 // uncomment it for debugging generation (like beam search etc) @@ -14,9 +15,11 @@ #if DUMP_CPU_TENSOR_LEVEL > 0 #define DUMP_CPU_TENSOR_INIT() onnxruntime::contrib::CpuTensorConsoleDumper cpu_dumper #define DUMP_CPU_TENSOR(...) cpu_dumper.Print(__VA_ARGS__) +#define DUMP_STRING(...) cpu_dumper.Print(::onnxruntime::MakeString(__VA_ARGS__)) #else #define DUMP_CPU_TENSOR_INIT() #define DUMP_CPU_TENSOR(...) +#define DUMP_STRING(...) #endif #if DUMP_CPU_TENSOR_LEVEL > 1 diff --git a/onnxruntime/contrib_ops/cpu/utils/dump_tensor.cc b/onnxruntime/contrib_ops/cpu/utils/dump_tensor.cc index 3a5deef35d6d6..87a9cd3965763 100644 --- a/onnxruntime/contrib_ops/cpu/utils/dump_tensor.cc +++ b/onnxruntime/contrib_ops/cpu/utils/dump_tensor.cc @@ -1,18 +1,38 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -#include #include "contrib_ops/cpu/utils/dump_tensor.h" +#include +#include +#include +#include #include "core/framework/print_tensor_utils.h" #include "contrib_ops/cpu/utils/debug_macros.h" +#include "core/platform/env_var_utils.h" namespace onnxruntime { namespace contrib { #if DUMP_CPU_TENSOR_LEVEL > 0 +// Environment variable to enable/disable dumping +constexpr const char* kEnableCpuTensorDumper = "ORT_ENABLE_CPU_DUMP"; + +// Environment variable to enable/disable dumping thread id +constexpr const char* kDumpThreadId = "ORT_DUMP_THREAD_ID"; + +// To avoid dumping at the same time from multiple threads +static std::mutex s_mutex; + +static bool s_output_thread_id = false; + template void DumpCpuTensor(const char* name, const T* tensor, int dim0, int dim1) { + std::unique_lock lock(s_mutex); + + if (s_output_thread_id) + std::cout << "Thread ID:" << std::this_thread::get_id() << std::endl; + if (nullptr != name) { std::cout << std::string(name) << std::endl; } @@ -26,6 +46,11 @@ void DumpCpuTensor(const char* name, const T* tensor, int dim0, int dim1) { template void DumpCpuTensor(const char* name, const T* tensor, int dim0, int dim1, int dim2) { + std::unique_lock lock(s_mutex); + + if (s_output_thread_id) + std::cout << "Thread ID:" << std::this_thread::get_id() << std::endl; + if (nullptr != name) { std::cout << std::string(name) << std::endl; } @@ -93,6 +118,21 @@ void DumpCpuTensor(const char* name, const Tensor& tensor) { DumpCpuTensor(nullptr, tensor, static_cast(num_rows), static_cast(row_size)); } +CpuTensorConsoleDumper::CpuTensorConsoleDumper() { + is_enabled_ = ParseEnvironmentVariableWithDefault(kEnableCpuTensorDumper, 1) != 0; + s_output_thread_id = ParseEnvironmentVariableWithDefault(kDumpThreadId, 0) != 0; +} + +void CpuTensorConsoleDumper::Print(const std::string& value) const { + if (!is_enabled_) + return; + + std::unique_lock lock(s_mutex); + if (s_output_thread_id) + std::cout << "Thread ID:" << std::this_thread::get_id() << std::endl; + std::cout << value << std::endl; +} + void CpuTensorConsoleDumper::Print(const char* name, const float* tensor, int dim0, int dim1) const { if (!is_enabled_) return; @@ -185,6 +225,8 @@ void CpuTensorConsoleDumper::Print(const char* name, const OrtValue& value) cons void CpuTensorConsoleDumper::Print(const char* name, int index, bool end_line) const { if (!is_enabled_) return; + + std::unique_lock lock(s_mutex); std::cout << std::string(name) << "[" << index << "]"; if (end_line) { @@ -196,6 +238,7 @@ void CpuTensorConsoleDumper::Print(const char* name, const std::string& value, b if (!is_enabled_) return; + std::unique_lock lock(s_mutex); std::cout << std::string(name) << "=" << value; if (end_line) { @@ -204,6 +247,12 @@ void CpuTensorConsoleDumper::Print(const char* name, const std::string& value, b } #else +CpuTensorConsoleDumper::CpuTensorConsoleDumper() { +} + +void CpuTensorConsoleDumper::Print(const std::string&) const { +} + void CpuTensorConsoleDumper::Print(const char*, const float*, int, int) const { } @@ -254,7 +303,6 @@ void CpuTensorConsoleDumper::Print(const char*, int, bool) const { void CpuTensorConsoleDumper::Print(const char*, const std::string&, bool) const { } - #endif } // namespace contrib diff --git a/onnxruntime/contrib_ops/cpu/utils/dump_tensor.h b/onnxruntime/contrib_ops/cpu/utils/dump_tensor.h index d902806fd0d18..f102eae6ec709 100644 --- a/onnxruntime/contrib_ops/cpu/utils/dump_tensor.h +++ b/onnxruntime/contrib_ops/cpu/utils/dump_tensor.h @@ -2,6 +2,7 @@ // Licensed under the MIT License. #pragma once +#include #include #include "core/framework/ort_value.h" #include "contrib_ops/cpu/utils/console_dumper.h" @@ -11,7 +12,7 @@ namespace contrib { class CpuTensorConsoleDumper : public IConsoleDumper { public: - CpuTensorConsoleDumper() = default; + CpuTensorConsoleDumper(); virtual ~CpuTensorConsoleDumper() {} void Print(const char* name, const float* tensor, int dim0, int dim1) const override; void Print(const char* name, const MLFloat16* tensor, int dim0, int dim1) const override; @@ -33,6 +34,14 @@ class CpuTensorConsoleDumper : public IConsoleDumper { void Print(const char* name, const OrtValue& value) const override; void Print(const char* name, int index, bool end_line) const override; void Print(const char* name, const std::string& value, bool end_line) const override; + + void Print(const std::string& value) const override; + + // Output a vector with a threshold for max number of elements to output. Default threshold 0 means no limit. + template + void Print(const char* name, const std::vector& vec, size_t max_count = 0) const { + this->Print(name, vec.data(), 1, static_cast(std::min(max_count, vec.size()))); + } }; } // namespace contrib diff --git a/onnxruntime/contrib_ops/cuda/utils/dump_cuda_tensor.cc b/onnxruntime/contrib_ops/cuda/utils/dump_cuda_tensor.cc index fb7af3cfdd54f..e10c2ec63fd51 100644 --- a/onnxruntime/contrib_ops/cuda/utils/dump_cuda_tensor.cc +++ b/onnxruntime/contrib_ops/cuda/utils/dump_cuda_tensor.cc @@ -202,6 +202,10 @@ void DumpGpuTensor(const char* name, const Tensor& tensor) { DumpGpuTensor(nullptr, tensor, static_cast(num_rows), static_cast(row_size)); } +void CudaTensorConsoleDumper::Print(const std::string& value) const { + std::cout << value << std::endl; +} + void CudaTensorConsoleDumper::Print(const char* name, const size_t* tensor, int dim0, int dim1) const { if (is_enabled_) DumpGpuTensor(name, tensor, dim0, dim1, true); @@ -325,6 +329,10 @@ void CudaTensorConsoleDumper::Print(const char* name, const std::string& value, } #else + +void CudaTensorConsoleDumper::Print(const std::string&) const { +} + void CudaTensorConsoleDumper::Print(const char*, const size_t*, int, int) const { } diff --git a/onnxruntime/contrib_ops/cuda/utils/dump_cuda_tensor.h b/onnxruntime/contrib_ops/cuda/utils/dump_cuda_tensor.h index 0f25e85bb97d7..6ad0ad9a67b75 100644 --- a/onnxruntime/contrib_ops/cuda/utils/dump_cuda_tensor.h +++ b/onnxruntime/contrib_ops/cuda/utils/dump_cuda_tensor.h @@ -46,6 +46,8 @@ class CudaTensorConsoleDumper : public onnxruntime::contrib::IConsoleDumper { void Print(const char* name, const OrtValue& value) const override; void Print(const char* name, int index, bool end_line) const override; void Print(const char* name, const std::string& value, bool end_line) const override; + + void Print(const std::string& value) const override; }; } // namespace cuda