use GetMayInplace in CreateKernelCreateInfo

microsoft · Mar 22, 2024 · 34dc5ec · 34dc5ec
1 parent 7e84ba0
commit 34dc5ec
Showing 1 changed file with 11 additions and 0 deletions.
diff --git 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);