diff --git a/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.h b/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.h index 984f5780bce..16b1ec392ec 100644 --- a/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.h +++ b/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.h @@ -61,6 +61,28 @@ typedef NS_ENUM(uint8_t, ExecuTorchShapeDynamism) { ExecuTorchShapeDynamismDynamicUnbound, } NS_SWIFT_NAME(ShapeDynamism); +/** + * Returns the size in bytes of the specified data type. + * + * @param dataType An ExecuTorchDataType value representing the tensor's element type. + * @return An NSInteger indicating the size in bytes. + */ +FOUNDATION_EXPORT +__attribute__((deprecated("This API is experimental."))) +NSInteger ExecuTorchSizeOfDataType(ExecuTorchDataType dataType) + NS_SWIFT_NAME(size(ofDataType:)); + +/** + * Computes the total number of elements in a tensor based on its shape. + * + * @param shape An NSArray of NSNumber objects, where each element represents a dimension size. + * @return An NSInteger equal to the product of the sizes of all dimensions. + */ +FOUNDATION_EXPORT +__attribute__((deprecated("This API is experimental."))) +NSInteger ExecuTorchElementCountOfShape(NSArray *shape) + NS_SWIFT_NAME(elementCount(ofShape:)); + /** * A tensor class for ExecuTorch operations. * diff --git a/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.mm b/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.mm index b198f073601..19851172b12 100644 --- a/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.mm +++ b/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.mm @@ -16,6 +16,18 @@ using namespace executorch::aten; using namespace executorch::extension; +NSInteger ExecuTorchSizeOfDataType(ExecuTorchDataType dataType) { + return elementSize(static_cast(dataType)); +} + +NSInteger ExecuTorchElementCountOfShape(NSArray *shape) { + NSInteger count = 1; + for (NSNumber *dimension in shape) { + count *= dimension.integerValue; + } + return count; +} + @implementation ExecuTorchTensor { TensorPtr _tensor; NSArray *_shape; diff --git a/extension/apple/ExecuTorch/__tests__/TensorTest.swift b/extension/apple/ExecuTorch/__tests__/TensorTest.swift index cc00d9fddfa..850ee10d007 100644 --- a/extension/apple/ExecuTorch/__tests__/TensorTest.swift +++ b/extension/apple/ExecuTorch/__tests__/TensorTest.swift @@ -11,6 +11,50 @@ import XCTest class TensorTest: XCTestCase { + func testElementCountOfShape() { + XCTAssertEqual(elementCount(ofShape: [2, 3, 4]), 24) + XCTAssertEqual(elementCount(ofShape: [5]), 5) + XCTAssertEqual(elementCount(ofShape: []), 1) + } + + func testSizeOfDataType() { + let expectedSizes: [DataType: Int] = [ + .byte: 1, + .char: 1, + .short: 2, + .int: 4, + .long: 8, + .half: 2, + .float: 4, + .double: 8, + .complexHalf: 4, + .complexFloat: 8, + .complexDouble: 16, + .bool: 1, + .qInt8: 1, + .quInt8: 1, + .qInt32: 4, + .bFloat16: 2, + .quInt4x2: 1, + .quInt2x4: 1, + .bits1x8: 1, + .bits2x4: 1, + .bits4x2: 1, + .bits8: 1, + .bits16: 2, + .float8_e5m2: 1, + .float8_e4m3fn: 1, + .float8_e5m2fnuz: 1, + .float8_e4m3fnuz: 1, + .uInt16: 2, + .uInt32: 4, + .uInt64: 8, + ] + for (dataType, expectedSize) in expectedSizes { + XCTAssertEqual(size(ofDataType: dataType), expectedSize, "Size for \(dataType) should be \(expectedSize)") + } + } + func testInitBytesNoCopy() { var data: [Float] = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0] let tensor = data.withUnsafeMutableBytes {