Skip to content

Commit 045b3a5

Browse files
shoumikhinshoumikhin
authored
Refine Tensor metadata for Swift.
Differential Revision: D76533892 Pull Request resolved: #11634
1 parent b3873ef commit 045b3a5

File tree

4 files changed

+78
-30
lines changed

4 files changed

+78
-30
lines changed

‎extension/apple/ExecuTorch/Exported/ExecuTorch+Module.swift‎

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,47 @@
88

99
@_exported import ExecuTorch
1010

11+
@available(*, deprecated, message: "This API is experimental.")
12+
public extension TensorMetadata {
13+
/// The size of each dimension.
14+
var shape: [Int] { __shape.map(\.intValue) }
15+
16+
/// The layout order of each dimension.
17+
var dimensionOrder: [Int] { __dimensionOrder.map(\.intValue) }
18+
}
19+
20+
@available(*, deprecated, message: "This API is experimental.")
21+
public extension MethodMetadata {
22+
/// The declared input tags.
23+
var inputValueTags: [ValueTag] {
24+
__inputValueTags.map { ValueTag(rawValue: $0.uint32Value)! }
25+
}
26+
27+
/// The declared output tags.
28+
var outputValueTags: [ValueTag] {
29+
__outputValueTags.map { ValueTag(rawValue: $0.uint32Value)! }
30+
}
31+
32+
/// A dictionary mapping each input index to its `TensorMetadata`.
33+
var inputTensorMetadata: [Int: TensorMetadata] {
34+
Dictionary(uniqueKeysWithValues:
35+
__inputTensorMetadata.map { (key, value) in (key.intValue, value) }
36+
)
37+
}
38+
39+
/// A dictionary mapping each output index to its `TensorMetadata`.
40+
var outputTensorMetadata: [Int: TensorMetadata] {
41+
Dictionary(uniqueKeysWithValues:
42+
__outputTensorMetadata.map { (key, value) in (key.intValue, value) }
43+
)
44+
}
45+
46+
/// The sizes of all memory-planned buffers as a native Swift array of `Int`.
47+
var memoryPlannedBufferSizes: [Int] {
48+
__memoryPlannedBufferSizes.map(\.intValue)
49+
}
50+
}
51+
1152
@available(*, deprecated, message: "This API is experimental.")
1253
public extension Module {
1354
/// Executes a specific method with the provided input values.

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

Lines changed: 15 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -20,10 +20,12 @@ __attribute__((deprecated("This API is experimental.")))
2020
@interface ExecuTorchTensorMetadata : NSObject
2121

2222
/** The size of each dimension. */
23-
@property (nonatomic, readonly) NSArray<NSNumber *> *shape;
23+
@property (nonatomic, readonly) NSArray<NSNumber *> *shape
24+
NS_REFINED_FOR_SWIFT;
2425

2526
/** The order in which dimensions are laid out. */
26-
@property (nonatomic, readonly) NSArray<NSNumber *> *dimensionOrder;
27+
@property (nonatomic, readonly) NSArray<NSNumber *> *dimensionOrder
28+
NS_REFINED_FOR_SWIFT;
2729

2830
/** The scalar type of each element in the tensor. */
2931
@property (nonatomic, readonly) ExecuTorchDataType dataType;
@@ -52,28 +54,33 @@ __attribute__((deprecated("This API is experimental.")))
5254
@property (nonatomic, readonly) NSString *name;
5355

5456
/** An array of ExecuTorchValueTag raw values, one per declared input. */
55-
@property (nonatomic, readonly) NSArray<NSNumber *> *inputValueTags;
57+
@property (nonatomic, readonly) NSArray<NSNumber *> *inputValueTags
58+
NS_REFINED_FOR_SWIFT;
5659

5760
/** An array of ExecuTorchValueTag raw values, one per declared output. */
58-
@property (nonatomic, readonly) NSArray<NSNumber *> *outputValueTags;
61+
@property (nonatomic, readonly) NSArray<NSNumber *> *outputValueTags
62+
NS_REFINED_FOR_SWIFT;
5963

6064
/**
6165
* Mapping from input-index to TensorMetadata.
6266
* Only present for those indices whose tag == .tensor
6367
*/
64-
@property (nonatomic, readonly) NSDictionary<NSNumber *, ExecuTorchTensorMetadata *> *inputTensorMetadatas;
68+
@property (nonatomic, readonly) NSDictionary<NSNumber *, ExecuTorchTensorMetadata *> *inputTensorMetadata
69+
NS_REFINED_FOR_SWIFT;
6570

6671
/**
6772
* Mapping from output-index to TensorMetadata.
6873
* Only present for those indices whose tag == .tensor
6974
*/
70-
@property (nonatomic, readonly) NSDictionary<NSNumber *, ExecuTorchTensorMetadata *> *outputTensorMetadatas;
75+
@property (nonatomic, readonly) NSDictionary<NSNumber *, ExecuTorchTensorMetadata *> *outputTensorMetadata
76+
NS_REFINED_FOR_SWIFT;
7177

7278
/** A list of attribute TensorsMetadata. */
73-
@property (nonatomic, readonly) NSArray<ExecuTorchTensorMetadata *> *attributeTensorMetadatas;
79+
@property (nonatomic, readonly) NSArray<ExecuTorchTensorMetadata *> *attributeTensorMetadata;
7480

7581
/** A list of memory-planned buffer sizes. */
76-
@property (nonatomic, readonly) NSArray<NSNumber *> *memoryPlannedBufferSizes;
82+
@property (nonatomic, readonly) NSArray<NSNumber *> *memoryPlannedBufferSizes
83+
NS_REFINED_FOR_SWIFT;
7784

7885
/** Names of all backends this method can run on. */
7986
@property (nonatomic, readonly) NSArray<NSString *> *backendNames;

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

Lines changed: 15 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -106,9 +106,9 @@ @implementation ExecuTorchMethodMetadata {
106106
NSString *_name;
107107
NSMutableArray<NSNumber *> *_inputValueTags;
108108
NSMutableArray<NSNumber *> *_outputValueTags;
109-
NSMutableDictionary<NSNumber *, ExecuTorchTensorMetadata *> *_inputTensorMetadatas;
110-
NSMutableDictionary<NSNumber *, ExecuTorchTensorMetadata *> *_outputTensorMetadatas;
111-
NSMutableArray<ExecuTorchTensorMetadata *> *_attributeTensorMetadatas;
109+
NSMutableDictionary<NSNumber *, ExecuTorchTensorMetadata *> *_inputTensorMetadata;
110+
NSMutableDictionary<NSNumber *, ExecuTorchTensorMetadata *> *_outputTensorMetadata;
111+
NSMutableArray<ExecuTorchTensorMetadata *> *_attributeTensorMetadata;
112112
NSMutableArray<NSNumber *> *_memoryPlannedBufferSizes;
113113
NSMutableArray<NSString *> *_backendNames;
114114
NSInteger _instructionCount;
@@ -127,9 +127,9 @@ - (nullable instancetype)initWithMethodMetadata:(const MethodMeta &)methodMeta
127127
_instructionCount = methodMeta.num_instructions();
128128
_inputValueTags = [[NSMutableArray alloc] initWithCapacity:inputCount];
129129
_outputValueTags = [[NSMutableArray alloc] initWithCapacity:outputCount];
130-
_inputTensorMetadatas = [NSMutableDictionary new];
131-
_outputTensorMetadatas = [NSMutableDictionary new];
132-
_attributeTensorMetadatas = [[NSMutableArray alloc] initWithCapacity:attributeCount];
130+
_inputTensorMetadata = [NSMutableDictionary new];
131+
_outputTensorMetadata = [NSMutableDictionary new];
132+
_attributeTensorMetadata = [[NSMutableArray alloc] initWithCapacity:attributeCount];
133133
_memoryPlannedBufferSizes = [[NSMutableArray alloc] initWithCapacity:memoryPlannedBufferCount];
134134
_backendNames = [[NSMutableArray alloc] initWithCapacity:backendCount];
135135

@@ -152,7 +152,7 @@ - (nullable instancetype)initWithMethodMetadata:(const MethodMeta &)methodMeta
152152
}
153153
return nil;
154154
}
155-
_inputTensorMetadatas[@(index)] = [[ExecuTorchTensorMetadata alloc] initWithTensorMetadata:tensorMetadataResult.get()];
155+
_inputTensorMetadata[@(index)] = [[ExecuTorchTensorMetadata alloc] initWithTensorMetadata:tensorMetadataResult.get()];
156156
}
157157
}
158158
for (NSInteger index = 0; index < outputCount; ++index) {
@@ -174,7 +174,7 @@ - (nullable instancetype)initWithMethodMetadata:(const MethodMeta &)methodMeta
174174
}
175175
return nil;
176176
}
177-
_outputTensorMetadatas[@(index)] = [[ExecuTorchTensorMetadata alloc] initWithTensorMetadata:tensorMetadataResult.get()];
177+
_outputTensorMetadata[@(index)] = [[ExecuTorchTensorMetadata alloc] initWithTensorMetadata:tensorMetadataResult.get()];
178178
}
179179
}
180180
for (NSInteger index = 0; index < attributeCount; ++index) {
@@ -185,7 +185,7 @@ - (nullable instancetype)initWithMethodMetadata:(const MethodMeta &)methodMeta
185185
}
186186
return nil;
187187
}
188-
[_attributeTensorMetadatas addObject:[[ExecuTorchTensorMetadata alloc] initWithTensorMetadata:result.get()]];
188+
[_attributeTensorMetadata addObject:[[ExecuTorchTensorMetadata alloc] initWithTensorMetadata:result.get()]];
189189
}
190190
for (NSInteger index = 0; index < memoryPlannedBufferCount; ++index) {
191191
auto result = methodMeta.memory_planned_buffer_size(index);
@@ -221,16 +221,16 @@ - (nullable instancetype)initWithMethodMetadata:(const MethodMeta &)methodMeta
221221
return _outputValueTags;
222222
}
223223

