Merge branch 'fix/metal-offset-views' into fix/metal-nonpacked-views

# Conflicts:
#	backends/apple/metal/runtime/shims/memory.cpp
#	backends/apple/metal/runtime/test/test_memory.cpp

Author: abdelaziz-mahdy

backends/apple/metal/runtime/shims/memory.cpp

@@ -63,6 +63,12 @@ std::unordered_map<void*, int32_t> not_own_handles;
 // one only drops the handle.
 std::unordered_set<Tensor*> retired_handles;
 
+// The address an empty view stands for, where that is not the address it is
+// bound at (aoti_torch__reinterpret_tensor binds empty views at their
+// parent's start). Views of it are offset from here, so they are checked
+// against the memory they really reach.
+std::unordered_map<Tensor*, void*> empty_view_address;
+
 // Size of each CPU allocation the runtime owns, so that views of it can be
 // bound into one Metal buffer over all of it (metal_register_cpu_view).
 std::unordered_map<void*, size_t> cpu_allocation_bytes;
@@ -762,6 +768,7 @@ AOTITorchError aoti_torch_delete_tensor_object(AOTITensorHandle tensor) {
 
   if (retired_handles.erase(tensor) != 0) {
     // Its memory is gone, and its address may be another allocation's now.
+    empty_view_address.erase(tensor);
     tensors.erase(it);
     return Error::Ok;
   }
@@ -770,6 +777,7 @@ AOTITorchError aoti_torch_delete_tensor_object(AOTITensorHandle tensor) {
   void* data_ptr = tensor_ptr->mutable_data_ptr();
   if (data_ptr == nullptr) {
     // A tensor with no memory (make_memoryless_tensor).
+    empty_view_address.erase(tensor);
     tensors.erase(it);
     return Error::Ok;
   }
@@ -801,6 +809,7 @@ AOTITorchError aoti_torch_delete_tensor_object(AOTITensorHandle tensor) {
       memory_to_n_tensor.erase(data_ptr);
     }
     metal_forget_strided_view(tensor);
+    empty_view_address.erase(tensor);
     tensors.erase(it);
     ET_LOG(
         Debug,
@@ -810,6 +819,7 @@ AOTITorchError aoti_torch_delete_tensor_object(AOTITensorHandle tensor) {
   ET_CHECK_OK_OR_RETURN_ERROR(release_memory(data_ptr));
 
   metal_forget_strided_view(tensor);
+  empty_view_address.erase(tensor);
   tensors.erase(it);
   ET_LOG(Debug, "aoti_torch_delete_tensor_object: successful");
   return Error::Ok;
@@ -1127,11 +1137,17 @@ AOTITorchError aoti_torch__reinterpret_tensor(
       InvalidArgument,
       "aoti_torch__reinterpret_tensor: sizes and strides must fit a tensor");
 
-  // Handle storage offset by adjusting the data pointer
+  // Handle storage offset by adjusting the data pointer. An empty `self` is
+  // offset from the address it stands for, not the one it is bound at.
+  void* origin = data_ptr;
+  auto empty_origin = empty_view_address.find(self);
+  if (empty_origin != empty_view_address.end()) {
+    origin = empty_origin->second;
+  }
   size_t element_size = dtype_to_element_size(dtype);
   void* adjusted_data = nullptr;
   ET_CHECK_OR_RETURN_ERROR(
-      offset_pointer(data_ptr, storage_offset, element_size, &adjusted_data),
+      offset_pointer(origin, storage_offset, element_size, &adjusted_data),
       InvalidArgument,
       "aoti_torch__reinterpret_tensor: storage_offset overflows");
 
@@ -1145,16 +1161,19 @@ AOTITorchError aoti_torch__reinterpret_tensor(
       "aoti_torch__reinterpret_tensor: sizes must not be negative");
 
   // An empty view reads nothing. Its offset can put it at the very end of its
-  // buffer, which may be where the next allocation starts; point it at the
-  // start of its parent instead, so that it is not taken for a view there.
+  // buffer, which may be where the next allocation starts; bind it at the
+  // start of its parent instead, so that it is not taken for a view there,
+  // and remember the address it stands for (empty_view_address).
   int64_t view_numel = 1;
   for (auto size : sizes) {
     ET_CHECK_OR_RETURN_ERROR(
         !__builtin_mul_overflow(view_numel, int64_t{size}, &view_numel),
         InvalidArgument,
         "aoti_torch__reinterpret_tensor: the number of elements overflows");
   }
-  if (view_numel == 0) {
+  const bool empty_view = view_numel == 0;
+  void* empty_address = adjusted_data;
+  if (empty_view) {
     adjusted_data = data_ptr;
   }
 
@@ -1366,6 +1385,9 @@ AOTITorchError aoti_torch__reinterpret_tensor(
     ET_DCHECK_MSG(recorded, "a packable strided view was not recorded");
     (void)recorded;
   }
+  if (empty_view && empty_address != adjusted_data) {
+    empty_view_address[tensor.get()] = empty_address;
+  }
   if (registered) {
     view_handles.insert(tensor.get());
   }
@@ -1463,6 +1485,11 @@ AOTITorchError aoti_torch_new_tensor_handle(
   *new_handle = tensor.get();
 
   metal_share_strided_view(orig_handle, tensor.get());
+  // A copy of an empty view stands for the same address.
+  auto empty_origin = empty_view_address.find(orig_handle);
+  if (empty_origin != empty_view_address.end()) {
+    empty_view_address[tensor.get()] = empty_origin->second;
+  }
 
   // If the original is a view into a Metal buffer, the new handle is one too,
   // and deleting either must leave the view registered for the other one.
@@ -1498,6 +1525,7 @@ void retire_handles_within(void* ptr, size_t nbytes) {
       // regions do not overlap Metal buffers.
       view_owner.erase(entry.first);
       metal_forget_strided_view(entry.first);
+      empty_view_address.erase(entry.first);
       retired_handles.insert(entry.first);
     }
   }
@@ -1545,6 +1573,7 @@ void cleanup_memory() {
   view_owner.clear();
   view_handles.clear();
   retired_handles.clear();
+  empty_view_address.clear();
   cpu_allocation_bytes.clear();
 
   // Clean up Metal resources

backends/apple/metal/runtime/test/test_memory.cpp

@@ -1991,6 +1991,46 @@ TEST_F(MetalGraphViewTest, FreeingABufferForgetsTheStridedViewsInIt) {
   ASSERT_EQ(aoti_torch_delete_tensor_object(weight), Error::Ok);
 }
 
+// An empty view is bound at its parent's start but stands for its own offset:
+// views of it are offset from there, so one past the end of the allocation is
+// refused, and one inside it lands where it belongs, also through a copy.
+TEST_F(MetalGraphViewTest, ViewsOfAnEmptyViewKeepItsOffset) {
+  const int64_t base_size = 8;
+  AOTITensorHandle base = nullptr;
+  ASSERT_EQ(
+      aoti_torch_empty_strided(
+          1, &base_size, &kStride, kFloat32, kDeviceMps, 0, &base),
+      Error::Ok);
+  auto* memory = static_cast<float*>(base->mutable_data_ptr());
+  const int64_t zero = 0;
+  const int64_t four = 4;
+
+  AOTITensorHandle at_end = nullptr;
+  ASSERT_EQ(
+      aoti_torch__reinterpret_tensor(
+          base, 1, &zero, &kStride, /*storage_offset=*/8, &at_end),
+      Error::Ok);
+  AOTITensorHandle past_end = nullptr;
+  EXPECT_NE(
+      aoti_torch__reinterpret_tensor(
+          at_end, 1, &four, &kStride, /*storage_offset=*/0, &past_end),
+      Error::Ok);
+
+  AOTITensorHandle middle = nullptr;
+  ASSERT_EQ(
+      aoti_torch__reinterpret_tensor(
+          base, 1, &zero, &kStride, /*storage_offset=*/4, &middle),
+      Error::Ok);
+  AOTITensorHandle copy = nullptr;
+  ASSERT_EQ(aoti_torch_new_tensor_handle(middle, &copy), Error::Ok);
+  AOTITensorHandle back_half = nullptr;
+  ASSERT_EQ(
+      aoti_torch__reinterpret_tensor(
+          copy, 1, &four, &kStride, /*storage_offset=*/0, &back_half),
+      Error::Ok);
+  EXPECT_EQ(back_half->const_data_ptr(), memory + 4);
+}
+
 // Memory a tensor owns goes with the tensor, not with aoti_torch_mps_free.
 TEST_F(MetalGraphViewTest, FreeingATensorsMemoryIsRefused) {
   const int64_t size = 4;