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();