Add APIs to fetch Method and Tensor metadata.

extension/apple/ExecuTorch/Exported/ExecuTorchModule.h

@@ -10,6 +10,82 @@
 
 NS_ASSUME_NONNULL_BEGIN
 
+/**
+ * Holds the static metadata for a single tensor: its shape, layout,
+ * element type, whether its memory was pre-planned by the runtime,
+ * and its debug name.
+ */
+NS_SWIFT_NAME(TensorMetadata)
+__attribute__((deprecated("This API is experimental.")))
+@interface ExecuTorchTensorMetadata : NSObject
+
+/** The size of each dimension. */
+@property (nonatomic, readonly) NSArray<NSNumber *> *shape;
+
+/** The order in which dimensions are laid out. */
+@property (nonatomic, readonly) NSArray<NSNumber *> *dimensionOrder;
+
+/** The scalar type of each element in the tensor. */
+@property (nonatomic, readonly) ExecuTorchDataType dataType;
+
+/** YES if the runtime pre-allocated memory for this tensor. */
+@property (nonatomic, readonly) BOOL isMemoryPlanned;
+
+/** The (optional) user-visible name of this tensor (may be empty) */
+@property (nonatomic, readonly) NSString *name;
+
++ (instancetype)new NS_UNAVAILABLE;
+- (instancetype)init NS_UNAVAILABLE;
+
+@end
+
+/**
+ * Encapsulates all of the metadata for a loaded method: its name,
+ * how many inputs/outputs/attributes it has, per-argument tags,
+ * per-tensor metadata, buffer sizes, backends, and instruction count.
+ */
+NS_SWIFT_NAME(MethodMetadata)
+__attribute__((deprecated("This API is experimental.")))
+@interface ExecuTorchMethodMetadata : NSObject
+
+/** The method’s name. */
+@property (nonatomic, readonly) NSString *name;
+
+/** An array of ExecuTorchValueTag raw values, one per declared input. */
+@property (nonatomic, readonly) NSArray<NSNumber *> *inputValueTags;
+
+/** An array of ExecuTorchValueTag raw values, one per declared output. */
+@property (nonatomic, readonly) NSArray<NSNumber *> *outputValueTags;
+
+/**
+ * Mapping from input-index to TensorMetadata.
+ * Only present for those indices whose tag == .tensor
+ */
+@property (nonatomic, readonly) NSDictionary<NSNumber *, ExecuTorchTensorMetadata *> *inputTensorMetadatas;
+
+/**
+ * Mapping from output-index to TensorMetadata.
+ * Only present for those indices whose tag == .tensor
+ */
+@property (nonatomic, readonly) NSDictionary<NSNumber *, ExecuTorchTensorMetadata *> *outputTensorMetadatas;
+
+/** A list of attribute TensorsMetadata. */
+@property (nonatomic, readonly) NSArray<ExecuTorchTensorMetadata *> *attributeTensorMetadatas;
+
+/** A list of memory-planned buffer sizes. */
+@property (nonatomic, readonly) NSArray<NSNumber *> *memoryPlannedBufferSizes;
+
+/** Names of all backends this method can run on. */
+@property (nonatomic, readonly) NSArray<NSString *> *backendNames;
+
+/** Total number of low-level instructions in this method’s body. */
+@property (nonatomic, readonly) NSInteger instructionCount;
+
++ (instancetype)new NS_UNAVAILABLE;
+- (instancetype)init NS_UNAVAILABLE;
+
+@end
+
 /**
  * Enum to define loading behavior.
  * Values can be a subset, but must numerically match exactly those defined in
@@ -115,6 +191,19 @@ __attribute__((deprecated("This API is experimental.")))
  */
 - (nullable NSSet<NSString *> *)methodNames:(NSError **)error;
 
