From 717ee3037b4d6f24f204b19c0317f620d34c1dcc Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Tue, 27 May 2025 15:01:18 -0700 Subject: [PATCH] Add APIs to fetch Method and Tensor metadata. Summary: . Differential Revision: D75482349 --- .../ExecuTorch/Exported/ExecuTorchModule.h | 89 ++++++++ .../ExecuTorch/Exported/ExecuTorchModule.mm | 194 ++++++++++++++++++ .../ExecuTorch/__tests__/ModuleTest.swift | 42 ++++ 3 files changed, 325 insertions(+) diff --git a/extension/apple/ExecuTorch/Exported/ExecuTorchModule.h b/extension/apple/ExecuTorch/Exported/ExecuTorchModule.h index 402ffa644ee..e21dc964416 100644 --- a/extension/apple/ExecuTorch/Exported/ExecuTorchModule.h +++ b/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 *shape; + +/** The order in which dimensions are laid out. */ +@property (nonatomic, readonly) NSArray *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 *inputValueTags; + +/** An array of ExecuTorchValueTag raw values, one per declared output. */ +@property (nonatomic, readonly) NSArray *outputValueTags; + +/** + * Mapping from input-index to TensorMetadata. + * Only present for those indices whose tag == .tensor + */ +@property (nonatomic, readonly) NSDictionary *inputTensorMetadatas; + +/** + * Mapping from output-index to TensorMetadata. + * Only present for those indices whose tag == .tensor + */ +@property (nonatomic, readonly) NSDictionary *outputTensorMetadatas; + +/** A list of attribute TensorsMetadata. */ +@property (nonatomic, readonly) NSArray *attributeTensorMetadatas; + +/** A list of memory-planned buffer sizes. */ +@property (nonatomic, readonly) NSArray *memoryPlannedBufferSizes; + +/** Names of all backends this method can run on. */ +@property (nonatomic, readonly) NSArray *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 *)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. * diff --git a/extension/apple/ExecuTorch/Exported/ExecuTorchModule.mm b/extension/apple/ExecuTorch/Exported/ExecuTorchModule.mm index 0c649382603..66e1fd057ee 100644 --- a/extension/apple/ExecuTorch/Exported/ExecuTorchModule.mm +++ b/extension/apple/ExecuTorch/Exported/ExecuTorchModule.mm @@ -9,6 +9,7 @@ #import "ExecuTorchModule.h" #import "ExecuTorchError.h" +#import "ExecuTorchUtils.h" #import #import @@ -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 *_shape; + NSArray *_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 *_inputValueTags; + NSMutableArray *_outputValueTags; + NSMutableDictionary *_inputTensorMetadatas; + NSMutableDictionary *_outputTensorMetadatas; + NSMutableArray *_attributeTensorMetadatas; + NSMutableArray *_memoryPlannedBufferSizes; + NSMutableArray *_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 *)inputValueTags { + return _inputValueTags; +} + +- (NSArray *)outputValueTags { + return _outputValueTags; +} + +- (NSDictionary *)inputTensorMetadatas { + return _inputTensorMetadatas; +} + +- (NSDictionary *)outputTensorMetadatas { + return _outputTensorMetadatas; +} + +- (NSArray *)attributeTensorMetadatas { + return _attributeTensorMetadatas; +} + +- (NSArray *)memoryPlannedBufferSizes { + return _memoryPlannedBufferSizes; +} + +- (NSArray *)backendNames { + return _backendNames; +} + +@end + @implementation ExecuTorchModule { std::unique_ptr _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 *)executeMethod:(NSString *)methodName withInputs:(NSArray *)values error:(NSError **)error { diff --git a/extension/apple/ExecuTorch/__tests__/ModuleTest.swift b/extension/apple/ExecuTorch/__tests__/ModuleTest.swift index b84919904ba..7ba5aacaa64 100644 --- a/extension/apple/ExecuTorch/__tests__/ModuleTest.swift +++ b/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) + } }