From 2f83999ff69b9e11d9691362a97218866619b657 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Wed, 26 Mar 2025 17:28:23 -0700 Subject: [PATCH] Overloads for the bytesNoCopy Tensor constrcutor. Summary: https://github.com/pytorch/executorch/issues/8366 Reviewed By: bsoyluoglu Differential Revision: D71902713 --- .../ExecuTorch/Exported/ExecuTorchTensor.h | 42 +++++++++++++++++++ .../ExecuTorch/Exported/ExecuTorchTensor.mm | 36 ++++++++++++++++ .../ExecuTorch/__tests__/TensorTest.swift | 12 +++++- 3 files changed, 89 insertions(+), 1 deletion(-) diff --git a/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.h b/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.h index 1c7e0ab8c8d..984f5780bce 100644 --- a/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.h +++ b/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.h @@ -156,6 +156,48 @@ __attribute__((deprecated("This API is experimental."))) dataType:(ExecuTorchDataType)dataType shapeDynamism:(ExecuTorchShapeDynamism)shapeDynamism; +/** + * Initializes a tensor without copying data using dynamic bound shape (default strides and dimension order). + * + * @param pointer A pointer to the 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. + */ +- (instancetype)initWithBytesNoCopy:(void *)pointer + shape:(NSArray *)shape + strides:(NSArray *)strides + dimensionOrder:(NSArray *)dimensionOrder + dataType:(ExecuTorchDataType)dataType; + +/** + * Initializes a tensor without copying data, with an explicit shape dynamism. + * + * @param pointer A pointer to the 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. + */ +- (instancetype)initWithBytesNoCopy:(void *)pointer + shape:(NSArray *)shape + dataType:(ExecuTorchDataType)dataType + shapeDynamism:(ExecuTorchShapeDynamism)shapeDynamism; + +/** + * Initializes a tensor without copying data, specifying only the shape and data type. + * + * @param pointer A pointer to the 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. + */ +- (instancetype)initWithBytesNoCopy:(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 8a0e7aea3dd..b198f073601 100644 --- a/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.mm +++ b/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.mm @@ -91,4 +91,40 @@ - (instancetype)initWithBytesNoCopy:(void *)pointer return [self initWithNativeInstance:&tensor]; } +- (instancetype)initWithBytesNoCopy:(void *)pointer + shape:(NSArray *)shape + strides:(NSArray *)strides + dimensionOrder:(NSArray *)dimensionOrder + dataType:(ExecuTorchDataType)dataType { + return [self initWithBytesNoCopy:pointer + shape:shape + strides:strides + dimensionOrder:dimensionOrder + dataType:dataType + shapeDynamism:ExecuTorchShapeDynamismDynamicBound]; +} + +- (instancetype)initWithBytesNoCopy:(void *)pointer + shape:(NSArray *)shape + dataType:(ExecuTorchDataType)dataType + shapeDynamism:(ExecuTorchShapeDynamism)shapeDynamism { + return [self initWithBytesNoCopy:pointer + shape:shape + strides:@[] + dimensionOrder:@[] + dataType:dataType + shapeDynamism:shapeDynamism]; +} + +- (instancetype)initWithBytesNoCopy:(void *)pointer + shape:(NSArray *)shape + dataType:(ExecuTorchDataType)dataType { + return [self initWithBytesNoCopy: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 f5c2ccdbeba..cc00d9fddfa 100644 --- a/extension/apple/ExecuTorch/__tests__/TensorTest.swift +++ b/extension/apple/ExecuTorch/__tests__/TensorTest.swift @@ -11,6 +11,16 @@ import XCTest class TensorTest: XCTestCase { - func test() { + func testInitBytesNoCopy() { + var data: [Float] = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0] + let tensor = data.withUnsafeMutableBytes { + Tensor(bytesNoCopy: $0.baseAddress!, shape: [2, 3], dataType: .float) + } + XCTAssertEqual(tensor.dataType, .float) + XCTAssertEqual(tensor.shape, [2, 3]) + XCTAssertEqual(tensor.strides, [3, 1]) + XCTAssertEqual(tensor.dimensionOrder, [0, 1]) + XCTAssertEqual(tensor.shapeDynamism, .dynamicBound) + XCTAssertEqual(tensor.count, 6) } }