+/**
+ * Retrieves full metadata for a particular method in the loaded module.
+ *
+ * This includes the method’s name, input/output value tags, tensor shapes
+ * and layouts, buffer sizes, backend support list, and instruction count.
+ *
+ * @param methodName A string representing the method name.
+ * @param error A pointer to an NSError pointer that is set if an error occurs.
+ * @return An ExecuTorchMethodMetadata object on success, or nil if the method isn’t found or a load error occurred.
+ */
+ - (nullable ExecuTorchMethodMetadata *)methodMetadata:(NSString *)methodName
+                                                 error:(NSError **)error;
+
 /**
  * Executes a specific method with the provided input values.
  *

extension/apple/ExecuTorch/Exported/ExecuTorchModule.mm

@@ -9,6 +9,7 @@
 #import "ExecuTorchModule.h"
 
 #import "ExecuTorchError.h"
+#import "ExecuTorchUtils.h"
 
 #import <executorch/extension/module/module.h>
 #import <executorch/extension/tensor/tensor.h>
@@ -62,6 +63,186 @@ static inline EValue toEValue(ExecuTorchValue *value) {
   return [ExecuTorchValue new];
 }
 
+@interface ExecuTorchTensorMetadata ()
+
+- (instancetype)initWithTensorMetadata:(const TensorInfo &)tensorInfo
+    NS_DESIGNATED_INITIALIZER;
+
+@end
+
+@implementation ExecuTorchTensorMetadata {
+  NSArray<NSNumber *> *_shape;
+  NSArray<NSNumber *> *_dimensionOrder;
+  ExecuTorchDataType _dataType;
+  BOOL _isMemoryPlanned;
+  NSString *_name;
+}
+
+- (instancetype)initWithTensorMetadata:(const TensorInfo &)tensorInfo {
+  self = [super init];
+  if (self) {
+    _shape = utils::toNSArray(tensorInfo.sizes());
+    _dimensionOrder = utils::toNSArray(tensorInfo.dim_order());
+    _dataType = (ExecuTorchDataType)tensorInfo.scalar_type();
+    _isMemoryPlanned = tensorInfo.is_memory_planned();
+    _name = [[NSString alloc] initWithBytes:tensorInfo.name().data()
+                                     length:tensorInfo.name().size()
+                                   encoding:NSUTF8StringEncoding];
+  }
+  return self;
+}
+
+@end
+
+@interface ExecuTorchMethodMetadata ()
+
+- (nullable instancetype)initWithMethodMetadata:(const MethodMeta &)methodMeta
+                                          error:(NSError **)error
+    NS_DESIGNATED_INITIALIZER;
+
+@end
+
+@implementation ExecuTorchMethodMetadata {
+  NSString *_name;
+  NSMutableArray<NSNumber *> *_inputValueTags;
+  NSMutableArray<NSNumber *> *_outputValueTags;
+  NSMutableDictionary<NSNumber *, ExecuTorchTensorMetadata *> *_inputTensorMetadatas;
+  NSMutableDictionary<NSNumber *, ExecuTorchTensorMetadata *> *_outputTensorMetadatas;
+  NSMutableArray<ExecuTorchTensorMetadata *> *_attributeTensorMetadatas;
+  NSMutableArray<NSNumber *> *_memoryPlannedBufferSizes;
+  NSMutableArray<NSString *> *_backendNames;
+  NSInteger _instructionCount;
+}
+
+- (nullable instancetype)initWithMethodMetadata:(const MethodMeta &)methodMeta
+                                          error:(NSError **)error {
+  self = [super init];
+  if (self) {
+    _name = @(methodMeta.name());
+    const NSInteger inputCount = methodMeta.num_inputs();
+    const NSInteger outputCount = methodMeta.num_outputs();
+    const NSInteger attributeCount = methodMeta.num_attributes();
+    const NSInteger memoryPlannedBufferCount = methodMeta.num_memory_planned_buffers();
+    const NSInteger backendCount = methodMeta.num_backends();
+    _instructionCount = methodMeta.num_instructions();
+    _inputValueTags = [NSMutableArray arrayWithCapacity:inputCount];
+    _outputValueTags = [NSMutableArray arrayWithCapacity:outputCount];
+    _inputTensorMetadatas = [NSMutableDictionary dictionary];
+    _outputTensorMetadatas = [NSMutableDictionary dictionary];
+    _attributeTensorMetadatas = [NSMutableArray arrayWithCapacity:attributeCount];
+    _memoryPlannedBufferSizes = [NSMutableArray arrayWithCapacity:memoryPlannedBufferCount];
+    _backendNames = [NSMutableArray arrayWithCapacity:backendCount];
+
+    for (NSInteger index = 0; index < inputCount; ++index) {
+      auto result = methodMeta.input_tag(index);
+      if (!result.ok()) {
+        if (error) {
+          *error = ExecuTorchErrorWithCode((ExecuTorchErrorCode)result.error());
+        }
+        return nil;
+      }
+      const auto inputValueTag = (ExecuTorchValueTag)result.get();
+      [_inputValueTags addObject:@(inputValueTag)];
+
+      if (inputValueTag == ExecuTorchValueTagTensor) {
+        auto tensorMetadataResult = methodMeta.input_tensor_meta(index);
+        if (!tensorMetadataResult.ok()) {
+          if (error) {
+            *error = ExecuTorchErrorWithCode((ExecuTorchErrorCode)tensorMetadataResult.error());
+          }
+          return nil;
+        }
+        _inputTensorMetadatas[@(index)] = [[ExecuTorchTensorMetadata alloc] initWithTensorMetadata:tensorMetadataResult.get()];
+      }
+    }
+    for (NSInteger index = 0; index < outputCount; ++index) {
+      auto result = methodMeta.output_tag(index);
+      if (!result.ok()) {
+        if (error) {
+          *error = ExecuTorchErrorWithCode((ExecuTorchErrorCode)result.error());
+        }
+        return nil;
+      }
+      const auto outputValueTag = (ExecuTorchValueTag)result.get();
+      [_outputValueTags addObject:@(outputValueTag)];
+
+      if (outputValueTag == ExecuTorchValueTagTensor) {
+        auto tensorMetadataResult = methodMeta.output_tensor_meta(index);
+        if (!tensorMetadataResult.ok()) {
+          if (error) {
+            *error = ExecuTorchErrorWithCode((ExecuTorchErrorCode)tensorMetadataResult.error());
+          }
+          return nil;
+        }
+        _outputTensorMetadatas[@(index)] = [[ExecuTorchTensorMetadata alloc] initWithTensorMetadata:tensorMetadataResult.get()];
+      }
+    }
+    for (NSInteger index = 0; index < attributeCount; ++index) {
+      auto result = methodMeta.attribute_tensor_meta(index);
+      if (!result.ok()) {
+        if (error) {
+          *error = ExecuTorchErrorWithCode((ExecuTorchErrorCode)result.error());
+        }
+        return nil;
+      }
+      [_attributeTensorMetadatas addObject:[[ExecuTorchTensorMetadata alloc] initWithTensorMetadata:result.get()]];
+    }
+    for (NSInteger index = 0; index < memoryPlannedBufferCount; ++index) {
+      auto result = methodMeta.memory_planned_buffer_size(index);
+      if (!result.ok()) {
+        if (error) {
+          *error = ExecuTorchErrorWithCode((ExecuTorchErrorCode)result.error());
+        }
+        return nil;
+      }
+      const auto memoryPlannedBufferSize = result.get();
+      [_memoryPlannedBufferSizes addObject:@(memoryPlannedBufferSize)];
+    }
+    for (NSInteger index = 0; index < backendCount; ++index) {
+      auto result = methodMeta.get_backend_name(index);
+      if (!result.ok()) {
+        if (error) {
+          *error = ExecuTorchErrorWithCode((ExecuTorchErrorCode)result.error());
+        }
+        return nil;
+      }
+      NSString *backendName = [NSString stringWithUTF8String:result.get()];
+      [_backendNames addObject:backendName];
+    }
+  }
+  return self;
+}
+
+- (NSArray<NSNumber *> *)inputValueTags {
+  return _inputValueTags;
+}
+
+- (NSArray<NSNumber *> *)outputValueTags {
+  return _outputValueTags;
+}
+
+- (NSDictionary<NSNumber *,ExecuTorchTensorMetadata *> *)inputTensorMetadatas {
+  return _inputTensorMetadatas;
+}
+
+- (NSDictionary<NSNumber *,ExecuTorchTensorMetadata *> *)outputTensorMetadatas {
+  return _outputTensorMetadatas;
+}
+
+- (NSArray<ExecuTorchTensorMetadata *> *)attributeTensorMetadatas {
+  return _attributeTensorMetadatas;
+}
+
+- (NSArray<NSNumber *> *)memoryPlannedBufferSizes {
+  return _memoryPlannedBufferSizes;
+}
+
+- (NSArray<NSString *> *)backendNames {
+  return _backendNames;
+}
+
+@end
+
 @implementation ExecuTorchModule {
   std::unique_ptr<Module> _module;
 }
@@ -134,6 +315,19 @@ - (BOOL)isMethodLoaded:(NSString *)methodName {
   return methods;
 }
 
+- (nullable ExecuTorchMethodMetadata *)methodMetadata:(NSString *)methodName
+                                                error:(NSError **)error {
+  const auto result = _module->method_meta(methodName.UTF8String);
+  if (!result.ok()) {
+    if (error) {
+      *error = ExecuTorchErrorWithCode((ExecuTorchErrorCode)result.error());
+    }
+    return nil;
+  }
+  return [[ExecuTorchMethodMetadata alloc] initWithMethodMetadata:result.get()
+                                                            error:error];
+}
+
 - (nullable NSArray<ExecuTorchValue *> *)executeMethod:(NSString *)methodName
                                             withInputs:(NSArray<ExecuTorchValue *> *)values
                                                  error:(NSError **)error {

extension/apple/ExecuTorch/__tests__/ModuleTest.swift

@@ -70,4 +70,46 @@ class ModuleTest: XCTestCase {
     XCTAssertNoThrow(outputs3 = try module.forward(inputs3))
     XCTAssertEqual(outputs3?.first?.tensor, Tensor([42.5], dataType: .float, shapeDynamism: .static))
   }
+
+  func testmethodMetadata() throws {
+    guard let modelPath = resourceBundle.path(forResource: "add", ofType: "pte") else {
+      XCTFail("Couldn't find the model file")
+      return
+    }
+    let module = Module(filePath: modelPath)
+    let methodMetadata = try module.methodMetadata("forward")
+    XCTAssertEqual(methodMetadata.name, "forward")
+    XCTAssertEqual(methodMetadata.inputValueTags.count, 2)
+    XCTAssertEqual(methodMetadata.outputValueTags.count, 1)
+
+    XCTAssertEqual(ValueTag(rawValue: methodMetadata.inputValueTags[0].uint32Value), .tensor)
+    let inputTensorMetadata1 = methodMetadata.inputTensorMetadatas[0]
+    XCTAssertEqual(inputTensorMetadata1?.shape, [1])
+    XCTAssertEqual(inputTensorMetadata1?.dimensionOrder, [0])
+    XCTAssertEqual(inputTensorMetadata1?.dataType, .float)
+    XCTAssertEqual(inputTensorMetadata1?.isMemoryPlanned, true)
+    XCTAssertEqual(inputTensorMetadata1?.name, "")
+
+    XCTAssertEqual(ValueTag(rawValue: methodMetadata.inputValueTags[1].uint32Value), .tensor)
+    let inputTensorMetadata2 = methodMetadata.inputTensorMetadatas[1]
+    XCTAssertEqual(inputTensorMetadata2?.shape, [1])
+    XCTAssertEqual(inputTensorMetadata2?.dimensionOrder, [0])
+    XCTAssertEqual(inputTensorMetadata2?.dataType, .float)
+    XCTAssertEqual(inputTensorMetadata2?.isMemoryPlanned, true)
+    XCTAssertEqual(inputTensorMetadata2?.name, "")
+
+    XCTAssertEqual(ValueTag(rawValue: methodMetadata.outputValueTags[0].uint32Value), .tensor)
+    let outputTensorMetadata = methodMetadata.outputTensorMetadatas[0]
+    XCTAssertEqual(outputTensorMetadata?.shape, [1])
+    XCTAssertEqual(outputTensorMetadata?.dimensionOrder, [0])
+    XCTAssertEqual(outputTensorMetadata?.dataType, .float)
+    XCTAssertEqual(outputTensorMetadata?.isMemoryPlanned, true)
+    XCTAssertEqual(outputTensorMetadata?.name, "")
+
+    XCTAssertEqual(methodMetadata.attributeTensorMetadatas.count, 0)
+    XCTAssertEqual(methodMetadata.memoryPlannedBufferSizes.count, 1)
+    XCTAssertEqual(methodMetadata.memoryPlannedBufferSizes[0], 48)
+    XCTAssertEqual(methodMetadata.backendNames.count, 0)
+    XCTAssertEqual(methodMetadata.instructionCount, 1)
+  }
 }