From 76b2b0c701c8dff21c3b57ee3a5ea27fba072e45 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Thu, 12 Jun 2025 21:48:39 -0700 Subject: [PATCH] Refine Tensor metadata for Swift. Summary: . Reviewed By: cccclai Differential Revision: D76533892 --- .../Exported/ExecuTorch+Module.swift | 41 +++++++++++++++++++ .../ExecuTorch/Exported/ExecuTorchModule.h | 23 +++++++---- .../ExecuTorch/Exported/ExecuTorchModule.mm | 30 +++++++------- .../ExecuTorch/__tests__/ModuleTest.swift | 14 +++---- 4 files changed, 78 insertions(+), 30 deletions(-) diff --git a/extension/apple/ExecuTorch/Exported/ExecuTorch+Module.swift b/extension/apple/ExecuTorch/Exported/ExecuTorch+Module.swift index 0cc2ffb0996..e097a9253de 100644 --- a/extension/apple/ExecuTorch/Exported/ExecuTorch+Module.swift +++ b/extension/apple/ExecuTorch/Exported/ExecuTorch+Module.swift @@ -8,6 +8,47 @@ @_exported import ExecuTorch +@available(*, deprecated, message: "This API is experimental.") +public extension TensorMetadata { + /// The size of each dimension. + var shape: [Int] { __shape.map(\.intValue) } + + /// The layout order of each dimension. + var dimensionOrder: [Int] { __dimensionOrder.map(\.intValue) } +} + +@available(*, deprecated, message: "This API is experimental.") +public extension MethodMetadata { + /// The declared input tags. + var inputValueTags: [ValueTag] { + __inputValueTags.map { ValueTag(rawValue: $0.uint32Value)! } + } + + /// The declared output tags. + var outputValueTags: [ValueTag] { + __outputValueTags.map { ValueTag(rawValue: $0.uint32Value)! } + } + + /// A dictionary mapping each input index to its `TensorMetadata`. + var inputTensorMetadata: [Int: TensorMetadata] { + Dictionary(uniqueKeysWithValues: + __inputTensorMetadata.map { (key, value) in (key.intValue, value) } + ) + } + + /// A dictionary mapping each output index to its `TensorMetadata`. + var outputTensorMetadata: [Int: TensorMetadata] { + Dictionary(uniqueKeysWithValues: + __outputTensorMetadata.map { (key, value) in (key.intValue, value) } + ) + } + + /// The sizes of all memory-planned buffers as a native Swift array of `Int`. + var memoryPlannedBufferSizes: [Int] { + __memoryPlannedBufferSizes.map(\.intValue) + } +} + @available(*, deprecated, message: "This API is experimental.") public extension Module { /// Executes a specific method with the provided input values. diff --git a/extension/apple/ExecuTorch/Exported/ExecuTorchModule.h b/extension/apple/ExecuTorch/Exported/ExecuTorchModule.h index d8f6d8b9bc5..0eafcca8cc7 100644 --- a/extension/apple/ExecuTorch/Exported/ExecuTorchModule.h +++ b/extension/apple/ExecuTorch/Exported/ExecuTorchModule.h @@ -20,10 +20,12 @@ __attribute__((deprecated("This API is experimental."))) @interface ExecuTorchTensorMetadata : NSObject /** The size of each dimension. */ -@property (nonatomic, readonly) NSArray *shape; +@property (nonatomic, readonly) NSArray *shape + NS_REFINED_FOR_SWIFT; /** The order in which dimensions are laid out. */ -@property (nonatomic, readonly) NSArray *dimensionOrder; +@property (nonatomic, readonly) NSArray *dimensionOrder + NS_REFINED_FOR_SWIFT; /** The scalar type of each element in the tensor. */ @property (nonatomic, readonly) ExecuTorchDataType dataType; @@ -52,28 +54,33 @@ __attribute__((deprecated("This API is experimental."))) @property (nonatomic, readonly) NSString *name; /** An array of ExecuTorchValueTag raw values, one per declared input. */ -@property (nonatomic, readonly) NSArray *inputValueTags; +@property (nonatomic, readonly) NSArray *inputValueTags + NS_REFINED_FOR_SWIFT; /** An array of ExecuTorchValueTag raw values, one per declared output. */ -@property (nonatomic, readonly) NSArray *outputValueTags; +@property (nonatomic, readonly) NSArray *outputValueTags + NS_REFINED_FOR_SWIFT; /** * Mapping from input-index to TensorMetadata. * Only present for those indices whose tag == .tensor */ -@property (nonatomic, readonly) NSDictionary *inputTensorMetadatas; +@property (nonatomic, readonly) NSDictionary *inputTensorMetadata + NS_REFINED_FOR_SWIFT; /** * Mapping from output-index to TensorMetadata. * Only present for those indices whose tag == .tensor */ -@property (nonatomic, readonly) NSDictionary *outputTensorMetadatas; +@property (nonatomic, readonly) NSDictionary *outputTensorMetadata + NS_REFINED_FOR_SWIFT; /** A list of attribute TensorsMetadata. */ -@property (nonatomic, readonly) NSArray *attributeTensorMetadatas; +@property (nonatomic, readonly) NSArray *attributeTensorMetadata; /** A list of memory-planned buffer sizes. */ -@property (nonatomic, readonly) NSArray *memoryPlannedBufferSizes; +@property (nonatomic, readonly) NSArray *memoryPlannedBufferSizes + NS_REFINED_FOR_SWIFT; /** Names of all backends this method can run on. */ @property (nonatomic, readonly) NSArray *backendNames; diff --git a/extension/apple/ExecuTorch/Exported/ExecuTorchModule.mm b/extension/apple/ExecuTorch/Exported/ExecuTorchModule.mm index 118b46540a8..7b0b15c00d0 100644 --- a/extension/apple/ExecuTorch/Exported/ExecuTorchModule.mm +++ b/extension/apple/ExecuTorch/Exported/ExecuTorchModule.mm @@ -106,9 +106,9 @@ @implementation ExecuTorchMethodMetadata { NSString *_name; NSMutableArray *_inputValueTags; NSMutableArray *_outputValueTags; - NSMutableDictionary *_inputTensorMetadatas; - NSMutableDictionary *_outputTensorMetadatas; - NSMutableArray *_attributeTensorMetadatas; + NSMutableDictionary *_inputTensorMetadata; + NSMutableDictionary *_outputTensorMetadata; + NSMutableArray *_attributeTensorMetadata; NSMutableArray *_memoryPlannedBufferSizes; NSMutableArray *_backendNames; NSInteger _instructionCount; @@ -127,9 +127,9 @@ - (nullable instancetype)initWithMethodMetadata:(const MethodMeta &)methodMeta _instructionCount = methodMeta.num_instructions(); _inputValueTags = [[NSMutableArray alloc] initWithCapacity:inputCount]; _outputValueTags = [[NSMutableArray alloc] initWithCapacity:outputCount]; - _inputTensorMetadatas = [NSMutableDictionary new]; - _outputTensorMetadatas = [NSMutableDictionary new]; - _attributeTensorMetadatas = [[NSMutableArray alloc] initWithCapacity:attributeCount]; + _inputTensorMetadata = [NSMutableDictionary new]; + _outputTensorMetadata = [NSMutableDictionary new]; + _attributeTensorMetadata = [[NSMutableArray alloc] initWithCapacity:attributeCount]; _memoryPlannedBufferSizes = [[NSMutableArray alloc] initWithCapacity:memoryPlannedBufferCount]; _backendNames = [[NSMutableArray alloc] initWithCapacity:backendCount]; @@ -152,7 +152,7 @@ - (nullable instancetype)initWithMethodMetadata:(const MethodMeta &)methodMeta } return nil; } - _inputTensorMetadatas[@(index)] = [[ExecuTorchTensorMetadata alloc] initWithTensorMetadata:tensorMetadataResult.get()]; + _inputTensorMetadata[@(index)] = [[ExecuTorchTensorMetadata alloc] initWithTensorMetadata:tensorMetadataResult.get()]; } } for (NSInteger index = 0; index < outputCount; ++index) { @@ -174,7 +174,7 @@ - (nullable instancetype)initWithMethodMetadata:(const MethodMeta &)methodMeta } return nil; } - _outputTensorMetadatas[@(index)] = [[ExecuTorchTensorMetadata alloc] initWithTensorMetadata:tensorMetadataResult.get()]; + _outputTensorMetadata[@(index)] = [[ExecuTorchTensorMetadata alloc] initWithTensorMetadata:tensorMetadataResult.get()]; } } for (NSInteger index = 0; index < attributeCount; ++index) { @@ -185,7 +185,7 @@ - (nullable instancetype)initWithMethodMetadata:(const MethodMeta &)methodMeta } return nil; } - [_attributeTensorMetadatas addObject:[[ExecuTorchTensorMetadata alloc] initWithTensorMetadata:result.get()]]; + [_attributeTensorMetadata addObject:[[ExecuTorchTensorMetadata alloc] initWithTensorMetadata:result.get()]]; } for (NSInteger index = 0; index < memoryPlannedBufferCount; ++index) { auto result = methodMeta.memory_planned_buffer_size(index); @@ -221,16 +221,16 @@ - (nullable instancetype)initWithMethodMetadata:(const MethodMeta &)methodMeta return _outputValueTags; } -- (NSDictionary *)inputTensorMetadatas { - return _inputTensorMetadatas; +- (NSDictionary *)inputTensorMetadata { + return _inputTensorMetadata; } -- (NSDictionary *)outputTensorMetadatas { - return _outputTensorMetadatas; +- (NSDictionary *)outputTensorMetadata { + return _outputTensorMetadata; } -- (NSArray *)attributeTensorMetadatas { - return _attributeTensorMetadatas; +- (NSArray *)attributeTensorMetadata { + return _attributeTensorMetadata; } - (NSArray *)memoryPlannedBufferSizes { diff --git a/extension/apple/ExecuTorch/__tests__/ModuleTest.swift b/extension/apple/ExecuTorch/__tests__/ModuleTest.swift index 156e373187f..0aaeaefbcd3 100644 --- a/extension/apple/ExecuTorch/__tests__/ModuleTest.swift +++ b/extension/apple/ExecuTorch/__tests__/ModuleTest.swift @@ -92,31 +92,31 @@ class ModuleTest: XCTestCase { XCTAssertEqual(methodMetadata.inputValueTags.count, 2) XCTAssertEqual(methodMetadata.outputValueTags.count, 1) - XCTAssertEqual(ValueTag(rawValue: methodMetadata.inputValueTags[0].uint32Value), .tensor) - let inputTensorMetadata1 = methodMetadata.inputTensorMetadatas[0] + XCTAssertEqual(methodMetadata.inputValueTags[0], .tensor) + let inputTensorMetadata1 = methodMetadata.inputTensorMetadata[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(methodMetadata.inputValueTags[1], .tensor) + let inputTensorMetadata2 = methodMetadata.inputTensorMetadata[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(methodMetadata.outputValueTags[0], .tensor) + let outputTensorMetadata = methodMetadata.outputTensorMetadata[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.attributeTensorMetadata.count, 0) XCTAssertEqual(methodMetadata.memoryPlannedBufferSizes.count, 1) XCTAssertEqual(methodMetadata.memoryPlannedBufferSizes[0], 48) XCTAssertEqual(methodMetadata.backendNames.count, 0)