From 102195fcccda42ed10758fb556545ebd3551cd89 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Wed, 26 Mar 2025 17:51:43 -0700 Subject: [PATCH] Overloads for the bytes Tensor constrcutor. Summary: https://github.com/pytorch/executorch/issues/8366 Reviewed By: metascroy Differential Revision: D71905631 --- .../ExecuTorch/Exported/ExecuTorchTensor.h | 42 +++++++++++++++++++ .../ExecuTorch/Exported/ExecuTorchTensor.mm | 36 ++++++++++++++++ .../ExecuTorch/__tests__/TensorTest.swift | 28 +++++++++++++ 3 files changed, 106 insertions(+) diff --git a/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.h b/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.h index fadbf137fad..627f28092b6 100644 --- a/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.h +++ b/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.h @@ -244,6 +244,48 @@ __attribute__((deprecated("This API is experimental."))) dataType:(ExecuTorchDataType)dataType shapeDynamism:(ExecuTorchShapeDynamism)shapeDynamism; +/** + * Initializes a tensor by copying bytes from the provided pointer with dynamic bound shape. + * + * @param pointer A pointer to the source data buffer. + * @param shape An NSArray of NSNumber objects representing the tensor's shape. + * @param strides An NSArray of NSNumber objects representing the tensor's strides. + * @param dimensionOrder An NSArray of NSNumber objects indicating the order of dimensions. + * @param dataType An ExecuTorchDataType value specifying the element type. + * @return An initialized ExecuTorchTensor instance with its own copy of the data. + */ +- (instancetype)initWithBytes:(const void *)pointer + shape:(NSArray *)shape + strides:(NSArray *)strides + dimensionOrder:(NSArray *)dimensionOrder + dataType:(ExecuTorchDataType)dataType; + +/** + * Initializes a tensor by copying bytes from the provided pointer, specifying shape, data type, and explicit shape dynamism. + * + * @param pointer A pointer to the source data buffer. + * @param shape An NSArray of NSNumber objects representing the tensor's shape. + * @param dataType An ExecuTorchDataType value specifying the element type. + * @param shapeDynamism An ExecuTorchShapeDynamism value indicating the shape dynamism. + * @return An initialized ExecuTorchTensor instance with its own copy of the data. + */ +- (instancetype)initWithBytes:(const void *)pointer + shape:(NSArray *)shape + dataType:(ExecuTorchDataType)dataType + shapeDynamism:(ExecuTorchShapeDynamism)shapeDynamism; + +/** + * Initializes a tensor by copying bytes from the provided pointer, specifying only the shape and data type. + * + * @param pointer A pointer to the source data buffer. + * @param shape An NSArray of NSNumber objects representing the tensor's shape. + * @param dataType An ExecuTorchDataType value specifying the element type. + * @return An initialized ExecuTorchTensor instance with its own copy of the data. + */ +- (instancetype)initWithBytes:(const void *)pointer + shape:(NSArray *)shape + dataType:(ExecuTorchDataType)dataType; + @end NS_ASSUME_NONNULL_END diff --git a/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.mm b/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.mm index 7f02320d2db..7f9a89742ea 100644 --- a/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.mm +++ b/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.mm @@ -164,4 +164,40 @@ - (instancetype)initWithBytes:(const void *)pointer return [self initWithNativeInstance:&tensor]; } +- (instancetype)initWithBytes:(const void *)pointer + shape:(NSArray *)shape + strides:(NSArray *)strides + dimensionOrder:(NSArray *)dimensionOrder + dataType:(ExecuTorchDataType)dataType { + return [self initWithBytes:pointer + shape:shape + strides:strides + dimensionOrder:dimensionOrder + dataType:dataType + shapeDynamism:ExecuTorchShapeDynamismDynamicBound]; +} + +- (instancetype)initWithBytes:(const void *)pointer + shape:(NSArray *)shape + dataType:(ExecuTorchDataType)dataType + shapeDynamism:(ExecuTorchShapeDynamism)shapeDynamism { + return [self initWithBytes:pointer + shape:shape + strides:@[] + dimensionOrder:@[] + dataType:dataType + shapeDynamism:shapeDynamism]; +} + +- (instancetype)initWithBytes:(const void *)pointer + shape:(NSArray *)shape + dataType:(ExecuTorchDataType)dataType { + return [self initWithBytes:pointer + shape:shape + strides:@[] + dimensionOrder:@[] + dataType:dataType + shapeDynamism:ExecuTorchShapeDynamismDynamicBound]; +} + @end diff --git a/extension/apple/ExecuTorch/__tests__/TensorTest.swift b/extension/apple/ExecuTorch/__tests__/TensorTest.swift index 850ee10d007..1d023abebe0 100644 --- a/extension/apple/ExecuTorch/__tests__/TensorTest.swift +++ b/extension/apple/ExecuTorch/__tests__/TensorTest.swift @@ -67,4 +67,32 @@ class TensorTest: XCTestCase { XCTAssertEqual(tensor.shapeDynamism, .dynamicBound) XCTAssertEqual(tensor.count, 6) } + + func testInitBytes() { + var data: [Double] = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0] + let tensor = data.withUnsafeMutableBytes { + Tensor(bytes: $0.baseAddress!, shape: [2, 3], dataType: .double) + } + XCTAssertEqual(tensor.dataType, .double) + XCTAssertEqual(tensor.shape, [2, 3]) + XCTAssertEqual(tensor.strides, [3, 1]) + XCTAssertEqual(tensor.dimensionOrder, [0, 1]) + XCTAssertEqual(tensor.shapeDynamism, .dynamicBound) + XCTAssertEqual(tensor.count, 6) + } + + func testWithCustomStridesAndDimensionOrder() { + let data: [Float] = [1.0, 2.0, 3.0, 4.0] + let tensor = Tensor( + bytes: data.withUnsafeBytes { $0.baseAddress! }, + shape: [2, 2], + strides: [1, 2], + dimensionOrder: [1, 0], + dataType: .float + ) + XCTAssertEqual(tensor.shape, [2, 2]) + XCTAssertEqual(tensor.strides, [1, 2]) + XCTAssertEqual(tensor.dimensionOrder, [1, 0]) + XCTAssertEqual(tensor.count, 4) + } }