From 641e5af05a0c72f5371abd1ca785f1668457007b Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <shoumikhin@meta.com>
Date: Wed, 26 Mar 2025 17:18:53 -0700
Subject: [PATCH] Tensor constructor to create with a raw pointer w/o copying.

Summary: https://github.com/pytorch/executorch/issues/8366

Reviewed By: bsoyluoglu

Differential Revision: D71902542
---
 .../ExecuTorch/Exported/ExecuTorchTensor.h    | 24 +++++++++++++++++++
 .../ExecuTorch/Exported/ExecuTorchTensor.mm   | 23 ++++++++++++++++++
 2 files changed, 47 insertions(+)

diff --git a/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.h b/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.h
index 3d54ac5977b..1c7e0ab8c8d 100644
--- a/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.h
+++ b/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.h
@@ -134,4 +134,28 @@ __attribute__((deprecated("This API is experimental.")))
 
 @end
 
+#pragma mark - BytesNoCopy Category
+
+@interface ExecuTorchTensor (BytesNoCopy)
+
+/**
+ * Initializes a tensor without copying the provided data.
+ *
+ * @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.
+ * @param shapeDynamism An ExecuTorchShapeDynamism value indicating whether the shape is static or dynamic.
+ * @return An initialized ExecuTorchTensor instance using the provided data buffer.
+ */
+- (instancetype)initWithBytesNoCopy:(void *)pointer
+                              shape:(NSArray<NSNumber *> *)shape
+                            strides:(NSArray<NSNumber *> *)strides
+                     dimensionOrder:(NSArray<NSNumber *> *)dimensionOrder
+                           dataType:(ExecuTorchDataType)dataType
+                      shapeDynamism:(ExecuTorchShapeDynamism)shapeDynamism;
+
+@end
+
 NS_ASSUME_NONNULL_END
diff --git a/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.mm b/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.mm
index ba4fe331cad..8a0e7aea3dd 100644
--- a/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.mm
+++ b/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.mm
@@ -13,6 +13,7 @@
 
 #import <executorch/extension/tensor/tensor.h>
 
+using namespace executorch::aten;
 using namespace executorch::extension;
 
 @implementation ExecuTorchTensor {
@@ -69,3 +70,25 @@ - (NSInteger)count {
 }
 
 @end
+
+@implementation ExecuTorchTensor (BytesNoCopy)
+
+- (instancetype)initWithBytesNoCopy:(void *)pointer
+                              shape:(NSArray<NSNumber *> *)shape
+                            strides:(NSArray<NSNumber *> *)strides
+                     dimensionOrder:(NSArray<NSNumber *> *)dimensionOrder
+                           dataType:(ExecuTorchDataType)dataType
+                      shapeDynamism:(ExecuTorchShapeDynamism)shapeDynamism {
+  ET_CHECK(pointer);
+  auto tensor = make_tensor_ptr(
+    utils::toVector<SizesType>(shape),
+    pointer,
+    utils::toVector<DimOrderType>(dimensionOrder),
+    utils::toVector<StridesType>(strides),
+    static_cast<ScalarType>(dataType),
+    static_cast<TensorShapeDynamism>(shapeDynamism)
+  );
+  return [self initWithNativeInstance:&tensor];
+}
+
+@end