use OrtCustomOp's new API GetMayInplace in CreateKernelCreateInfo (#20037)

### Description
<!-- Describe your changes. -->
use OrtCustomOp's new API GetMayInplace in CreateKernelCreateInfo to
hook the inplace map of custom ops


### Motivation and Context
<!-- - Why is this change required? What problem does it solve?
- If it fixes an open issue, please link to the issue here. -->
This PR is to use OrtCustomOp's new API GetMayInplace in
CreateKernelCreateInfo to hook the inplace map of custom ops

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);
 };
 
 /*

include/onnxruntime/core/session/onnxruntime_cxx_api.h

@@ -2301,6 +2301,9 @@ struct CustomOpBase : OrtCustomOp {
     OrtCustomOp::GetEndVersion = [](const OrtCustomOp* this_) {
       return static_cast<const TOp*>(this_)->end_ver_;
     };
+
+    OrtCustomOp::GetMayInplace = nullptr;
+    OrtCustomOp::ReleaseMayInplace = nullptr;
   }
 
   // Default implementation of GetExecutionProviderType that returns nullptr to default to the CPU provider

include/onnxruntime/core/session/onnxruntime_lite_custom_op.h

@@ -862,6 +862,9 @@ struct OrtLiteCustomOp : public OrtCustomOp {
       auto self = reinterpret_cast<const OrtLiteCustomOp*>(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
+}  // namespace Ort

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;
 };
 

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<OpKernel>& out) -> Status {
     out = std::make_unique<CustomOpKernel>(info, *op);

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<size_t>(2));
-  free(input_index);
-  free(output_index);
+  mock_gqa.ReleaseMayInplace(input_index, output_index);
 }