From ba8a00f62f96ce93eeae2216353a33832b610742 Mon Sep 17 00:00:00 2001 From: Valery Chernov Date: Wed, 8 Feb 2023 22:02:20 +0400 Subject: [PATCH] [TVM EP] Support zero copying TVM EP output tensor to ONNX Runtime output tensor (#12593) **Description**: Support new feature of TVM Virtual Machine (method `set_outputs`) on TVM Execution Provider side. It allows to avoid excess copying from TVM EP output tensor to ONNX Runtime one **Motivation and Context** Tests with multiple output topologies and big output tensors shows that there is overheads spent on copying from TVM EP to ONNX Runtime. Returning output(s) on preallocated memory for VirtualMachine was implemented on TVM side. **Details** `set_output_zero_copy` provider option for TVM EP switches on/off this feature. It is true by default. The feature works for both GraphExecutor and VirtualMachine from TVM. --------- Co-authored-by: Valery Chernov --- cmake/external/tvm.cmake | 2 +- onnxruntime/core/providers/tvm/tvm_api.cc | 49 ++++++++++----- onnxruntime/core/providers/tvm/tvm_api.h | 2 + .../core/providers/tvm/tvm_ep_options.cc | 4 ++ .../core/providers/tvm/tvm_ep_options.h | 1 + onnxruntime/core/providers/tvm/tvm_runner.cc | 5 +- .../core/providers/tvm/tvm_runner_impl.cc | 47 +++++++++++---- .../core/providers/tvm/tvm_runner_impl.h | 60 ++++++++++++++----- 8 files changed, 123 insertions(+), 47 deletions(-) diff --git a/cmake/external/tvm.cmake b/cmake/external/tvm.cmake index 1e224a2dad..93049c8b85 100644 --- a/cmake/external/tvm.cmake +++ b/cmake/external/tvm.cmake @@ -21,4 +21,4 @@ if (onnxruntime_USE_TVM) set(tvm_INCLUDE_DIRS ${tvm_SOURCE_DIR}/include) -endif() \ No newline at end of file +endif() diff --git a/onnxruntime/core/providers/tvm/tvm_api.cc b/onnxruntime/core/providers/tvm/tvm_api.cc index 2e841f38ff..fc46d0a90e 100644 --- a/onnxruntime/core/providers/tvm/tvm_api.cc +++ b/onnxruntime/core/providers/tvm/tvm_api.cc @@ -33,8 +33,7 @@ TvmModule TVMCompile(const TvmEPOptions& options, const std::string& onnx_txt, const std::string& model_path, int opset, - const TVMTensorShapes& input_shapes) -{ + const TVMTensorShapes& input_shapes) { ::tvm::Array shapes; for (size_t i = 0; i < input_shapes.size(); ++i) { @@ -203,8 +202,7 @@ TvmModule TVMSoCompile(const TvmEPOptions& options) { void TVMSetInputs(TvmModule& mod, std::vector& inds, - std::vector& inputs) -{ + std::vector& inputs) { TvmPackedFunc set_input = mod.GetFunction("set_input", false); TvmPackedFunc set_input_zero_copy = mod.GetFunction("set_input_zero_copy", false); for (size_t i = 0; i < inds.size(); ++i) { @@ -218,8 +216,7 @@ void TVMSetInputs(TvmModule& mod, void TVM_VM_SetInputs(TvmModule& mod, std::vector& inds, - std::vector& inputs) -{ + std::vector& inputs) { size_t num_total_args = inputs.size() + 1; std::vector tvm_values(num_total_args); std::vector tvm_type_codes(num_total_args); @@ -235,9 +232,33 @@ void TVM_VM_SetInputs(TvmModule& mod, set_input.CallPacked(::tvm::runtime::TVMArgs(tvm_values.data(), tvm_type_codes.data(), int(num_total_args)), &rv); } +void TVMSetOutputsZeroCopy(TvmModule& mod, + std::vector& outputs) { + TvmPackedFunc set_output = mod.GetFunction("set_output_zero_copy", false); + for (size_t i = 0; i < outputs.size(); ++i) { + set_output(i, &outputs[i]); + } +} + +void TVM_VM_SetOutputsZeroCopy(TvmModule& mod, + std::vector& outputs) { + size_t num_total_args = outputs.size() + 1; + std::vector tvm_values(num_total_args); + std::vector tvm_type_codes(num_total_args); + tvm_rt::TVMArgsSetter setter(tvm_values.data(), tvm_type_codes.data()); + const std::string func_name = "main"; + setter(0, func_name.c_str()); + for (size_t k = 0; k < num_total_args - 1; ++k) { + setter(k+1, &outputs[k]); + } + + TvmPackedFunc set_output = mod.GetFunction("set_outputs", false); + tvm_rt::TVMRetValue rv; + set_output.CallPacked(tvm_rt::TVMArgs(tvm_values.data(), tvm_type_codes.data(), num_total_args), &rv); +} + void TVMGetOutputs(TvmModule& mod, - std::vector& outputs) -{ + std::vector& outputs) { TvmPackedFunc get_output = mod.GetFunction("get_output", false); for (size_t i = 0; i < outputs.size(); ++i) { get_output(i, &outputs[i]); @@ -245,8 +266,7 @@ void TVMGetOutputs(TvmModule& mod, } void TVM_VM_GetOutputs(TvmModule& mod, - std::vector& outputs) -{ + std::vector& outputs) { TvmPackedFunc get_output = mod.GetFunction("get_output", false); for (size_t i = 0; i < outputs.size(); ++i) { // TODO(vvchernov): think about improvement of memory management @@ -256,8 +276,7 @@ void TVM_VM_GetOutputs(TvmModule& mod, } void TVMGetOutputShapes(TvmModule& mod, - TVMTensorShapes& output_shapes) -{ + TVMTensorShapes& output_shapes) { size_t size = output_shapes.size(); TvmPackedFunc get_output = mod.GetFunction("get_output", false); for (size_t i = 0; i < size; ++i) { @@ -272,15 +291,13 @@ void TVMGetOutputShapes(TvmModule& mod, } } -void TVMRun(TvmModule& mod) -{ +void TVMRun(TvmModule& mod) { TvmPackedFunc run = mod.GetFunction("run", false); ORT_ENFORCE(run != nullptr, "Unable to retrieve graph executor run."); run(); } -void TVM_VM_Run(TvmModule& mod) -{ +void TVM_VM_Run(TvmModule& mod) { TvmPackedFunc run = mod.GetFunction("invoke", false); ORT_ENFORCE(run != nullptr, "Unable to retrieve virtual machine invoke."); run("main"); diff --git a/onnxruntime/core/providers/tvm/tvm_api.h b/onnxruntime/core/providers/tvm/tvm_api.h index 810334231a..a245b6398d 100644 --- a/onnxruntime/core/providers/tvm/tvm_api.h +++ b/onnxruntime/core/providers/tvm/tvm_api.h @@ -24,6 +24,8 @@ namespace tvm { void TVMSetInputs(TvmModule& mod, std::vector& inds, std::vector& inputs); void TVM_VM_SetInputs(TvmModule& mod, std::vector& inds, std::vector& inputs); + void TVMSetOutputsZeroCopy(TvmModule& mod, std::vector& outputs); + void TVM_VM_SetOutputsZeroCopy(TvmModule& mod, std::vector& outputs); void TVMGetOutputs(TvmModule& mod, std::vector& outputs); void TVM_VM_GetOutputs(TvmModule& mod, std::vector& outputs); void TVMGetOutputShapes(TvmModule& mod, diff --git a/onnxruntime/core/providers/tvm/tvm_ep_options.cc b/onnxruntime/core/providers/tvm/tvm_ep_options.cc index 1d701dc11d..91ed836b60 100644 --- a/onnxruntime/core/providers/tvm/tvm_ep_options.cc +++ b/onnxruntime/core/providers/tvm/tvm_ep_options.cc @@ -23,6 +23,7 @@ constexpr const char* kTarget = "target"; constexpr const char* kTargetHost = "target_host"; constexpr const char* kOptLevel = "opt_level"; constexpr const char* kFreezeWeights = "freeze_weights"; +constexpr const char* kSetOutputZeroCopy = "set_output_zero_copy"; constexpr const char* kToNHWC = "to_nhwc"; constexpr const char* kTuningFilePath = "tuning_file_path"; constexpr const char* kTuningType = "tuning_type"; @@ -38,6 +39,7 @@ static const std::unordered_set valid_keys { std::string{kTargetHost}, std::string{kOptLevel}, std::string{kFreezeWeights}, + std::string{kSetOutputZeroCopy}, std::string{kToNHWC}, std::string{kTuningFilePath}, std::string{kTuningType}, @@ -124,6 +126,7 @@ TvmEPOptions TvmEPOptionsHelper::FromProviderOptions(const ProviderOptions& pr_o .AddAssignmentToReference(tvm::provider_option_names::kTargetHost, options.target_host) .AddAssignmentToReference(tvm::provider_option_names::kOptLevel, options.opt_level) .AddAssignmentToReference(tvm::provider_option_names::kFreezeWeights, options.freeze_weights) + .AddAssignmentToReference(tvm::provider_option_names::kSetOutputZeroCopy, options.set_output_zero_copy) .AddAssignmentToReference(tvm::provider_option_names::kToNHWC, options.to_nhwc) .AddAssignmentToReference(tvm::provider_option_names::kTuningFilePath, options.tuning_file_path) .AddAssignmentToReference(tvm::provider_option_names::kTuningType, options.tuning_type) @@ -261,6 +264,7 @@ std::ostream& operator<<(std::ostream& out, const TvmEPOptions& options) { "target_host: " << options.target_host << "\n" << "opt level: " << options.opt_level << "\n" << "freeze weights: " << options.freeze_weights << "\n" << + "set_output_zero_copy: " << options.set_output_zero_copy << "\n" << "tuning file path: " << options.tuning_file_path << "\n" << "tuning type: " << options.tuning_type << "\n" << "convert layout to NHWC: " << options.to_nhwc << "\n" << diff --git a/onnxruntime/core/providers/tvm/tvm_ep_options.h b/onnxruntime/core/providers/tvm/tvm_ep_options.h index b9810b2734..f8444a090b 100644 --- a/onnxruntime/core/providers/tvm/tvm_ep_options.h +++ b/onnxruntime/core/providers/tvm/tvm_ep_options.h @@ -41,6 +41,7 @@ struct TvmEPOptions { unsigned int opt_level{tvm::default_opt_level}; bool freeze_weights = true; bool to_nhwc = false; + bool set_output_zero_copy = true; std::string tuning_file_path{""}; std::string tuning_type{tvm::default_tuning_type}; std::string input_names_str{""}; diff --git a/onnxruntime/core/providers/tvm/tvm_runner.cc b/onnxruntime/core/providers/tvm/tvm_runner.cc index 8d2797668d..afeabba477 100644 --- a/onnxruntime/core/providers/tvm/tvm_runner.cc +++ b/onnxruntime/core/providers/tvm/tvm_runner.cc @@ -18,8 +18,9 @@ TVMRunner::TVMRunner(const TvmEPOptions& options, runner_ = getTVMRunnerImpl(mod, options, inputs_info, output_tensors); } -common::Status TVMRunner::operator()(FunctionState state, const OrtApi* api, OrtKernelContext* context) { - return runner_->run(api, context); +common::Status TVMRunner::operator()(FunctionState state, const OrtApi* /*api*/, OrtKernelContext* context) { + Ort::KernelContext ctx(context); + return runner_->run(ctx); } } // namespace tvm diff --git a/onnxruntime/core/providers/tvm/tvm_runner_impl.cc b/onnxruntime/core/providers/tvm/tvm_runner_impl.cc index 6b22f7eaf5..c9a79ffbd0 100644 --- a/onnxruntime/core/providers/tvm/tvm_runner_impl.cc +++ b/onnxruntime/core/providers/tvm/tvm_runner_impl.cc @@ -18,9 +18,11 @@ std::shared_ptr getTVMRunnerImpl(const std::shared_ptr& m const std::vector output_tensors) { const std::string& name = options.executor; if (name == "graph") { - return std::make_shared(mod, inputs_info, options.output_shapes, output_tensors); + return std::make_shared(mod, inputs_info, options.output_shapes, + output_tensors, options.set_output_zero_copy); } else if (name == "vm") { - return std::make_shared(mod, inputs_info, options.output_shapes, output_tensors); + return std::make_shared(mod, inputs_info, options.output_shapes, + output_tensors, options.set_output_zero_copy); } return nullptr; } @@ -30,10 +32,12 @@ std::shared_ptr getTVMRunnerImpl(const std::shared_ptr& m RunnerImpl::RunnerImpl(const std::shared_ptr& mod, const InputsInfoMap& inputs_info, const TVMTensorShapes output_shapes, - const std::vector output_tensors) : mod_(mod), - inputs_info_(inputs_info), - output_shapes_(output_shapes), - output_tensors_(output_tensors) { + const std::vector output_tensors, + bool set_output_zero_copy) : mod_(mod), + inputs_info_(inputs_info), + output_shapes_(output_shapes), + output_tensors_(output_tensors), + set_output_zero_copy_(set_output_zero_copy) { } void RunnerImpl::convert_input_tensors2dl_tensors(Ort::KernelContext& context, @@ -88,7 +92,9 @@ void RunnerImpl::add_device_type_data2output_tensors(Ort::KernelContext& context GERunnerImpl::GERunnerImpl(const std::shared_ptr& mod, const InputsInfoMap& inputs_info, const TVMTensorShapes output_shapes, - const std::vector output_tensors) : RunnerImpl(mod, inputs_info, output_shapes, output_tensors) { + const std::vector output_tensors, + bool set_output_zero_copy) : + RunnerImpl(mod, inputs_info, output_shapes, output_tensors, set_output_zero_copy) { } void GERunnerImpl::set_input(Ort::KernelContext& context) { @@ -103,8 +109,15 @@ void GERunnerImpl::connect_output_tensors2ort(Ort::KernelContext& context) { add_device_type_data2output_tensors(context); } -void GERunnerImpl::run_and_get_output() { +void GERunnerImpl::set_output_zero_copy() { + tvm::TVMSetOutputsZeroCopy(*mod_, output_tensors_); +} + +void GERunnerImpl::run() { tvm::TVMRun(*mod_); +} + +void GERunnerImpl::get_outputs() { tvm::TVMGetOutputs(*mod_, output_tensors_); } @@ -113,7 +126,9 @@ void GERunnerImpl::run_and_get_output() { VMRunnerImpl::VMRunnerImpl(const std::shared_ptr& mod, const InputsInfoMap& inputs_info, const TVMTensorShapes output_shapes, - const std::vector output_tensors) : RunnerImpl(mod, inputs_info, output_shapes, output_tensors) { + const std::vector output_tensors, + bool set_output_zero_copy) : + RunnerImpl(mod, inputs_info, output_shapes, output_tensors, set_output_zero_copy) { } void VMRunnerImpl::set_input(Ort::KernelContext& context) { @@ -125,20 +140,28 @@ void VMRunnerImpl::set_input(Ort::KernelContext& context) { } void VMRunnerImpl::connect_output_tensors2ort(Ort::KernelContext& context) { - if (!probe_infer_) { + // TODO(vvchernov): try to find more flexible solution + if(!probe_infer_) { infer_once_to_get_output_shapes(); } add_device_type_data2output_tensors(context); } -void VMRunnerImpl::run_and_get_output() { +void VMRunnerImpl::set_output_zero_copy() { + tvm::TVM_VM_SetOutputsZeroCopy(*mod_, output_tensors_); +} + +void VMRunnerImpl::run() { tvm::TVM_VM_Run(*mod_); +} + +void VMRunnerImpl::get_outputs() { tvm::TVM_VM_GetOutputs(*mod_, output_tensors_); } void VMRunnerImpl::infer_once_to_get_output_shapes() { - tvm::TVM_VM_Run(*mod_); + run(); size_t num_outputs = output_tensors_.size(); // TODO(vvchernov): check it output_shapes_.resize(num_outputs); diff --git a/onnxruntime/core/providers/tvm/tvm_runner_impl.h b/onnxruntime/core/providers/tvm/tvm_runner_impl.h index 6fcddddd8a..ff04561d44 100644 --- a/onnxruntime/core/providers/tvm/tvm_runner_impl.h +++ b/onnxruntime/core/providers/tvm/tvm_runner_impl.h @@ -24,22 +24,43 @@ public: RunnerImpl(const std::shared_ptr& mod, const InputsInfoMap& inputs_info, const TVMTensorShapes output_shapes, - const std::vector tensors_outputs); + const std::vector tensors_outputs, + bool set_output_zero_copy); virtual ~RunnerImpl() = default; - virtual common::Status run(const OrtApi* /* api */, OrtKernelContext* context) { + virtual common::Status run(Ort::KernelContext& context) { + common::Status res; + if (set_output_zero_copy_) { + res = run_without_output_copying(context); + } else { + res = run_with_output_copying(context); + } + return res; + } - Ort::KernelContext ctx{context}; - set_input(ctx); - connect_output_tensors2ort(ctx); - run_and_get_output(); + virtual common::Status run_without_output_copying(Ort::KernelContext& context) { + set_input(context); + connect_output_tensors2ort(context); + set_output_zero_copy(); + run(); return Status::OK(); } - virtual void set_input(Ort::KernelContext& ctx) = 0; + virtual common::Status run_with_output_copying(Ort::KernelContext& context) { + set_input(context); + connect_output_tensors2ort(context); + run(); + get_outputs(); + + return Status::OK(); + } + + virtual void set_input(Ort::KernelContext& context) = 0; virtual void connect_output_tensors2ort(Ort::KernelContext& context) = 0; - virtual void run_and_get_output() = 0; + virtual void set_output_zero_copy() = 0; + virtual void run() = 0; + virtual void get_outputs() = 0; protected: void convert_input_tensors2dl_tensors(Ort::KernelContext& context, @@ -52,6 +73,7 @@ protected: InputsInfoMap inputs_info_; TVMTensorShapes output_shapes_; std::vector output_tensors_; + bool set_output_zero_copy_; }; @@ -61,12 +83,15 @@ public: GERunnerImpl(const std::shared_ptr& mod, const InputsInfoMap& inputs_info, const TVMTensorShapes output_shapes, - const std::vector tensors_outputs); + const std::vector tensors_outputs, + bool set_output_zero_copy); virtual ~GERunnerImpl() = default; - void set_input(Ort::KernelContext& context) override final; - void connect_output_tensors2ort(Ort::KernelContext& context) override final; - void run_and_get_output() override final; + void set_input(Ort::KernelContext& context) final; + void connect_output_tensors2ort(Ort::KernelContext& context) final; + void set_output_zero_copy() final; + void run() final; + void get_outputs() final; }; @@ -76,12 +101,15 @@ public: VMRunnerImpl(const std::shared_ptr& mod, const InputsInfoMap& inputs_info, const TVMTensorShapes output_shapes, - const std::vector tensors_outputs); + const std::vector tensors_outputs, + bool set_output_zero_copy); virtual ~VMRunnerImpl() = default; - void set_input(Ort::KernelContext& context) override final; - void connect_output_tensors2ort(Ort::KernelContext& context) override final; - void run_and_get_output() override final; + void set_input(Ort::KernelContext& context) final; + void connect_output_tensors2ort(Ort::KernelContext& context) final; + void set_output_zero_copy() final; + void run() final; + void get_outputs() final; private: void infer_once_to_get_output_shapes();