Introduce unload method API to Module. (pytorch#13364)

Summary: .

Differential Revision: D80149419

Author: shoumikhin

extension/apple/ExecuTorch/Exported/ExecuTorchModule.h

@@ -187,6 +187,14 @@ __attribute__((deprecated("This API is experimental.")))
  */
 - (BOOL)isMethodLoaded:(NSString *)methodName NS_SWIFT_NAME(isLoaded(_:));
 
+/**
+ * Unloads a method and releases its native resources and planned buffers.
+ *
+ * @param methodName The method to unload.
+ * @return YES if the method was unloaded; NO if it was not loaded at all.
+ */
+- (BOOL)unloadMethod:(NSString *)methodName NS_SWIFT_NAME(unload(_:));
+
 /**
  * Retrieves the set of method names available in the loaded program.
  *

extension/apple/ExecuTorch/Exported/ExecuTorchModule.mm

@@ -304,6 +304,13 @@ - (BOOL)isMethodLoaded:(NSString *)methodName {
   return _module->is_method_loaded(methodName.UTF8String);
 }
 
+- (BOOL)unloadMethod:(NSString *)methodName {
+  const auto didUnload = _module->unload_method(methodName.UTF8String);
+  [_inputs removeObjectForKey:methodName];
+  [_outputs removeObjectForKey:methodName];
+  return didUnload;
+}
+
 - (nullable NSSet<NSString *> *)methodNames:(NSError **)error {
   const auto result = _module->method_names();
   if (!result.ok()) {

extension/apple/ExecuTorch/__tests__/ModuleTest.swift

@@ -171,4 +171,26 @@ class ModuleTest: XCTestCase {
 
     XCTAssertThrowsError(try module.setInputs(Tensor<Float>([1])))
   }
+
+  func testUnloadMethod() {
+    guard let modelPath = resourceBundle.path(forResource: "add", ofType: "pte") else {
+      XCTFail("Couldn't find the model file")
+      return
+    }
+    let module = Module(filePath: modelPath)
+    XCTAssertNoThrow(try module.load("forward"))
+    XCTAssertTrue(module.isLoaded("forward"))
+
+    XCTAssertNoThrow(try module.setInputs(Tensor<Float>([1]), Tensor<Float>([2])))
+    XCTAssertEqual(try module.forward(), Tensor<Float>([3]))
+
+    XCTAssertTrue(module.unload("forward"))
+    XCTAssertFalse(module.isLoaded("forward"))
+    XCTAssertFalse(module.unload("forward"))
+
+    XCTAssertThrowsError(try module.forward())
+    XCTAssertTrue(module.isLoaded("forward"))
+    XCTAssertNoThrow(try module.setInputs(Tensor<Float>([2]), Tensor<Float>([3])))
+    XCTAssertEqual(try module.forward(), Tensor<Float>([5]))
+  }
 }