Skip to content

Commit 58e99b9

Browse files
shoumikhinshoumikhin
authored
Add APIs to fetch Method and Tensor metadata. (#11160)
1 parent a8a7e1b commit 58e99b9

File tree

3 files changed

+325
-0
lines changed

3 files changed

+325
-0
lines changed

‎extension/apple/ExecuTorch/Exported/ExecuTorchModule.h‎

Lines changed: 89 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,82 @@
1010

1111
NS_ASSUME_NONNULL_BEGIN
1212

13+
/**
14+
* Holds the static metadata for a single tensor: its shape, layout,
15+
* element type, whether its memory was pre-planned by the runtime,
16+
* and its debug name.
17+
*/
18+
NS_SWIFT_NAME(TensorMetadata)
19+
__attribute__((deprecated("This API is experimental.")))
20+
@interface ExecuTorchTensorMetadata : NSObject
21+
22+
/** The size of each dimension. */
23+
@property (nonatomic, readonly) NSArray<NSNumber *> *shape;
24+
25+
/** The order in which dimensions are laid out. */
26+
@property (nonatomic, readonly) NSArray<NSNumber *> *dimensionOrder;
27+
28+
/** The scalar type of each element in the tensor. */
29+
@property (nonatomic, readonly) ExecuTorchDataType dataType;
30+
31+
/** YES if the runtime pre-allocated memory for this tensor. */
32+
@property (nonatomic, readonly) BOOL isMemoryPlanned;
33+
34+
/** The (optional) user-visible name of this tensor (may be empty) */
35+
@property (nonatomic, readonly) NSString *name;
36+
37+
+ (instancetype)new NS_UNAVAILABLE;
38+
- (instancetype)init NS_UNAVAILABLE;
39+
40+
@end
41+
42+
/**
43+
* Encapsulates all of the metadata for a loaded method: its name,
44+
* how many inputs/outputs/attributes it has, per-argument tags,
45+
* per-tensor metadata, buffer sizes, backends, and instruction count.
46+
*/
47+
NS_SWIFT_NAME(MethodMetadata)
48+
__attribute__((deprecated("This API is experimental.")))
49+
@interface ExecuTorchMethodMetadata : NSObject
50+
51+
/** The method’s name. */
52+
@property (nonatomic, readonly) NSString *name;
53+
54+
/** An array of ExecuTorchValueTag raw values, one per declared input. */
55+
@property (nonatomic, readonly) NSArray<NSNumber *> *inputValueTags;
56+
57+
/** An array of ExecuTorchValueTag raw values, one per declared output. */
58+
@property (nonatomic, readonly) NSArray<NSNumber *> *outputValueTags;
59+
60+
/**
61+
* Mapping from input-index to TensorMetadata.
62+
* Only present for those indices whose tag == .tensor
63+
*/
64+
@property (nonatomic, readonly) NSDictionary<NSNumber *, ExecuTorchTensorMetadata *> *inputTensorMetadatas;
65+
66+
/**
67+
* Mapping from output-index to TensorMetadata.
68+
* Only present for those indices whose tag == .tensor
69+
*/
70+
@property (nonatomic, readonly) NSDictionary<NSNumber *, ExecuTorchTensorMetadata *> *outputTensorMetadatas;
71+
72+
/** A list of attribute TensorsMetadata. */
73+
@property (nonatomic, readonly) NSArray<ExecuTorchTensorMetadata *> *attributeTensorMetadatas;
74+
75+
/** A list of memory-planned buffer sizes. */
76+
@property (nonatomic, readonly) NSArray<NSNumber *> *memoryPlannedBufferSizes;
77+
78+
/** Names of all backends this method can run on. */
79+
@property (nonatomic, readonly) NSArray<NSString *> *backendNames;
80+
81+
/** Total number of low-level instructions in this method’s body. */
82+
@property (nonatomic, readonly) NSInteger instructionCount;
83+
84+
+ (instancetype)new NS_UNAVAILABLE;
85+
- (instancetype)init NS_UNAVAILABLE;
86+
87+
@end
88+
1389
/**
1490
* Enum to define loading behavior.
1591
* Values can be a subset, but must numerically match exactly those defined in
@@ -115,6 +191,19 @@ __attribute__((deprecated("This API is experimental.")))
115191
*/
116192
- (nullable NSSet<NSString *> *)methodNames:(NSError **)error;
117193

194+
/**
195+
* Retrieves full metadata for a particular method in the loaded module.
196+
*
197+
* This includes the method’s name, input/output value tags, tensor shapes
198+
* and layouts, buffer sizes, backend support list, and instruction count.
199+
*
200+
* @param methodName A string representing the method name.
201+
* @param error A pointer to an NSError pointer that is set if an error occurs.
202+
* @return An ExecuTorchMethodMetadata object on success, or nil if the method isn’t found or a load error occurred.
203+
*/
204+
- (nullable ExecuTorchMethodMetadata *)methodMetadata:(NSString *)methodName
205+
error:(NSError **)error;
206+
118207
/**
119208
* Executes a specific method with the provided input values.
120209
*

‎extension/apple/ExecuTorch/Exported/ExecuTorchModule.mm‎

Lines changed: 194 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
#import "ExecuTorchModule.h"
1010

1111
#import "ExecuTorchError.h"
12+
#import "ExecuTorchUtils.h"
1213

1314
#import <executorch/extension/module/module.h>
1415
#import <executorch/extension/tensor/tensor.h>
@@ -62,6 +63,186 @@ static inline EValue toEValue(ExecuTorchValue *value) {
6263
return [ExecuTorchValue new];
6364
}
6465

66+
@interface ExecuTorchTensorMetadata ()
67+
68+
- (instancetype)initWithTensorMetadata:(const TensorInfo &)tensorInfo
69+
NS_DESIGNATED_INITIALIZER;
70+
71+
@end
72+
73+
@implementation ExecuTorchTensorMetadata {
74+
NSArray<NSNumber *> *_shape;
75+
NSArray<NSNumber *> *_dimensionOrder;
76+
ExecuTorchDataType _dataType;
77+
BOOL _isMemoryPlanned;
78+
NSString *_name;
79+
}
80+
81+
- (instancetype)initWithTensorMetadata:(const TensorInfo &)tensorInfo {
82+
self = [super init];
83+
if (self) {
84+
_shape = utils::toNSArray(tensorInfo.sizes());
85+
_dimensionOrder = utils::toNSArray(tensorInfo.dim_order());
86+
_dataType = (ExecuTorchDataType)tensorInfo.scalar_type();
87+
_isMemoryPlanned = tensorInfo.is_memory_planned();
88+
_name = [[NSString alloc] initWithBytes:tensorInfo.name().data()
89+
length:tensorInfo.name().size()
90+
encoding:NSUTF8StringEncoding];
91+
}
92+
return self;
93+
}
94+
95+
@end
96+
97+
@interface ExecuTorchMethodMetadata ()
98+
99+
- (nullable instancetype)initWithMethodMetadata:(const MethodMeta &)methodMeta
100+
error:(NSError **)error
101+
NS_DESIGNATED_INITIALIZER;
102+
103+
@end
104+
105+
@implementation ExecuTorchMethodMetadata {
106+
NSString *_name;
107+
NSMutableArray<NSNumber *> *_inputValueTags;
108+
NSMutableArray<NSNumber *> *_outputValueTags;
109+
NSMutableDictionary<NSNumber *, ExecuTorchTensorMetadata *> *_inputTensorMetadatas;
110+
NSMutableDictionary<NSNumber *, ExecuTorchTensorMetadata *> *_outputTensorMetadatas;
111+
NSMutableArray<ExecuTorchTensorMetadata *> *_attributeTensorMetadatas;
112+
NSMutableArray<NSNumber *> *_memoryPlannedBufferSizes;
113+
NSMutableArray<NSString *> *_backendNames;
114+
NSInteger _instructionCount;
115+
}
116+
117+
- (nullable instancetype)initWithMethodMetadata:(const MethodMeta &)methodMeta
118+
error:(NSError **)error {
119+
self = [super init];
120+
if (self) {
121+
_name = @(methodMeta.name());
122+
const NSInteger inputCount = methodMeta.num_inputs();
123+
const NSInteger outputCount = methodMeta.num_outputs();
124+
const NSInteger attributeCount = methodMeta.num_attributes();
125+
const NSInteger memoryPlannedBufferCount = methodMeta.num_memory_planned_buffers();
126+
const NSInteger backendCount = methodMeta.num_backends();
127+
_instructionCount = methodMeta.num_instructions();
128+
_inputValueTags = [NSMutableArray arrayWithCapacity:inputCount];
129+
_outputValueTags = [NSMutableArray arrayWithCapacity:outputCount];
130+
_inputTensorMetadatas = [NSMutableDictionary dictionary];
131+
_outputTensorMetadatas = [NSMutableDictionary dictionary];
132+
_attributeTensorMetadatas = [NSMutableArray arrayWithCapacity:attributeCount];
133+
_memoryPlannedBufferSizes = [NSMutableArray arrayWithCapacity:memoryPlannedBufferCount];
134+
_backendNames = [NSMutableArray arrayWithCapacity:backendCount];
135+
136+
for (NSInteger index = 0; index < inputCount; ++index) {
137+
auto result = methodMeta.input_tag(index);
138+
if (!result.ok()) {
139+
if (error) {
140+
*error = ExecuTorchErrorWithCode((ExecuTorchErrorCode)result.error());
141+
}
142+
return nil;
143+
}
144+
const auto inputValueTag = (ExecuTorchValueTag)result.get();
145+
[_inputValueTags addObject:@(inputValueTag)];
146+
147+
if (inputValueTag == ExecuTorchValueTagTensor) {
148+
auto tensorMetadataResult = methodMeta.input_tensor_meta(index);
149+
if (!tensorMetadataResult.ok()) {
150+
if (error) {
151+
*error = ExecuTorchErrorWithCode((ExecuTorchErrorCode)tensorMetadataResult.error());
152+
}
153+
return nil;
154+
}
155+
_inputTensorMetadatas[@(index)] = [[ExecuTorchTensorMetadata alloc] initWithTensorMetadata:tensorMetadataResult.get()];
156+
}
157+
}
158+
for (NSInteger index = 0; index < outputCount; ++index) {
159+
auto result = methodMeta.output_tag(index);
160+
if (!result.ok()) {
161+
if (error) {
162+
*error = ExecuTorchErrorWithCode((ExecuTorchErrorCode)result.error());
163+
}
164+
return nil;
165+
}
166+
const auto outputValueTag = (ExecuTorchValueTag)result.get();
167+
[_outputValueTags addObject:@(outputValueTag)];
168+
169+
if (outputValueTag == ExecuTorchValueTagTensor) {
170+
auto tensorMetadataResult = methodMeta.output_tensor_meta(index);
171+
if (!tensorMetadataResult.ok()) {
172+
if (error) {
173+
*error = ExecuTorchErrorWithCode((ExecuTorchErrorCode)tensorMetadataResult.error());
174+
}
175+
return nil;
176+
}
177+
_outputTensorMetadatas[@(index)] = [[ExecuTorchTensorMetadata alloc] initWithTensorMetadata:tensorMetadataResult.get()];
178+
}
179+
}
180+
for (NSInteger index = 0; index < attributeCount; ++index) {
181+
auto result = methodMeta.attribute_tensor_meta(index);
182+
if (!result.ok()) {
183+
if (error) {
184+
*error = ExecuTorchErrorWithCode((ExecuTorchErrorCode)result.error());
185+
}
186+
return nil;
187+
}
188+
[_attributeTensorMetadatas addObject:[[ExecuTorchTensorMetadata alloc] initWithTensorMetadata:result.get()]];
189+
}
190+
for (NSInteger index = 0; index < memoryPlannedBufferCount; ++index) {
191+
auto result = methodMeta.memory_planned_buffer_size(index);
192+
if (!result.ok()) {
193+
if (error) {
194+
*error = ExecuTorchErrorWithCode((ExecuTorchErrorCode)result.error());
195+
}
196+
return nil;
197+
}
198+
const auto memoryPlannedBufferSize = result.get();
199+
[_memoryPlannedBufferSizes addObject:@(memoryPlannedBufferSize)];
200+
}
201+
for (NSInteger index = 0; index < backendCount; ++index) {
202+
auto result = methodMeta.get_backend_name(index);
203+
if (!result.ok()) {
204+
if (error) {
205+
*error = ExecuTorchErrorWithCode((ExecuTorchErrorCode)result.error());
206+
}
207+
return nil;
208+
}
209+
NSString *backendName = [NSString stringWithUTF8String:result.get()];
210+
[_backendNames addObject:backendName];
211+
}
212+
}
213+
return self;
214+
}
215+
216+
- (NSArray<NSNumber *> *)inputValueTags {
217+
return _inputValueTags;
218+
}
219+
220+
- (NSArray<NSNumber *> *)outputValueTags {
221+
return _outputValueTags;
222+
}
223+
224+
- (NSDictionary<NSNumber *,ExecuTorchTensorMetadata *> *)inputTensorMetadatas {
225+
return _inputTensorMetadatas;
226+
}
227+
228+
- (NSDictionary<NSNumber *,ExecuTorchTensorMetadata *> *)outputTensorMetadatas {
229+
return _outputTensorMetadatas;
230+
}
231+
232+
- (NSArray<ExecuTorchTensorMetadata *> *)attributeTensorMetadatas {
233+
return _attributeTensorMetadatas;
234+
}
235+
236+
- (NSArray<NSNumber *> *)memoryPlannedBufferSizes {
237+
return _memoryPlannedBufferSizes;
238+
}
239+
240+
- (NSArray<NSString *> *)backendNames {
241+
return _backendNames;
242+
}
243+
244+
@end
245+
65246
@implementation ExecuTorchModule {
66247
std::unique_ptr<Module> _module;
67248
}
@@ -134,6 +315,19 @@ - (BOOL)isMethodLoaded:(NSString *)methodName {
134315
return methods;
135316
}
136317

318+
- (nullable ExecuTorchMethodMetadata *)methodMetadata:(NSString *)methodName
319+
error:(NSError **)error {
320+
const auto result = _module->method_meta(methodName.UTF8String);
321+
if (!result.ok()) {
322+
if (error) {
323+
*error = ExecuTorchErrorWithCode((ExecuTorchErrorCode)result.error());
324+
}
325+
return nil;
326+
}
327+
return [[ExecuTorchMethodMetadata alloc] initWithMethodMetadata:result.get()
328+
error:error];
329+
}
330+
137331
- (nullable NSArray<ExecuTorchValue *> *)executeMethod:(NSString *)methodName
138332
withInputs:(NSArray<ExecuTorchValue *> *)values
139333
error:(NSError **)error {

‎extension/apple/ExecuTorch/__tests__/ModuleTest.swift‎

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -70,4 +70,46 @@ class ModuleTest: XCTestCase {
7070
XCTAssertNoThrow(outputs3 = try module.forward(inputs3))
7171
XCTAssertEqual(outputs3?.first?.tensor, Tensor([42.5], dataType: .float, shapeDynamism: .static))
7272
}
73+
74+
func testmethodMetadata() throws {
75+
guard let modelPath = resourceBundle.path(forResource: "add", ofType: "pte") else {
76+
XCTFail("Couldn't find the model file")
77+
return
78+
}
79+
let module = Module(filePath: modelPath)
80+
let methodMetadata = try module.methodMetadata("forward")
81+
XCTAssertEqual(methodMetadata.name, "forward")
82+
XCTAssertEqual(methodMetadata.inputValueTags.count, 2)
83+
XCTAssertEqual(methodMetadata.outputValueTags.count, 1)
84+
85+
XCTAssertEqual(ValueTag(rawValue: methodMetadata.inputValueTags[0].uint32Value), .tensor)
86+
let inputTensorMetadata1 = methodMetadata.inputTensorMetadatas[0]
87+
XCTAssertEqual(inputTensorMetadata1?.shape, [1])
88+
XCTAssertEqual(inputTensorMetadata1?.dimensionOrder, [0])
89+
XCTAssertEqual(inputTensorMetadata1?.dataType, .float)
90+
XCTAssertEqual(inputTensorMetadata1?.isMemoryPlanned, true)
91+
XCTAssertEqual(inputTensorMetadata1?.name, "")
92+
93+
XCTAssertEqual(ValueTag(rawValue: methodMetadata.inputValueTags[1].uint32Value), .tensor)
94+
let inputTensorMetadata2 = methodMetadata.inputTensorMetadatas[1]
95+
XCTAssertEqual(inputTensorMetadata2?.shape, [1])
96+
XCTAssertEqual(inputTensorMetadata2?.dimensionOrder, [0])
97+
XCTAssertEqual(inputTensorMetadata2?.dataType, .float)
98+
XCTAssertEqual(inputTensorMetadata2?.isMemoryPlanned, true)
99+
XCTAssertEqual(inputTensorMetadata2?.name, "")
100+
101+
XCTAssertEqual(ValueTag(rawValue: methodMetadata.outputValueTags[0].uint32Value), .tensor)
102+
let outputTensorMetadata = methodMetadata.outputTensorMetadatas[0]
103+
XCTAssertEqual(outputTensorMetadata?.shape, [1])
104+
XCTAssertEqual(outputTensorMetadata?.dimensionOrder, [0])
105+
XCTAssertEqual(outputTensorMetadata?.dataType, .float)
106+
XCTAssertEqual(outputTensorMetadata?.isMemoryPlanned, true)
107+
XCTAssertEqual(outputTensorMetadata?.name, "")
108+
109+
XCTAssertEqual(methodMetadata.attributeTensorMetadatas.count, 0)
110+
XCTAssertEqual(methodMetadata.memoryPlannedBufferSizes.count, 1)
111+
XCTAssertEqual(methodMetadata.memoryPlannedBufferSizes[0], 48)
112+
XCTAssertEqual(methodMetadata.backendNames.count, 0)
113+
XCTAssertEqual(methodMetadata.instructionCount, 1)
114+
}
73115
}

0 commit comments

Comments
 (0)