up

Author: metascroy

Package.swift

@@ -18,6 +18,7 @@
 // https://pytorch.org/executorch/main/using-executorch-ios
 
 import PackageDescription
+import Foundation
 
 let debug_suffix = "_debug"
 let dependencies_suffix = "_with_dependencies"
@@ -126,22 +127,24 @@ for (key, value) in products {
   packageTargets.append(target)
 }
 
-// Test fixtures. add_coreml.pte is generated at CI time by
-// extension/apple/ExecuTorch/__tests__/resources/generate_coreml_test_models.py
-// (invoked by scripts/build_apple_frameworks.sh before `swift test`). It is
-// gitignored, so include it in test resources only when present so that
-// `swift test` runs on dev machines without CoreML python deps don't fail
-// at the SwiftPM resolve stage.
+// Test fixtures. add_coreml.pte and add_mul_coreml.pte are generated at CI
+// time by extension/apple/ExecuTorch/__tests__/resources/generate_coreml_test_models.py
+// (invoked by scripts/build_apple_frameworks.sh before `swift test`). They
+// are gitignored, so include them in test resources only when present so
+// that `swift test` runs on dev machines without CoreML python deps don't
+// fail at the SwiftPM resolve stage.
 let testResourcesDir = "extension/apple/ExecuTorch/__tests__/resources"
 var testResources: [Resource] = [.copy("resources/add.pte")]
 if FileManager.default.fileExists(atPath: "\(testResourcesDir)/add_coreml.pte") {
   testResources.append(.copy("resources/add_coreml.pte"))
 }
+if FileManager.default.fileExists(atPath: "\(testResourcesDir)/add_mul_coreml.pte") {
+  testResources.append(.copy("resources/add_mul_coreml.pte"))
+}
 
 // SwiftPM resources must live under the target's path, so the ObjC test
 // target uses symlinks to the canonical resources directory. The symlinks
