From 34dc5ec48f791d93aed383a8c553a0968b5d25c1 Mon Sep 17 00:00:00 2001
From: jslhcl <jslhcl@gmail.com>
Date: Fri, 22 Mar 2024 14:30:47 -0700
Subject: [PATCH] use GetMayInplace in CreateKernelCreateInfo

---
 onnxruntime/core/session/custom_ops.cc | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/onnxruntime/core/session/custom_ops.cc b/onnxruntime/core/session/custom_ops.cc
index 513aafcdadb7d..64a29da9d2e1e 100644
--- a/onnxruntime/core/session/custom_ops.cc
+++ b/onnxruntime/core/session/custom_ops.cc
@@ -865,6 +865,17 @@ KernelCreateInfo CreateKernelCreateInfo(const std::string& domain, const OrtCust
     def_builder.Provider(onnxruntime::kCpuExecutionProvider);
   }
 
+  if (op->version >= 18) {
+    int* input_index = nullptr;
+    int* output_index = nullptr;
+    size_t len = op->GetMayInplace(&input_index, &output_index);
+    if (len > 0) {
+      for (int i = 0; i < len; i++) def_builder.MayInplace(input_index[i], output_index[i]);
+      free(input_index);
+      free(output_index);
+    }
+  }
+
   KernelCreateFn kernel_create_fn = [op](FuncManager&, const OpKernelInfo& info,
                                          std::unique_ptr<OpKernel>& out) -> Status {
     out = std::make_unique<CustomOpKernel>(info, *op);