Drop _loadedBackendOptions ivar from Apple bindings (pytorch#19680)

Differential Revision: D105100368

Pull Request resolved: pytorch#19680

Author: metascroy

extension/apple/ExecuTorch/Exported/ExecuTorchModule.mm

@@ -252,27 +252,6 @@ @implementation ExecuTorchModule {
   std::unique_ptr<Module> _module;
   NSMutableDictionary<NSString *, NSMutableArray<ExecuTorchValue *> *> *_inputs;
   NSMutableDictionary<NSString *, NSMutableArray<ExecuTorchValue *> *> *_outputs;
-  // Strong reference to the most recently passed BackendOptionsMap. The
-  // C++ Module borrows a pointer into the map's underlying C++ storage and
-  // dereferences it during lazy load_method calls (triggered by forward),
-  // so the ObjC wrapper must keep it alive. ARC handles the lifetime.
-  //
-  // INVARIANT: this ivar is only ever overwritten with another non-nil
-  // BackendOptionsMap, and never reset to nil while `_module` is alive.
-  // Resetting to nil would release the C++ map while `_module` still holds
-  // a borrowed pointer into it.
-  //
-  // THREAD SAFETY: like the rest of `ExecuTorchModule`, write access here
-  // is not thread-safe. The ARC retain/release on assignment is non-atomic
-  // for direct ivars; serialize `loadWithOptions:` calls externally if you
-  // share a `Module` across threads.
-  //
-  // TODO: remove this ivar once the C++ Module owns its LoadBackendOptionsMap
-  // by value (today it borrows a raw pointer). With owned options the ObjC
-  // wrapper has nothing to retain, the thread-safety caveat above goes
-  // away, and -loadMethod:options: / -loadWithOptions: stop needing a
-  // custom lifetime contract between the bindings and the C++ layer.
-  ExecuTorchBackendOptionsMap *_loadedBackendOptions;
 }
 
 - (instancetype)initWithFilePath:(NSString *)filePath
@@ -354,25 +333,10 @@ - (BOOL)loadWithOptions:(ExecuTorchBackendOptionsMap *)options
            verification:(ExecuTorchVerification)verification
                   error:(NSError **)error {
   NSParameterAssert(options);
-  // Retain the options object so the C++ borrowed pointer it contains stays
-  // valid for the lifetime of any methods loaded with these options.
-  // (Methods load lazily during forward(), so the borrow may outlive this
-  // call.) See ExecuTorchBackendOptionsMap.h for the lifetime contract.
-  //
-  // No rollback on failure: Module::load updates its backend_options_ raw
-  // pointer BEFORE attempting load_internal, so after a failed call the
-  // C++ side already references `options`. The ObjC retain therefore
-  // always matches what C++ points at, even on the failure path — a
-  // two-phase commit here would instead leave C++ pointing at a map the
-  // wrapper no longer retains. See:
-  // https://github.com/pytorch/executorch/blob/6412f55a54dd3ce1f4ed220a3e96ee19b8f37967/extension/module/module.cpp#L192-L197
-  //
-  // TODO: once Module::load is made transactional (i.e. it only commits
-  // `backend_options_` after load_internal succeeds), replace the
-  // unconditional assignment below with a proper two-phase commit that
-  // only overwrites _loadedBackendOptions on success. This removes the
-  // "match C++'s unconditional write" workaround documented above.
-  _loadedBackendOptions = options;
+  // Module deep-copies the LoadBackendOptionsMap on the C++ side, so we
+  // do not need to retain `options` past this call. ARC releases the
+  // wrapper when the parameter goes out of scope and the Module's owned
+  // copy keeps lazy load_method paths working.
   const auto errorCode = _module->load(*[options cppMap],
                                         static_cast<Program::Verification>(verification));
   if (errorCode != Error::Ok) {
@@ -395,22 +359,10 @@ - (BOOL)loadMethod:(NSString *)methodName
            options:(ExecuTorchBackendOptionsMap *)options
              error:(NSError **)error {
   NSParameterAssert(options);
-  // Do NOT assign to _loadedBackendOptions here. Module::load_method
-  // consumes `backend_options` synchronously within this call — it is
-  // passed through to program_->load_method and is not cached on the
-  // Module. Only Module::load(backend_options, ...) stores the pointer
-  // (via backend_options_). ARC keeps `options` alive for the call
-  // duration via the parameter, so no ivar retention is needed here.
-  // See:
-  //   load_method: https://github.com/pytorch/executorch/blob/6412f55a54dd3ce1f4ed220a3e96ee19b8f37967/extension/module/module.cpp#L353-L409
-  //   load (stores backend_options_): https://github.com/pytorch/executorch/blob/6412f55a54dd3ce1f4ed220a3e96ee19b8f37967/extension/module/module.cpp#L195
-  //
-  // Overwriting _loadedBackendOptions would release any map previously
-  // installed by -loadWithOptions:, but the C++ Module's backend_options_
-  // raw pointer would still reference that released map's storage — a
-  // use-after-free on the next lazy load_method. The XCTest
-  // testMixedLoadWithOptionsAndLoadMethodWithOptionsOnMultiMethodModel
-  // pins this invariant via a weak reference.
+  // load_method consumes `options` synchronously: the cppMap pointer is
+  // passed through to program_->load_method and is not cached on the C++
+  // Module. ARC keeps `options` alive for the duration of this call via
+  // the parameter, so no extra retention is needed here.
   const auto errorCode = _module->load_method(methodName.UTF8String,
                                                /*planned_memory=*/nullptr,
                                                /*event_tracer=*/nullptr,

extension/apple/ExecuTorch/__tests__/ModuleTest.swift

@@ -525,22 +525,23 @@ class ModuleTest: XCTestCase {
   }
 
   // Mixed sequence on a multi-method delegated model:
-  //   1. load(optionsA)            — installs optionsA; the C++ Module
-  //      stores a raw pointer into optionsA's storage and the ObjC
-  //      wrapper retains optionsA via _loadedBackendOptions.
+  //   1. load(optionsA)              — installs optionsA. The C++ Module
+  //      deep-copies optionsA at load time; the Apple binding does NOT
+  //      retain the input map. After the autoreleasepool drains, the
+  //      caller's optionsA is dropped — the C++ Module's internal copy
+  //      keeps the configuration alive.
   //   2. load("mul", options: optionsB) — loads "mul" explicitly with
-  //      optionsB, synchronously. Must NOT release optionsA (doing so
-  //      would leave _module->backend_options_ dangling).
-  //   3. forward(inputs)           — triggers a lazy load_method on
-  //      "forward" which falls back to the stored pointer (into optionsA).
+  //      optionsB; the C++ Module deep-copies optionsB.
+  //   3. forward(inputs)             — triggers a lazy load_method on
+  //      "forward" which consults the deep-copied optionsA stored
+  //      inside the C++ Module.
   //
-  // The XCTAssertNotNil(weakA) after step 2 is the deterministic check:
-  // a buggy loadMethod:options: that assigns `_loadedBackendOptions =
-  // options` releases optionsA's last strong ref there, weakA becomes
-  // nil, and the assertion fails independent of heap layout. With the
-  // correct implementation weakA stays non-nil. The forward/execute
-  // assertions additionally verify the positive path end-to-end.
-  func testMixedLoadWithOptionsAndLoadMethodWithOptionsOnMultiMethodModel() throws {
+  // The XCTAssertNil(weakA) after step 1 is the deterministic check:
+  // it proves the Apple binding does not silently retain the input map
+  // (a regression would re-introduce a strong reference). The forward
+  // assertions verify the positive path: even though the caller has
+  // dropped optionsA, the deep-copied configuration is still applied.
+  func testMultiMethodLoadAndForwardWithBackendOptions() throws {
     let modelPath = try requireFixture("add_mul_coreml", ofType: "pte")
     let module = Module(filePath: modelPath)
 
@@ -552,7 +553,9 @@ class ModuleTest: XCTestCase {
       weakA = optionsA
       try module.load(options: optionsA)
     }
-    XCTAssertNotNil(weakA, "Module must retain optionsA after load(optionsA)")
+    XCTAssertNil(weakA,
+      "Apple binding must not retain the input BackendOptionsMap — " +
+      "the C++ Module deep-copies it at load time")
 
     try autoreleasepool {
       let optionsB = try BackendOptionsMap(options: [
@@ -561,11 +564,9 @@ class ModuleTest: XCTestCase {
       try module.load("mul", options: optionsB)
     }
     XCTAssertTrue(module.isLoaded("mul"))
-    XCTAssertNotNil(weakA,
-      "load(\"mul\", options: optionsB) must not release optionsA — " +
-      "_module->backend_options_ still points into its storage")
 
-    // Lazy load_method("forward") must still see valid optionsA storage.
+    // Lazy load_method("forward") must still apply the configuration
+    // from the now-released optionsA, via the C++ Module's deep copy.
     let inputs: [Tensor<Float>] = [Tensor([2]), Tensor([3])]
     let addOuts: [Value] = try module.forward(inputs)
     XCTAssertEqual(addOuts.first?.tensor(), Tensor([Float(5)]))

extension/apple/ExecuTorch/__tests__/ObjC/ModuleTestObjC.m

@@ -15,6 +15,55 @@ @interface ModuleTestObjC : XCTestCase
 
 @implementation ModuleTestObjC
 
+// Pins the deep-copy contract: -loadWithOptions: must not retain the
+// options wrapper. The C++ Module deep-copies the LoadBackendOptionsMap
+// (and its BackendOption arrays) into Module-owned storage, so dropping
+// the ObjC wrapper after loadWithOptions: returns must be safe and a
+// subsequent forward() must continue to work using the Module's owned
+// copy. If a future refactor reverts to the borrowed-pointer design and
+// reintroduces an ivar to retain the wrapper, this test fails (the
+// weak reference stays non-nil).
+- (void)testLoadWithOptionsDoesNotRetainOptions {
+  NSString *modelPath = [self requireFixture:@"add_coreml" ofType:@"pte"];
+  if (!modelPath) return;
+  NSError *error = nil;
+  ExecuTorchModule *module = [[ExecuTorchModule alloc] initWithFilePath:modelPath];
+
+  __weak ExecuTorchBackendOptionsMap *weakOptions = nil;
+  @autoreleasepool {
+    ExecuTorchBackendOptionsMap *options = [ExecuTorchBackendOptionsMap mapWithOptions:@{
+      @"CoreMLBackend": @[
+        [ExecuTorchBackendOption optionWithKey:@"compute_unit" stringValue:@"cpu_only"],
+      ],
+    } error:&error];
+    XCTAssertNotNil(options, @"%@", error);
+    weakOptions = options;
+    XCTAssertTrue([module loadWithOptions:options error:&error], @"%@", error);
+  }
+  // Local + autoreleased refs have drained. With deep-copy, the wrapper
+  // can be reclaimed and the C++ Module keeps its own copy of the
+  // backend options for any future lazy load_method calls.
+  XCTAssertNil(weakOptions,
+      @"loadWithOptions: must not retain the map (Module deep-copies). "
+      @"See module.cpp Module::load(LoadBackendOptionsMap, ...) deep-copy.");
+
+  // Forward must still execute against the Module-owned options copy.
+  ExecuTorchTensor *one =
+      [[ExecuTorchTensor alloc] initWithScalars:@[@1.0f] dataType:ExecuTorchDataTypeFloat];
+  NSArray<ExecuTorchValue *> *outputs =
+      [module forwardWithTensors:@[one, one] error:&error];
+  XCTAssertNotNil(outputs, @"%@", error);
+
+  __block float result = NAN;
+  [outputs.firstObject.tensorValue
+      bytesWithHandler:^(const void *bytes, NSInteger count, ExecuTorchDataType dt) {
+    if (dt == ExecuTorchDataTypeFloat && count >= 1) {
+      result = ((const float *)bytes)[0];
+    }
+  }];
+  XCTAssertEqual(result, 2.0f);
+}
+
 - (NSBundle *)resourceBundle {
 #if SWIFT_PACKAGE
   return SWIFTPM_MODULE_BUNDLE;