224-
- (NSDictionary<NSNumber *,ExecuTorchTensorMetadata *> *)inputTensorMetadatas {
225-
return _inputTensorMetadatas;
224+
- (NSDictionary<NSNumber *,ExecuTorchTensorMetadata *> *)inputTensorMetadata {
225+
return _inputTensorMetadata;
226226
}
227227

228-
- (NSDictionary<NSNumber *,ExecuTorchTensorMetadata *> *)outputTensorMetadatas {
229-
return _outputTensorMetadatas;
228+
- (NSDictionary<NSNumber *,ExecuTorchTensorMetadata *> *)outputTensorMetadata {
229+
return _outputTensorMetadata;
230230
}
231231

232-
- (NSArray<ExecuTorchTensorMetadata *> *)attributeTensorMetadatas {
233-
return _attributeTensorMetadatas;
232+
- (NSArray<ExecuTorchTensorMetadata *> *)attributeTensorMetadata {
233+
return _attributeTensorMetadata;
234234
}
235235

236236
- (NSArray<NSNumber *> *)memoryPlannedBufferSizes {

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

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -92,31 +92,31 @@ class ModuleTest: XCTestCase {
9292
XCTAssertEqual(methodMetadata.inputValueTags.count, 2)
9393
XCTAssertEqual(methodMetadata.outputValueTags.count, 1)
9494

95-
XCTAssertEqual(ValueTag(rawValue: methodMetadata.inputValueTags[0].uint32Value), .tensor)
96-
let inputTensorMetadata1 = methodMetadata.inputTensorMetadatas[0]
95+
XCTAssertEqual(methodMetadata.inputValueTags[0], .tensor)
96+
let inputTensorMetadata1 = methodMetadata.inputTensorMetadata[0]
9797
XCTAssertEqual(inputTensorMetadata1?.shape, [1])
9898
XCTAssertEqual(inputTensorMetadata1?.dimensionOrder, [0])
9999
XCTAssertEqual(inputTensorMetadata1?.dataType, .float)
100100
XCTAssertEqual(inputTensorMetadata1?.isMemoryPlanned, true)
101101
XCTAssertEqual(inputTensorMetadata1?.name, "")
102102

103-
XCTAssertEqual(ValueTag(rawValue: methodMetadata.inputValueTags[1].uint32Value), .tensor)
104-
let inputTensorMetadata2 = methodMetadata.inputTensorMetadatas[1]
103+
XCTAssertEqual(methodMetadata.inputValueTags[1], .tensor)
104+
let inputTensorMetadata2 = methodMetadata.inputTensorMetadata[1]
105105
XCTAssertEqual(inputTensorMetadata2?.shape, [1])
106106
XCTAssertEqual(inputTensorMetadata2?.dimensionOrder, [0])
107107
XCTAssertEqual(inputTensorMetadata2?.dataType, .float)
108108
XCTAssertEqual(inputTensorMetadata2?.isMemoryPlanned, true)
109109
XCTAssertEqual(inputTensorMetadata2?.name, "")
110110

111-
XCTAssertEqual(ValueTag(rawValue: methodMetadata.outputValueTags[0].uint32Value), .tensor)
112-
let outputTensorMetadata = methodMetadata.outputTensorMetadatas[0]
111+
XCTAssertEqual(methodMetadata.outputValueTags[0], .tensor)
112+
let outputTensorMetadata = methodMetadata.outputTensorMetadata[0]
113113
XCTAssertEqual(outputTensorMetadata?.shape, [1])
114114
XCTAssertEqual(outputTensorMetadata?.dimensionOrder, [0])
115115
XCTAssertEqual(outputTensorMetadata?.dataType, .float)
116116
XCTAssertEqual(outputTensorMetadata?.isMemoryPlanned, true)
117117
XCTAssertEqual(outputTensorMetadata?.name, "")
118118

119-
XCTAssertEqual(methodMetadata.attributeTensorMetadatas.count, 0)
119+
XCTAssertEqual(methodMetadata.attributeTensorMetadata.count, 0)
120120
XCTAssertEqual(methodMetadata.memoryPlannedBufferSizes.count, 1)
121121
XCTAssertEqual(methodMetadata.memoryPlannedBufferSizes[0], 48)
122122
XCTAssertEqual(methodMetadata.backendNames.count, 0)

0 commit comments

Comments
 (0)