-// themselves are gitignored and (re)created by scripts/build_apple_frameworks.sh
-// (CI) and swift_play/run.sh (local).
+// themselves are gitignored and (re)created by scripts/build_apple_frameworks.sh.
 let objcTestsDir = "extension/apple/ExecuTorch/__tests__/ObjC"
 var objcTestResources: [Resource] = []
 if FileManager.default.fileExists(atPath: "\(objcTestsDir)/add.pte") {
@@ -150,6 +153,9 @@ if FileManager.default.fileExists(atPath: "\(objcTestsDir)/add.pte") {
 if FileManager.default.fileExists(atPath: "\(objcTestsDir)/add_coreml.pte") {
   objcTestResources.append(.copy("add_coreml.pte"))
 }
+if FileManager.default.fileExists(atPath: "\(objcTestsDir)/add_mul_coreml.pte") {
+  objcTestResources.append(.copy("add_mul_coreml.pte"))
+}
 
 let testLinkerSettings: [LinkerSetting] = [
   .unsafeFlags([

extension/apple/ExecuTorch/Exported/ExecuTorchBackendOptionsMap.mm

@@ -71,7 +71,11 @@ Error buildBackendOptionsMap(
           break;
         case ExecuTorchBackendOptionTypeInteger:
           // The C++ runtime stores integer option values as 32-bit `int`.
-          // Reject anything that would silently narrow.
+          // Reject anything that would silently narrow. On Apple's current
+          // 64-bit-only targets NSInteger is int64_t, so this check is
+          // meaningful; on a hypothetical 32-bit build NSInteger would be
+          // int32_t and the comparison would be tautological (still correct,
+          // just never trips).
           if (opt.intValue < INT_MIN || opt.intValue > INT_MAX) {
             return Error::InvalidArgument;
           }
@@ -138,6 +142,13 @@ - (nullable instancetype)initWithOptions:(NSDictionary<NSString *, NSArray<Execu
     return nil;
   }
   _storage = std::move(storage);
+  // Move-assignment order matters: `_storage` is moved first so the Spans
+  // inside `map` (which point at each inner vector's heap buffer) survive
+  // into `_storage`. std::vector's move preserves data() pointers, so the
+  // Spans remain valid. This relies on `LoadBackendOptionsMap`'s move
+  // being a shallow member-wise move that does not recompute span
+  // pointers; if that ever changes, the move order here would need to be
+  // revisited. The end-to-end CoreML-delegated test exercises this path.
   _map = std::move(map);
   // Snapshot the input as a shallow-immutable dictionary, then also copy
   // each value array immutably. Combined with ExecuTorchBackendOption

extension/apple/ExecuTorch/Exported/ExecuTorchModule.mm

@@ -375,13 +375,14 @@ - (BOOL)loadMethod:(NSString *)methodName
            options:(ExecuTorchBackendOptionsMap *)options
              error:(NSError **)error {
   NSParameterAssert(options);
-  // Retain the options for the Module's lifetime, matching -loadWithOptions:.
-  // Today Module::load_method consumes the pointer synchronously and does
-  // not cache it, so retention is not strictly required for this call.
-  // We retain uniformly so the ObjC ownership rule is the same for both
-  // entry points: any BackendOptionsMap handed to the Module stays alive
-  // as long as the Module does.
-  _loadedBackendOptions = options;
+  // Do NOT assign to _loadedBackendOptions here. Module::load_method
+  // consumes `backend_options` synchronously within this call and does
+  // not cache it (see module.cpp); ARC keeps `options` alive for the
+  // call duration via the parameter. Overwriting _loadedBackendOptions
+  // would release any map previously installed by -loadWithOptions:,
+  // but the C++ Module's `backend_options_` raw pointer (set by
+  // -loadWithOptions:) would still reference that released map's
+  // storage — a use-after-free on the next lazy load_method.
   const auto errorCode = _module->load_method(methodName.UTF8String,
                                                /*planned_memory=*/nullptr,
                                                /*event_tracer=*/nullptr,

extension/apple/ExecuTorch/__tests__/ModuleTest.swift

@@ -374,4 +374,74 @@ class ModuleTest: XCTestCase {
       XCTAssertEqual(outs.first?.tensor(), Tensor([Float(2)]))
     }
   }
+
+  // Covers the Module.load(_:options:) Swift wrapper over the ObjC
+  // loadMethod:options: bridge. Loads a specific method with per-delegate
+  // backend options, then executes it.
+  func testLoadMethodWithBackendOptionsThenExecuteOnCoreMLDelegatedModel() throws {
+    let modelPath = try requireFixture("add_coreml", ofType: "pte")
+    let module = Module(filePath: modelPath)
+    let options = try BackendOptionsMap(options: [
+      "CoreMLBackend": [
+        BackendOption("compute_unit", "cpu_and_gpu"),
+      ]
+    ])
+    try module.load("forward", options: options)
+    XCTAssertTrue(module.isLoaded("forward"))
+
+    let inputs: [Tensor<Float>] = [Tensor([1]), Tensor([1])]
+    let outputs: [Value] = try module.forward(inputs)
+    XCTAssertEqual(outputs.first?.tensor(), Tensor([Float(2)]))
+  }
+
+  // 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.
+  //   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).
+  //
+  // 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 {
+    let modelPath = try requireFixture("add_mul_coreml", ofType: "pte")
+    let module = Module(filePath: modelPath)
+
+    weak var weakA: BackendOptionsMap?
+    try autoreleasepool {
+      let optionsA = try BackendOptionsMap(options: [
+        "CoreMLBackend": [BackendOption("compute_unit", "cpu_only")]
+      ])
+      weakA = optionsA
+      try module.load(optionsA)
+    }
+    XCTAssertNotNil(weakA, "Module must retain optionsA after load(optionsA)")
+
+    try autoreleasepool {
+      let optionsB = try BackendOptionsMap(options: [
+        "CoreMLBackend": [BackendOption("compute_unit", "cpu_and_gpu")]
+      ])
+      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.
+    let inputs: [Tensor<Float>] = [Tensor([2]), Tensor([3])]
+    let addOuts: [Value] = try module.forward(inputs)
+    XCTAssertEqual(addOuts.first?.tensor(), Tensor([Float(5)]))
+
+    // "mul" was loaded explicitly with optionsB and should compute 2 * 3.
+    let mulOuts: [Value] = try module.execute("mul", inputs)
+    XCTAssertEqual(mulOuts.first?.tensor(), Tensor([Float(6)]))
+  }
 }

extension/apple/ExecuTorch/__tests__/ObjC/.gitignore

@@ -1,6 +1,6 @@
-# Resource symlinks created at build time by scripts/build_apple_frameworks.sh
-# (CI) and swift_play/run.sh (local). Pointing at ../resources/*.pte so the
-# ObjC test target's .copy(...) can find them under its own path per SwiftPM
-# rules.
+# Resource symlinks created at build time by scripts/build_apple_frameworks.sh.
+# Pointing at ../resources/*.pte so the ObjC test target's .copy(...) can
+# find them under its own path per SwiftPM rules.
 add.pte
 add_coreml.pte
+add_mul_coreml.pte

extension/apple/ExecuTorch/__tests__/resources/.gitignore

@@ -1,2 +1,3 @@
 # Generated by generate_coreml_test_models.py at CI time, not committed.
 add_coreml.pte
+add_mul_coreml.pte

extension/apple/ExecuTorch/__tests__/resources/generate_coreml_test_models.py

@@ -1,19 +1,22 @@
 """Generate CoreML-delegated test fixtures for the Swift/ObjC bindings.
 
 Currently produces:
-  - add_coreml.pte: a CoreML-delegated tensor-add model whose forward(x, y)
-    returns x + y. Used by ModuleTest.testLoadWithBackendOptionsThenExecuteOnCoreML
-    to exercise the BackendOptionsMap lifetime path end-to-end against a
-    delegated model (add.pte has no delegates, so the per-delegate option
-    lookup path is only exercised with a delegated fixture like this one).
+  - add_coreml.pte: a single-method CoreML-delegated tensor-add model whose
+    forward(x, y) returns x + y.
+  - add_mul_coreml.pte: a two-method CoreML-delegated model exposing
+    forward(x, y) = x + y and mul(x, y) = x * y. Used to exercise mixed
+    load(options:) / load(_:options:) sequences where one method is loaded
+    explicitly with its own options and another is loaded lazily, so the
+    C++ Module's stored backend_options_ is consulted during the lazy path.
+    A non-delegated or single-method fixture does not reach that code path.
 
 Usage:
     python extension/apple/ExecuTorch/__tests__/resources/generate_coreml_test_models.py
 
 This script is invoked by scripts/build_apple_frameworks.sh in CI before
-`swift test` so the fixture is always present in CI runs. The output .pte is
-gitignored; local developers who want to run the CoreML-dependent tests should
-run this script once.
+`swift test` so the fixtures are always present in CI runs. The output .pte
+files are gitignored; local developers who want to run the CoreML-dependent
+tests should run this script once.
 """
 
 import os
@@ -30,22 +33,38 @@ def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
         return x + y
 
 
-def main() -> None:
-    model = AddModule().eval()
-    example_inputs = (torch.tensor([1.0]), torch.tensor([1.0]))
+class MulModule(nn.Module):
+    def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+        return x * y
 
-    ep = torch.export.export(model, example_inputs)
-    lowered = to_edge_transform_and_lower(
-        ep,
-        partitioner=[CoreMLPartitioner()],
-    )
-    exec_program = lowered.to_executorch()
 
-    out_path = os.path.join(os.path.dirname(__file__), "add_coreml.pte")
+def _write_pte(exec_program, filename: str) -> None:
+    out_path = os.path.join(os.path.dirname(__file__), filename)
     with open(out_path, "wb") as f:
         exec_program.write_to_file(f)
     print(f"Wrote {out_path} ({os.path.getsize(out_path)} bytes)")
 
 
+def main() -> None:
+    example_inputs = (torch.tensor([1.0]), torch.tensor([1.0]))
+
+    # Single-method add model.
+    ep_add = torch.export.export(AddModule().eval(), example_inputs)
+    add_only = to_edge_transform_and_lower(
+        ep_add,
+        partitioner=[CoreMLPartitioner()],
+    ).to_executorch()
+    _write_pte(add_only, "add_coreml.pte")
+
+    # Two-method model: forward (add) and mul. Both are CoreML-delegated so
+    # each has its own per-delegate option set to query at load time.
+    ep_mul = torch.export.export(MulModule().eval(), example_inputs)
+    add_mul = to_edge_transform_and_lower(
+        {"forward": ep_add, "mul": ep_mul},
+        partitioner=[CoreMLPartitioner()],
+    ).to_executorch()
+    _write_pte(add_mul, "add_mul_coreml.pte")
+
+
 if __name__ == "__main__":
     main()

scripts/build_apple_frameworks.sh

@@ -342,8 +342,9 @@ python3 extension/apple/ExecuTorch/__tests__/resources/generate_coreml_test_mode
 
 # SwiftPM requires resources to live under the test target's path. The ObjC
 # test target shares fixtures with the Swift one via symlinks.
-ln -sfn ../resources/add.pte         extension/apple/ExecuTorch/__tests__/ObjC/add.pte
-ln -sfn ../resources/add_coreml.pte  extension/apple/ExecuTorch/__tests__/ObjC/add_coreml.pte
+ln -sfn ../resources/add.pte             extension/apple/ExecuTorch/__tests__/ObjC/add.pte
+ln -sfn ../resources/add_coreml.pte      extension/apple/ExecuTorch/__tests__/ObjC/add_coreml.pte
+ln -sfn ../resources/add_mul_coreml.pte  extension/apple/ExecuTorch/__tests__/ObjC/add_mul_coreml.pte
 
 echo "Running tests"