diff --git a/include/onnxruntime/core/session/onnxruntime_c_api.h b/include/onnxruntime/core/session/onnxruntime_c_api.h index 2a7b17977f..c4fb0d3a83 100644 --- a/include/onnxruntime/core/session/onnxruntime_c_api.h +++ b/include/onnxruntime/core/session/onnxruntime_c_api.h @@ -4732,8 +4732,12 @@ struct OrtCustomOp { // Callers will provide 2 raw int* and pass in their address, this function will fill these 2 arrays // when return, output (*output_index)[i] may reuse the input (*input_index[i]). // The return value is the size of these 2 arrays. - // Callers are responsible to delete these 2 arrays after use. + // Callers are responsible to delete these 2 arrays after use by calling OrtCustomOp::ReleaseMayInplace(). size_t(ORT_API_CALL* GetMayInplace)(_Out_ int** input_index, _Out_ int** output_index); + + // Release the pointer input_index and output_index allocated from GetMayInplace() function. + // If GetMayInplace() is defined, this function MUST be defined as well. + void(ORT_API_CALL* ReleaseMayInplace)(_Frees_ptr_opt_ int* input_index, _Frees_ptr_opt_ int* output_index); }; /* diff --git a/include/onnxruntime/core/session/onnxruntime_cxx_api.h b/include/onnxruntime/core/session/onnxruntime_cxx_api.h index 60540514fb..5f2e0a470a 100644 --- a/include/onnxruntime/core/session/onnxruntime_cxx_api.h +++ b/include/onnxruntime/core/session/onnxruntime_cxx_api.h @@ -2301,6 +2301,9 @@ struct CustomOpBase : OrtCustomOp { OrtCustomOp::GetEndVersion = [](const OrtCustomOp* this_) { return static_cast(this_)->end_ver_; }; + + OrtCustomOp::GetMayInplace = nullptr; + OrtCustomOp::ReleaseMayInplace = nullptr; } // Default implementation of GetExecutionProviderType that returns nullptr to default to the CPU provider diff --git a/include/onnxruntime/core/session/onnxruntime_lite_custom_op.h b/include/onnxruntime/core/session/onnxruntime_lite_custom_op.h index 0c0af16d4e..896893e986 100644 --- a/include/onnxruntime/core/session/onnxruntime_lite_custom_op.h +++ b/include/onnxruntime/core/session/onnxruntime_lite_custom_op.h @@ -862,6 +862,9 @@ struct OrtLiteCustomOp : public OrtCustomOp { auto self = reinterpret_cast(op); return self->end_ver_; }; + + OrtCustomOp::GetMayInplace = {}; + OrtCustomOp::ReleaseMayInplace = {}; } const std::string op_name_; @@ -1111,4 +1114,4 @@ OrtLiteCustomOp* CreateLiteCustomOp(const char* op_name, } } // namespace Custom -} // namespace Ort \ No newline at end of file +} // namespace Ort diff --git a/onnxruntime/core/framework/feeds_fetches_manager.h b/onnxruntime/core/framework/feeds_fetches_manager.h index 31c00d65ce..c2c1be64f3 100644 --- a/onnxruntime/core/framework/feeds_fetches_manager.h +++ b/onnxruntime/core/framework/feeds_fetches_manager.h @@ -73,6 +73,8 @@ struct FeedsFetchesInfo { struct MLValueCopyInfo { OrtDevice source_device{}; OrtDevice target_device{}; // default is CPU + + // if all the consume ops are from the same stream, this variable is the stream index; otherwise -1 int unique_stream_index_consumes_it = -1; }; diff --git a/onnxruntime/core/session/custom_ops.cc b/onnxruntime/core/session/custom_ops.cc index 513aafcdad..cc23d0822c 100644 --- a/onnxruntime/core/session/custom_ops.cc +++ b/onnxruntime/core/session/custom_ops.cc @@ -865,6 +865,16 @@ KernelCreateInfo CreateKernelCreateInfo(const std::string& domain, const OrtCust def_builder.Provider(onnxruntime::kCpuExecutionProvider); } + if (op->version >= 18 && op->GetMayInplace != nullptr) { + int* input_index = nullptr; + int* output_index = nullptr; + size_t len = op->GetMayInplace(&input_index, &output_index); + if (len > 0) { + for (size_t i = 0; i < len; i++) def_builder.MayInplace(input_index[i], output_index[i]); + op->ReleaseMayInplace(input_index, output_index); + } + } + KernelCreateFn kernel_create_fn = [op](FuncManager&, const OpKernelInfo& info, std::unique_ptr& out) -> Status { out = std::make_unique(info, *op); diff --git a/onnxruntime/test/shared_lib/test_inference.cc b/onnxruntime/test/shared_lib/test_inference.cc index 68f5d866a9..5dd5fabd26 100644 --- a/onnxruntime/test/shared_lib/test_inference.cc +++ b/onnxruntime/test/shared_lib/test_inference.cc @@ -4035,6 +4035,10 @@ struct MockGQA : public OrtCustomOp { (*output_index)[1] = 2; return ret; }; + OrtCustomOp::ReleaseMayInplace = [](int* input_index, int* output_index) { + free(input_index); + free(output_index); + }; } }; @@ -4050,6 +4054,5 @@ TEST(CApiTest, OrtCustomOp_GetInPlace) { ASSERT_EQ(output_index[0], 1); ASSERT_EQ(output_index[1], 2); ASSERT_EQ(len, static_cast(2)); - free(input_index); - free(output_index); + mock_gqa.ReleaseMayInplace(input_index, output_index); }