[objc] Add session options register custom ops with function pointer API (#16603)

objectivec/include/onnxruntime.h

@@ -5,8 +5,9 @@
 // the headers below can also be imported individually
 
 #import "ort_coreml_execution_provider.h"
-#import "ort_xnnpack_execution_provider.h"
+#import "ort_custom_op_registration.h"
 #import "ort_enums.h"
 #import "ort_env.h"
 #import "ort_session.h"
 #import "ort_value.h"
+#import "ort_xnnpack_execution_provider.h"

objectivec/include/ort_custom_op_registration.h

@@ -0,0 +1,23 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+/** C API type forward declaration. */
+struct OrtStatus;
+
+/** C API type forward declaration. */
+struct OrtApiBase;
+
+/** C API type forward declaration. */
+struct OrtSessionOptions;
+
+/**
+ * Pointer to a custom op registration function that uses the ONNX Runtime C API.
+ *
+ * The signature is defined in the ONNX Runtime C API:
+ * https://github.com/microsoft/onnxruntime/blob/67f4cd54fab321d83e4a75a40efeee95a6a17079/include/onnxruntime/core/session/onnxruntime_c_api.h#L697
+ *
+ * This is a low-level type intended for interoperating with libraries which provide such a function for custom op
+ * registration, such as [ONNX Runtime Extensions](https://github.com/microsoft/onnxruntime-extensions).
+ */
+typedef struct OrtStatus* (*ORTCAPIRegisterCustomOpsFnPtr)(struct OrtSessionOptions* /*options*/,
+                                                           const struct OrtApiBase* /*api*/);

objectivec/include/ort_session.h

@@ -3,6 +3,7 @@
 
 #import <Foundation/Foundation.h>
 
+#import "ort_custom_op_registration.h"
 #import "ort_enums.h"
 
 NS_ASSUME_NONNULL_BEGIN
@@ -196,12 +197,19 @@ NS_ASSUME_NONNULL_BEGIN
  * Available since 1.14.
  *
  * The registration function must have the signature:
- *    OrtStatus* (*fn)(OrtSessionOptions* options, const OrtApiBase* api);
+ *    `OrtStatus* (*fn)(OrtSessionOptions* options, const OrtApiBase* api);`
+ *
+ * The signature is defined in the ONNX Runtime C API:
+ * https://github.com/microsoft/onnxruntime/blob/67f4cd54fab321d83e4a75a40efeee95a6a17079/include/onnxruntime/core/session/onnxruntime_c_api.h#L697
  *
  * See https://onnxruntime.ai/docs/reference/operators/add-custom-op.html for more information on custom ops.
  * See https://github.com/microsoft/onnxruntime/blob/342a5bf2b756d1a1fc6fdc582cfeac15182632fe/onnxruntime/test/testdata/custom_op_library/custom_op_library.cc#L115
  * for an example of a custom op library registration function.
  *
+ * @note The caller must ensure that `registrationFuncName` names a valid function that is visible to the native ONNX
+ * Runtime code and has the correct signature.
+ * They must ensure that the function does what they expect it to do because this method will just call it.
+ *
  * @param registrationFuncName The name of the registration function to call.
  * @param error Optional error information set if an error occurs.
  * @return Whether the registration function was successfully called.
@@ -209,6 +217,32 @@ NS_ASSUME_NONNULL_BEGIN
 - (BOOL)registerCustomOpsUsingFunction:(NSString*)registrationFuncName
                                  error:(NSError**)error;
 
+/**
+ * Registers custom ops for use with `ORTSession`s using this SessionOptions by calling the specified function
+ * pointed to by `registerCustomOpsFn`.
+ *
+ * Available since 1.16.
+ *
+ * The registration function must have the signature:
+ *    `OrtStatus* (*fn)(OrtSessionOptions* options, const OrtApiBase* api);`
+ *
+ * The signature is defined in the ONNX Runtime C API:
+ * https://github.com/microsoft/onnxruntime/blob/67f4cd54fab321d83e4a75a40efeee95a6a17079/include/onnxruntime/core/session/onnxruntime_c_api.h#L697
+ *
+ * See https://onnxruntime.ai/docs/reference/operators/add-custom-op.html for more information on custom ops.
+ * See https://github.com/microsoft/onnxruntime/blob/342a5bf2b756d1a1fc6fdc582cfeac15182632fe/onnxruntime/test/testdata/custom_op_library/custom_op_library.cc#L115
+ * for an example of a custom op library registration function.
+ *
+ * @note The caller must ensure that `registerCustomOpsFn` is a valid function pointer and has the correct signature.
+ * They must ensure that the function does what they expect it to do because this method will just call it.
+ *
+ * @param registerCustomOpsFn A pointer to the registration function to call.
+ * @param error Optional error information set if an error occurs.
+ * @return Whether the registration function was successfully called.
+ */
+- (BOOL)registerCustomOpsUsingFunctionPointer:(ORTCAPIRegisterCustomOpsFnPtr)registerCustomOpsFn
+                                        error:(NSError**)error;
+
 @end
 
 /**

objectivec/ort_session.mm

@@ -309,6 +309,19 @@ NS_ASSUME_NONNULL_BEGIN
   ORT_OBJC_API_IMPL_CATCH_RETURNING_BOOL(error)
 }
 
+- (BOOL)registerCustomOpsUsingFunctionPointer:(ORTCAPIRegisterCustomOpsFnPtr)registerCustomOpsFn
+                                        error:(NSError**)error {
+  try {
+    if (!registerCustomOpsFn) {
+      ORT_CXX_API_THROW("registerCustomOpsFn must not be null", ORT_INVALID_ARGUMENT);
+    }
+    Ort::ThrowOnError((*registerCustomOpsFn)(static_cast<OrtSessionOptions*>(*_sessionOptions),
+                                             OrtGetApiBase()));
+    return YES;
+  }
+  ORT_OBJC_API_IMPL_CATCH_RETURNING_BOOL(error)
+}
+
 #pragma mark - Internal
 
 - (Ort::SessionOptions&)CXXAPIOrtSessionOptions {

objectivec/test/ort_session_test.mm

@@ -238,6 +238,27 @@ NS_ASSUME_NONNULL_BEGIN
                                                   error:&err];
   ORTAssertNullableResultSuccessful(session, err);
 }
+
+static bool gDummyRegisterCustomOpsFnCalled = false;
+
+static OrtStatus* _Nullable DummyRegisterCustomOpsFn(OrtSessionOptions* /*session_options*/,
+                                                     const OrtApiBase* /*api*/) {
+  gDummyRegisterCustomOpsFnCalled = true;
+  return nullptr;
+}
+
+- (void)testRegisterCustomOpsUsingFunctionPointer {
+  NSError* err = nil;
+  ORTSessionOptions* sessionOptions = [ORTSessionTest makeSessionOptions];
+
+  gDummyRegisterCustomOpsFnCalled = false;
+  BOOL registerResult = [sessionOptions registerCustomOpsUsingFunctionPointer:&DummyRegisterCustomOpsFn
+                                                                        error:&err];
+  ORTAssertBoolResultSuccessful(registerResult, err);
+
+  XCTAssertEqual(gDummyRegisterCustomOpsFnCalled, true);
+}
+
 @end
 
 NS_ASSUME_NONNULL_END